执行摘要

• 一句话：移除 DSv4 sparse MLA 的 q-head 填充，高 TP 下省计算
• 推荐动作：值得精读。虽然改动只有 1 个文件 35 行，但它展示了如何与上游内核版本演进同步、通过共享辅助函数统一两个注意力类的填充逻辑，并明确标注了对 flashinfer 版本的强依赖。对于维护 DeepSeek 模型或关注 FlashInfer 集成性能的工程师有参考价值。

功能与动机

上游 flashinfer-ai/flashinfer#3545 将 trtllm_batch_decode_sparse_mla_dsv4 的 guard 从 {64,128} 放宽为 {8,16,32,64,128}（SM100 与 SM120 两条 decode 路径均支持）。此前 vLLM 为兼容旧内核，将每个 rank 的 query head 数向上填充到内核支持值，SM100 上 TP4/TP8/TP16 的填充后 head count 恒为 64（32→64、16→64、8→64），白白浪费计算。因此本 PR 在 flashinfer>=0.6.14 前提下移除 padding，仅保留对 sub-8 head count 的兜底向上取整到 8。

实现拆解

1. 新增共享填充辅助函数：在 vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py 顶部定义模块级常量 _SPARSE_MLA_SUPPORTED_Q_HEADS = (8, 16, 32, 64, 128) 和辅助函数 _pad_to_supported_q_heads(num_heads)。该函数按升序遍历支持集合，返回第一个不小于 num_heads 的值；若超出 128 则抛出 ValueError，错误信息统一描述为 h_q in {8, 16, 32, 64, 128}。
2. 改造 SM100 注意力类的填充逻辑：DeepseekV4FlashInferMLAAttention.get_padded_num_q_heads 原来对 num_heads > 128 抛错、否则返回 64 if num_heads <= 64 else 128，现替换为一行委托 return _pad_to_supported_q_heads(num_heads)，使 TP1/2/4/8/16 对应的 128/64/32/16/8 全部走原始值。
3. 改造 SM120 注意力类的填充逻辑：DeepseekV4FlashInferSM120Attention.get_padded_num_q_heads 原来分四段填充到 16/32/64/128，同样替换为委托共享 helper，行为对齐到 SM100，并支持 8 作为最小合法值。
4. 配套测试与验证：本 PR 未新增自动化测试文件，但作者在 body 中报告了在 flashinfer==0.6.14 下对内核两个 validator（SM100 seq_lens/cum_seq_lens_q 与 SM120 swa_topk_lens/extra_sparse_indices）的验证结果，以及 _pad_to_supported_q_heads 映射的单元级检查（TP1→128、TP2→64、TP4→32、TP8→16、TP16→8、num_heads=4 → 8、num_heads=256 → ValueError）。Lint/type 检查通过 pre-commit run --files ... 完成。

关键文件：

• vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py（模块 注意力层；类别 source；类型 data-contract；符号 _pad_to_supported_q_heads, get_padded_num_q_heads）: 唯一变更文件，集中实现了 sparse MLA q-head padding 的移除与共享化。

关键符号：_pad_to_supported_q_heads, DeepseekV4FlashInferMLAAttention.get_padded_num_q_heads, DeepseekV4FlashInferSM120Attention.get_padded_num_q_heads

关键源码片段

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

唯一变更文件，集中实现了 sparse MLA q-head padding 的移除与共享化。

# Sparse MLA h_q 支持的原生集合，flashinfer>=0.6.14 后 SM100/SM120 双路径均接受。
_SPARSE_MLA_SUPPORTED_Q_HEADS = (8, 16, 32, 64, 128)
​
​
def _pad_to_supported_q_heads(num_heads: int) -> int:
    """返回能容纳 num_heads 的最小支持 h_q。
​
    TP1/2/4/8/16 分别对应 128/64/32/16/8，全部无需填充；
    仅 sub-8 的 per-rank 计数（如 TP > 16 时）仍向上取整到 8。
    """
    for supported in _SPARSE_MLA_SUPPORTED_Q_HEADS:
        if num_heads <= supported:
            return supported
    raise ValueError(
        f"DeepseekV4 FlashInfer MLA Sparse does not support {num_heads} heads "
        "(sparse MLA kernel requires h_q in {8, 16, 32, 64, 128})."
    )
​
​
class DeepseekV4FlashInferMLAAttention(DeepseekV4Attention):
    """FlashInfer TRTLLM-gen sparse MLA 注意力层（SM100 DeepSeek V4）。"""
​
    @classmethod
    def get_padded_num_q_heads(cls, num_heads: int) -> int:
        # 委托共享 helper，TP1/2/4/8/16 不再被填充到 64。
        return _pad_to_supported_q_heads(num_heads)
​
​
class DeepseekV4FlashInferSM120Attention(DeepseekV4Attention):
    """FlashInfer sparse MLA 注意力层（SM120 DeepSeek V4）。"""
​
    @classmethod
    def get_padded_num_q_heads(cls, num_heads: int) -> int:
        # SM120 同样委托共享 helper，行为与 SM100 对齐。
        return _pad_to_supported_q_heads(num_heads)

评论区精华

唯一人工 review 来自 yewentao256（APPROVED），指出参照 #48537 可预期 6%+ 的 TTFT 提升；claude[bot] 自动评论说明 fork PR 禁用自动 review。PR body 中详细讨论了与 #47669 的依赖关系及非重复性说明。

• TTFT 性能提升预期 (performance): 合入者认可该变更的性能价值，将其与 #48537 的性能收益关联。

风险与影响

• 风险：主要风险是版本耦合：仓库当前 pin 的 flashinfer-python==0.6.13 内核仍强制 num_heads ∈ {64,128}，若本 PR 单独先行合入，SM100 上 TP4/TP8/TP16 会立即触发 ValueError。其次，_pad_to_supported_q_heads 对 num_heads=256 等超限值抛错的行为与原 SM100 的 >128 抛错、SM120 的 >128 抛错一致，但错误信息统一为集合描述，可能影响依赖该异常文本的外部调用。此外，本 PR 未附带自动化测试，sub-8 填充和超限路径的覆盖依赖手工验证。
• 影响：影响范围限定在使用 DeepSeek V4 模型 + FlashInfer sparse MLA 注意力后端且运行在 NVIDIA SM100/SM120 硬件上的服务。对用户而言，高 TP（TP4/8/16）下 decode 阶段 TTFT 有可观提升（参照 #48537 约 6%+），且 TP1/TP2 的 head count 不再被无谓填充。对团队而言，需要与 #47669 协同合入，并确认 CI 中 flashinfer 版本升级后行为一致；由于没有新增测试，后续需要补充针对 _pad_to_supported_q_heads 的单元测试。
• 风险标记：依赖 flashinfer>=0.6.14, 回归风险, 缺少自动化测试

关联脉络

• PR #47669 Bump flashinfer version to 0.6.14: 本 PR 的运行时依赖，需在其之后合入或堆叠；#47669 只做版本 bump 不涉及 padding 移除。