let page-first work on kernel - #11
Merged
zhangfei829 merged 5 commits intoJul 9, 2026
Merged
Conversation
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.
(cherry picked from commit ad737c3)
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). (cherry picked from commit c17ff90)
…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. (cherry picked from commit 2032ba2)
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.
(cherry picked from commit 21243ba)
…n ROCm (sgl-project#28473)" This reverts commit 9b8c411.
zhangfei829
merged commit Jul 9, 2026
bf5877d
into
perf/hicache-l2-l1-opt-rocm
68 of 78 checks passed
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request enables ROCm (HIP) support for the HiCache JIT kernels and staged write-back kernels. It introduces non-temporal load/store operations using native Clang vector types under the USE_ROCM guard in hicache.cuh, updates TensorMatcher to support ROCm device types, enables JIT flags on HIP in memory_pool_host.py, and removes the ROCm fallback logic in server_args.py since the page_first layout is now fully supported on ROCm. There are no review comments, so I have no feedback to provide.
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.
TianDi101
pushed a commit
that referenced
this pull request
Jul 15, 2026
* [ROCm] Enable JIT staged HiCache write-back and fix CPU-index crash
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.
(cherry picked from commit ad737c3)
* [ROCm] Accept ROCm device types in staged HiCache write-back matchers
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).
(cherry picked from commit c17ff90)
* [ROCm] Accept ROCm device types in hicache.cuh load/store kernel matchers
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.
(cherry picked from commit 2032ba2)
* [ROCm] Vectorize HiCache nontemporal load/store via native vector types
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.
(cherry picked from commit 21243ba)
* Revert "[AMD] Fall back to layer_first layout for kernel write-back on ROCm (sgl-project#28473)"
This reverts commit 9b8c411.
---------
Co-authored-by: AMD-yanfeiwang <yanfei.wang@amd.com>
TianDi101
pushed a commit
that referenced
this pull request
Jul 15, 2026
* [ROCm] Enable JIT staged HiCache write-back and fix CPU-index crash
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.
(cherry picked from commit ad737c3)
* [ROCm] Accept ROCm device types in staged HiCache write-back matchers
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).
(cherry picked from commit c17ff90)
* [ROCm] Accept ROCm device types in hicache.cuh load/store kernel matchers
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.
(cherry picked from commit 2032ba2)
* [ROCm] Vectorize HiCache nontemporal load/store via native vector types
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.
(cherry picked from commit 21243ba)
* Revert "[AMD] Fall back to layer_first layout for kernel write-back on ROCm (sgl-project#28473)"
This reverts commit 9b8c411.
---------
Co-authored-by: AMD-yanfeiwang <yanfei.wang@amd.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Cherry-pick sgl-project#28534's ROCm hicache commits (author preserved) and revert sgl-project#28473 so
ROCm runs the same
page_firstkernel HiCache path as CUDA.sgl-project#28473 unconditionally forces
page_first + kernel -> layer_firston ROCm, whichpermanently disables the JIT staged write-back path. sgl-project#28534 makes that path build
and run on ROCm (hipcc +
can_use_jit/can_use_write_back_jiton HIP + CPU-indexcrash fix), so sgl-project#28473's fallback is now redundant and reverted.
Changes: 5 files, +96/-37 (hicache.cuh / staged_write_back.cuh ROCm guards;
memory_pool_host
_is_cuda or _is_hip; cache_controller index gating; server_args -14).Verified on MI355X/ROCm7.2, DeepSeek-R1-0528-MXFP4-th, page_first+kernel: layout
stays page_first, no crash, load-back BW scales with CU quota (~3→46 GiB/s, CU 2→32).
CI States
Latest PR Test (Base): ❌ Run #28999860254
Latest PR Test (Extra): ❌ Run #28999860132