[AMD] Enable JIT staged HiCache write-back and fix CPU-index crash by AMD-yanfeiwang · Pull Request #28534 · sgl-project/sglang

AMD-yanfeiwang · 2026-06-17T15:01:12Z

On ROCm, page_first + kernel HiCache write-back crashes on the first prefill with RuntimeError: Destination indices must be a CUDA tensor. This PR fixes the underlying cause so ROCm runs the same page_first + kernel JIT staged write-back path as CUDA, instead of the #28473 layer_first fallback.

Root cause

HiCacheController.start_writing() keeps host_indices on the CPU for the kernel io-backend + page_first layout, assuming the staged JIT write-back kernel (which stages through device memory and accepts a CPU destination index) will consume them. That JIT path is gated behind _is_cuda, so on ROCm it is disabled and the code falls back to the plain transfer_kv_all_layer_mla_lf_pf C++ kernel, whose launcher asserts dst_indices.is_cuda(). CPU host_indices -> assert -> all TP scheduler ranks crash -> prefill dies.

Why fix the cause instead of the #28473 `layer_first` fallback

#28473 works around the same crash by forcing hicache_mem_layout = "layer_first" on ROCm whenever page_first + kernel is requested. That keeps CI green, but it:

disables the JIT staged write-back path on AMD entirely — ROCm permanently runs the older layer_first path and never benefits from the staged kernel that CUDA uses;
diverges ROCm from CUDA at the config layer, so the two platforms exercise different layouts/kernels and have to be maintained separately;
leaves the actual page_first + kernel ROCm path broken.

This PR fixes the underlying cause so ROCm keeps the same path as CUDA: platform parity, retained staged write-back bandwidth, vectorized non-temporal device transfers, and no second code path to maintain.

Once this lands, #28473 should be reverted — its ROCm layer_first fallback in ServerArgs._resolve_layout_io_compatibility() is no longer needed (and would otherwise keep short-circuiting ROCm away from the now-correct page_first + kernel JIT path).

Modifications

hicache.cuh: guard the NVIDIA-only PTX ld/st.global.L1::no_allocate helpers (load_nc/store_nc) behind #ifndef USE_ROCM and provide ROCm equivalents using non-temporal loads/stores, so the JIT HiCache module also builds with hipcc. The ROCm paths use a single __builtin_nontemporal_{load,store} over Clang ext_vector_type(2/4) (with __builtin_bit_cast) so the vectorized global_{load,store}_dwordx{2,4} non-temporal ops are deterministic instead of relying on the LoadStoreVectorizer.
memory_pool_host.py: allow can_use_jit on HIP, not only CUDA, so ROCm uses the same staged write-back path as CUDA.
cache_controller.py: only keep host_indices on CPU when the staged JIT kernel is actually available (can_use_jit); otherwise move them to the device. This makes the kernel io-backend correct on any backend where the JIT kernel is unavailable.
staged_write_back.cuh / hicache.cuh device matchers: accept kDLROCM (device) / kDLROCMHost (pinned host) alongside kDLCUDA/kDLCUDAHost, mirroring the other JIT kernels (clamp_position, kvcache, resolve_future_token_ids).

Reproduction

Note: main currently contains #28473's ROCm page_first + kernel -> layer_first fallback in ServerArgs._resolve_layout_io_compatibility(). Revert it first; otherwise ROCm silently switches to layer_first and this path is never exercised.

2P1D ROCm deployment with page_first + kernel HiCache write-back (--hicache-io-backend defaults to kernel, the JIT staged write-back path this PR fixes). Core launch commands (TP=8, Kimi-K2.6-MXFP4):

Prefill

SGLANG_USE_AITER=1 SGLANG_AITER_MLA_PERSIST=1 AITER_MXFP4_MOE_SF=1 \
python3 -m sglang.launch_server \
  --model-path /models/amd/Kimi-K2.6-MXFP4 \
  --served-model-name Kimi-K2.6-MXFP4 \
  --tool-call-parser kimi_k2 --reasoning-parser kimi_k2 \
  --chat-template /models/amd/Kimi-K2.6-MXFP4/chat_template.jinja \
  --tp-size 8 --page-size 64 \
  --context-length 262144 --kv-cache-dtype bf16 \
  --attention-backend aiter \
  --mem-fraction-static 0.8 --max-running-requests 128 \
  --chunked-prefill-size 16384 \
  --cuda-graph-bs $(seq 1 128) --cuda-graph-max-bs 128 \
  --trust-remote-code \
  --disaggregation-mode prefill \
  --disaggregation-transfer-backend mori \
  --disaggregation-bootstrap-port 8998 \
  --disaggregation-ib-device ionic_0,ionic_1,ionic_2,ionic_3,ionic_4,ionic_5,ionic_6,ionic_7 \
  --enable-hierarchical-cache --hicache-size 192 \
  --hicache-mem-layout page_first --hicache-write-policy write_through \
  --enable-metrics --enable-cache-report \
  --host 0.0.0.0 --port 30020

Decode

SGLANG_USE_AITER=1 SGLANG_AITER_MLA_PERSIST=1 AITER_MXFP4_MOE_SF=1 \
python3 -m sglang.launch_server \
  --model-path /models/amd/Kimi-K2.6-MXFP4 \
  --served-model-name Kimi-K2.6-MXFP4 \
  --tool-call-parser kimi_k2 --reasoning-parser kimi_k2 \
  --chat-template /models/amd/Kimi-K2.6-MXFP4/chat_template.jinja \
  --tp-size 8 --page-size 64 \
  --context-length 262144 --kv-cache-dtype bf16 \
  --attention-backend aiter \
  --mem-fraction-static 0.85 --max-running-requests 128 \
  --chunked-prefill-size 8192 \
  --cuda-graph-bs $(seq 1 128) --cuda-graph-max-bs 128 \
  --num-continuous-decode-steps 4 \
  --trust-remote-code \
  --disaggregation-mode decode \
  --disaggregation-transfer-backend mori \
  --disaggregation-bootstrap-port 19100 \
  --disaggregation-ib-device ionic_0,ionic_1,ionic_2,ionic_3,ionic_4,ionic_5,ionic_6,ionic_7 \
  --enable-metrics --enable-cache-report \
  --host 0.0.0.0 --port 30030

