Refactor Sparse MLA SM120

CLAUDE.md

@@ -583,7 +583,7 @@ Used by `flashinfer.trace` / `fi_trace`.
 | `FLASHINFER_AUTOTUNER_LOAD_FROM_FILE` | `0` | `flashinfer/autotuner/autotuner.py` | `1` loads previously serialized autotune results from disk instead of re-running the search. |
 | `FLASHINFER_DIST_AWARE_AUTOTUNE` | `0` | `flashinfer/fused_moe/da_config.py` | `1` enables experimental distribution-aware autotune and kernel dispatch (TRT-LLM MoE only). |
 | `FLASHINFER_DA_DISTRIBUTIONS` | built-in distribution catalog | `flashinfer/fused_moe/da_config.py` | Comma-separated training distributions used by the experimental TRT-LLM distribution-aware MoE autotuner. |
-| `FLASHINFER_AUTOTUNE_DIR` | unset | `flashinfer/mla/_sparse_mla_sm120.py`, `flashinfer/comm/pcie_ipc_tuning.py` | Override the disk path for AutoTuner cache files (MLA, and the PCIe IPC all-reduce). Falls back to `FLASHINFER_WORKSPACE_DIR` when unset. |
+| `FLASHINFER_AUTOTUNE_DIR` | unset | `flashinfer/mla/_sparse_mla_sm120_cpb.py`, `flashinfer/comm/pcie_ipc_tuning.py` | Override the disk path for tuning cache files (sparse-MLA SM120 cpb calibration constants, and the PCIe IPC all-reduce). Falls back to `FLASHINFER_WORKSPACE_DIR` when unset. |
 | `FLASHINFER_AUTOTUNE_TIMER` | unset (auto) | `flashinfer/autotuner/autotuner.py` | Selects the autotuner's per-tactic timer: `globaltimer` forces the GPU `%globaltimer` register, `cuda_event` forces `cudaEvent`, unset/anything-else auto-detects (uses `%globaltimer` only when Confidential Computing is detected). Under CC `cudaEventElapsedTime` is unreliable (can go negative), so the globaltimer path keeps tactic ranking stable. |
 | `FLASHINFER_CUTILE_AUTOTUNE_DISABLED` | `0` | `flashinfer/quantization/kernels/cutile/rope_quantize_fp8_cutile.py` | Non-zero skips exhaustive cuTile RoPE-FP8 tuning and uses the built-in token-count heuristic. |
 | `FLASHINFER_CONFIDENTIAL_COMPUTE` | unset | `flashinfer/utils.py` | Override NVIDIA Confidential Computing (CC) auto-detection used by `is_confidential_compute()` (which drives the autotuner timer above): `1` forces CC, `0` forces non-CC. Useful for CI or hosts without `pynvml`. |

benchmarks/bench_sparse_mla_sm120.py

@@ -307,11 +307,15 @@ def fn():
     return ms * 1e3, kv_bw_gbps, tflops
 
 
-def bench_sparse_mla_sm120_dsv3_2(num_heads, num_tokens, with_sink=False, seed=0):
+def bench_sparse_mla_sm120_dsv3_2(
+    num_heads, num_tokens, with_sink=False, seed=0, kv_scale_format="arbitrary_fp32"
+):
     """Returns (median_us, kv_bw_gbps, attn_tflops) for DSv3.2.
 
     Fixed: topk=2048, page_block_size=64 (= _DECODE_DSV3_2_PAGE_BLOCK_SIZE),
-    d_qk=576, d_v=512.
+    d_qk=576, d_v=512. ``kv_scale_format`` picks the power-of-2 inline-scale
+    path (``"auto"``) or the GLM arbitrary-FP32 path (software fold with
+    dual-limb weights).
     """
     torch.manual_seed(seed)
     device = torch.device("cuda")
