[Perf] Fix moe `reduce_scatter` perf regression by removing additional comm, 5% E2E throughput gain back. (#48763)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Jul 25, 2026 at 12:36 UTC
7fe6d3c76b6f44b54ba629e9accf572adfe86f6a
2 files changed
+13
-37
vllm/model_executor/models/deepseek_mtp.py
+3
-25
@@ -44,25 +44,6 @@ from .utils import (
44
)
45
46
47
-def _restore_full_token_layout_if_needed(
48
- hidden_states: torch.Tensor,
49
- residual: torch.Tensor,
50
- num_tokens: int,
51
- is_sequence_parallel: bool = False,
52
-) -> tuple[torch.Tensor, torch.Tensor]:
53
- """Restore full token rows for the MTP proposer after SP MoE layers."""
54
- if not is_sequence_parallel and hidden_states.shape[0] == num_tokens:
55
- return hidden_states, residual
56
-
57
- combined_states = torch.cat([hidden_states, residual], dim=-1)
58
- combined_states = tensor_model_parallel_all_gather(combined_states, 0)
59
- combined_states = combined_states[:num_tokens]
60
- hidden_states, residual = combined_states.split(
61
- [hidden_states.shape[-1], residual.shape[-1]], dim=-1
62
- )
63
- return hidden_states, residual
64
-
65
-
47
class SharedHead(nn.Module):
48
def __init__(
49
self,
@@ -142,13 +123,10 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
123
hidden_states=hidden_states,
124
residual=None,
125
)
145
- hidden_states, residual = _restore_full_token_layout_if_needed(
146
- hidden_states,
147
- residual,
148
- positions.shape[0],
149
- is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe,
150
- )
126
hidden_states = residual + hidden_states # pre-final-norm (logits hidden)
127
+ if self.mtp_block.use_sequence_parallel_moe:
128
+ hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
129
+ hidden_states = hidden_states[: positions.shape[0]]
130
# Recycle the post-final-norm hidden into the next draft step.
131
# compute_logits applies shared_head (== final norm) to the pre-norm
132
# element, so logits and the recycle each get exactly one final-norm.
vllm/models/deepseek_v32/nvidia/mtp.py
+10
-12
@@ -8,7 +8,10 @@ import torch.nn as nn
8
9
from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.config import VllmConfig
11
-from vllm.distributed import tensor_model_parallel_all_reduce
11
+from vllm.distributed import (
12
+ tensor_model_parallel_all_gather,
13
+ tensor_model_parallel_all_reduce,
14
+)
15
from vllm.model_executor.layers.fused_moe import (
16
fused_moe_make_expert_params_mapping,
17
)
@@ -24,10 +27,7 @@ from vllm.model_executor.model_loader.weight_utils import (
27
default_weight_loader,
28
maybe_remap_kv_scale_name,
29
)
27
-from vllm.model_executor.models.deepseek_mtp import (
28
- SharedHead,
29
- _restore_full_token_layout_if_needed,
30
-)
30
+from vllm.model_executor.models.deepseek_mtp import SharedHead
31
from vllm.model_executor.models.deepseek_v2 import (
32
DeepseekV2MixtureOfExperts,
33
DeepseekV2MoE,
@@ -95,13 +95,8 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
95
hidden_states, residual = self.mtp_block(
96
positions=positions, hidden_states=hidden_states, residual=None
97
)
98
- hidden_states, residual = _restore_full_token_layout_if_needed(
99
- hidden_states,
100
- residual,
101
- positions.shape[0],
102
- is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe,
103
- )
104
- if not self.mtp_block.use_sequence_parallel_moe:
98
+ is_sequence_parallel = self.mtp_block.use_sequence_parallel_moe
99
+ if not is_sequence_parallel:
100
# Without sequence parallelism, the MoE output is left un-reduced.
101
hidden_states = tensor_model_parallel_all_reduce(hidden_states)
102
# Recycle the POST-final-norm hidden into the next draft step. The
@@ -115,6 +110,9 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
110
# the legacy proposer (model_returns_tuple is True for the
111
# DeepSeekMTPModel architecture).
112
hidden_states, _ = self.shared_head.norm(hidden_states, residual)
113
+ if is_sequence_parallel:
114
+ hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
115
+ hidden_states = hidden_states[: positions.shape[0]]
116
return hidden_states, hidden_states
117
118