@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 38a466e7b6

[DSV4] Implement Sequence Parallelism (#46789)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

Woosuk Kwon committed Aug 1, 2026 at 10:22 UTC 38a466e7b6e087d67c35e7f924c04c245423c99f

7 files changed +138 -29

tests/models/kimi_k3/test_sequence_parallel.py

+19 -1

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import math

5 from types import MethodType, SimpleNamespace

6 from unittest.mock import Mock

@@ -9,9 +10,9 @@ import torch

10 from torch import nn

12 from vllm.config import ParallelConfig

13 +from vllm.models.common.ops import sequence_parallel as sp_ops

14 from vllm.models.kimi_k3.nvidia import model as kimi_model

15 from vllm.models.kimi_k3.nvidia import mtp as kimi_mtp

14 -from vllm.models.kimi_k3.nvidia.ops import sequence_parallel as sp_ops

16 from vllm.platforms import current_platform

@@ -316,6 +317,23 @@ def test_sp_reduce_scatter_uses_custom_kernel_after_padding(monkeypatch):

317 fallback.assert_not_called()

318

319

320 +@pytest.mark.parametrize("shape", [(3,), (3, 2, 2)])

321 +def test_sp_shard_pads_only_the_token_axis(monkeypatch, shape):

322 + hidden_states = torch.arange(math.prod(shape), dtype=torch.float32).view(shape)

323 + monkeypatch.setattr(

324 + sp_ops,

325 + "get_tensor_model_parallel_world_size",

326 + lambda: 2,

327 + )

328 + monkeypatch.setattr(sp_ops, "get_tensor_model_parallel_rank", lambda: 1)

329 +

330 + output = sp_ops.sp_shard(hidden_states)

331 +

332 + padding = hidden_states.new_zeros((1, *shape[1:]))

333 + expected = torch.cat([hidden_states, padding])[2:]

334 + torch.testing.assert_close(output, expected)

335 +

336 +

337 def test_sp_collectives_fall_back_without_custom_kernel(monkeypatch):

338 hidden_states = torch.arange(4, dtype=torch.float32).view(2, 2)

339 monkeypatch.setattr(

vllm/models/common/ops/sequence_parallel.py renamed

+4 -9

@@ -12,17 +12,12 @@ from vllm.distributed import (

12 )

15 -def _custom_collective(

16 - name: str,

17 - x: torch.Tensor,

18 -) -> torch.Tensor | None:

15 +def _custom_collective(name: str, x: torch.Tensor) -> torch.Tensor | None:

16 device_communicator = get_tp_group().device_communicator

17 if device_communicator is None:

18 return None

19 collective = getattr(device_communicator, name, None)

23 - if collective is None:

24 - return None

25 - return collective(x)

20 + return None if collective is None else collective(x)

23 def sp_all_gather(x: torch.Tensor) -> torch.Tensor:

@@ -45,12 +40,12 @@ def sp_reduce_scatter(x: torch.Tensor) -> torch.Tensor:

42 def sp_shard(x: torch.Tensor) -> torch.Tensor:

48 - assert x.ndim == 2

43 tp_size = get_tensor_model_parallel_world_size()

44 tp_rank = get_tensor_model_parallel_rank()

45 sp_pad = (-x.shape[0]) % tp_size

46 if sp_pad > 0:

53 - x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad))

47 + pad = (0, 0) * (x.ndim - 1) + (0, sp_pad)

48 + x = torch.nn.functional.pad(x, pad)

49 chunk = x.shape[0] // tp_size

50 return x[tp_rank * chunk : (tp_rank + 1) * chunk]

vllm/models/deepseek_v4/nvidia/dspark.py

+30 -4

@@ -15,11 +15,13 @@ import regex as re

15 import torch

16 import torch.nn as nn

18 +import vllm.envs as envs

19 from vllm.config import VllmConfig, get_current_vllm_config

20 from vllm.distributed import (

21 get_tensor_model_parallel_rank,

22 get_tensor_model_parallel_world_size,

23 )

24 +from vllm.forward_context import get_forward_context, is_forward_context_available

25 from vllm.logger import init_logger

