Fix sparse MLA cache layout validation and eight-head decode

csrc/sparse_mla_sm120.cu

@@ -69,8 +69,8 @@ inline ModelType resolve_model_type(int d_qk, int64_t model_type) {
   if (d_qk == 512) {
     const auto mt = static_cast<ModelType>(
         model_type == kAuto ? static_cast<int64_t>(ModelType::DSV4) : model_type);
-    TVM_FFI_ICHECK(mt == ModelType::DSV4 || mt == ModelType::GLM53_NOPE)
-        << "d_qk=512 supports model_type auto, DSV4, or GLM53_NOPE; got " << model_type;
+    TVM_FFI_ICHECK(mt == ModelType::DSV4 || mt == ModelType::GLM53_NOPE || mt == ModelType::DSV4_1)
+        << "d_qk=512 supports model_type auto, DSV4, GLM53_NOPE, or DSV4_1; got " << model_type;
     return mt;
   }
   if (d_qk == 1088) {
@@ -81,8 +81,8 @@ inline ModelType resolve_model_type(int d_qk, int64_t model_type) {
     return mt;
   }
   TVM_FFI_ICHECK(false) << "Unsupported d_qk=" << d_qk
-                        << "; expected 576 (DSV3_2/GLM_NSA), 512 (DSV4/GLM53_NOPE) or 1088 "
-                           "(DOTS3_SWA)";
+                        << "; expected 576 (DSV3_2/GLM_NSA), 512 (DSV4/GLM53_NOPE/DSV4_1) "
+                           "or 1088 (DOTS3_SWA)";
   return ModelType::DSV4;
 }
 
@@ -106,18 +106,27 @@ struct PagedKVLayout {
 inline PagedKVLayout parse_paged_kv_layout(const TensorView& kv, int bpt, bool inline_scale,
                                            const char* name) {
   const size_t elem_bytes = static_cast<size_t>(kv.dtype().bits / 8);
+  TVM_FFI_ICHECK_EQ(kv.stride(-1), 1) << name << " last dim must be contiguous";
+  // Bulk gathers require aligned addresses, including sliced cache origins
+  // and the first row of each block.
+  TVM_FFI_ICHECK_EQ(reinterpret_cast<uintptr_t>(kv.data_ptr()) % 16, 0)
+      << name << " data pointer must be 16B-aligned (cp.async.bulk requirement)";
+  const size_t block_stride = static_cast<size_t>(kv.stride(0)) * elem_bytes;
+  TVM_FFI_ICHECK_EQ(block_stride % 16, 0)
+      << name << " block stride must be 16B-aligned (cp.async.bulk requirement)";
   if (kv.ndim() == 2) {
     const size_t block_bytes = static_cast<size_t>(kv.size(1)) * elem_bytes;
     TVM_FFI_ICHECK_EQ(block_bytes % static_cast<size_t>(bpt), 0)
         << name << " 2D block width " << block_bytes
         << " is not divisible by bytes_per_token=" << bpt;
     // A flat 2D block carries no row padding to infer, so the row advance is
     // exactly bytes_per_token.
-    return {static_cast<int>(block_bytes / static_cast<size_t>(bpt)), block_bytes,
+    TVM_FFI_ICHECK_GE(block_stride, block_bytes)
+        << name << " block stride is smaller than the packed block width";
+    return {static_cast<int>(block_bytes / static_cast<size_t>(bpt)), block_stride,
             static_cast<size_t>(bpt)};
   }
   auto row_advance = [&](int64_t token_axis) {
-    TVM_FFI_ICHECK_EQ(kv.stride(-1), 1) << name << " last dim must be contiguous";
     const size_t bytes = static_cast<size_t>(kv.size(-1)) * elem_bytes;
     if (inline_scale) {
       TVM_FFI_ICHECK_GE(bytes, static_cast<size_t>(bpt))
@@ -136,6 +145,9 @@ inline PagedKVLayout parse_paged_kv_layout(const TensorView& kv, int bpt, bool i
     if (inline_scale) {
       TVM_FFI_ICHECK_EQ(advance % 16, 0) << name << " token-axis stride " << advance
                                          << " is not 16B-aligned (cp.async.bulk requirement)";
+    } else {
+      TVM_FFI_ICHECK_EQ(advance, static_cast<size_t>(bpt))
+          << name << " footer-scale rows must stay packed (token-axis stride == " << bpt << ")";
     }
     return advance;
   };
@@ -359,6 +371,9 @@ void SparseMlaSm120PagedAttention(
     case ModelType::DOTS3_SWA:
       mt_name = "DOTS3_SWA";
       break;
+    case ModelType::DSV4_1:
+      mt_name = "DSV4_1";
+      break;
     case ModelType::DSV4:
       break;
   }

csrc/sparse_mla_sm120_decode_dsv3_2.cu

@@ -9,7 +9,7 @@
 // Supports the V32-family dispatch grid: dedicated instantiations at
 //   num_heads ∈ {8, 16, 32, 64, 128}
 // plus one runtime-H instantiation (any num_heads <= 128 off the grid) and
-// GLM53_NOPE dedicated 32/64 + runtime-H. topk is a runtime argument — one
+// GLM53_NOPE dedicated 8/32/64 + runtime-H. topk is a runtime argument — one
 // instantiation serves every indices-row width.
 
 #include <cuda_runtime.h>
@@ -189,6 +189,9 @@ bool launch_sparse_mla_decode_dsv3_2(ModelType mt, int num_heads, int topk, int
   // indexer window. The TP1 (64-head) and TP2 (32-head) shapes keep
   // dedicated instantiations; any other shard rides the runtime-H fallback.
   if (mt == ModelType::GLM53_NOPE) {
+    // The public scratch allocator uses eight rows for H=8. A dedicated
+    // instantiation preserves that ABI instead of the runtime-H padded stride.
+    DSV3_2_DISPATCH_MT(ModelType::GLM53_NOPE, 8)
     DSV3_2_DISPATCH_MT(ModelType::GLM53_NOPE, 32)
     DSV3_2_DISPATCH_MT(ModelType::GLM53_NOPE, 64)
     DSV3_2_DISPATCH_RT_MT(ModelType::GLM53_NOPE)

csrc/sparse_mla_sm120_decode_dsv4.cu

@@ -43,26 +43,26 @@ static bool launch_decode_dsv4_impl(int num_heads, int topk, const bf16* Q, cons
   // Dynamic smem layout (FP8 XV, double-buffered KV). Measured on sm_120
   // against a 101376 B per-block opt-in cap:
   //
-  //   term                                        DSV4       DOTS3_SWA
-  //                                            (BI=64,W=8)  (BI=32,W=4)
-  //   sm_q_rope    HPB * D_ROPE * 2B               2048         2048
-  //   sm_q_fp8     HPB * Q_NOPE_STRIDE             7424        16640
-  //   sm_q_sc      HPB * NUM_SCALES * 4B            448          512
-  //   sm_kv_fp8    2 * BI * KV_SMEM_STRIDE        59392        66560
-  //   sm_kv_sc     2 * BI * SCALE_BYTES_PER_TOKEN  1024          512
-  //   sm_kv_rope   2 * BI * D_ROPE * 2B           16384         8192
-  //   mbar + pad                                     48           48
-  //   sm_reduce    2 * N_WARPS * HPB * 4           1024          512
-  //   sm_w_head_sc N_V_CHUNKS * HPB * 4             448          512
-  //   sm_w_fp8 x2  2 * HPB * (BI + 16)             2560         1536
-  //   dynamic total                               90800        97072
+  //   term                                        DSV4       DOTS3_SWA   DSV4_1
+  //                                            (BI=64,W=8)  (BI=32,W=4)  (BI=64,W=8)
+  //   sm_q_rope    HPB * D_ROPE * 2B               2048         2048          0
+  //   sm_q_fp8     HPB * Q_NOPE_STRIDE             7424        16640       8448
+  //   sm_q_sc      HPB * NUM_SCALES * 4B            448          512       1024
+  //   sm_kv_fp8    2 * BI * KV_SMEM_STRIDE        59392        66560      67584
+  //   sm_kv_sc     2 * BI * SCALE_BYTES_PER_TOKEN  1024          512       2048
+  //   sm_kv_rope   2 * BI * D_ROPE * 2B           16384         8192          0
+  //   mbar + pad                                     48           48         48
+  //   sm_reduce    2 * N_WARPS * HPB * 4           1024          512       1024
+  //   sm_w_head_sc N_V_CHUNKS * HPB * 4             448          512       1024
+  //   sm_w_fp8     2 * XV_FOLD * HPB * (BI + 16)   2560         1536       5120
+  //   dynamic total                               90800        97072      86296
   // Static smem (kernel-side), sm_p_full = HPB * BI * 2B:
-  //   DSV4 2048 B; DOTS3_SWA 0 (V_HAS_ROPE=false makes the bf16 P dead).
-  //   grand total                                 92848        97072
+  //   DSV4 2048 B; DOTS3_SWA/DSV4_1 0 (V_HAS_ROPE=false makes the bf16 P dead).
+  //   grand total                                 92848        97072      86296
   //
   // DOTS3_SWA leaves ~4.2 KB spare. BI=64 for it needs 173872 B and the driver
-  // rejects the opt-in outright. Both configs run 1 block/SM.
-  constexpr int N_V_CHUNKS_LAUNCH = KV::D_NOPE / KV::QUANT_TILE;  // DSV4 7, DOTS3_SWA 8
+  // rejects the opt-in outright. All configs run 1 block/SM.
+  constexpr int N_V_CHUNKS_LAUNCH = KV::D_NOPE / KV::QUANT_TILE;  // DSV4 7, DOTS3_SWA 8, DSV4_1 16
   constexpr int DYN_SMEM_BYTES =
       HPB * KV::D_ROPE * (int)sizeof(bf16)                            // sm_q_rope
       + HPB * KV::Q_NOPE_STRIDE                                       // sm_q_fp8
@@ -74,7 +74,7 @@ static bool launch_decode_dsv4_impl(int num_heads, int topk, const bf16* Q, cons
       + 4 * (int)sizeof(uint64_t)                                     // mbar_full+empty
       + 2 * Cfg::N_WARPS * HPB * (int)sizeof(float)                   // sm_reduce
       + N_V_CHUNKS_LAUNCH * HPB * (int)sizeof(float)                  // sm_w_head_sc
-      + 2 * HPB * (Cfg::BI + 16);                                     // sm_w_fp8 ×2 (vc parity)
+      + 2 * Cfg::XV_FOLD * HPB * (Cfg::BI + 16);  // sm_w_fp8 ×2 parities × XV_FOLD
 
   auto kernel = sparse_mla_decode_dsv4_kernel<MT, NUM_HEADS, PAGE_BLOCK_SIZE>;
   CUDA_CHECK_BOOL(
@@ -165,7 +165,9 @@ bool launch_sparse_mla_decode_dsv4(
     int extra_topk, int pbs_extra, size_t stride_extra_kv_block, int chunks_per_block_override,
     float sm_scale, size_t stride_kv_block, size_t stride_indices_token,
     size_t stride_extra_indices_token, size_t stride_out_lse, cudaStream_t stream) {
-  if (mt != ModelType::DSV4 && mt != ModelType::DOTS3_SWA) return false;
+  if (mt != ModelType::DSV4 && mt != ModelType::DOTS3_SWA && mt != ModelType::DSV4_1) {
+    return false;
+  }
   // DOTS3_SWA has no dual-cache instantiation; the planner never routes one
   // here, and the launcher rejects it so a direct FFI caller cannot silently
   // run an untested path.
@@ -219,6 +221,16 @@ bool launch_sparse_mla_decode_dsv4(
   DECODE_DISPATCH(ModelType::DOTS3_SWA, 32)
   DECODE_DISPATCH(ModelType::DOTS3_SWA, 64)
   DECODE_DISPATCH_RT(ModelType::DOTS3_SWA)
+  // DSV4_1 (DeepSeek-V4.1): all-FP8 512-wide K, 16B UE8M0 footer. Same dual-
+  // cache capability as DSV4 (vLLM routes the SWA cache as main + compressed
+  // as extra); the 32-wide quant groups run the pair-folded XV
+  // (DecodeTileCfg::XV_FOLD=2) on the standard 8-warp tile.
+  DECODE_DISPATCH(ModelType::DSV4_1, 8)
+  DECODE_DISPATCH(ModelType::DSV4_1, 16)
+  DECODE_DISPATCH(ModelType::DSV4_1, 32)
+  DECODE_DISPATCH(ModelType::DSV4_1, 64)
+  DECODE_DISPATCH(ModelType::DSV4_1, 128)
+  DECODE_DISPATCH_RT(ModelType::DSV4_1)
 #undef DECODE_DISPATCH_RT
 #undef DECODE_DISPATCH
   return false;

csrc/sparse_mla_sm120_jit_binding.cu

@@ -53,25 +53,32 @@ struct PagedKVLayout {
   int stride_kv_row;
 };
 
-// inline_scale: the model stores scales inside the row (DSV3_2 / GLM_NSA /
-// GLM53_NOPE) and gathers whole rows with cp.async.bulk, so the row advance
-// must be 16B-aligned. Footer-scale models (DSV4 / DOTS3_SWA) address data
-// rows by the packed data stride and skip the check (584 % 16 != 0 is legal
-// there).
+// Inline-scale rows may be padded, with a 16B-aligned row advance.
+// Footer-scale caches must keep data and scale sections packed. Their data
+// rows use an aligned stride (e.g. 576B for DSV4), separate from the total
+// payload per token (584B including footer scales). Both families require
+// 16B-aligned cache origins and block strides for cp.async.bulk.
 inline PagedKVLayout parse_paged_kv_layout(const TensorView& kv, int bpt, bool inline_scale,
                                            const char* name) {
   const size_t elem_bytes = static_cast<size_t>(kv.dtype().bits / 8);
+  TVM_FFI_ICHECK_EQ(kv.stride(-1), 1) << name << " last dim must be contiguous";
+  TVM_FFI_ICHECK_EQ(reinterpret_cast<uintptr_t>(kv.data_ptr()) % 16, 0)
+      << name << " data pointer must be 16B-aligned (cp.async.bulk requirement)";
+  const size_t block_stride = static_cast<size_t>(kv.stride(0)) * elem_bytes;
+  TVM_FFI_ICHECK_EQ(block_stride % 16, 0)
+      << name << " block stride must be 16B-aligned (cp.async.bulk requirement)";
   if (kv.ndim() == 2) {
     const size_t block_bytes = static_cast<size_t>(kv.size(1)) * elem_bytes;
     TVM_FFI_ICHECK_EQ(block_bytes % static_cast<size_t>(bpt), 0)
         << name << " 2D block width " << block_bytes
         << " is not divisible by bytes_per_token=" << bpt;
     // A flat 2D block carries no row padding to infer, so the row advance is
     // exactly bytes_per_token.
-    return {static_cast<int>(block_bytes / static_cast<size_t>(bpt)), block_bytes, bpt};
+    TVM_FFI_ICHECK_GE(block_stride, block_bytes)
+        << name << " block stride is smaller than the packed block width";
+    return {static_cast<int>(block_bytes / static_cast<size_t>(bpt)), block_stride, bpt};
   }
   auto row_advance = [&](int64_t token_axis) {
-    TVM_FFI_ICHECK_EQ(kv.stride(-1), 1) << name << " last dim must be contiguous";
     const size_t bytes = static_cast<size_t>(kv.size(-1)) * elem_bytes;
     TVM_FFI_ICHECK_GE(bytes, static_cast<size_t>(bpt))
         << name << " row width " << bytes << " is smaller than bytes_per_token=" << bpt;
@@ -120,7 +127,7 @@ void SparseMlaSm120DecodeDsv4(TensorView q, TensorView kv_cache, TensorView indi
                               Optional<TensorView> topk_length, Optional<TensorView> attn_sink,
                               Optional<TensorView> extra_kv_cache,
                               Optional<TensorView> extra_indices,
-                              Optional<TensorView> extra_topk_length,
+                              Optional<TensorView> extra_topk_length, int64_t model_type,
                               int64_t chunks_per_block_override) {
   TVM_FFI_ICHECK_EQ(q.ndim(), 3) << "q must be [T, H, D_QK]";
   TVM_FFI_ICHECK_GE(kv_cache.ndim(), 2);
@@ -151,11 +158,18 @@ void SparseMlaSm120DecodeDsv4(TensorView q, TensorView kv_cache, TensorView indi
       << "indices leading dimension must match num_tokens";
   const int topk = static_cast<int>(indices.size(-1));
   const int d_qk = static_cast<int>(q.size(2));
-  // This kernel serves the footer-scale model types. d_qk selects between them:
-  // 512 -> DSV4, 1088 -> DOTS3_SWA (sliding-window family, d_v 1024).
+  // This kernel serves the footer-scale model types. model_type is the
+  // explicit selector from the Python planner; -1 keeps the legacy width
+  // inference (512 -> DSV4, 1088 -> DOTS3_SWA). Width alone cannot separate
+  // DSV4 from DSV4_1 (both are d_qk=512), so DSV4_1 is only reachable
+  // explicitly.
   TVM_FFI_ICHECK(d_qk == 512 || d_qk == 1088)
-      << "decode-dsv4 supports d_qk 512 (DSV4) or 1088 (DOTS3_SWA); got " << d_qk;
-  const ModelType mt = (d_qk == 512) ? ModelType::DSV4 : ModelType::DOTS3_SWA;
+      << "decode-dsv4 supports d_qk 512 (DSV4/DSV4_1) or 1088 (DOTS3_SWA); got " << d_qk;
+  const ModelType mt = model_type == -1 ? ((d_qk == 512) ? ModelType::DSV4 : ModelType::DOTS3_SWA)
+                                        : static_cast<ModelType>(model_type);
+  TVM_FFI_ICHECK((d_qk == 512 && (mt == ModelType::DSV4 || mt == ModelType::DSV4_1)) ||
+                 (d_qk == 1088 && mt == ModelType::DOTS3_SWA))
+      << "decode-dsv4 model_type mismatch: d_qk=" << d_qk << " model_type=" << model_type;
   // DOTS3_SWA's sliding window (513 candidates, DecodeTileCfg::WINDOW) needs an
   // indices buffer at least that wide; a narrower one can never name the full
   // window. Report it here so the message names the actual constraint.
@@ -165,7 +179,7 @@ void SparseMlaSm120DecodeDsv4(TensorView q, TensorView kv_cache, TensorView indi
       << topk;
   TVM_FFI_ICHECK(mt != ModelType::DOTS3_SWA || !extra_kv_cache.has_value())
       << "decode-dsv4 (dots3_swa) has no dual-cache form; extra_kv_cache is "
-         "DSV4-only";
+         "DSV4/DSV4_1-only";
 
   // topk_length is optional for DOTS3_SWA: DecodeTileCfg<DOTS3_SWA>::WINDOW caps
   // the per-token candidate count inside the kernel, so omitting it costs

csrc/sparse_mla_sm120_prefill.cu

@@ -32,7 +32,8 @@
 //     GLM_NSA / GLM53_NOPE), num_heads 64 / 128, single cache
 //   - SG (single-group, 16 heads/CTA): V32 family num_heads 8 / 16;
 //     DOTS3_SWA num_heads {8, 16, 32, 64} — SG-only, its D_NOPE=1024 does not
-//     fit the MG layout
+//     fit the MG layout; DSV4_1 num_heads {8, 16, 32, 64} — also SG-only, its
+//     32-wide quant groups floor the MG XV warp split to zero tiles
 //   - MG (multi-group, 32 heads/CTA): V32 family num_heads >= 32; DSV4
 //     num_heads {8..128}
 //   - MG_DUAL: dual-cache MG variants (DSV4 only)
@@ -388,6 +389,42 @@ inline bool dispatch_dots3_swa_sg(int num_heads, int topk, int page_block_size,
 #undef DISPATCH_DOTS3_SWA_SG
 }
 
+// DSV4_1 is SG-only for the same structural reason as DOTS3_SWA, but with the
+// warp split driven by its 32-wide quant groups (V_CHUNK=32 floors
+// NT_PER_WARP_XV to 0 for an 8-warp XV) rather than by smem capacity. Any
+// runtime topk made of whole index tiles is served; the binding enforces
+// topk % 64 == 0. TP1..TP8 shards of the 64-head layer ride REPLICATE_H.
+inline bool dispatch_dsv4_1_sg(int num_heads, int topk, int page_block_size, const bf16* Q,
+                               const uint8_t* KV, const int32_t* indices, const float* attn_sink,
+                               bf16* output, float* out_lse, float sm_scale, int num_tokens,
+                               size_t stride_kv_block, size_t stride_out_lse,
+                               const int* topk_length_ptr, cudaStream_t stream) {
+  if (page_block_size != 64) return false;
+
+#define DISPATCH_DSV4_1_SG(NH)                                                                 \
+  launch_prefill_sg<ModelType::DSV4_1, ComputeMode::FP8, NH, 64>(                              \
+      Q, KV, indices, attn_sink, output, out_lse, sm_scale, num_tokens, topk, stride_kv_block, \
+      stride_out_lse, topk_length_ptr, stream)
+
+  switch (num_heads) {
+    case 8:
+      DISPATCH_DSV4_1_SG(8);
+      return true;
+    case 16:
+      DISPATCH_DSV4_1_SG(16);
+      return true;
+    case 32:
+      DISPATCH_DSV4_1_SG(32);
+      return true;
+    case 64:
+      DISPATCH_DSV4_1_SG(64);
+      return true;
+    default:
+      return false;
+  }
+#undef DISPATCH_DSV4_1_SG
+}
+
 inline bool dispatch_dsv4_single(int num_heads, int topk, int page_block_size, const bf16* Q,
                                  const uint8_t* KV, const int32_t* indices, const float* attn_sink,
                                  bf16* output, float* out_lse, float sm_scale, int num_tokens,
@@ -576,6 +613,11 @@ bool sparse_mla_prefill_dispatch(ModelType mt, PrefillVariant variant, int num_h
                                      attn_sink, output, out_lse, sm_scale, num_tokens,
                                      stride_kv_block, stride_out_lse, topk_length, stream);
       }
+      if (mt == ModelType::DSV4_1) {
+        return dispatch_dsv4_1_sg(num_heads, topk, page_block_size, Q, KV_cache, indices, attn_sink,
+                                  output, out_lse, sm_scale, num_tokens, stride_kv_block,
+                                  stride_out_lse, topk_length, stream);
+      }
       DISPATCH_V32(dispatch_v32_sg);
     }
     case PrefillVariant::MG: {