# PR #48429 完整报告

- 仓库：`vllm-project/vllm`
- 标题：[BugFix] Restore full tokens for Qwen MTP When MoE SP
- 合并时间：2026-07-13 13:29
- 原文链接：http://prhub.com.cn/vllm-project/vllm/pull/48429

---

# 执行摘要

- 一句话：修复 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 视图。

```python
# 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 完全对等的修复。

```python
# 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 视图损坏。