[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>

tests/model_executor/layers/test_mla_short_prefill_indexer.py

@@ -8,6 +8,7 @@ import torch
 
 import vllm.model_executor.layers.sparse_attn_indexer as sparse_indexer
 from vllm.config import CUDAGraphMode
+from vllm.models.deepseek_v32.attention import DeepseekV32Attention
 from vllm.v1.attention.backends.mla.indexer import DeepseekV32IndexerMetadata
 
 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)
 
 @pytest.mark.parametrize(
     "batch_kind",
-    ["short", "threshold_mismatch", "force_mqa", "mla_decode", "capture", "full"],
+    [
+        "short",
+        "threshold_mismatch",
+        "mqa_only_layer",
+        "force_mqa",
+        "mla_decode",
+        "capture",
+        "full",
+    ],
 )
 def test_short_prefill_updates_k_cache_before_scoring_decision(
     monkeypatch: pytest.MonkeyPatch,
@@ -132,6 +141,11 @@ def test_short_prefill_updates_k_cache_before_scoring_decision(
     topk_indices = torch.full((7, 2048), 17, dtype=torch.int32)
 
     def run_indexer():
+        if batch_kind == "mqa_only_layer":
+            assert not DeepseekV32Attention.supports_dense_mha_prefill
+            dense_mha_layer = ""
+        else:
+            dense_mha_layer = MLA_LAYER
         return sparse_indexer.sparse_attn_indexer(
             hidden_states,
             INDEXER_LAYER,
@@ -149,7 +163,7 @@ def test_short_prefill_updates_k_cache_before_scoring_decision(
             topk_indices,
             False,
             False,
-            MLA_LAYER,
+            dense_mha_layer,
         )
 
     if should_skip:

vllm/model_executor/layers/attention/mla_attention.py

@@ -390,6 +390,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
     3. Return the output tensor.
     """
 
+    supports_dense_mha_prefill: ClassVar[bool] = True
+
     def __init__(
         self,
         num_heads: int,
@@ -549,9 +551,11 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         compilation_config.static_forward_context[prefix] = self
 
         self.prefill_backend: MLAPrefillBackend | None
-        if self.impl.is_sparse and not self.impl.supports_dense_mha_prefill:
+        if self.impl.is_sparse and not (
+            self.impl.supports_dense_mha_prefill and self.supports_dense_mha_prefill
+        ):
             logger.warning_once(
-                "Sparse MLA impl has no dense-MHA prefill path; using the top-k "
+                "Sparse MLA layer has no dense-MHA prefill path; using the top-k "
                 "MQA path only."
             )
             self.prefill_backend = None

vllm/models/deepseek_v32/attention.py

@@ -160,6 +160,7 @@ class DeepseekV32Attention(MLAAttention):
     indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer
 
     require_fp8_kv_cache: bool = True
+    supports_dense_mha_prefill = False
 
     def __init__(
         self,
@@ -265,6 +266,7 @@ class DeepseekV32Attention(MLAAttention):
             and not skip_topk
             and not self.use_pcp
             and current_platform.is_cuda()
+            and self.supports_dense_mha_prefill
         )
         self._dense_mha_metadata_layer_name = (
             self.layer_name if enable_short_prefill_scoring_skip else ""