karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[Bugfix] Enable FlashAttention MLA prefill for Mistral Small 4 head dims (#48855)

Browse Source 
Signed-off-by: juliendenize <julien.denize@mistral.ai>
Co-authored-by: Nicolò Lucchesi <nlucches@redhat.com>
This commit is contained in:
Julien Denize
2026-07-17 18:07:00 +08:00
committed by GitHub
co-authored by Nicolò Lucchesi
parent d5b1ec2684
commit ce4bdcbda4
3 changed files with 10 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/attention_backends.md
+1 -1
View File
@@ -205,7 +205,7 @@ hardware and configuration.
| Backend | Description | Dtypes | Compute Cap. | Notes |
| ------- | ----------- | ------ | ------------ | ----- |
| `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) |
| `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) |
| `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 |
| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
| `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
View File
@@ -125,8 +125,8 @@ class TestGetMLAPrefillBackend:
@pytest.mark.parametrize(
("qk_nope_head_dim", "v_head_dim"),
[(128, 128), (192, 256)],
ids=["deepseek", "glm"],
[(128, 128), (192, 256), (64, 128)],
ids=["deepseek", "glm", "mistral_s4"],
)
def test_auto_selection_on_hopper(self, qk_nope_head_dim: int, v_head_dim: int):
try:
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+7 -2
View File
@@ -63,11 +63,16 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
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 == dims_deepseek
return mla_dimensions in [dims_deepseek, dims_mistral_s4]
else:
return mla_dimensions in [dims_deepseek, dims_glm]
return mla_dimensions in [dims_deepseek, dims_glm, dims_mistral_s4]
def __init__(
self,