[Attention] Enable FA4 sparse MQA decode on Blackwell (#55866)

Signed-off-by: Simon Veitner <sveitner@redhat.com>
Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

cmake/external_projects/vllm_flash_attn.cmake

@@ -39,7 +39,7 @@ else()
   FetchContent_Declare(
           vllm-flash-attn
           GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
-          GIT_TAG 9cd61de38763d712bb6ce56e2a02cc2bf718c89f
+          GIT_TAG b273a02621d3c94bc747df67204f591b67f38fae
           GIT_PROGRESS TRUE
           # Don't share the vllm-flash-attn build between build types
           BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn

docs/design/attention_backends.md

@@ -147,7 +147,7 @@ Priority is **1 = highest** (tried first).
 
 --8<-- "gen:priority-mla"
 
-> **\*** 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.
+> **\*** For sparse MLA, FP8 KV cache always prefers `FLASHINFER_MLA_SPARSE`. With BF16 KV cache, `FLASH_ATTN_MLA_SPARSE_FA4` (then `FLASHINFER_MLA_SPARSE`) is preferred for low query-head counts (<= 16) on SM 10.x, while `FLASHMLA_SPARSE` is preferred otherwise.
 >
 > **Note:** ROCm and CPU platforms have their own selection logic. See the platform-specific documentation for details.
 

tests/kernels/attention/test_flashinfer_mla_decode.py

@@ -175,7 +175,15 @@ def test_flashinfer_trtllm_sparse_mla_decode_without_rope(dtype, mode, monkeypat
             backend, "prepare_sparse_mla_safe_lengths", unexpected_helper
         )
 
-    out, lse = impl._run_mqa_kernel(query.squeeze(1), kv_cache, indices, counts)
+    out, lse = impl._forward_mqa_kernel(
+        query.squeeze(1),
+        kv_cache,
+        indices,
+        counts,
+        layer=None,
+        block_size=block_size,
+        is_decode=True,
+    )
     reference = trtllm_batch_decode_with_kv_cache_mla(
         query=query,
         kv_cache=kv_cache.unsqueeze(1),

tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py

@@ -238,6 +238,7 @@ def _sparse_order(kv_cache_dtype, head_size, num_heads=32, capability=SM100):
     )
     sparse = {
         "FLASH_ATTN_MLA_SPARSE",
+        "FLASH_ATTN_MLA_SPARSE_FA4",
         "FLASHMLA_SPARSE",
         "FLASHINFER_MLA_SPARSE",
         "FLASHINFER_MLA_SPARSE_SM90",
@@ -276,3 +277,36 @@ def test_sm90_nope_fp8_ds_mla_resolves_to_flashmla():
     for name in order[: order.index("FLASHMLA_SPARSE")]:
         backend_cls = AttentionBackendEnum[name].get_class()
         assert not backend_cls.supports_kv_cache_dtype("fp8_ds_mla"), name
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA backends")
+def test_sm100_576_bf16_low_heads_falls_back_from_fa4(monkeypatch):
+    """With FA4 unusable, low-head bf16 selection lands on FlashInfer."""
+    import vllm.v1.attention.backends.mla.flashattn_mla_sparse as fa4_sparse
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.v1.attention.backends.registry import AttentionBackendEnum
+
+    monkeypatch.setattr(fa4_sparse, "_fa4_cute_mla_available", lambda: "no FA4")
+
+    def reasons(name):
+        backend_cls = AttentionBackendEnum[name].get_class()
+        return backend_cls.validate_configuration(
+            head_size=576,
+            dtype=torch.bfloat16,
+            kv_cache_dtype="auto",
+            block_size=64,
+            use_mla=True,
+            has_sink=False,
+            use_sparse=True,
+            use_mm_prefix=False,
+            use_per_head_quant_scales=False,
+            device_capability=SM100,
+            attn_type="decoder",
+        )
+
+    order = _sparse_order("auto", 576, num_heads=16)
+    with set_current_vllm_config(VllmConfig()):
+        # The selector takes the first candidate with no rejection reasons.
+        selected = next(name for name in order if not reasons(name))
+    assert order[0] == "FLASH_ATTN_MLA_SPARSE_FA4", order
+    assert selected == "FLASHINFER_MLA_SPARSE"

tests/v1/attention/test_sparse_mla_backends.py

@@ -34,6 +34,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.hisparse.worker import (
 )
 from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.attention import mla_attention
+from vllm.model_executor.layers.attention import (
+    sparse_mla_attention as sparse_mla_module,
+)
 from vllm.model_executor.layers.attention.mla_attention import _use_masked_mha
 from vllm.model_executor.layers.attention.sparse_mla_attention import (
     GLOBAL_TOPK_MASK_MAX_BYTES,
@@ -46,6 +49,7 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import (
 )
 from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.platforms import current_platform
+from vllm.platforms.interface import DeviceCapability
 
 # This module's end-to-end sparse-backend correctness test supports CUDA and
 # ROCm.  The remaining tests exercise CUDA-specific FlashMLA/FlashInfer APIs
@@ -71,8 +75,11 @@ if current_platform.is_cuda():
         flashattn_mla_sparse as flashattn_sparse_module,
     )
     from vllm.v1.attention.backends.mla.flashattn_mla_sparse import (
+        FlashAttnMLASparseFA4Backend,
+        FlashAttnMLASparseFA4Impl,
         FlashAttnMLASparseImpl,
         FlashAttnMLASparseMetadataBuilder,
+        _fa4_cute_mla_available,
     )
 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
     FlashInferMLASparseImpl,
@@ -227,6 +234,7 @@ def test_nope_flashinfer_sparse_mla_uses_model_scale(monkeypatch):
         req_id_per_token=torch.zeros(1, dtype=torch.int32),
         block_table=torch.zeros((1, 1), dtype=torch.int32),
         block_size=32,
+        num_decode_tokens=1,
     )
     recorded_scale = None
 
@@ -1401,6 +1409,7 @@ def test_masked_mha_workspace_fits_single_request_boundary(max_query_len, expect
         ("FLASHMLA_SPARSE", 8, 112 * 1024),
         ("FLASHINFER_MLA_SPARSE", 4, 36 * 1024),
         ("FLASHINFER_MLA_SPARSE", 8, 64 * 1024),
+        ("FLASH_ATTN_MLA_SPARSE_FA4", 4, 36 * 1024),
     ],
 )
 def test_masked_mha_workspace_guards_long_routing_policy(
@@ -3569,56 +3578,6 @@ def test_hisparse_keeps_speculative_rows_as_decodes():
     assert builder.reorder_batch_threshold == 4
 
 
-def test_flashinfer_hisparse_decode_runs_batched_attention():
-    device = torch.device("cpu")
-    num_tokens = 6
-    q = torch.randn(num_tokens, 2, 4, device=device)
-    topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=device)
-    physical_topk = topk + 10
-    valid_counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=device)
-    kernel_shapes: list[torch.Size] = []
-
-    def convert_decode(self, *args, **kwargs):  # noqa: ARG001
-        return physical_topk, valid_counts
-
-    def prepare_kernel(self, *args, **kwargs):  # noqa: ARG001
-        pass
-
-    def run_kernel(self, q, cache, indices, counts):  # noqa: ARG001
-        kernel_shapes.append(q.shape)
-        return q[..., :1], None
-
-    cache_handle = SimpleNamespace(
-        runtime=SimpleNamespace(
-            hot=SimpleNamespace(attention_cache=torch.empty(1, device=device))
-        ),
-    )
-    index_group = object.__new__(HiSparseMLAIndexGroup)
-    index_group.caches = [cache_handle]
-    index_group.convert_logical_to_physical_topk = MethodType(
-        convert_decode, index_group
-    )
-    impl = object.__new__(FlashInferMLASparseImpl)
-    impl.topk_indices_buffer = topk
-    impl.index_group = index_group
-    impl.index_group_index = 0
-    impl._prepare_mqa_kernel = MethodType(prepare_kernel, impl)
-    impl._run_mqa_kernel = MethodType(run_kernel, impl)
-    metadata = SimpleNamespace(num_decode_tokens=num_tokens)
-
-    output, lse = FlashInferMLASparseImpl.forward_mqa(
-        impl,
-        q,
-        torch.empty(1, device=device),
-        metadata,
-        SimpleNamespace(),
-    )
-
-    assert kernel_shapes == [q.shape]
-    assert output.shape == (num_tokens, 2, 1)
-    assert lse is None
-
-
 @pytest.mark.skipif(
     not current_platform.is_device_capability_family(90),
     reason="FlashAttention MLA requires Hopper",
@@ -3690,7 +3649,7 @@ def test_flashattn_shared_indices_match_independent_conversion(
 
     convert = triton_convert_req_index_to_global_index
     convert_spy = MagicMock(wraps=convert)
-    for module in (index_group_module, flashattn_sparse_module):
+    for module in (index_group_module, sparse_mla_module):
         monkeypatch.setattr(
             module, "triton_convert_req_index_to_global_index", convert_spy
         )
@@ -3735,48 +3694,19 @@ def test_flashattn_shared_indices_match_independent_conversion(
                 return_valid_counts=True,
             )
             for impl, cache, actual in zip(impls, caches, outputs):
-                expected = impl._run_mqa_kernel(
-                    q_nope, q_rope, cache, physical, counts, block_size
+                expected, _ = impl._forward_mqa_kernel(
+                    (q_nope, q_rope),
+                    cache,
+                    physical,
+                    counts,
+                    layer=None,
+                    block_size=block_size,
+                    is_decode=num_decode_tokens == num_tokens,
                 )
                 torch.testing.assert_close(actual, expected, rtol=0, atol=0)
     stream.synchronize()
 
 
-def test_flashattn_hisparse_decode_uses_index_group():
-    num_tokens = 4
-    q_nope = torch.empty(num_tokens, 2, 3, device=DEVICE_TYPE)
-    q_rope = torch.empty(num_tokens, 2, 1, device=DEVICE_TYPE)
-    topk = torch.zeros(num_tokens, 4, dtype=torch.int32, device=DEVICE_TYPE)
-    physical = topk.clone()
-    counts = torch.full((num_tokens,), 4, dtype=torch.int32, device=DEVICE_TYPE)
-    index_group = object.__new__(HiSparseMLAIndexGroup)
-    index_group.convert_logical_to_physical_topk = MagicMock(
-        return_value=(physical, counts)
-    )
-    index_group.physical_kv_cache = MagicMock(
-        return_value=torch.empty(1, device=DEVICE_TYPE)
-    )
-    impl = object.__new__(FlashAttnMLASparseImpl)
-    impl.topk_indices_buffer = topk
-    impl.index_group = index_group
-    impl.index_group_index = 0
-    impl._run_mqa_kernel = MagicMock(return_value=q_nope[..., :1])
-    metadata = SimpleNamespace(num_decode_tokens=num_tokens, block_size=64)
-
-    output, lse = FlashAttnMLASparseImpl.forward_mqa(
-        impl,
-        (q_nope, q_rope),
-        torch.empty(1, device=DEVICE_TYPE),
-        metadata,
-        SimpleNamespace(),
-    )
-
-    assert output.shape == (num_tokens, 2, 1)
-    assert lse is None
-    index_group.convert_logical_to_physical_topk.assert_called_once()
-    impl._run_mqa_kernel.assert_called_once()
-
-
 def test_flashinfer_sm120_hisparse_decode_uses_index_group():
     num_tokens = 4
     q = torch.empty(num_tokens, 2, 4, device=DEVICE_TYPE)
@@ -3811,6 +3741,131 @@ def test_flashinfer_sm120_hisparse_decode_uses_index_group():
     impl._run_mqa_kernel.assert_called_once()
 
 
+def test_sparse_mqa_plain_lane_wiring(monkeypatch):
+    """Under DCP the plain lane converts with the DCP filter at the flat stride."""
+    import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla
+
+    num_tokens, topk_tokens, block_size, head_dim = 3, 8, 4, 6
+    # Flat row stride (8) differs from block size (4): the wrong one misaddresses.
+    kv_cache = torch.zeros(2, 8, head_dim)[:, :block_size]
+    q = torch.zeros(num_tokens, 2, head_dim)
+    metadata = SimpleNamespace(
+        num_decode_tokens=num_tokens,
+        block_size=block_size,
+        block_table=torch.zeros((num_tokens, 1), dtype=torch.int32),
+        req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32),
+        cp_kv_cache_interleave_size=2,
+    )
+    index_filter = MagicMock(return_value=(None, None))
+    monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", index_filter)
+
+    impl = object.__new__(FlashInferMLASparseImpl)
+    impl.topk_indices_buffer = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32)
+    impl.index_group = MagicMock()
+    impl.index_group_index = 0
+    impl.dcp_world_size = 2
+    impl.dcp_rank = 1
+    impl._convert_logical_to_physical_topk = MagicMock()
+    impl._forward_mqa_kernel = MagicMock(return_value=(q, None))
+
+    impl.forward_mqa(q, kv_cache, metadata, None)
+
+    impl._convert_logical_to_physical_topk.assert_not_called()
+    assert index_filter.call_args.kwargs["BLOCK_STRIDE_ROWS"] == 8
+    impl._forward_mqa_kernel.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    "num_decode_tokens,all_resident,split_q",
+    [
+        pytest.param(0, True, False, id="resident"),
+        pytest.param(0, False, False, id="staged"),
+        pytest.param(2, True, True, id="mixed"),
+        pytest.param(5, True, False, id="decode"),
+    ],
+)
+def test_sparse_mqa_hisparse_lanes(
+    monkeypatch, num_decode_tokens, all_resident, split_q
+):
+    """HiSparse splits a batch into at most one decode and one prefill call."""
+    import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla
+
+    num_tokens, topk_tokens, block_size, head_dim = 5, 8, 4, 6
+    num_prefill = num_tokens - num_decode_tokens
+    # Distinct block counts, so the KV tensor a call is handed names its lane.
+    host = torch.zeros(2, block_size, head_dim, dtype=torch.bfloat16)
+    hot = torch.zeros(6, block_size, head_dim, dtype=torch.bfloat16)
+    staged = torch.zeros(3, block_size, head_dim, dtype=torch.bfloat16)
+    q = torch.arange(num_tokens, dtype=torch.float32).view(num_tokens, 1, 1)
+    metadata = SimpleNamespace(
+        num_decode_tokens=num_decode_tokens, block_size=block_size
+    )
+
+    def converted(rows, fill):
+        return (
+            torch.full((rows, topk_tokens), fill, dtype=torch.int32),
+            torch.full((rows,), rows, dtype=torch.int32),
+        )
+
+    index_group = object.__new__(HiSparseMLAIndexGroup)
+    index_group.physical_kv_cache = MagicMock(return_value=hot)
+    index_group.cache = MagicMock(
+        return_value=SimpleNamespace(all_context_pages_resident=all_resident)
+    )
+    # The fill value names the converter that produced a call's indices.
+    index_group.convert_logical_to_physical_topk = MagicMock(
+        return_value=converted(num_decode_tokens, 1)
+        if num_decode_tokens
+        else converted(num_tokens, 2)
+    )
+    index_group.stage_prefill_rows = MagicMock(return_value=(staged, None, None))
+    calls: list[tuple] = []
+
+    def convert_staged(*args, **kwargs):
+        # Decode rows are consumed before a later conversion can reuse their buffer.
+        assert len(calls) == min(num_decode_tokens, 1)
+        return converted(num_prefill, 3)
+
+    monkeypatch.setattr(
+        sparse_mla,
+        "triton_convert_req_index_to_global_index",
+        MagicMock(side_effect=convert_staged),
+    )
+
+    def hook(rows, cache, indices, counts, *, is_decode, **kwargs):
+        if split_q:
+            assert rows[0].tolist() == rows[1].tolist()
+            rows = rows[0]
+        calls.append((rows.flatten().tolist(), cache.shape[0], indices[0, 0].item()))
+        return torch.full((rows.shape[0], 1, 1), 1.0 if is_decode else 2.0), None
+
+    impl = object.__new__(FlashInferMLASparseImpl)
+    impl.topk_indices_buffer = torch.zeros(num_tokens, topk_tokens, dtype=torch.int32)
+    impl.index_group = index_group
+    impl.index_group_index = 0
+    impl.dcp_world_size = 1
+    impl._forward_mqa_kernel = hook
+
+    out, lse = impl.forward_mqa((q, q) if split_q else q, host, metadata, None)
+
+    q_rows = q.flatten().tolist()
+    expected = []
+    if num_decode_tokens:
+        # Decode reads the hot cache through the group's conversion.
+        expected.append((q_rows[:num_decode_tokens], hot.shape[0], 1))
+    if num_decode_tokens == 0 and all_resident:
+        expected.append((q_rows, hot.shape[0], 2))
+    elif num_prefill:
+        expected.append((q_rows[num_decode_tokens:], staged.shape[0], 3))
+    assert calls == expected
+    if not num_prefill:
+        # A decode-only batch returns before it looks up the prefill cache.
+        index_group.cache.assert_not_called()
+    assert lse is None
+    # A lane handed the whole batch would concatenate to more rows than this.
+    assert out.flatten().tolist() == [1.0] * num_decode_tokens + [2.0] * num_prefill
+
+
 def test_hisparse_resident_prefill_uses_attention_block_stride():
     """A batch beyond the swap-row capacity must bypass swap_in.
 
@@ -3887,6 +3942,250 @@ def test_hisparse_mixed_single_token_batch_is_not_decode_only():
     assert not cache.decode_batch
 
 
+def _fa4_impl(topk_indices_buffer, **fields):
+    impl = object.__new__(FlashAttnMLASparseFA4Impl)
+    impl.__dict__.update(
+        kv_lora_rank=512,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=64,
+        scale=1.0,
+        kv_cache_dtype="auto",
+        dcp_world_size=1,
+        dcp_rank=0,
+        need_to_return_lse_for_decode=False,
+        is_nope_mla=False,
+        _workspace_buffer=None,
+        bmm1_scale=None,
+        bmm2_scale=None,
+        # A non-index-producing layer does not wait on the indexer's ready event.
+        index_group=SparseMLAIndexGroupBuilder(topk_indices_buffer).register_layer(
+            False
+        )[0],
+        index_group_index=0,
+    )
+    impl.__dict__.update(fields)
+    # topk_indices_buffer is a data descriptor, so it needs its own setter.
+    impl.topk_indices_buffer = topk_indices_buffer
+    return impl
+
+
+def _fa4_inputs(counts, *, num_heads=16, topk=128, num_blocks=4, device=DEVICE_TYPE):
+    """Decode inputs whose token ``i`` has ``counts[i]`` valid rows, -1 padded."""
+    device = torch.device(device)
+    block_size, kv_lora_rank, rope_dim = 64, 512, 64
+    num_tokens = len(counts)
+    torch.manual_seed(0)
+
+    def values(*shape):
+        return torch.rand(shape, dtype=torch.bfloat16, device=device) - 0.5
+
+    topk_indices = torch.full((num_tokens, topk), -1, dtype=torch.int32, device=device)
+    for tok, count in enumerate(counts):
+        topk_indices[tok, :count] = torch.randperm(
+            num_blocks * block_size, device=device
+        )[:count].to(torch.int32)
+
+    return SimpleNamespace(
+        kv_lora_rank=kv_lora_rank,
+        rope_dim=rope_dim,
+        block_size=block_size,
+        topk=topk,
+        kv_cache=values(num_blocks, block_size, kv_lora_rank + rope_dim),
+        topk_indices=topk_indices,
+        valid_counts=torch.tensor(counts, dtype=torch.int32, device=device),
+        ql_nope=values(num_tokens, num_heads, kv_lora_rank),
+        q_pe=values(num_tokens, num_heads, rope_dim),
+        scale=(kv_lora_rank + rope_dim) ** -0.5,
+        metadata=SimpleNamespace(
+            req_id_per_token=torch.zeros(num_tokens, dtype=torch.int32, device=device),
+            block_table=torch.arange(num_blocks, dtype=torch.int32, device=device).view(
+                1, num_blocks
+            ),
+            block_size=block_size,
+            num_decode_tokens=num_tokens,
+            cp_kv_cache_interleave_size=1,
+        ),
+    )
+
+
+def _record_fa4_kernels(monkeypatch, num_heads, *, device):
+    """Replace both lanes with recorders whose outputs identify the lane."""
+    import flashinfer.decode
+
+    fa4_calls: list[dict] = []
+    trtllm_calls: list[dict] = []
+
+    def lse(rows):
+        return torch.arange(rows * num_heads, dtype=torch.float32, device=device).view(
+            rows, num_heads
+        )
+
+    def fake_fa4(**kwargs):
+        fa4_calls.append(kwargs)
+        rows = kwargs["q"].shape[0]
+        out = torch.ones(rows, num_heads, 512, dtype=torch.bfloat16, device=device)
+        return (out, lse(rows)) if kwargs.get("return_softmax_lse") else out
+
+    def fake_trtllm(**kwargs):
+        trtllm_calls.append(kwargs)
+        rows = kwargs["query"].shape[0]
+        out = torch.full(
+            (rows, 1, num_heads, 512), 2.0, dtype=torch.bfloat16, device=device
+        )
+        return (out, lse(rows)) if kwargs.get("return_lse") else out
+
+    monkeypatch.setattr(flashattn_sparse_module, "flash_attn_varlen_func", fake_fa4)
+    monkeypatch.setattr(
+        flashinfer.decode, "trtllm_batch_decode_with_kv_cache_mla", fake_trtllm
+    )
+    return SimpleNamespace(fa4=fa4_calls, trtllm=trtllm_calls, lse=lse)
+
+
+@pytest.mark.parametrize(
+    "num_heads,return_lse", [(128, False), (16, True)], ids=["heads128", "heads16_lse"]
+)
+def test_fa4_sparse_decode_kernel_correctness(num_heads, return_lse):
+    if not current_platform.is_device_capability_family(100):
+        pytest.skip("FA4 sparse MLA requires SM 10.x")
+    if (reason := _fa4_cute_mla_available()) is not None:
+        pytest.skip(reason)
+    counts = [0, 1, 127, 128, 129, 2047, 2048]
+    inputs = _fa4_inputs(counts, num_heads=num_heads, topk=2048, num_blocks=64)
+    impl = _fa4_impl(
+        topk_indices_buffer=inputs.topk_indices,
+        scale=inputs.scale,
+        need_to_return_lse_for_decode=return_lse,
+    )
+    q = (inputs.ql_nope, inputs.q_pe)
+
+    with torch.inference_mode():
+        out, lse = impl.forward_mqa(q, inputs.kv_cache, inputs.metadata, None)
+        if not return_lse:
+            assert lse is None
+        else:
+            out_fused, lse_fused = impl.forward_mqa(
+                torch.cat(q, dim=-1), inputs.kv_cache, inputs.metadata, None
+            )
+            assert lse.shape == (len(counts), num_heads)
+            assert torch.equal(out, out_fused) and torch.equal(lse, lse_fused)
+
+    kv_flat = inputs.kv_cache.view(-1, inputs.kv_lora_rank + inputs.rope_dim)
+    keys = kv_flat[inputs.topk_indices.clamp(min=0).long()].float()
+    scores = torch.einsum("thd,tkd->thk", torch.cat(q, dim=-1).float(), keys)
+    scores = (scores * inputs.scale).masked_fill(
+        (inputs.topk_indices < 0)[:, None, :], float("-inf")
+    )
+    # The all-sentinel row softmaxes to NaN here; the kernel writes (0, -inf).
+    probs = torch.softmax(scores, dim=-1).nan_to_num()
+    torch.testing.assert_close(
+        out.float(),
+        torch.einsum("thk,tkd->thd", probs, keys[..., : inputs.kv_lora_rank]),
+        rtol=0.01,
+        atol=0.01,
+    )
+    if return_lse:
+        torch.testing.assert_close(
+            lse, torch.logsumexp(scores, dim=-1), rtol=0, atol=5e-3
+        )
+
+
+@pytest.mark.parametrize("is_decode", [True, False], ids=["decode", "prefill"])
+def test_fa4_sparse_hook_routes_lanes(monkeypatch, is_decode):
+    """The hook runs a decode batch on FA4 and every other batch on trtllm-gen."""
+    import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla
+
+    counts = [1, 2, 3, 4, 5, 6, 7]
+    num_tokens, dcp = len(counts), 2
+    # Under DCP the hook gets the heads of every DCP rank, not just its own.
+    kernel_heads = 16 * dcp
+    inputs = _fa4_inputs(counts, num_heads=kernel_heads, device="cpu")
+    calls = _record_fa4_kernels(monkeypatch, kernel_heads, device="cpu")
+    # DCP without PCP (the gate rejects the pair) always fuses the query.
+    q = torch.cat((inputs.ql_nope, inputs.q_pe), dim=-1)
+    impl = _fa4_impl(
+        # One row longer than the batch: the varlen-scalar slices become observable.
+        topk_indices_buffer=inputs.topk_indices.new_zeros(num_tokens + 1, inputs.topk),
+        dcp_world_size=dcp,
+        need_to_return_lse_for_decode=True,
+    )
+    impl._convert_logical_to_physical_topk = MagicMock()
+    dcp_filter = MagicMock(return_value=(inputs.topk_indices, inputs.valid_counts))
+    monkeypatch.setattr(sparse_mla, "triton_filter_and_convert_dcp_index", dcp_filter)
+    inputs.metadata.num_decode_tokens = num_tokens if is_decode else 1
+
+    with torch.inference_mode():
+        out, lse = impl.forward_mqa(q, inputs.kv_cache, inputs.metadata, None)
+
+    dcp_filter.assert_called_once()
+    impl._convert_logical_to_physical_topk.assert_not_called()
+    assert out.shape == (num_tokens, kernel_heads, 512)
+    if is_decode:
+        (fa4,) = calls.fa4
+        assert calls.trtllm == []
+        assert fa4["return_softmax_lse"] is True
+        assert fa4["cu_seqlens_q"].tolist() == list(range(num_tokens + 1))
+        assert (out == 1.0).all()
+        # FA4's LSE is already natural-log.
+        torch.testing.assert_close(lse, calls.lse(num_tokens), rtol=0, atol=0)
+    else:
+        (trtllm,) = calls.trtllm
+        assert calls.fa4 == []
+        assert trtllm["return_lse"] is True
+        assert (out == 2.0).all()
+        # trtllm-gen's LSE is base 2; the DCP merge expects natural log.
+        expected = calls.lse(num_tokens) * math.log(2.0)
+        torch.testing.assert_close(lse, expected, rtol=0, atol=1e-5)
+
+
+@pytest.mark.parametrize(
+    "config,expected",
+    [
+        ({}, None),
+        ({"heads": 128}, None),
+        ({"heads": 12}, "heads"),
+        ({"dcp": 2}, None),
+        ({"dcp": 8}, "heads"),
+        ({"index_topk": 192}, "index_topk"),
+        ({"rope": 32}, "(kv_lora_rank, qk_rope_head_dim)"),
+        ({"flashinfer": False}, "FlashInfer"),
+    ],
+)
+def test_fa4_sparse_gate(monkeypatch, config, expected):
+    """``validate_configuration`` on a fixed SM100: one row per gate reason."""
+    has_flashinfer = config.get("flashinfer", True)
+    monkeypatch.setattr("vllm.utils.flashinfer.has_flashinfer", lambda: has_flashinfer)
+    monkeypatch.setattr(
+        flashattn_sparse_module, "_fa4_cute_mla_available", lambda: None
+    )
+    dcp = config.get("dcp", 1)
+    vllm_config = _build_sparse_dcp_vllm_config(config.get("heads", 16), dcp)
+    hf_text_config = vllm_config.model_config.hf_text_config
+    hf_text_config.index_topk = config.get("index_topk", 128)
+    hf_text_config.qk_rope_head_dim = config.get("rope", 64)
+
+    with set_current_vllm_config(vllm_config):
+        reasons = FlashAttnMLASparseFA4Backend.validate_configuration(
+            head_size=576,
+            dtype=torch.bfloat16,
+            kv_cache_dtype="auto",
+            block_size=64,
+            use_mla=True,
+            has_sink=False,
+            use_sparse=True,
+            use_mm_prefix=False,
+            use_per_head_quant_scales=False,
+            device_capability=DeviceCapability(10, 0),
+            attn_type="decoder",
+            use_dcp=dcp > 1,
+        )
+
+    if expected is None:
+        assert reasons == []
+    else:
+        (reason,) = reasons
+        assert expected in reason
+
+
 def test_flashmla_cache_dtype_aliases_use_ds_layout():
     # kv-cache dtype aliases are canonicalized to fp8_ds_mla before the layer
     # stores kv_cache_dtype, so they cannot bypass the gate.

tests/v1/attention/test_sparse_mla_kv_cache_layout.py

@@ -13,15 +13,12 @@ from vllm.model_executor.layers.attention.mla_attention import MLAAttention
 from vllm.model_executor.models.deepseek_v2 import DeepseekV32IndexerCache
 from vllm.platforms import current_platform
 from vllm.v1.attention.backend import MultipleOf
-from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
-    FlashInferMLASparseTRTLLMBackend,
-)
-from vllm.v1.attention.backends.mla.flashmla_sparse import FlashMLASparseBackend
 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerBackend
 from vllm.v1.attention.backends.mla.sparse_utils import (
     _CONVERT_REQ_INDEX_TO_GLOBAL_INDEX_KERNEL,
     flat_kv_row_view,
 )
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts
 from vllm.v1.core.kv_cache_planning import get_kv_cache_config_from_groups
 from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs
@@ -49,17 +46,27 @@ pytestmark = pytest.mark.skip_global_cleanup
 @pytest.mark.parametrize(
     "backend,cache_dtype,sm100,stride_alignment",
     [
-        (FlashInferMLASparseTRTLLMBackend, "auto", True, 32 * 1152),
-        (FlashInferMLASparseTRTLLMBackend, "fp8", True, 32 * 576),
-        (FlashMLASparseBackend, "auto", True, 1152),
-        (FlashMLASparseBackend, "fp8_ds_mla", False, MultipleOf(64)),
-        (FlashMLASparseBackend, "fp8_ds_mla", True, 656),
-        (FlashMLASparseBackend, "nvfp4_ds_mla", True, 352),
+        (AttentionBackendEnum.FLASHINFER_MLA_SPARSE, "auto", True, 32 * 1152),
+        (AttentionBackendEnum.FLASHINFER_MLA_SPARSE, "fp8", True, 32 * 576),
+        pytest.param(
+            AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE_FA4,
+            "auto",
+            True,
+            64 * 1152,
+            marks=pytest.mark.skipif(
+                not current_platform.is_cuda(), reason="CUDA only"
+            ),
+        ),
+        (AttentionBackendEnum.FLASHMLA_SPARSE, "auto", True, 1152),
+        (AttentionBackendEnum.FLASHMLA_SPARSE, "fp8_ds_mla", False, MultipleOf(64)),
+        (AttentionBackendEnum.FLASHMLA_SPARSE, "fp8_ds_mla", True, 656),
+        (AttentionBackendEnum.FLASHMLA_SPARSE, "nvfp4_ds_mla", True, 352),
     ],
 )
 def test_allocation_and_warmup_follow_addressing_mode(
     monkeypatch, layout, backend, cache_dtype, sm100, stride_alignment, device
 ):
+    backend = backend.get_class()
     monkeypatch.setattr(
         current_platform, "is_device_capability_family", lambda family: sm100
     )

tests/v1/e2e/general/test_hisparse.py

@@ -83,13 +83,20 @@ def _get_hisparse_worker(runner: VllmRunner) -> HiSparseConnectorWorker:
     not current_platform.is_cuda(), reason="HiSparse requires NVIDIA CUDA"
 )
 @pytest.mark.parametrize(
-    "with_offloading", [False, True], ids=["standalone", "offload"]
+    "with_offloading,attention_backend",
+    [
+        pytest.param(False, None, id="standalone"),
+        pytest.param(True, None, id="offload"),
+        # Offloading does not touch the attention kernels: FA4 runs standalone.
+        pytest.param(False, "FLASH_ATTN_MLA_SPARSE_FA4", id="standalone-fa4"),
+    ],
 )
 @fork_new_process_for_each_test
 def test_hisparse_spill_and_prefix_restore(
     monkeypatch: pytest.MonkeyPatch,
     vllm_runner: type[VllmRunner],
     with_offloading: bool,
+    attention_backend: str | None,
 ):
     """Spilled prefixes restore and FULL-graph decode writes reach host KV.
 
