[Attention] Add FlashInfer Sparse MLA backend

docs/design/attention_backends.md

@@ -128,6 +128,7 @@ Priority is **1 = highest** (tried first).
 | 4 | `FLASHMLA` |
 | 5 | `TRITON_MLA` |
 | 6 | `FLASHMLA_SPARSE` |
+| 7 | `FLASHINFER_MLA_SPARSE` |
 
 **Ampere/Hopper (SM 8.x-9.x):**
 
@@ -203,6 +204,7 @@ configuration.
 |---------|--------|-----------|-------------|------------|------|--------|-----------|-----------------|--------------|
 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | Decoder | 10.x |
 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | Decoder | 10.x |
+| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 10.x |
 | `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | Decoder | 9.x-10.x |
 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | Decoder | 9.x-10.x |
 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | Decoder | 9.x |

tests/v1/attention/test_sparse_mla_backends.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for the FlashMLA sparse backend utilities."""
+"""Unit tests for the sparse MLA backends and utilities."""
 
 import math
 from types import MethodType, SimpleNamespace
@@ -25,6 +25,9 @@
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.platforms import current_platform
 from vllm.utils.math_utils import cdiv
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
+    FlashInferMLASparseBackend,
+)
 from vllm.v1.attention.backends.mla.flashmla_sparse import (
     FlashMLASparseBackend,
     triton_convert_req_index_to_global_index,
@@ -156,32 +159,40 @@ def _quantize_dequantize_fp8_ds_mla(
     return dequant_kv_c, dequant_k_pe
 
 
+@pytest.mark.parametrize(
+    "backend_cls",
+    [FlashMLASparseBackend, FlashInferMLASparseBackend],
+    ids=["FlashMLA", "FlashInfer"],
+)
 @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys()))
-@pytest.mark.parametrize("kv_cache_dtype", ["fp8_ds_mla", "auto"])
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_ds_mla"])
 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
-@pytest.mark.skipif(
-    torch.cuda.get_device_capability() < (9, 0),
-    reason="FlashMLASparseBackend requires CUDA 9.0 or higher",
-)
 def test_sparse_backend_decode_correctness(
     default_vllm_config,
     dist_init,
+    backend_cls,
     batch_name,
     kv_cache_dtype,
     tensor_parallel_size,
     workspace_init,
 ):
-    if current_platform.is_rocm():
-        pytest.skip("ROCm does not support fp8_ds_mla data type for kv cache.")
+    if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:
+        pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")
 
-    if not torch.cuda.is_available():
-        pytest.skip("CUDA is required for sparse MLA decode test")
+    if backend_cls == FlashMLASparseBackend:
+        ok, reason = flashmla.is_flashmla_sparse_supported()
+        if not ok:
+            pytest.skip(reason)
+    elif backend_cls == FlashInferMLASparseBackend:
+        if not current_platform.has_device_capability(100):
+            pytest.skip("FlashInferMLASparseBackend requires SM 10.0 or higher")
+
+    batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
+    use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"
 
     device = torch.device("cuda")
     dtype = torch.bfloat16
 
-    batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]
-
     # Model hyper-parameters (kept intentionally small for the unit test)
     total_num_heads = 128
     # Compute per-rank heads for simulated TP
@@ -268,22 +279,22 @@ def test_sparse_backend_decode_correctness(
         kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
         k_pe_full = torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
 
-        # SM100 (Blackwell) uses float -> e8m0 -> bf16 scale conversion
-        # which truncates scales to powers of 2. Simulate this in reference.
-        is_sm100 = torch.cuda.get_device_capability()[0] >= 10
-        kv_c_full, k_pe_full = _quantize_dequantize_fp8_ds_mla(
-            kv_c_full,
-            k_pe_full.squeeze(1),
-            block_size=vllm_config.cache_config.block_size,
-            scale=kv_cache_scale,
-            simulate_sm100_e8m0_scales=is_sm100,
-        )
+        if use_fp8_ds_mla_quantization:
+            is_sm100 = torch.cuda.get_device_capability()[0] >= 10
+            kv_c_full, k_pe_squeezed = _quantize_dequantize_fp8_ds_mla(
+                kv_c_full,
+                k_pe_full.squeeze(1),
+                block_size=block_size,
+                scale=kv_cache_scale,
+                simulate_sm100_e8m0_scales=is_sm100,
+            )
+            k_pe_full = k_pe_squeezed.unsqueeze(1)
 
         q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)
         ql_nope = torch.einsum("qnh,lnh->qnl", q_nope, W_UK)
         q_mqa = torch.cat([ql_nope, q_pe], dim=-1)
 
-        k_mqa = torch.cat([kv_c_full, k_pe_full], dim=-1)
+        k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)
         k_mqa = k_mqa.unsqueeze(1).expand(-1, num_heads, -1)
         v_mqa = kv_c_full.unsqueeze(1).expand(-1, num_heads, -1)
 
@@ -334,11 +345,11 @@ def test_sparse_backend_decode_correctness(
         num_blocks=vllm_config.cache_config.num_gpu_blocks,
         common_attn_metadata=common_attn_metadata,
         randomize_blocks=False,
-        kv_cache_dtype=vllm_config.cache_config.cache_dtype,
+        kv_cache_dtype=kv_cache_dtype if use_fp8_ds_mla_quantization else "auto",
         scale=kv_cache_scale,
     )
 
-    builder_cls = FlashMLASparseBackend.get_builder_cls()
+    builder_cls = backend_cls.get_builder_cls()
     builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
     metadata = builder.build(
         common_prefix_len=0, common_attn_metadata=common_attn_metadata
@@ -362,15 +373,11 @@ def test_sparse_backend_decode_correctness(
         causal_mask, debug_indices, torch.full_like(debug_indices, -1)
     )
 
-    # FlashMLASparseImpl now reads top-k indices from the indexer-provided
+    # Sparse backends read top-k indices from the indexer-provided
     # buffer, so emulate that contract with a simple namespace mock.
     debug_indices = debug_indices.expand(metadata.num_actual_tokens, -1).clone()
     mock_indexer = SimpleNamespace(topk_indices_buffer=debug_indices)
 
-    ok, reason = flashmla.is_flashmla_sparse_supported()
-    if not ok:
-        pytest.skip(reason)
-
     kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1)
     kv_b_proj_weight = kv_b_proj_weight.view(
         kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)
@@ -383,7 +390,7 @@ def test_sparse_backend_decode_correctness(
     ).to(device=device, dtype=dtype)
     mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous())
 
-    impl_cls = FlashMLASparseBackend.get_impl_cls()
+    impl_cls = backend_cls.get_impl_cls()
     with set_current_vllm_config(vllm_config):
         impl = impl_cls(
             num_heads=num_heads,
@@ -430,7 +437,7 @@ def test_sparse_backend_decode_correctness(
 
     # FP8 quantization introduces some error, but should be within reasonable bounds
     # BF16 (auto) should be very accurate, FP8 allows slightly more tolerance
-    if kv_cache_dtype == "fp8_ds_mla":
+    if kv_cache_dtype.startswith("fp8"):
         torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.05, atol=0.05)
     else:
         torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01)

vllm/platforms/cuda.py

@@ -56,6 +56,7 @@ def _get_backend_priorities(
                 AttentionBackendEnum.FLASHMLA,
                 AttentionBackendEnum.TRITON_MLA,
                 AttentionBackendEnum.FLASHMLA_SPARSE,
+                AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
             ]
         else:
             return [
@@ -179,6 +180,7 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
             use_flashmla = False
             use_cutlass_mla = False
             use_flashinfer_mla = False
+            use_flashinfer_mla_sparse = False
 
             from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
 
@@ -187,6 +189,16 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
                 hf_text_config = model_config.hf_text_config
                 qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)
                 if (
+                    cls.is_device_capability_family(100)
+                    and use_sparse
+                    and qk_nope_head_dim == 128
+                ):
+                    # Blackwell + sparse => Use FlashInfer MLA Sparse
+                    use_flashinfer_mla_sparse = True
+                    vllm_config.attention_config.backend = (
+                        AttentionBackendEnum.FLASHINFER_MLA_SPARSE
+                    )
+                elif (
                     cls.is_device_capability_family(100)
                     and not use_sparse
                     and qk_nope_head_dim == 128
@@ -214,6 +226,9 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
                 use_flashmla = backend == AttentionBackendEnum.FLASHMLA
                 use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA
                 use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA
+                use_flashinfer_mla_sparse = (
+                    backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE
+                )
 
             if (
                 use_flashmla
@@ -239,8 +254,18 @@ def check_and_update_config(cls, vllm_config: "VllmConfig") -> None:
                     "Forcing kv cache block size to 64 for FlashInferMLA backend."
                 )
 
+            if use_flashinfer_mla_sparse and cache_config.block_size != 64:
+                cache_config.block_size = 64
+                logger.info(
+                    "Forcing kv cache block size to 64 for FlashInferMLASparse backend."
+                )
+
             # TODO(Chen): remove this hacky code
-            if use_sparse and cache_config.block_size != 64:
+            if (
+                use_sparse
+                and not use_flashinfer_mla_sparse
+                and cache_config.block_size != 64
+            ):
                 cache_config.block_size = 64
                 logger.info(
                     "Forcing kv cache block size to 64 for FlashMLASparse backend."