执行摘要

• 一句话：修复 Qwen MTP 在 MoE SP 下的 token 损坏
• 推荐动作：建议精读本 PR，理解 _all_gather_hidden_and_residual 的修复模式。该模式可作为未来 PR #47006 类似变更影响 MTP 路径时的参考样板。同时建议为 MoE SP + MTP 组合添加自动化回归测试。

功能与动机

PR #47006 优化了 Qwen MoE 通信，将 all-reduce 替换为 reduce-scatter，但导致 Qwen MTP 在 MoE SP 开启时 CUDA assert 崩溃。原因是 MTP 直接调用 Qwen 的 DecoderLayer（而非完整模型），而 #47006 将 gather 移到 attention 前，使得 MTP 调用时共享了错误的 token 视图。

实现拆解

1. 导入新函数：在 qwen3_5_mtp.py 和 qwen3_next_mtp.py 中，从 vllm.model_executor.models.qwen3_next 导入 _all_gather_hidden_and_residual。
2. 缓存当前层引用：将 self.layers[current_step_idx] 赋值给 mtp_layer 变量，避免多次索引并便于后续访问层属性。
3. 条件性 all-gather：在 norm 之前，检查 mtp_layer.use_attn_reduce_scatter_for_moe，若为 True（即 MoE SP 已启用），则调用 _all_gather_hidden_and_residual 恢复完整的 hidden_states 和 residual 张量。_all_gather_hidden_and_residual 函数负责在 tensor-parallel 组内 all-gather 缩减的 token 维度，还原成完整序列视图。

关键文件：

• vllm/model_executor/models/qwen3_5_mtp.py（模块 模型层；类别 source；类型 data-contract；符号 forward, Qwen3_5MultiTokenPredictor）: Qwen3.5 MTP 模型文件，修复核心改动：在 norm 前调用 _all_gather_hidden_and_residual 恢复完整 token 视图。
• vllm/model_executor/models/qwen3_next_mtp.py（模块 模型层；类别 source；类型 data-contract；符号 forward, Qwen3NextMultiTokenPredictor）: Qwen3Next MTP 模型文件，与 qwen3_5_mtp.py 完全对等的修复。

关键符号：Qwen3_5MultiTokenPredictor.forward, Qwen3NextMultiTokenPredictor.forward

关键源码片段

vllm/model_executor/models/qwen3_5_mtp.py

Qwen3.5 MTP 模型文件，修复核心改动：在 norm 前调用 _all_gather_hidden_and_residual 恢复完整 token 视图。

# qwen3_5_mtp.py - 在 MTP forward 中根据层属性恢复完整 token 视图
# 新增导入
from vllm.model_executor.models.qwen3_next import (
    QwenNextMixtureOfExperts,
    _all_gather_hidden_and_residual, # 新增：用于 gather 缩减的 token 维度
    _is_shared_expert_fse_compatible,
)
​
# forward 方法中的新增逻辑（位于 norm 之前）
if mtp_layer.use_attn_reduce_scatter_for_moe:
    # 当 MoE 使用 reduce-scatter 时，MTP 层输出中的 hidden_states 和 residual
    # 仅在部分 token 维度上有效（经过 scatter）。此调用在 tensor-parallel 组内
    # all-gather 完整的 token 视图，确保后续 norm 得到正确的全局表示。
    hidden_states, residual = _all_gather_hidden_and_residual(
        hidden_states,
        residual,
        positions.shape[-1], # 完整序列长度
        self.config.hidden_size, # 完整的 hidden size
    )
hidden_states, _ = self.norm(hidden_states, residual)
return hidden_states

vllm/model_executor/models/qwen3_next_mtp.py

Qwen3Next MTP 模型文件，与 qwen3_5_mtp.py 完全对等的修复。

# qwen3_next_mtp.py - 与 qwen3_5_mtp.py 完全相同的修复模式
# 新增导入
from vllm.model_executor.models.qwen3_next import (
    Qwen3NextDecoderLayer,
    Qwen3NextModel,
    Qwen3NextRMSNorm,
    QwenNextMixtureOfExperts,
    _all_gather_hidden_and_residual, # 新增
)
​
# forward 方法中的新增逻辑
if mtp_layer.use_attn_reduce_scatter_for_moe:
    hidden_states, residual = _all_gather_hidden_and_residual(
        hidden_states,
        residual,
        positions.shape[-1],
        self.config.hidden_size,
    )
hidden_states, _ = self.norm(hidden_states, residual)
return hidden_states

评论区精华

PR 无直接 review 评论。自动审核机器人 claude[bot] 仅为 fork PR 留了一条自动化提示。维护者 ZJY0516 直接批准。没有公开的设计辩论。

• 暂无高价值评论线程

风险与影响

• 风险：
  • 回归风险低：改动范围小（2 个文件，共 +20/-2），仅在 MoE SP 条件分支下增加 all-gather 调用，对非 SP 路径无影响。
  • 缺失测试覆盖：PR 未附带对应单元测试或集成测试，仅提供了手动验证命令。未来若修改 _all_gather_hidden_and_residual 签名或行为，无自动化测试保障。
  • 依赖紧耦合：use_attn_reduce_scatter_for_moe 属性依赖 DecoderLayer 的实现细节，若该属性重命名或语义变更，本修复将静默失效。
• 影响：
  • 用户：修复 Qwen3.5 和 Qwen3Next 使用 MTP 投机解码 + MoE SP 时的崩溃，使该组合正常工作。lm_eval gsm8k 测试显示准确率恢复至 ~85.8%。
  • 系统：对非 MTP 或非 MoE SP 配置无影响。
  • 团队：确认了架构层设计决策（gather 前置）对 MTP 路径的副作用，需在后续类似优化中同步考虑。
  • 风险标记：核心路径变更, 缺少测试覆盖

关联脉络

• PR #47006 [Perf][Qwen] Replace MOE all-reduce with reduce-scatter: 本 PR 正是为了修复 #47006 引入的回归。#47006 将 MoE 的 all-reduce 替换为 reduce-scatter，导致 MTP 路径中 token 视图损坏。