[Perf][GLM-5.3-Flash] Dense/masked-MHA sparse prefill for the NoPE (256, 0, 256) layout + skip the NoPE K concat (#55738)
Signed-off-by: Jared Wen <jaredwen@inferact.ai> Signed-off-by: Jared Wen <w13431838023@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jared Wen committed
Sep 14, 2026 at 17:36 UTC
238cb2b191630e8b02965adc772e3236d26936fe
6 files changed
+96
tests/v1/attention/test_mla_backends.py
+20
@@ -25,6 +25,7 @@ from vllm.config.vllm import set_current_vllm_config
25
from vllm.model_executor.layers.attention import mla_attention as mla_attention_module
26
from vllm.model_executor.layers.attention.mla_attention import (
27
MLAAttention,
28
+ MLACommonBaseImpl,
29
QueryLenSupport,
30
_DecodeConcatQuantFP8,
31
_use_masked_mha,
@@ -166,6 +167,25 @@ def test_glm5_flashinfer_masked_mha_routing(
167
)
168
169
170
+@pytest.mark.parametrize("qk_rope_head_dim", [64, 0], ids=["rope", "nope"])
171
+def test_concat_k_nope_k_pe_matches_torch_cat(qk_rope_head_dim):
172
+ """The K concat used by the MLA prefill context loop must equal torch.cat of
173
+ k_nope with the broadcast k_pe; with no RoPE part it returns k_nope itself
174
+ instead of allocating and copying."""
175
+ torch.manual_seed(0)
176
+ num_tokens, num_heads, qk_nope_head_dim = 5, 4, 256
177
+ k_nope = torch.randn(num_tokens, num_heads, qk_nope_head_dim, dtype=torch.bfloat16)
178
+ k_pe = torch.randn(num_tokens, 1, qk_rope_head_dim, dtype=torch.bfloat16)
179
+ impl = SimpleNamespace(_use_flashinfer_concat_mla_k=False)
180
+
181
+ k = MLACommonBaseImpl._concat_k_nope_k_pe(impl, k_nope, k_pe)
182
+
183
+ expected = torch.cat([k_nope, k_pe.expand(-1, num_heads, -1)], dim=-1)
184
+ assert k.shape == (num_tokens, num_heads, qk_nope_head_dim + qk_rope_head_dim)
185
+ torch.testing.assert_close(k, expected, rtol=0, atol=0)
186
+ assert (k.data_ptr() == k_nope.data_ptr()) == (qk_rope_head_dim == 0)
187
+
188
+
189
def test_masked_mha_routing_is_dimension_specific():
190
assert _use_masked_mha(
191
backend_name="FLASHMLA_SPARSE",
tests/v1/attention/test_mla_prefill_selector.py
+34
@@ -313,6 +313,40 @@ class TestBackendValidation:
313
)
314
assert invalid_reasons == []
315
316
+ def test_flash_attn_accepts_glm53_flash_nope_dimensions(self):
317
+ """(256, 0, 256) runs the same FA kernels as GLM-5's (192, 64, 256); a
318
+ RoPE-carrying 320-wide query does not."""
319
+ try:
320
+ from vllm.v1.attention.backends.mla.prefill.flash_attn import (
321
+ FlashAttnPrefillBackend,
322
+ )
323
+ except ImportError:
324
+ pytest.skip("MLA prefill backend not available")
325
+ return
326
+
327
+ capability = DeviceCapability(major=10, minor=0)
328
+
329
+ def validate(qk_rope_head_dim: int) -> list[str]:
330
+ selector_config = MLAPrefillSelectorConfig(
331
+ dtype=torch.bfloat16,
332
+ mla_dimensions=MLADimensions(
333
+ qk_nope_head_dim=256,
334
+ qk_rope_head_dim=qk_rope_head_dim,
335
+ v_head_dim=256,
336
+ ),
337
+ )
338
+ with patch.object(
339
+ FlashAttnPrefillBackend, "is_available", return_value=True
340
+ ):
341
+ return FlashAttnPrefillBackend.validate_configuration(
342
+ capability, selector_config
343
+ )
344
+
345
+ assert validate(qk_rope_head_dim=0) == []
346
+ invalid_reasons = validate(qk_rope_head_dim=64)
347
+ assert len(invalid_reasons) == 1
348
+ assert "supported MLA dimensions" in invalid_reasons[0]
349
+
350
351
@pytest.mark.skipif(
352
not current_platform.is_cuda_alike(),
tests/v1/attention/test_sparse_mla_backends.py
+29
@@ -32,6 +32,7 @@ from vllm.model_executor.layers.attention.sparse_mla_attention import (
32
GLOBAL_TOPK_MASK_MAX_BYTES,
33
SparseMLACommonImpl,
34
SparseMLAPrefillMetadata,
35
+ _is_masked_mha_available,
36
_masked_mha_workspace_fits,
37
_topk_mask_shape,
38
_use_dense_mha_prefill,
@@ -1165,6 +1166,33 @@ def test_masked_mha_workspace_guards_long_routing_policy(
1166
)
1167
1168
1169
+@pytest.mark.parametrize(
1170
+ ("model_dims", "kv_cache_dtype", "fa_version", "expected"),
1171
+ [
1172
+ pytest.param((128, 512, 128, 64, 128), "auto", 4, True, id="deepseek_v32"),
1173
+ pytest.param((64, 512, 192, 64, 256), "auto", 4, True, id="glm5"),
1174
+ pytest.param((64, 512, 256, 0, 256), "auto", 4, True, id="glm53_flash_nope"),
1175
+ pytest.param((64, 512, 256, 0, 256), "fp8", 4, False, id="quantized_kv"),
1176
+ pytest.param((64, 512, 256, 0, 256), "auto", 3, False, id="no_fa4"),
1177
+ pytest.param((64, 512, 256, 64, 256), "auto", 4, False, id="rope_320"),
1178
+ pytest.param((128, 512, 256, 0, 256), "auto", 4, False, id="wrong_heads"),
1179
+ ],
1180
+)
1181
+def test_is_masked_mha_available_model_dims(
1182
+ monkeypatch, model_dims, kv_cache_dtype, fa_version, expected
1183
+):
1184
+ """The allow-list gates masked MHA per exact model geometry: the DeepSeek-V3.2,
1185
+ GLM-5 and NoPE GLM-5.3-Flash layouts on an SM100-family GPU with FA4 and an
1186
+ unquantized KV cache, nothing else."""
1187
+ import vllm.model_executor.layers.attention.sparse_mla_attention as mod
1188
+
1189
+ monkeypatch.setattr(
1190
+ mod.current_platform, "is_device_capability_family", lambda family: True
1191
+ )
1192
+ monkeypatch.setattr(mod, "get_flash_attn_version", lambda **kwargs: fa_version)
1193
+ assert _is_masked_mha_available(*model_dims, kv_cache_dtype) is expected
1194
+
1195
+
1196
def test_masked_mha_workspace_fits_accounts_for_batch_and_context():
1197
"""Request count and context chunk length are independent multipliers."""
1198
base = dict(batch_size=2, max_query_len=2048, max_context_chunk_seq_len=2048)
@@ -1198,6 +1226,7 @@ PREFILL_BATCH_SPECS = {
1226
[
1227
pytest.param(128, 128, 64, 128, id="deepseek_hd192_v128"),
1228
pytest.param(64, 192, 64, 256, id="glm5_hd256_v256"),
1229
+ pytest.param(64, 256, 0, 256, id="glm53_flash_nope_hd256_v256"),
1230
],
1231
)
1232
def test_sparse_backend_prefill_correctness(
vllm/model_executor/layers/attention/mla_attention.py
+4
@@ -2871,6 +2871,10 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2871
Returns:
2872
Tensor of shape [..., nope_dim + pe_dim]
2873
"""
2874
+ if k_pe.shape[-1] == 0:
2875
+ # NoPE MLA: nothing to append, so no copy either.
2876
+ return k_nope
2877
+
2878
k = torch.empty(
2879
(*k_nope.shape[:-1], k_nope.shape[-1] + k_pe.shape[-1]),
2880
dtype=k_nope.dtype,
vllm/model_executor/layers/attention/sparse_mla_attention.py
+2
@@ -111,6 +111,8 @@ def _is_masked_mha_available(
111
if model_dims not in (
112
(128, 512, 128, 64, 128),
113
(64, 512, 192, 64, 256),
114
+ # GLM-5.3-Flash: NoPE, qk_head_dim 256 == the (192, 64, 256) kernel.
115
+ (64, 512, 256, 0, 256),
116
):
117
return False
118
qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+7
@@ -327,6 +327,13 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
327
qk_rope_head_dim=64,
328
v_head_dim=128,
329
),
330
+ # GLM5Next NoPE layout: qk_head_dim 256 + 0 and v_head_dim 256 run
331
+ # the same kernels as the (192, 64, 256) DeepSeek-V3.2 layout.
332
+ MLADimensions(
333
+ qk_nope_head_dim=256,
334
+ qk_rope_head_dim=0,
335
+ v_head_dim=256,
336
+ ),
337
]
338
339
def __init__(