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

Co-authored-by: Duyi-Wang <duyi.wang@amd.com>

python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh

@@ -37,44 +37,86 @@ inline constexpr auto get_mem_package() {
 template <int kUnit>
 using PackageType = decltype(get_mem_package<kUnit>());
 
+// NVIDIA exposes an explicit "do not allocate in L1" cache hint via PTX. ROCm
+// has no equivalent PTX, but non-temporal (streaming) loads/stores express the
+// same intent for one-shot HiCache write-back traffic that should not pollute
+// the cache. Guard the PTX behind USE_ROCM so the JIT module also compiles with
+// hipcc; see python/sglang/jit_kernel/utils.py for the ROCm build flags.
+#ifdef USE_ROCM
+// Native Clang vector types so a single __builtin_nontemporal_{load,store} maps
+// to one vectorized global_{load,store}_dwordx{2,4}. Issuing N independent
+// 32-bit nontemporal ops instead leaves merging to the LoadStoreVectorizer,
+// which is not guaranteed and may drop the nontemporal hint, throttling HiCache
+// bandwidth. uint2/uint4 already carry 8B/16B alignment matching the vector
+// types, so the pointer reinterpret_casts stay correctly aligned.
+typedef uint32_t native_uint2 __attribute__((ext_vector_type(2)));
+typedef uint32_t native_uint4 __attribute__((ext_vector_type(4)));
+#endif
+
 SGL_DEVICE uint1 load_nc(const uint1* __restrict__ src) {
+#ifndef USE_ROCM
   uint32_t tmp;
   asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src));
   return uint1{tmp};
+#else
+  return uint1{__builtin_nontemporal_load(&src->x)};
+#endif
 }
 
 SGL_DEVICE uint2 load_nc(const uint2* __restrict__ src) {
+#ifndef USE_ROCM
   uint32_t tmp0, tmp1;
   asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src));
   return uint2{tmp0, tmp1};
+#else
+  native_uint2 tmp = __builtin_nontemporal_load(reinterpret_cast<const native_uint2*>(src));
+  return __builtin_bit_cast(uint2, tmp);
+#endif
 }
 
 SGL_DEVICE uint4 load_nc(const uint4* __restrict__ src) {
+#ifndef USE_ROCM
   uint32_t tmp0, tmp1, tmp2, tmp3;
   asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];"
                : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3)
                : "l"(src));
   return uint4{tmp0, tmp1, tmp2, tmp3};
+#else
+  native_uint4 tmp = __builtin_nontemporal_load(reinterpret_cast<const native_uint4*>(src));
+  return __builtin_bit_cast(uint4, tmp);
+#endif
 }
 
 SGL_DEVICE void store_nc(uint1* __restrict__ dst, const uint1& value) {
+#ifndef USE_ROCM
   uint32_t tmp = value.x;
   asm volatile("st.global.L1::no_allocate.b32 [%0],%1;" ::"l"(dst), "r"(tmp));
+#else
+  __builtin_nontemporal_store(value.x, &dst->x);
+#endif
 }
 
 SGL_DEVICE void store_nc(uint2* __restrict__ dst, const uint2& value) {
+#ifndef USE_ROCM
   uint32_t tmp0 = value.x;
   uint32_t tmp1 = value.y;
   asm volatile("st.global.L1::no_allocate.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1));
+#else
+  __builtin_nontemporal_store(__builtin_bit_cast(native_uint2, value), reinterpret_cast<native_uint2*>(dst));
+#endif
 }
 
 SGL_DEVICE void store_nc(uint4* __restrict__ dst, const uint4& value) {
+#ifndef USE_ROCM
   uint32_t tmp0 = value.x;
   uint32_t tmp1 = value.y;
   uint32_t tmp2 = value.z;
   uint32_t tmp3 = value.w;
   asm volatile(
       "st.global.L1::no_allocate.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3));
+#else
+  __builtin_nontemporal_store(__builtin_bit_cast(native_uint4, value), reinterpret_cast<native_uint4*>(dst));
+#endif
 }
 
 }  // namespace details
@@ -256,18 +298,18 @@ struct HiCacheKernel {
     TensorMatcher({-1, D})  //
         .with_strides({N, 1})
         .with_dtype(cache_dtype)
-        .with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
+        .with_device<kDLGPU, kDLGPUHost, kDLCPU>()
         .verify(k_cache_src)
         .verify(v_cache_src);
     TensorMatcher({-1, D})  //
         .with_strides({M, 1})
         .with_dtype(cache_dtype)
-        .with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
+        .with_device<kDLGPU, kDLGPUHost, kDLCPU>()
         .verify(k_cache_dst)
         .verify(v_cache_dst);
     TensorMatcher({L})  //
         .with_dtype<int32_t, int64_t>(indices_dtype)
-        .with_device<kDLCUDA>(indices_device)
+        .with_device<kDLGPU>(indices_device)
         .verify(indices_src)
         .verify(indices_dst);
 
@@ -323,14 +365,14 @@ struct HiCacheKernel {
 
     TensorMatcher({N})  //
         .with_dtype<uint64_t>()
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLGPU>(device_)
         .verify(k_ptr_src)
         .verify(v_ptr_src)
         .verify(k_ptr_dst)
         .verify(v_ptr_dst);
     TensorMatcher({L})  //
         .with_dtype<int32_t, int64_t>(dtype_)
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLGPU>(device_)
         .verify(indices_src)
         .verify(indices_dst);
 
@@ -381,16 +423,16 @@ struct HiCacheKernel {
     TensorMatcher({-1, D})  //
         .with_strides({N, 1})
         .with_dtype(cache_dtype)
-        .with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
+        .with_device<kDLGPU, kDLGPUHost, kDLCPU>()
         .verify(cache_src);
     TensorMatcher({-1, D})  //
         .with_strides({M, 1})
         .with_dtype(cache_dtype)
-        .with_device<kDLCUDA, kDLCUDAHost, kDLCPU>()
+        .with_device<kDLGPU, kDLGPUHost, kDLCPU>()
         .verify(cache_dst);
     TensorMatcher({L})  //
         .with_dtype<int32_t, int64_t>(indices_dtype)
-        .with_device<kDLCUDA>(indices_device)
+        .with_device<kDLGPU>(indices_device)
         .verify(indices_src)
         .verify(indices_dst);
 
@@ -441,12 +483,12 @@ struct HiCacheKernel {
 
     TensorMatcher({N})  //
         .with_dtype<uint64_t>()
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLGPU>(device_)
         .verify(ptr_src)
         .verify(ptr_dst);
     TensorMatcher({L})  //
         .with_dtype<int32_t, int64_t>(dtype_)
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLGPU>(device_)
         .verify(indices_src)
         .verify(indices_dst);
 

python/sglang/jit_kernel/csrc/kvcacheio/staged_write_back.cuh

@@ -210,41 +210,41 @@ struct HiCacheStagedWriteBackKernel {
 
     TensorMatcher({T, N, D})  //
         .with_dtype(cache_dtype)
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLGPU>(device_)
         .verify(staging_k);
     if constexpr (!kIsMLA) {
       TensorMatcher({T, N, D})  //
           .with_dtype(cache_dtype)
-          .with_device<kDLCUDA>(device_)
+          .with_device<kDLGPU>(device_)
           .verify(staging_v);
     }
     TensorMatcher({-1, N, D})  //
         .with_dtype(cache_dtype)
-        .with_device<kDLCPU, kDLCUDAHost>()
+        .with_device<kDLCPU, kDLGPUHost>()
         .verify(k_cache_dst);
     if constexpr (!kIsMLA) {
       TensorMatcher({-1, N, D})  //
           .with_dtype(cache_dtype)
-          .with_device<kDLCPU, kDLCUDAHost>()
+          .with_device<kDLCPU, kDLGPUHost>()
           .verify(v_cache_dst);
     }
     TensorMatcher({N})  //
         .with_dtype<uint64_t>()
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLGPU>(device_)
         .verify(k_ptr_src);
     if constexpr (!kIsMLA) {
       TensorMatcher({N})  //
           .with_dtype<uint64_t>()
-          .with_device<kDLCUDA>(device_)
+          .with_device<kDLGPU>(device_)
           .verify(v_ptr_src);
     }
     TensorMatcher({P})  //
         .with_dtype<int32_t, int64_t>(indices_dtype)
-        .with_device<kDLCUDA>(device_)
+        .with_device<kDLGPU>(device_)
         .verify(page_indices_src);
     TensorMatcher({T})  //
         .with_dtype<int64_t>(dst_indices_dtype)
-        .with_device<kDLCPU, kDLCUDAHost>()
+        .with_device<kDLCPU, kDLGPUHost>()
         .verify(dst_indices_cpu);
 
     RuntimeCheck(page_size > 0, "HiCache staged relayout: page_size must be positive");

python/sglang/jit_kernel/include/sgl_kernel/utils.cuh

@@ -89,8 +89,10 @@ using fp32x4_t = float4;
 // DLPack device type for the current platform
 #ifndef USE_ROCM
 inline constexpr auto kDLGPU = kDLCUDA;
+inline constexpr auto kDLGPUHost = kDLCUDAHost;
 #else
 inline constexpr auto kDLGPU = kDLROCM;
+inline constexpr auto kDLGPUHost = kDLROCMHost;
 #endif
 
 namespace device {

python/sglang/srt/managers/cache_controller.py

@@ -675,7 +675,14 @@ class HiCacheController:
             return
 
         op = CacheOperation.merge_ops(self.write_queue)
-        # Page-first write-back JIT kernels can keep destination host indices on CPU.
+        # Kernel write-back keeps host indices on CPU only for page_first AND only
+        # when the staged JIT write-back kernel is available (it stages through
+        # device memory and accepts CPU destination indices). Otherwise we fall back
+        # to the plain transfer kernel, whose CUDA/HIP implementation requires
+        # device-resident destination indices -- so the indices must be moved to the
+        # device first. Without the can_use_write_back_jit check this crashes on
+        # backends where the JIT kernel is unavailable, with
+        # "Destination indices must be a CUDA tensor".
         if (
             self.io_backend == "kernel"
             and self.mem_pool_host.layout == "page_first"

python/sglang/srt/mem_cache/memory_pool_host.py

@@ -94,7 +94,11 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
             device,
             allocator_type,
         )
-        self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel(
+        # The JIT HiCache kernels also build with hipcc (ROCm): the PTX-only
+        # helpers in hicache.cuh are guarded by USE_ROCM and the staged
+        # write-back kernel has a ROCm path, so enable them on HIP too. This
+        # keeps the ROCm write-back path consistent with CUDA.
+        self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel(
             element_size=self.kv_cache_dim * self.dtype.itemsize
         )
 
@@ -214,7 +218,11 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache):
         if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
             return
 
-        self.can_use_write_back_jit = _is_cuda and can_use_write_back_jit_kernel(
+        # The staged write-back JIT kernel builds with hipcc and has a ROCm
+        # path, so enable it on HIP too (consistent with the CUDA path).
+        self.can_use_write_back_jit = (
+            _is_cuda or _is_hip
+        ) and can_use_write_back_jit_kernel(
             element_size=self.kv_cache_dim * self.dtype.itemsize,
         )
         if not self.can_use_write_back_jit:

python/sglang/srt/mem_cache/pool_host/mha.py

@@ -88,7 +88,11 @@ class MHATokenToKVPoolHost(HostKVCache):
             allocator_type,
         )
         self.element_dim = self.device_pool.head_num * self.device_pool.head_dim
-        self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel(
+        # The JIT HiCache kernels also build with hipcc (ROCm): the PTX-only
+        # helpers in hicache.cuh are guarded by USE_ROCM and the staged
+        # write-back kernel has a ROCm path, so enable them on HIP too. This
+        # keeps the ROCm write-back path consistent with CUDA.
+        self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel(
             element_size=self.element_dim * self.dtype.itemsize
         )
 
@@ -170,7 +174,11 @@ class MHATokenToKVPoolHost(HostKVCache):
         if self.layout != "page_first" or (_is_npu or _is_xpu or _is_mps):
             return
 
-        self.can_use_write_back_jit = _is_cuda and can_use_write_back_jit_kernel(
+        # The staged write-back JIT kernel builds with hipcc and has a ROCm
+        # path, so enable it on HIP too (consistent with the CUDA path).
+        self.can_use_write_back_jit = (
+            _is_cuda or _is_hip
+        ) and can_use_write_back_jit_kernel(
             element_size=self.element_dim * self.dtype.itemsize,
         )
         if not self.can_use_write_back_jit:

test/registered/jit/test_hicache_page_first_write_back.py

@@ -0,0 +1,259 @@
+"""Unit tests for the page_first + ``kernel`` JIT HiCache write-back / load path.
+
+This file specifically exercises the JIT staged write-back and load kernels that
+accept a CPU-resident destination index and stage through device memory
+(``staged_write_back.cuh`` / ``hicache.cuh``). Unlike ``test_hicache.py`` (which
+is registered CUDA-only), this file is also registered for the AMD PR-CI kernel
+suite so the ROCm/HIP build and execution of those kernels are validated on AMD
+hardware, not just CUDA.
+"""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.jit_kernel.hicache import can_use_write_back_jit_kernel
+from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool
+from sglang.srt.mem_cache.memory_pool_host import MLATokenToKVPoolHost
+from sglang.srt.mem_cache.pool_host.common import (
+    ALLOC_MEMORY_FUNCS,
+    alloc_with_pin_memory,
+)
+from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
+from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
+from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
+
+register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+register_amd_ci(est_time=30, stage="jit-kernel-unit", runner_config="amd")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available()
+    or is_npu()
+    or is_xpu()
+    or not (is_cuda() or is_hip()),
+    reason="HiCache JIT write-back tests require CUDA/ROCm.",
+)
+
+DEVICE = "cuda"
+PAGE_SIZE = 1 if is_hip() else 16
+NUM_LAYERS = 2
+MHA_ELEMENT_DIMS = [128, 512]
+MLA_ELEMENT_DIMS = [576]
+# Include counts around and above the staging capacity so both the single-pass
+# and the multi-chunk staged relayout branches are exercised.
+PAGE_COUNTS = [1, 64, 65, 129]
+
+
+def _token_indices_for_pages(
+    pages: torch.Tensor,
+    device: str = DEVICE,
+    dtype: torch.dtype = torch.int64,
+) -> torch.Tensor:
+    parts = [
+        torch.arange(
+            int(page) * PAGE_SIZE,
+            (int(page) + 1) * PAGE_SIZE,
+            device=device,
+            dtype=dtype,
+        )
+        for page in pages.tolist()
+    ]
+    return torch.cat(parts, dim=0)
+
+
+def _pinned_host_pool(host_pool_cls, **kwargs):
+    original_alloc = ALLOC_MEMORY_FUNCS[DEVICE]
+    ALLOC_MEMORY_FUNCS[DEVICE] = alloc_with_pin_memory
+    try:
+        return host_pool_cls(
+            host_to_device_ratio=2.0,
+            host_size=0,
+            page_size=PAGE_SIZE,
+            pin_memory=True,
+            device="cpu",
+            **kwargs,
+        )
+    finally:
+        ALLOC_MEMORY_FUNCS[DEVICE] = original_alloc
+
+
+def _fill_with_offset(tensor: torch.Tensor, offset: int) -> None:
+    data = torch.arange(
+        tensor.numel(), device=tensor.device, dtype=tensor.dtype
+    ).view_as(tensor)
+    tensor.copy_(data + offset)
+
+
+def _assert_pages_equal(host_ref, device_ref, host_pages, device_pages) -> None:
+    for host_page, device_page in zip(host_pages.tolist(), device_pages.tolist()):
+        host_start = host_page * PAGE_SIZE
+        device_start = device_page * PAGE_SIZE
+        assert torch.equal(
+            host_ref[host_start : host_start + PAGE_SIZE].cpu(),
+            device_ref[device_start : device_start + PAGE_SIZE].cpu(),
+        )
+
+
+def _run_mha(element_dim: int, page_count: int) -> None:
+    pool_size = PAGE_SIZE * (page_count + 8)
+    device_pool = MHATokenToKVPool(
+        size=pool_size,
+        page_size=PAGE_SIZE,
+        head_num=element_dim // 128,
+        head_dim=128,
+        dtype=torch.bfloat16,
+        layer_num=NUM_LAYERS,
+        device=DEVICE,
+        enable_memory_saver=False,
+    )
+    host_pool = _pinned_host_pool(
+        MHATokenToKVPoolHost, device_pool=device_pool, layout="page_first"
+    )
+    assert can_use_write_back_jit_kernel(
+        element_size=element_dim * host_pool.dtype.itemsize,
+    )
+    # page_first + kernel staged write-back JIT path must be enabled.
+    assert host_pool.can_use_write_back_jit
+
+    for layer_id in range(NUM_LAYERS):
+        _fill_with_offset(device_pool.k_buffer[layer_id], layer_id)
+        _fill_with_offset(device_pool.v_buffer[layer_id], layer_id + 100)
+
+    device_pages = torch.arange(2, 2 + page_count, device=DEVICE, dtype=torch.int64)
+    host_pages = torch.arange(page_count, 0, -1, dtype=torch.int64)
+    device_indices = _token_indices_for_pages(device_pages)
+    # host_indices stay on the CPU: this is the case the staged JIT kernel must
+    # accept (kDLCPU / kDLGPUHost destination indices).
+    host_indices = _token_indices_for_pages(host_pages, device="cpu")
+    assert not host_indices.is_cuda
+
+    host_pool.backup_from_device_all_layer(
+        device_pool, host_indices, device_indices, "kernel"
+    )
+    torch.cuda.synchronize()
+
+    for layer_id in range(NUM_LAYERS):
+        _assert_pages_equal(
+            host_pool.k_data_refs[layer_id],
+            device_pool.k_buffer[layer_id],
+            host_pages,
+            device_pages,
+        )
+        _assert_pages_equal(
+            host_pool.v_data_refs[layer_id],
+            device_pool.v_buffer[layer_id],
+            host_pages,
+            device_pages,
+        )
+
+    # Load path (prefix-cache hit): exercises the hicache.cuh load matchers.
+    if not host_pool.can_use_jit:
+        return
+    for layer_id in range(NUM_LAYERS):
+        device_pool.k_buffer[layer_id].zero_()
+        device_pool.v_buffer[layer_id].zero_()
+
+    load_pages = torch.arange(1, 1 + page_count, device=DEVICE, dtype=torch.int64)
+    load_indices = _token_indices_for_pages(load_pages)
+    host_indices_device = host_indices.to(DEVICE)
+    for layer_id in range(NUM_LAYERS):
+        host_pool.load_to_device_per_layer(
+            device_pool, host_indices_device, load_indices, layer_id, "kernel"
+        )
+    torch.cuda.synchronize()
+
+    for layer_id in range(NUM_LAYERS):
+        _assert_pages_equal(
+            host_pool.k_data_refs[layer_id],
+            device_pool.k_buffer[layer_id],
+            host_pages,
+            load_pages,
+        )
+        _assert_pages_equal(
+            host_pool.v_data_refs[layer_id],
+            device_pool.v_buffer[layer_id],
+            host_pages,
+            load_pages,
+        )
+
+
+def _run_mla(element_dim: int, page_count: int) -> None:
+    pool_size = PAGE_SIZE * (page_count + 8)
+    device_pool = MLATokenToKVPool(
+        size=pool_size,
+        page_size=PAGE_SIZE,
+        kv_lora_rank=element_dim - 64,
+        qk_rope_head_dim=64,
+        dtype=torch.bfloat16,
+        layer_num=NUM_LAYERS,
+        device=DEVICE,
+        enable_memory_saver=False,
+    )
+    host_pool = _pinned_host_pool(
+        MLATokenToKVPoolHost, device_pool=device_pool, layout="page_first"
+    )
+    assert can_use_write_back_jit_kernel(
+        element_size=element_dim * host_pool.dtype.itemsize,
+    )
+    assert host_pool.can_use_write_back_jit
+
+    for layer_id in range(NUM_LAYERS):
+        _fill_with_offset(device_pool.kv_buffer[layer_id], layer_id)
+
+    device_pages = torch.arange(2, 2 + page_count, device=DEVICE, dtype=torch.int64)
+    host_pages = torch.arange(page_count, 0, -1, dtype=torch.int64)
+    device_indices = _token_indices_for_pages(device_pages)
+    host_indices = _token_indices_for_pages(host_pages, device="cpu")
+    assert not host_indices.is_cuda
+
+    host_pool.backup_from_device_all_layer(
+        device_pool, host_indices, device_indices, "kernel"
+    )
+    torch.cuda.synchronize()
+
+    for layer_id in range(NUM_LAYERS):
+        _assert_pages_equal(
+            host_pool.data_refs[layer_id],
+            device_pool.kv_buffer[layer_id],
+            host_pages,
+            device_pages,
+        )
+
+    if not host_pool.can_use_jit:
+        return
+    for layer_id in range(NUM_LAYERS):
+        device_pool.kv_buffer[layer_id].zero_()
+
+    load_pages = torch.arange(1, 1 + page_count, device=DEVICE, dtype=torch.int64)
+    load_indices = _token_indices_for_pages(load_pages)
+    host_indices_device = host_indices.to(DEVICE)
+    for layer_id in range(NUM_LAYERS):
+        host_pool.load_to_device_per_layer(
+            device_pool, host_indices_device, load_indices, layer_id, "kernel"
+        )
+    torch.cuda.synchronize()
+
+    for layer_id in range(NUM_LAYERS):
+        _assert_pages_equal(
+            host_pool.data_refs[layer_id],
+            device_pool.kv_buffer[layer_id],
+            host_pages,
+            load_pages,
+        )
+
+
+@pytest.mark.parametrize("element_dim", MHA_ELEMENT_DIMS)
+@pytest.mark.parametrize("page_count", PAGE_COUNTS)
+def test_page_first_staged_write_back_mha(element_dim: int, page_count: int) -> None:
+    _run_mha(element_dim, page_count)
+
+
+@pytest.mark.parametrize("element_dim", MLA_ELEMENT_DIMS)
+@pytest.mark.parametrize("page_count", PAGE_COUNTS)
+def test_page_first_staged_write_back_mla(element_dim: int, page_count: int) -> None:
+    _run_mla(element_dim, page_count)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))