[ROCm] Widen the HiCache JIT copy rounds and enable the K-only host pool

[jiejingzhangamd]

Depends on #35233. That PR makes registered host pointers usable from the
transfer kernels on ROCm; this one is the part that decides how much of the KV
cache can actually take that path, and how fast it moves once it does. An
earlier revision of this PR carried its own pointer fix (pin_memory allocation,
an unregister guard, and an io-backend downgrade for storage-owned memory);
those are dropped in favour of #35233, which solves the same problem more
generally and does not have to give up the kernel backend for storage-owned
pools.

What is left here

Copy rounds narrower than 128 B. The JIT transfer kernels tiled each element
in 128 B rounds only, so an element size 128 does not divide could not use them
at all. MLA's fp8 row is 576 B. pick_group_bytes() now takes the widest round
of 128/64/32/16 that both divides the element and splits across the lanes into a
package the hardware has (4, 8 or 16 B).

A block quota that suits ROCm. 2 on CUDA, 16 on ROCm. The per-layer read path
loses bandwidth below that, which is also why the narrow rounds are worth their
extra instantiations only on ROCm.

MHATokenToKOnlyPoolHost can use the JIT path. can_use_jit accepted CUDA
only; the kernels are not CUDA-specific.

CUDA is unchanged

Not "verified by inspection" — the narrow rounds are behind #ifdef USE_ROCM,
and the #else branch reduces to group_fits(bytes, num_threads, 128), whose
three conditions are the two original static_asserts plus the package-size
constraint PackageType already imposed. DEFAULT_BLOCK_QUOTA and GROUP_BYTES
are gated the same way in Python. Non-ROCm builds generate the code they did
before.

Tests

test/registered/unit/mem_cache/test_hicache_copy_rounds.py, CPU-only, in
base-a-test-cpu. kvcache/hicache.py screens element sizes with a Python copy
of pick_group_bytes(), and the two have to answer alike: a size the screen
wrongly admits fails a static_assert at JIT compile time, and one it wrongly
rejects silently loses the JIT path. The test sweeps element sizes against the
rule on both platforms and names 576 B as the boundary the narrow rounds exist
for. 5 cases, 963 subtests.

The lane count the screen derives from is a literal 32 mirroring
device::kWarpThreads; kvcacheio/hicache.cuh now static_asserts that
constant, so a wavefront-size change fails at the rule instead of silently
disagreeing with the screen.

Measurements

TODO: re-measure on top of #35233. The numbers in the previous revision of this
description were taken with the pointer fix that has since been dropped, so they
do not describe this change.

ruff (F401,F821,UP037), ruff-format and isort clean.

---

CI States

Latest PR Test (Base): ✅ Run #35396160706
Latest PR Test (Extra): ✅ Run #35396160406
Latest PR Test (AMD ROCm 10): ❌ Run #35396160667

[jiejingzhangamd]

Could a maintainer add the run-ci label? The red checks here are all pr-gate and the *-finish aggregators failing on Missing required label 'run-ci', so no test has actually executed yet.

What I ran locally in the meantime, on MI300X:

• test/registered/unit/server_args — 261 passed
• test/registered/unit/mem_cache — 520 passed
• AMD HiCache JIT transfer tests — 59 passed
• A GPU transfer matrix over the five host pool classes and both memory layouts — 23 passed, 3 unsupported, no faults and no data mismatches

The 46 new cases in this PR are all CPU-only and land in the base-a-test-cpu suite.

[HaiShaw]

/tag-and-rerun-ci

[HaiShaw]

cc @Lzy17 @amd-danli103

[Duyi-Wang]

@AMD-yanfeiwang Parts of this change overlap with #35233

[jiejingzhangamd]

/rerun-failed-ci

[jiejingzhangamd]

@AMD-yanfeiwang Parts of this change overlap with #35233

@Duyi-Wang You're right, and I've rescoped this. The version you commented on carried a pointer-domain fix of its own: allocating ROCm host pools with pin_memory (hipHostMalloc) instead of cudaHostRegister, plus a server-args downgrade of --hicache-io-backend kernel to direct when a storage backend owns the host memory. Both are gone as of the latest push. #37152 is now four commits — the narrow copy rounds in pick_group_bytes() so MLA's 576 B fp8 row can reach the read JIT, the ROCm DEFAULT_BLOCK_QUOTA, the MHATokenToKOnlyPoolHost JIT gate, and a CPU test. I'll update the description to match.

