Sparse MLA on Hopper: Use SGLang's kernel for the sparse mla low latency runs

Signed-off-by: Alexander Matveev <amatveev@redhat.com>
2026-04-17 16:51:32 +00:00
9 changed files with 331 additions and 320 deletions
@@ -66,7 +66,7 @@ TORCH_LIBRARY_FRAGMENT(_cutlass_fa3_C, m) {

 // Python module initialization for _cutlass_fa3_C
 PyMODINIT_FUNC PyInit__cutlass_fa3_C() {
-  static struct PyModuleDef module = {
-      PyModuleDef_HEAD_INIT, "_cutlass_fa3_C", nullptr, 0, nullptr};
+  static struct PyModuleDef module = {PyModuleDef_HEAD_INIT, "_cutlass_fa3_C",
+                                      nullptr, 0, nullptr};
  return PyModule_Create(&module);
 }
@@ -24,35 +24,22 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> mha_fwd(
    at::Tensor q,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
    at::Tensor k,  // (b_k, s_k, h_k, d) or (total_k, h_k, d) or paged
    at::Tensor v,  // (b_k, s_k, h_k, dv) or (total_k, h_k, dv) or paged
-    std::optional<at::Tensor> k_new_,
-    std::optional<at::Tensor> v_new_,
-    std::optional<at::Tensor> q_v_,     // MLA value projection query
-    std::optional<at::Tensor> out_,
-    std::optional<at::Tensor> cu_seqlens_q_,
+    std::optional<at::Tensor> k_new_, std::optional<at::Tensor> v_new_,
+    std::optional<at::Tensor> q_v_,  // MLA value projection query
+    std::optional<at::Tensor> out_, std::optional<at::Tensor> cu_seqlens_q_,
    std::optional<at::Tensor> cu_seqlens_k_,
    std::optional<at::Tensor> cu_seqlens_k_new_,
-    std::optional<at::Tensor> seqused_q_,
-    std::optional<at::Tensor> seqused_k_,
-    std::optional<int64_t> max_seqlen_q_,
-    std::optional<int64_t> max_seqlen_k_,
+    std::optional<at::Tensor> seqused_q_, std::optional<at::Tensor> seqused_k_,
+    std::optional<int64_t> max_seqlen_q_, std::optional<int64_t> max_seqlen_k_,
    std::optional<at::Tensor> page_table_,
    std::optional<at::Tensor> kv_batch_idx_,
-    std::optional<at::Tensor> leftpad_k_,
-    std::optional<at::Tensor> rotary_cos_,
+    std::optional<at::Tensor> leftpad_k_, std::optional<at::Tensor> rotary_cos_,
    std::optional<at::Tensor> rotary_sin_,
    std::optional<at::Tensor> seqlens_rotary_,
-    std::optional<at::Tensor> q_descale_,
-    std::optional<at::Tensor> k_descale_,
-    std::optional<at::Tensor> v_descale_,
-    std::optional<double> softmax_scale_,
-    bool is_causal,
-    int64_t window_size_left,
-    int64_t window_size_right,
-    int64_t attention_chunk,
-    double softcap,
-    bool is_rotary_interleaved,
-    std::optional<at::Tensor> scheduler_metadata_,
-    int64_t num_splits,
-    std::optional<bool> pack_gqa_,
-    int64_t sm_margin,
+    std::optional<at::Tensor> q_descale_, std::optional<at::Tensor> k_descale_,
+    std::optional<at::Tensor> v_descale_, std::optional<double> softmax_scale_,
+    bool is_causal, int64_t window_size_left, int64_t window_size_right,
+    int64_t attention_chunk, double softcap, bool is_rotary_interleaved,
+    std::optional<at::Tensor> scheduler_metadata_, int64_t num_splits,
+    std::optional<bool> pack_gqa_, int64_t sm_margin,
    std::optional<const at::Tensor>& sinks_);
@@ -38,16 +38,16 @@
 template <typename T>
 struct pytorch_library_compatible_type {
  using type = T;
-  static T convert_from_type(T arg) {
-    return arg;
-  }
+  static T convert_from_type(T arg) { return arg; }
 };

 template <typename T>
-using pytorch_library_compatible_type_t = typename pytorch_library_compatible_type<T>::type;
+using pytorch_library_compatible_type_t =
+    typename pytorch_library_compatible_type<T>::type;

 template <typename T>
-T convert_from_pytorch_compatible_type(pytorch_library_compatible_type_t<T> arg) {
+T convert_from_pytorch_compatible_type(
+    pytorch_library_compatible_type_t<T> arg) {
  return pytorch_library_compatible_type<T>::convert_from_type(arg);
 }

@@ -68,7 +68,8 @@ struct pytorch_library_compatible_type<c10::optional<T>&> {
 template <typename T>
 struct pytorch_library_compatible_type<c10::optional<T>> {
  using type = c10::optional<pytorch_library_compatible_type_t<T>>;
-  static c10::optional<pytorch_library_compatible_type_t<T>> convert_from_type(c10::optional<T> arg) {
+  static c10::optional<pytorch_library_compatible_type_t<T>> convert_from_type(
+      c10::optional<T> arg) {
    return arg;
  }
 };
@@ -77,8 +78,10 @@ struct pytorch_library_compatible_type<c10::optional<T>> {
 template <>
 struct pytorch_library_compatible_type<c10::optional<const at::Tensor>&> {
  using type = const c10::optional<at::Tensor>&;
-  static c10::optional<const at::Tensor>& convert_from_type(const c10::optional<at::Tensor>& arg) {
-    return const_cast<c10::optional<const at::Tensor>&>(reinterpret_cast<const c10::optional<const at::Tensor>&>(arg));
+  static c10::optional<const at::Tensor>& convert_from_type(
+      const c10::optional<at::Tensor>& arg) {
+    return const_cast<c10::optional<const at::Tensor>&>(
+        reinterpret_cast<const c10::optional<const at::Tensor>&>(arg));
  }
 };

@@ -87,8 +90,10 @@ template <>
 struct pytorch_library_compatible_type<int> {
  using type = int64_t;
  static int convert_from_type(int64_t arg) {
-    TORCH_CHECK(arg <= std::numeric_limits<int>::max(), "int64_t value is too large to be converted to int");
-    TORCH_CHECK(arg >= std::numeric_limits<int>::min(), "int64_t value is too small to be converted to int");
+    TORCH_CHECK(arg <= std::numeric_limits<int>::max(),
+                "int64_t value is too large to be converted to int");
+    TORCH_CHECK(arg >= std::numeric_limits<int>::min(),
+                "int64_t value is too small to be converted to int");
    return arg;
  }
 };
@@ -98,8 +103,8 @@ template <>
 struct pytorch_library_compatible_type<float> {
  using type = double;
  static float convert_from_type(double arg) {
-    TORCH_CHECK(
-        std::abs(arg) <= std::numeric_limits<float>::max(), "double value is too large to be converted to float");
+    TORCH_CHECK(std::abs(arg) <= std::numeric_limits<float>::max(),
+                "double value is too large to be converted to float");
    return arg;
  }
 };
@@ -132,16 +132,6 @@ Priority is **1 = highest** (tried first).
 | 6 | `FLASHINFER_MLA_SPARSE`**\*** |
 | 7 | `FLASHMLA_SPARSE` |

-**Ampere/Hopper (SM 8.x-9.x):**
-
-| Priority | Backend |
-| -------- | ------- |
-| 1 | `FLASH_ATTN_MLA` |
-| 2 | `FLASHMLA` |
-| 3 | `FLASHINFER_MLA` |
-| 4 | `TRITON_MLA` |
-| 5 | `FLASHMLA_SPARSE` |
-
 > **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASHINFER_MLA_SPARSE` is preferred for low query-head counts (<= 16), while `FLASHMLA_SPARSE` is preferred otherwise.
 >
 > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
@@ -209,6 +199,7 @@ configuration.

 | Backend | Dtypes | KV Dtypes | Block Sizes | Head Sizes | Sink | Sparse | MM Prefix | DCP | Attention Types | Compute Cap. |
 | ------- | ------ | --------- | ----------- | ---------- | ---- | ------ | --------- | --- | --------------- | ------------ |
+| `CUTLASS_FA3_MLA_SPARSE` | bf16 | `auto` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x |
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |
@@ -15,6 +15,7 @@ Tests cover:
  7. num_splits variants
  8. softcap parameter
 """
+
 import pytest
 import torch
 import torch.nn.functional as F
@@ -46,9 +47,7 @@ def _make_kv_cache(kv_pool_size: int, dtype=torch.bfloat16, device="cuda"):
    k_rope_cache = torch.randn(
        kv_pool_size, 1, 1, HEADDIM_QK, dtype=dtype, device=device
    )
-    v_cache = torch.randn(
-        kv_pool_size, 1, 1, HEADDIM_V, dtype=dtype, device=device
-    )
+    v_cache = torch.randn(kv_pool_size, 1, 1, HEADDIM_V, dtype=dtype, device=device)
    return k_rope_cache, v_cache


@@ -222,9 +221,9 @@ def test_fa3_mla_correctness_vs_reference():
        # SDPA expects: Q[batch, heads, seq_q, dim], K[batch, heads, seq_k, dim]
        # Q: [1, N, 1, 576], K: [1, N, topk, 576], V: [1, N, topk, 512]
        ref_out = F.scaled_dot_product_attention(
-            q_full.unsqueeze(0).unsqueeze(2),                    # [1, N, 1, 576]
-            k_expanded.transpose(0, 1).unsqueeze(0),             # [1, N, topk, 576]
-            v_expanded.transpose(0, 1).unsqueeze(0),             # [1, N, topk, 512]
+            q_full.unsqueeze(0).unsqueeze(2),  # [1, N, 1, 576]
+            k_expanded.transpose(0, 1).unsqueeze(0),  # [1, N, topk, 576]
+            v_expanded.transpose(0, 1).unsqueeze(0),  # [1, N, topk, 512]
            scale=SOFTMAX_SCALE,
        )  # -> [1, N, 1, 512]
        ref_out = ref_out.squeeze(0).squeeze(1)  # [N, 512]
@@ -251,12 +250,10 @@ def test_fa3_mla_with_invalid_indices():
    k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device)

    # Mix of valid and -1 entries
-    page_table = torch.full(
-        (T, topk_total), -1, dtype=torch.int32, device=device
-    )
-    page_table[0, :topk_valid] = torch.randperm(
-        kv_pool_size, device=device
-    )[:topk_valid].to(torch.int32)
+    page_table = torch.full((T, topk_total), -1, dtype=torch.int32, device=device)
+    page_table[0, :topk_valid] = torch.randperm(kv_pool_size, device=device)[
+        :topk_valid
+    ].to(torch.int32)

    cache_seqlens = torch.tensor([topk_valid], dtype=torch.int32, device=device)
    cu_seqlens_q, cu_seqlens_k = _make_cu_seqlens(cache_seqlens, device=device)
@@ -304,16 +301,14 @@ def test_fa3_mla_variable_seqlens(seq_lens):
    q_nope = torch.randn(T, NUM_HEADS, HEADDIM_V, dtype=torch.bfloat16, device=device)
    k_rope_cache, v_cache = _make_kv_cache(kv_pool_size, device=device)

-    page_table = torch.full(
-        (T, topk), 0, dtype=torch.int32, device=device
-    )
+    page_table = torch.full((T, topk), 0, dtype=torch.int32, device=device)
    actual_seqlens = []
    for i, sl in enumerate(seq_lens):
        actual_topk = min(sl, topk)
        actual_seqlens.append(actual_topk)
-        page_table[i, :actual_topk] = torch.randperm(
-            kv_pool_size, device=device
-        )[:actual_topk].to(torch.int32)
+        page_table[i, :actual_topk] = torch.randperm(kv_pool_size, device=device)[
+            :actual_topk
+        ].to(torch.int32)

    cache_seqlens = torch.tensor(actual_seqlens, dtype=torch.int32, device=device)
    cu_seqlens_q, cu_seqlens_k = _make_cu_seqlens(cache_seqlens, device=device)
@@ -448,9 +443,9 @@ def test_fa3_mla_prefill_short_sequence(seq_len):
    for i in range(T):
        num_valid = i + 1  # causal: token i sees positions 0..i
        valid_counts.append(num_valid)
-        page_table[i, :num_valid] = torch.randperm(
-            kv_pool_size, device=device
-        )[:num_valid].to(torch.int32)
+        page_table[i, :num_valid] = torch.randperm(kv_pool_size, device=device)[
+            :num_valid
+        ].to(torch.int32)

    cache_seqlens = torch.tensor(valid_counts, dtype=torch.int32, device=device)
    cu_seqlens_q, _ = _make_cu_seqlens(cache_seqlens, device=device)
@@ -498,9 +493,9 @@ def test_fa3_mla_page_table_minus1_clamped():
    valid_per_token = [10, 50, 100, 200]
    for i in range(T):
        nv = valid_per_token[i]
-        page_table[i, :nv] = torch.randperm(
-            kv_pool_size, device=device
-        )[:nv].to(torch.int32)
+        page_table[i, :nv] = torch.randperm(kv_pool_size, device=device)[:nv].to(
+            torch.int32
+        )

    page_table = page_table.clamp(min=0)
    cache_seqlens = torch.tensor(valid_per_token, dtype=torch.int32, device=device)
@@ -626,9 +621,9 @@ def test_fa3_valid_counts_correctness():
    page_table_raw = torch.full((T, topk), -1, dtype=torch.int32, device=device)
    for i in range(T):
        nv = valid_per_token[i]
-        page_table_raw[i, :nv] = torch.randperm(
-            kv_pool_size, device=device
-        )[:nv].to(torch.int32)
+        page_table_raw[i, :nv] = torch.randperm(kv_pool_size, device=device)[:nv].to(
+            torch.int32
+        )

    page_table = page_table_raw.clamp(min=0)
    cache_seqlens = torch.tensor(valid_per_token, dtype=torch.int32, device=device)
@@ -662,16 +657,23 @@ def test_fa3_valid_counts_correctness():
        v_exp = v_gathered.unsqueeze(1).expand(-1, NUM_HEADS, -1)
        q_full = torch.cat([q_n, q_r], dim=-1)

-        ref_out = F.scaled_dot_product_attention(
-            q_full.unsqueeze(0).unsqueeze(2),
-            k_exp.transpose(0, 1).unsqueeze(0),
-            v_exp.transpose(0, 1).unsqueeze(0),
-            scale=SOFTMAX_SCALE,
-        ).squeeze(0).squeeze(1)
+        ref_out = (
+            F.scaled_dot_product_attention(
+                q_full.unsqueeze(0).unsqueeze(2),
+                k_exp.transpose(0, 1).unsqueeze(0),
+                v_exp.transpose(0, 1).unsqueeze(0),
+                scale=SOFTMAX_SCALE,
+            )
+            .squeeze(0)
+            .squeeze(1)
+        )

        torch.testing.assert_close(
-            out[b].float(), ref_out.float(), rtol=0.02, atol=0.02,
-            msg=f"Token {b} (valid={nv}) mismatch"
+            out[b].float(),
+            ref_out.float(),
+            rtol=0.02,
+            atol=0.02,
+            msg=f"Token {b} (valid={nv}) mismatch",
        )


@@ -738,9 +740,9 @@ def test_fa3_mla_mixed_batch_prefill_decode():
    page_table = torch.zeros((T, topk), dtype=torch.int32, device=device)
    for i in range(T):
        nv = valid_per_token[i]
-        page_table[i, :nv] = torch.randperm(
-            kv_pool_size, device=device
-        )[:nv].to(torch.int32)
+        page_table[i, :nv] = torch.randperm(kv_pool_size, device=device)[:nv].to(
+            torch.int32
+        )

    cache_seqlens = torch.tensor(valid_per_token, dtype=torch.int32, device=device)
    cu_seqlens_q = torch.arange(0, T + 1, dtype=torch.int32, device=device)
@@ -786,8 +788,11 @@ def test_fa3_mla_kv_cache_write_then_read():

    # 1) Create BF16 KV cache and write known values
    cache = torch.zeros(
-        num_blocks, block_size, head_size,
-        dtype=torch.bfloat16, device=device,
+        num_blocks,
+        block_size,
+        head_size,
+        dtype=torch.bfloat16,
+        device=device,
    )
    kv_c_normed = torch.randn(T, kv_lora_rank, dtype=torch.bfloat16, device=device)
    k_pe = torch.randn(T, 1, qk_rope_head_dim, dtype=torch.bfloat16, device=device)
@@ -796,9 +801,14 @@ def test_fa3_mla_kv_cache_write_then_read():
    k_scale = torch.ones(1, dtype=torch.float32, device=device)

    from vllm import _custom_ops as ops
+
    ops.concat_and_cache_mla(
-        kv_c_normed, k_pe.squeeze(1), cache, slot_mapping,
-        kv_cache_dtype="auto", scale=k_scale,
+        kv_c_normed,
+        k_pe.squeeze(1),
+        cache,
+        slot_mapping,
+        kv_cache_dtype="auto",
+        scale=k_scale,
    )

    # 2) Reshape cache for FA3 (page_size=1 format)
@@ -808,8 +818,12 @@ def test_fa3_mla_kv_cache_write_then_read():
    k_rope = kv_flat[:, kv_lora_rank:].reshape(S, 1, 1, qk_rope_head_dim)

    # 3) Create Q and page_table simulating causal prefill
-    q_rope = torch.randn(T, NUM_HEADS, qk_rope_head_dim, dtype=torch.bfloat16, device=device)
-    q_nope = torch.randn(T, NUM_HEADS, kv_lora_rank, dtype=torch.bfloat16, device=device)
+    q_rope = torch.randn(
+        T, NUM_HEADS, qk_rope_head_dim, dtype=torch.bfloat16, device=device
+    )
+    q_nope = torch.randn(
+        T, NUM_HEADS, kv_lora_rank, dtype=torch.bfloat16, device=device
+    )

    # Causal page_table: token i sees slots 0..i, rest padded with 0
    topk = 128
@@ -848,27 +862,34 @@ def test_fa3_mla_kv_cache_write_then_read():
    for b in range(T):
        nv = valid_per_token[b]
        # Gather the actual written KV from cache slots 0..nv-1
-        k_gathered = k_rope[:nv].squeeze(1).squeeze(1)    # [nv, 64]
-        v_gathered = c_kv[:nv].squeeze(1).squeeze(1)      # [nv, 512]
+        k_gathered = k_rope[:nv].squeeze(1).squeeze(1)  # [nv, 64]
+        v_gathered = c_kv[:nv].squeeze(1).squeeze(1)  # [nv, 512]

-        q_r = q_rope[b]     # [N, 64]
-        q_n = q_nope[b]     # [N, 512]
+        q_r = q_rope[b]  # [N, 64]
+        q_n = q_nope[b]  # [N, 512]
        k_full = torch.cat([v_gathered, k_gathered], dim=-1)  # [nv, 576]

        k_exp = k_full.unsqueeze(1).expand(-1, NUM_HEADS, -1)
        v_exp = v_gathered.unsqueeze(1).expand(-1, NUM_HEADS, -1)
        q_full = torch.cat([q_n, q_r], dim=-1)

-        ref_out = F.scaled_dot_product_attention(
-            q_full.unsqueeze(0).unsqueeze(2),
-            k_exp.transpose(0, 1).unsqueeze(0),
-            v_exp.transpose(0, 1).unsqueeze(0),
-            scale=SOFTMAX_SCALE,
-        ).squeeze(0).squeeze(1)
+        ref_out = (
+            F.scaled_dot_product_attention(
+                q_full.unsqueeze(0).unsqueeze(2),
+                k_exp.transpose(0, 1).unsqueeze(0),
+                v_exp.transpose(0, 1).unsqueeze(0),
+                scale=SOFTMAX_SCALE,
+            )
+            .squeeze(0)
+            .squeeze(1)
+        )

        torch.testing.assert_close(
-            out[b].float(), ref_out.float(), rtol=0.02, atol=0.02,
-            msg=f"E2E token {b} (valid={nv}): FA3 vs SDPA mismatch"
+            out[b].float(),
+            ref_out.float(),
+            rtol=0.02,
+            atol=0.02,
+            msg=f"E2E token {b} (valid={nv}): FA3 vs SDPA mismatch",
        )


@@ -931,15 +952,22 @@ def test_fa3_small_batch_correctness(batch_size):
    k_full = torch.cat([v_gathered, k_gathered], dim=-1)
    k_exp = k_full.unsqueeze(1).expand(-1, NUM_HEADS, -1)
    v_exp = v_gathered.unsqueeze(1).expand(-1, NUM_HEADS, -1)
-    ref_out = F.scaled_dot_product_attention(
-        q_full.unsqueeze(0).unsqueeze(2),
-        k_exp.transpose(0, 1).unsqueeze(0),
-        v_exp.transpose(0, 1).unsqueeze(0),
-        scale=SOFTMAX_SCALE,
-    ).squeeze(0).squeeze(1)
+    ref_out = (
+        F.scaled_dot_product_attention(
+            q_full.unsqueeze(0).unsqueeze(2),
+            k_exp.transpose(0, 1).unsqueeze(0),
+            v_exp.transpose(0, 1).unsqueeze(0),
+            scale=SOFTMAX_SCALE,
+        )
+        .squeeze(0)
+        .squeeze(1)
+    )
    torch.testing.assert_close(
-        out[0].float(), ref_out.float(), rtol=0.02, atol=0.02,
-        msg=f"FA3 small batch (bs={batch_size}) token 0 mismatch"
+        out[0].float(),
+        ref_out.float(),
+        rtol=0.02,
+        atol=0.02,
+        msg=f"FA3 small batch (bs={batch_size}) token 0 mismatch",
    )


@@ -951,6 +979,7 @@ def test_flashmla_sparse_fallback_available():
    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
        _flashmla_sparse_available,
    )
+
    # On SM90 with FlashMLA compiled, the fallback should be available
    # This test may skip if FlashMLA is not compiled (non-standard build)
    if not _flashmla_sparse_available:
@@ -982,8 +1011,8 @@ def test_flashmla_bf16_fallback_correctness(batch_size):
    uses topk=2048 which satisfies this. We use topk=256 for tractability.
    """
    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
-        _flashmla_sparse_available,
        _FLASHMLA_SM90_HEAD_PADDING,
+        _flashmla_sparse_available,
    )

    if not _flashmla_sparse_available:
@@ -1002,8 +1031,9 @@ def test_flashmla_bf16_fallback_correctness(batch_size):

    # Create combined KV cache: [kv_pool_size, 1, 576]
    # Layout: [kv_c_normed(512) | k_pe(64)]
-    kv_combined = torch.randn(kv_pool_size, 1, HEADDIM_V + HEADDIM_QK,
-                              dtype=torch.bfloat16, device=device)
+    kv_combined = torch.randn(
+        kv_pool_size, 1, HEADDIM_V + HEADDIM_QK, dtype=torch.bfloat16, device=device
+    )

    # Create page table with unique random indices per token
    page_table = torch.stack(
@@ -1019,7 +1049,7 @@ def test_flashmla_bf16_fallback_correctness(batch_size):

    # 2) Pad heads to 64
    padded_heads = _FLASHMLA_SM90_HEAD_PADDING
-    if NUM_HEADS < padded_heads:
+    if padded_heads > NUM_HEADS:
        q_padded = q_concat.new_zeros((T, padded_heads, q_concat.shape[-1]))
        q_padded[:, :NUM_HEADS, :] = q_concat
        q_concat = q_padded
@@ -1029,11 +1059,11 @@ def test_flashmla_bf16_fallback_correctness(batch_size):

    # 4) Call FlashMLA BF16 sparse prefill kernel
    output = flash_mla_sparse_fwd(
-        q_concat,                  # [T, padded_heads, 576]
-        kv_combined,               # [kv_pool_size, 1, 576]
-        indices,                   # [T, 1, topk]
-        SOFTMAX_SCALE,             # 192**-0.5
-        d_v=HEADDIM_V,             # 512
+        q_concat,  # [T, padded_heads, 576]
+        kv_combined,  # [kv_pool_size, 1, 576]
+        indices,  # [T, 1, topk]
+        SOFTMAX_SCALE,  # 192**-0.5
+        d_v=HEADDIM_V,  # 512
        topk_length=valid_counts,  # [T]
    )[0]

@@ -1043,16 +1073,20 @@ def test_flashmla_bf16_fallback_correctness(batch_size):
    assert output.shape == (T, NUM_HEADS, HEADDIM_V), (
        f"Output shape mismatch: {output.shape} vs expected ({T}, {NUM_HEADS}, {HEADDIM_V})"
    )
-    assert not output.isnan().any(), f"NaN in FlashMLA fallback output for bs={batch_size}"
-    assert not output.isinf().any(), f"Inf in FlashMLA fallback output for bs={batch_size}"
+    assert not output.isnan().any(), (
+        f"NaN in FlashMLA fallback output for bs={batch_size}"
+    )
+    assert not output.isinf().any(), (
+        f"Inf in FlashMLA fallback output for bs={batch_size}"
+    )

    # 6) Verify correctness against SDPA reference for first 4 tokens
    for b in range(min(4, T)):
        idx = page_table[b]  # [topk]
        # Gather KV from combined cache: [topk, 1, 576] -> split
        kv_gathered = kv_combined[idx].squeeze(1)  # [topk, 576]
-        v_gathered = kv_gathered[:, :HEADDIM_V]    # [topk, 512] (kv_c_normed)
-        k_gathered = kv_gathered                    # [topk, 576] (full key)
+        v_gathered = kv_gathered[:, :HEADDIM_V]  # [topk, 512] (kv_c_normed)
+        k_gathered = kv_gathered  # [topk, 576] (full key)

        # Full Q for this token
        q_full = torch.cat([q_nope[b], q_rope[b]], dim=-1)  # [N, 576]
@@ -1061,16 +1095,23 @@ def test_flashmla_bf16_fallback_correctness(batch_size):
        k_expanded = k_gathered.unsqueeze(1).expand(-1, NUM_HEADS, -1)  # [topk, N, 576]
        v_expanded = v_gathered.unsqueeze(1).expand(-1, NUM_HEADS, -1)  # [topk, N, 512]

-        ref_out = F.scaled_dot_product_attention(
-            q_full.unsqueeze(0).unsqueeze(2),                    # [1, N, 1, 576]
-            k_expanded.transpose(0, 1).unsqueeze(0),             # [1, N, topk, 576]
-            v_expanded.transpose(0, 1).unsqueeze(0),             # [1, N, topk, 512]
-            scale=SOFTMAX_SCALE,
-        ).squeeze(0).squeeze(1)  # [N, 512]
+        ref_out = (
+            F.scaled_dot_product_attention(
+                q_full.unsqueeze(0).unsqueeze(2),  # [1, N, 1, 576]
+                k_expanded.transpose(0, 1).unsqueeze(0),  # [1, N, topk, 576]
+                v_expanded.transpose(0, 1).unsqueeze(0),  # [1, N, topk, 512]
+                scale=SOFTMAX_SCALE,
+            )
+            .squeeze(0)
+            .squeeze(1)
+        )  # [N, 512]

        torch.testing.assert_close(
-            output[b].float(), ref_out.float(), rtol=0.02, atol=0.02,
-            msg=f"FlashMLA fallback (bs={batch_size}) token {b} mismatch"
+            output[b].float(),
+            ref_out.float(),
+            rtol=0.02,
+            atol=0.02,
+            msg=f"FlashMLA fallback (bs={batch_size}) token {b} mismatch",
        )


@@ -1092,8 +1133,8 @@ def test_fa3_vs_flashmla_cross_kernel_consistency(batch_size):
    kernels agree at the exact gating threshold.
    """
    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
-        _flashmla_sparse_available,
        _FLASHMLA_SM90_HEAD_PADDING,
+        _flashmla_sparse_available,
    )

    if not _flashmla_sparse_available:
@@ -1112,8 +1153,9 @@ def test_fa3_vs_flashmla_cross_kernel_consistency(batch_size):

    # Create combined KV cache in BF16: [kv_pool_size, 576]
    # This flat format is what both kernels see after reshaping
-    kv_flat = torch.randn(kv_pool_size, HEADDIM_V + HEADDIM_QK,
-                          dtype=torch.bfloat16, device=device)
+    kv_flat = torch.randn(
+        kv_pool_size, HEADDIM_V + HEADDIM_QK, dtype=torch.bfloat16, device=device
+    )

    # Create page table with unique random indices per token
    page_table = torch.stack(
@@ -1147,7 +1189,7 @@ def test_fa3_vs_flashmla_cross_kernel_consistency(batch_size):
    q_concat = torch.cat([q_nope, q_rope], dim=-1)  # [T, N, 576]

    padded_heads = _FLASHMLA_SM90_HEAD_PADDING
-    if NUM_HEADS < padded_heads:
+    if padded_heads > NUM_HEADS:
        q_padded = q_concat.new_zeros((T, padded_heads, q_concat.shape[-1]))
        q_padded[:, :NUM_HEADS, :] = q_concat
        q_concat = q_padded
@@ -1175,7 +1217,9 @@ def test_fa3_vs_flashmla_cross_kernel_consistency(batch_size):
    # Cross-kernel comparison with slightly relaxed tolerance
    # (different kernels may have different accumulation order)
    torch.testing.assert_close(
-        fa3_out.float(), flashmla_out.float(),
-        rtol=0.03, atol=0.03,
-        msg=f"FA3 vs FlashMLA mismatch at bs={batch_size}"
+        fa3_out.float(),
+        flashmla_out.float(),
+        rtol=0.03,
+        atol=0.03,
+        msg=f"FA3 vs FlashMLA mismatch at bs={batch_size}",
    )
@@ -8,6 +8,7 @@ Tests verify:
  - KV cache write/read consistency
  - Backend registration and selection
 """
+
 import pytest
 import torch

@@ -133,9 +134,7 @@ def test_kv_cache_write_read_consistency():

    # Verify consistency
    torch.testing.assert_close(c_kv_read, kv_c_normed, rtol=1e-3, atol=1e-3)
-    torch.testing.assert_close(
-        k_rope_read, k_pe.squeeze(1), rtol=1e-3, atol=1e-3
-    )
+    torch.testing.assert_close(k_rope_read, k_pe.squeeze(1), rtol=1e-3, atol=1e-3)


 def test_kv_cache_dtype_auto():
@@ -151,8 +150,7 @@ def test_kv_cache_dtype_auto():
    from vllm import _custom_ops as ops

    ops.concat_and_cache_mla(
-        kv_c, k_pe.squeeze(1), cache, slot_mapping,
-        kv_cache_dtype="auto", scale=k_scale
+        kv_c, k_pe.squeeze(1), cache, slot_mapping, kv_cache_dtype="auto", scale=k_scale
    )

    assert cache.dtype == torch.bfloat16
@@ -163,9 +161,6 @@ def test_kv_cache_dtype_auto():

 def test_empty_kv_cache():
    """Verify do_kv_cache_update handles empty cache gracefully."""
-    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
-        CutlassFA3MLASparseImpl,
-    )

    kv_cache = torch.empty(0, device="cuda")
    # Should return without error (numel() == 0 check)
@@ -203,9 +198,9 @@ def test_triton_convert_valid_counts():
    block_size = 64

    req_id = torch.zeros(T, dtype=torch.int32, device=device)
-    block_table = torch.arange(
-        num_blocks, dtype=torch.int32, device=device
-    ).unsqueeze(0)  # [1, num_blocks]
+    block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
+        0
+    )  # [1, num_blocks]

    # Create topk_indices with varying valid entries per token
    topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
@@ -214,12 +209,17 @@ def test_triton_convert_valid_counts():
        nv = expected_valid[i]
        # Use indices within the valid range
        topk_indices[i, :nv] = torch.randint(
-            0, num_blocks * block_size, (nv,),
-            dtype=torch.int32, device=device,
+            0,
+            num_blocks * block_size,
+            (nv,),
+            dtype=torch.int32,
+            device=device,
        )

    global_idx, valid_counts = triton_convert_req_index_to_global_index(
-        req_id, block_table, topk_indices,
+        req_id,
+        block_table,
+        topk_indices,
        BLOCK_SIZE=block_size,
        NUM_TOPK_TOKENS=topk,
        return_valid_counts=True,
@@ -267,9 +267,7 @@ def test_prefill_cache_seqlens_vs_valid_counts():
    seg_lens = np.diff(starts)  # [4]
    seq_lens_np = np.array([seq_len], dtype=np.int32)

-    per_tok_seqlens = np.minimum(
-        np.repeat(seq_lens_np, seg_lens), topk
-    )  # [4, 4, 4, 4]
+    per_tok_seqlens = np.minimum(np.repeat(seq_lens_np, seg_lens), topk)  # [4, 4, 4, 4]

    # This is what the metadata builder produces:
    assert all(per_tok_seqlens == 4), (
@@ -288,9 +286,7 @@ def test_prefill_cache_seqlens_vs_valid_counts():
    )

    req_id = torch.zeros(T, dtype=torch.int32, device=device)
-    block_table = torch.arange(
-        32, dtype=torch.int32, device=device
-    ).unsqueeze(0)
+    block_table = torch.arange(32, dtype=torch.int32, device=device).unsqueeze(0)

    topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
    for i in range(T):
@@ -298,7 +294,9 @@ def test_prefill_cache_seqlens_vs_valid_counts():
        topk_indices[i, :nv] = torch.arange(nv, dtype=torch.int32, device=device)

    _, valid_counts = triton_convert_req_index_to_global_index(
-        req_id, block_table, topk_indices,
+        req_id,
+        block_table,
+        topk_indices,
        BLOCK_SIZE=64,
        NUM_TOPK_TOKENS=topk,
        return_valid_counts=True,
@@ -321,7 +319,8 @@ def test_global_idx_clamp_safety():
    # Create a page_table with -1 entries
    page_table = torch.tensor(
        [[5, 10, -1, -1], [3, -1, -1, -1]],
-        dtype=torch.int32, device=device,
+        dtype=torch.int32,
+        device=device,
    )

    # Clamp -1 to 0
@@ -330,7 +329,8 @@ def test_global_idx_clamp_safety():
    # Verify
    expected = torch.tensor(
        [[5, 10, 0, 0], [3, 0, 0, 0]],
-        dtype=torch.int32, device=device,
+        dtype=torch.int32,
+        device=device,
    )
    assert torch.equal(clamped, expected), (
        f"Clamped page_table doesn't match expected: {clamped} vs {expected}"
@@ -351,7 +351,8 @@ def test_inplace_clamp_no_negative_indices():
    # Create a global_idx tensor with -1 entries
    global_idx = torch.tensor(
        [[100, 200, -1, -1, -1], [50, -1, -1, -1, -1]],
-        dtype=torch.int32, device=device,
+        dtype=torch.int32,
+        device=device,
    )

    # In-place clamp
@@ -390,9 +391,9 @@ def test_full_fix_flow_valid_counts_and_clamp():
    block_size = 64

    req_id = torch.zeros(T, dtype=torch.int32, device=device)
-    block_table = torch.arange(
-        num_blocks, dtype=torch.int32, device=device
-    ).unsqueeze(0)
+    block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).unsqueeze(
+        0
+    )

    # Simulate causal prefill: token i has (i+1) valid entries
    topk_indices = torch.full((T, topk), -1, dtype=torch.int32, device=device)
@@ -403,7 +404,9 @@ def test_full_fix_flow_valid_counts_and_clamp():

    # Step 1: Convert with valid counts
    global_idx, valid_counts = triton_convert_req_index_to_global_index(
-        req_id, block_table, topk_indices,
+        req_id,
+        block_table,
+        topk_indices,
        BLOCK_SIZE=block_size,
        NUM_TOPK_TOKENS=topk,
        return_valid_counts=True,
@@ -456,6 +459,7 @@ def test_full_fix_flow_valid_counts_and_clamp():
 def _make_mock_vllm_config(max_tokens=512):
    """Create a mock VllmConfig for metadata builder tests."""
    from unittest.mock import MagicMock
+
    vllm_config = MagicMock()
    vllm_config.scheduler_config.max_num_batched_tokens = max_tokens
    vllm_config.speculative_config = None
@@ -473,7 +477,6 @@ def test_metadata_builder_cuda_graph_padding():
    This happens when num_actual_tokens=32 (padded for CUDA graph)
    but only 31 real tokens exist (one request completed mid-batch).
    """
-    import numpy as np
    from unittest.mock import MagicMock

    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
@@ -569,8 +572,7 @@ def test_metadata_builder_cuda_graph_padding():
    # Verify cu_seqlens_q is [0, 1, 2, ..., padded_T] (always correct)
    for i in range(padded_T + 1):
        assert metadata.cu_seqlens_q[i].item() == i, (
-            f"cu_seqlens_q[{i}] should be {i}, "
-            f"got {metadata.cu_seqlens_q[i].item()}"
+            f"cu_seqlens_q[{i}] should be {i}, got {metadata.cu_seqlens_q[i].item()}"
        )

    # Verify cu_seqlens_k is monotonically non-decreasing
@@ -589,11 +591,11 @@ def test_metadata_builder_cuda_graph_padding():
@pytest.mark.parametrize(
    "real_tokens,padded_T",
    [
-        (1, 2),      # minimal padding
-        (3, 32),     # large padding gap
-        (7, 8),      # small batch
-        (15, 16),    # medium batch
-        (31, 32),    # the exact crash scenario
+        (1, 2),  # minimal padding
+        (3, 32),  # large padding gap
+        (7, 8),  # small batch
+        (15, 16),  # medium batch
+        (31, 32),  # the exact crash scenario
        (100, 104),  # larger padding gap
    ],
 )
@@ -604,7 +606,6 @@ def test_metadata_builder_cuda_graph_padding_various(real_tokens, padded_T):
    - query_start_loc_cpu has num_reqs_padded+1 entries (with padded suffix)
    - seq_lens_cpu has num_reqs_padded entries (with stale padding entries)
    """
-    import numpy as np
    from unittest.mock import MagicMock

    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
@@ -691,7 +692,6 @@ def test_metadata_builder_cuda_graph_padding_various(real_tokens, padded_T):

 def test_metadata_builder_no_padding():
    """Verify build() still works correctly when T == actual_tokens (no padding)."""
-    import numpy as np
    from unittest.mock import MagicMock

    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
@@ -759,7 +759,6 @@ def test_metadata_builder_mixed_prefill_decode_with_padding():
    This tests a more complex scenario: 2 decode tokens + 3 prefill tokens
    from 3 requests, padded from 5 to 8 tokens.
    """
-    import numpy as np
    from unittest.mock import MagicMock

    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
@@ -833,11 +832,11 @@ def test_metadata_builder_mixed_prefill_decode_with_padding():
    assert metadata.req_id_per_token[7].item() == 0

    # Verify cache_seqlens
-    assert metadata.cache_seqlens[0].item() == 100   # req0 seq_len
-    assert metadata.cache_seqlens[1].item() == 200   # req1 seq_len
-    assert metadata.cache_seqlens[2].item() == 3     # req2 seq_len
-    assert metadata.cache_seqlens[3].item() == 3     # req2 seq_len
-    assert metadata.cache_seqlens[4].item() == 3     # req2 seq_len
+    assert metadata.cache_seqlens[0].item() == 100  # req0 seq_len
+    assert metadata.cache_seqlens[1].item() == 200  # req1 seq_len
+    assert metadata.cache_seqlens[2].item() == 3  # req2 seq_len
+    assert metadata.cache_seqlens[3].item() == 3  # req2 seq_len
+    assert metadata.cache_seqlens[4].item() == 3  # req2 seq_len
    # Padding (default = 1)
    assert metadata.cache_seqlens[5].item() >= 1
    assert metadata.cache_seqlens[6].item() >= 1
@@ -855,7 +854,6 @@ def test_metadata_builder_zero_real_tokens():
    This edge case can occur during CUDA graph warmup or capture where
    dummy batches may have zero real tokens but T > 0 (padded size).
    """
-    import numpy as np
    from unittest.mock import MagicMock

    from vllm.v1.attention.backends.mla.cutlass_fa3_sparse import (
@@ -899,9 +897,7 @@ def test_metadata_builder_zero_real_tokens():
        query_start_loc_cpu, dtype=torch.int32, device=device
    )
    cm.slot_mapping = torch.zeros(padded_T, dtype=torch.int64, device=device)
-    cm.block_table_tensor = torch.zeros(
-        1, 4, dtype=torch.int32, device=device
-    )
+    cm.block_table_tensor = torch.zeros(1, 4, dtype=torch.int32, device=device)

    # Should NOT raise any errors
    metadata = builder.build(
@@ -964,9 +960,9 @@ def test_forward_mqa_has_fa3_and_fallback_methods():
        CutlassFA3MLASparseImpl,
    )

-    assert hasattr(CutlassFA3MLASparseImpl, '_forward_fa3'), (
+    assert hasattr(CutlassFA3MLASparseImpl, "_forward_fa3"), (
        "CutlassFA3MLASparseImpl should have _forward_fa3 method"
    )
-    assert hasattr(CutlassFA3MLASparseImpl, '_forward_flashmla_bf16_fallback'), (
+    assert hasattr(CutlassFA3MLASparseImpl, "_forward_flashmla_bf16_fallback"), (
        "CutlassFA3MLASparseImpl should have _forward_flashmla_bf16_fallback method"
    )
@@ -20,8 +20,10 @@ This simplifies metadata building and CUDA graph support.
 Backend priority: Highest for SM90 with kv_cache_dtype="auto".
 Graceful fallback to FlashMLA Sparse when FP8 cache requested or non-SM90.
 """
+
 from __future__ import annotations

+import logging
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, ClassVar

@@ -39,8 +41,6 @@ from vllm.v1.attention.backend import (
 )
 from vllm.v1.attention.ops.cutlass_fa3 import is_cutlass_fa3_available

-import logging
-
 logger = logging.getLogger(__name__)

 # Maximum batch size (number of tokens) for which CUTLASS FA3 is used.
@@ -58,6 +58,7 @@ _FLASHMLA_SM90_HEAD_PADDING = 64
 _flashmla_sparse_available = False
 try:
    from vllm.v1.attention.ops.flashmla import flash_mla_sparse_fwd
+
    _flashmla_sparse_available = True
 except (ImportError, Exception):
    pass
@@ -113,9 +114,7 @@ class CutlassFA3MLASparseBackend(AttentionBackend):
        return True

    @classmethod
-    def supports_compute_capability(
-        cls, capability: DeviceCapability
-    ) -> bool:
+    def supports_compute_capability(cls, capability: DeviceCapability) -> bool:
        return capability.major == 9

    @staticmethod
@@ -161,9 +160,7 @@ class CutlassFA3MLASparseBackend(AttentionBackend):
            use_non_causal=use_non_causal,
        )
        if not is_cutlass_fa3_available():
-            invalid.append(
-                "_cutlass_fa3_C not available (requires CUDA >= 12.4, SM90)"
-            )
+            invalid.append("_cutlass_fa3_C not available (requires CUDA >= 12.4, SM90)")
        return invalid


@@ -189,7 +186,7 @@ class CutlassFA3MLASparseMetadata(AttentionMetadata):
    num_actual_tokens: int
    query_start_loc: torch.Tensor
    slot_mapping: torch.Tensor
-    block_table: torch.Tensor      # [num_reqs, max_blocks_per_req] int32
+    block_table: torch.Tensor  # [num_reqs, max_blocks_per_req] int32
    req_id_per_token: torch.Tensor  # [T] int32

    block_size: int = 64
@@ -197,8 +194,8 @@ class CutlassFA3MLASparseMetadata(AttentionMetadata):

    # FA3-specific metadata (pre-allocated for CUDA graph safety)
    cache_seqlens: torch.Tensor | None = None  # [T] int32
-    cu_seqlens_q: torch.Tensor | None = None   # [T+1] int32
-    cu_seqlens_k: torch.Tensor | None = None   # [T+1] int32
+    cu_seqlens_q: torch.Tensor | None = None  # [T+1] int32
+    cu_seqlens_k: torch.Tensor | None = None  # [T+1] int32

    # For MLAAttention.forward_impl() routing: sparse -> all MQA
    # Setting num_decodes = num_reqs ensures all tokens go through
@@ -223,9 +220,7 @@ class CutlassFA3MLASparseMetadataBuilder(
    - Uses in-place .copy_() for buffer updates (safe for CUDA graph replay)
    """

-    _cudagraph_support: ClassVar[AttentionCGSupport] = (
-        AttentionCGSupport.UNIFORM_BATCH
-    )
+    _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

    def __init__(
        self,
@@ -244,9 +239,7 @@ class CutlassFA3MLASparseMetadataBuilder(

        # Pre-allocate GPU buffers (persist across CUDA graph replays).
        # These are updated in-place via .copy_() before each replay.
-        self.req_id_buf = torch.zeros(
-            max_tokens, dtype=torch.int32, device=device
-        )
+        self.req_id_buf = torch.zeros(max_tokens, dtype=torch.int32, device=device)
        self.cache_seqlens_buf = torch.ones(
            max_tokens, dtype=torch.int32, device=device
        )
@@ -277,9 +270,7 @@ class CutlassFA3MLASparseMetadataBuilder(
        seg_lens = np.diff(starts)

        # req_id_per_token: map each token -> request index
-        req_ids = np.repeat(
-            np.arange(len(seg_lens), dtype=np.int32), seg_lens
-        )
+        req_ids = np.repeat(np.arange(len(seg_lens), dtype=np.int32), seg_lens)
        # CUDA graph padding fix: T = cm.num_actual_tokens may include
        # padding tokens (e.g., T=32 when only 31 real tokens exist).
        # The computed req_ids array has sum(seg_lens) elements which
@@ -304,9 +295,7 @@ class CutlassFA3MLASparseMetadataBuilder(
        # from the index conversion kernel, which correctly reflects
        # the number of valid KV entries per token.
        seq_lens_np = np.asarray(cm.seq_lens_cpu, dtype=np.int32)
-        per_tok_seqlens = np.minimum(
-            np.repeat(seq_lens_np, seg_lens), self.topk_tokens
-        )
+        per_tok_seqlens = np.minimum(np.repeat(seq_lens_np, seg_lens), self.topk_tokens)
        # Same CUDA graph padding fix: zero-fill then copy actual data.
        # Default to 1 (safe minimum seqlen for FA3 kernel).
        self.cache_seqlens_buf.fill_(1)
@@ -349,9 +338,7 @@ class CutlassFA3MLASparseMetadataBuilder(
 # ─── Implementation ───────────────────────────────────────────────────


-class CutlassFA3MLASparseImpl(
-    SparseMLAAttentionImpl[CutlassFA3MLASparseMetadata]
-):
+class CutlassFA3MLASparseImpl(SparseMLAAttentionImpl[CutlassFA3MLASparseMetadata]):
    """CUTLASS FA3 sparse MLA attention implementation.

    This implementation replaces the FlashMLA C sparse_attn_fwd_kernel
@@ -386,24 +373,22 @@ class CutlassFA3MLASparseImpl(
        qk_rope_head_dim: int = 64,
        qk_head_dim: int = 192,
        v_head_dim: int = 128,
-        kv_b_proj: "ColumnParallelLinear | None" = None,
+        kv_b_proj: ColumnParallelLinear | None = None,
        indexer: object | None = None,
        q_pad_num_heads: int | None = None,
        **kwargs,
    ) -> None:
-        self.num_heads = num_heads          # 16 (per GPU for TP=8)
-        self.head_size = head_size          # 576 (kv_lora_rank + qk_rope_head_dim)
-        self.scale = float(scale)           # 192**-0.5
-        self.num_kv_heads = num_kv_heads    # 1 (MQA)
+        self.num_heads = num_heads  # 16 (per GPU for TP=8)
+        self.head_size = head_size  # 576 (kv_lora_rank + qk_rope_head_dim)
+        self.scale = float(scale)  # 192**-0.5
+        self.num_kv_heads = num_kv_heads  # 1 (MQA)
        self.kv_cache_dtype = kv_cache_dtype  # "auto" (maps to BF16)
        self.kv_lora_rank = kv_lora_rank
        self.qk_rope_head_dim = qk_rope_head_dim
        self.softmax_scale = scale
        self.topk_tokens = 2048
        self.num_splits = 0  # auto; CUDA-graph safe (deterministic per bs)
-        self.logits_soft_cap = (
-            float(logits_soft_cap) if logits_soft_cap else 0.0
-        )
+        self.logits_soft_cap = float(logits_soft_cap) if logits_soft_cap else 0.0

        # The indexer provides topk_indices_buffer shared across layers
        assert indexer is not None, (
@@ -450,8 +435,8 @@ class CutlassFA3MLASparseImpl(
        if isinstance(q, tuple):
            ql_nope, q_pe = q  # [T, N, 512], [T, N, 64]
        else:
-            ql_nope = q[..., : self.kv_lora_rank]   # [T, N, 512]
-            q_pe = q[..., self.kv_lora_rank :]       # [T, N, 64]
+            ql_nope = q[..., : self.kv_lora_rank]  # [T, N, 512]
+            q_pe = q[..., self.kv_lora_rank :]  # [T, N, 64]
        T = ql_nope.shape[0]

        # 2) Convert topk_indices -> global cache slot indices
@@ -461,11 +446,11 @@ class CutlassFA3MLASparseImpl(
        )

        global_idx, valid_counts = triton_convert_req_index_to_global_index(
-            attn_metadata.req_id_per_token,     # [T] int32
-            attn_metadata.block_table,           # [R, max_blocks] int32
-            self.topk_indices_buffer[:T],        # [T, 2048] int32
+            attn_metadata.req_id_per_token,  # [T] int32
+            attn_metadata.block_table,  # [R, max_blocks] int32
+            self.topk_indices_buffer[:T],  # [T, 2048] int32
            BLOCK_SIZE=attn_metadata.block_size,  # 64
-            NUM_TOPK_TOKENS=self.topk_tokens,    # 2048
+            NUM_TOPK_TOKENS=self.topk_tokens,  # 2048
            return_valid_counts=True,
        )
        # global_idx: [T, 2048] int32 — flat cache slot IDs
@@ -488,13 +473,20 @@ class CutlassFA3MLASparseImpl(
        use_fa3 = (T <= MAX_BATCH_SIZE_FOR_FA3) or not _flashmla_sparse_available
        if use_fa3:
            attn_out = self._forward_fa3(
-                ql_nope, q_pe, kv_c_and_k_pe_cache,
-                global_idx, cache_seqlens, attn_metadata,
+                ql_nope,
+                q_pe,
+                kv_c_and_k_pe_cache,
+                global_idx,
+                cache_seqlens,
+                attn_metadata,
            )
        else:
            attn_out = self._forward_flashmla_bf16_fallback(
-                ql_nope, q_pe, kv_c_and_k_pe_cache,
-                global_idx, cache_seqlens,
+                ql_nope,
+                q_pe,
+                kv_c_and_k_pe_cache,
+                global_idx,
+                cache_seqlens,
            )

        # Output: [T, N, 512] — already 3D
@@ -502,11 +494,11 @@ class CutlassFA3MLASparseImpl(

    def _forward_fa3(
        self,
-        ql_nope: torch.Tensor,          # [T, N, 512]
-        q_pe: torch.Tensor,             # [T, N, 64]
+        ql_nope: torch.Tensor,  # [T, N, 512]
+        q_pe: torch.Tensor,  # [T, N, 64]
        kv_c_and_k_pe_cache: torch.Tensor,
-        global_idx: torch.Tensor,       # [T, 2048]
-        cache_seqlens: torch.Tensor,    # [T]
+        global_idx: torch.Tensor,  # [T, 2048]
+        cache_seqlens: torch.Tensor,  # [T]
        attn_metadata: CutlassFA3MLASparseMetadata,
    ) -> torch.Tensor:
        """CUTLASS FA3 kernel path — fast for small batch sizes (bs<=16).
@@ -529,16 +521,16 @@ class CutlassFA3MLASparseImpl(
        from vllm.v1.attention.ops.cutlass_fa3 import flash_attn_with_kvcache

        attn_out = flash_attn_with_kvcache(
-            q=q_pe,                                     # [T, N, 64]
-            k_cache=k_rope,                             # [S, 1, 1, 64]
-            v_cache=c_kv,                               # [S, 1, 1, 512]
-            qv=ql_nope,                                 # [T, N, 512]
-            page_table=global_idx,                      # [T, 2048]
-            cache_seqlens=cache_seqlens,                # [T]
-            cu_seqlens_q=attn_metadata.cu_seqlens_q,    # [T+1]
+            q=q_pe,  # [T, N, 64]
+            k_cache=k_rope,  # [S, 1, 1, 64]
+            v_cache=c_kv,  # [S, 1, 1, 512]
+            qv=ql_nope,  # [T, N, 512]
+            page_table=global_idx,  # [T, 2048]
+            cache_seqlens=cache_seqlens,  # [T]
+            cu_seqlens_q=attn_metadata.cu_seqlens_q,  # [T+1]
            cu_seqlens_k_new=None,
            max_seqlen_q=1,
-            softmax_scale=self.softmax_scale,           # 192**-0.5
+            softmax_scale=self.softmax_scale,  # 192**-0.5
            causal=True,
            window_size=(-1, -1),
            softcap=self.logits_soft_cap,
@@ -548,11 +540,11 @@ class CutlassFA3MLASparseImpl(

    def _forward_flashmla_bf16_fallback(
        self,
-        ql_nope: torch.Tensor,          # [T, N, 512]
-        q_pe: torch.Tensor,             # [T, N, 64]
+        ql_nope: torch.Tensor,  # [T, N, 512]
+        q_pe: torch.Tensor,  # [T, N, 64]
        kv_c_and_k_pe_cache: torch.Tensor,
-        global_idx: torch.Tensor,       # [T, 2048]
-        cache_seqlens: torch.Tensor,    # [T]
+        global_idx: torch.Tensor,  # [T, 2048]
+        cache_seqlens: torch.Tensor,  # [T]
    ) -> torch.Tensor:
        """FlashMLA BF16 sparse prefill fallback — for larger batch sizes.

@@ -579,10 +571,8 @@ class CutlassFA3MLASparseImpl(

        # 2) Pad heads to _FLASHMLA_SM90_HEAD_PADDING (64 for SM90)
        padded_heads = _FLASHMLA_SM90_HEAD_PADDING
-        if N < padded_heads:
-            q_padded = q_concat.new_zeros(
-                (T, padded_heads, q_concat.shape[-1])
-            )
+        if padded_heads > N:
+            q_padded = q_concat.new_zeros((T, padded_heads, q_concat.shape[-1]))
            q_padded[:, :N, :] = q_concat
            q_concat = q_padded

@@ -602,12 +592,12 @@ class CutlassFA3MLASparseImpl(
        #    from processing these clamped entries, saving compute and ensuring
        #    correctness.
        output = flash_mla_sparse_fwd(
-            q_concat,                  # [T, padded_heads, 576]
-            kv,                        # [S, 1, 576]
-            indices,                   # [T, 1, 2048]
-            self.softmax_scale,        # 192**-0.5
-            d_v=self.kv_lora_rank,     # 512
-            topk_length=cache_seqlens, # [T] valid entry counts
+            q_concat,  # [T, padded_heads, 576]
+            kv,  # [S, 1, 576]
+            indices,  # [T, 1, 2048]
+            self.softmax_scale,  # 192**-0.5
+            d_v=self.kv_lora_rank,  # 512
+            topk_length=cache_seqlens,  # [T] valid entry counts
        )[0]  # extract output tensor from (output, max_logits, lse) tuple

        # 6) Unpad heads: (T, padded_heads, 512) -> (T, N, 512)
@@ -74,8 +74,7 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
        "vllm.v1.attention.backends.mla.flashmla_sparse.FlashMLASparseBackend"
    )
    CUTLASS_FA3_MLA_SPARSE = (
-        "vllm.v1.attention.backends.mla.cutlass_fa3_sparse."
-        "CutlassFA3MLASparseBackend"
+        "vllm.v1.attention.backends.mla.cutlass_fa3_sparse.CutlassFA3MLASparseBackend"
    )
    FLASH_ATTN_MLA = "vllm.v1.attention.backends.mla.flashattn_mla.FlashAttnMLABackend"
    NO_ATTENTION = "vllm.v1.attention.backends.no_attention.NoAttentionBackend"
@@ -17,7 +17,6 @@ The FA3 kernel supports MLA (Multi-head Latent Attention) with:

 Source: https://github.com/sgl-project/sgl-attn (commit bcf72ccc)
 """
-from typing import Optional

 import torch

@@ -45,34 +44,34 @@ def flash_attn_with_kvcache(
    q: torch.Tensor,
    k_cache: torch.Tensor,
    v_cache: torch.Tensor,
-    k: Optional[torch.Tensor] = None,
-    v: Optional[torch.Tensor] = None,
-    qv: Optional[torch.Tensor] = None,
-    rotary_cos: Optional[torch.Tensor] = None,
-    rotary_sin: Optional[torch.Tensor] = None,
-    cache_seqlens: Optional[torch.Tensor] = None,
-    cache_batch_idx: Optional[torch.Tensor] = None,
-    cache_leftpad: Optional[torch.Tensor] = None,
-    page_table: Optional[torch.Tensor] = None,
-    cu_seqlens_q: Optional[torch.Tensor] = None,
-    cu_seqlens_k_new: Optional[torch.Tensor] = None,
-    max_seqlen_q: Optional[int] = None,
-    rotary_seqlens: Optional[torch.Tensor] = None,
-    q_descale: Optional[torch.Tensor] = None,
-    k_descale: Optional[torch.Tensor] = None,
-    v_descale: Optional[torch.Tensor] = None,
-    softmax_scale: Optional[float] = None,
+    k: torch.Tensor | None = None,
+    v: torch.Tensor | None = None,
+    qv: torch.Tensor | None = None,
+    rotary_cos: torch.Tensor | None = None,
+    rotary_sin: torch.Tensor | None = None,
+    cache_seqlens: torch.Tensor | None = None,
+    cache_batch_idx: torch.Tensor | None = None,
+    cache_leftpad: torch.Tensor | None = None,
+    page_table: torch.Tensor | None = None,
+    cu_seqlens_q: torch.Tensor | None = None,
+    cu_seqlens_k_new: torch.Tensor | None = None,
+    max_seqlen_q: int | None = None,
+    rotary_seqlens: torch.Tensor | None = None,
+    q_descale: torch.Tensor | None = None,
+    k_descale: torch.Tensor | None = None,
+    v_descale: torch.Tensor | None = None,
+    softmax_scale: float | None = None,
    causal: bool = False,
    window_size: tuple[int, int] = (-1, -1),
-    attention_chunk: Optional[int] = None,
+    attention_chunk: int | None = None,
    softcap: float = 0.0,
    rotary_interleaved: bool = True,
-    scheduler_metadata: Optional[torch.Tensor] = None,
+    scheduler_metadata: torch.Tensor | None = None,
    num_splits: int = 0,
-    pack_gqa: Optional[bool] = None,
+    pack_gqa: bool | None = None,
    sm_margin: int = 0,
    return_softmax_lse: bool = False,
-    sinks: Optional[torch.Tensor] = None,
+    sinks: torch.Tensor | None = None,
 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """CUTLASS FA3 attention with paged KV cache for MLA.

@@ -124,40 +123,40 @@ def flash_attn_with_kvcache(
    attention_chunk_val = 0 if attention_chunk is None else int(attention_chunk)

    out, softmax_lse, *rest = torch.ops._cutlass_fa3_C.fwd.default(
-        q,                                           # 0: q
-        k_cache,                                     # 1: k (paged KV cache)
-        v_cache,                                     # 2: v (paged KV cache)
-        k,                                           # 3: k_new
-        v,                                           # 4: v_new
-        qv,                                          # 5: q_v (MLA NoPE query)
-        None,                                        # 6: out buffer
-        cu_seqlens_q,                                # 7: cu_seqlens_q
-        None,                                        # 8: cu_seqlens_k
-        cu_seqlens_k_new,                            # 9: cu_seqlens_k_new
-        None,                                        # 10: seqused_q
-        cache_seqlens,                               # 11: seqused_k
-        max_seqlen_q,                                # 12: max_seqlen_q
-        None,                                        # 13: max_seqlen_k
-        page_table,                                  # 14: page_table
-        cache_batch_idx,                             # 15: kv_batch_idx
-        cache_leftpad,                               # 16: leftpad_k
-        rotary_cos,                                  # 17: rotary_cos
-        rotary_sin,                                  # 18: rotary_sin
-        rotary_seqlens,                              # 19: seqlens_rotary
-        q_descale,                                   # 20: q_descale
-        k_descale,                                   # 21: k_descale
-        v_descale,                                   # 22: v_descale
-        softmax_scale,                               # 23: softmax_scale
-        causal,                                      # 24: is_causal
-        window_size[0],                              # 25: window_size_left
-        window_size[1],                              # 26: window_size_right
-        attention_chunk_val,                         # 27: attention_chunk
-        softcap,                                     # 28: softcap
-        rotary_interleaved,                          # 29: is_rotary_interleaved
-        scheduler_metadata,                          # 30: scheduler_metadata
-        num_splits,                                  # 31: num_splits
-        pack_gqa,                                    # 32: pack_gqa
-        sm_margin,                                   # 33: sm_margin
-        sinks,                                       # 34: sinks
+        q,  # 0: q
+        k_cache,  # 1: k (paged KV cache)
+        v_cache,  # 2: v (paged KV cache)
+        k,  # 3: k_new
+        v,  # 4: v_new
+        qv,  # 5: q_v (MLA NoPE query)
+        None,  # 6: out buffer
+        cu_seqlens_q,  # 7: cu_seqlens_q
+        None,  # 8: cu_seqlens_k
+        cu_seqlens_k_new,  # 9: cu_seqlens_k_new
+        None,  # 10: seqused_q
+        cache_seqlens,  # 11: seqused_k
+        max_seqlen_q,  # 12: max_seqlen_q
+        None,  # 13: max_seqlen_k
+        page_table,  # 14: page_table
+        cache_batch_idx,  # 15: kv_batch_idx
+        cache_leftpad,  # 16: leftpad_k
+        rotary_cos,  # 17: rotary_cos
+        rotary_sin,  # 18: rotary_sin
+        rotary_seqlens,  # 19: seqlens_rotary
+        q_descale,  # 20: q_descale
+        k_descale,  # 21: k_descale
+        v_descale,  # 22: v_descale
+        softmax_scale,  # 23: softmax_scale
+        causal,  # 24: is_causal
+        window_size[0],  # 25: window_size_left
+        window_size[1],  # 26: window_size_right
+        attention_chunk_val,  # 27: attention_chunk
+        softcap,  # 28: softcap
+        rotary_interleaved,  # 29: is_rotary_interleaved
+        scheduler_metadata,  # 30: scheduler_metadata
+        num_splits,  # 31: num_splits
+        pack_gqa,  # 32: pack_gqa
+        sm_margin,  # 33: sm_margin
+        sinks,  # 34: sinks
    )
    return (out, softmax_lse) if return_softmax_lse else out