执行摘要

• 一句话：MLA decode q-heads 填充至64以加速FlashMLA kernel
• 推荐动作：值得精读。该 PR 展示了如何通过理解底层 kernel 的 specialization 条件，以极小的代码改动（+21/-5）获得显著的性能提升。核心思路——“在满足 kernel 约束最小化的前提下填充而非总是填充到最大值”——可推广到其他类似场景。同时，_attn_sink_local 的懒初始化模式避免了 CUDA graph 内的冗余操作，是面向 GPU 图捕获性能优化的良好实践。

功能与动机

DeepSeek-V4-Pro 在使用 attention tensor parallelism 时，每个 rank 的 query head 数只有 n_heads // attn_tp（如 TP=4 时为 32），但 MQALayer.forward 将其填充至全局 n_heads（128），导致 FlashMLA 调度了慢的 fwd_for_small_topk::head128 kernel。FlashMLA 只对 h_q 为 {64, 128} 进行了专门优化，填充至 64 可触发约 2x 更便宜的 decode::head64 变体。

实现拆解

1. 调整 __init__ 中 attn_sink 的缓存策略（python/sglang/srt/models/deepseek_v4.py）：新增 self._attn_sink_local 属性，当 attn_tp_size == 1 时直接指向 self.attn_sink，否则为 None，为后续懒初始化做好准备。

2. 修改 forward 中的头部填充逻辑：当 self.tp_size > 1 时，计算 padded_num_heads——若 n_local_heads <= 64 则使用 64，否则使用 n_heads（保持向后兼容）。创建形状为 [batch, padded_num_heads, head_dim] 的 q_padded，并将 tp_slice 设置为 slice(0, n_local_heads)，使 q_out 指向该 rank 的有效头部区域。

3. 懒初始化缓存的 _attn_sink_local：在第一次 forward 调用时（此时权重已加载），为当前 rank 构建 padded sink 张量：分配 padded_num_heads 长度的零张量，并从 self.attn_sink 中复制属于该 rank 的 n_local_heads 个元素。后续 forward 重用该缓存，避免 decode CUDA graph 内部每层重复 fill+cop y 操作。

4. 更新下游调用：将 flash attention 和 unified-kv-triton 路径中的 self.attn_sink 替换为 self._attn_sink_local，确保使用已正确填充的 sink。

关键文件：

• python/sglang/srt/models/deepseek_v4.py（模块 模型层；类别 source；类型 data-contract；符号 MQALayer.init, MQALayer.forward）: 核心变更文件，修改了 MQALayer 中 query 头部填充逻辑以及 attn_sink 的缓存方式，直接影响 DeepSeek-V4 的 MLA decode 注意力性能。

关键符号：MQALayer.init, MQALayer.forward

关键源码片段

python/sglang/srt/models/deepseek_v4.py

核心变更文件，修改了 MQALayer 中 query 头部填充逻辑以及 attn_sink 的缓存方式，直接影响 DeepSeek-V4 的 MLA decode 注意力性能。

def forward(self, x, positions, forward_batch, x_quant=None):
    # ... 前置判断和 enable_multi_stream 逻辑略 ...
    tp_slice, q_padded, q_out = slice(None), None, None
    if self.tp_size > 1:
        # FlashMLA 的 fp8 sparse decode kernel 仅对 h_q = {64, 128} 做了特化。
        # 当本地 head 数 ≤ 64 时，填充到 64 以触发更快的 decode::head64 变体；
        # attn_sink 同步裁剪并填充到对应大小。
        padded_num_heads = 64 if self.n_local_heads <= 64 else self.n_heads
        q_padded = x.new_empty(x.shape[0], padded_num_heads, self.head_dim)
        tp_slice = slice(0, self.n_local_heads)
        q_out = q_padded[:, tp_slice, :]
​
        # 懒初始化 per-rank 的 attn_sink 缓存：仅首次 forward 时构建一次，
        # 避免在 decode CUDA graph 内每层重复执行 fill+cop y。
        if self._attn_sink_local is None:
            rank = self.tp_rank
            sink = self.attn_sink.new_zeros(padded_num_heads)
            sink[: self.n_local_heads] = self.attn_sink[
                rank * self.n_local_heads : (rank + 1) * self.n_local_heads
            ]
            self._attn_sink_local = sink
    # ... 后续 multi-stream / forward_prepare 分支 ...
    # 调用 attention 后端时使用 self._attn_sink_local 而非 self.attn_sink

评论区精华

Fridge003 在 review 中提出：“未来我们可能为 dsv4 实现不同的 attention kernel，因此当不使用 flashmla 内核时，可能仍需要原始的填充逻辑。”由于当前逻辑仅在 n_local_heads > 64 时保留 self.n_heads 作为回退，而此条件对于 TP 场景通常不成立（n_heads=128, TP=4 时为 32），开发者未立即修改，但这一设计权衡值得在未来迭代中关注。

• 未来 attention kernel 兼容性 (design): 当前逻辑通过条件 n_local_heads <= 64 保留了回退路径（n_heads），但该条件在 TP 场景下几乎总是成立；未来需注意在引入新 kernel 时调整或移除该优化。

风险与影响

• 风险：风险较低。修改范围集中于单个文件，且通过 n_local_heads <= 64 的条件保留了向后兼容路径：当未来使用不支持 head64 的 attention kernel 时，只需增大或移除该条件即可回退到原始 n_heads 填充。主要风险在于：
  • _attn_sink_local 的懒初始化假设第一次 forward 时权重已加载（self.attn_sink 已初始化），该假设在标准流程中成立，但若出现异构执行顺序（如 partial graph capture）需验证。
  • 没有新增测试文件，回归风险由 e2e 准确率测试覆盖（GSM8K、AIME25）。
  • 影响：正面性能影响：对 DeepSeek-V4 在 B300 上使用 attention TP 的场景，decode attention kernel 加速约 3.4x，端到端总吞吐提升 2-7%（取决于并发数）。TTFT 降低约 100ms，TPOT 略微改善。对单 rank 无 attention TP 的场景无影响。

准确性无退化：GSM8K (96.2%→96.3%) 和 AIME25 (97.5%) 准确率均无统计显著变化。

代码可维护性：增加约 15 行逻辑，但通过注释清晰说明了 padding 原因和缓存策略，维护成本低。

• 风险标记：未新增测试文件, 依赖 FlashMLA kernel specialize 条件

关联脉络

• PR #27986 [dsv4] Prewarm MHC prenorm kernel at startup: 同属 DeepSeek-V4 性能优化系列，皆涉及 MLA 相关 kernel 启动和预热。
• PR #28073 fix: Fix DSR1 perf regression due to unnecessarily falling back to triton gemm: 同为 B300 Blackwell 平台 fp8 kernel 调度优化，体现对底层 kernel specialization 的精细利用。