26 from vllm.model_executor.kernels.mhc.tilelang import (

27 hc_head_fused_kernel_tilelang,

@@ -40,10 +42,16 @@ from vllm.model_executor.models.qwen3_dspark import (

42 DSparkMarkovHead,

43 )

44 from vllm.model_executor.models.utils import maybe_prefix

45 +from vllm.models.common.ops.sequence_parallel import (

46 + sp_all_gather,

47 + sp_padding_mask,

48 + sp_shard,

49 +)

51 from .model import (

52 DeepseekV4DecoderLayer,

53 DeepseekV4Model,

54 + _use_sequence_parallel,

55 make_deepseek_v4_expert_params_mapping,

56 )

@@ -66,6 +74,7 @@ class DSparkDeepseekV4Model(nn.Module):

74 self.rms_norm_eps = config.rms_norm_eps

75 self.num_hidden_layers = config.num_hidden_layers

76 self.target_layer_ids = tuple(config.dspark_target_layer_ids)

77 + self.use_sequence_parallel = _use_sequence_parallel(vllm_config)

79 self.num_dspark_layers = getattr(config, "n_mtp_layers", None) or 3

@@ -86,12 +95,19 @@ class DSparkDeepseekV4Model(nn.Module):

95 )

96 self.main_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

98 + self.topk_indices_buffer = torch.empty(

99 + vllm_config.scheduler_config.max_num_batched_tokens,

100 + config.index_topk,

101 + dtype=torch.int32,

102 + )

103 +

104 current_vllm_config = get_current_vllm_config()

105 self.layers = nn.ModuleList(

106 [

107 DeepseekV4DecoderLayer(

108 current_vllm_config,

109 prefix=maybe_prefix(prefix, f"layers.{self.num_hidden_layers + i}"),

110 + topk_indices_buffer=self.topk_indices_buffer,

111 )

112 for i in range(self.num_dspark_layers)

113 ]

@@ -172,6 +188,15 @@ class DSparkDeepseekV4Model(nn.Module):

188 ) -> torch.Tensor:

189 if inputs_embeds is None:

190 inputs_embeds = self.embed_input_ids(input_ids)

191 + full_num_tokens = positions.shape[0]

192 + if self.use_sequence_parallel:

193 + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():

194 + forward_context = get_forward_context()

195 + forward_context.is_padding = sp_padding_mask(

196 + forward_context.is_padding, inputs_embeds

197 + )

198 + inputs_embeds = sp_shard(inputs_embeds)

199 + input_ids = sp_shard(input_ids)

200 # Expand to hc_mult copies for hyper-connections ([T, H] -> [T, hc, H]).

201 hidden_states = inputs_embeds.unsqueeze(-2).repeat(1, self.hc_mult, 1)

202

@@ -186,6 +211,8 @@ class DSparkDeepseekV4Model(nn.Module):

211 residual,

212 )

213 hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)

214 + if self.use_sequence_parallel:

215 + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]

216 # hc_head reduces the hc copies; return the PRE-norm head hidden

217 hidden_states = hc_head_fused_kernel_tilelang(

218 hidden_states,

@@ -279,10 +306,9 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):

306 self.draft_model_config = vllm_config.speculative_config.draft_model_config

307 self.config = self.draft_model_config.hf_config

308 self.quant_config = vllm_config.quant_config

282 - self.pad_shared_expert = (

283 - getattr(self.quant_config, "weight_block_size", None) is not None

284 - and not vllm_config.parallel_config.use_sequence_parallel_moe

285 - )

309 + self.pad_shared_expert = getattr(

310 + self.quant_config, "weight_block_size", None

311 + ) is not None and not _use_sequence_parallel(vllm_config)

312 self.model = DSparkDeepseekV4Model(

313 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

314 )

vllm/models/deepseek_v4/nvidia/model.py

+53 -4

@@ -65,6 +65,12 @@ from vllm.model_executor.models.utils import (

65 maybe_prefix,

66 )

67 from vllm.model_executor.utils import set_weight_attrs

68 +from vllm.models.common.ops.sequence_parallel import (

69 + sp_all_gather,

70 + sp_padding_mask,

71 + sp_reduce_scatter,

72 + sp_shard,

73 +)

74 from vllm.models.deepseek_v4.attention import DeepseekV4Attention

75 from vllm.models.deepseek_v4.eager_scratch import DeepseekV4EagerScratchPool

76 from vllm.models.deepseek_v4.nvidia.flashinfer_sparse import (

@@ -515,6 +521,7 @@ class DeepseekV4MoE(nn.Module):

521 self,

522 vllm_config: VllmConfig,

523 prefix: str = "",

524 + use_sequence_parallel: bool = False,

525 ):

526 super().__init__()

527

@@ -522,6 +529,7 @@ class DeepseekV4MoE(nn.Module):

529 config = vllm_config.model_config.hf_config

530 quant_config = vllm_config.quant_config

531 self.prefix = prefix

532 + self.use_sequence_parallel = use_sequence_parallel

533 self.use_mega_moe = (

534 vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"

535 )

@@ -595,6 +603,7 @@ class DeepseekV4MoE(nn.Module):

603 swiglu_limit=self.swiglu_limit,

604 quant_config=quant_config,

605 reduce_results=self.use_mega_moe,

606 + is_sequence_parallel=use_sequence_parallel,

607 prefix=f"{prefix}.shared_experts",

608 )

609

@@ -689,6 +698,7 @@ class DeepseekV4MoE(nn.Module):

698 router_logits_dtype=torch.float32,

699 enable_eplb=parallel_config.enable_eplb,

700 num_redundant_experts=eplb_config.num_redundant_experts,

701 + is_sequence_parallel=self.use_sequence_parallel,

702 )

703

704 def forward(

@@ -792,6 +802,17 @@ def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:

802 return DeepseekV4FlashMLAAttention

803

804

805 +def _use_sequence_parallel(vllm_config: VllmConfig) -> bool:

806 + parallel_config = vllm_config.parallel_config

807 + use_mega_moe = vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"

808 + return (

809 + parallel_config.pipeline_parallel_size == 1

810 + and parallel_config.enable_expert_parallel

811 + and parallel_config.tensor_parallel_size > 1

812 + and (use_mega_moe or parallel_config.data_parallel_size > 1)

813 + )

814 +

815 +

816 class DeepseekV4DecoderLayer(nn.Module):

817 def __init__(

818 self,

@@ -805,6 +826,7 @@ class DeepseekV4DecoderLayer(nn.Module):

826

827 config = vllm_config.model_config.hf_config

828 self.hidden_size = config.hidden_size

829 + self.use_sequence_parallel = _use_sequence_parallel(vllm_config)

830

831 self.rms_norm_eps = config.rms_norm_eps

832 self.attn = _select_dsv4_attn_cls(vllm_config)(

@@ -814,7 +836,13 @@ class DeepseekV4DecoderLayer(nn.Module):

836 aux_stream_list=aux_stream_list,

837 eager_scratch_pool=eager_scratch_pool,

838 )

817 - self.ffn = DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn")

839 + if self.use_sequence_parallel:

840 + self.attn.wo_b.reduce_results = False

841 + self.ffn = DeepseekV4MoE(

842 + vllm_config,

843 + prefix=f"{prefix}.ffn",

844 + use_sequence_parallel=self.use_sequence_parallel,

845 + )

846

847 self.attn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)

848 self.ffn_norm = RMSNorm(self.hidden_size, self.rms_norm_eps)

@@ -932,8 +960,12 @@ class DeepseekV4DecoderLayer(nn.Module):

960 norm_eps=attn_norm_eps,

961 )

962

935 - # attn_norm is fused into mhc_pre_tilelang / mhc_fused_post_pre above.

963 + if self.use_sequence_parallel:

964 + x = sp_all_gather(x)[: positions.shape[0]]

965 +

966 x = self.attn(positions, x, None)

967 + if self.use_sequence_parallel:

968 + x = sp_reduce_scatter(x)

969

970 ffn_norm_weight = self.ffn_norm.weight.data

971 ffn_norm_eps = self.ffn_norm.variance_epsilon

@@ -972,6 +1004,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

