From c192145830bf240f3fd3eedda6b6c111474f0d8b Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 30 Jul 2026 00:26:10 -0700 Subject: [PATCH] [Kernel] Fuse KV-cache writes for asymmetric K/V (head_dim != v_head_dim) (#32813) --- .../kernels/jit/csrc/elementwise/kvcache.cuh | 185 ++++++++++++---- python/sglang/kernels/ops/kvcache/kvcache.py | 43 ++-- python/sglang/srt/mem_cache/memory_pool.py | 28 ++- .../benchmark/kvcache/bench_store_cache.py | 28 +++ .../kernels/ops/kvcache/test_store_cache.py | 87 ++++++++ .../mem_cache/test_asymmetric_mha_pool.py | 202 ++++++++++++++++++ 6 files changed, 511 insertions(+), 62 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_asymmetric_mha_pool.py diff --git a/python/sglang/kernels/jit/csrc/elementwise/kvcache.cuh b/python/sglang/kernels/jit/csrc/elementwise/kvcache.cuh index 6be2f04c2..2f8be6aae 100644 --- a/python/sglang/kernels/jit/csrc/elementwise/kvcache.cuh +++ b/python/sglang/kernels/jit/csrc/elementwise/kvcache.cuh @@ -21,7 +21,9 @@ struct StoreKVCacheParams { const void* __restrict__ indices; int64_t stride_k_bytes; int64_t stride_v_bytes; - int64_t stride_cache_bytes; + // Independent slot strides: head_dim != v_head_dim gives K and V different row widths. + int64_t stride_k_cache_bytes; + int64_t stride_v_cache_bytes; int64_t stride_indices; uint32_t batch_size; int64_t size_limit; @@ -31,9 +33,37 @@ struct StoreKVCacheParams { constexpr uint32_t kNumWarps = 4; constexpr uint32_t kThreadsPerBlock = kNumWarps * device::kWarpThreads; +/** + * \brief How a warp vectorizes one row of kElementBytes: the widest aligned + * vector type it can use, and how many full loop iterations that takes. + * Shared by the interleaved and single-row copies so the two cannot drift. + * kElementBytes == 0 is a valid (empty) plan, so a zero-width tail can be + * queried before being branched away. + */ +template +struct RowVecPlan { + static constexpr int64_t kAlignment = (kElementBytes % (16 * device::kWarpThreads) == 0) ? 16 + : kElementBytes % (8 * device::kWarpThreads) == 0 ? 8 + : kElementBytes % (4 * device::kWarpThreads) == 0 ? 4 + : kElementBytes % 4 == 0 ? 4 + : 0; + + static_assert(kAlignment > 0, "Element size must be multiple of 4 bytes"); + + using vec_t = device::AlignedStorage; + static constexpr int64_t kLoopBytes = sizeof(vec_t) * device::kWarpThreads; + static constexpr int64_t kLoopCount = kElementBytes / kLoopBytes; + static constexpr int64_t kElementCount = kElementBytes / sizeof(vec_t); + static constexpr bool kHasEpilogue = kLoopCount * kLoopBytes < kElementBytes; +}; + /** * \brief Use a single warp to copy key and value data from source to destination. * Each thread in the warp copies a portion of the data in a coalesced manner. + * Both loads are issued before either store: the two rows live in different + * tensors, and the params' __restrict__ does not survive into the kernel body, + * so the compiler cannot prove k_dst and v_src disjoint and will not sink the + * V load past the K store on its own. * \tparam kElementBytes The size of each key/value element in bytes. * \param k_src Pointer to the source key data. * \param v_src Pointer to the source value data. @@ -47,17 +77,9 @@ SGL_DEVICE void copy_kv_warp( void* __restrict__ k_dst, void* __restrict__ v_dst) { using namespace device; - constexpr int64_t kAlignment = (kElementBytes % (16 * kWarpThreads) == 0) ? 16 - : kElementBytes % (8 * kWarpThreads) == 0 ? 8 - : kElementBytes % (4 * kWarpThreads) == 0 ? 4 - : kElementBytes % 4 == 0 ? 4 - : 0; - - static_assert(kAlignment > 0, "Element size must be multiple of 4 bytes"); - - using vec_t = AlignedStorage; - constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads; - constexpr auto kLoopCount = kElementBytes / kLoopBytes; + using plan_t = RowVecPlan; + using vec_t = typename plan_t::vec_t; + constexpr auto kLoopCount = plan_t::kLoopCount; const auto gmem = tile::Memory::warp(); @@ -70,8 +92,8 @@ SGL_DEVICE void copy_kv_warp( } // handle the epilogue if any - if constexpr (kLoopCount * kLoopBytes < kElementBytes) { - if (gmem.in_bound(kElementBytes / sizeof(vec_t), kLoopCount)) { + if constexpr (plan_t::kHasEpilogue) { + if (gmem.in_bound(plan_t::kElementCount, kLoopCount)) { const auto k = gmem.load(k_src, kLoopCount); const auto v = gmem.load(v_src, kLoopCount); gmem.store(k_dst, k, kLoopCount); @@ -80,24 +102,100 @@ SGL_DEVICE void copy_kv_warp( } } +/** + * \brief Use a single warp to copy one row from source to destination. + * Serves the width by which asymmetric K/V rows differ, which has no counterpart + * row to interleave with. + * \tparam kElementBytes The size of the row in bytes. + * \param src Pointer to the source data. + * \param dst Pointer to the destination data. + */ +template +SGL_DEVICE void copy_row_warp(const void* __restrict__ src, void* __restrict__ dst) { + using namespace device; + using plan_t = RowVecPlan; + using vec_t = typename plan_t::vec_t; + constexpr auto kLoopCount = plan_t::kLoopCount; + + const auto gmem = tile::Memory::warp(); + +#pragma unroll kLoopCount + for (int64_t i = 0; i < kLoopCount; ++i) { + gmem.store(dst, gmem.load(src, i), i); + } + + // handle the epilogue if any + if constexpr (plan_t::kHasEpilogue) { + if (gmem.in_bound(plan_t::kElementCount, kLoopCount)) { + gmem.store(dst, gmem.load(src, kLoopCount), kLoopCount); + } + } +} + +/** + * \brief Copy a K row of kKBytes and a V row of kVBytes with one warp. + * The overlapping prefix goes through the interleaved copy; only the width by + * which the rows differ is left as a serial tail. Equal widths degenerate to a + * single interleaved copy with no tail. + */ +template +SGL_DEVICE void copy_kv_rows_warp( + const void* __restrict__ k_src, + const void* __restrict__ v_src, + void* __restrict__ k_dst, + void* __restrict__ v_dst) { + using namespace device; + constexpr auto kCommon = kKBytes < kVBytes ? kKBytes : kVBytes; + constexpr auto kTail = (kKBytes < kVBytes ? kVBytes : kKBytes) - kCommon; + + // The interleaved copy indexes BOTH rows with kCommon's vector width, so that + // width must divide each row's split offset -- the narrower row's alignment + // does not imply the wider one's (e.g. 512 picks 16B, but 516 is not 16B + // aligned). The tail's own width must likewise divide its kCommon start. + // Whatever these gates admit is alignment-safe for the strides too, since a + // stride is a whole multiple of its split size. + constexpr auto kTailOrCommon = kTail == 0 ? kCommon : kTail; + constexpr auto kCommonAlign = RowVecPlan::kAlignment; + constexpr auto kTailAlign = RowVecPlan::kAlignment; + constexpr bool kCanInterleave = + kKBytes % kCommonAlign == 0 && kVBytes % kCommonAlign == 0 && kCommon % kTailAlign == 0; + + if constexpr (kCanInterleave) { + copy_kv_warp(k_src, v_src, k_dst, v_dst); + if constexpr (kTail > 0) { + if constexpr (kKBytes > kVBytes) { + copy_row_warp(pointer::offset(k_src, kCommon), pointer::offset(k_dst, kCommon)); + } else { + copy_row_warp(pointer::offset(v_src, kCommon), pointer::offset(v_dst, kCommon)); + } + } + } else { + copy_row_warp(k_src, k_dst); + copy_row_warp(v_src, v_dst); + } +} + /** * \brief Kernel to store key-value pairs into the KV cache. * Each element is split into multiple parts to allow parallel memory copy. - * \tparam kElementBytes The size of each key/value element in bytes. + * \tparam kKElementBytes The size of each key element in bytes. + * \tparam kVElementBytes The size of each value element in bytes. Differs from + * kKElementBytes for asymmetric KV (head_dim != v_head_dim). * \tparam kSplit The number of warps that handle each element. * \tparam kUsePDL Whether to use PDL feature. * \tparam T The data type of the indices (`int32_t` or `int64_t`). */ -template +template __global__ void store_kvcache(const __grid_constant__ StoreKVCacheParams params) { using namespace device; - constexpr auto kSplitSize = kElementBytes / kSplit; + constexpr auto kKSplitSize = kKElementBytes / kSplit; + constexpr auto kVSplitSize = kVElementBytes / kSplit; const uint32_t warp_id = blockIdx.x * kNumWarps + threadIdx.x / kWarpThreads; const uint32_t item_id = warp_id / kSplit; const uint32_t split_id = warp_id % kSplit; const auto& [ k_input, v_input, k_cache, v_cache, indices, // ptr - stride_k, stride_v, stride_cache, stride_indices, batch_size, // size + stride_k, stride_v, stride_k_cache, stride_v_cache, stride_indices, batch_size, // size size_limit, reserved_skip_index // bounds and reserved sink ] = params; if (item_id >= batch_size) return; @@ -109,36 +207,37 @@ __global__ void store_kvcache(const __grid_constant__ StoreKVCacheParams params) // A stale/OOB slot id would cause an illegal memory access in the store below; // fail fast at the culprit instead. always-on (kvcache JIT compiles without NDEBUG). assert(index >= 0 && index < size_limit); - const auto k_src = pointer::offset(k_input, item_id * stride_k, split_id * kSplitSize); - const auto v_src = pointer::offset(v_input, item_id * stride_v, split_id * kSplitSize); - const auto k_dst = pointer::offset(k_cache, index * stride_cache, split_id * kSplitSize); - const auto v_dst = pointer::offset(v_cache, index * stride_cache, split_id * kSplitSize); + const auto k_src = pointer::offset(k_input, item_id * stride_k, split_id * kKSplitSize); + const auto v_src = pointer::offset(v_input, item_id * stride_v, split_id * kVSplitSize); + const auto k_dst = pointer::offset(k_cache, index * stride_k_cache, split_id * kKSplitSize); + const auto v_dst = pointer::offset(v_cache, index * stride_v_cache, split_id * kVSplitSize); if (index != reserved_skip_index) { - copy_kv_warp(k_src, v_src, k_dst, v_dst); + copy_kv_rows_warp(k_src, v_src, k_dst, v_dst); } PDLTriggerSecondary(); } -template +template struct StoreKVCacheKernel { - static_assert(kElementBytes > 0 && kElementBytes % 4 == 0); + static_assert(kKElementBytes > 0 && kKElementBytes % 4 == 0); + static_assert(kVElementBytes > 0 && kVElementBytes % 4 == 0); template - static constexpr auto store_kernel = store_kvcache; + static constexpr auto store_kernel = store_kvcache; template static auto get_kernel(const int num_split) { using namespace host; - // only apply split optimization when element size is aligned - if constexpr (kElementBytes % (4 * 128) == 0) { + // only apply split optimization when both element sizes are aligned + if constexpr (kKElementBytes % (4 * 128) == 0 && kVElementBytes % (4 * 128) == 0) { if (num_split == 4) return store_kernel<4, T>; } - if constexpr (kElementBytes % (2 * 128) == 0) { + if constexpr (kKElementBytes % (2 * 128) == 0 && kVElementBytes % (2 * 128) == 0) { if (num_split == 2) return store_kernel<2, T>; } if (num_split == 1) return store_kernel<1, T>; - Panic("Unsupported num_split {} for element size {}", num_split, kElementBytes); + Panic("Unsupported num_split {} for element sizes k={} v={}", num_split, kKElementBytes, kVElementBytes); } static void @@ -152,31 +251,37 @@ struct StoreKVCacheKernel { const int64_t reserved_skip_index) { using namespace host; auto B = SymbolicSize{"batch_size"}; - auto D = SymbolicSize{"element_size"}; + auto DK = SymbolicSize{"k_element_size"}; + auto DV = SymbolicSize{"v_element_size"}; auto KS = SymbolicSize{"k_stride"}; auto VS = SymbolicSize{"v_stride"}; - auto S = SymbolicSize{"cache_stride"}; + auto SK = SymbolicSize{"k_cache_stride"}; + auto SV = SymbolicSize{"v_cache_stride"}; auto I = SymbolicSize{"indices_stride"}; auto dtype = SymbolicDType{}; auto device = SymbolicDevice{}; auto indice_dtype = SymbolicDType{}; device.set_options(); - TensorMatcher({B, D}) // + TensorMatcher({B, DK}) // .with_strides({KS, 1}) .with_dtype(dtype) .with_device(device) .verify(k); - TensorMatcher({B, D}) // + TensorMatcher({B, DV}) // .with_strides({VS, 1}) .with_dtype(dtype) .with_device(device) .verify(v); - TensorMatcher({-1, D}) // - .with_strides({S, 1}) + TensorMatcher({-1, DK}) // + .with_strides({SK, 1}) + .with_dtype(dtype) + .with_device(device) + .verify(k_cache); + TensorMatcher({-1, DV}) // + .with_strides({SV, 1}) .with_dtype(dtype) .with_device(device) - .verify(k_cache) .verify(v_cache); TensorMatcher({B}) // .with_strides({I}) @@ -186,7 +291,8 @@ struct StoreKVCacheKernel { const int64_t dtype_size = dtype_bytes(dtype.unwrap()); const uint32_t num_elements = static_cast(B.unwrap()); - RuntimeCheck(kElementBytes == dtype_size * D.unwrap()); + RuntimeCheck(kKElementBytes == dtype_size * DK.unwrap()); + RuntimeCheck(kVElementBytes == dtype_size * DV.unwrap()); const auto params = StoreKVCacheParams{ .k = k.data_ptr(), @@ -196,7 +302,8 @@ struct StoreKVCacheKernel { .indices = indices.data_ptr(), .stride_k_bytes = KS.unwrap() * dtype_size, .stride_v_bytes = VS.unwrap() * dtype_size, - .stride_cache_bytes = S.unwrap() * dtype_size, + .stride_k_cache_bytes = SK.unwrap() * dtype_size, + .stride_v_cache_bytes = SV.unwrap() * dtype_size, .stride_indices = I.unwrap(), .batch_size = static_cast(B.unwrap()), .size_limit = size_limit, diff --git a/python/sglang/kernels/ops/kvcache/kvcache.py b/python/sglang/kernels/ops/kvcache/kvcache.py index f4534d5e3..98a012f6b 100644 --- a/python/sglang/kernels/ops/kvcache/kvcache.py +++ b/python/sglang/kernels/ops/kvcache/kvcache.py @@ -18,8 +18,8 @@ if TYPE_CHECKING: @cache_once -def _jit_kvcache_module(row_bytes: int) -> Module: - args = make_cpp_args(row_bytes, is_arch_support_pdl()) +def _jit_kvcache_module(k_row_bytes: int, v_row_bytes: int) -> Module: + args = make_cpp_args(k_row_bytes, v_row_bytes, is_arch_support_pdl()) return load_jit( "kvcache", *args, @@ -29,20 +29,25 @@ def _jit_kvcache_module(row_bytes: int) -> Module: @cache_once -def can_use_store_cache(size: int) -> bool: +def can_use_store_cache(k_row_bytes: int, v_row_bytes: int = 0) -> bool: + """Whether the JIT store_cache kernel can serve these row widths. + v_row_bytes=0 means symmetric, i.e. it defaults to k_row_bytes.""" logger = logging.getLogger(__name__) - if size % 4 != 0: - logger.warning( - f"Unsupported row_bytes={size} for JIT KV-Cache kernel:" - " must be multiple of 4" - ) - return False + v_row_bytes = v_row_bytes or k_row_bytes + for name, size in (("k_row_bytes", k_row_bytes), ("v_row_bytes", v_row_bytes)): + if size % 4 != 0: + logger.warning( + f"Unsupported {name}={size} for JIT KV-Cache kernel:" + " must be multiple of 4" + ) + return False try: - _jit_kvcache_module(size) + _jit_kvcache_module(k_row_bytes, v_row_bytes) return True except Exception as e: logger.warning( - f"Failed to load JIT KV-Cache kernel " f"with row_bytes={size}: {e}" + f"Failed to load JIT KV-Cache kernel with " + f"k_row_bytes={k_row_bytes} v_row_bytes={v_row_bytes}: {e}" ) return False @@ -56,6 +61,7 @@ def store_cache( indices: torch.Tensor, *, row_bytes: int = 0, + v_row_bytes: int = 0, num_split: int = 0, # can be tuned for performance size_limit: int = 0, reserved_skip_index: int = 0, @@ -64,10 +70,13 @@ def store_cache( Args: k (torch.Tensor): Key tensor of shape (batch_size, H * D). - v (torch.Tensor): Value tensor of shape (batch_size, H * D). + v (torch.Tensor): Value tensor of shape (batch_size, H * Dv). k_cache (torch.Tensor): Key cache tensor of shape (num_pages, H * D). - v_cache (torch.Tensor): Value cache tensor of shape (num_pages, H * D). + v_cache (torch.Tensor): Value cache tensor of shape (num_pages, H * Dv). indices (torch.Tensor): Indices tensor of shape (batch_size,). + row_bytes (int): Key row width in bytes. Inferred from k when 0. + v_row_bytes (int): Value row width in bytes; differs from row_bytes for + asymmetric KV (head_dim != v_head_dim). Inferred from v when 0. size_limit (int): Valid slot bound (cache row count = real slots + the reserved padding slot); an index outside [0, size_limit) fails fast (device assert) instead of an illegal memory access. Defaults to the @@ -77,11 +86,13 @@ def store_cache( pass -1 to disable skipping. """ row_bytes = row_bytes or k.shape[-1] * k.element_size() - module = _jit_kvcache_module(row_bytes) + v_row_bytes = v_row_bytes or v.shape[-1] * v.element_size() + module = _jit_kvcache_module(row_bytes, v_row_bytes) if num_split <= 0: - if row_bytes % 2048 == 0: + # A split must divide BOTH rows, so require the alignment on each. + if row_bytes % 2048 == 0 and v_row_bytes % 2048 == 0: num_split = 4 - elif row_bytes % 1024 == 0: + elif row_bytes % 1024 == 0 and v_row_bytes % 1024 == 0: num_split = 2 else: num_split = 1 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 000491f84..66ef8e472 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -149,21 +149,26 @@ def _set_kv_buffer_impl( device_module: Any, size_limit: int, alt_stream: Optional[torch.cuda.Stream] = None, - same_kv_dim: bool = True, + v_row_dim: Optional[int] = None, # head_num * v_head_dim; defaults to row_dim ) -> None: + v_row_dim = row_dim if v_row_dim is None else v_row_dim row_bytes = row_dim * store_dtype.itemsize - if (_is_cuda or _is_hip) and same_kv_dim and can_use_store_cache(row_bytes): + v_row_bytes = v_row_dim * store_dtype.itemsize + if (_is_cuda or _is_hip) and can_use_store_cache(row_bytes, v_row_bytes): return store_cache( k.view(-1, row_dim), - v.view(-1, row_dim), + v.view(-1, v_row_dim), k_cache.view(-1, row_dim), - v_cache.view(-1, row_dim), + v_cache.view(-1, v_row_dim), indices, row_bytes=row_bytes, + v_row_bytes=v_row_bytes, size_limit=size_limit, ) - if _is_cpu and _cpu_has_amx_support: + # store_cache_cpu takes a single row_dim for both K and V, so it only serves + # equal-width rows; asymmetric KV falls through to the naive path below. + if _is_cpu and _cpu_has_amx_support and v_row_dim == row_dim: return torch.ops.sgl_kernel.store_cache_cpu( k, v, @@ -1819,7 +1824,7 @@ class MHATokenToKVPool(KVCache): # for store_cache JIT kernel self.row_dim = self.head_num * self.head_dim - self.same_kv_dim = self.head_dim == self.v_head_dim + self.v_row_dim = self.head_num * self.v_head_dim def _init_kv_copy_and_warmup(self): # Zero-layer pool (e.g. all-SWA model's full sub-pool) has no buffers. @@ -2386,7 +2391,7 @@ class MHATokenToKVPool(KVCache): # dummy tokens write there); valid index range is [0, size + page_size). size_limit=self.size + self.page_size, alt_stream=self.alt_stream, - same_kv_dim=self.same_kv_dim, + v_row_dim=self.v_row_dim, ) def _quantized_scales(self, global_layer_id: int, k_scale, v_scale): @@ -2712,6 +2717,15 @@ class MHATokenToKVPool(KVCache): ) return + # The tiled kernel takes one ROW_BYTES for both tensors, so an asymmetric V + # row would be written at K's width and bleed into the next slot. Only this + # path needs the gate; the non-CUDA branch above handles both widths. + if self.v_row_dim != self.row_dim: + raise NotImplementedError( + "prefix-valid commit requires equal-width K/V rows, got " + f"head_dim={self.head_dim} v_head_dim={self.v_head_dim}." + ) + _set_kv_buffer_prefix_valid_impl( cache_k, cache_v, diff --git a/test/registered/kernels/benchmark/kvcache/bench_store_cache.py b/test/registered/kernels/benchmark/kvcache/bench_store_cache.py index 94564441b..ec9f78725 100644 --- a/test/registered/kernels/benchmark/kvcache/bench_store_cache.py +++ b/test/registered/kernels/benchmark/kvcache/bench_store_cache.py @@ -72,5 +72,33 @@ def benchmark(batch_size: int, item_size: int, impl: str): ) +# Asymmetric K/V (head_dim != v_head_dim). The item_size sweep above drives both +# rows from one value, so it never reaches the split-prefix-plus-tail path. +# 192/128 and 384/256 both go live in a single MiMoV2 TP=4 deployment, whose +# layers carry either 1 or 2 kv heads per rank; the reversed and wide pairs cover +# a V-side tail and a num_split > 1 shape. +ASYM_ITEM_SIZES = [(192, 128), (384, 256), (128, 192), (1024, 512)] + + +@marker.parametrize("k_item,v_item", ASYM_ITEM_SIZES, [(192, 128), (1024, 512)]) +@marker.parametrize("batch_size", [2**n for n in range(0, 15)], [16]) +@marker.benchmark("impl", ["jit", "torch_compile", "torch_streams"]) +def benchmark_asymmetric(batch_size: int, k_item: int, v_item: int, impl: str): + torch.manual_seed(42) + k = create_random(batch_size, k_item) + k_cache = create_empty(CACHE_SIZE, k_item) + v = create_random(batch_size, v_item) + v_cache = create_empty(CACHE_SIZE, v_item) + indices = torch.randperm(CACHE_SIZE, device=DEFAULT_DEVICE)[:batch_size] + return marker.do_bench( + FN_MAP[impl], + input_args=(k, v, k_cache, v_cache, indices), + graph_clone_args=(0, 1, 4), # not need to clone cache, which is large + memory_args=(k, v, indices), # k_cache / v_cache excluded + memory_output=(k, v), # inplace write, size = k + v + ) + + if __name__ == "__main__": benchmark.run() + benchmark_asymmetric.run() diff --git a/test/registered/kernels/ops/kvcache/test_store_cache.py b/test/registered/kernels/ops/kvcache/test_store_cache.py index 957185ae4..a804a1ca6 100644 --- a/test/registered/kernels/ops/kvcache/test_store_cache.py +++ b/test/registered/kernels/ops/kvcache/test_store_cache.py @@ -174,11 +174,98 @@ def test_store_cache_num_split( assert torch.all(v_cache[indices] == v) +# Asymmetric K/V (head_dim != v_head_dim): different row widths AND cache strides. +# MiMoV2 is 192/128. Both orderings, since nothing may assume K is the wider one. +ASYM_DIM_PAIRS = get_ci_test_range( + [(192, 128), (128, 192), (1024, 512), (512, 1024), (96, 64), (2048, 1024)], + [(192, 128), (512, 1024)], +) + + +# The kernel is a byte copier specialized on (k_row_bytes, v_row_bytes) -- no dtype +# in its template args -- so equal-itemsize dtypes share one instantiation. bf16 and +# fp32 are the two distinct itemsizes; fp16 would just re-run the bf16 one. +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("k_dim,v_dim", ASYM_DIM_PAIRS) +def test_store_cache_asymmetric(k_dim: int, v_dim: int, dtype: torch.dtype) -> None: + batch_size = 128 + k = torch.randn((batch_size, k_dim), dtype=dtype, device=DEVICE) + v = torch.randn((batch_size, v_dim), dtype=dtype, device=DEVICE) + k_cache = torch.randn((SMALL_CACHE, k_dim), dtype=dtype, device=DEVICE) + v_cache = torch.randn((SMALL_CACHE, v_dim), dtype=dtype, device=DEVICE) + k_before, v_before = k_cache.clone(), v_cache.clone() + indices = torch.randperm(SMALL_CACHE - 1, device=DEVICE)[:batch_size] + 1 + + store_cache(k, v, k_cache, v_cache, indices) + + assert torch.all(k_cache[indices] == k) + assert torch.all(v_cache[indices] == v) + # Applying K's stride to V (or vice versa) would corrupt neighbouring slots, + # which the target-slot assertions above cannot see. + untouched = torch.ones(SMALL_CACHE, dtype=torch.bool, device=DEVICE) + untouched[indices] = False + assert torch.all(k_cache[untouched] == k_before[untouched]) + assert torch.all(v_cache[untouched] == v_before[untouched]) + + +def _valid_asym_num_splits(k_dim: int, v_dim: int, dtype: torch.dtype) -> list: + """num_split values valid for BOTH rows; a split must divide each of them.""" + k_bytes, v_bytes = k_dim * dtype.itemsize, v_dim * dtype.itemsize + splits = [1] + if k_bytes % (2 * 128) == 0 and v_bytes % (2 * 128) == 0: + splits.append(2) + if k_bytes % (4 * 128) == 0 and v_bytes % (4 * 128) == 0: + splits.append(4) + return splits + + +def _default_num_split(k_dim: int, v_dim: int, dtype: torch.dtype) -> int: + """Mirrors the heuristic in store_cache(); the default is already exercised + by test_store_cache_asymmetric, which does not pass num_split.""" + k_bytes, v_bytes = k_dim * dtype.itemsize, v_dim * dtype.itemsize + if k_bytes % 2048 == 0 and v_bytes % 2048 == 0: + return 4 + if k_bytes % 1024 == 0 and v_bytes % 1024 == 0: + return 2 + return 1 + + +# Only splits the default heuristic would NOT pick: the split gate is two-sided +# (K and V must both align), so the off-default branches are what needs pinning. +_ASYM_NUM_SPLIT_CASES = [ + (_k, _v, _ns) + for _k, _v in ASYM_DIM_PAIRS + for _ns in _valid_asym_num_splits(_k, _v, DTYPE) + if _ns != _default_num_split(_k, _v, DTYPE) +] + + +@pytest.mark.parametrize("k_dim,v_dim,num_split", _ASYM_NUM_SPLIT_CASES) +def test_store_cache_asymmetric_num_split( + k_dim: int, v_dim: int, num_split: int +) -> None: + batch_size = 128 + k = torch.randn((batch_size, k_dim), dtype=DTYPE, device=DEVICE) + v = torch.randn((batch_size, v_dim), dtype=DTYPE, device=DEVICE) + k_cache = torch.randn((SMALL_CACHE, k_dim), dtype=DTYPE, device=DEVICE) + v_cache = torch.randn((SMALL_CACHE, v_dim), dtype=DTYPE, device=DEVICE) + indices = torch.randperm(SMALL_CACHE - 1, device=DEVICE)[:batch_size] + 1 + + store_cache(k, v, k_cache, v_cache, indices, num_split=num_split) + + assert torch.all(k_cache[indices] == k) + assert torch.all(v_cache[indices] == v) + + def test_can_use_store_cache() -> None: assert can_use_store_cache(128) assert can_use_store_cache(256) assert can_use_store_cache(1024) assert can_use_store_cache(2048) + # asymmetric widths, and the documented default (v falls back to k) + assert can_use_store_cache(384, 256) + assert can_use_store_cache(256, 384) + assert can_use_store_cache(1024, 0) == can_use_store_cache(1024) if __name__ == "__main__": diff --git a/test/registered/unit/mem_cache/test_asymmetric_mha_pool.py b/test/registered/unit/mem_cache/test_asymmetric_mha_pool.py new file mode 100644 index 000000000..4f79fed3b --- /dev/null +++ b/test/registered/unit/mem_cache/test_asymmetric_mha_pool.py @@ -0,0 +1,202 @@ +"""Device-side tests for MHATokenToKVPool with asymmetric KV (head_dim != v_head_dim). + +Covers the wiring the kernel-level tests cannot see: that the pool derives +``v_row_dim`` from ``v_head_dim`` and threads it into the fused store_cache kernel. +A mis-wired ``v_row_dim`` still writes the right bytes into the right K slots, so +the untouched-slot assertions are what pin the V width and stride down. + +Skipped on CPU -- the fused path is CUDA/HIP only. + + python -m pytest test/registered/unit/mem_cache/test_asymmetric_mha_pool.py -v +""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.kernels.ops.kvcache.kvcache import can_use_store_cache +from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool +from sglang.test.ci.ci_register import register_cuda_ci + +_HAS_CUDA = torch.cuda.is_available() + +register_cuda_ci(est_time=20, stage="base-b", runner_config="1-gpu-small") + +DTYPE = torch.bfloat16 +HEAD_NUM = 2 +POOL_SIZE = 63 # buffers get POOL_SIZE + page_size rows +NUM_WRITES = 16 + +# (head_dim, v_head_dim). Both orderings, since nothing may assume K is wider. +# The last pair is wide enough for the split heuristic to pick num_split=2. +ASYM_DIM_PAIRS = [(192, 128), (128, 192), (512, 256)] + + +def _build_pool(head_dim: int, v_head_dim: int) -> MHATokenToKVPool: + return MHATokenToKVPool( + size=POOL_SIZE, + page_size=1, + dtype=DTYPE, + head_num=HEAD_NUM, + head_dim=head_dim, + v_head_dim=v_head_dim, + layer_num=1, + device="cuda", + enable_memory_saver=False, + enable_alt_stream=False, + ) + + +@unittest.skipUnless(_HAS_CUDA, "fused store_cache path requires CUDA") +class TestAsymmetricMHAPoolRowDims(unittest.TestCase): + def test_v_row_dim_tracks_v_head_dim(self): + for head_dim, v_head_dim in ASYM_DIM_PAIRS: + with self.subTest(head_dim=head_dim, v_head_dim=v_head_dim): + pool = _build_pool(head_dim, v_head_dim) + self.assertEqual(pool.row_dim, HEAD_NUM * head_dim) + self.assertEqual(pool.v_row_dim, HEAD_NUM * v_head_dim) + + def test_v_row_dim_defaults_to_row_dim_when_symmetric(self): + pool = _build_pool(128, 128) + self.assertEqual(pool.v_row_dim, pool.row_dim) + + def test_swa_dims_override_row_dims(self): + # A hybrid sliding-window model builds a second pool through the swa_* + # parameters, which override head_num/head_dim/v_head_dim wholesale. Both + # of MiMoV2's pools are asymmetric, so v_row_dim has to follow + # swa_v_head_dim rather than the full pool's v_head_dim. + pool = MHATokenToKVPool( + size=POOL_SIZE, + page_size=1, + dtype=DTYPE, + head_num=HEAD_NUM, + head_dim=512, + v_head_dim=256, + swa_head_num=1, + swa_head_dim=192, + swa_v_head_dim=128, + layer_num=1, + device="cuda", + enable_memory_saver=False, + enable_alt_stream=False, + ) + self.assertEqual(pool.row_dim, 1 * 192) + self.assertEqual(pool.v_row_dim, 1 * 128) + + def test_swa_v_head_dim_falls_back_to_v_head_dim(self): + # swa_v_head_dim omitted: head_dim comes from swa_head_dim but v_head_dim + # does not, so the two are read from different sources. Pinned because a + # pool built this way is asymmetric in a way neither config states. + pool = MHATokenToKVPool( + size=POOL_SIZE, + page_size=1, + dtype=DTYPE, + head_num=HEAD_NUM, + head_dim=512, + v_head_dim=256, + swa_head_num=1, + swa_head_dim=192, + layer_num=1, + device="cuda", + enable_memory_saver=False, + enable_alt_stream=False, + ) + self.assertEqual(pool.row_dim, 1 * 192) + self.assertEqual(pool.v_row_dim, 1 * 256) + + +@unittest.skipUnless(_HAS_CUDA, "fused store_cache path requires CUDA") +class TestAsymmetricMHAPoolSetKVBuffer(unittest.TestCase): + """set_kv_buffer round-trip through the fused kernel, per dim pair.""" + + def _run_roundtrip(self, head_dim: int, v_head_dim: int): + pool = _build_pool(head_dim, v_head_dim) + k_buf, v_buf = pool.k_buffer[0], pool.v_buffer[0] + self.assertEqual(tuple(k_buf.shape[1:]), (HEAD_NUM, head_dim)) + self.assertEqual(tuple(v_buf.shape[1:]), (HEAD_NUM, v_head_dim)) + + itemsize = pool.store_dtype.itemsize + self.assertTrue( + can_use_store_cache(pool.row_dim * itemsize, pool.v_row_dim * itemsize), + "fused store_cache unavailable; the naive fallback is also correct, so " + "this test would pass without covering anything", + ) + + # Seed every slot so an over-wide V write shows up on a slot never targeted. + k_buf.copy_(torch.randn_like(k_buf)) + v_buf.copy_(torch.randn_like(v_buf)) + k_before, v_before = k_buf.clone(), v_buf.clone() + + # Slot 0 is the reserved padding slot store_cache skips; target [1, num_slots). + num_slots = k_buf.shape[0] + loc = torch.randperm(num_slots - 1, device="cuda")[:NUM_WRITES] + 1 + cache_k = torch.randn( + (NUM_WRITES, HEAD_NUM, head_dim), dtype=DTYPE, device="cuda" + ) + cache_v = torch.randn( + (NUM_WRITES, HEAD_NUM, v_head_dim), dtype=DTYPE, device="cuda" + ) + + pool.set_kv_buffer(SimpleNamespace(layer_id=0), loc, cache_k, cache_v) + + self.assertTrue(torch.equal(k_buf[loc], cache_k), "K target slots") + self.assertTrue(torch.equal(v_buf[loc], cache_v), "V target slots") + + untouched = torch.ones(num_slots, dtype=torch.bool, device="cuda") + untouched[loc] = False + self.assertTrue( + torch.equal(k_buf[untouched], k_before[untouched]), + "K bled outside its target slots", + ) + self.assertTrue( + torch.equal(v_buf[untouched], v_before[untouched]), + "V bled outside its target slots (wrong row width or stride)", + ) + + def test_asymmetric_roundtrip(self): + for head_dim, v_head_dim in ASYM_DIM_PAIRS: + with self.subTest(head_dim=head_dim, v_head_dim=v_head_dim): + self._run_roundtrip(head_dim, v_head_dim) + + def test_symmetric_roundtrip_unchanged(self): + self._run_roundtrip(128, 128) + + +@unittest.skipUnless(_HAS_CUDA, "prefix-valid tiled kernel requires CUDA") +class TestAsymmetricPrefixValidGuard(unittest.TestCase): + """set_kv_buffer_prefix_valid's tiled kernel takes one row width for both + tensors, so it must refuse asymmetric KV rather than truncate V.""" + + def _call_prefix_valid(self, pool, head_dim, v_head_dim): + rows = 2 + loc_2d = torch.tensor([[1, 2]], dtype=torch.int64, device="cuda") + commit_lens = torch.tensor([rows], dtype=torch.int32, device="cuda") + cache_k = torch.randn((rows, HEAD_NUM, head_dim), dtype=DTYPE, device="cuda") + cache_v = torch.randn((rows, HEAD_NUM, v_head_dim), dtype=DTYPE, device="cuda") + pool.set_kv_buffer_prefix_valid( + SimpleNamespace(layer_id=0, k_scale=None, v_scale=None), + loc_2d, + commit_lens, + cache_k, + cache_v, + ) + + def test_rejects_asymmetric(self): + for head_dim, v_head_dim in ASYM_DIM_PAIRS: + with self.subTest(head_dim=head_dim, v_head_dim=v_head_dim): + pool = _build_pool(head_dim, v_head_dim) + with self.assertRaises(NotImplementedError): + self._call_prefix_valid(pool, head_dim, v_head_dim) + + def test_accepts_symmetric(self): + # The guard must not tighten the equal-width path it already served. + pool = _build_pool(128, 128) + self._call_prefix_valid(pool, 128, 128) + expected = torch.arange(1, 3, device="cuda") + self.assertTrue(torch.any(pool.k_buffer[0][expected] != 0)) + self.assertTrue(torch.any(pool.v_buffer[0][expected] != 0)) + + +if __name__ == "__main__": + unittest.main()