diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh index 2b5228e9b..aa0038bca 100644 --- a/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh +++ b/python/sglang/kernels/jit/csrc/kvcacheio/hicache.cuh @@ -17,6 +17,20 @@ namespace sglang { namespace device { +// Logical threads collaborating on one copied element. This is not the +// hardware warp/wavefront size: on CDNA wave64, one wavefront contains two +// logically independent 32-thread copy groups. The transfer kernels use no shuffle, +// ballot, barrier, shared memory, or other cross-lane communication. +inline constexpr uint32_t kCopyGroupThreads = 32; + +template +inline constexpr uint32_t copy_lanes_per_worker() { + static_assert(kUnroll > 0, "unroll must be positive"); + static_assert(kUnroll <= kCopyGroupThreads, "unroll cannot exceed the logical copy-group width"); + static_assert(kCopyGroupThreads % kUnroll == 0, "unroll must divide the logical copy-group width"); + return kCopyGroupThreads / kUnroll; +} + namespace details { template @@ -40,6 +54,31 @@ inline constexpr auto get_mem_package() { template using PackageType = decltype(get_mem_package()); +// A worker copies one element in rounds of `group` bytes, each lane moving +// group / lanes_per_worker bytes as one vector package. That quotient has to +// be a package size the hardware supports. +inline constexpr bool group_fits(int64_t bytes, uint32_t lanes_per_worker, uint32_t group) { + if (group % lanes_per_worker != 0 || bytes % static_cast(group) != 0) { + return false; + } + const uint32_t package = group / lanes_per_worker; + return package == 4 || package == 8 || package == 16; +} + +inline constexpr uint32_t pick_group_bytes(int64_t bytes, uint32_t lanes_per_worker) { + // The narrow rounds only pay off against the raised ROCm block quota, so CUDA + // keeps the original 128 B requirement and generates the same code as before. +#ifdef USE_ROCM + return group_fits(bytes, lanes_per_worker, 128) ? 128u + : group_fits(bytes, lanes_per_worker, 64) ? 64u + : group_fits(bytes, lanes_per_worker, 32) ? 32u + : group_fits(bytes, lanes_per_worker, 16) ? 16u + : 0u; +#else + return group_fits(bytes, lanes_per_worker, 128) ? 128u : 0u; +#endif +} + // 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 @@ -124,40 +163,40 @@ SGL_DEVICE void store_nc(uint4* __restrict__ dst, const uint4& value) { } // namespace details -template +template SGL_DEVICE auto load_vec(const void* __restrict__ src) { - static_assert(kBytes % 128 == 0, "kBytes must be multiple of 128 bytes"); - static_assert(128 % kNumThreads == 0, "kNumThreads must divide 128 bytes"); - constexpr uint32_t kLoopCount = kBytes / 128; - using Package = details::PackageType<128 / kNumThreads>; + constexpr uint32_t kGroupBytes = details::pick_group_bytes(kBytes, kLanesPerWorker); + static_assert(kGroupBytes != 0, "no 4/8/16 B package tiles kBytes across the worker lanes"); + constexpr uint32_t kLoopCount = kBytes / kGroupBytes; + using Package = details::PackageType; using Storage = details::LocalStorage; const auto src_packed = static_cast(src); - const auto lane_id = threadIdx.x % kNumThreads; + const auto lane_id = threadIdx.x % kLanesPerWorker; Storage vec; #pragma unroll kLoopCount for (uint32_t i = 0; i < kLoopCount; ++i) { - const auto j = i * kNumThreads + lane_id; + const auto j = i * kLanesPerWorker + lane_id; vec.data[i] = details::load_nc(&src_packed[j]); } return vec; } -template +template SGL_DEVICE void store_vec(void* __restrict__ dst, const Storage& vec) { using Package = std::decay_t; - constexpr uint32_t kBytesPerLoop = sizeof(Package) * kNumThreads; + constexpr uint32_t kBytesPerLoop = sizeof(Package) * kLanesPerWorker; constexpr uint32_t kLoopCount = kBytes / kBytesPerLoop; static_assert(kBytes % kBytesPerLoop == 0, "Invalid Storage configuration"); const auto dst_packed = static_cast(dst); - const auto lane_id = threadIdx.x % kNumThreads; + const auto lane_id = threadIdx.x % kLanesPerWorker; #pragma unroll kLoopCount for (uint32_t i = 0; i < kLoopCount; ++i) { - const auto j = i * kNumThreads + lane_id; + const auto j = i * kLanesPerWorker + lane_id; details::store_nc(&dst_packed[j], vec.data[i]); } } @@ -188,11 +227,10 @@ template < bool kIsMLA = false> SGL_HICACHE_KERNEL void hicache_transfer_per_layer(const __grid_constant__ HicacheKernelParams params) { using namespace device; - static_assert(kBlockSize % kWarpThreads == 0); - static_assert(kWarpThreads % kUnroll == 0); + static_assert(kBlockSize % kCopyGroupThreads == 0); - constexpr uint32_t kNumThreads = kWarpThreads / kUnroll; - constexpr uint32_t kWorkersPerBlock = kBlockSize / kNumThreads; + constexpr uint32_t kLanesPerWorker = copy_lanes_per_worker(); + constexpr uint32_t kWorkersPerBlock = kBlockSize / kLanesPerWorker; constexpr uint32_t kNumWorkers = kWorkersPerBlock * kBlockQuota; const auto& [ @@ -201,24 +239,24 @@ SGL_HICACHE_KERNEL void hicache_transfer_per_layer(const __grid_constant__ Hicac kv_cache_src_stride, kv_cache_dst_stride, length, _ // metadata ] = params; - const uint32_t work_id = blockIdx.x * kWorkersPerBlock + threadIdx.x / kNumThreads; + const uint32_t work_id = blockIdx.x * kWorkersPerBlock + threadIdx.x / kLanesPerWorker; for (uint32_t i = work_id; i < length; i += kNumWorkers) { const auto pos_src = static_cast(indices_src)[i]; const auto pos_dst = static_cast(indices_dst)[i]; const auto src_k = pointer::offset(k_cache_src, pos_src * kv_cache_src_stride); const auto dst_k = pointer::offset(k_cache_dst, pos_dst * kv_cache_dst_stride); - const auto vec_k = load_vec(src_k); + const auto vec_k = load_vec(src_k); // Both loads are issued before either store: the compiler cannot prove // dst_k and src_v disjoint, so it will not hoist the V load on its own. std::decay_t vec_v; if constexpr (!kIsMLA) { const auto src_v = pointer::offset(v_cache_src, pos_src * kv_cache_src_stride); - vec_v = load_vec(src_v); + vec_v = load_vec(src_v); } - store_vec(dst_k, vec_k); + store_vec(dst_k, vec_k); if constexpr (!kIsMLA) { const auto dst_v = pointer::offset(v_cache_dst, pos_dst * kv_cache_dst_stride); - store_vec(dst_v, vec_v); + store_vec(dst_v, vec_v); } } } @@ -235,11 +273,10 @@ SGL_HICACHE_KERNEL void hicache_transfer_all_layer(const __grid_constant__ Hicac using src_ptr_t = const void*; using dst_ptr_t = void*; - static_assert(kBlockSize % kWarpThreads == 0); - static_assert(kWarpThreads % kUnroll == 0); + static_assert(kBlockSize % kCopyGroupThreads == 0); - constexpr uint32_t kNumThreads = kWarpThreads / kUnroll; - constexpr uint32_t kWorkersPerBlock = kBlockSize / kNumThreads; + constexpr uint32_t kLanesPerWorker = copy_lanes_per_worker(); + constexpr uint32_t kWorkersPerBlock = kBlockSize / kLanesPerWorker; constexpr uint32_t kNumWorkers = kWorkersPerBlock * kBlockQuota; const auto& [ @@ -248,7 +285,7 @@ SGL_HICACHE_KERNEL void hicache_transfer_all_layer(const __grid_constant__ Hicac kv_cache_src_stride, kv_cache_dst_stride, length, num_layers // metadata ] = params; - const uint32_t work_id = blockIdx.x * kWorkersPerBlock + threadIdx.x / kNumThreads; + const uint32_t work_id = blockIdx.x * kWorkersPerBlock + threadIdx.x / kLanesPerWorker; for (uint32_t i = work_id; i < length; i += kNumWorkers) { const auto pos_src = static_cast(indices_src)[i]; const auto pos_dst = static_cast(indices_dst)[i]; @@ -257,20 +294,20 @@ SGL_HICACHE_KERNEL void hicache_transfer_all_layer(const __grid_constant__ Hicac const auto k_cache_dst = static_cast(k_ptr_dst)[layer]; const auto src_k = pointer::offset(k_cache_src, pos_src * kv_cache_src_stride); const auto dst_k = pointer::offset(k_cache_dst, pos_dst * kv_cache_dst_stride); - const auto vec_k = load_vec(src_k); + const auto vec_k = load_vec(src_k); // Both loads are issued before either store: the compiler cannot prove // dst_k and src_v disjoint, so it will not hoist the V load on its own. std::decay_t vec_v; if constexpr (!kIsMLA) { const auto v_cache_src = static_cast(v_ptr_src)[layer]; const auto src_v = pointer::offset(v_cache_src, pos_src * kv_cache_src_stride); - vec_v = load_vec(src_v); + vec_v = load_vec(src_v); } - store_vec(dst_k, vec_k); + store_vec(dst_k, vec_k); if constexpr (!kIsMLA) { const auto v_cache_dst = static_cast(v_ptr_dst)[layer]; const auto dst_v = pointer::offset(v_cache_dst, pos_dst * kv_cache_dst_stride); - store_vec(dst_v, vec_v); + store_vec(dst_v, vec_v); } } } @@ -341,7 +378,7 @@ struct HiCacheKernel { const auto kv_cache_dst_stride = static_cast(M.unwrap() * dtype_size); const auto use_int32 = indices_dtype.unwrap().bits == 32; - constexpr auto kWorkersPerBlock = kBlockSize / (device::kWarpThreads / kUnroll); + constexpr auto kWorkersPerBlock = kBlockSize / device::copy_lanes_per_worker(); const auto num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota); const auto params = HicacheKernelParams{ .k_cache_dst = k_cache_dst_ptr, @@ -398,7 +435,7 @@ struct HiCacheKernel { const auto use_int32 = dtype_.unwrap().bits == 32; const auto device = device_.unwrap(); - constexpr auto kWorkersPerBlock = kBlockSize / (device::kWarpThreads / kUnroll); + constexpr auto kWorkersPerBlock = kBlockSize / device::copy_lanes_per_worker(); const auto num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota); const auto params = HicacheKernelParams{ .k_cache_dst = k_cache_dst_ptr, @@ -461,7 +498,7 @@ struct HiCacheKernel { const auto cache_dst_stride = static_cast(M.unwrap() * dtype_size); const auto use_int32 = indices_dtype.unwrap().bits == 32; - constexpr auto kWorkersPerBlock = kBlockSize / (device::kWarpThreads / kUnroll); + constexpr auto kWorkersPerBlock = kBlockSize / device::copy_lanes_per_worker(); const auto num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota); const auto params = HicacheKernelParams{ .k_cache_dst = cache_dst_ptr, @@ -511,7 +548,7 @@ struct HiCacheKernel { const auto use_int32 = dtype_.unwrap().bits == 32; const auto device = device_.unwrap(); - constexpr auto kWorkersPerBlock = kBlockSize / (device::kWarpThreads / kUnroll); + constexpr auto kWorkersPerBlock = kBlockSize / device::copy_lanes_per_worker(); const auto num_blocks = std::min(div_ceil(length, kWorkersPerBlock), kBlockQuota); const auto params = HicacheKernelParams{ .k_cache_dst = cache_dst_ptr, diff --git a/python/sglang/kernels/ops/kvcache/hicache.py b/python/sglang/kernels/ops/kvcache/hicache.py index 6844b34e0..8e454390a 100644 --- a/python/sglang/kernels/ops/kvcache/hicache.py +++ b/python/sglang/kernels/ops/kvcache/hicache.py @@ -20,6 +20,14 @@ _is_hip = is_hip_runtime() # ROCm needs a wider grid to saturate mapped-host transfers; CUDA keeps the legacy quota. DEFAULT_BLOCK_QUOTA = 32 if _is_hip else 2 +# Logical copy-group width; this is not the hardware warp/wavefront size. +COPY_GROUP_THREADS = 32 + +# Copy-round widths, widest first. The narrow rounds admit element sizes 128 +# does not divide, such as MLA's 576 B fp8 row, but only pay off against the +# ROCm quota above, so CUDA keeps the original 128 B requirement. +GROUP_BYTES = (128, 64, 32, 16) if _is_hip else (128,) + @cache_once def _jit_hicache_module(*, element_size: int, unroll: int, block_quota: int) -> Module: @@ -80,11 +88,11 @@ def can_use_hicache_jit_kernel( block_quota: int | None = None, # can be tuned for less interference ) -> bool: logger = logging.getLogger(__name__) - if element_size % 128 != 0: + unroll = unroll or _default_unroll(element_size) + if not _tiles_across_lanes(element_size, unroll): logger.warning(f"Unsupported {element_size = } for JIT HiCache kernel") return False try: - unroll = unroll or _default_unroll(element_size) block_quota = block_quota or DEFAULT_BLOCK_QUOTA _jit_hicache_module( element_size=element_size, @@ -121,6 +129,20 @@ def can_use_write_back_jit_kernel( return False +def _tiles_across_lanes(element_size: int, unroll: int) -> bool: + """Mirror of pick_group_bytes() in kvcacheio/hicache.cuh.""" + if unroll <= 0 or unroll > COPY_GROUP_THREADS or COPY_GROUP_THREADS % unroll != 0: + return False + + lanes_per_worker = COPY_GROUP_THREADS // unroll + return any( + group % lanes_per_worker == 0 + and element_size % group == 0 + and group // lanes_per_worker in (4, 8, 16) + for group in GROUP_BYTES + ) + + def _default_unroll(element_size: int) -> int: if element_size <= 512: return 4 diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index 74ddba1ab..03e1629b1 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/python/sglang/srt/mem_cache/pool_host/mha.py @@ -782,7 +782,7 @@ class MHATokenToKOnlyPoolHost(HostKVCache): self.lock = threading.RLock() self.clear() - self.can_use_jit = _is_cuda and can_use_hicache_jit_kernel( + self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel( element_size=self.token_stride_size ) self.k_device_ptrs = torch.tensor( diff --git a/test/registered/kernels/ops/kvcache/test_hicache.py b/test/registered/kernels/ops/kvcache/test_hicache.py index 8dbf47cbe..2bc89d94f 100644 --- a/test/registered/kernels/ops/kvcache/test_hicache.py +++ b/test/registered/kernels/ops/kvcache/test_hicache.py @@ -12,9 +12,10 @@ from sglang.srt.mem_cache.pool_host.common import ( from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu -from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci register_cuda_ci(est_time=12, stage="base-b", runner_config="1-gpu-large") +register_amd_ci(est_time=12, stage="jit-kernel-unit", runner_config="amd") pytestmark = pytest.mark.skipif( not torch.cuda.is_available() diff --git a/test/registered/unit/mem_cache/test_hicache_copy_rounds.py b/test/registered/unit/mem_cache/test_hicache_copy_rounds.py new file mode 100644 index 000000000..4893ddc05 --- /dev/null +++ b/test/registered/unit/mem_cache/test_hicache_copy_rounds.py @@ -0,0 +1,70 @@ +"""The copy-round screen in kvcache/hicache.py must agree with pick_group_bytes() +in kvcacheio/hicache.cuh: a size the screen admits has to compile, and a size the +kernel cannot tile has to be turned away before the JIT ever sees it.""" + +import unittest +from unittest import mock + +from sglang.kernels.ops.kvcache import hicache +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _screen(element_size: int, unroll: int, *, is_hip: bool) -> bool: + """Run _tiles_across_lanes() as the given platform would see it.""" + group_bytes = (128, 64, 32, 16) if is_hip else (128,) + with mock.patch.object(hicache, "GROUP_BYTES", group_bytes): + return hicache._tiles_across_lanes(element_size, unroll) + + +class TestHiCacheCopyRounds(unittest.TestCase): + def test_mla_fp8_row_is_admitted_only_on_rocm(self): + # 576 B is MLA's fp8 row and the reason the narrow rounds exist: 128 does + # not divide it, so the 128-only CUDA screen has to keep turning it away. + unroll = hicache._default_unroll(576) + self.assertEqual(unroll, 2) + self.assertTrue(_screen(576, unroll=unroll, is_hip=True)) + self.assertFalse(_screen(576, unroll=unroll, is_hip=False)) + + def test_a_size_128_divides_is_admitted_everywhere(self): + self.assertTrue(_screen(512, unroll=4, is_hip=True)) + self.assertTrue(_screen(512, unroll=4, is_hip=False)) + + def test_rocm_admits_no_more_than_the_kernel_can_tile(self): + # The screen mirrors group_fits(): the round must divide the element and + # split across lanes into a package the hardware has (4, 8 or 16 B). + for element_size in range(16, 1300, 4): + for unroll in (1, 2, 4, 8, 16, 32): + lanes_per_worker = hicache.COPY_GROUP_THREADS // unroll + expected = any( + group % lanes_per_worker == 0 + and element_size % group == 0 + and group // lanes_per_worker in (4, 8, 16) + for group in (128, 64, 32, 16) + ) + with self.subTest(element_size=element_size, unroll=unroll): + self.assertEqual( + _screen(element_size, unroll, is_hip=True), expected + ) + + def test_an_odd_size_is_turned_away_on_both(self): + # 100 B is divisible by no round, so no lane split can cover it. + self.assertFalse(_screen(100, unroll=4, is_hip=True)) + self.assertFalse(_screen(100, unroll=4, is_hip=False)) + + def test_invalid_unroll_is_turned_away(self): + for unroll in (-1, 0, 3, 5, 7, 33, 64): + with self.subTest(unroll=unroll): + self.assertFalse(_screen(576, unroll=unroll, is_hip=True)) + + def test_default_unrolls_use_expected_logical_worker_widths(self): + cases = ((128, 8), (576, 16), (1152, 32)) + for element_size, expected_lanes in cases: + with self.subTest(element_size=element_size): + unroll = hicache._default_unroll(element_size) + self.assertEqual(hicache.COPY_GROUP_THREADS // unroll, expected_lanes) + + +if __name__ == "__main__": + unittest.main()