@@ -351,7 +355,7 @@ def bench_sparse_mla_sm120_dsv3_2(num_heads, num_tokens, with_sink=False, seed=0
     runner = _SparseMLAPagedAttentionRunner(
         max_num_tokens=num_tokens,
         max_num_heads=num_heads,
-        kv_scale_format="arbitrary_fp32",
+        kv_scale_format=kv_scale_format,
         device=device,
     )
 
@@ -513,3 +517,18 @@ def fn():
         print(
             f"{h:>10}  {2048:>6}  {t:>11}  {lat_us:>10.1f}  {kvbw:>13.1f}  {tfl:>12.2f}"
         )
+
+    # DSv3.2 prefill: topk fixed at 2048, num_tokens > 64. Sweep is num_heads x
+    # num_tokens x kv_scale_format; 64/128 heads run the swapAB kernel.
+    dsv3_2_prefill_configs = [(h, t) for t in (128, 512) for h in (64, 128)]
+
+    for fmt in ("auto", "arbitrary_fp32"):
+        print()
+        print(f"DSv3.2 prefill path (num_tokens > 64, kv_scale_format={fmt}):")
+        print(header)
+        print("-" * len(header))
+        for h, t in dsv3_2_prefill_configs:
+            lat_us, kvbw, tfl = bench_sparse_mla_sm120_dsv3_2(h, t, kv_scale_format=fmt)
+            print(
+                f"{h:>10}  {2048:>6}  {t:>11}  {lat_us:>10.1f}  {kvbw:>13.1f}  {tfl:>12.2f}"
+            )

csrc/sparse_mla_sm120.cu

@@ -45,14 +45,15 @@ using tvm::ffi::Optional;
 namespace flashinfer::sparse_mla_sm120 {
 
 // Forward declaration (defined in sparse_mla_sm120_prefill.cu).
-bool sparse_mla_prefill_dispatch(ModelType mt, int num_heads, int topk, int page_block_size,
-                                 int topk_extra, int extra_page_block_size, const bf16* Q,
-                                 const uint8_t* KV_cache, const int32_t* indices,
+bool sparse_mla_prefill_dispatch(ModelType mt, PrefillVariant variant, int num_heads, int topk,
+                                 int page_block_size, int topk_extra, int extra_page_block_size,
+                                 const bf16* Q, const uint8_t* KV_cache, const int32_t* indices,
                                  const uint8_t* extra_KV_cache, const int32_t* extra_indices,
                                  bf16* output, float* out_lse, float sm_scale, int num_tokens,
                                  size_t stride_kv_block, size_t extra_stride_kv_block,
-                                 const float* attn_sink, const int* topk_length,
-                                 const int* extra_topk_length, cudaStream_t stream);
+                                 size_t stride_out_lse, const float* attn_sink,
+                                 const int* topk_length, const int* extra_topk_length,
+                                 cudaStream_t stream);
 
 namespace {
 
@@ -68,26 +69,26 @@ 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)
-        << "d_qk=512 supports only model_type auto or DSV4; got " << 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;
+    return mt;
+  }
+  if (d_qk == 1088) {
+    const auto mt = static_cast<ModelType>(
+        model_type == kAuto ? static_cast<int64_t>(ModelType::DOTS3_SWA) : model_type);
+    TVM_FFI_ICHECK(mt == ModelType::DOTS3_SWA)
+        << "d_qk=1088 supports only model_type auto or DOTS3_SWA; got " << model_type;
     return mt;
   }
   TVM_FFI_ICHECK(false) << "Unsupported d_qk=" << d_qk
-                        << "; expected 576 (DSV3_2/GLM_NSA) or 512 (DSV4)";
+                        << "; expected 576 (DSV3_2/GLM_NSA), 512 (DSV4/GLM53_NOPE) or 1088 "
+                           "(DOTS3_SWA)";
   return ModelType::DSV4;
 }
 