Router (native PD-disaggregation router)

python3 -m sglang_router.launch_router \
  --pd-disaggregation \
  --prefill http://<PREFILL_IP>:30020 8998 \
  --decode  http://<DECODE_IP>:30030 \
  --prefill-policy cache_aware --decode-policy round_robin \
  --disable-circuit-breaker \
  --host 0.0.0.0 --port 8100

Then send multi-turn / prefix-reusing traffic at the router (:8100). Before this PR, on ROCm the first prefill HiCache write-back crashes with RuntimeError: Destination indices must be a CUDA tensor, and the first prefix-cache hit crashes with Tensor match failed ... device=rocm:N. With this PR both paths run cleanly.

Accuracy / functional tests

JIT HiCache module compiles and loads with hipcc on gfx942 and gfx950 (ROCm 7.2).
ISA check on gfx950: the ROCm load_nc/store_nc emit single global_load_dwordx2/x4 and global_store_dwordx2/x4 with the non-temporal (nt) flag.
End-to-end on a 2P1D Kimi-K2.6-MXFP4 disaggregated deployment (TP=8, HiCache page_first + kernel, write_through) running the AgentX v0.3 agentic trace replay at concurrency 64: warmup + 900s profiling complete (328 requests, 0 errors) with no Destination indices must be a CUDA tensor and no Tensor match failed ... device=rocm:N crashes (both crashed pre-fix on the first write-back / first prefix-cache-hit load).

Speed tests and profiling

