[Kernel] Fuse KV-cache writes for asymmetric K/V (head_dim != v_head_dim) (#32813)

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 <int64_t kElementBytes>
+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<uint32_t, kAlignment / 4>;
+  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
+  using plan_t = RowVecPlan<kElementBytes>;
-                                 : kElementBytes % (8 * kWarpThreads) == 0  ? 8
+  using vec_t = typename plan_t::vec_t;
-                                 : kElementBytes % (4 * kWarpThreads) == 0  ? 4
+  constexpr auto kLoopCount = plan_t::kLoopCount;
-                                 : kElementBytes % 4 == 0                   ? 4
-                                                                            : 0;
-
-  static_assert(kAlignment > 0, "Element size must be multiple of 4 bytes");
-
-  using vec_t = AlignedStorage<uint32_t, kAlignment / 4>;
-  constexpr auto kLoopBytes = sizeof(vec_t) * kWarpThreads;
-  constexpr auto kLoopCount = kElementBytes / kLoopBytes;
 
   const auto gmem = tile::Memory<vec_t>::warp();
 
@@ -70,8 +92,8 @@ SGL_DEVICE void copy_kv_warp(
   }
 
   // handle the epilogue if any
-  if constexpr (kLoopCount * kLoopBytes < kElementBytes) {
+  if constexpr (plan_t::kHasEpilogue) {
-    if (gmem.in_bound(kElementBytes / sizeof(vec_t), kLoopCount)) {
+    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 <int64_t kElementBytes>
+SGL_DEVICE void copy_row_warp(const void* __restrict__ src, void* __restrict__ dst) {
+  using namespace device;
+  using plan_t = RowVecPlan<kElementBytes>;
+  using vec_t = typename plan_t::vec_t;
+  constexpr auto kLoopCount = plan_t::kLoopCount;
+
+  const auto gmem = tile::Memory<vec_t>::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 <int64_t kKBytes, int64_t kVBytes>
+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<kCommon>::kAlignment;
+  constexpr auto kTailAlign = RowVecPlan<kTailOrCommon>::kAlignment;
+  constexpr bool kCanInterleave =
+      kKBytes % kCommonAlign == 0 && kVBytes % kCommonAlign == 0 && kCommon % kTailAlign == 0;
+
+  if constexpr (kCanInterleave) {
+    copy_kv_warp<kCommon>(k_src, v_src, k_dst, v_dst);
+    if constexpr (kTail > 0) {
+      if constexpr (kKBytes > kVBytes) {
+        copy_row_warp<kTail>(pointer::offset(k_src, kCommon), pointer::offset(k_dst, kCommon));
+      } else {
+        copy_row_warp<kTail>(pointer::offset(v_src, kCommon), pointer::offset(v_dst, kCommon));
+      }
+    }
+  } else {
+    copy_row_warp<kKBytes>(k_src, k_dst);
+    copy_row_warp<kVBytes>(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 <int64_t kElementBytes, int kSplit, bool kUsePDL, typename T>
+template <int64_t kKElementBytes, int64_t kVElementBytes, int kSplit, bool kUsePDL, typename T>
 __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 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 * kSplitSize);
+  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_cache, split_id * kSplitSize);
+  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_cache, split_id * kSplitSize);
+  const auto v_dst = pointer::offset(v_cache, index * stride_v_cache, split_id * kVSplitSize);
 
   if (index != reserved_skip_index) {
-    copy_kv_warp<kSplitSize>(k_src, v_src, k_dst, v_dst);
+    copy_kv_rows_warp<kKSplitSize, kVSplitSize>(k_src, v_src, k_dst, v_dst);
   }
   PDLTriggerSecondary<kUsePDL>();
 }
 
-template <int64_t kElementBytes, bool kUsePDL>
+template <int64_t kKElementBytes, int64_t kVElementBytes, bool kUsePDL>
 struct StoreKVCacheKernel {
-  static_assert(kElementBytes > 0 && kElementBytes % 4 == 0);
+  static_assert(kKElementBytes > 0 && kKElementBytes % 4 == 0);
+  static_assert(kVElementBytes > 0 && kVElementBytes % 4 == 0);
 
   template <int kSplit, typename T>
-  static constexpr auto store_kernel = store_kvcache<kElementBytes, kSplit, kUsePDL, T>;
+  static constexpr auto store_kernel = store_kvcache<kKElementBytes, kVElementBytes, kSplit, kUsePDL, T>;
 
   template <typename T>
   static auto get_kernel(const int num_split) {
     using namespace host;
-    // only apply split optimization when element size is aligned
+    // only apply split optimization when both element sizes are aligned
-    if constexpr (kElementBytes % (4 * 128) == 0) {
+    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<kDLCUDA, kDLROCM>();
 
-    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})  //
+    TensorMatcher({-1, DK})  //
-        .with_strides({S, 1})
+        .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<uint32_t>(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<uint32_t>(B.unwrap()),
         .size_limit = size_limit,

python/sglang/kernels/ops/kvcache/kvcache.py

@@ -18,8 +18,8 @@ if TYPE_CHECKING:
 
 
 @cache_once
-def _jit_kvcache_module(row_bytes: int) -> Module:
+def _jit_kvcache_module(k_row_bytes: int, v_row_bytes: int) -> Module:
-    args = make_cpp_args(row_bytes, is_arch_support_pdl())
+    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:
+    v_row_bytes = v_row_bytes or k_row_bytes
-        logger.warning(
+    for name, size in (("k_row_bytes", k_row_bytes), ("v_row_bytes", v_row_bytes)):
-            f"Unsupported row_bytes={size} for JIT KV-Cache kernel:"
+        if size % 4 != 0:
-            " must be multiple of 4"
+            logger.warning(
-        )
+                f"Unsupported {name}={size} for JIT KV-Cache kernel:"
-        return False
+                " 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

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,

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()

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__":

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()