1004 self.use_mega_moe = (

1005 vllm_config.kernel_config.moe_backend == "deep_gemm_mega_moe"

1006 )

1007 + self.use_sequence_parallel = _use_sequence_parallel(vllm_config)

1008 if self.use_mega_moe and not vllm_config.parallel_config.enable_expert_parallel:

1009 raise NotImplementedError(

1010 "DeepSeek V4 MegaMoE currently requires expert parallel. "

@@ -1114,6 +1147,16 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

1147 if self.use_mega_moe:

1148 input_ids = input_ids.to(torch.int64)

1149

1150 + full_num_tokens = positions.shape[0]

1151 + if self.use_sequence_parallel:

1152 + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():

1153 + forward_context = get_forward_context()

1154 + forward_context.is_padding = sp_padding_mask(

1155 + forward_context.is_padding, hidden_states

1156 + )

1157 + hidden_states = sp_shard(hidden_states)

1158 + input_ids = sp_shard(input_ids)

1159 +

1160 residual, post_mix, res_mix = None, None, None

1161 aux_hidden_states: list[torch.Tensor] = []

1162 final_aux_recon: torch.Tensor | None = None # avoid duplicate mhc_post call

@@ -1134,7 +1177,10 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

1177 aux_recon = mhc_post_tilelang(

1178 hidden_states, residual, post_mix, res_mix

1179 )

1137 - aux_hidden_states.append(aux_recon.mean(dim=1))

1180 + aux_hidden_state = aux_recon.mean(dim=1)

1181 + if self.use_sequence_parallel:

1182 + aux_hidden_state = sp_all_gather(aux_hidden_state)[:full_num_tokens]

1183 + aux_hidden_states.append(aux_hidden_state)

1184 final_aux_recon = aux_recon

1185 if layer is not None:

1186 # Reuse if the last layer was captured as an aux hidden state

@@ -1148,6 +1194,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

1194 if not get_pp_group().is_last_rank:

1195 return IntermediateTensors({"hidden_states": hidden_states})

1196

1197 + if self.use_sequence_parallel:

1198 + hidden_states = sp_all_gather(hidden_states)[:full_num_tokens]

1199 +

1200 if self._mtp_hidden_buffer is not None:

1201 num_tokens = hidden_states.shape[0]

1202 self._mtp_hidden_buffer[:num_tokens].copy_(hidden_states.flatten(1))

@@ -1194,7 +1243,7 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):

1243 # ranks land on the zero pad). SP / unquantized ones need no padding.

1244 pad_shared_expert = (

1245 getattr(self.quant_config, "weight_block_size", None) is not None

1197 - and not self.parallel_config.use_sequence_parallel_moe

1246 + and not self.use_sequence_parallel

1247 )

1248

1249 for name, loaded_weight in weights:

vllm/models/deepseek_v4/nvidia/mtp.py

+21 -4

@@ -18,11 +18,13 @@ import regex as re

18 import torch

19 import torch.nn as nn

21 +import vllm.envs as envs

22 from vllm.config import VllmConfig

23 from vllm.distributed import (

24 get_tensor_model_parallel_rank,

25 get_tensor_model_parallel_world_size,

26 )

27 +from vllm.forward_context import get_forward_context, is_forward_context_available

28 from vllm.logger import init_logger

29 from vllm.model_executor.kernels.mhc.tilelang import (

30 hc_head_fused_kernel_tilelang,

@@ -44,6 +46,11 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader

46 from vllm.model_executor.models.deepseek_mtp import SharedHead

47 from vllm.model_executor.models.deepseek_v2 import get_spec_layer_idx_from_weight_name

48 from vllm.model_executor.models.utils import maybe_prefix

49 +from vllm.models.common.ops.sequence_parallel import (

50 + sp_all_gather,

51 + sp_padding_mask,

52 + sp_shard,

53 +)

54 from vllm.models.deepseek_v4.common.ops import (

55 fused_mtp_input_rmsnorm,

56 mtp_shared_head_rmsnorm,

@@ -53,6 +60,7 @@ from vllm.sequence import IntermediateTensors

60 from .model import (

61 DeepseekV4DecoderLayer,

62 DeepseekV4Model,

63 + _use_sequence_parallel,

64 make_deepseek_v4_expert_params_mapping,

65 )

@@ -158,6 +166,14 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):

166 self.enorm.variance_epsilon,

167 self.hc_mult,

168 )

