[Bugfix][MLA] Do not use Dense MHA for GLM-5.2 (#52512)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
Woosuk Kwon committed
Aug 18, 2026 at 11:20 UTC
01e56caaf2b2d6f62be5fd78e0a7733fc9c9ed5f
3 files changed
+24
-4
tests/model_executor/layers/test_mla_short_prefill_indexer.py
+16
-2
@@ -8,6 +8,7 @@ import torch
8
9
import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer
10
from vllm.config import CUDAGraphMode
11
+from vllm.models.deepseek_v32.attention import DeepseekV32Attention
12
from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
13
14
INDEXER_LAYER = "model.layers.0.self_attn.indexer.k_cache"
@@ -45,7 +46,15 @@ def make_mla_metadata(*, use_dense_mha: bool = True, num_decode_tokens: int = 0)
46
47
@pytest.mark.parametrize(
48
"batch_kind",
48
- ["short", "threshold_mismatch", "force_mqa", "mla_decode", "capture", "full"],
49
+ [
50
+ "short",
51
+ "threshold_mismatch",
52
+ "mqa_only_layer",
53
+ "force_mqa",
54
+ "mla_decode",
55
+ "capture",
56
+ "full",
57
+ ],
58
)
59
def test_short_prefill_updates_k_cache_before_scoring_decision(
60
monkeypatch: pytest.MonkeyPatch,
@@ -132,6 +141,11 @@ def test_short_prefill_updates_k_cache_before_scoring_decision(
141
topk_indices = torch.full((7, 2048), 17, dtype=torch.int32)
142
143
def run_indexer():
144
+ if batch_kind == "mqa_only_layer":
145
+ assert not DeepseekV32Attention.supports_dense_mha_prefill
146
+ dense_mha_layer = ""
147
+ else:
148
+ dense_mha_layer = MLA_LAYER
149
return sparse_indexer.sparse_attn_indexer(
150
hidden_states,
151
INDEXER_LAYER,
@@ -149,7 +163,7 @@ def test_short_prefill_updates_k_cache_before_scoring_decision(
163
topk_indices,
164
False,
165
False,
152
- MLA_LAYER,
166
+ dense_mha_layer,
167
)
168
169
if should_skip:
vllm/model_executor/layers/attention/mla_attention.py
+6
-2
@@ -390,6 +390,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
390
3. Return the output tensor.
391
"""
392
393
+ supports_dense_mha_prefill: ClassVar[bool] = True
394
+
395
def __init__(
396
self,
397
num_heads: int,
@@ -549,9 +551,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):
551
compilation_config.static_forward_context[prefix] = self
552
553
self.prefill_backend: MLAPrefillBackend | None
552
- if self.impl.is_sparse and not self.impl.supports_dense_mha_prefill:
554
+ if self.impl.is_sparse and not (
555
+ self.impl.supports_dense_mha_prefill and self.supports_dense_mha_prefill
556
+ ):
557
logger.warning_once(
554
- "Sparse MLA impl has no dense-MHA prefill path; using the top-k "
558
+ "Sparse MLA layer has no dense-MHA prefill path; using the top-k "
559
"MQA path only."
560
)
561
self.prefill_backend = None
vllm/models/deepseek_v32/attention.py
+2
@@ -160,6 +160,7 @@ class DeepseekV32Attention(MLAAttention):
160
indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer
161
162
require_fp8_kv_cache: bool = True
163
+ supports_dense_mha_prefill = False
164
165
def __init__(
166
self,
@@ -265,6 +266,7 @@ class DeepseekV32Attention(MLAAttention):
266
and not skip_topk
267
and not self.use_pcp
268
and current_platform.is_cuda()
269
+ and self.supports_dense_mha_prefill
270
)
271
self._dense_mha_metadata_layer_name = (
272
self.layer_name if enable_short_prefill_scoring_skip else ""