[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>

vllm/models/deepseek_v32/nvidia/mtp.py

@@ -9,7 +9,6 @@ import torch.nn as nn
 import vllm.envs as envs
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
-from vllm.distributed import tensor_model_parallel_all_reduce
 from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.model_executor.layers.fused_moe import (
     fused_moe_make_expert_params_mapping,
@@ -37,6 +36,7 @@ from vllm.model_executor.models.utils import (
     get_pp_missing_layer_names,
     maybe_prefix,
 )
+from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm
 from vllm.models.common.ops.sequence_parallel import (
     sp_all_gather,
     sp_padding_mask,
@@ -119,9 +119,6 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
         hidden_states, residual = self.mtp_block(
             positions=positions, hidden_states=hidden_states, residual=None
         )
-        if not is_sequence_parallel:
-            # Without sequence parallelism, the MoE output is left un-reduced.
-            hidden_states = tensor_model_parallel_all_reduce(hidden_states)
         # Recycle the POST-final-norm hidden into the next draft step. The
         # residual-add is fused into the final RMSNorm so it is computed
         # exactly once, and the result is returned for both tuple positions:
@@ -132,9 +129,15 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
         # is understood by both the V2 speculator (isinstance-tuple check) and
         # the legacy proposer (model_returns_tuple is True for the
         # DeepSeekMTPModel architecture).
-        hidden_states, _ = self.shared_head.norm(hidden_states, residual)
         if is_sequence_parallel:
+            hidden_states, _ = self.shared_head.norm(hidden_states, residual)
             hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]]
+        else:
+            # The MoE output is left un-reduced; fuse its all-reduce into the
+            # final norm, as the main model does at layer boundaries.
+            hidden_states, _ = fused_allreduce_rms_norm(
+                hidden_states, residual, self.shared_head.norm
+            )
         return hidden_states, hidden_states
 
 
@@ -212,6 +215,23 @@ class DeepseekV32MultiTokenPredictor(nn.Module):
         # second RMSNorm.
         return self.logits_processor(mtp_layer.shared_head.head, hidden_states)
 
+    def get_top_tokens(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        """Greedy draft token ids via per-rank argmax over the vocab shard.
+
+        Saves the full-vocab all-gather ``compute_logits`` does; same tokens.
+        Name is fixed by the protocol the proposer probes for
+        (``use_local_argmax_reduction``).
+        """
+        current_step_idx = spec_step_idx % self.num_mtp_layers
+        mtp_layer = self.layers[str(self.mtp_start_layer_idx + current_step_idx)]
+        return self.logits_processor.get_top_tokens(
+            mtp_layer.shared_head.head, hidden_states
+        )
+
 
 class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
@@ -262,6 +282,14 @@ class DeepseekV32MTP(nn.Module, DeepseekV2MixtureOfExperts):
     ) -> torch.Tensor | None:
         return self.model.compute_logits(hidden_states, spec_step_idx)
 
+    def get_top_tokens(
+        self,
+        hidden_states: torch.Tensor,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        """See ``DeepseekV32MultiTokenPredictor.get_top_tokens``."""
+        return self.model.get_top_tokens(hidden_states, spec_step_idx)
+
     def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
         spec_layer_weight_names = [
             "embed_tokens",