169 + if self.mtp_block.use_sequence_parallel:

170 + if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():

171 + forward_context = get_forward_context()

172 + forward_context.is_padding = sp_padding_mask(

173 + forward_context.is_padding, inputs_embeds

174 + )

175 + inputs_embeds = sp_shard(inputs_embeds)

176 + previous_hidden_states = sp_shard(previous_hidden_states)

177 hidden_states = self.h_proj(previous_hidden_states) + self.e_proj(

178 inputs_embeds

179 ).unsqueeze(-2)

@@ -165,6 +181,8 @@ class DeepSeekV4MultiTokenPredictorLayer(nn.Module):

181 positions=positions, x=hidden_states, input_ids=None

182 )

183 hidden_states = mhc_post_tilelang(hidden_states, residual, post_mix, res_mix)

184 + if self.mtp_block.use_sequence_parallel:

185 + hidden_states = sp_all_gather(hidden_states)[: positions.shape[0]]

186 # Return the flat pre-hc_head residual so it can be re-fed as the

187 # next spec step's `previous_hidden_states` when

188 # num_speculative_tokens > 1. hc_head is deferred to compute_logits.

@@ -266,10 +284,9 @@ class DeepSeekV4MTP(nn.Module):

284 super().__init__()

285 self.config = vllm_config.model_config.hf_config

286 self.quant_config = vllm_config.quant_config

269 - self.pad_shared_expert = (

270 - getattr(self.quant_config, "weight_block_size", None) is not None

271 - and not vllm_config.parallel_config.use_sequence_parallel_moe

272 - )

287 + self.pad_shared_expert = getattr(

288 + self.quant_config, "weight_block_size", None

289 + ) is not None and not _use_sequence_parallel(vllm_config)

290 self.model = DeepSeekV4MultiTokenPredictor(

291 vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

292 )

vllm/models/kimi_k3/nvidia/model.py

+6 -6

@@ -88,6 +88,12 @@ from vllm.model_executor.models.utils import (

88 maybe_prefix,

89 )

90 from vllm.model_executor.models.vision import is_vit_use_data_parallel

91 +from vllm.models.common.ops.sequence_parallel import (

92 + sp_all_gather,

93 + sp_padding_mask,

94 + sp_reduce_scatter,

95 + sp_shard,

96 +)

97 from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts

98 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs

99 from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention

@@ -96,12 +102,6 @@ from vllm.models.kimi_k3.nvidia.low_latency_gemm import (

102 )

103 from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention

104 from vllm.models.kimi_k3.nvidia.ops import attn_res

99 -from vllm.models.kimi_k3.nvidia.ops.sequence_parallel import (

100 - sp_all_gather,

101 - sp_padding_mask,

102 - sp_reduce_scatter,

103 - sp_shard,

104 -)

105 from vllm.multimodal import MULTIMODAL_REGISTRY

106 from vllm.multimodal.inputs import NestedTensors

107 from vllm.platforms import current_platform

vllm/models/kimi_k3/nvidia/mtp.py

+5 -1

@@ -27,6 +27,11 @@ from vllm.model_executor.model_loader.weight_utils import (

27 maybe_remap_kv_scale_name,

28 )

29 from vllm.model_executor.models.utils import get_pp_missing_layer_names, maybe_prefix

30 +from vllm.models.common.ops.sequence_parallel import (

31 + sp_all_gather,

32 + sp_padding_mask,

33 + sp_shard,

34 +)

35 from vllm.sequence import IntermediateTensors

36 from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig

@@ -38,7 +43,6 @@ from .model import (

43 get_spec_layer_idx_from_weight_name,

44 make_kimi_k3_mega_moe_expert_params_mapping,

45 )

41 -from .ops.sequence_parallel import sp_all_gather, sp_padding_mask, sp_shard

47 logger = init_logger(__name__)