执行摘要

• 一句话：为稀疏 MLA 短序列新增密集 MHA 预填充路径
• 推荐动作：建议阅读 mla_attention.py 中的路由逻辑和 sparse_mla_attention.py 中的共享实现，理解如何在 vLLM 中为注意力后端添加条件优化路径。特别关注 use_mha 的决策条件以及 SparseMLACommonMetadataBuilder 的 workspace 分配策略。由于本 PR 是系列中的第 1 部分，关注后续 PR 中如何扩展此机制至 masked MHA。

功能与动机

来自 PR 描述：'Adds a fast path when all prefills in a batch have seq_len <= topk, routing to dense MHA instead of sparse MQA (subject to a q_len threshold) to obtain a substantial speedup.' 同时这是实现 masked MHA 的前置步骤 (参见 https://github.com/vllm-project/vllm/pull/34744)。

实现拆解

1. 提取共享组件：在 sparse_mla_attention.py 中新增 SparseMLACommonMetadataBuilder 和 SparseMLACommonImpl，将原来分散在各个稀疏后端中的元数据构建和 forward_mha 逻辑统一管理，包括 DCP 感知的 workspace 分配。

2. 修改路由决策：在 mla_attention.py 的 MLACommonImpl.__init__ 中，构建 prefill backend 时增加异常处理，当无可用 prefill backend 时仅使用 sparse MQA；在 forward_impl 中新增 use_mha 判断条件：is_sparse、prefill_backend is not None、prefill_max_seq_len <= topk_tokens 且未强制 sparse_mla_force_mqa。

3. 重构稀疏后端：FlashInfer、FlashMLA、FlashAttn 等稀疏 MLA 后端迁移为继承自 SparseMLACommonMetadataBuilder 和 SparseMLACommonImpl，移除重复代码（约 80-90 行/后端），新增 forward_mha 实现（通常直接调用父类 SparseMLACommonImpl.forward_mha）。

4. 简化基类：删除 backend.py 中的 SparseMLAAttentionImpl 抽象类，将 is_sparse 标志整合到 MLAAttentionImpl 中，减少继承层次。

5. 配套测试和基准：新增 test_sparse_mla_backends.py 中的 test_sparse_backend_prefill_correctness 测试（覆盖 short dense MHA 和 short context dense MHA 场景），修改基准测试框架以支持 dense MHA 基准，新增配置文件 mla_sparse_mha_vs_mqa.yaml。

关键文件：

• vllm/model_executor/layers/attention/sparse_mla_attention.py（模块 模型执行器；类别 source；类型 core-logic；符号 SparseMLACommonMetadataBuilder, init, determine_chunked_prefill_workspace_size, _build_req_id_per_token）: 新增的核心文件，定义了共享的元数据构建器和 forward_mha 实现，是所有稀疏 MLA 后端的基础。
• vllm/model_executor/layers/attention/mla_attention.py（模块 模型执行器；类别 source；类型 data-contract；符号 build_mla_chunked_context_metadata, MLACommonImpl, MLACommonBaseImpl, fused_output_quant_supported）: 核心路由逻辑所在文件，修改了初始化（prefill backend 降级）和 forward_impl（use_mha 决策）
• vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py（模块 V1 注意力；类别 source；类型 refactor；符号 get_impl_cls, build, FlashInferMLASparseImpl）: FlashInfer 稀疏 MLA 后端，进行了重大重构以继承新基类并启用密集 MHA。
• vllm/v1/attention/backend.py（模块 V1 注意力；类别 source；类型 refactor；符号 SparseMLAAttentionImpl, MLAAttentionImpl, AttentionImplBase）: 删除了 SparseMLAAttentionImpl 抽象类，减少了继承层次，将 is_sparse 标志整合到 MLAAttentionImpl。
• tests/v1/attention/test_sparse_mla_backends.py（模块 测试；类别 test；类型 test-coverage；符号 test_sparse_backend_prefill_correctness）: 新增了针对密集 MHA 路径的测试用例，验证正确性。

关键符号：SparseMLACommonMetadataBuilder.init, SparseMLACommonMetadataBuilder.determine_chunked_prefill_workspace_size, SparseMLACommonImpl.forward_mha, MLACommonImpl.init, MLACommonImpl.forward_impl, FlashInferMLASparseImpl.init

关键源码片段

vllm/model_executor/layers/attention/sparse_mla_attention.py

新增的核心文件，定义了共享的元数据构建器和 forward_mha 实现，是所有稀疏 MLA 后端的基础。

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
​
class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
    """稀疏 MLA 公共元数据构建器 —— 所有稀疏 MLA 后端共享此基类。"""
​
    @staticmethod
    def determine_chunked_prefill_workspace_size(
        vllm_config: "VllmConfig",
    ) -> int:
        """根据配置确定 chunked-prefill workspace 的行数。
​
        该 workspace 用于在 DCP 模式下收集所有 KV shard，以便运行密集 MHA。
        返回值为所有 rank 中最大的 sequence 长度（受 topk 限制），
        确保能容纳完整的密集 KV。
        """
        scheduler_config = vllm_config.scheduler_config
        cache_config = vllm_config.cache_config
        model_config = vllm_config.model_config
        topk_tokens = model_config.hf_config.index_topk
​
        workspace_size = min(
            max(
                8 * model_config.max_model_len, # 保守上限
                4 * scheduler_config.max_num_seqs * cache_config.block_size,
            ),
            64 * 1024, # 绝对上限，防止过度分配
            scheduler_config.max_num_seqs * topk_tokens,
        )
        return max(
            workspace_size,
            scheduler_config.max_num_seqs * cache_config.block_size,
        )

vllm/model_executor/layers/attention/mla_attention.py

核心路由逻辑所在文件，修改了初始化（prefill backend 降级）和 forward_impl（use_mha 决策）

# 在 MLACommonImpl.__init__ 方法中
# 构建 prefill backend；稀疏 MLA 在没有对应 backend 时允许降级到纯 MQA 模式
self.prefill_backend: MLAPrefillBackend | None
try:
    prefill_backend_cls = get_mla_prefill_backend(vllm_config)
except ValueError:
    # 如果无法获取预填充后端，仅当当前 impl 是稀疏且用户未显式指定 prefill backend 时才降级
    if (
        not self.impl.is_sparse
        or vllm_config.attention_config.mla_prefill_backend is not None
    ):
        raise # 非稀疏或显式指定时失败
    logger.warning_once(
        "No MLA prefill backend supports this model; "
        "sparse MLA will use the top-k MQA path only "
        "(no dense-MHA prefill)."
    )
    self.prefill_backend = None
else:
    self.prefill_backend = prefill_backend_cls(
        num_heads=self.num_heads,
        scale=self.scale,
        kv_lora_rank=self.kv_lora_rank,
        qk_nope_head_dim=self.qk_nope_head_dim,
        qk_rope_head_dim=self.qk_rope_head_dim,
        v_head_dim=self.v_head_dim,
        vllm_config=vllm_config,
    )

评论区精华

LucasWilkinson 担心该变更可能破坏稀疏 DCP 支持（关联 PR #46076），MatthewBonanni 确认在 commit 9c09be2e 中添加了 DCP 支持并通过了验证。

LucasWilkinson 对测试覆盖提出疑问，建议将新增的 PREFILL_BATCH_SPECS 合并到现有的 SPARSE_BACKEND_BATCH_SPECS 中，MatthewBonanni 解释现有测试仅覆盖 forward_mqa，新测试专为 forward_mha 设计。

LucasWilkinson 对 backend.py 中的文档字符串变更提出 nit，认为过于详细反而造成混淆，MatthewBonanni 在后续 commit 中清理了该变更。

MatthewBonanni 在 self-review 中指出 workspace 大小计算可以更保守，但当前 PR 仅路由 max_seq_len <= topk 的情况，后续 masked MHA 实现将进一步限制。

• DCP 兼容性验证 (correctness): 已通过 commit 9c09be2e 添加 DCP 支持并测试验证。
• 测试覆盖策略 (testing): 作者澄清差异，测试保持独立。
• 文档字符串简洁性 (documentation): 已简化。
• workspace 大小保守性 (design): 当前实现是保守的，留待后续优化。

风险与影响

• 风险：
  • 性能回归风险：use_mha 路由条件依赖于 prefill_max_seq_len 的正确计算和 topk_tokens 的准确比对，若计算有误可能导致短序列错误地走稀疏 MQA 路径，失去优化机会；或长序列错误地走密集 MHA 导致性能下降。当前条件明确且受 sparse_mla_force_mqa 配置保护。
  • 稀疏 DCP 兼容性：首次实现未正确处理 DCP 配置，已在 commit 9c09be2e 中修复并验证。
  • 跨后端一致性：多个稀疏后端（flashinfer, flashmla, flashattn）需要一致的新元数据字段，若某后端未正确更新可能导致运行时错误。
  • 测试覆盖不完整：密集 MHA 路径的测试仅覆盖了两种 batch spec，缺少对混合序列长度、多 batch 及不同 topk 配置的覆盖。
• 影响：
  • 用户影响：使用 DeepSeek-V3、GLM 等稀疏 MLA 模型的用户在短序列工作负载下将获得显著吞吐提升（6.8%-24.5%），TTFT 和 TPOT 均有改善。
  • 系统影响：统一元数据构建器降低了未来添加新注意后端的成本；基类简化使注意力框架更清晰。
  • 团队影响：本 PR 是稀疏 MLA 优化系列的第一部分，后续将有 masked MHA 实现，需要密切关注此 PR 引入的接口变更。
  • 风险标记：核心注意力路径变更, 稀疏 DCP 兼容性, 多个后端重构, 短序列性能敏感

关联脉络

• PR #34744 : 本 PR 从该 PR 拆分而来，body 中有明确链接。
• PR #46076 : LucasWilkinson 要求验证不会破坏该 PR 引入的稀疏 DCP 支持。