[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>
Matthew Bonanni committed
Jul 27, 2026 at 14:25 UTC
bf2b45b5d6a991f47702d4d87c81bb0ef8619a65
7 files changed
+66
-34
benchmarks/attention_benchmarks/benchmark.py
+4
@@ -1358,6 +1358,10 @@ def main():
1358
profile_memory=args.profile_memory,
1359
warmup_ms=args.warmup_ms,
1360
prefill_backend=pb,
1361
+ kv_lora_rank=args.kv_lora_rank,
1362
+ qk_nope_head_dim=args.qk_nope_head_dim,
1363
+ qk_rope_head_dim=args.qk_rope_head_dim,
1364
+ v_head_dim=args.v_head_dim,
1365
)
1366
1367
result = run_benchmark(config)
docs/design/attention_backends.md
+3
-1
@@ -153,7 +153,9 @@ hardware and configuration.
153
154
> **‡** Automatic selection tries FlashAttention first. On Blackwell
155
> (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then
156
-> TokenSpeed MLA. On other GPUs, only FlashAttention is considered.
156
+> TokenSpeed MLA; for (qk_nope_head_dim=192, qk_rope_head_dim=64,
157
+> v_head_dim=256) TRT-LLM Ragged is tried before FlashAttention. On other
158
+> GPUs, only FlashAttention is considered.
159
160
### Decode Backends
161
tests/v1/attention/test_mla_prefill_selector.py
+7
-7
@@ -163,7 +163,7 @@ class TestGetMLAPrefillBackend:
163
class TestAutoSelectMLAPrefillBackend:
164
"""Tests for fallback and error paths in auto-selection."""
165
166
- def test_blackwell_glm_dimensions_fall_back_to_trtllm(self):
166
+ def test_blackwell_glm_dimensions_use_trtllm(self):
167
capability = DeviceCapability(major=10, minor=0)
168
selector_config = MLAPrefillSelectorConfig(
169
dtype=torch.bfloat16,
@@ -184,11 +184,6 @@ class TestAutoSelectMLAPrefillBackend:
184
with (
185
patch("vllm.platforms.current_platform") as mock_platform,
186
patch.object(flash_attn_cls, "is_available", return_value=True),
187
- patch(
188
- "vllm.v1.attention.backends.mla.prefill.flash_attn."
189
- "get_flash_attn_version",
190
- return_value=4,
191
- ),
187
patch.object(trtllm_cls, "validate_configuration", return_value=[]),
188
):
189
# Force the non-ROCm priority on the Blackwell.
@@ -300,7 +295,12 @@ class TestROCmAiterFAPrefillSelection:
295
with patch("vllm.platforms.current_platform") as mock_platform:
296
mock_platform.is_rocm.return_value = True
297
priorities = _get_mla_prefill_backend_priorities(
303
- DeviceCapability(major=9, minor=5)
298
+ DeviceCapability(major=9, minor=5),
299
+ MLADimensions(
300
+ qk_nope_head_dim=128,
301
+ qk_rope_head_dim=64,
302
+ v_head_dim=128,
303
+ ),
304
)
305
306
assert priorities == [
vllm/v1/attention/backends/fa_utils.py
+17
-5
@@ -134,6 +134,7 @@ def get_flash_attn_version(
134
head_size: int | None = None,
135
head_size_v: int | None = None,
136
has_sinks: bool = False,
137
+ requires_local_attention: bool = False,
138
) -> int | None:
139
if current_platform.is_xpu():
140
return 2
@@ -230,17 +231,28 @@ def get_flash_attn_version(
231
)
232
fa_version = 2
233
234
+ if (
235
+ fa_version == 4
236
+ and device_capability.major >= 10
237
+ and head_size == 256
238
+ and requires_local_attention
239
+ ):
240
+ logger.warning_once(
241
+ "FA4 on Blackwell does not support local attention with "
242
+ "head_size=256, defaulting to FA version 2."
243
+ )
244
+ fa_version = 2
245
+
246
# FA4 on SM100 (Blackwell) has TMEM capacity limits that restrict
234
- # supported head dimensions.
235
- # See: https://github.com/Dao-AILab/flash-attention/issues/1959
236
- # Exception: hdim 192 is supported for MLA's diff-headdim case
237
- # (qk=192, v=128), added upstream in commits 1a15733e/1b36ab19.
247
+ # supported head dimensions to ≤128, with exceptions for 256 and 192/128 (MLA
248
+ # prefill). Development of symmetric 192, 384, and 512 support is being tracked
249
+ # in https://github.com/Dao-AILab/flash-attention/issues/2456
250
if (
251
fa_version == 4
252
and device_capability.major >= 10
253
and head_size is not None
254
and head_size > 128
243
- and head_size != 192
255
+ and not (head_size == 256 or (head_size == 192 and head_size_v == 128))
256
):
257
logger.warning_once(
258
"FA4 on Blackwell does not support head_size=%d due to TMEM "
vllm/v1/attention/backends/flash_attn.py
+1
@@ -782,6 +782,7 @@ class FlashAttentionImpl(AttentionImpl):
782
self.attn_type = attn_type
783
self.vllm_flash_attn_version = get_flash_attn_version(
784
requires_alibi=alibi_slopes is not None,
785
+ requires_local_attention=sliding_window is not None,
786
head_size=head_size,
787
has_sinks=sinks is not None,
788
)
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+17
-20
@@ -288,26 +288,23 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
288
289
@classmethod
290
def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool:
291
- dims_deepseek = MLADimensions(
292
- qk_nope_head_dim=128,
293
- qk_rope_head_dim=64,
294
- v_head_dim=128,
295
- )
296
- dims_glm = MLADimensions(
297
- qk_nope_head_dim=192,
298
- qk_rope_head_dim=64,
299
- v_head_dim=256,
300
- )
301
- dims_mistral_s4 = MLADimensions(
302
- qk_nope_head_dim=64,
303
- qk_rope_head_dim=64,
304
- v_head_dim=128,
305
- )
306
- fa_version = get_flash_attn_version()
307
- if fa_version == 4:
308
- return mla_dimensions in [dims_deepseek, dims_mistral_s4]
309
- else:
310
- return mla_dimensions in [dims_deepseek, dims_glm, dims_mistral_s4]
291
+ return mla_dimensions in [
292
+ MLADimensions(
293
+ qk_nope_head_dim=128,
294
+ qk_rope_head_dim=64,
295
+ v_head_dim=128,
296
+ ),
297
+ MLADimensions(
298
+ qk_nope_head_dim=192,
299
+ qk_rope_head_dim=64,
300
+ v_head_dim=256,
301
+ ),
302
+ MLADimensions(
303
+ qk_nope_head_dim=64,
304
+ qk_rope_head_dim=64,
305
+ v_head_dim=128,
306
+ ),
307
+ ]
308
309
def __init__(
310
self,
vllm/v1/attention/backends/mla/prefill/selector.py
+17
-1
@@ -47,11 +47,13 @@ class MLAPrefillSelectorConfig(NamedTuple):
47
48
def _get_mla_prefill_backend_priorities(
49
device_capability: DeviceCapability,
50
+ mla_dimensions: MLADimensions,
51
) -> list[MLAPrefillBackendEnum]:
52
"""Get MLA prefill backend priorities based on device capability.
53
54
Args:
55
device_capability: The device's compute capability.
56
+ mla_dimensions: The model's MLA head dimensions.
57
58
Returns:
59
List of backends in priority order (highest priority first).
@@ -65,6 +67,17 @@ def _get_mla_prefill_backend_priorities(
67
]
68
69
if device_capability.major == 10: # Blackwell
70
+ if mla_dimensions == MLADimensions(
71
+ qk_nope_head_dim=192,
72
+ qk_rope_head_dim=64,
73
+ v_head_dim=256,
74
+ ):
75
+ return [
76
+ MLAPrefillBackendEnum.TRTLLM_RAGGED,
77
+ MLAPrefillBackendEnum.FLASH_ATTN,
78
+ MLAPrefillBackendEnum.FLASHINFER,
79
+ MLAPrefillBackendEnum.TOKENSPEED_MLA,
80
+ ]
81
return [
82
MLAPrefillBackendEnum.FLASH_ATTN,
83
MLAPrefillBackendEnum.TRTLLM_RAGGED,
@@ -157,7 +170,10 @@ def _auto_select_mla_prefill_backend(
170
Returns:
171
The selected prefill backend class.
172
"""
160
- priorities = _get_mla_prefill_backend_priorities(device_capability)
173
+ priorities = _get_mla_prefill_backend_priorities(
174
+ device_capability,
175
+ selector_config.mla_dimensions,
176
+ )
177
all_invalid_reasons: dict[str, list[str]] = {}
178
179
for backend_enum in priorities: