[Attention] Enable FA4 sparse MQA decode on Blackwell

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 unexpected_helper(*args, **kwargs):
             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"