-inline int bytes_per_token(ModelType mt) {
-  switch (mt) {
-    case ModelType::DSV3_2:
-    case ModelType::GLM_NSA:
-      return 656;
-    case ModelType::DSV4:
-      return 584;
-  }
-  TVM_FFI_ICHECK(false) << "Unsupported sparse MLA model type";
-  return 0;
-}
+// Output width. Every DeepSeek-family model absorbs V into the 512-wide
+// latent; DOTS3_SWA's is 1024 (its 1024-wide latent, rope excluded).
+inline int d_v_for(ModelType mt) { return mt == ModelType::DOTS3_SWA ? 1024 : 512; }
 
 struct PagedKVLayout {
   int page_block_size;
@@ -105,15 +106,19 @@ inline PagedKVLayout parse_paged_kv_layout(const TensorView& kv, int bpt, const
   }
   if (kv.ndim() == 3) {
     TVM_FFI_ICHECK_EQ(kv.size(-1), bpt)
-        << name << " 3D form must be [num_pages, page_block_size, " << bpt << "]";
+        << name << " 3D form must be [num_pages, page_block_size, " << bpt
+        << "]; prefill requires tightly packed rows — padded-row KV caches are "
+           "decode-only (the decode-v32 kernel honors stride_kv_row)";
     return {static_cast<int>(kv.size(1)), static_cast<size_t>(kv.stride(0)) * elem_bytes};
   }
   TVM_FFI_ICHECK_EQ(kv.ndim(), 4) << name << " must be 2D [num_pages, page_bytes], 3D "
                                   << "[num_pages, page_block_size, bytes_per_token], HND "
                                   << "[num_pages, 1, page_block_size, bytes_per_token], or NHD "
                                   << "[num_pages, page_block_size, 1, bytes_per_token]";
   TVM_FFI_ICHECK_EQ(kv.size(-1), bpt)
-      << name << " last dim must be bytes_per_token=" << bpt << ", got " << kv.size(-1);
+      << name << " last dim must be bytes_per_token=" << bpt << ", got " << kv.size(-1)
+      << "; prefill requires tightly packed rows — padded-row KV caches are "
+         "decode-only (the decode-v32 kernel honors stride_kv_row)";
   if (kv.size(1) == 1) {
     return {static_cast<int>(kv.size(2)), static_cast<size_t>(kv.stride(0)) * elem_bytes};
   }
@@ -147,8 +152,10 @@ void SparseMlaSm120PagedAttention(
     TensorView kv_cache,  // packed paged FP8; HND/NHD/3D/2D forms accepted
     TensorView indices,   // [num_tokens, topk] or [num_tokens, 1, topk] int32 (-1 = skip)
     TensorView output,    // [num_tokens, num_heads, d_v] bf16 — in-place
-    TensorView out_lse,   // [num_tokens, num_heads] f32 — in-place
-    double sm_scale, int64_t model_type,
+    // [num_tokens, num_heads] f32 — in-place; row-strided views (a column
+    // slice of a wider buffer) are supported, the head dim stays contiguous.
+    TensorView out_lse, double sm_scale, int64_t model_type,
+    int64_t variant,                         // PrefillVariant; planner-selected
     Optional<TensorView> topk_length,        // [num_tokens] int32, optional
     Optional<TensorView> attn_sink,          // [num_heads] f32, optional
     Optional<TensorView> extra_kv_cache,     // optional dual cache
@@ -157,13 +164,18 @@ void SparseMlaSm120PagedAttention(
 {
   // ── Input validation ───────────────────────────────────────────────
   CHECK_INPUT_AND_TYPE(q, dl_bfloat16);
-  // kv_cache: CUDA + last-dim contiguous only; padded block stride is OK.
+  // kv_cache: CUDA + last-dim contiguous only. Padded block strides are OK
+  // for footer-scale models (DSV4 / DOTS3_SWA); inline-scale models are
+  // checked below.
   CHECK_CUDA(kv_cache);
   CHECK_LAST_DIM_CONTIGUOUS(kv_cache);
   CHECK_INPUT_TYPE(kv_cache, dl_uint8);
   CHECK_INPUT_AND_TYPE(indices, dl_int32);
   CHECK_INPUT_AND_TYPE(output, dl_bfloat16);
-  CHECK_INPUT_AND_TYPE(out_lse, dl_float32);
+  // out_lse may be a row-strided column slice of a wider buffer; checked
+  // against its real strides below instead of requiring full contiguity.
+  CHECK_CUDA(out_lse);
+  CHECK_INPUT_TYPE(out_lse, dl_float32);
 
   CHECK_DIM(3, q);
 
@@ -174,18 +186,45 @@ void SparseMlaSm120PagedAttention(
   const ModelType mt = resolve_model_type(d_qk, model_type);
   const PagedKVLayout kv_layout = parse_paged_kv_layout(kv_cache, bytes_per_token(mt), "kv_cache");
   const int page_block_size = kv_layout.page_block_size;
+  // Inline-scale models (DSV3_2 / GLM_NSA / GLM53_NOPE) are addressed by the
+  // prefill kernels as a flat token array (prefill_kv_entry_base), so a
+  // padded block stride would be silently misread; only footer-scale models
+  // honor stride_kv_block. Padded strides remain a decode-path capability.
+  if (mt == ModelType::DSV3_2 || mt == ModelType::GLM_NSA || mt == ModelType::GLM53_NOPE) {
+    TVM_FFI_ICHECK_EQ(kv_layout.stride_kv_block,
+                      static_cast<size_t>(page_block_size) * bytes_per_token(mt))
+        << "prefill for inline-scale KV caches (DSv3.2/GLM) requires densely packed blocks "
+           "(stride_kv_block == page_block_size * bytes_per_token); padded block strides are "
+           "decode-only";
+  }
 
   TVM_FFI_ICHECK_GT(num_heads, 0);
   TVM_FFI_ICHECK_LE(num_heads, 128);
   TVM_FFI_ICHECK_GT(topk, 0);
+  // The prefill kernels issue whole BI=64-wide index tiles (the tail tile is
+  // not masked), so the indices row width must be a multiple of 64.
+  TVM_FFI_ICHECK_EQ(topk % 64, 0)
+      << "sparse-MLA SM120 prefill requires topk % 64 == 0 (BI=64: index tiles "
+         "are issued whole, the tail tile is not masked); got topk="
+      << topk;
+  // DOTS3_SWA's candidate list is the 513-wide sliding window; the buffer
+  // must fit it (the kernel clamps the scan to the window, not the row).
+  TVM_FFI_ICHECK(mt != ModelType::DOTS3_SWA || topk >= 513)
+      << "sparse-MLA SM120 DOTS3_SWA prefill requires topk >= 513 (the "
+         "513-wide sliding window must fit the indices buffer); got topk="
+      << topk;
   TVM_FFI_ICHECK_GT(page_block_size, 0);
-  TVM_FFI_ICHECK_EQ(output.ndim(), 3) << "output must be [num_tokens, num_heads, 512]";
+  const int d_v = d_v_for(mt);
+  TVM_FFI_ICHECK_EQ(output.ndim(), 3) << "output must be [num_tokens, num_heads, d_v]";
   TVM_FFI_ICHECK_EQ(output.size(0), num_tokens);
   TVM_FFI_ICHECK_EQ(output.size(1), num_heads);
-  TVM_FFI_ICHECK_EQ(output.size(2), 512) << "SM120 sparse-MLA requires d_v == 512";
+  TVM_FFI_ICHECK_EQ(output.size(2), d_v)
+      << "SM120 sparse-MLA requires d_v == " << d_v << " for this model type";
   TVM_FFI_ICHECK_EQ(out_lse.ndim(), 2) << "out_lse must be [num_tokens, num_heads]";
   TVM_FFI_ICHECK_EQ(out_lse.size(0), num_tokens);
   TVM_FFI_ICHECK_EQ(out_lse.size(1), num_heads);
+  TVM_FFI_ICHECK_EQ(out_lse.stride(-1), 1) << "out_lse last dimension must be contiguous";
+  const size_t stride_out_lse = static_cast<size_t>(out_lse.stride(0));
 
   if (topk_length.has_value()) {
     const auto& tl = topk_length.value();
@@ -256,25 +295,43 @@ void SparseMlaSm120PagedAttention(
   const auto LSE_ptr = static_cast<float*>(out_lse.data_ptr());
 
   // Decode (num_tokens <= 64) is dispatched by Python directly through the
-  // standalone decode-dsv3_2 / decode-dsv4 entry points. The orchestrator
-  // only handles prefill.
-  TVM_FFI_ICHECK_GT(num_tokens, 64)
-      << "Decode (num_tokens <= 64) must go through sparse_mla_sm120_decode_dsv3_2 "
-         "or sparse_mla_sm120_decode_dsv4; got num_tokens="
-      << num_tokens;
+  // standalone decode-dsv3_2 / decode-dsv4 entry points. The prefill kernels
+  // launch one CTA group per query token and guard s_i >= num_tokens, so any
+  // num_tokens >= 1 is valid here.
+  TVM_FFI_ICHECK_GE(num_tokens, 1)
+      << "prefill requires num_tokens >= 1; got num_tokens=" << num_tokens;
+
+  TVM_FFI_ICHECK(variant >= static_cast<int64_t>(PrefillVariant::SG) &&
+                 variant <= static_cast<int64_t>(PrefillVariant::SWAPAB))
+      << "variant must be a PrefillVariant (1=SG, 2=MG, 3=MG_DUAL, 4=SWAPAB); got " << variant;
 
   const bool ok = sparse_mla_prefill_dispatch(
-      mt, num_heads, topk, page_block_size, extra_topk, extra_page_block_size, Q_ptr, KV_ptr,
-      idx_ptr, extra_kv_ptr, extra_idx_ptr, O_ptr, LSE_ptr, static_cast<float>(sm_scale),
-      num_tokens, kv_layout.stride_kv_block, extra_stride_kv_block, attn_sink_ptr, tl_ptr, etl_ptr,
-      stream);
+      mt, static_cast<PrefillVariant>(variant), num_heads, topk, page_block_size, extra_topk,
+      extra_page_block_size, Q_ptr, KV_ptr, idx_ptr, extra_kv_ptr, extra_idx_ptr, O_ptr, LSE_ptr,
+      static_cast<float>(sm_scale), num_tokens, kv_layout.stride_kv_block, extra_stride_kv_block,
+      stride_out_lse, attn_sink_ptr, tl_ptr, etl_ptr, stream);
+  const char* mt_name = "DSV4";
+  switch (mt) {
+    case ModelType::DSV3_2:
+      mt_name = "DSV3_2";
+      break;
+    case ModelType::GLM_NSA:
+      mt_name = "GLM_NSA";
+      break;
+    case ModelType::GLM53_NOPE:
+      mt_name = "GLM53_NOPE";
+      break;
+    case ModelType::DOTS3_SWA:
+      mt_name = "DOTS3_SWA";
+      break;
+    case ModelType::DSV4:
+      break;
+  }
   TVM_FFI_ICHECK(ok) << "Unsupported sparse-MLA prefill configuration: "
-                     << "model="
-                     << (mt == ModelType::DSV3_2 ? "DSV3_2"
-                                                 : (mt == ModelType::GLM_NSA ? "GLM_NSA" : "DSV4"))
-                     << " num_heads=" << num_heads << " topk=" << topk
+                     << "model=" << mt_name << " num_heads=" << num_heads << " topk=" << topk
                      << " page_block_size=" << page_block_size << " topk_extra=" << extra_topk
-                     << " extra_page_block_size=" << extra_page_block_size;
+                     << " extra_page_block_size=" << extra_page_block_size
+                     << " variant=" << variant;
 }
 
 }  // namespace flashinfer::sparse_mla_sm120