[AMD] Remove ROCm page_first+kernel -> layer_first HiCache fallback (follow-up to #28534)

[AMD-yanfeiwang]

Motivation

#28534 enabled the JIT staged HiCache write-back path on ROCm:

• can_use_jit / can_use_write_back_jit are now allowed on HIP (not only CUDA) in the MHA/MLA host pools;
• hicache.cuh / staged_write_back.cuh build and run with hipcc (PTX-only helpers guarded by USE_ROCM, ROCm device-type matchers, vectorized non-temporal load/store).

As called out in #28534's description, once that landed, the ROCm layer_first fallback added by #28473 in ServerArgs._resolve_layout_io_compatibility() should be reverted.

What this PR does

Removes the ROCm-only block that rewrites page_first + kernel to layer_first:

# The page_first kernel write-back relies on the CUDA-only JIT staged
# kernel. On ROCm it falls back to a kernel that requires CUDA index
# tensors and crashes on host write-back, so use layer_first there.
if (
    self.hicache_mem_layout == "page_first"
    and self.hicache_io_backend == "kernel"
    and is_hip()
):
    self.hicache_mem_layout = "layer_first"
    logger.warning(...)

Why it must be removed (not just cleanup)

The block runs at ServerArgs resolution time, i.e. before the host KV pool is created. On ROCm it forces any page_first + kernel request to layer_first, so ROCm can never reach the page_first + kernel JIT staged write-back path that #28534 just made correct. The two are directly contradictory: #28534 says "ROCm can use the JIT staged write-back", while this fallback says "ROCm always downgrades to layer_first". Leaving it in makes #28534's ROCm work dead code.

After this change, ROCm exercises the same page_first + kernel write-back path as CUDA (platform parity, retained staged write-back bandwidth), which is what #28534 validated on MI355X (gfx950, ROCm 7.2).

Test plan

• is_hip import still used elsewhere in server_args.py (no dead import); file byte-compiles.
• HiCache JIT unit tests (test/registered/jit/test_hicache.py, test_hicache_page_first_write_back.py) pass on MI355X (ROCm 7.2) with the page_first + kernel JIT staged write-back path enabled (validated for [AMD] Enable JIT staged HiCache write-back and fix CPU-index crash #28534).
• CI: AMD jit-kernel-unit + HiCache serving suites.
• Optional: end-to-end --enable-hierarchical-cache --hicache-mem-layout page_first --hicache-io-backend kernel launch on ROCm no longer silently downgrades to layer_first.

---

CI States

Latest PR Test (Base): ❌ Run #29231811696
Latest PR Test (Extra): ❌ Run #29231811508

[gemini-code-assist]

Code Review

This pull request removes the ROCm-specific fallback logic in _resolve_layout_io_compatibility that reverted the memory layout from page_first to layer_first when using the kernel IO backend. I have no feedback to provide as there are no review comments.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[yctseng0211]

passed all related tests on AMD, LGTM

[HaiShaw]

/tag-and-rerun-ci

[HaiShaw]

AMD only change