执行摘要

• 一句话：V2 下 MTP 投机解码打通流水线并行
• 推荐动作：值得精读。这是 MTP 与 pipeline parallelism 组合的补齐之作，也是典型的「单个功能背后藏着 6 个独立 bug」的案例。重点学习三个设计决策：SharedTopkIndicesBuffer 以 property 懒解析替代构造快照，规避 nn.Module 遍历与普通对象之间的可见性差异；Qwen3.5 draft head 按输入（intermediate_tensors is None）而非 rank 分支；tied 顶层权重在 PP 下必须由 draft 自持加载。对要新增 MTP 模型或改 sparse MLA 后端的工程师，本 PR 是必读样板。

功能与动机

MTP speculative decoding 在 V2 model runner 上无法与 pipeline parallelism 组合使用：#50514 落地了通用的 PP spec-decode draft-token 传输，但 MTP 特有的缺口仍然存在。PR body 明确指出主要障碍是引擎直接拒绝构建 draft："NotImplementedError: Pipeline parallelism is not supported for this model"，以及 Qwen3.5 draft head 因按 rank 分支而跳过 fc 投影导致接受率仅约 1%，DeepSeek MTP 在 PP 下用随机初始化的 embed_tokens 嵌入导致接受率 32.69% 对 93.51% 的塌缩。

实现拆解

实现按 6 个独立修复点拆解，整体围绕「让 MTP draft 在最后一个 PP stage 上可构建、可正确加载、可正确读取共享缓冲区」展开：

1. 为 MTP 模型声明 SupportsPP 并提供 make_empty_intermediate_tensors：DeepSeekMTP（vllm/model_executor/models/deepseek_mtp.py）、Qwen3_5MTP（vllm/model_executor/models/qwen3_5_mtp.py）及 DeepseekV32MTP 的 nvidia/amd 变体（vllm/models/deepseek_v32/{nvidia,amd}/mtp.py）均继承 SupportsPP 接口并在 __init__ 中注册 make_empty_intermediate_tensors_factory(["hidden_states", "residual"], hidden_size)。draft 模型整体只跑在最后一个 PP stage，不消费中间张量，但接口要求提供工厂，因此这是最小必要改动。

2. 修正 Qwen3.5 draft head 的分支条件：Qwen3_5MultiTokenPredictor.forward 原来以 get_pp_group().is_first_rank 决定是否执行 fc 投影，而 draft 构建在末段 rank 上恒为 False，导致投影被跳过、头读到无数据的 intermediate 路径。改为按 intermediate_tensors is None 分支——对流水线化的模块等价，对「整个 drafter 落在单 rank」的模型才是正确判定。

3. DeepSeek MTP 在 PP 下加载自己的 embed_tokens：checkpoint 中顶层 tied embed_tokens.weight 的 spec_layer 为 None，原 load_weights 会直接跳过；无 PP 时 draft 可与 target 共用，但 PP 下 target 的副本是 PPMissingLayer。修复在 spec_layer is None 分支内，对名字含 embed_tokens 的参数显式调用 weight_loader 加载。

4. 补齐 DeepseekV32MTP 的遗漏：该模型从 DeepSeekMTP fork 出来后既没带走 SupportsPP 也没带走 embed_tokens 加载逻辑，而 speculative.py 会把 deepseek_v32/glm_moe_dsa 路由到这里（即 GLM-5.2 与 DeepSeek-V3.2 路径）。nvidia 与 amd 两份实现同步恢复。

5. sparse MLA 后端引入 SharedTopkIndicesBuffer mixin：LLMBaseProposer.load_model 在 impl 构造完成之后才把 draft 的 indexer 重指向 target 的 topk_indices_buffer，而 named_modules() 能走到 indexer（nn.Module）却走不到 AttentionImplBase 对象，导致 impl 快照到无人写入的旧缓冲区。mixin 将缓冲区改为按需解析（property 每次从 indexer 读取），显式赋值优先于 indexer。SparseMLACommonImpl 以及非子类的 FlashInferMLASparseSM120Impl、ROCMAiterMLASparseImpl、XPUMLASparseImpl 全部接入。

