[Perf] Optimize DeepSeek V3.2 sequence parallelism (#51434)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Woosuk Kwon committed
Aug 7, 2026 at 16:51 UTC
e644c8cd8c734bf3b5e662a4bc363cfa8524d821
3 files changed
+222
-77
tests/models/deepseek_v32/test_sequence_parallel.py
new
+148
@@ -0,0 +1,148 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+from unittest.mock import Mock
6
+
7
+import torch
8
+from torch import nn
9
+
10
+from vllm.models.deepseek_v32.nvidia import model as deepseek_v32_model
11
+from vllm.models.deepseek_v32.nvidia import mtp as deepseek_v32_mtp
12
+
13
+
14
+class _IdentityNorm(nn.Module):
15
+ def __init__(self, hidden_size: int = 2) -> None:
16
+ super().__init__()
17
+ self.weight = nn.Parameter(torch.ones(hidden_size), requires_grad=False)
18
+ self.variance_epsilon = 1e-5
19
+
20
+ def forward(
21
+ self,
22
+ hidden_states: torch.Tensor,
23
+ residual: torch.Tensor | None = None,
24
+ ):
25
+ if residual is None:
26
+ return hidden_states
27
+ return hidden_states, residual
28
+
29
+
30
+class _RecordingModule(nn.Module):
31
+ def __init__(self) -> None:
32
+ super().__init__()
33
+ self.num_tokens = 0
34
+
35
+ def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor:
36
+ self.num_tokens = hidden_states.shape[0]
37
+ return hidden_states
38
+
39
+
40
+class _RecordingProjection(nn.Module):
41
+ def __init__(self) -> None:
42
+ super().__init__()
43
+ self.weight = nn.Parameter(torch.ones(2, 4), requires_grad=False)
44
+ self.num_tokens = 0
45
+
46
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
47
+ self.num_tokens = hidden_states.shape[0]
48
+ return hidden_states[:, :2]
49
+
50
+
51
+class _SequenceParallelMTPBlock:
52
+ use_sequence_parallel = True
53
+
54
+ def __call__(
55
+ self,
56
+ *,
57
+ positions: torch.Tensor,
58
+ hidden_states: torch.Tensor,
59
+ residual: torch.Tensor | None,
60
+ ):
61
+ assert residual is None
62
+ return hidden_states * 2, hidden_states * 3
63
+
64
+
65
+def _mock_sequence_parallel_collectives(monkeypatch, module):
66
+ monkeypatch.setattr(
67
+ module,
68
+ "sp_reduce_scatter",
69
+ lambda tensor: tensor.chunk(2, dim=0)[0],
70
+ raising=False,
71
+ )
72
+ monkeypatch.setattr(
73
+ module,
74
+ "sp_shard",
75
+ lambda tensor: torch.nn.functional.pad(tensor, (0, 0, 0, 1))[:2],
76
+ raising=False,
77
+ )
78
+ monkeypatch.setattr(
79
+ module,
80
+ "sp_all_gather",
81
+ lambda tensor: torch.cat([tensor, tensor], dim=0),
82
+ )
83
+
84
+
85
+def test_decoder_layer_keeps_dense_states_sequence_sharded(monkeypatch):
86
+ layer = object.__new__(deepseek_v32_model.DeepseekV32DecoderLayer)
87
+ nn.Module.__init__(layer)
88
+ layer.use_sequence_parallel = True
89
+ layer.input_layernorm = _IdentityNorm()
90
+ layer.post_attention_layernorm = _IdentityNorm()
91
+ layer.self_attn = _RecordingModule()
92
+ layer.mlp = _RecordingModule()
93
+
94
+ _mock_sequence_parallel_collectives(monkeypatch, deepseek_v32_model)
95
+
96
+ positions = torch.arange(3)
97
+ full_hidden_states = torch.arange(6, dtype=torch.float32).view(3, 2)
98
+ hidden_states = deepseek_v32_model.sp_shard(full_hidden_states)
99
+ hidden_states, residual = layer(positions, hidden_states, residual=None)
100
+
101
+ assert hidden_states.shape == residual.shape == (2, 2)
102
+ assert layer.self_attn.num_tokens == 3
103
+ assert layer.mlp.num_tokens == 2
104
+
105
+ hidden_states, residual = layer(positions, hidden_states, residual)
106
+
107
+ assert hidden_states.shape == residual.shape == (2, 2)
108
+ assert layer.self_attn.num_tokens == 3
109
+ assert layer.mlp.num_tokens == 2
110
+
111
+
112
+def test_mtp_projects_sequence_shard_and_restores_full_output(monkeypatch):
113
+ layer = object.__new__(deepseek_v32_mtp.DeepseekV32MultiTokenPredictorLayer)
114
+ nn.Module.__init__(layer)
115
+ layer.enorm = _IdentityNorm()
116
+ layer.hnorm = _IdentityNorm()
117
+ layer.eh_proj = _RecordingProjection()
118
+ layer._eh_plan = None
119
+ object.__setattr__(layer, "mtp_block", _SequenceParallelMTPBlock())
120
+ norm = Mock(
121
+ side_effect=lambda hidden_states, residual: (hidden_states + residual, None)
122
+ )
123
+ object.__setattr__(layer, "shared_head", SimpleNamespace(norm=norm))
124
+
125
+ monkeypatch.setattr(
126
+ deepseek_v32_mtp,
127
+ "fused_eh_norm",
128
+ lambda positions, inputs_embeds, previous_hidden_states, *args: torch.cat(
129
+ [inputs_embeds, previous_hidden_states], dim=-1
130
+ ),
131
+ )
132
+ monkeypatch.setattr(deepseek_v32_mtp, "run_glm52_plan", lambda *args: None)
133
+ _mock_sequence_parallel_collectives(monkeypatch, deepseek_v32_mtp)
134
+
135
+ inputs_embeds = torch.arange(6, dtype=torch.float32).view(3, 2)
136
+ hidden_states, recycled_hidden_states = layer(
137
+ input_ids=torch.zeros(3, dtype=torch.long),
138
+ positions=torch.arange(3),
139
+ previous_hidden_states=torch.zeros_like(inputs_embeds),
140
+ inputs_embeds=inputs_embeds,
141
+ )
142
+
143
+ sharded_states = torch.nn.functional.pad(inputs_embeds, (0, 0, 0, 1))[:2]
144
+ expected = torch.cat([sharded_states * 5, sharded_states * 5])[:3]
145
+ assert layer.eh_proj.num_tokens == 2
146
+ torch.testing.assert_close(hidden_states, expected)
147
+ torch.testing.assert_close(recycled_hidden_states, expected)
148
+ norm.assert_called_once()
vllm/models/deepseek_v32/nvidia/model.py
+57
-70
@@ -6,12 +6,10 @@ from itertools import islice
6
7
import torch
8
9
+import vllm.envs as envs
10
from vllm.config import VllmConfig
10
-from vllm.distributed import (
11
- get_pp_group,
12
- tensor_model_parallel_all_gather,
13
- tensor_model_parallel_reduce_scatter,
14
-)
11
+from vllm.distributed import get_pp_group
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,
15
)
@@ -36,27 +34,20 @@ from vllm.model_executor.models.utils import (
34
is_pp_missing_parameter,
35
make_empty_intermediate_tensors_factory,
36
make_layers,
39
- sequence_parallel_chunk,
37
)
38
from vllm.models.common.ops.fused_allreduce_rms_norm import fused_allreduce_rms_norm
39
+from vllm.models.common.ops.sequence_parallel import (
40
+ sp_all_gather,
41
+ sp_padding_mask,
42
+ sp_reduce_scatter,
43
+ sp_shard,
44
+)
45
from vllm.models.deepseek_v32.attention import DeepseekV32Attention
46
from vllm.sequence import IntermediateTensors
47
48
from .glm52_low_latency_gemm import enable_glm52_low_latency_gemm
49
50
48
-def _all_gather_sp_states(
49
- hidden_states: torch.Tensor,
50
- residual: torch.Tensor,
51
- num_tokens: int,
52
-) -> tuple[torch.Tensor, torch.Tensor]:
53
- # combine hidden_states and residual and all gather once
54
- combined_states = torch.cat([hidden_states, residual], dim=-1)
55
- combined_states = tensor_model_parallel_all_gather(combined_states, 0)[:num_tokens]
56
- hidden_states, residual = combined_states.chunk(2, dim=-1)
57
- return hidden_states, residual.contiguous()
58
-
59
-
51
class DeepseekV32DecoderLayer(torch.nn.Module):
52
def __init__(
53
self,
@@ -77,6 +68,10 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
68
layer_idx = int(prefix.split(sep=".")[-1])
69
self.layer_idx = layer_idx
70
self.use_mha = False
71
+ self.use_sequence_parallel = (
72
+ parallel_config.use_sequence_parallel_moe
73
+ and parallel_config.pipeline_parallel_size == 1
74
+ )
75
76
self.self_attn = DeepseekV32Attention(
77
vllm_config=vllm_config,
@@ -90,9 +85,8 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
85
and layer_idx >= config.first_k_dense_replace
86
and layer_idx % moe_layer_freq == 0
87
):
93
- # Defer the MoE cross-rank all-reduce; it is fused into the next
94
- # layer's input_layernorm (or the final norm) via
95
- # fused_allreduce_rms_norm.
88
+ # Keep the MoE output un-reduced. Non-SP fuses its all-reduce into
89
+ # the next layer norm; SP keeps the token shard local.
90
self.mlp = DeepseekV2MoE(
91
config=config,
92
parallel_config=parallel_config,
@@ -109,13 +103,8 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
103
quant_config=quant_config,
104
prefix=f"{prefix}.mlp",
105
reduce_results=False,
106
+ is_sequence_parallel=self.use_sequence_parallel,
107
)
113
- self.use_sequence_parallel_moe = (
114
- parallel_config.use_sequence_parallel_moe
115
- and parallel_config.pipeline_parallel_size == 1
116
- and isinstance(self.mlp, DeepseekV2MoE)
117
- )
118
- self.tp_size = parallel_config.tensor_parallel_size
108
self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
109
self.post_attention_layernorm = RMSNorm(
110
config.hidden_size, eps=config.rms_norm_eps
@@ -129,17 +118,11 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
118
residual: torch.Tensor | None,
119
) -> tuple[torch.Tensor, torch.Tensor]:
120
full_num_tokens = positions.shape[0]
132
- input_is_sequence_parallel = (
133
- self.use_sequence_parallel_moe
134
- and residual is not None
135
- and hidden_states.shape[0] != full_num_tokens
136
- )
121
122
if residual is None:
139
- # First layer: hidden_states is the (already reduced) embedding.
123
residual = hidden_states
124
hidden_states = self.input_layernorm(hidden_states)
142
- elif input_is_sequence_parallel:
125
+ elif self.use_sequence_parallel:
126
hidden_states, residual = self.input_layernorm(hidden_states, residual)
127
else:
128
# The previous layer's MLP/MoE output is left un-reduced; fuse its
@@ -147,20 +130,13 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
130
hidden_states, residual = fused_allreduce_rms_norm(
131
hidden_states, residual, self.input_layernorm
132
)
150
- if input_is_sequence_parallel:
151
- hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
152
- hidden_states = hidden_states[:full_num_tokens]
133
+ if self.use_sequence_parallel:
134
+ hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
135
136
# self_attn's o_proj runs reduce_results=False; reduce before RMSNorm.
137
hidden_states = self.self_attn(positions=positions, hidden_states=hidden_states)
156
- if self.use_sequence_parallel_moe:
157
- # small trick using minus, eg. -17 % 8 = 7
158
- sp_pad = (-hidden_states.shape[0]) % self.tp_size
159
- # pad if not divisible by world size
160
- hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad))
161
- hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0)
162
- if not input_is_sequence_parallel:
163
- residual = sequence_parallel_chunk(residual)
138
+ if self.use_sequence_parallel:
139
+ hidden_states = sp_reduce_scatter(hidden_states)
140
hidden_states, residual = self.post_attention_layernorm(
141
hidden_states, residual
142
)
@@ -169,9 +145,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module):
145
hidden_states, residual, self.post_attention_layernorm
146
)
147
172
- # MLP/MoE runs un-reduced; its all-reduce is fused into the next layer's
173
- # input_layernorm (or the model's final norm).
174
- if self.use_sequence_parallel_moe:
148
+ if self.use_sequence_parallel and isinstance(self.mlp, DeepseekV2MoE):
149
hidden_states = self.mlp(hidden_states, already_sequence_parallel=True)
150
else:
151
hidden_states = self.mlp(hidden_states)
@@ -192,6 +166,11 @@ class DeepseekV32Model(torch.nn.Module):
166
self.device = current_platform.device_type
167
168
self.vocab_size = config.vocab_size
169
+ parallel_config = vllm_config.parallel_config
170
+ self.use_sequence_parallel = (
171
+ parallel_config.use_sequence_parallel_moe
172
+ and parallel_config.pipeline_parallel_size == 1
173
+ )
174
# DSA is always sparse (has index_topk); allocate the shared top-k
175
# buffer the indexer writes and the sparse MLA backend reads.
176
self.is_v32 = True
@@ -231,9 +210,7 @@ class DeepseekV32Model(torch.nn.Module):
210
)
211
212
self.aux_hidden_state_layers = tuple[int, ...]()
234
- self.num_redundant_experts = (
235
- vllm_config.parallel_config.eplb_config.num_redundant_experts
236
- )
213
+ self.num_redundant_experts = parallel_config.eplb_config.num_redundant_experts
214
215
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
216
return self.embed_tokens(input_ids)
@@ -257,39 +234,49 @@ class DeepseekV32Model(torch.nn.Module):
234
hidden_states = intermediate_tensors["hidden_states"]
235
residual = intermediate_tensors["residual"]
236
260
- aux_hidden_states = []
237
full_num_tokens = positions.shape[0]
238
+ if self.use_sequence_parallel:
239
+ if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
240
+ forward_context = get_forward_context()
241
+ forward_context.is_padding = sp_padding_mask(
242
+ forward_context.is_padding, hidden_states
243
+ )
244
+ hidden_states = sp_shard(hidden_states)
245
+ assert residual is None, "Currently, SP is not supported with PP"
246
+
247
+ aux_hidden_states = []
248
for idx, layer in enumerate(
249
islice(self.layers, self.start_layer, self.end_layer),
250
start=self.start_layer,
251
):
266
- if (
267
- hidden_states.shape[0] != full_num_tokens
268
- and not layer.use_sequence_parallel_moe
269
- ):
270
- hidden_states, residual = _all_gather_sp_states(
271
- hidden_states, residual, full_num_tokens
272
- )
252
if idx in self.aux_hidden_state_layers:
274
- aux_hidden_state = hidden_states + residual
275
- if aux_hidden_state.shape[0] != full_num_tokens:
276
- aux_hidden_state = tensor_model_parallel_all_gather(
277
- aux_hidden_state, 0
278
- )
279
- aux_hidden_state = aux_hidden_state[:full_num_tokens]
280
- aux_hidden_states.append(aux_hidden_state)
253
+ aux_hidden_states.append(
254
+ hidden_states if residual is None else hidden_states + residual
255
+ )
256
hidden_states, residual = layer(positions, hidden_states, residual)
257
258
if not get_pp_group().is_last_rank:
259
+ assert not self.use_sequence_parallel, (
260
+ "Currently, SP is not supported with PP"
261
+ )
262
return IntermediateTensors(
263
{"hidden_states": hidden_states, "residual": residual}
264
)
265
288
- if hidden_states.shape[0] != full_num_tokens:
289
- hidden_states, residual = _all_gather_sp_states(
290
- hidden_states, residual, full_num_tokens
291
- )
266
+ if self.use_sequence_parallel:
267
hidden_states, _ = self.norm(hidden_states, residual)
268
+ if aux_hidden_states:
269
+ hidden_size = hidden_states.shape[-1]
270
+ packed_hidden_states = torch.cat(
271
+ [hidden_states, *aux_hidden_states], dim=-1
272
+ )
273
+ packed_hidden_states = sp_all_gather(packed_hidden_states)
274
+ packed_hidden_states = packed_hidden_states[:full_num_tokens]
275
+ hidden_states, *aux_hidden_states = packed_hidden_states.split(
276
+ hidden_size, dim=-1
277
+ )
278
+ else:
279
+ hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]
280
else:
281
hidden_states, _ = fused_allreduce_rms_norm(
282
hidden_states, residual, self.norm
vllm/models/deepseek_v32/nvidia/mtp.py
+17
-7
@@ -6,12 +6,11 @@ from collections.abc import Callable, Iterable
6
import torch
7
import torch.nn as nn
8
9
+import vllm.envs as envs
10
from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.config import VllmConfig
11
-from vllm.distributed import (
12
- tensor_model_parallel_all_gather,
13
- tensor_model_parallel_all_reduce,
14
-)
12
+from vllm.distributed import tensor_model_parallel_all_reduce
13
+from vllm.forward_context import get_forward_context, is_forward_context_available
14
from vllm.model_executor.layers.fused_moe import (
15
fused_moe_make_expert_params_mapping,
16
)
@@ -38,6 +37,11 @@ from vllm.model_executor.models.utils import (
37
get_pp_missing_layer_names,
38
maybe_prefix,
39
)
40
+from vllm.models.common.ops.sequence_parallel import (
41
+ sp_all_gather,
42
+ sp_padding_mask,
43
+ sp_shard,
44
+)
45
from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
46
from vllm.platforms import current_platform
47
from vllm.sequence import IntermediateTensors
@@ -101,13 +105,20 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
105
self.hnorm.weight,
106
self.enorm.variance_epsilon,
107
)
108
+ is_sequence_parallel = self.mtp_block.use_sequence_parallel
109
+ if is_sequence_parallel:
110
+ if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
111
+ forward_context = get_forward_context()
112
+ forward_context.is_padding = sp_padding_mask(
113
+ forward_context.is_padding, eh_input
114
+ )
115
+ eh_input = sp_shard(eh_input)
116
hidden_states = run_glm52_plan(self._eh_plan, eh_input, self.eh_proj.weight)
117
if hidden_states is None:
118
hidden_states = self.eh_proj(eh_input)
119
hidden_states, residual = self.mtp_block(
120
positions=positions, hidden_states=hidden_states, residual=None
121
)
110
- is_sequence_parallel = self.mtp_block.use_sequence_parallel_moe
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)
@@ -123,8 +134,7 @@ class DeepseekV32MultiTokenPredictorLayer(nn.Module):
134
# DeepSeekMTPModel architecture).
135
hidden_states, _ = self.shared_head.norm(hidden_states, residual)
136
if is_sequence_parallel:
126
- hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
127
- hidden_states = hidden_states[: positions.shape[0]]
137
+ hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]]
138
return hidden_states, hidden_states
139
140