From d74619b37337753ed9f0ec9c51c5f010c59c63ea Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Thu, 9 Jul 2026 16:22:37 +0800 Subject: [PATCH] [AMD] Enable JIT staged HiCache write-back and fix CPU-index crash (#28534) Co-authored-by: Duyi-Wang --- .../jit_kernel/csrc/kvcacheio/hicache.cuh | 62 ++++- .../csrc/kvcacheio/staged_write_back.cuh | 16 +- .../jit_kernel/include/sgl_kernel/utils.cuh | 2 + .../sglang/srt/managers/cache_controller.py | 9 +- .../sglang/srt/mem_cache/memory_pool_host.py | 12 +- python/sglang/srt/mem_cache/pool_host/mha.py | 12 +- .../jit/test_hicache_page_first_write_back.py | 259 ++++++++++++++++++ 7 files changed, 349 insertions(+), 23 deletions(-) create mode 100644 test/registered/jit/test_hicache_page_first_write_back.py diff --git a/python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh b/python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh index 7abb5ab19..ee16fba0a 100644 --- a/python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh +++ b/python/sglang/jit_kernel/csrc/kvcacheio/hicache.cuh @@ -37,44 +37,86 @@ inline constexpr auto get_mem_package() { template using PackageType = decltype(get_mem_package()); +// 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(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(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(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(dst)); +#endif } } // namespace details @@ -256,18 +298,18 @@ struct HiCacheKernel { TensorMatcher({-1, D}) // .with_strides({N, 1}) .with_dtype(cache_dtype) - .with_device() + .with_device() .verify(k_cache_src) .verify(v_cache_src); TensorMatcher({-1, D}) // .with_strides({M, 1}) .with_dtype(cache_dtype) - .with_device() + .with_device() .verify(k_cache_dst) .verify(v_cache_dst); TensorMatcher({L}) // .with_dtype(indices_dtype) - .with_device(indices_device) + .with_device(indices_device) .verify(indices_src) .verify(indices_dst); @@ -323,14 +365,14 @@ struct HiCacheKernel { TensorMatcher({N}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(k_ptr_src) .verify(v_ptr_src) .verify(k_ptr_dst) .verify(v_ptr_dst); TensorMatcher({L}) // .with_dtype(dtype_) - .with_device(device_) + .with_device(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() + .with_device() .verify(cache_src); TensorMatcher({-1, D}) // .with_strides({M, 1}) .with_dtype(cache_dtype) - .with_device() + .with_device() .verify(cache_dst); TensorMatcher({L}) // .with_dtype(indices_dtype) - .with_device(indices_device) + .with_device(indices_device) .verify(indices_src) .verify(indices_dst); @@ -441,12 +483,12 @@ struct HiCacheKernel { TensorMatcher({N}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(ptr_src) .verify(ptr_dst); TensorMatcher({L}) // .with_dtype(dtype_) - .with_device(device_) + .with_device(device_) .verify(indices_src) .verify(indices_dst); diff --git a/python/sglang/jit_kernel/csrc/kvcacheio/staged_write_back.cuh b/python/sglang/jit_kernel/csrc/kvcacheio/staged_write_back.cuh index 0e0a08715..38074167e 100644 --- a/python/sglang/jit_kernel/csrc/kvcacheio/staged_write_back.cuh +++ b/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(device_) + .with_device(device_) .verify(staging_k); if constexpr (!kIsMLA) { TensorMatcher({T, N, D}) // .with_dtype(cache_dtype) - .with_device(device_) + .with_device(device_) .verify(staging_v); } TensorMatcher({-1, N, D}) // .with_dtype(cache_dtype) - .with_device() + .with_device() .verify(k_cache_dst); if constexpr (!kIsMLA) { TensorMatcher({-1, N, D}) // .with_dtype(cache_dtype) - .with_device() + .with_device() .verify(v_cache_dst); } TensorMatcher({N}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(k_ptr_src); if constexpr (!kIsMLA) { TensorMatcher({N}) // .with_dtype() - .with_device(device_) + .with_device(device_) .verify(v_ptr_src); } TensorMatcher({P}) // .with_dtype(indices_dtype) - .with_device(device_) + .with_device(device_) .verify(page_indices_src); TensorMatcher({T}) // .with_dtype(dst_indices_dtype) - .with_device() + .with_device() .verify(dst_indices_cpu); RuntimeCheck(page_size > 0, "HiCache staged relayout: page_size must be positive"); diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh index 25cc89aab..bd2d9ab52 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh +++ b/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 { diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index 29f497c54..f23edf60b 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/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" diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 0ce1e059d..43b6af776 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index af28daeae..4017f4da3 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/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: diff --git a/test/registered/jit/test_hicache_page_first_write_back.py b/test/registered/jit/test_hicache_page_first_write_back.py new file mode 100644 index 000000000..55dd6ca82 --- /dev/null +++ b/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"]))