6. 多模态 draft gather 按 rank 收窄：vllm/v1/worker/gpu/model_runner.py 中 gather_mm_embeddings 的触发条件追加 self.model_state.supports_mm_inputs，避免末段 PP rank 上无 encoder_runner 时的 AttributeError。

7. 测试与配套：tests/v1/attention/test_sparse_mla_backends.py 新增 3 个用例覆盖缓冲区重指向回归、显式赋值优先级、以及整个 SparseMLACommonImpl 家族的懒解析约束；tests/v1/e2e/spec_decode/test_mtp_parallel_load.py 启用了 PP=2 的 MTP 并行加载测试；vllm/models/deepseek_v32/nvidia/model.py 将原来持有在局部变量中的 buffer 暴露到模型属性上，使 proposer 的共享 block 对 DSA 路径生效。

关键文件：

• vllm/model_executor/layers/attention/sparse_mla_attention.py（模块 注意力层；类别 source；类型 data-contract；符号 SharedTopkIndicesBuffer, init_topk_indices_buffer, topk_indices_buffer, SparseMLACommonImpl）: 新增 SharedTopkIndicesBuffer mixin，将 topk 缓冲区从构造快照改为按需解析，是修复「draft 读取无人写入缓冲区」的核心数据契约变更。
• vllm/model_executor/models/deepseek_mtp.py（模块 MTP 模型；类别 source；类型 data-contract；符号 DeepSeekMTP, forward, load_weights）: DeepSeekMTP 声明 SupportsPP，并在 load_weights 中为 PP 场景显式加载顶层 tied embed_tokens，修复接受率 32.69% -> 93.51%。
• vllm/models/deepseek_v32/amd/mtp.py（模块 MTP 模型；类别 source；类型 data-contract；符号 DeepseekV32MTP, forward, load_weights）: DeepseekV32MTP 的 ROCm 变体同时补回 SupportsPP 与 embed_tokens 加载，覆盖 GLM-5.2 / DeepSeek-V3.2 的 DSA 路径。
• vllm/models/deepseek_v32/nvidia/mtp.py（模块 MTP 模型；类别 source；类型 data-contract；符号 DeepseekV32MTP, forward, load_weights）: DeepseekV32MTP 的 NVIDIA 变体同样补回 SupportsPP 与 embed_tokens 加载；配合 model.py 暴露 topk buffer 以启用 DSA 的缓冲区共享。
• vllm/model_executor/models/qwen3_5_mtp.py（模块 MTP 模型；类别 source；类型 data-contract；符号 Qwen3_5MTP, Qwen3_5MultiTokenPredictor.forward, make_empty_intermediate_tensors）: Qwen3.5 MTP draft head 改按 intermediate_tensors is None 分支，修复末段 PP rank 上 fc 投影被跳过导致的 ~1% 接受率；同时声明 SupportsPP。
• vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py（模块 MLA 后端；类别 source；类型 core-logic；符号 ROCMAiterMLASparseImpl）: ROCMAiterMLASparseImpl 接入 SharedTopkIndicesBuffer，替换原构造时快照逻辑。
• vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py（模块 MLA 后端；类别 source；类型 core-logic；符号 FlashInferMLASparseSM120Impl）: FlashInferMLASparseSM120Impl 作为非 SparseMLACommonImpl 子类，同样接入 SharedTopkIndicesBuffer。
• vllm/v1/attention/backends/mla/xpu_mla_sparse.py（模块 MLA 后端；类别 source；类型 core-logic；符号 XPUMLASparseImpl）: XPUMLASparseImpl 作为非 SparseMLACommonImpl 子类，同样接入 SharedTopkIndicesBuffer。
• vllm/v1/worker/gpu/model_runner.py（模块 运行器；类别 source；类型 core-logic；符号 sample_tokens）: 多模态 draft 的 MM embedding gather 增加 model_state.supports_mm_inputs 条件，避免末段 PP rank 无 encoder_runner 时 warmup 崩溃。
• vllm/models/deepseek_v32/nvidia/model.py（模块 模型层；类别 source；类型 data-contract）: 将原保存在局部变量中的 topk_indices_buffer 暴露到 DeepseekV32Model 属性上，使 DSA 路径能参与 proposer 的缓冲区共享。
• tests/v1/attention/test_sparse_mla_backends.py（模块 单元测试；类别 test；类型 test-coverage；符号 test_sparse_impl_observes_repointed_indexer_buffer, test_explicit_topk_buffer_supersedes_indexer, test_sparse_mla_common_impl_resolves_buffer_lazily）: 新增 3 个单测覆盖缓冲区重指向回归、显式赋值优先级、SparseMLACommonImpl 家族懒解析约束。
• tests/v1/e2e/spec_decode/test_mtp_parallel_load.py（模块 单元测试；类别 test；类型 test-coverage）: 启用 PP=2 的 MTP 并行加载测试，验证 MTP+PP 组合可构建可加载。

关键符号：SharedTopkIndicesBuffer.init_topk_indices_buffer, SharedTopkIndicesBuffer.topk_indices_buffer, DeepSeekMTP.load_weights, DeepSeekMTP.forward, DeepseekV32MTP.load_weights, Qwen3_5MultiTokenPredictor.forward, sample_tokens

关键源码片段

vllm/model_executor/layers/attention/sparse_mla_attention.py

新增 SharedTopkIndicesBuffer mixin，将 topk 缓冲区从构造快照改为按需解析，是修复「draft 读取无人写入缓冲区」的核心数据契约变更。

class SharedTopkIndicesBuffer:
    """为 sparse MLA 实现解析共享的 top-k 索引缓冲区。
​
    indexer 拥有缓冲区，但 `LLMBaseProposer.load_model` 在 impl 构造完成
    之后才会把 draft 的 indexer 重指向 target 的缓冲区，因此这里必须在每次
    读取时解析，而不是在 `__init__` 里快照。骨干网络的 skip-topk 层没有
    indexer，会显式传入缓冲区。
    """
​
    _indexer: object | None = None
    _topk_indices_buffer: torch.Tensor | None = None
​
    def init_topk_indices_buffer(
        self,
        indexer: object | None,
        topk_indices_buffer: torch.Tensor | None,
    ) -> None:
        self._indexer = indexer
        self._topk_indices_buffer = topk_indices_buffer
​
    @property
    def topk_indices_buffer(self) -> torch.Tensor | None:
        # 每次读取都从 indexer 上现取，保持与重指向后的缓冲区同步。
        if self._indexer is not None:
            return self._indexer.topk_indices_buffer # type: ignore[attr-defined]
        return self._topk_indices_buffer
​
    @topk_indices_buffer.setter
    def topk_indices_buffer(self, buffer: torch.Tensor | None) -> None:
        # 显式赋值优先于 indexer：设置后直接切换到显式缓冲区，
        # 这是 proposer 模块遍历和 skip-topk 层赋值共用的一条路径。
        self._indexer = None
        self._topk_indices_buffer = buffer
​
​
class SparseMLACommonImpl(MLACommonBaseImpl[T], SharedTopkIndicesBuffer, Generic[T]):
    """Sparse MLA 基类，包含 dense 与 masked-MHA 两条 prefill 路径。"""
​
    is_sparse = True
​
    def __init__(
        self,
        ...
        indexer: object | None = None,
        topk_indices_buffer: torch.Tensor | None = None,
        ...
    ) -> None:
        super().__init__(...)
        # 不再直接快照 indexer 的缓冲区，而是交给 mixin 做延迟解析，
        # 以便后续 proposer 重指向 indexer 时对 impl 立即生效。
        self.init_topk_indices_buffer(indexer, topk_indices_buffer)
        ...

vllm/model_executor/models/deepseek_mtp.py

DeepSeekMTP 声明 SupportsPP，并在 load_weights 中为 PP 场景显式加载顶层 tied embed_tokens，修复接受率 32.69% -> 93.51%。

@support_torch_compile
class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts, SupportsPP):
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        self.config = vllm_config.model_config.hf_config
        self.quant_config = vllm_config.quant_config
        self.model = DeepSeekMultiTokenPredictor(
            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
        )
        self.set_moe_parameters()
        # PP 接口要求提供中间张量工厂；draft 只跑在末段 rank，
        # 不消费中间张量，但引擎仍会按 SupportsPP 检查该工厂。
        self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
            ["hidden_states", "residual"], self.config.hidden_size
        )
​
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
        ...
        for name, loaded_weight in weights:
            if "rotary_emb.inv_freq" in name:
                continue
            spec_layer = get_spec_layer_idx_from_weight_name(self.config, name)
            if spec_layer is None:
                # 顶层 tied embed_tokens 没有可重写的 spec 层；PP 下 target
                # 的副本在 draft 所在 stage 上是 PPMissingLayer，无法借用，
                # 所以 draft 必须加载自己的副本，否则会用随机权重嵌入。
                param = params_dict.get(name) if "embed_tokens" in name else None
                if param is not None:
                    weight_loader = getattr(
                        param, "weight_loader", default_weight_loader
                    )
                    weight_loader(param, loaded_weight)
                    loaded_params.add(name)
                continue
            ...

vllm/model_executor/models/qwen3_5_mtp.py

Qwen3.5 MTP draft head 改按 intermediate_tensors is None 分支，修复末段 PP rank 上 fc 投影被跳过导致的 ~1% 接受率；同时声明 SupportsPP。

class Qwen3_5MultiTokenPredictor(nn.Module):
    ...
    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
        hidden_states: torch.Tensor,
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
        spec_step_idx: int = 0,
    ) -> torch.Tensor:
        # 按输入分支而不是按 rank 分支：drafter 整体构建在最后一个 PP stage 上，
        # 此时 get_pp_group().is_first_rank 恒为 False，按 rank 分支会让 fc 投影
        # 被跳过，头落到读取空中间张量的错误路径上。
        if intermediate_tensors is None:
            if inputs_embeds is None:
                inputs_embeds = self.embed_input_ids(input_ids)
            assert hidden_states.shape[-1] == inputs_embeds.shape[-1]
            inputs_embeds = self.pre_fc_norm_embedding(inputs_embeds)
            hidden_states = self.pre_fc_norm_hidden(hidden_states)
            hidden_states = torch.cat([inputs_embeds, hidden_states], dim=-1)
            hidden_states = self.fc(hidden_states)
            residual = None
        else:
            hidden_states = intermediate_tensors["hidden_states"]
            residual = intermediate_tensors["residual"]
        ...

评论区精华

评审与社区讨论中最有价值的几轮交锋：

• intermediate_tensors plumbing 是否该删：yewentao256 提出整条新增的 intermediate-tensor 传输没有真实数据流，且引发 fused multi-step TypeError，建议连参数、转发逻辑和 Qwen3.5 里的零填充张量一起移除。eastwood-c 随后按建议删除，相关四个 v1/worker/gpu/spec_decode/ 文件恢复与上游逐字节一致。
• 与 #50514 的重叠：njhill 指出 #50514 准备先行落地，eastwood-c 随即砍掉自己 PR 里的 draft-token relay，只保留 MTP 特有修复，并拆出 #55745 的 record_stream 守护作为前置依赖。
• Draft 质量塌缩的自我发现：eastwood-c 在合并 main 后发现接受率回归到修复前水平，定位到 embed_tokens 未加载，用 dirty working tree 找回遗漏改动，并给出 32.69% → 93.51% 的量化对比。
• NCCL 死锁的外部复现：hackxixi 在 Qwen4Exp（第四种架构）上独立复现同一死锁，确认 broadcast() 发送 [num_reqs, 1] 而 receive() 预定 [num_reqs, max_sample_len] 时 NCCL 不协商元素数、接收永不完成，用 pad 到 max_sample_len 得到与本 PR 相同的修复。

• 测试取舍：yewentao256 认为 test_pp_utils.py 新增测试意义不大可移除；zixi-qi 建议启用 test_mtp_parallel_load.py 中的 PP 测试，njhill 直接代为合入。

  • 是否移除整条 intermediate_tensors plumbing (design): eastwood-c 按建议删除，四个 v1/worker/gpu/spec_decode/ 文件恢复为与上游逐字节一致；保留 SupportsPP 与 make_empty_intermediate_tensors 工厂。
  • 与 #50514 的职责重叠与 rebase (design): eastwood-c 砍掉自有 relay 改为依赖 #50514，只保留 MTP 特有修复，并拆出 #55745 的 record_stream 守护作为前置。
  • embed_tokens 未加载导致接受率塌缩 (correctness): 在 load_weights 的 spec_layer is None 分支中显式加载 embed_tokens；接受率从 32.69% 恢复到 93.51%。
  • NCCL 死锁：broadcast 宽度不一致 (correctness): 确认与本 PR 的 sampled-token 分发修复一致，将源 padding 到 max_sample_len 即可解决。
  • test_pp_utils.py 新增测试的意义 (testing): 随着 intermediate_tensors plumbing 移除与 #50514 落地，该测试不再保留在最终改动集中。
  • 启用 PP=2 MTP 并行加载测试 (testing): njhill 提交 bb07636 启用了该测试。

  • 55745 record_stream 守护依赖 (correctness): 拆出 #55745 单行修复作为本 PR 前置依赖，合入后从本 PR 中移除该 commit。

风险与影响

• 风险：风险集中在四处：

1. sparse MLA 缓冲区语义变更：topk_indices_buffer 从构造时快照变为每次读取时解析，影响 SparseMLACommonImpl 全部子类（FLASHMLA、flashinfer sm90 等）以及 sm120、ROCm、XPU 三个独立实现。若未来有代码在构造后、重指向前提前读取缓冲区，行为会与旧版不同；不过按 PR 设计这恰恰是修复方向。
2. 权重加载路径变更：DeepSeekMTP/DeepseekV32MTP.load_weights 现在会把顶层 embed_tokens 显式加载到 draft。对逐 spec-layer 存 embed_tokens 的 checkpoint 不受影响（它们的名字命中 spec_layer 分支），但对其他误命中 embed_tokens 名字的权重需警惕重复加载。
3. Qwen3.5 forward 分支条件等价性：intermediate_tensors is None 与 is_first_rank 对「整个 drafter 在同一 rank」的模型等价，但依赖调用方始终传入或始终不传 intermediate tensors 的约定，未来若出现 drafter 跨 PP stage 拆分，此判定需要重新审视。
4. 硬件验证面：ROCm 与 XPU 后端仅做了代码接入，测试主要跑在 CUDA；rocm_aiter_mla_sparse.py 的测试文件头部仍保留 CUDA-only skip，ROCm 侧缺少针对缓冲区重指向的真实硬件回归。
   - 影响：影响面集中在「PP>1 + MTP」组合场景：

• 用户侧：Qwen3.5/3.6、GLM-5.2、DeepSeek-V3.2 等 MTP 模型在 --pipeline-parallel-size > 1 下可直接启用 --speculative-config '{"method":"mtp"}'，接受率从不可用/1%/32% 恢复到 83%~96%，并随 K=1/2/3 平滑退化；max_tokens=2048 下 PP=1 与 PP=2 的接受率几乎一致（1.9375 vs 1.9354~1.9380）。
• 系统侧：sparse MLA 注意力在「draft 共享 target 缓冲区」场景下的正确性被修正，DSA 路径的 topk_indices_buffer 暴露方式与 ROCm 对齐。
• 团队侧：确立了 MTP draft 模型在 PP 下的四类标准姿势——SupportsPP 声明、embed_tokens 自持、按输入而非 rank 分支、共享缓冲区懒解析；后续其他 MTP 家族接入可照此模板。
• 风险标记：跨平台注意力后端改动, 权重加载路径变更, 前置依赖已合并, ROCm/XPU 硬件验证有限

关联脉络

• PR #50514 Generic PP spec-decode transport: 本 PR 依赖其提供的 PP draft-token relay；早期自研版本已砍掉，最终只保留 MTP 特有修复。
• PR #55745 Record stream guard in broadcast_drafts: 一行的 record_stream 守护，修复 async launch 下广播 gather 越界的 IndexKernel assert，本 PR 的必经前置。
• PR #47686 PP broadcast pack: eastwood-c 堆叠在本 PR 之上的优化：将 PPHandler 每步约 3 个 Broadcast_RING_LL kernel 合并为单个 op，与 46994 同源演进。
• PR #54709 qwen4_exp draft head fix: 与修复点 2 同一种 bug（draft head 分支条件错误），采用相同的一行修复，本 PR 为聚焦而刻意排除。
• PR #55081 SupportsPP on MTP families: 与修复点 4 的 SupportsPP 部分重叠；若先合入则本 PR 可丢弃该半段。
• PR #39704 V1 runner MTP PP support: PR body 明确引用为先例：为 V1 runner 添加 MTP SupportsPP 的镜像实现。