[Attention] Integrate FlashAttention 4 SM100 headdim 256 support (#42669)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

benchmarks/attention_benchmarks/benchmark.py

@@ -1358,6 +1358,10 @@ def main():
                             profile_memory=args.profile_memory,
                             warmup_ms=args.warmup_ms,
                             prefill_backend=pb,
+                            kv_lora_rank=args.kv_lora_rank,
+                            qk_nope_head_dim=args.qk_nope_head_dim,
+                            qk_rope_head_dim=args.qk_rope_head_dim,
+                            v_head_dim=args.v_head_dim,
                         )
 
                         result = run_benchmark(config)

docs/design/attention_backends.md

@@ -153,7 +153,9 @@ hardware and configuration.
 
 > **‡** Automatic selection tries FlashAttention first. On Blackwell
 > (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then
-> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.
+> TokenSpeed MLA; for (qk_nope_head_dim=192, qk_rope_head_dim=64,
+> v_head_dim=256) TRT-LLM Ragged is tried before FlashAttention. On other
+> GPUs, only FlashAttention is considered.
 
 ### Decode Backends
 

tests/v1/attention/test_mla_prefill_selector.py

@@ -163,7 +163,7 @@ class TestGetMLAPrefillBackend:
 class TestAutoSelectMLAPrefillBackend:
     """Tests for fallback and error paths in auto-selection."""
 
-    def test_blackwell_glm_dimensions_fall_back_to_trtllm(self):
+    def test_blackwell_glm_dimensions_use_trtllm(self):
         capability = DeviceCapability(major=10, minor=0)
         selector_config = MLAPrefillSelectorConfig(
             dtype=torch.bfloat16,
@@ -184,11 +184,6 @@ class TestAutoSelectMLAPrefillBackend:
         with (
             patch("vllm.platforms.current_platform") as mock_platform,
             patch.object(flash_attn_cls, "is_available", return_value=True),
-            patch(
-                "vllm.v1.attention.backends.mla.prefill.flash_attn."
-                "get_flash_attn_version",
-                return_value=4,
-            ),
             patch.object(trtllm_cls, "validate_configuration", return_value=[]),
         ):
             # Force the non-ROCm priority on the Blackwell.
@@ -300,7 +295,12 @@ class TestROCmAiterFAPrefillSelection:
         with patch("vllm.platforms.current_platform") as mock_platform:
             mock_platform.is_rocm.return_value = True
             priorities = _get_mla_prefill_backend_priorities(
-                DeviceCapability(major=9, minor=5)
+                DeviceCapability(major=9, minor=5),
+                MLADimensions(
+                    qk_nope_head_dim=128,
+                    qk_rope_head_dim=64,
+                    v_head_dim=128,
+                ),
             )
 
         assert priorities == [

vllm/v1/attention/backends/fa_utils.py

@@ -134,6 +134,7 @@ def get_flash_attn_version(
     head_size: int | None = None,
     head_size_v: int | None = None,
     has_sinks: bool = False,
+    requires_local_attention: bool = False,
 ) -> int | None:
     if current_platform.is_xpu():
         return 2
@@ -230,17 +231,28 @@ def get_flash_attn_version(
             )
             fa_version = 2
 
+        if (
+            fa_version == 4
+            and device_capability.major >= 10
+            and head_size == 256
+            and requires_local_attention
+        ):
+            logger.warning_once(
+                "FA4 on Blackwell does not support local attention with "
+                "head_size=256, defaulting to FA version 2."
+            )
+            fa_version = 2
+
         # FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict
-        # supported head dimensions.
-        # See: https://github.com/Dao-AILab/flash-attention/issues/1959
-        # Exception: hdim 192 is supported for MLA's diff-headdim case
-        # (qk=192, v=128), added upstream in commits 1a15733e/1b36ab19.
+        # supported head dimensions to ≤128, with exceptions for 256 and 192/128 (MLA
+        # prefill). Development of symmetric 192, 384, and 512 support is being tracked
+        # in https://github.com/Dao-AILab/flash-attention/issues/2456
         if (
             fa_version == 4
             and device_capability.major >= 10
             and head_size is not None
             and head_size > 128
-            and head_size != 192
+            and not (head_size == 256 or (head_size == 192 and head_size_v == 128))
         ):
             logger.warning_once(
                 "FA4 on Blackwell does not support head_size=%d due to TMEM "

vllm/v1/attention/backends/flash_attn.py

@@ -782,6 +782,7 @@ class FlashAttentionImpl(AttentionImpl):
         self.attn_type = attn_type
         self.vllm_flash_attn_version = get_flash_attn_version(
             requires_alibi=alibi_slopes is not None,
+            requires_local_attention=sliding_window is not None,
             head_size=head_size,
             has_sinks=sinks is not None,
         )

vllm/v1/attention/backends/mla/prefill/flash_attn.py

@@ -288,26 +288,23 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
 
     @classmethod
     def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool:
-        dims_deepseek = MLADimensions(
-            qk_nope_head_dim=128,
-            qk_rope_head_dim=64,
-            v_head_dim=128,
-        )
-        dims_glm = MLADimensions(
-            qk_nope_head_dim=192,
-            qk_rope_head_dim=64,
-            v_head_dim=256,
-        )
-        dims_mistral_s4 = MLADimensions(
-            qk_nope_head_dim=64,
-            qk_rope_head_dim=64,
-            v_head_dim=128,
-        )
-        fa_version = get_flash_attn_version()
-        if fa_version == 4:
-            return mla_dimensions in [dims_deepseek, dims_mistral_s4]
-        else:
-            return mla_dimensions in [dims_deepseek, dims_glm, dims_mistral_s4]
+        return mla_dimensions in [
+            MLADimensions(
+                qk_nope_head_dim=128,
+                qk_rope_head_dim=64,
+                v_head_dim=128,
+            ),
+            MLADimensions(
+                qk_nope_head_dim=192,
+                qk_rope_head_dim=64,
+                v_head_dim=256,
+            ),
+            MLADimensions(
+                qk_nope_head_dim=64,
+                qk_rope_head_dim=64,
+                v_head_dim=128,
+            ),
+        ]
 
     def __init__(
         self,

vllm/v1/attention/backends/mla/prefill/selector.py

@@ -47,11 +47,13 @@ class MLAPrefillSelectorConfig(NamedTuple):
 
 def _get_mla_prefill_backend_priorities(
     device_capability: DeviceCapability,
+    mla_dimensions: MLADimensions,
 ) -> list[MLAPrefillBackendEnum]:
     """Get MLA prefill backend priorities based on device capability.
 
     Args:
         device_capability: The device's compute capability.
+        mla_dimensions: The model's MLA head dimensions.
 
     Returns:
         List of backends in priority order (highest priority first).
@@ -65,6 +67,17 @@ def _get_mla_prefill_backend_priorities(
         ]
 
     if device_capability.major == 10:  # Blackwell
+        if mla_dimensions == MLADimensions(
+            qk_nope_head_dim=192,
+            qk_rope_head_dim=64,
+            v_head_dim=256,
+        ):
+            return [
+                MLAPrefillBackendEnum.TRTLLM_RAGGED,
+                MLAPrefillBackendEnum.FLASH_ATTN,
+                MLAPrefillBackendEnum.FLASHINFER,
+                MLAPrefillBackendEnum.TOKENSPEED_MLA,
+            ]
         return [
             MLAPrefillBackendEnum.FLASH_ATTN,
             MLAPrefillBackendEnum.TRTLLM_RAGGED,
@@ -157,7 +170,10 @@ def _auto_select_mla_prefill_backend(
     Returns:
         The selected prefill backend class.
     """
-    priorities = _get_mla_prefill_backend_priorities(device_capability)
+    priorities = _get_mla_prefill_backend_priorities(
+        device_capability,
+        selector_config.mla_dimensions,
+    )
     all_invalid_reasons: dict[str, list[str]] = {}
 
     for backend_enum in priorities: