执行摘要

• 一句话：修复 GLM-5.2 短 prefill 误走 dense MHA 导致输出退化
• 推荐动作：值得精读。这是一个典型的低改动、高影响正确性修复：改动仅 24 行，但涉及 sparse MLA 的 prefill 路由契约。建议关注两个设计点：一是用类级能力声明把后端能力与模型层能力解耦，避免 wrapper 未实现却由 indexer 错误宣告 dense 路径；二是在补实现与关能力之间选择后者，回归既有 MQA 路径，是控制复杂度的好例子。若你在维护其他 MLA 模型或 sparse attention 后端，应检查自己的模型是否也需要显式声明 supports_dense_mha_prefill。

功能与动机

PR body 明确指出核心动机：Fix incorrect short-prefill dispatch in the NVIDIA DeepSeek-V3.2 attention wrapper. The generic sparse-MLA metadata path may select dense MHA when a prefill's sequence length is at most index_topk. It then allows the sparse indexer to skip top-k scoring. DeepseekV32Attention, however, only invokes forward_mqa; it does not execute the dense-MHA backend。也就是说后端能力与模型层实际执行路径不一致，导致 GLM-5.2 输出退化为重复 token。作者在评论中补充设计取舍：This PR essentially reverts #51298. I chose to just disable dense MHA instead of adding support for it, since the complexity outweighs the performance gain.

实现拆解

实现拆解

1. 基类能力声明与判定收紧（vllm/model_executor/layers/attention/mla_attention.py）
   在 MLAAttention 上新增类级常量 supports_dense_mha_prefill = True，并把 __init__ 中 prefill backend 的实例化条件从仅看后端 self.impl.supports_dense_mha_prefill 收紧为后端与模型层同时支持：当 self.impl.is_sparse 且任一方不支持时，self.prefill_backend 置为 None，sparse MLA 的 prefill 全部回退到 top-k MQA 路径。默认值 True 保证未显式声明的模型行为不变。

2. 模型层显式禁用（vllm/models/deepseek_v32/attention.py）
   DeepseekV32Attention 声明 supports_dense_mha_prefill = False，因为其前向只调用 forward_mqa、没有 dense-MHA 分支。同时在 enable_short_prefill_scoring_skip 的推导中并入该标志，使 _dense_mha_metadata_layer_name 保持空字符串，indexer 无法获知可跳过的 dense 层，从而始终执行 top-k scoring 并回填 top-k 缓冲。

3. 测试配套（tests/model_executor/layers/test_mla_short_prefill_indexer.py）
   参数化列表新增 mqa_only_layer 分支：断言 DeepseekV32Attention.supports_dense_mha_prefill 为 False，并以空 dense_mha_layer 调用 sparse_attn_indexer，验证其进入 scoring（触发 ScoringReached）且 top-k 缓冲被重置为全 -1；原有 short 与 threshold_mismatch 分支继续验证非 MQA-only 模型的 skip 行为。focused pytest 结果 7 passed，pre-commit 全部通过。

4. 真实模型验证（手动，未入 CI）
   PR body 报告在 nvidia/GLM-5.2-NVFP4、TP=4、FP8 E4M3 KV cache、eager 执行下，串行探测 97～1,001 token 共 19 个边界长度全部通过精确检索；对照实验表明仅添加 sparse_mla_force_mqa=true 也能让全部用例通过，因果确认了路由不一致是根因。

关键文件：

• vllm/model_executor/layers/attention/mla_attention.py（模块 注意力层；类别 source；类型 core-logic；符号 MLAAttention, MLAAttention.supports_dense_mha_prefill, MLAAttention.init）: 核心修复文件：在 MLAAttention 基类新增模型层能力声明 supports_dense_mha_prefill，并将 prefill backend 的实例化条件收紧为后端与模型层同时支持，是本次路由修复的公共契约改动。
• vllm/models/deepseek_v32/attention.py（模块 模型适配层；类别 source；类型 data-contract；符号 DeepseekV32Attention, DeepseekV32Attention.supports_dense_mha_prefill, DeepseekV32Attention.init）: 模型层修复：DeepseekV32Attention 显式声明 MQA-only，并在 enable_short_prefill_scoring_skip 中纳入该标志，阻止 indexer 在短 prefill 时跳过 top-k scoring。
• tests/model_executor/layers/test_mla_short_prefill_indexer.py（模块 索引器测试；类别 test；类型 test-coverage；符号 test_short_prefill_updates_k_cache_before_scoring_decision）: 测试配套：新增 mqa_only_layer 分支验证 MQA-only 模型在短 prefill 时必须进入 indexer scoring，防止能力声明被误用。

关键符号：MLAAttention.init, MLAAttention.supports_dense_mha_prefill, DeepseekV32Attention.init, DeepseekV32Attention.supports_dense_mha_prefill, test_short_prefill_updates_k_cache_before_scoring_decision

关键源码片段

vllm/model_executor/layers/attention/mla_attention.py

核心修复文件：在 MLAAttention 基类新增模型层能力声明 supports_dense_mha_prefill，并将 prefill backend 的实例化条件收紧为后端与模型层同时支持，是本次路由修复的公共契约改动。

class MLAAttention(nn.Module, AttentionLayerBase):
    """Multi-Head Latent Attention layer.
​
    NOTE: Please read the comment at the top of the file before trying to
    understand this class.
    """
​
    # 模型层是否具备 dense-MHA prefill 执行能力。该声明与后端实现
    # self.impl.supports_dense_mha_prefill 是两个独立的 capability，
    # 只有两者同时为 True，sparse MLA 才会启用 dense-MHA prefill 路径，
    # 并允许 indexer 跳过 top-k scoring。
    supports_dense_mha_prefill: ClassVar[bool] = True
​
    def __init__(self, ...):
        ...
        # 只有当 sparse MLA 后端与模型层都声明支持 dense-MHA prefill 时，
        # 才实例化 prefill backend；否则强制回退到 top-k MQA 路径。
        self.prefill_backend: MLAPrefillBackend | None
        if self.impl.is_sparse and not (
            self.impl.supports_dense_mha_prefill and self.supports_dense_mha_prefill
        ):
            logger.warning_once(
                "Sparse MLA layer has no dense-MHA prefill path; using the top-k "
                "MQA path only."
            )
            self.prefill_backend = None
        else:
            # 正常实例化后端指定的 prefill backend。
            ...

vllm/models/deepseek_v32/attention.py

模型层修复：DeepseekV32Attention 显式声明 MQA-only，并在 enable_short_prefill_scoring_skip 中纳入该标志，阻止 indexer 在短 prefill 时跳过 top-k scoring。

class DeepseekV32Attention(MLAAttention):
    indexer: "DeepseekV32Indexer | None"
    indexer_cls: "type[DeepseekV32Indexer]" = DeepseekV32Indexer
​
    require_fp8_kv_cache: bool = True
    # NVIDIA DeepSeek-V3.2 wrapper 只实现 forward_mqa，没有 dense-MHA 分支，
    # 因此显式声明不支持 dense-MHA prefill，让 prefill 保持走 top-k MQA。
    supports_dense_mha_prefill = False
​
    def __init__(self, vllm_config, config, prefix, ...):
        ...
        # 只有模型层支持 dense-MHA prefill 时，才允许 indexer 在短 prefill
        # 场景跳过 top-k scoring。对 MQA-only 模型必须保持空字符串，
        # 让 indexer 始终执行 top-k scoring，避免 MQA 读取全 -1 的 top-k 缓冲。
        # （is_mtp_layer、skip_topk、use_pcp 来自外层上下文变量。）
        self.skip_topk = False
        enable_short_prefill_scoring_skip = (
            not is_mtp_layer
            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 ""
        )

评论区精华

本次 review 没有形成实质讨论线程（仅有 bot 提示），唯一有价值的设计说明来自作者 WoosukKwon 的 issue 评论：

This PR essentially reverts #51298. I chose to just disable dense MHA instead of adding support for it, since the complexity outweighs the performance gain.

这里传达了两个决策信息：一是本 PR 与 #51298 是方向性回退关系；二是作者在补全 dense-MHA 分支与关闭该能力之间选择了后者，理由是维护复杂度大于性能收益。配合 PR body 中的因果对照实验，这是一个相当完整的正确性修复论证。

• 为什么禁用 dense MHA 而不是补实现？ (design): 接受 disable 方案：DeepseekV32Attention 显式声明 MQA-only，sparse MLA 的 prefill 全部保持走 top-k MQA，不新增 dense-MHA 分支。

风险与影响

• 风险：技术风险
• 性能回退：GLM-5.2 的短 prefill（长度 ≤ index_topk）从可能的 dense-MHA 路径回到 top-k MQA，prefill 计算量上升；PR body 未报告性能数据，属于已知取舍。
• 能力契约需双端同步：impl.supports_dense_mha_prefill 与模型层 supports_dense_mha_prefill 是两套独立标志，新增 MLA 模型或改造 wrapper 时若遗漏声明，可能重演同类路由错误。
• 公共路径改动：MLAAttention.__init__ 是 MLA 家族共享逻辑，不过默认 True 保持旧行为，风险集中在显式声明 False 的子类。
• 测试盲区：单测只覆盖一个 MQA-only 样本（DeepseekV32Attention），真实模型的 19 个边界用例是人工验证，未纳入 CI，后续改动可能回归。
• 影响：影响范围
• 用户侧：修复 GLM-5.2 在特定 prefill 长度下输出退化为重复 token 的严重正确性问题，且无需 sparse_mla_force_mqa=true 变通参数。
• 系统侧：sparse MLA 的 prefill 路由能力从后端单一判定升级为后端 + 模型层双重判定，为后续平台差异化（SM120、PCP、DSA）留下统一钩子。
• 团队侧：改动量小、语义清晰，便于 review 与回滚；但需要维护者长期维护双 flag 的一致性。
• 风险标记：核心路径变更, 能力契约需双端同步, 真实模型验证未入 CI, 潜在性能回退

关联脉络

• PR #51298 （标题未提供，讨论中提及）: 作者评论指出本 PR 实质回退该 PR 引入的模型层 dense-MHA prefill 能力，改为显式禁用。
• PR #51395 （标题未提供，PR body 提及）: 该 PR 只对 SM120 sparse-MLA 后端禁用 dense prefill，与本 PR 的模型层能力禁用互补。
• PR #52046 （标题未提供，PR body 提及）: 该 PR 为 PCP 路径增加 dense-MHA 分支，而本 PR 修复的是非 PCP SM100 路径。
• PR #49790 （标题未提供，PR body 提及）: 该 PR 将 DSA 架构路由到 NVIDIA 实现，但未解决本次 dispatch 不一致。