@@ -101,6 +108,8 @@ def test_hisparse_spill_and_prefix_restore(
     capability = current_platform.get_device_capability()
     if capability is None or capability.major < 9:
         pytest.skip("Sparse MLA requires Hopper or newer")
+    if attention_backend and not current_platform.is_device_capability_family(100):
+        pytest.skip("FA4 sparse MLA requires SM 10.x")
 
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
     monkeypatch.setenv("VLLM_DEEP_GEMM_WARMUP", "skip")
@@ -142,6 +151,7 @@ def test_hisparse_spill_and_prefix_restore(
             )
         ),
         kv_transfer_config=kv_transfer_config,
+        attention_backend=attention_backend,
         block_size=64,
         max_model_len=320,
         max_num_batched_tokens=1024,

vllm/model_executor/layers/attention/mla_attention.py

@@ -661,6 +661,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 "FLASHMLA_SPARSE",
                 "FLASHINFER_MLA_SPARSE",
                 "FLASHINFER_MLA_SPARSE_SM120",
+                "FLASH_ATTN_MLA_SPARSE_FA4",
                 "DEEPSEEK_V32_INDEXER",
             ):
                 from vllm.v1.attention.backends.mla.compressor_utils import (
@@ -751,7 +752,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
 
     def _uses_flat_kv_cache(self) -> bool:
         backend = self.attn_backend.get_name()
-        return backend == "FLASHINFER_MLA_SPARSE" or (
+        return backend in ("FLASHINFER_MLA_SPARSE", "FLASH_ATTN_MLA_SPARSE_FA4") or (
             backend == "FLASHMLA_SPARSE"
             and self.kv_cache_dtype not in ("fp8_ds_mla", "nvfp4_ds_mla")
         )
@@ -1857,6 +1858,9 @@ def _use_masked_mha(
     seq_len: int,
     has_context: bool,
 ) -> bool:
+    if backend_name == "FLASH_ATTN_MLA_SPARSE_FA4":
+        # Its prefill batches run on FlashInfer's kernel, so it shares those rows.
+        backend_name = "FLASHINFER_MLA_SPARSE"
     thresholds = (
         _MASKED_MHA_THRESHOLDS.get((qk_head_dim, v_head_dim), {})
         .get(backend_name, {})

vllm/model_executor/layers/attention/sparse_mla_attention.py

@@ -36,12 +36,18 @@ from vllm.utils.torch_utils import (
     is_quantized_kv_cache,
     np_to_pinned_tensor,
 )
-from vllm.v1.attention.backend import AttentionMetadataBuilder
+from vllm.v1.attention.backend import AttentionLayer, AttentionMetadataBuilder
 from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
 from vllm.v1.attention.backends.mla.index_group import (
+    HiSparseMLAIndexGroup,
     SparseMLAIndexGroup,
     SparseMLAIndexGroupBuilder,
 )
+from vllm.v1.attention.backends.mla.sparse_utils import (
+    flat_kv_row_view,
+    triton_convert_req_index_to_global_index,
+    triton_filter_and_convert_dcp_index,
+)
 from vllm.v1.attention.backends.utils import split_decodes_and_prefills
 from vllm.v1.attention.ops.dcp import MLADCPManager
 from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
@@ -630,6 +636,9 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
     """Sparse MLA base with dense and masked-MHA prefill paths."""
 
     is_sparse = True
+    # Reuse the index group's top-k conversion only for uniform decode batches
+    # of multi-layer groups; other batches convert per layer.
+    group_conversion_decode_only: ClassVar[bool] = False
 
     def __init__(
         self,
@@ -728,6 +737,154 @@ class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic
             return_valid_counts=return_valid_counts,
         )
 
+    def forward_mqa(
+        self,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        kv_c_and_k_pe_cache: torch.Tensor,
+        attn_metadata: T,
+        layer: AttentionLayer,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        assert self.topk_indices_buffer is not None
+        num_tokens = q[0].shape[0] if isinstance(q, tuple) else q.shape[0]
+        topk = self.topk_indices_buffer[:num_tokens]
+        num_decode_tokens = attn_metadata.num_decode_tokens
+        block_size = attn_metadata.block_size
+        index_group = self.index_group
+        if not isinstance(index_group, HiSparseMLAIndexGroup):
+            _, block_stride_rows = flat_kv_row_view(kv_c_and_k_pe_cache, block_size)
+            if self.dcp_world_size > 1:
+                physical_topk, valid_counts = triton_filter_and_convert_dcp_index(
+                    attn_metadata.req_id_per_token[:num_tokens],
+                    attn_metadata.block_table,
+                    topk,
+                    dcp_size=self.dcp_world_size,
+                    dcp_rank=self.dcp_rank,
+                    cp_kv_cache_interleave_size=attn_metadata.cp_kv_cache_interleave_size,
+                    BLOCK_SIZE=block_size,
+                    BLOCK_STRIDE_ROWS=block_stride_rows,
+                    NUM_TOPK_TOKENS=topk.shape[1],
+                    return_valid_counts=True,
+                )
+            elif self.group_conversion_decode_only and not (
+                index_group is not None
+                and index_group.num_layers > 1
+                and num_decode_tokens >= num_tokens
+            ):
+                physical_topk, valid_counts = triton_convert_req_index_to_global_index(
+                    attn_metadata.req_id_per_token[:num_tokens],
+                    attn_metadata.block_table,
+                    topk,
+                    BLOCK_SIZE=block_size,
+                    BLOCK_STRIDE_ROWS=block_stride_rows,
+                    NUM_TOPK_TOKENS=topk.shape[1],
+                    return_valid_counts=True,
+                )
+            else:
+                physical_topk, valid_counts = self._convert_logical_to_physical_topk(
+                    topk,
+                    attn_metadata,
+                    block_stride_rows=block_stride_rows,
+                    return_valid_counts=True,
+                )
+            return self._forward_mqa_kernel(
+                q,
+                kv_c_and_k_pe_cache,
+                physical_topk,
+                valid_counts,
+                layer=layer,
+                block_size=block_size,
+                is_decode=num_decode_tokens >= num_tokens,
+            )
+
+        assert self.dcp_world_size == 1
+        layer_index = self.index_group_index
+        hot_cache = index_group.physical_kv_cache(layer_index).view(
+            kv_c_and_k_pe_cache.dtype
+        )
+        assert hot_cache.is_contiguous(), "HiSparse hot buffer must be contiguous"
+        decode_out = decode_lse = None
+        if num_decode_tokens:
+            physical_topk, valid_counts = index_group.convert_logical_to_physical_topk(
+                layer_index,
+                topk[:num_decode_tokens],
+                attn_metadata,
+                block_stride_rows=None,
+                return_valid_counts=True,
+            )
+            q_decode = (
+                (q[0][:num_decode_tokens], q[1][:num_decode_tokens])
+                if isinstance(q, tuple)
+                else q[:num_decode_tokens]
+            )
+            decode_out, decode_lse = self._forward_mqa_kernel(
+                q_decode,
+                hot_cache,
+                physical_topk,
+                valid_counts,
+                layer=layer,
+                block_size=block_size,
+                is_decode=True,
+            )
+            if num_decode_tokens >= num_tokens:
+                return decode_out, decode_lse
+
+        cache = index_group.cache(layer_index)
+        if num_decode_tokens == 0 and cache.all_context_pages_resident:
+            physical_topk, valid_counts = index_group.convert_logical_to_physical_topk(
+                layer_index,
+                topk,
+                attn_metadata,
+                block_stride_rows=None,
+                return_valid_counts=True,
+            )
+            prefill_cache = hot_cache
+        else:
+            prefill_cache, block_table, req_ids = index_group.stage_prefill_rows(
+                layer_index, kv_c_and_k_pe_cache, attn_metadata
+            )
+            physical_topk, valid_counts = triton_convert_req_index_to_global_index(
+                req_ids,
+                block_table,
+                topk[num_decode_tokens:],
+                BLOCK_SIZE=block_size,
+                NUM_TOPK_TOKENS=topk.shape[1],
+                return_valid_counts=True,
+            )
+        q_prefill = (
+            (q[0][num_decode_tokens:], q[1][num_decode_tokens:])
+            if isinstance(q, tuple)
+            else q[num_decode_tokens:]
+        )
+        prefill_out, prefill_lse = self._forward_mqa_kernel(
+            q_prefill,
+            prefill_cache,
+            physical_topk,
+            valid_counts,
+            layer=layer,
+            block_size=block_size,
+            is_decode=False,
+        )
+        if decode_out is None:
+            return prefill_out, prefill_lse
+        out = torch.cat((decode_out, prefill_out))
+        if decode_lse is None:
+            return out, None
+        assert prefill_lse is not None
+        return out, torch.cat((decode_lse, prefill_lse))
+
+    def _forward_mqa_kernel(
+        self,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        kv_cache: torch.Tensor,
+        topk_indices: torch.Tensor,
+        valid_counts: torch.Tensor,
+        *,
+        layer: AttentionLayer,
+        block_size: int,
+        is_decode: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        raise NotImplementedError
+
     @staticmethod
     def masked_mha_workspace_fits(prefill: MLACommonPrefillMetadata) -> bool:
         """Whether this prefill batch's top-k masks fit the workspace."""

vllm/platforms/cuda.py

@@ -105,9 +105,10 @@ def _get_backend_priorities(
                 ]
             else:
                 # BF16 KV Cache
-                # Prefer FlashInfer at low head counts (FlashMLA uses padding)
+                # Prefer FA4, then FlashInfer at low head counts (FlashMLA uses padding)
                 if num_heads is not None and num_heads <= 16:
                     sparse_backends = [
+                        AttentionBackendEnum.FLASH_ATTN_MLA_SPARSE_FA4,
                         AttentionBackendEnum.FLASHINFER_MLA_SPARSE,
                         AttentionBackendEnum.FLASHMLA_SPARSE,
                     ]

vllm/v1/attention/backends/mla/flashattn_mla_sparse.py

@@ -1,7 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import math
 from dataclasses import dataclass
+from functools import cache
 from typing import Any, ClassVar
 
 import torch
@@ -22,11 +24,14 @@ from vllm.v1.attention.backend import (
     MultipleOf,
 )
 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
-from vllm.v1.attention.backends.mla.index_group import HiSparseMLAIndexGroup
+from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
+    FlashInferMLASparseImpl,
+    FlashInferMLASparseTRTLLMBackend,
+    FlashInferMLASparseTRTLLMMetadataBuilder,
+)
 from vllm.v1.attention.backends.mla.sparse_utils import (
     align_blocks_to_rows,
     flat_kv_row_view,
-    triton_convert_req_index_to_global_index,
 )
 from vllm.v1.attention.ops.metadata import compute_token_to_req_indices
 from vllm.v1.kv_cache_interface import AttentionSpec
@@ -153,6 +158,8 @@ class FlashAttnMLASparseMetadataBuilder(
 
 
 class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
+    group_conversion_decode_only: ClassVar[bool] = True
+
     def __init__(
         self,
         num_heads: int,
@@ -205,144 +212,23 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
         )
         self.supports_quant_query_input = False
 
-    def forward_mqa(
+    def _forward_mqa_kernel(
         self,
         q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
-        kv_c_and_k_pe_cache: torch.Tensor,
-        attn_metadata: FlashAttnMLASparseMetadata,
+        kv_cache: torch.Tensor,
+        topk_indices: torch.Tensor,
+        valid_counts: torch.Tensor,
+        *,
         layer: AttentionLayer,
+        block_size: int,
+        is_decode: bool,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         if not isinstance(q, tuple):
             raise NotImplementedError(
                 "FlashAttnMLASparseImpl expects split (q_nope, q_rope) input."
             )
         q_nope, q_rope = q
-        num_actual_toks = q_rope.shape[0]
-
-        assert self.topk_indices_buffer is not None
-        topk_indices = self.topk_indices_buffer[:num_actual_toks]
-        index_group = self.index_group
-        if isinstance(index_group, HiSparseMLAIndexGroup):
-            num_decode_tokens = attn_metadata.num_decode_tokens
-            outputs = []
-            if num_decode_tokens:
-                physical_topk, valid_counts = (
-                    index_group.convert_logical_to_physical_topk(
-                        self.index_group_index,
-                        topk_indices[:num_decode_tokens],
-                        attn_metadata,
-                        block_stride_rows=None,
-                        return_valid_counts=True,
-                    )
-                )
-                outputs.append(
-                    self._run_mqa_kernel(
-                        q_nope[:num_decode_tokens],
-                        q_rope[:num_decode_tokens],
-                        index_group.physical_kv_cache(self.index_group_index).view(
-                            kv_c_and_k_pe_cache.dtype
-                        ),
-                        physical_topk,
-                        valid_counts,
-                        attn_metadata.block_size,
-                    )
-                )
-            if num_decode_tokens < num_actual_toks:
-                cache = index_group.cache(self.index_group_index)
-                if num_decode_tokens == 0 and cache.all_context_pages_resident:
-                    physical_topk, valid_counts = (
-                        index_group.convert_logical_to_physical_topk(
-                            self.index_group_index,
-                            topk_indices,
-                            attn_metadata,
-                            block_stride_rows=None,
-                            return_valid_counts=True,
-                        )
-                    )
-                    prefill_cache = index_group.physical_kv_cache(
-                        self.index_group_index
-                    ).view(kv_c_and_k_pe_cache.dtype)
-                else:
-                    prefill_cache, block_table, req_ids = (
-                        index_group.stage_prefill_rows(
-                            self.index_group_index,
-                            kv_c_and_k_pe_cache,
-                            attn_metadata,
-                        )
-                    )
-                    physical_topk, valid_counts = (
-                        triton_convert_req_index_to_global_index(
-                            req_ids,
-                            block_table,
-                            topk_indices[num_decode_tokens:],
-                            BLOCK_SIZE=attn_metadata.block_size,
-                            NUM_TOPK_TOKENS=topk_indices.shape[1],
-                            return_valid_counts=True,
-                        )
-                    )
-                outputs.append(
-                    self._run_mqa_kernel(
-                        q_nope[num_decode_tokens:],
-                        q_rope[num_decode_tokens:],
-                        prefill_cache,
-                        physical_topk,
-                        valid_counts,
-                        attn_metadata.block_size,
-                    )
-                )
-            return torch.cat(outputs) if len(outputs) > 1 else outputs[0], None
-
-        kv_rows, block_stride_rows = flat_kv_row_view(
-            kv_c_and_k_pe_cache, attn_metadata.block_size
-        )
-        if (
-            index_group is not None
-            and index_group.num_layers > 1
-            and attn_metadata.num_decode_tokens == num_actual_toks
-        ):
-            topk_indices, valid_counts = self._convert_logical_to_physical_topk(
-                topk_indices,
-                attn_metadata,
-                block_stride_rows=block_stride_rows,
-                return_valid_counts=True,
-            )
-        else:
-            topk_indices, valid_counts = triton_convert_req_index_to_global_index(
-                attn_metadata.req_id_per_token[:num_actual_toks],
-                attn_metadata.block_table,
-                topk_indices,
-                BLOCK_SIZE=attn_metadata.block_size,
-                BLOCK_STRIDE_ROWS=block_stride_rows,
-                NUM_TOPK_TOKENS=topk_indices.shape[1],
-                return_valid_counts=True,
-            )
-        return (
-            self._run_mqa_kernel(
-                q_nope,
-                q_rope,
-                kv_rows,
-                topk_indices,
-                valid_counts,
-                attn_metadata.block_size,
-                cache_is_flat=True,
-            ),
-            None,
-        )
-
-    def _run_mqa_kernel(
-        self,
-        q_nope: torch.Tensor,
-        q_rope: torch.Tensor,
-        kv_cache: torch.Tensor,
-        topk_indices: torch.Tensor,
-        valid_counts: torch.Tensor,
-        block_size: int,
-        *,
-        cache_is_flat: bool = False,
-    ) -> torch.Tensor:
-        kv_rows = (
-            kv_cache if cache_is_flat else flat_kv_row_view(kv_cache, block_size)[0]
-        )
+        kv_rows, _ = flat_kv_row_view(kv_cache, block_size)
 
         cu_seqlens_q = self.cu_seqlens_q_buffer[: q_rope.shape[0] + 1]
         v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1)
@@ -370,4 +256,167 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
             causal=True,
             fa_version=3,
         )
-        return out
+        return out, None
+
+
+def _fa4_cute_mla_available() -> str | None:
+    """Reason the FA4 cute-DSL MLA entry point is unusable, or None."""
+    from vllm.vllm_flash_attn import flash_attn_interface as fa
+
+    if not fa.is_fa_version_supported(4):
+        return f"FA4 unavailable: {fa.fa_version_unsupported_reason(4)}"
+    try:
+        from vllm.vllm_flash_attn.cute.interface import _flash_attn_fwd  # noqa: F401
+    except Exception as e:  # cute-DSL / cutlass import chain
+        return f"FA4 cute-DSL entry point failed to import: {e!r}"
+    return None
+
+
+class FlashAttnMLASparseFA4Backend(FlashInferMLASparseTRTLLMBackend):
+    """FA4 for is_decode kernel calls; trtllm-gen for the rest."""
+
+    supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
+    # The qv kernel asserts every descale is None: BF16 KV cache only.
+    supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "bfloat16"]
+
+    @staticmethod
+    def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]:
+        return [MultipleOf(64)]
+
+    @staticmethod
+    def get_name() -> str:
+        return "FLASH_ATTN_MLA_SPARSE_FA4"
+
+    @staticmethod
+    def get_builder_cls() -> type["FlashAttnMLASparseFA4MetadataBuilder"]:
+        return FlashAttnMLASparseFA4MetadataBuilder
+
+    @staticmethod
+    def get_impl_cls() -> type[MLAAttentionImpl[Any]]:
+        return FlashAttnMLASparseFA4Impl
+
+    @classmethod
+    def get_supported_head_sizes(cls) -> list[int]:
+        # 512 latent + 64 RoPE; split back apart for the qv kernel.
+        return [576]
+
+    @classmethod
+    def supports_combination(cls, *args: Any, **kwargs: Any) -> str | None:
+        from vllm.config import get_current_vllm_config_or_none
+        from vllm.utils.flashinfer import has_flashinfer
+
+        if not has_flashinfer():
+            return "FA4 sparse MLA runs its prefill batches on FlashInfer's kernel"
+        if (reason := _fa4_cute_mla_available()) is not None:
+            return reason
+        vllm_config = get_current_vllm_config_or_none()
+        if vllm_config is None:
+            return None
+        if vllm_config.model_config is not None:
+            hf_config = vllm_config.model_config.hf_text_config
+            topk = getattr(hf_config, "index_topk", None)
+            # gather_kv_indices' last dim must be a whole number of n-tiles.
+            if topk is None or topk % 128 != 0:
+                return f"FA4 sparse MLA requires index_topk % 128 == 0, got {topk}"
+            dims = (
+                getattr(hf_config, "kv_lora_rank", None),
+                getattr(hf_config, "qk_rope_head_dim", None),
+            )
+            if dims != (512, 64):
+                return (
+                    "FA4 sparse MLA requires (kv_lora_rank, qk_rope_head_dim) "
+                    f"== (512, 64), got {dims}"
+                )
+            # Under DCP the query is all-gathered: the kernel sees heads * dcp_size.
+            dcp_size = vllm_config.parallel_config.decode_context_parallel_size
+            num_heads = vllm_config.model_config.get_num_attention_heads(
+                vllm_config.parallel_config
+            )
+            # 128 gathered heads take FA's prefill kernel, unvalidated under DCP.
+            head_counts = (8, 16, 32, 64) if dcp_size > 1 else (8, 16, 32, 64, 128)
+            if num_heads * dcp_size not in head_counts:
+                return (
+                    f"FA4 sparse MLA requires {head_counts} gathered query heads, "
+                    f"got num_heads={num_heads} * dcp_size={dcp_size}"
+                )
+        return super().supports_combination(*args, **kwargs)
+
+
+class FlashAttnMLASparseFA4MetadataBuilder(FlashInferMLASparseTRTLLMMetadataBuilder):
+    """The trtllm-gen builder, for the prefill lane's DCP workspace pre-size.
+
+    Its UNIFORM_BATCH support and varlen decode bound also fit FA4, which runs
+    one query row per token and bakes the row count into its varlen scalars.
+    """
+
+
+@cache
+def _fa4_varlen_scalars(
+    num_tokens: int, num_kv_rows: int, device: torch.device
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Per-token varlen scalars; never evict, CUDA graphs bake in the pointers."""
+    cu_seqlens_q = torch.arange(num_tokens + 1, dtype=torch.int32, device=device)
+    cu_seqlens_k = torch.zeros(num_tokens + 1, dtype=torch.int32, device=device)
+    seqused_k = torch.full((num_tokens,), num_kv_rows, dtype=torch.int32, device=device)
+    return cu_seqlens_q, cu_seqlens_k, seqused_k
+
+
+class FlashAttnMLASparseFA4Impl(FlashInferMLASparseImpl):
+    # FA4 emits a natural-log LSE; prefill converts trtllm-gen's log2 LSE.
+    lse_base_on_e: bool = True
+
+    def _forward_mqa_kernel(
+        self,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+        kv_cache: torch.Tensor,
+        topk_indices: torch.Tensor,
+        valid_counts: torch.Tensor,
+        *,
+        layer: AttentionLayer,
+        block_size: int,
+        is_decode: bool,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if not is_decode:
+            out, lse = super()._forward_mqa_kernel(
+                q,
+                kv_cache,
+                topk_indices,
+                valid_counts,
+                layer=layer,
+                block_size=block_size,
+                is_decode=is_decode,
+            )
+            return out, None if lse is None else lse * math.log(2.0)
+        if isinstance(q, tuple):
+            ql_nope, q_pe = q
+        else:
+            # 16B-aligned halves of the fused query; FA4 reads them in place.
+            ql_nope, q_pe = q[..., : self.kv_lora_rank], q[..., self.kv_lora_rank :]
+        kv_rows, _ = flat_kv_row_view(kv_cache, block_size)
+        num_tokens = q_pe.shape[0]
+        num_kv_rows = kv_rows.shape[0]
+        assert self.topk_indices_buffer is not None
+        cu_seqlens_q, cu_seqlens_k, seqused_k = _fa4_varlen_scalars(
+            self.topk_indices_buffer.shape[0], num_kv_rows, kv_rows.device
+        )
+        # An all-sentinel row yields (0, -inf), the DCP merge identity; no masking.
+        kernel_out = flash_attn_varlen_func(
+            q=q_pe,
+            k=kv_rows[:, self.kv_lora_rank :].unsqueeze(1),
+            v=kv_rows[:, : self.kv_lora_rank].unsqueeze(1),
+            q_v=ql_nope,
+            max_seqlen_q=1,
+            cu_seqlens_q=cu_seqlens_q[: num_tokens + 1],
+            max_seqlen_k=num_kv_rows,
+            cu_seqlens_k=cu_seqlens_k[: num_tokens + 1],
+            seqused_k=seqused_k[:num_tokens],
+            gather_kv_indices=topk_indices,
+            # Part of FA4's compile key: pass it unconditionally.
+            gather_kv_valid_length=valid_counts,
+            softmax_scale=self.scale,
+            # Causality is in the top-k list; causal=True would clamp flat rows.
+            causal=False,
+            fa_version=4,
+            return_softmax_lse=self.need_to_return_lse_for_decode,
+        )
+        return kernel_out if self.need_to_return_lse_for_decode else (kernel_out, None)

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -34,8 +34,6 @@ from vllm.v1.attention.backends.mla.sparse_utils import (
     align_blocks_to_rows,
     flat_kv_row_view,
     prepare_sparse_mla_safe_lengths,
-    triton_convert_req_index_to_global_index,
-    triton_filter_and_convert_dcp_index,
 )
 from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec
 
@@ -129,14 +127,12 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
             and parallel_config.decode_context_parallel_size > 1
         ):
             return (
-                "FLASHINFER_MLA_SPARSE does not support combined PCP+DCP; "
+                f"{cls.get_name()} does not support combined PCP+DCP; "
                 "use FLASHMLA_SPARSE, which gathers each DCP KV shard before "
                 "running the rank-local PCP prefill queries"
             )
         if kv_cache_dtype == "fp8_ds_mla":
-            return (
-                "FLASHINFER_MLA_SPARSE SM10 does not support fp8_ds_mla kv-cache dtype"
-            )
+            return f"{cls.get_name()} SM10 does not support fp8_ds_mla kv-cache dtype"
 
         # FlashInfer MLA sparse SM10 kernel requires qk_nope_head_dim in [128, 192].
         if vllm_config.model_config is not None:
@@ -492,135 +488,6 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
         # for query and kv_cache (mixed bf16+fp8 is not supported).
         self.supports_quant_query_input = True
 
-    def forward_mqa(
-        self,
-        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
-        kv_c_and_k_pe_cache: torch.Tensor,
-        attn_metadata: FlashInferMLASparseMetadata,
-        layer: AttentionLayer,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        if isinstance(q, tuple):
-            ql_nope, q_pe = q
-            if q_pe.shape[-1] == 0 and ql_nope.is_contiguous():
-                q = ql_nope
-            else:
-                q = torch.cat(q, dim=-1)
-
-        num_actual_toks = q.shape[0]
-
-        assert self.topk_indices_buffer is not None
-        topk_indices = self.topk_indices_buffer[:num_actual_toks]
-
-        self._prepare_mqa_kernel(layer, q.device)
-
-        index_group = self.index_group
-        if isinstance(index_group, HiSparseMLAIndexGroup):
-            num_decode_tokens = attn_metadata.num_decode_tokens
-            decode_out: torch.Tensor | None = None
-            decode_lse: torch.Tensor | None = None
-            if num_decode_tokens > 0:
-                physical_topk, valid_counts = (
-                    index_group.convert_logical_to_physical_topk(
-                        self.index_group_index,
-                        topk_indices[:num_decode_tokens],
-                        attn_metadata,
-                        block_stride_rows=None,
-                        return_valid_counts=True,
-                    )
-                )
-                decode_out, decode_lse = self._run_mqa_kernel(
-                    q[:num_decode_tokens],
-                    index_group.physical_kv_cache(self.index_group_index).view(
-                        kv_c_and_k_pe_cache.dtype
-                    ),
-                    physical_topk,
-                    valid_counts,
-                )
-                if num_decode_tokens == num_actual_toks:
-                    return decode_out, decode_lse
-
-            cache = index_group.cache(self.index_group_index)
-            if num_decode_tokens == 0 and cache.all_context_pages_resident:
-                physical_topk, valid_counts = (
-                    index_group.convert_logical_to_physical_topk(
-                        self.index_group_index,
-                        topk_indices,
-                        attn_metadata,
-                        block_stride_rows=None,
-                        return_valid_counts=True,
-                    )
-                )
-                return self._run_mqa_kernel(
-                    q,
-                    index_group.physical_kv_cache(self.index_group_index).view(
-                        kv_c_and_k_pe_cache.dtype
-                    ),
-                    physical_topk,
-                    valid_counts,
-                )
-
-            prefill_cache, block_table, req_ids = index_group.stage_prefill_rows(
-                self.index_group_index, kv_c_and_k_pe_cache, attn_metadata
-            )
-            prefill_indices, prefill_lens = triton_convert_req_index_to_global_index(
-                req_ids,
-                block_table,
-                topk_indices[num_decode_tokens:],
-                BLOCK_SIZE=attn_metadata.block_size,
-                NUM_TOPK_TOKENS=topk_indices.shape[1],
-                return_valid_counts=True,
-            )
-            prefill_out, prefill_lse = self._run_mqa_kernel(
-                q[num_decode_tokens:],
-                prefill_cache,
-                prefill_indices,
-                prefill_lens,
-            )
-            if decode_out is None:
-                return prefill_out, prefill_lse
-            output = torch.cat((decode_out, prefill_out))
-            if decode_lse is None:
-                return output, None
-            assert prefill_lse is not None
-            return output, torch.cat((decode_lse, prefill_lse))
-
-        kv_rows, block_stride_rows = flat_kv_row_view(
-            kv_c_and_k_pe_cache, attn_metadata.block_size
-        )
-
-        if self.dcp_world_size > 1:
-            topk_indices_physical, seq_lens = triton_filter_and_convert_dcp_index(
-                attn_metadata.req_id_per_token[:num_actual_toks],
-                attn_metadata.block_table,
-                topk_indices,
-                dcp_size=self.dcp_world_size,
-                dcp_rank=self.dcp_rank,
-                cp_kv_cache_interleave_size=(attn_metadata.cp_kv_cache_interleave_size),
-                BLOCK_SIZE=attn_metadata.block_size,
-                BLOCK_STRIDE_ROWS=block_stride_rows,
-                NUM_TOPK_TOKENS=topk_indices.shape[1],
-                return_valid_counts=True,
-            )
-        else:
-            topk_indices_physical, seq_lens = self._convert_logical_to_physical_topk(
-                topk_indices,
-                attn_metadata,
-                block_stride_rows=block_stride_rows,
-                return_valid_counts=True,
-            )
-
-        return self._run_mqa_kernel(
-            q,
-            # Block sizes and strides have been aligned to 32 rows.
-            kv_rows.view(
-                -1,
-                math.gcd(block_stride_rows, attn_metadata.block_size, 64),
-                kv_rows.shape[-1],
-            ),
-            topk_indices_physical,
-            seq_lens,
-        )
-
     def _prepare_mqa_kernel(
         self,
         layer: AttentionLayer,
@@ -648,7 +515,6 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
         kv_cache = runtime.hot.attention_cache
         num_tokens = runtime.max_num_reqs
         topk_tokens = self.topk_indices_buffer.shape[1]
-        self._prepare_mqa_kernel(layer, kv_cache.device)
 
         q_dtype = (
             current_platform.fp8_dtype()
@@ -679,22 +545,48 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
             dtype=torch.int32,
             device=kv_cache.device,
         )
-        self._run_mqa_kernel(q, kv_cache, topk_indices, seq_lens)
+        self._forward_mqa_kernel(
+            q,
+            kv_cache,
+            topk_indices,
+            seq_lens,
+            layer=layer,
+            block_size=runtime.hot.block_size,
+            is_decode=True,
+        )
 
-    def _run_mqa_kernel(
+    def _forward_mqa_kernel(
         self,
-        q: torch.Tensor,
+        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
         kv_cache: torch.Tensor,
         topk_indices: torch.Tensor,
         seq_lens: torch.Tensor,
+        *,
+        layer: AttentionLayer,
+        block_size: int,
+        is_decode: bool,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if isinstance(q, tuple):
+            ql_nope, q_pe = q
+            if q_pe.shape[-1] == 0 and ql_nope.is_contiguous():
+                q = ql_nope
+            else:
+                q = torch.cat(q, dim=-1)
+
+        self._prepare_mqa_kernel(layer, q.device)
         assert self._workspace_buffer is not None
         assert self.bmm1_scale is not None
         assert self.bmm2_scale is not None
 
         from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
 
-        kv_cache = kv_cache.view(q.dtype)
+        kv_rows, block_stride_rows = flat_kv_row_view(
+            kv_cache.view(q.dtype), block_size
+        )
+        # Block sizes and strides have been aligned to 32 rows.
+        kv_cache = kv_rows.view(
+            -1, math.gcd(block_stride_rows, block_size, 64), kv_rows.shape[-1]
+        )
 
         # Single-token sparse decode. trtllm-gen requires the q_len_per_request
         # dim, but the sparse attention mask is fully per-token (each query token

vllm/v1/attention/backends/registry.py

@@ -125,6 +125,10 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
     FLASH_ATTN_MLA_SPARSE = (
         "vllm.v1.attention.backends.mla.flashattn_mla_sparse.FlashAttnMLASparseBackend"
     )
+    FLASH_ATTN_MLA_SPARSE_FA4 = (
+        "vllm.v1.attention.backends.mla.flashattn_mla_sparse."
+        "FlashAttnMLASparseFA4Backend"
+    )
     MINIMAX_M3_SPARSE = (
         "vllm.models.minimax_m3.common.sparse_attention.MiniMaxM3SparseBackend"
     )

vllm/vllm_flash_attn/flash_attn_interface.py

@@ -214,6 +214,8 @@ def flash_attn_varlen_func(
     aux_tensors=None,
     aux_tensor_leading_dims=None,
     dynamic_causal: "torch.Tensor | None" = None,
+    gather_kv_indices=None,
+    gather_kv_valid_length=None,
 ):
     """dropout_p should be set to 0.0 during evaluation
     Supports multi-query and grouped-query attention (MQA/GQA) by passing in K, V with fewer heads
@@ -262,9 +264,10 @@ def flash_attn_varlen_func(
            testing only. The returned probabilities are not guaranteed to be correct
            (they might not have the right scaling).
         seqused_k: (batch_size,), dtype torch.int32. Actual number of key tokens per
-           sequence. Mutually exclusive with cu_seqlens_k, and required when
+           sequence. Mutually exclusive with cu_seqlens_k except on FA4, which reads
+           cu_seqlens_k as the KV offset and seqused_k as the length. Required when
            block_table is given.
-        q_v: optional value tensor fused with q. FA3 only.
+        q_v: optional value tensor fused with q. FA3 and FA4 only.
         block_table: (batch_size, max_blocks_per_seq), dtype torch.int32. Paged KV
            cache block table. Requires seqused_k.
         return_softmax_lse: bool. Whether to also return softmax_lse.
@@ -289,18 +292,21 @@ def flash_attn_varlen_func(
         aux_tensors: auxiliary tensors consumed by mask_mod. FA4 only.
         aux_tensor_leading_dims: leading dimensions of each entry in aux_tensors.
         dynamic_causal: optional per-sequence causal offsets. FA4 only.
+        gather_kv_indices: (total_q, k) int32 KV rows per query token, `-1` padded,
+            k a multiple of 128. Requires q_v, excludes block_table. FA4 only.
+        gather_kv_valid_length: (total_q,) int32, real entries per row. FA4 only.
 
     Return:
         out: (total, nheads, headdim).
-        softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen). The
-            logsumexp of each row of the matrix QK^T * scaling (e.g., log of the softmax
-            normalization factor).
+        softmax_lse [optional, if return_softmax_lse=True]: (nheads, total_q_seqlen), or
+            (total_q, nheads) on FA4's q_v path. The logsumexp of each row of the matrix
+            QK^T * scaling (e.g., log of the softmax normalization factor).
 
     """
     assert cu_seqlens_k is not None or seqused_k is not None, (
         "cu_seqlens_k or seqused_k must be provided"
     )
-    assert cu_seqlens_k is None or seqused_k is None, (
+    assert fa_version == 4 or cu_seqlens_k is None or seqused_k is None, (
         "cu_seqlens_k and seqused_k cannot be provided at the same time"
     )
     assert block_table is None or seqused_k is not None, (
@@ -311,6 +317,9 @@ def flash_attn_varlen_func(
         f"Fused FP8 output (output_scale) is only supported by FA4, "
         f"got fa_version={fa_version}"
     )
+    assert fa_version == 4 or (
+        gather_kv_indices is None and gather_kv_valid_length is None
+    ), f"gather_kv_* is only supported by FA4, got fa_version={fa_version}"
 
     if softmax_scale is None:
         softmax_scale = q.shape[-1] ** (-0.5)
@@ -446,12 +455,15 @@ def flash_attn_varlen_func(
             q,
             k,
             v,
+            qv=q_v,
             cu_seqlens_q=cu_seqlens_q,
             cu_seqlens_k=cu_seqlens_k,
             seqused_k=seqused_k,
             max_seqlen_q=max_seqlen_q,
             max_seqlen_k=max_seqlen_k,
             page_table=block_table,
+            gather_kv_indices=gather_kv_indices,
+            gather_kv_valid_length=gather_kv_valid_length,
             softmax_scale=softmax_scale,
             causal=causal,
             dynamic_causal=dynamic_causal,