To check that the staged write-back path is worth keeping on ROCm (rather than degrading to layer_first per #28473), I ran the equivalent of the #21631 write-back micro-benchmark on AMD MI355X / gfx950, ROCm 7.2, comparing the installed sgl_kernel LF->PF kernels against the JIT staged LF->PF kernels on the same host-destination write-back workload.

Setup mirrors #21631: dtype=bf16, page_size=64, batch_pages=64, total_pages=128, timing via triton.testing.do_bench(warmup=5, rep=25). Each kernel is fed indices in the residency it requires (baseline *_lf_pf: device dst_indices; staged: pinned-host dst_indices); correctness is verified with torch.testing.assert_close for every row before timing. MHA bandwidth counts K + V; MLA counts a single buffer. speedup = staged_jit / sgl_kernel_lf_pf. ROCm has no cudaMemcpyBatchAsync equivalent, so staged_write_back.cuh uses the non-batch ROCm fallback (device relayout into staging + per-page async H2D copies).

Per-shape write-back microbenchmark (MHA / MLA)

MHA (`batch_pages=64`)

num_layers	element_dim	item_bytes	`sgl_kernel *_lf_pf` GiB/s	`jit *_staged_lf_pf` GiB/s	speedup
16	256	512	21.12	25.59	1.21x
16	512	1024	15.43	34.15	2.21x
16	1024	2048	21.82	39.42	1.81x
16	2048	4096	24.15	42.01	1.74x
24	256	512	15.73	30.63	1.95x
24	512	1024	18.67	36.83	1.97x
24	1024	2048	18.45	42.05	2.28x
24	2048	4096	23.75	47.85	2.02x
32	256	512	18.14	34.14	1.88x
32	512	1024	20.76	41.45	2.00x
32	1024	2048	24.36	45.08	1.85x
32	2048	4096	23.24	48.85	2.10x
40	256	512	17.68	36.90	2.09x
40	512	1024	18.92	43.06	2.28x
40	1024	2048	24.32	47.20	1.94x
40	2048	4096	25.28	49.67	1.96x
48	256	512	20.20	38.79	1.92x
48	512	1024	22.60	44.16	1.95x
48	1024	2048	20.72	47.99	2.32x
48	2048	4096	23.36	49.98	2.14x
56	256	512	19.73	40.25	2.04x
56	512	1024	21.87	45.11	2.06x
56	1024	2048	24.10	48.48	2.01x
56	2048	4096	25.97	50.39	1.94x
64	256	512	16.90	41.41	2.45x
64	512	1024	22.47	45.55	2.03x
64	1024	2048	17.16	48.87	2.85x
64	2048	4096	23.00	50.46	2.19x
72	256	512	16.50	42.33	2.56x
72	512	1024	23.52	46.78	1.99x
72	1024	2048	24.36	49.49	2.03x
72	2048	4096	25.00	50.78	2.03x
80	256	512	20.55	43.25	2.10x
80	512	1024	17.14	47.45	2.77x
80	1024	2048	22.13	49.73	2.25x
80	2048	4096	23.53	48.88	2.08x

MLA (`batch_pages=64`)

num_layers	element_dim	item_bytes	`sgl_kernel *_lf_pf` GiB/s	`jit *_staged_lf_pf` GiB/s	speedup
16	256	512	24.03	25.68	1.07x
16	512	1024	26.65	34.33	1.29x
16	1024	2048	26.32	41.56	1.58x
16	2048	4096	27.37	45.65	1.67x
24	256	512	24.34	30.86	1.27x
24	512	1024	25.63	38.98	1.52x
24	1024	2048	27.63	44.16	1.60x
24	2048	4096	29.27	48.12	1.64x
32	256	512	24.95	34.31	1.38x
32	512	1024	25.86	41.59	1.61x
32	1024	2048	27.12	45.47	1.68x
32	2048	4096	30.48	48.90	1.60x
40	256	512	24.60	37.03	1.51x
40	512	1024	26.59	43.11	1.62x
40	1024	2048	28.56	47.08	1.65x
40	2048	4096	28.22	49.57	1.76x
48	256	512	23.51	38.96	1.66x
48	512	1024	25.18	44.20	1.76x
48	1024	2048	29.75	47.89	1.61x
48	2048	4096	28.92	49.90	1.73x
56	256	512	24.16	40.42	1.67x
56	512	1024	26.59	44.80	1.68x
56	1024	2048	28.29	48.44	1.71x
56	2048	4096	28.57	50.31	1.76x
64	256	512	24.46	41.56	1.70x
64	512	1024	25.73	45.25	1.76x
64	1024	2048	29.66	48.71	1.64x
64	2048	4096	29.34	50.45	1.72x
72	256	512	24.15	42.29	1.75x
72	512	1024	26.93	46.53	1.73x
72	1024	2048	27.08	49.21	1.82x
72	2048	4096	27.22	50.69	1.86x
80	256	512	23.32	43.10	1.85x
80	512	1024	27.90	47.00	1.68x
80	1024	2048	27.94	49.71	1.78x
80	2048	4096	29.98	46.94	1.57x

Takeaway: on gfx950 the *_lf_pf kernels top out around 15-30 GiB/s, while the staged JIT path reaches ~25-50 GiB/s (close to the same peak CUDA achieves). That is a 1.21x-2.85x speedup for MHA and 1.07x-1.86x for MLA. The gap is larger than on CUDA (where #21631 reported roughly parity, ~0.87-1.25x, because the CUDA *_lf_pf baseline already runs near peak), so keeping ROCm on the staged page_first + kernel path is a clear win over the layer_first fallback.

CI States

Latest PR Test (Base): ⏳ Run #28997089073
Latest PR Test (Extra): ⏳ Run #28997088947

gemini-code-assist

Code Review

This pull request introduces ROCm/HIP support for the JIT HiCache kernels by guarding NVIDIA-specific PTX instructions and utilizing non-temporal builtins on ROCm. It also updates the cache controller and host memory pool to enable and check JIT kernel availability on HIP platforms, preventing crashes during write-back. The review feedback highlights critical performance and correctness improvements: first, using native Clang vector types (ext_vector_type) for uint2 and uint4 loads and stores on ROCm to ensure proper compiler vectorization; second, handling cases where self.mem_pool_host is a HostPoolGroup by checking its underlying anchor pool for JIT support so that the optimization is not silently disabled.

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.

gemini-code-assist · 2026-06-17T15:02:47Z

+        if (
+            self.io_backend == "kernel"
+            and self.mem_pool_host.layout == "page_first"
+            and getattr(self.mem_pool_host, "can_use_jit", False)
+        ):


When self.mem_pool_host is a HostPoolGroup (which is used for grouped pools like DSA/MLA models), it does not directly expose the can_use_jit attribute. As a result, getattr(self.mem_pool_host, "can_use_jit", False) will return False, silently disabling the JIT staged write-back optimization for these models.

We should also check the can_use_jit attribute of the underlying anchor pool if self.mem_pool_host is a HostPoolGroup.

Suggested change

if (

self.io_backend == "kernel"

and self.mem_pool_host.layout == "page_first"

and getattr(self.mem_pool_host, "can_use_jit", False)

):

if (

self.io_backend == "kernel"

and self.mem_pool_host.layout == "page_first"

and (

getattr(self.mem_pool_host, "can_use_jit", False)

or (

hasattr(self.mem_pool_host, "anchor_entry")

and getattr(self.mem_pool_host.anchor_entry.host_pool, "can_use_jit", False)

)

)

):

gemini-code-assist · 2026-06-18T03:25:41Z

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

On ROCm, page_first + kernel HiCache write-back crashed on the first prefill with: RuntimeError: Destination indices must be a CUDA tensor Root cause: HiCacheController.start_writing() keeps host_indices on the CPU for the kernel io-backend + page_first layout, assuming the staged JIT write-back kernel (which stages through device memory and accepts a CPU destination index) will consume them. That JIT path is gated behind `_is_cuda`, so on ROCm it is disabled and the code falls back to the plain `transfer_kv_all_layer_mla_lf_pf` C++ kernel, whose launcher asserts `dst_indices.is_cuda()`. CPU host_indices -> assert -> all TP scheduler ranks crash -> prefill dies. Fixes: - hicache.cuh: guard the NVIDIA-only PTX `ld/st.global.L1::no_allocate` helpers (load_nc/store_nc) behind `#ifndef USE_ROCM` and provide ROCm equivalents using non-temporal loads/stores, so the JIT HiCache module also builds with hipcc. The staged write-back kernel already has a USE_ROCM path. Verified the module compiles and loads on gfx942/ROCm 7.2. - memory_pool_host.py: allow `can_use_jit` on HIP, not only CUDA, so ROCm uses the same staged write-back path as CUDA. - cache_controller.py: only keep host_indices on CPU when the staged JIT kernel is actually available (`can_use_jit`); otherwise move them to the device as before. This makes the kernel io-backend correct on any backend where the JIT kernel is unavailable, independent of the change above.

The staged write-back kernel's TensorMatcher checks hard-coded kDLCUDA / kDLCUDAHost, so on ROCm the device-resident tensors (staging, layer ptrs, page indices) and pinned host buffers fail verification with "Tensor match failed ... device=rocm:N" at staged_write_back.cuh. Accept kDLROCM for device tensors and kDLROCMHost for host tensors, mirroring the kDLCUDA/kDLROCM pattern already used by the other JIT kernels (clamp_position, kvcache, resolve_future_token_ids).

…hers The non-staged HiCache JIT kernels (load host->device, write store) in hicache.cuh have the same CUDA-only TensorMatcher device checks as the staged kernel. These are exercised on prefix-cache hits (load path), so warmup (write-only) passed but profiling with resumed sessions crashed prefill with "Tensor match failed ... device=rocm:N at hicache.cuh". Add kDLROCM (device) / kDLROCMHost (pinned host) to all device matchers, same as the staged_write_back.cuh fix.

Express the ROCm load_nc/store_nc paths as a single __builtin_nontemporal_ {load,store} over Clang ext_vector_type(2/4) instead of N independent 32-bit ops. This makes the vectorized global_{load,store}_dwordx{2,4} with the nontemporal hint deterministic rather than relying on the LoadStoreVectorizer to merge per-scalar accesses. Use __builtin_bit_cast to convert between uintN and the native vector type to avoid strict-aliasing UB.

HaiShaw · 2026-07-05T20:39:36Z

/tag-and-rerun-ci

HaiShaw · 2026-07-06T05:26:32Z

@amd-bot ci-status

amd-bot · 2026-07-06T05:32:00Z

@HaiShaw

CI Status for PR #28534

Merge verdict: Likely safe to merge. All 14 executed CI failures are pre-existing perf-threshold flakes, infra library-load errors, or unrelated tests — none are caused by this PR's changes. PR CI is complete (185/185 checks finished, nothing pending/fast-fail-skipped). The changed dispatch logic + CPU-index fix is exercised and green on the NVIDIA CPU suite; however, the AMD-hardware JIT-kernel enablement (the PR's headline feature) is not directly exercised by any AMD PR-CI test.

Warning

Coverage is partial. The changed Python (cache_controller.py, memory_pool_host.py) — including the CPU-index crash fix — is covered by test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py (suite base-a-test-cpu, ✅ green), and the shared CUDA kernels by test/registered/jit/test_hicache.py (suite base-b-kernel-unit, NVIDIA, ✅ green). But test_hicache.py is registered CUDA-only (register_cuda_ci), so the ROCm/HIP compilation of the modified hicache.cuh / staged_write_back.cuh and the AMD write-back enablement are not run on AMD hardware in PR CI. Green does not verify the AMD JIT-kernel path. Author should run the HiCache JIT suite on an AMD (MI3xx) box before merge.

Changed files: python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh (+52/-10), .../staged_write_back.cuh (+8/-8), python/sglang/srt/managers/cache_controller.py (+8/-1), python/sglang/srt/mem_cache/memory_pool_host.py (+28/-4)

Executed CI failure attribution: AMD: 7 failures (0 related) · Others (NVIDIA/NPU/XPU): 7 failures (0 related) · 6 aggregator *-finish jobs collapsed into their upstream causes

AMD Executed Failures

Job	Test File	Test Function	Error	Related?	Why
stage-b-1gpu-large (0)	`test/registered/perf/test_bench_serving_1gpu_part1.py`	`test_online_lora_latency`	`AssertionError: median_e2e_latency_ms >= 2400`	🟢	Perf-threshold flake on LoRA serving; HiCache not enabled in this bench
stage-b-2gpu-large (0)	`test/registered/perf/test_bench_serving_2gpu.py`	`test_pp_offline_throughput_default_decode`	`AssertionError: 5148.4 not greater than 6700`	🟢	PP throughput perf-threshold flake; unrelated to write-back
stage-b-1gpu-small (7)	`test/registered/perf/test_vlm_perf_5090.py`	`test_vlm_online_latency`	`AssertionError: 156.36 not less than 150`	🟢	VLM TTFT perf-threshold flake (6ms over)
stage-b-8gpu-mi35x-disagg	`test/registered/amd/disaggregation/test_disaggregation_pp.py`	N/A (server startup)	`ConnectionRefusedError: port 11200`	🟢	Disagg server failed to come up; unrelated code path
extra-a-2gpu-large-amd	`test/registered/kv_canary/test_self_e2e_pp_perturb.py`	N/A	`AssertionError: -1 != 200 (ReadTimeout port 11000)`	🟢	PP perturbation e2e read-timeout; unrelated
stage-c-8gpu-mi35x (0)	N/A (runner setup)	N/A	`FAILED: VRAM cleanup unsuccessful after 3 attempts`	🟢	Infra: pre-test VRAM cleanup, before any test ran
stage-b-1gpu-small-mi35x	N/A (runner setup)	N/A	`FAILED: VRAM cleanup unsuccessful after 3 attempts`	🟢	Infra: same runner cleanup failure

Other Executed Failures

Job	Test File	Test Function	Error	Related?	Why
jit-kernel-b200-test	`test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py`	`test_ltx2_qknorm_split_rope_matches_torch_exactly`	`AssertionError: assert False`	🟢	LTX2 diffusion qk-norm/rope kernel; different `jit_kernel` module, not `kvcacheio`
base-c-8gpu-h20 (0) · (1)	N/A (import)	N/A	`ImportError: libnvshmem_host.so.3`	🟢	Infra: missing shared lib at import, before tests
extra-a-1gpu-large (1)	`test/registered/perf/test_bench_serving_1gpu_part2.py`	N/A (server startup)	`ConnectionRefusedError`	🟢	Server-startup/perf flake; unrelated
stage-b-4-npu-a3	NPU perf suite	`Qwen3-32B-w4a4-LAOS`	`AssertionError: 868.0 not >= 1000`	🟢	NPU quantized-model throughput perf-threshold flake
multimodal-gen-2-npu-a3	`test/server/ascend/test_server_2_npu.py`	`test_diffusion_generation[qwen_image_t2i_2npu]`	`Consistency Check FAILED`	🟢	NPU diffusion consistency; unrelated to KV-cache
stage-b-1gpu-xpu	N/A (import)	N/A	`OSError: libavutil.so.60 / torchcodec`	🟢	Infra: XPU torchcodec/ffmpeg lib load failure

Details / what to do before merge

No action needed on the 14 red jobs — every one is a pre-existing perf-threshold flake (test_bench_serving_*, test_vlm_perf_5090, NPU throughput), an infra library-load failure (libnvshmem, libavutil, VRAM cleanup), or an unrelated test (test_ltx2_qknorm_split_rope, NPU diffusion). None touch HiCache write-back, cache_controller, or memory_pool_host.
Close the coverage gap (the important one): the CPU-index fix + dispatch logic is verified green on base-a-test-cpu, and the shared kernels on NVIDIA base-b-kernel-unit, but no AMD PR-CI job exercises the ROCm build of the modified .cuh files or the AMD write-back enablement. Before merge, run test/registered/jit/test_hicache.py (and ideally an end-to-end --enable-hierarchical-cache run) on AMD MI3xx hardware to confirm the HIP-compiled staged write-back path works — a green PR-CI run does not cover it.

Generated by amd-bot using Claude Code CLI

HaiShaw

Can we add an unit test case file as well?

HaiShaw · 2026-07-06T05:42:15Z

        .with_strides({N, 1})
        .with_dtype(cache_dtype)
-        .with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
+        .with_device<kDLCUDA, kDLROCM, kDLCUDAHost, kDLROCMHost, kDLCPU>()


Can we use KDLGPU instead?

Addressed your comments and added unit tests. PTAL.

HaiShaw · 2026-07-06T05:43:49Z

@amd-bot ci-status

amd-bot · 2026-07-06T05:49:34Z

@HaiShaw

CI Status for PR #28534

Merge verdict: No executed CI failure is attributable to this PR — all 14 real failures are pre-existing perf regressions, infra/runner problems (VRAM-cleanup, missing .so libs), or unrelated tests on other backends (NPU/XPU/diffusion). The CUDA HiCache staged write-back kernel this PR touches is directly tested and green. However, the PR's primary value — enabling the JIT staged write-back kernel on ROCm/HIP — is not exercised by any PR-CI test (the AMD kernel tests are nightly-only), so a green run does not verify the ROCm path. Safe to merge only after the ROCm path is validated.

Caution

This PR's headline change (enabling can_use_jit / can_use_write_back_jit on HIP in memory_pool_host.py, and the ROCm paths in hicache.cuh / staged_write_back.cuh) is not run by any PR-CI test. The dedicated AMD kernel tests — test/registered/jit/test_kvcacheio_asymmetric.py and test/registered/amd/test_deepseek_r1_hicache_mi35x.py — are both nightly=True, and test/registered/jit/test_hicache.py is CUDA-only. The AMD JIT unit job (jit-kernel-unit-test-amd, green) runs unrelated kernels (activation, rmsnorm, …). Before merge, run the AMD nightly kernel suite nightly-amd-kernel-1-gpu (which includes test_kvcacheio_asymmetric.py) on this branch to actually validate the ROCm staged write-back path.

Changed files: jit_kernel/csrc/kvcacheio/hicache.cuh (+52/-10), staged_write_back.cuh (+8/-8), srt/managers/cache_controller.py (+8/-1), srt/mem_cache/memory_pool_host.py (+28/-4)

What IS covered (green): CUDA test_hicache.py::test_hicache_page_first_staged_write_back_mha[*] and test_hicache_transfer_* all PASSED in jit-kernel-unit-test — this exercises the exact CUDA staged-write-back path plus the cache_controller.py gating logic. CUDA HiCache serving tests (test_hicache_storage*.py, base-b/base-c/extra-a) also passed.

Executed CI failure attribution: AMD: 7 failures (0 related) · Others: 7 failures (0 related). 6 additional *-finish jobs are cascade aggregators, not independent failures.

AMD Executed Failures

Job	Test File	Test Function	Error	Related?	Why
stage-b-1gpu-small-amd (7)	`test/registered/perf/test_vlm_perf_5090.py`	perf suite	`returned exit code 1` (perf threshold)	🟢	VLM perf benchmark; PR touches KV-cache IO, not VLM perf
stage-b-1gpu-large-amd (0)	`test/registered/perf/test_bench_serving_1gpu_part1.py`	perf suite	`returned exit code 1`	🟢	serving perf benchmark, unrelated to write-back kernel
stage-b-2gpu-large-amd (0)	`test/registered/perf/test_bench_serving_2gpu.py`	perf suite	`returned exit code 1`	🟢	serving perf benchmark
extra-a-2gpu-large-amd	`test/registered/kv_canary/test_self_e2e_pp_perturb.py`	PP perturb e2e	`AssertionError: -1 != 200` ReadTimeout on `127.0.0.1:11000`	🟢	server read-timeout (PP e2e); no HiCache kernel path in trace
stage-b-large-8gpu-mi35x-disagg-amd	`test/registered/amd/disaggregation/test_disaggregation_pp.py`	PP disagg	`ClientConnectorError` / `Internal Server Error`	🟢	disagg PP server-connect failure; unrelated to page_first write-back
stage-b-1gpu-small-amd-mi35x	N/A (infra)	N/A	`VRAM cleanup unsuccessful after 3 attempts` (zombie KFD, ROCm/aiter#2061)	🟢	runner infra — needs node reboot
stage-c-large-8gpu-amd-mi35x	N/A (infra)	N/A	`VRAM cleanup unsuccessful after 3 attempts`	🟢	runner infra — same zombie-KFD class

Other Executed Failures

Job	Test File	Test Function	Error	Related?	Why
call-jit-kernel-tests / jit-kernel-b200-test	`test/registered/jit/diffusion/test_ltx2_qknorm_split_rope.py`	`test_ltx2_qknorm_split_rope_matches_torch_exactly`, `..._custom_op_torch_compile_fullgraph`	`AssertionError: assert False` (`torch.equal(q_ref, q_out)`)	🟢	Same workflow but a different JIT kernel (diffusion qk-norm split-rope), not the kvcacheio/hicache kernels this PR edits
base-c-test-8-gpu-h20 (0)	N/A (env)	N/A	`ImportError: libnvshmem_host.so.3: cannot open shared object file`	🟢	missing shared lib on runner (env/infra)
base-c-test-8-gpu-h20 (1)	N/A (env)	N/A	`ImportError: libnvshmem_host.so.3`	🟢	same env issue
extra-a-1gpu-large (1)	`test/registered/perf/test_bench_serving_1gpu_part2.py`	perf suite	`returned exit code 1`	🟢	serving perf benchmark
stage-b-test-4-npu-a3	`test/registered/ascend/basic_function/quant/test_npu_w4a4_quantization.py`	NPU quant	`returned exit code 1`	🟢	NPU W4A4 quantization, unrelated backend/path
multimodal-gen-test-2-npu-a3	`sglang/multimodal_gen/test/server/ascend/test_server_2_npu.py`	`test_diffusion_generation[qwen_image_t2i_2npu, wan2_2_t2v_14b_w8a8_2npu]`	Diffusion testcase failed check(s)	🟢	NPU diffusion serving, unrelated
stage-b-test-1-gpu-xpu	`test/registered/xpu/test_triton_attention_backend.py`	XPU triton attn	`OSError: libavutil.so.60: cannot open shared object file` (torchcodec)	🟢	XPU dependency/env issue, unrelated

Details / what to do before merge

Close the ROCm coverage gap (the important one): run the AMD nightly kernel suite on this branch so test_kvcacheio_asymmetric.py (AMD) exercises the newly-HIP-enabled can_use_write_back_jit path; optionally the nightly-amd-8-gpu-mi35x-deepseek-r1-hicache suite for an e2e HiCache check. The CUDA equivalent (test_hicache.py) already passed, so the kernel logic is sound on CUDA — the concern is strictly the ROCm compile/execution path added here.
All 14 executed failures are unrelated — a mix of perf-threshold benchmarks, runner infra (VRAM zombie-KFD ROCm/aiter#2061, missing libnvshmem/libavutil), and other-backend tests. No triage needed on this PR's behalf, though the libnvshmem/VRAM-cleanup runner issues are worth flagging to infra independently.
The 6 *-finish failures are fast-fail aggregators reflecting the above; no separate investigation.

Generated by amd-bot using Claude Code CLI

…ck unit test Address review feedback on the HiCache JIT device-type matchers: - Replace the explicit `kDLCUDA, kDLROCM` / `kDLCUDAHost, kDLROCMHost` enumerations in hicache.cuh and staged_write_back.cuh with the platform-conditional `kDLGPU` / `kDLGPUHost` aliases, matching the convention already used by the other JIT kernels. Add the missing `kDLGPUHost` alias next to `kDLGPU` in utils.cuh so the host side is symmetric. - Add test/registered/jit/test_hicache_page_first_write_back.py covering the page_first + `kernel` staged write-back (D2H) and load (H2D) roundtrip for MHA and MLA across page counts around the staging capacity. Unlike test_hicache.py (CUDA-only), this file is also registered for the AMD PR-CI kernel suite so the ROCm/HIP build and execution of the modified kernels are validated on AMD hardware. Verified passing on MI355X (gfx950, ROCm 7.2).

… API The new test imported ALLOC_MEMORY_FUNCS / alloc_with_pin_memory from memory_pool_host and asserted host_pool.can_use_jit, which broke on the merged main: those helpers now live in sglang.srt.mem_cache.pool_host.common and the JIT flag was split into can_use_jit (load/transfer) and can_use_write_back_jit (staged write-back). Align the imports and assert can_use_write_back_jit for the write-back path, mirroring test_hicache.py. Verified passing on MI355X (gfx950, ROCm 7.2).

Duyi-Wang · 2026-07-07T08:06:00Z

@amd-bot ci-status

amd-bot · 2026-07-07T08:11:13Z

@Duyi-Wang

CI Status for PR #28534

Merge verdict: No executed CI failure is attributable to this PR — the changed hicache code IS exercised and green (new test_hicache_page_first_write_back.py = 13/13 passed on AMD; existing test_hicache.py staged write-back = 47/47 passed on NVIDIA). However, PR CI is not yet complete: several AMD shards (stage-c-test-large-8-gpu-amd ×4, stage-b-test-1-gpu-small-amd shard 12) are still queued. Do not merge on "green" until those finish. All 14 executed failures are pre-existing/unrelated (sparse-MLA kernel, CPU build, quant timeout, FP4 accuracy, NUMA/infra, NPU/XPU backends).

Warning

Changed paths are covered by test_hicache_page_first_write_back.py (new, ran on jit-kernel-unit-test-amd) and test_hicache.py::test_hicache_page_first_staged_write_back_* (ran on jit-kernel-unit-test) — both green. But required AMD downstream shards are still running/queued; a couple of AMD stages have not reported yet. Green ≠ verified until stage-c-test-large-8-gpu-amd and stage-b-test-1-gpu-small-amd (12) complete.

Changed files: jit_kernel/csrc/kvcacheio/hicache.cuh (+52/-10), staged_write_back.cuh (+8/-8), include/sgl_kernel/utils.cuh (+2/-0), srt/managers/cache_controller.py (+8/-1), srt/mem_cache/memory_pool_host.py (+28/-4), test/registered/jit/test_hicache_page_first_write_back.py (+261, new)

Executed CI failure attribution: AMD: 3 failures (0 related) · Others: 11 failures (0 related) · plus 3 fast-fail cascade finish jobs · AMD stage-c (×4) + stage-b-1gpu-small (12) still pending

AMD Executed Failures

Job	Test File	Test Function	Error	Related?	Why
stage-b-1gpu-small-amd-mi35x	`test/registered/quant/test_quark_mxfp4.py`	N/A (suite)	`exit 255` after 1200s timeout	🟢	Quark MXFP4 quant path; PR touches only hicache kvcacheio kernels + hicache Python
stage-b-2gpu-large-amd	`test/registered/model_loading/test_load_weights_from_remote_instance.py`	N/A (suite)	`exit 255`	🟢	Remote weight-loading path; unrelated to hicache write-back
stage-c-large-8-gpu-amd-mi35x	N/A	N/A	`Run test timed out after 60 minutes`	🟢	Job-level timeout (infra); no hicache test implicated

Other Executed Failures

Job	Test File	Test Function	Error	Related?	Why
call-jit-kernel-tests / jit-kernel-unit-test	`test/registered/jit/test_sparse_mla_q8kv8_prefill_sm90.py`	`test_sparse_mla_q8kv8_prefill_corner_cases[512-False-65-256-592]`	`AssertionError: assert tensor(False)` (isfinite)	🟢	Sparse-MLA prefill kernel; PR's own hicache tests in this same job passed
build-test	`sgl-kernel/csrc/cpu/norm.cpp:391`	N/A (build)	`error: incomplete type` → CMake/ninja build failed	🟢	CPU (arm64) build of `sgl-kernel`; PR touches `jit_kernel/csrc`, not `sgl-kernel/csrc/cpu`
base-c-test-4-gpu-b200 (2)	`test/registered/models_e2e/*` (FP4)	N/A (accuracy)	`AssertionError: 0.934 not >= 0.935`	🟢	Marginal FP4 accuracy threshold; no hicache path
base-c-test-4-gpu-gb300 (0/1)	`test/registered/utils/test_numa_utils.py`	N/A	`ValueError: invalid literal for int() ... b'kB'`; scheduler died at init (exit -3)	🟢	NUMA/memory-parse infra failure at startup; unrelated to PR
stage-b-test-1-gpu-xpu	XPU suite	N/A	job failure	🟢	XPU backend; PR is AMD/CUDA hicache
stage-b-test-4-npu-a3 · single-node-poc ascend	NPU/Ascend suites	N/A	job failure	🟢	NPU backend; unrelated
pr-test-finish · finish · pr-test-npu-finish	N/A	N/A	aggregator	🟢	Fast-fail cascade rollup of the above; not independent failures

Details / what to do before merge

Wait for AMD to finish. stage-c-test-large-8-gpu-amd (shards 0-3) and stage-b-test-1-gpu-small-amd (12) are queued. The AMD signal is incomplete until they report — do not merge on the current partial green.
Coverage is satisfied for the changed code: the new CPU-index / page-first staged write-back test passed 13/13 on AMD, and the existing staged-write-back suite passed 47/47 on NVIDIA, directly exercising hicache.cuh / staged_write_back.cuh and the memory_pool_host.py / cache_controller.py write-back path.
No action needed on the red X's — all are pre-existing failures on unrelated code paths (sparse-MLA kernel, CPU sgl-kernel build, quant timeout, FP4 accuracy, NUMA infra, NPU/XPU backends). These are not introduced by this PR; if they block the finish gate, they should be addressed on main/separately (or use bypass-fastfail sparingly if full AMD signal is needed).

Generated by amd-bot using Claude Code CLI

Resolve conflict in memory_pool_host.py from main's mem_cache refactor (MHA host pool moved to pool_host/mha.py; new pools added). Reapply the PR's HIP enablement so MHA and MLA host pools keep can_use_jit / can_use_write_back_jit on ROCm: - memory_pool_host.py: MLATokenToKVPoolHost (_is_cuda -> _is_cuda or _is_hip) - pool_host/mha.py: MHATokenToKVPoolHost (_is_cuda -> _is_cuda or _is_hip) Newly-added pools (DeepSeekV4/Mamba/DSA) are left CUDA-only (out of this PR's scope). Also update the new unit test import for MHATokenToKVPoolHost's new location (pool_host.mha). Verified: HiCache JIT tests pass on MI355X (ROCm 7.2).

…28534) Co-authored-by: Duyi-Wang <duyi.wang@amd.com> (cherry picked from commit d74619b)

…follow-up to #28534) (#30622)

…follow-up to sgl-project#28534) (sgl-project#30622)

…gl-project#28534) Co-authored-by: Duyi-Wang <duyi.wang@amd.com>

…follow-up to sgl-project#28534) (sgl-project#30622)

github-actions Bot added hicache Hierarchical Caching for SGLang jit-kernel labels Jun 17, 2026

gemini-code-assist Bot reviewed Jun 17, 2026

View reviewed changes

This was referenced Jun 17, 2026

[CI Monitor] Daily Report - 2026-06-17 bingxche/sglang-ci-bot#107

Open

[CI Monitor] Daily Report - 2026-06-18 bingxche/sglang-ci-bot#108

Open

AMD-yanfeiwang force-pushed the rocm/hicache-jit-staged-writeback branch 2 times, most recently from eeec323 to 09d02f9 Compare June 18, 2026 02:12

AMD-yanfeiwang marked this pull request as ready for review June 18, 2026 03:25

AMD-yanfeiwang requested review from BBuf, DarkSharpness, HydraQYH, Ying1123, alphabetc1, celve, hanming-lu, hnyls2002, hzh0425, ispobock, merrymercy, xiezhq-hermann, yizhang2077 and yuan-luo as code owners June 18, 2026 03:25

AMD-yanfeiwang changed the title ~~[ROCm] Enable JIT staged HiCache write-back and fix CPU-index crash~~ [AMD] Enable JIT staged HiCache write-back and fix CPU-index crash Jun 18, 2026

yctseng0211 added the run-ci label Jun 18, 2026

AMD-yanfeiwang added 4 commits June 19, 2026 06:09

AMD-yanfeiwang force-pushed the rocm/hicache-jit-staged-writeback branch from 8485afd to 21243ba Compare June 19, 2026 06:15

HaiShaw added the bypass-fastfail label Jul 5, 2026

Merge branch 'main' into rocm/hicache-jit-staged-writeback

7dad0e1

HaiShaw reviewed Jul 6, 2026

View reviewed changes

AMD-yanfeiwang added 2 commits July 6, 2026 07:33

This was referenced Jul 6, 2026

[CI Monitor] Daily Report - 2026-07-06 bingxche/sglang-ci-bot#127

Open

[CI Monitor] Daily Report - 2026-07-07 bingxche/sglang-ci-bot#128

Open

Merge branch 'main' into rocm/hicache-jit-staged-writeback

d33fbea

This was referenced Jul 8, 2026

[CI Monitor] Daily Report - 2026-07-08 bingxche/sglang-ci-bot#129

Open

[CI Monitor] Daily Report - 2026-07-09 bingxche/sglang-ci-bot#130

Open

zhangfei829 mentioned this pull request Jul 9, 2026

let page-first work on kernel maning00/sglang#11

Merged

HaiShaw approved these changes Jul 9, 2026

View reviewed changes

HaiShaw merged commit d74619b into sgl-project:main Jul 9, 2026
121 of 180 checks passed

AMD-yanfeiwang mentioned this pull request Jul 9, 2026

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

Merged

4 tasks

michaelzhang-ai mentioned this pull request Jul 10, 2026

[AMD] Cherry-pick AMD fixes into release/v0.5.15 #30722

Closed

3 tasks

michaelzhang-ai pushed a commit that referenced this pull request Jul 10, 2026

[AMD] Enable JIT staged HiCache write-back and fix CPU-index crash (#…

b38810f

…28534) Co-authored-by: Duyi-Wang <duyi.wang@amd.com> (cherry picked from commit d74619b)

HaiShaw pushed a commit that referenced this pull request Jul 15, 2026

[AMD] Remove ROCm page_first+kernel -> layer_first HiCache fallback (…

a3194d3

…follow-up to #28534) (#30622)

roopaksrivastav pushed a commit to roopaksrivastav/sglang that referenced this pull request Jul 16, 2026

[AMD] Remove ROCm page_first+kernel -> layer_first HiCache fallback (…

b3c243b

…follow-up to sgl-project#28534) (sgl-project#30622)

llying-001 mentioned this pull request Aug 4, 2026

Add HiCache JIT test and benchmark for ROCm/HIP CI support #30350

Closed

5 tasks

dorado269 mentioned this pull request Aug 6, 2026

docs(patches): give every patch a schema-checked upstream status record AMD-AGI/Infera#95

Closed

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

[AMD] Enable JIT staged HiCache write-back and fix CPU-index crash (s…

4e012ad

…gl-project#28534) Co-authored-by: Duyi-Wang <duyi.wang@amd.com>

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

[AMD] Remove ROCm page_first+kernel -> layer_first HiCache fallback (…

69f4c19

…follow-up to sgl-project#28534) (sgl-project#30622)

Uh oh!

Conversation

AMD-yanfeiwang commented Jun 17, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Root cause

Why fix the cause instead of the #28473 layer_first fallback

Modifications

Reproduction

Prefill

Decode

Router (native PD-disaggregation router)

Accuracy / functional tests

Speed tests and profiling

MHA (batch_pages=64)

MLA (batch_pages=64)

CI States

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot Jun 17, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot commented Jun 18, 2026

Uh oh!

HaiShaw commented Jul 5, 2026

Uh oh!

HaiShaw commented Jul 6, 2026

Uh oh!

amd-bot commented Jul 6, 2026

CI Status for PR #28534

AMD Executed Failures

Other Executed Failures

Details / what to do before merge

Uh oh!

HaiShaw left a comment

Choose a reason for hiding this comment

Uh oh!

HaiShaw Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

AMD-yanfeiwang Jul 6, 2026

Choose a reason for hiding this comment

Uh oh!

HaiShaw commented Jul 6, 2026

Uh oh!

amd-bot commented Jul 6, 2026

CI Status for PR #28534

AMD Executed Failures

Other Executed Failures

Details / what to do before merge

Uh oh!

Duyi-Wang commented Jul 7, 2026

Uh oh!

amd-bot commented Jul 7, 2026

CI Status for PR #28534

AMD Executed Failures

Other Executed Failures

Details / what to do before merge

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

AMD-yanfeiwang commented Jun 17, 2026 •

edited by github-actions Bot

Loading

Why fix the cause instead of the #28473 `layer_first` fallback

MHA (`batch_pages=64`)

MLA (`batch_pages=64`)