#35233 is the better fix for the part that overlapped. Translating to the device alias at the launch boundary is more general than changing how we allocate: it covers registered buffers as well as hipHostMalloc'd ones, so the storage-owned allocators can use the kernel backend instead of being downgraded, and it also covers HiSparse, Mamba and the DSV4/DSA tables, plus the CUDA/WSL2 case @seokwoosong reproduced. My own probe agrees with its premise — hipHostGetDevicePointer succeeds for register flags 0 through 3 and returns a valid device address that differs from the host VA. I read that as "don't hand registered memory to the kernel"; reading it as "use the alias" was the better call.

One caveat on that, which I've raised on #35233 rather than here: the alias's offset arithmetic is only valid within a single registration, and _cuda_host_register registers in multiple chunks for the page_first layouts once a buffer exceeds SGLANG_HICACHE_HOST_REGISTER_CHUNK_GB. layer_first is single-registration and safe, and it's what #35233's tests cover. So I'd say #35233 is the right mechanism and needs one guard before the storage-owned path can be relied on, not that the downgrade should come back.

On the two approaches performing about the same: for the default allocator, agreed, and I have nothing that separates them. For the storage-owned case they aren't equivalent, because the downgrade meant running direct rather than kernel. On MI355X with the scattered page indices a host pool hands out once it has evicted anything, direct measures 0.3–2.4 GB/s against 39–56 for the JIT; with contiguous arange indices it climbs back to ~50 and looks level, which is why a short benchmark hides it. End to end on GLM-5.2 at concurrency 16 the kernel backend led by about 25%. Caveat in the honest direction: those runs used the default allocator rather than a storage backend, so with L3 I/O in the loop the end-to-end share of that gain would be smaller. Either way it argues for #35233's approach over the downgrade, so it isn't a reason to keep it.

Two places the PRs still touch:

• DEFAULT_BLOCK_QUOTA in kvcache/hicache.py — this PR sets 16 on ROCm, [AMD] Fix registered HiCache host pointer aliases #35233 sets 32, same line. Raised on [AMD] Fix registered HiCache host pointer aliases #35233.
• The MHATokenToKOnlyPoolHost gate commit depends on the pointer-domain fix I just removed. Under the default page_first layout it hands k_data_ptrs to the JIT to dereference on the GPU, so on ROCm it needs [AMD] Fix registered HiCache host pointer aliases #35233's alias to be correct. Happy to mark this PR blocked on [AMD] Fix registered HiCache host pointer aliases #35233, or to drop that commit and re-land it afterwards — whichever you'd prefer.

[jiejingzhangamd]

@HaiShaw add accuracy reading and performance reading in 3 way, for baseline, baseline + #35233, and baseline + #35233 +#37152

baseline docker image:rocm/sgl-dev:v0.5.18-rocm724-mi35x-20260901
Setup: sglang.test.run_eval --eval-name gsm8k --num-examples 200 --num-threads 14 --num-shots 5; AgentX AIPERF_EXPERIMENTAL_FAST=1, conc 14 then 2, simulated acc 3.61, 20 min profile.

Accuracy. GSM8K 200-shot (5-shot, real MTP) is in-band across arms:

• baseline 0.970,
• +[AMD] Fix registered HiCache host pointer aliases #35233 0.960,
• +[AMD] Fix registered HiCache host pointer aliases #35233+[ROCm] Widen the HiCache JIT copy rounds and enable the K-only host pool #37152 0.965.
• Speculative accept length is also matched (mean 4.46 / 4.49 / 4.48).
• AgentX (conc=14 then 2) had 0 errors on all three arms.

Speed. TP4/EP4 GLM-5.2-MXFP4, 8×MI355X, HiCache on. Same recipe besides HiCache backend:

arm	HiCache	AgentX c14 out	c14 TTFT/ITL p50	AgentX c2 out
pre-#35233	direct / page_first_direct	291 tok/s	931 ms / 15.3 ms	129 tok/s
+#35233	kernel / page_first	338 tok/s (+16%)	742 ms / 13.3 ms	130 tok/s
+#35233+#37152	kernel / page_first	338 tok/s	759 ms / 13.5 ms	130 tok/s

c14 also: GPU cache hit ~80% on all three; per-GPU total ~11.5k (direct) vs ~12.3k (kernel).

Reading. #35233 (host-pointer alias → kernel HiCache) is the win vs stock direct at AgentX conc=14. #37152 (ROCm tiled copy rounds 128/64/32/16 for GLM MLA fp8 576 B rows, which previously silently fell back from the 128 B-only JIT tile) is accuracy-neutral and speed-neutral on top of #35233 for this workload — as expected if #35233 already keeps the kernel path engaged. Conc=2 is HiCache-light and matches (~130 tok/s).

[jiejingzhangamd]

Updated this: made the 32-thread copy group explicit (it is not a hardware warp), added compile-time checks, and added AMD CI coverage for the 576B path. Could you take another look?

[xiezhq-hermann]

@DarkSharpness can you help do a final pass on the kernel changes?