[Feature] Migrate moe sp support to non-torch compiled path for GLM5.2 (#47881)
Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Wentao Ye committed
Jul 15, 2026 at 19:33 UTC
4238b011a76629104b00ecf999d5390b94a1d289
4 files changed
+92
-16
vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py
+3
-1
@@ -122,7 +122,9 @@ class TrtLlmFp8ExpertsBase:
122
return (
123
not moe_parallel_config.use_all2all_kernels
124
or moe_parallel_config.use_ag_rs_all2all_kernels
125
- ) and not moe_parallel_config.enable_eplb
125
+ ) and not (
126
+ moe_parallel_config.enable_eplb or moe_parallel_config.is_sequence_parallel
127
+ )
128
129
130
class TrtLlmFp8ExpertsModular(TrtLlmFp8ExpertsBase, mk.FusedMoEExpertsModular):
vllm/model_executor/models/deepseek_mtp.py
+3
-1
@@ -45,9 +45,10 @@ def _restore_full_token_layout_if_needed(
45
hidden_states: torch.Tensor,
46
residual: torch.Tensor,
47
num_tokens: int,
48
+ is_sequence_parallel: bool = False,
49
) -> tuple[torch.Tensor, torch.Tensor]:
50
"""Restore full token rows for the MTP proposer after SP MoE layers."""
50
- if hidden_states.shape[0] == num_tokens:
51
+ if not is_sequence_parallel and hidden_states.shape[0] == num_tokens:
52
return hidden_states, residual
53
54
combined_states = torch.cat([hidden_states, residual], dim=-1)
@@ -142,6 +143,7 @@ class DeepSeekMultiTokenPredictorLayer(nn.Module):
143
hidden_states,
144
residual,
145
positions.shape[0],
146
+ is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe,
147
)
148
hidden_states = residual + hidden_states # pre-final-norm (logits hidden)
149
# Recycle the post-final-norm hidden into the next draft step.
vllm/models/deepseek_v32/nvidia/model.py
+82
-10
@@ -7,7 +7,11 @@ from itertools import islice
7
import torch
8
9
from vllm.config import VllmConfig
10
-from vllm.distributed import get_pp_group
10
+from vllm.distributed import (
11
+ get_pp_group,
12
+ tensor_model_parallel_all_gather,
13
+ tensor_model_parallel_reduce_scatter,
14
+)
15
from vllm.model_executor.layers.fused_moe import (
16
fused_moe_make_expert_params_mapping,
17
)
@@ -32,6 +36,7 @@ from vllm.model_executor.models.utils import (
36
is_pp_missing_parameter,
37
make_empty_intermediate_tensors_factory,
38
make_layers,
39
+ sequence_parallel_chunk,
40
)
41
from vllm.sequence import IntermediateTensors
42
@@ -39,6 +44,18 @@ from .attention import DeepseekV32Attention
44
from .fused_ops import fused_allreduce_rms_norm
45
46
47
+def _all_gather_sp_states(
48
+ hidden_states: torch.Tensor,
49
+ residual: torch.Tensor,
50
+ num_tokens: int,
51
+) -> tuple[torch.Tensor, torch.Tensor]:
52
+ # combine hidden_states and residual and all gather once
53
+ combined_states = torch.cat([hidden_states, residual], dim=-1)
54
+ combined_states = tensor_model_parallel_all_gather(combined_states, 0)[:num_tokens]
55
+ hidden_states, residual = combined_states.chunk(2, dim=-1)
56
+ return hidden_states, residual.contiguous()
57
+
58
+
59
class DeepseekV32DecoderLayer(torch.nn.Module):
60
def __init__(
61
self,
@@ -92,6 +109,12 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
109
prefix=f"{prefix}.mlp",
110
reduce_results=False,
111
)
112
+ self.use_sequence_parallel_moe = (
113
+ parallel_config.use_sequence_parallel_moe
114
+ and parallel_config.pipeline_parallel_size == 1
115
+ and isinstance(self.mlp, DeepseekV2MoE)
116
+ )
117
+ self.tp_size = parallel_config.tensor_parallel_size
118
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
119
self.post_attention_layernorm = RMSNorm(
120
config.hidden_size, eps=config.rms_norm_eps
@@ -104,25 +127,53 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
127
hidden_states: torch.Tensor,
128
residual: torch.Tensor | None,
129
) -> tuple[torch.Tensor, torch.Tensor]:
130
+ full_num_tokens = positions.shape[0]
131
+ input_is_sequence_parallel = (
132
+ self.use_sequence_parallel_moe
133
+ and residual is not None
134
+ and hidden_states.shape[0] != full_num_tokens
135
+ )
136
+
137
if residual is None:
138
# First layer: hidden_states is the (already reduced) embedding.
139
residual = hidden_states
140
hidden_states = self.input_layernorm(hidden_states)
141
+ elif input_is_sequence_parallel:
142
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
143
else:
144
# The previous layer's MLP/MoE output is left un-reduced; fuse its
145
# all-reduce into this input_layernorm.
146
hidden_states, residual = fused_allreduce_rms_norm(
147
hidden_states, residual, self.input_layernorm
148
)
117
- # self_attn's o_proj runs reduce_results=False; fuse its all-reduce with
118
- # the post-attention RMSNorm.
149
+ if input_is_sequence_parallel:
150
+ hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
151
+ hidden_states = hidden_states[:full_num_tokens]
152
+
153
+ # self_attn's o_proj runs reduce_results=False; reduce before RMSNorm.
154
hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
120
- hidden_states, residual = fused_allreduce_rms_norm(
121
- hidden_states, residual, self.post_attention_layernorm
122
- )
155
+ if self.use_sequence_parallel_moe:
156
+ # small trick using minus, eg. -17 % 8 = 7
157
+ sp_pad = (-hidden_states.shape[0]) % self.tp_size
158
+ # pad if not divisible by world size
159
+ hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad))
160
+ hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0)
161
+ if not input_is_sequence_parallel:
162
+ residual = sequence_parallel_chunk(residual)
163
+ hidden_states, residual = self.post_attention_layernorm(
164
+ hidden_states, residual
165
+ )
166
+ else:
167
+ hidden_states, residual = fused_allreduce_rms_norm(
168
+ hidden_states, residual, self.post_attention_layernorm
169
+ )
170
+
171
# MLP/MoE runs un-reduced; its all-reduce is fused into the next layer's
172
# input_layernorm (or the model's final norm).
125
- hidden_states = self.mlp(hidden_states)
173
+ if self.use_sequence_parallel_moe:
174
+ hidden_states = self.mlp(hidden_states, already_sequence_parallel=True)
175
+ else:
176
+ hidden_states = self.mlp(hidden_states)
177
return hidden_states, residual
178
179
@@ -206,12 +257,26 @@ class DeepseekV32Model(torch.nn.Module):
257
residual = intermediate_tensors["residual"]
258
259
aux_hidden_states = []
260
+ full_num_tokens = positions.shape[0]
261
for idx, layer in enumerate(
262
islice(self.layers, self.start_layer, self.end_layer),
263
start=self.start_layer,
264
):
265
+ if (
266
+ hidden_states.shape[0] != full_num_tokens
267
+ and not layer.use_sequence_parallel_moe
268
+ ):
269
+ hidden_states, residual = _all_gather_sp_states(
270
+ hidden_states, residual, full_num_tokens
271
+ )
272
if idx in self.aux_hidden_state_layers:
214
- aux_hidden_states.append(hidden_states + residual)
273
+ aux_hidden_state = hidden_states + residual
274
+ if aux_hidden_state.shape[0] != full_num_tokens:
275
+ aux_hidden_state = tensor_model_parallel_all_gather(
276
+ aux_hidden_state, 0
277
+ )
278
+ aux_hidden_state = aux_hidden_state[:full_num_tokens]
279
+ aux_hidden_states.append(aux_hidden_state)
280
hidden_states, residual = layer(positions, hidden_states, residual)
281
282
if not get_pp_group().is_last_rank:
@@ -219,8 +284,15 @@ class DeepseekV32Model(torch.nn.Module):
284
{"hidden_states": hidden_states, "residual": residual}
285
)
286
222
- # Last layer's MoE output is un-reduced; fuse its all-reduce into norm.
223
- hidden_states, _ = fused_allreduce_rms_norm(hidden_states, residual, self.norm)
287
+ if hidden_states.shape[0] != full_num_tokens:
288
+ hidden_states, residual = _all_gather_sp_states(
289
+ hidden_states, residual, full_num_tokens
290
+ )
291
+ hidden_states, _ = self.norm(hidden_states, residual)
292
+ else:
293
+ hidden_states, _ = fused_allreduce_rms_norm(
294
+ hidden_states, residual, self.norm
295
+ )
296
if len(aux_hidden_states) > 0:
297
return hidden_states, aux_hidden_states
298
return hidden_states
vllm/models/deepseek_v32/nvidia/mtp.py
+4
-4
@@ -96,11 +96,11 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
96
hidden_states,
97
residual,
98
positions.shape[0],
99
+ is_sequence_parallel=self.mtp_block.use_sequence_parallel_moe,
100
)
100
- # mtp_block's MoE output is left un-reduced (skip_final_all_reduce); the
101
- # main model fuses that all-reduce into the next norm, but here the
102
- # recycle hidden is consumed directly, so reduce it now.
103
- hidden_states = tensor_model_parallel_all_reduce(hidden_states)
101
+ if not self.mtp_block.use_sequence_parallel_moe:
102
+ # Without sequence parallelism, the MoE output is left un-reduced.
103
+ hidden_states = tensor_model_parallel_all_reduce(hidden_states)
104
# Recycle the POST-final-norm hidden into the next draft step. The
105
# residual-add is fused into the final RMSNorm so it is computed
106
# exactly once, and the result is returned for both tuple positions: