执行摘要

• 一句话：Dots3 NOTE 运行时优化，吞吐提升约 9.5%
• 推荐动作：值得精读。该 PR 展示了将自定义 MLA metadata builder 接入 CUDA Graph uniform-batch 的标准模式，以及如何在 MTP 投机解码中融合 embedding 与 RMSNorm。建议关注 _pad_dense_mlp_weight 的模块级化与 replicated_embed 的传播方式，这两个设计点可复用到其他 NOTE 类模型。

功能与动机

Dots3 NOTE 此前在 Model Runner V1 下运行，MTP 投机解码路径无法被 CUDA Graph 捕获，GPU 上存在大量 graph 之外的调度与 kernel 启动开销。PR body 指出优化目标是让语言模型和视觉编码器的运行时路径更高效，其中 decode 路径的剩余 GPU 工作应被压缩到 LM head 与 sampling/metadata 操作。

实现拆解

1. 默认启用 Model Runner V2：在 vllm/config/vllm.py 的 DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES 中加入 Dots3NoteForCausalLM 与 Dots3NoteMTPModel，使这两个架构在 NVIDIA 平台默认走 V2 Model Runner 并可打断 CUDA Graph。
2. 为 MLA metadata builder 声明 uniform-batch CUDA Graph 支持：在 vllm/models/dots3_note/nvidia/attention.py 中给 Dots3NoteMLAMetadataBuilder 增加 _cudagraph_support = AttentionCGSupport.UNIFORM_BATCH，使 decode 路径可被 CUDA Graph 捕获。
3. 融合 MTP embedding 与 RMSNorm：在 mtp.py 的 Dots3NoteMultiTokenPredictorLayer.forward 中新增 embed_table 分支，调用 fused_embed_eh_norm；Dots3NoteMultiTokenPredictor 记录 replicated_embed 状态，并在词表复制时跳过独立 embedding 计算；Dots3NoteModel 同样记录 replicated_embed。
4. 统一 dense MLP 权重 padding 实现：把 Dots3NoteModel._pad_dense_mlp_weight 提升为模块级函数 _pad_dense_mlp_weight(name, loaded_weight, weight_block_size)，MTP 的 Dots3NoteMTP._adapt_weights 改为复用该函数，消除两处重复实现。
5. 统一视觉 RMSNorm 实现：删除 vision.py 中重复定义的 RMSNorm，统一改用 vision_attention.py 导出的 VisionRMSNorm（含 extra_repr），并将 fallback 提示从 print 改为 logger.warning_once。

关键文件：

• vllm/models/dots3_note/nvidia/mtp.py（模块 模型层；类别 source；类型 data-contract；符号 Dots3NoteMultiTokenPredictorLayer, Dots3NoteMultiTokenPredictor, Dots3NoteMTP, _adapt_weights）: MTP 投机解码路径的核心改动：新增 fused_embed_eh_norm 分支、replicated_embed 判断，并复用模块级 _pad_dense_mlp_weight。这是 CUDA Graph 捕获和 embedding 融合的关键。
• vllm/models/dots3_note/nvidia/model.py（模块 模型层；类别 source；类型 data-contract；符号 _pad_dense_mlp_weight, Dots3NoteModel, _adapt_weights）: backbone 模型的权重 padding 方法被提炼为模块级 _pad_dense_mlp_weight，并新增 replicated_embed 判断；是 MTP 融合路径的配套基础。
• vllm/models/dots3_note/nvidia/vision_attention.py（模块 模型层；类别 source；类型 data-contract；符号 VisionRMSNorm, resolve_attn_implementation）: 视觉 RMSNorm 统一为 VisionRMSNorm，并加入 all 导出；fallback 提示从 print 改为 logger.warning_once，影响视觉编码器调试体验。
• vllm/config/vllm.py（模块 配置层；类别 source；类型 core-logic；符号 DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES）: 将 Dots3NoteForCausalLM 与 Dots3NoteMTPModel 加入默认可打断 CUDA Graph 架构集合，是 Model Runner V2 默认启用的开关。
• vllm/models/dots3_note/nvidia/attention.py（模块 模型层；类别 source；类型 data-contract；符号 Dots3NoteMLAMetadataBuilder）: 为 Dots3NoteMLAMetadataBuilder 声明 AttentionCGSupport.UNIFORM_BATCH，这是 decode 路径能被 CUDA Graph 捕获的直接条件。
• vllm/models/dots3_note/nvidia/vision.py（模块 模型层；类别 source；类型 data-contract；符号 RMSNorm, DotsPatchEmbed, MoEVisionBlock, DotsMoEVitModel）: 删除重复的 RMSNorm 定义，改用 VisionRMSNorm；是视觉编码器侧代码清理与统一。

关键符号：_pad_dense_mlp_weight, Dots3NoteMultiTokenPredictorLayer.forward, Dots3NoteMultiTokenPredictor.forward, Dots3NoteMTP._adapt_weights, Dots3NoteModel._adapt_weights, VisionRMSNorm.forward, resolve_attn_implementation

关键源码片段

vllm/models/dots3_note/nvidia/mtp.py

MTP 投机解码路径的核心改动：新增 fused_embed_eh_norm 分支、replicated_embed 判断，并复用模块级 _pad_dense_mlp_weight。这是 CUDA Graph 捕获和 embedding 融合的关键。

# vllm/models/dots3_note/nvidia/mtp.py
​
class Dots3NoteMultiTokenPredictorLayer(nn.Module):
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        previous_hidden_states: torch.Tensor,
        inputs_embeds: torch.Tensor | None = None,
        embed_table: torch.Tensor | None = None,
        spec_step_index: int = 0,
    ) -> torch.Tensor:
        del spec_step_index
        if embed_table is not None:
            # 融合 embedding 查找 + E/H 两个 RMSNorm，减少 graph 外 kernel 启动
            eh_input = fused_embed_eh_norm(
                positions,
                input_ids,
                embed_table,
                previous_hidden_states,
                self.enorm.weight,
                self.hnorm.weight,
                self.enorm.variance_epsilon,
            )
        else:
            assert inputs_embeds is not None
            eh_input = fused_eh_norm(
                positions,
                inputs_embeds,
                previous_hidden_states,
                self.enorm.weight,
                self.hnorm.weight,
                self.enorm.variance_epsilon,
            )
        hidden_states = self.eh_proj(eh_input)[0]
        hidden_states, residual = self.mtp_block(
            positions=positions, hidden_states=hidden_states, residual=None
        )
        is_sequence_parallel = self.mtp_block.use_sequence_parallel_moe
        if not is_sequence_parallel:
            hidden_states = tensor_model_parallel_all_reduce(hidden_states)
        hidden_states, _ = self.shared_head.norm(hidden_states, residual)
        if is_sequence_parallel:
            hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
            hidden_states = hidden_states[: positions.shape[0]]
        return hidden_states

评论区精华

Review 讨论较少，主要是 jeejeelee 批准合并。提交历史中 jeejeelee 处理了 merge conflict。claude[bot] 因 fork 仓库未启用自动 review。核心权衡在于 Dots3NoteMultiTokenPredictorLayer.forward 中 embedding 路径的处理：新增 embed_table 参数后，inputs_embeds is None 时不再调用 self.embed_input_ids，而是直接对 input_ids 做 torch.minimum(input_ids, self.max_token_id) 后交给父类处理，这保留了 embed_input_ids 的边界裁剪语义。

• 暂无高价值评论线程

风险与影响

• 风险：
  1. Model Runner V2 默认切换风险：vllm/config/vllm.py 中把 Dots3NoteForCausalLM 加入默认 MRV2 架构列表，若用户环境对 MRV2 支持不完整，可能遇到行为差异；但 ROCm 平台已有专门排除逻辑。
  2. CUDA Graph uniform-batch 限制：Dots3NoteMLAMetadataBuilder 声明 UNIFORM_BATCH 后，非 uniform batch（如 mixed prefill/decode）可能不被 CUDA Graph 捕获，需确认调度器在 decode 阶段始终使用 uniform batch。
  3. embedding 融合边界条件：fused_embed_eh_norm 依赖 embed_table 在词表复制时可用；若 TP 配置下词表被切分，replicated_embed 为 False，会走旧路径，但需要确认 embed_table 的传递逻辑在 pipeline parallel 下正确。
  4. 视觉 RMSNorm 重构：VisionRMSNorm 与旧 RMSNorm 数值行为一致（fp32 reduce 后 cast 回原 dtype），但 extra_repr 的 shape 显示格式变化可能影响状态字典或 repr 输出。
  5. 权重 padding 共享：_pad_dense_mlp_weight 现在是模块级函数，签名从 (name, loaded_weight) 变为 (name, loaded_weight, weight_block_size)，若其他模块仍按旧签名调用会报错。
     - 影响：对 Dots3 NOTE 用户：默认路径切换后吞吐提升约 9.47%、延迟降低约 8.68%；对 MTP 投机解码场景收益更明显。对系统：CUDA Graph 捕获范围扩大可能增加显存占用，但减少了 kernel 启动开销。对团队：为其他 DeepSeek 变体模型（如 DeepseekV32）提供了 MTP + MRV2 + CUDA Graph 的参考实现模式。
     - 风险标记：核心路径变更, 默认行为切换

关联脉络

• 暂无明显关联 PR