[Bugfix] Enable FlashAttention MLA prefill for Mistral Small 4 head dims (#48855)
Signed-off-by: juliendenize <julien.denize@mistral.ai> Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
Julien Denize committed
Jul 17, 2026 at 12:07 UTC
ce4bdcbda4f3c2da67b2fe241e80937ccdcc1fc6
3 files changed
+10
-5
docs/design/attention_backends.md
+1
-1
@@ -205,7 +205,7 @@ hardware and configuration.
205
206
| Backend | Description | Dtypes | Compute Cap. | Notes |
207
| ------- | ----------- | ------ | ------------ | ----- |
208
-| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) |
208
+| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=64, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) |
209
| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only |
210
| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
211
| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
tests/v1/attention/test_mla_prefill_selector.py
+2
-2
@@ -125,8 +125,8 @@ class TestGetMLAPrefillBackend:
125
126
@pytest.mark.parametrize(
127
("qk_nope_head_dim", "v_head_dim"),
128
- [(128, 128), (192, 256)],
129
- ids=["deepseek", "glm"],
128
+ [(128, 128), (192, 256), (64, 128)],
129
+ ids=["deepseek", "glm", "mistral_s4"],
130
)
131
def test_auto_selection_on_hopper(self, qk_nope_head_dim: int, v_head_dim: int):
132
try:
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+7
-2
@@ -63,11 +63,16 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
63
qk_rope_head_dim=64,
64
v_head_dim=256,
65
)
66
+ dims_mistral_s4 = MLADimensions(
67
+ qk_nope_head_dim=64,
68
+ qk_rope_head_dim=64,
69
+ v_head_dim=128,
70
+ )
71
fa_version = get_flash_attn_version()
72
if fa_version == 4:
68
- return mla_dimensions == dims_deepseek
73
+ return mla_dimensions in [dims_deepseek, dims_mistral_s4]
74
else:
70
- return mla_dimensions in [dims_deepseek, dims_glm]
75
+ return mla_dimensions in [dims_deepseek, dims_glm, dims_mistral_s4]
76
77
def __init__(
78
self,