[Spec Decode][Perf] Fuse the MTP trailing all-reduce; local-argmax draft tokens (#49793)
Signed-off-by: Peiyuan Zhou <peiyuanzhou1994@gmail.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
xiaozhoupy committed
Aug 15, 2026 at 08:07 UTC
7b544ecb524b6a5847465b671db2df9b8560ad93
1 file changed
+33
-5
vllm/models/deepseek_v32/nvidia/mtp.py
+33
-5
@@ -9,7 +9,6 @@ import torch.nn as nn
9
import vllm.envs as envs
10
from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.config import VllmConfig
12
-from vllm.distributed import tensor_model_parallel_all_reduce
12
from vllm.forward_context import get_forward_context, is_forward_context_available
13
from vllm.model_executor.layers.fused_moe import (
14
fused_moe_make_expert_params_mapping,
@@ -37,6 +36,7 @@ from vllm.model_executor.models.utils import (
36
get_pp_missing_layer_names,
37
maybe_prefix,
38
)
39
+from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm
40
from vllm.models.common.ops.sequence_parallel import (
41
sp_all_gather,
42
sp_padding_mask,
@@ -119,9 +119,6 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
119
hidden_states, residual = self.mtp_block(
120
positions=positions, hidden_states=hidden_states, residual=None
121
)
122
- if not is_sequence_parallel:
123
- # Without sequence parallelism, the MoE output is left un-reduced.
124
- hidden_states = tensor_model_parallel_all_reduce(hidden_states)
122
# Recycle the POST-final-norm hidden into the next draft step. The
123
# residual-add is fused into the final RMSNorm so it is computed
124
# exactly once, and the result is returned for both tuple positions:
@@ -132,9 +129,15 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
129
# is understood by both the V2 speculator (isinstance-tuple check) and
130
# the legacy proposer (model_returns_tuple is True for the
131
# DeepSeekMTPModel architecture).
135
- hidden_states, _ = self.shared_head.norm(hidden_states, residual)
132
if is_sequence_parallel:
133
+ hidden_states, _ = self.shared_head.norm(hidden_states, residual)
134
hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]]
135
+ else:
136
+ # The MoE output is left un-reduced; fuse its all-reduce into the
137
+ # final norm, as the main model does at layer boundaries.
138
+ hidden_states, _ = fused_allreduce_rms_norm(
139
+ hidden_states, residual, self.shared_head.norm
140
+ )
141
return hidden_states, hidden_states
142
143
@@ -212,6 +215,23 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
215
# second RMSNorm.
216
return self.logits_processor(mtp_layer.shared_head.head, hidden_states)
217
218
+ def get_top_tokens(
219
+ self,
220
+ hidden_states: torch.Tensor,
221
+ spec_step_idx: int = 0,
222
+ ) -> torch.Tensor:
223
+ """Greedy draft token ids via per-rank argmax over the vocab shard.
224
+
225
+ Saves the full-vocab all-gather ``compute_logits`` does; same tokens.
226
+ Name is fixed by the protocol the proposer probes for
227
+ (``use_local_argmax_reduction``).
228
+ """
229
+ current_step_idx = spec_step_idx % self.num_mtp_layers
230
+ mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
231
+ return self.logits_processor.get_top_tokens(
232
+ mtp_layer.shared_head.head, hidden_states
233
+ )
234
+
235
236
class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
237
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -262,6 +282,14 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
282
) -> torch.Tensor | None:
283
return self.model.compute_logits(hidden_states, spec_step_idx)
284
285
+ def get_top_tokens(
286
+ self,
287
+ hidden_states: torch.Tensor,
288
+ spec_step_idx: int = 0,
289
+ ) -> torch.Tensor:
290
+ """See ``DeepseekV32MultiTokenPredictor.get_top_tokens``."""
291
+ return self.model.get_top_tokens(hidden_states, spec_step_idx)
292
+
293
def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
294
spec_layer_weight_names = [
295
"embed_tokens",