[perf] wire FA and FlashMLA for sm90 GLM5Next NoPE SparseMLA (#55385)
Signed-off-by: JaredforReal <w13431838023@gmail.com> Co-authored-by: Kimi Code <noreply@moonshot.cn> Co-authored-by: Leoyzen <leoyzen@gmail.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>
Jared Wen committed
Sep 22, 2026 at 17:03 UTC
91d7324cb19d301c72d849e457221ee8dd645024
4 files changed
+225
-5
cmake/external_projects/vllm_flash_attn.cmake
+1
-1
@@ -39,7 +39,7 @@ else()
39
FetchContent_Declare(
40
vllm-flash-attn
41
GIT_REPOSITORY https://github.com/vllm-project/flash-attention.git
42
- GIT_TAG 506341a143fcabd4bb79052a7605ada727d6b3f5
42
+ GIT_TAG 9cd61de38763d712bb6ce56e2a02cc2bf718c89f
43
GIT_PROGRESS TRUE
44
# Don't share the vllm-flash-attn build between build types
45
BINARY_DIR ${CMAKE_BINARY_DIR}/vllm-flash-attn
tests/v1/attention/test_flashmla_nope_sm90_backend_selection.py
new
+176
@@ -0,0 +1,176 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""CPU tests for the FlashMLA sparse backend NoPE-512 (GLM5Next) gates.
4
+
5
+GLM5Next is rope-free MLA: qk_nope_head_dim=256, qk_rope_head_dim=0,
6
+v_head_dim=256, kv_lora_rank=512 -> head_size 512. These tests pin the
7
+selection contract only (no GPU required):
8
+
9
+* ``get_supported_head_sizes`` advertises 512 so the backend is a candidate;
10
+* ``supports_combination`` admits exactly two shapes for 512 -- a quantized
11
+ DS-MLA cache format (zero-padded 576/656B envelope) and a rope-free bf16
12
+ cache on SM90 -- and rejects everything else;
13
+* adding 512 must NOT reorder the SM100 priority list for bf16, which is why
14
+ the bf16 arm is restricted to SM90.
15
+"""
16
+
17
+from types import SimpleNamespace
18
+
19
+import pytest
20
+import torch
21
+
22
+from vllm.platforms.interface import DeviceCapability
23
+from vllm.v1.attention.backends.mla.flashmla_sparse import (
24
+ QUANTIZED_DS_MLA_CACHE_FORMATS,
25
+ FlashMLASparseBackend,
26
+)
27
+
28
+SM90 = DeviceCapability(major=9, minor=0)
29
+SM100 = DeviceCapability(major=10, minor=0)
30
+
31
+BF16_CACHE_DTYPES = (None, "auto", "bfloat16", "float16")
32
+QUANTIZED_CACHE_DTYPES = ("fp8_ds_mla", "nvfp4_ds_mla")
33
+
34
+
35
+def _supports_combination(head_size, kv_cache_dtype, capability):
36
+ return FlashMLASparseBackend.supports_combination(
37
+ head_size,
38
+ torch.bfloat16,
39
+ kv_cache_dtype,
40
+ 64,
41
+ use_mla=True,
42
+ has_sink=False,
43
+ use_sparse=True,
44
+ use_mm_prefix=False,
45
+ device_capability=capability,
46
+ )
47
+
48
+
49
+@pytest.fixture
50
+def rope_free_model(monkeypatch):
51
+ """Patch get_current_vllm_config so the model looks like GLM5Next NoPE."""
52
+ import vllm.config as cfg
53
+
54
+ monkeypatch.setattr(
55
+ cfg,
56
+ "get_current_vllm_config",
57
+ lambda: SimpleNamespace(
58
+ model_config=SimpleNamespace(
59
+ hf_text_config=SimpleNamespace(qk_rope_head_dim=0)
60
+ )
61
+ ),
62
+ )
63
+
64
+
65
+@pytest.fixture
66
+def rope_carrying_model(monkeypatch):
67
+ """Patch get_current_vllm_config with a rope-carrying (576-style) model."""
68
+ import vllm.config as cfg
69
+
70
+ monkeypatch.setattr(
71
+ cfg,
72
+ "get_current_vllm_config",
73
+ lambda: SimpleNamespace(
74
+ model_config=SimpleNamespace(
75
+ hf_text_config=SimpleNamespace(qk_rope_head_dim=64)
76
+ )
77
+ ),
78
+ )
79
+
80
+
81
+def test_supported_head_sizes_include_512():
82
+ assert FlashMLASparseBackend.get_supported_head_sizes() == [576, 512]
83
+
84
+
85
+def test_quantized_ds_mla_formats_are_the_envelope_set():
86
+ assert frozenset({"fp8_ds_mla", "nvfp4_ds_mla"}) == QUANTIZED_DS_MLA_CACHE_FORMATS
87
+
88
+
89
+@pytest.mark.parametrize("kv_cache_dtype", QUANTIZED_CACHE_DTYPES)
90
+def test_nope_512_quantized_rejected_here(kv_cache_dtype):
91
+ # Quantized DS-MLA NoPE-512 is served by the zero-padded 576/656B
92
+ # envelope, wired up in a separate change; this backend must reject it.
93
+ reason = _supports_combination(512, kv_cache_dtype, SM90)
94
+ assert reason is not None
95
+
96
+
97
+@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
98
+def test_nope_512_bf16_sm90_rope_free_accepted(kv_cache_dtype, rope_free_model):
99
+ assert _supports_combination(512, kv_cache_dtype, SM90) is None
100
+
101
+
102
+@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
103
+def test_nope_512_bf16_sm90_rope_carrying_rejected(kv_cache_dtype, rope_carrying_model):
104
+ reason = _supports_combination(512, kv_cache_dtype, SM90)
105
+ assert reason is not None and "rope-free" in reason
106
+
107
+
108
+@pytest.mark.parametrize("kv_cache_dtype", BF16_CACHE_DTYPES)
109
+def test_nope_512_bf16_sm100_rejected(kv_cache_dtype, rope_free_model):
110
+ # The SM90 restriction is what keeps the SM100 priority order unchanged.
111
+ reason = _supports_combination(512, kv_cache_dtype, SM100)
112
+ assert reason is not None
113
+
114
+
115
+@pytest.mark.parametrize("kv_cache_dtype", ("fp8", "fp8_e4m3"))
116
+def test_nope_512_plain_fp8_rejected(kv_cache_dtype, rope_free_model):
117
+ reason = _supports_combination(512, kv_cache_dtype, SM90)
118
+ assert reason is not None
119
+
120
+
121
+@pytest.mark.parametrize("capability", [SM90, SM100], ids=["sm90", "sm100"])
122
+@pytest.mark.parametrize(
123
+ "kv_cache_dtype",
124
+ ("auto", "bfloat16", "fp8", "fp8_ds_mla", "nvfp4_ds_mla"),
125
+)
126
+def test_576_unchanged_for_all_dtypes(kv_cache_dtype, capability):
127
+ # DeepSeek 576 must be untouched by the NoPE-512 gates, except for the
128
+ # pre-existing SM100-only nvfp4 rule.
129
+ reason = _supports_combination(576, kv_cache_dtype, capability)
130
+ if kv_cache_dtype == "nvfp4_ds_mla" and capability.major != 10:
131
+ assert reason is not None and "SM100" in reason
132
+ else:
133
+ assert reason is None, reason
134
+
135
+
136
+def _sparse_order(kv_cache_dtype, head_size, num_heads=32, capability=SM100):
137
+ try:
138
+ from vllm.platforms.cuda import _get_backend_priorities
139
+ except ImportError as e: # pragma: no cover - depends on host
140
+ pytest.skip(f"CUDA platform unavailable on this host: {e}")
141
+
142
+ priorities = _get_backend_priorities(
143
+ use_mla=True,
144
+ device_capability=capability,
145
+ num_heads=num_heads,
146
+ kv_cache_dtype=kv_cache_dtype,
147
+ head_size=head_size,
148
+ )
149
+ sparse = {
150
+ "FLASH_ATTN_MLA_SPARSE",
151
+ "FLASHMLA_SPARSE",
152
+ "FLASHINFER_MLA_SPARSE",
153
+ "FLASHINFER_MLA_SPARSE_SM90",
154
+ }
155
+ return [b.name for b in priorities if b.name in sparse]
156
+
157
+
158
+@pytest.mark.parametrize("num_heads", [32, 64], ids=["tp2", "tp1"])
159
+def test_sm100_bf16_512_priority_unchanged(num_heads):
160
+ """Adding head_size 512 must not reorder the SM100 bf16 sparse list.
161
+
162
+ On SM100 with >16 heads the hardcoded bf16 order is
163
+ [FLASHMLA_SPARSE, FLASHINFER_MLA_SPARSE] (FlashMLA first). head_size 512
164
+ must leave this untouched (identical to the 576 order); the guard that
165
+ keeps bf16-512 from *serving* through FlashMLA on SM100 lives in
166
+ supports_combination (gated to SM90), tested separately.
167
+ """
168
+ order512 = _sparse_order("auto", 512, num_heads=num_heads)
169
+ order576 = _sparse_order("auto", 576, num_heads=num_heads)
170
+ assert order512 == order576, order512
171
+ assert order512[0] == "FLASHMLA_SPARSE", order512
172
+
173
+
174
+def test_sm100_576_priority_unchanged():
175
+ order = _sparse_order("auto", 576, num_heads=32)
176
+ assert order[0] == "FLASHMLA_SPARSE", order
vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
+9
-1
@@ -309,8 +309,16 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
309
cu_seqlens_q = torch.arange(
310
0, q_rope.shape[0] + 1, dtype=torch.int32, device=q_rope.device
311
)
312
- k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1)
312
v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1)
313
+ if self.qk_rope_head_dim == 0:
314
+ # FA3's QV path requires the 64-wide Q/K specialization. For NoPE
315
+ # MLA, zero Q preserves QV-only attention scores while providing a
316
+ # valid TMA shape.
317
+ _FA3_QV_HEAD_DIM = 64
318
+ q_rope = q_nope.new_zeros(*q_nope.shape[:-1], _FA3_QV_HEAD_DIM)
319
+ k_cache = v_cache[..., :_FA3_QV_HEAD_DIM]
320
+ else:
321
+ k_cache = kv_rows[:, self.kv_lora_rank :].unsqueeze(1).unsqueeze(1)
322
323
out = flash_attn_varlen_func(
324
q=q_rope,
vllm/v1/attention/backends/mla/flashmla_sparse.py
+39
-3
@@ -156,8 +156,8 @@ class FlashMLASparseBackend(AttentionBackend):
156
157
@classmethod
158
def get_supported_head_sizes(cls) -> list[int]:
159
- # DeepSeek V3.2 layout: 512 NoPE + 64 RoPE = 576.
160
- return [576]
159
+ # DeepSeek V3.2: 512 NoPE + 64 RoPE = 576; GLM5Next NoPE: 512.
160
+ return [576, 512]
161
162
@classmethod
163
def is_mla(cls) -> bool:
@@ -184,6 +184,37 @@ class FlashMLASparseBackend(AttentionBackend):
184
use_mm_prefix: bool,
185
device_capability: DeviceCapability,
186
) -> str | None:
187
+ if head_size == 512:
188
+ # GLM5Next NoPE (qk_rope_head_dim == 0, kv_lora_rank == 512) has
189
+ # head_size 512 and is served here by the direct 512-wide bf16
190
+ # cache on SM90. Quantized DS-MLA caches need the zero-padded
191
+ # 576/656B envelope, which is wired up separately; plain fp8,
192
+ # SM100 bf16, and rope-carrying 512 models must fall through to
193
+ # FlashInfer/TRITON.
194
+ if (
195
+ kv_cache_dtype in (None, "auto", "bfloat16", "float16")
196
+ and device_capability.major == 9
197
+ ):
198
+ # Direct bf16 NoPE-512 is only correct for rope-free models.
199
+ # Precedent for reading hf_text_config in supports_combination:
200
+ # flashinfer_mla_sparse.py.
201
+ from vllm.config import get_current_vllm_config
202
+
203
+ vllm_config = get_current_vllm_config()
204
+ if vllm_config.model_config is not None:
205
+ hf_text_config = vllm_config.model_config.hf_text_config
206
+ if getattr(hf_text_config, "qk_rope_head_dim", 64) != 0:
207
+ return (
208
+ "FLASHMLA_SPARSE supports head_size 512 only for "
209
+ "rope-free (NoPE) models"
210
+ )
211
+ else:
212
+ return (
213
+ "FLASHMLA_SPARSE supports head_size 512 only with bf16 "
214
+ "kv-cache on SM90 (NoPE), got "
215
+ f"kv_cache_dtype={kv_cache_dtype}, "
216
+ f"capability={device_capability}"
217
+ )
218
if kv_cache_dtype == "nvfp4_ds_mla" and device_capability.major != 10:
219
return (
220
f"FLASHMLA_SPARSE only supports the {kv_cache_dtype} kv-cache "
@@ -1332,7 +1363,12 @@ class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
1363
*self.workspace_specs
1364
)
1365
q = q_concat_buffer[: ql_nope.shape[0]]
1335
- ops.concat_mla_q(ql_nope, q_pe, q)
1366
+ if q_pe.size(-1) == 0:
1367
+ # NoPE (GLM5Next): concat_mla_q requires rope_dim == 64,
1368
+ # copy directly into the head-padded buffer instead.
1369
+ q[:, : ql_nope.shape[1]].copy_(ql_nope)
1370
+ else:
1371
+ ops.concat_mla_q(ql_nope, q_pe, q)
1372
else:
1373
actual_num_heads = q.shape[1]
1374