@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c05d75aaa9

[Bugfix][Refactor] Keep Qwen3Next layer boundaries sequence parallel (#50685)

Signed-off-by: kunzh <RimeK20@outlook.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com>

kzwrime committed Aug 14, 2026 at 11:58 UTC c05d75aaa95cf89f547503044c1921625905085d

4 files changed +75 -111

vllm/model_executor/models/interns2_mobius.py

+22 -35

@@ -43,7 +43,7 @@ from .qwen3_5 import (

43 Qwen3_5RMSNorm,

44 )

45 from .qwen3_5_mtp import Qwen3_5MoeMTP

46 -from .qwen3_next import Qwen3NextAttention, _all_gather_hidden_and_residual

46 +from .qwen3_next import Qwen3NextAttention

47 from .qwen3_vl import (

48 Qwen3_VisionTransformer,

49 Qwen3VLDummyInputsBuilder,

@@ -233,11 +233,6 @@ class InternS2MobiusDecoderLayer(nn.Module):

233 meta_mlp: nn.ModuleList,

234 ) -> tuple[torch.Tensor, torch.Tensor]:

235 full_num_tokens = positions.shape[-1]

236 - input_is_sequence_parallel = (

237 - self.use_attn_reduce_scatter_for_moe

238 - and residual is not None

239 - and hidden_states.shape[0] != full_num_tokens

240 - )

236

237 if residual is None:

238 residual = hidden_states

@@ -245,7 +240,7 @@ class InternS2MobiusDecoderLayer(nn.Module):

240 else:

241 hidden_states, residual = self.input_layernorm(hidden_states, residual)

242

248 - if input_is_sequence_parallel:

243 + if self.use_attn_reduce_scatter_for_moe:

244 hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

245 hidden_states = hidden_states[:full_num_tokens]

246

@@ -266,8 +261,6 @@ class InternS2MobiusDecoderLayer(nn.Module):

261 sp_pad = (-hidden_states.shape[0]) % tp_world_size

262 hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad))

263 hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0)

269 - if not input_is_sequence_parallel:

270 - residual = sequence_parallel_chunk(residual)

264

265 hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)

266 block_idx = self.layer_idx % self.num_blocks

@@ -343,6 +336,10 @@ class InternS2MobiusModel(nn.Module, EagleModelMixin):

336 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

337 return self.embed_tokens(input_ids)

338

339 + @property

340 + def use_sequence_parallel(self) -> bool:

341 + return self.layers[self.start_layer].use_attn_reduce_scatter_for_moe

342 +

343 def forward(

344 self,

345 input_ids: torch.Tensor | None,

@@ -363,36 +360,21 @@ class InternS2MobiusModel(nn.Module, EagleModelMixin):

360 residual = intermediate_tensors["residual"]

361

362 full_num_tokens = positions.shape[-1]

363 + if self.use_sequence_parallel:

364 + hidden_states = sequence_parallel_chunk(hidden_states)

365 + assert residual is None

366 +

367 aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

368 for layer_idx, layer in enumerate(

369 islice(self.layers, self.start_layer, self.end_layer),

370 start=self.start_layer,

371 ):

371 - if (

372 - hidden_states.shape[0] != full_num_tokens

373 - and not layer.use_attn_reduce_scatter_for_moe

374 - ):

375 - hidden_states, residual = _all_gather_hidden_and_residual(

376 - hidden_states,

377 - residual,

378 - full_num_tokens,

379 - self.config.hidden_size,

380 - )

372 hidden_states, residual = layer(

373 positions=positions,

374 hidden_states=hidden_states,

375 residual=residual,

376 meta_mlp=self.meta_mlp,

377 )

387 - if (layer_idx + 1) in self.aux_hidden_state_layers and hidden_states.shape[

388 - 0

389 - ] != full_num_tokens:

390 - hidden_states, residual = _all_gather_hidden_and_residual(

391 - hidden_states,

392 - residual,

393 - full_num_tokens,

394 - self.config.hidden_size,

395 - )

378 self._maybe_add_hidden_state(

379 aux_hidden_states,

380 layer_idx + 1,

@@ -404,14 +386,19 @@ class InternS2MobiusModel(nn.Module, EagleModelMixin):

386 return IntermediateTensors(

387 {"hidden_states": hidden_states, "residual": residual}

388 )

407 - if hidden_states.shape[0] != full_num_tokens:

408 - hidden_states, residual = _all_gather_hidden_and_residual(

409 - hidden_states,

410 - residual,

411 - full_num_tokens,

412 - self.config.hidden_size,

413 - )

389 hidden_states, _ = self.norm(hidden_states, residual)

390 + if self.use_sequence_parallel:

391 + if aux_hidden_states:

392 + hidden_size = hidden_states.shape[-1]

393 + hidden_states = torch.cat([hidden_states, *aux_hidden_states], dim=-1)

394 + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

395 + hidden_states = hidden_states[:full_num_tokens]

396 + hidden_states, *aux_hidden_states = hidden_states.split(

397 + hidden_size, dim=-1

398 + )

399 + else:

400 + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

401 + hidden_states = hidden_states[:full_num_tokens]

402 if aux_hidden_states:

403 return hidden_states, aux_hidden_states

404 return hidden_states

vllm/model_executor/models/qwen3_5_mtp.py

+9 -10

@@ -10,7 +10,7 @@ from torch import nn

10 from vllm._aiter_ops import rocm_aiter_ops

11 from vllm.compilation.decorators import support_torch_compile

12 from vllm.config import VllmConfig, get_current_vllm_config

13 -from vllm.distributed.parallel_state import get_pp_group

13 +from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather

14 from vllm.logger import init_logger

15 from vllm.model_executor.layers.linear import ColumnParallelLinear

16 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -26,9 +26,9 @@ from vllm.model_executor.models.qwen3_5 import (

26 )

27 from vllm.model_executor.models.qwen3_next import (

28 QwenNextMixtureOfExperts,

29 - _all_gather_hidden_and_residual,

29 _is_shared_expert_fse_compatible,

30 )

31 +from vllm.model_executor.models.utils import sequence_parallel_chunk

32 from vllm.sequence import IntermediateTensors

33 from vllm.transformers_utils.configs.qwen3_5 import Qwen3_5TextConfig

34 from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeTextConfig

@@ -161,6 +161,10 @@ class Qwen3_5MultiTokenPredictor(nn.Module):

161

162 current_step_idx = spec_step_idx % self.num_mtp_layers

163 mtp_layer = self.layers[current_step_idx]

164 + if mtp_layer.use_attn_reduce_scatter_for_moe:

165 + assert hidden_states.shape[0] == positions.shape[-1]

166 + hidden_states = sequence_parallel_chunk(hidden_states)

167 + assert residual is None

168 hidden_states, residual = mtp_layer(

169 positions=positions,

170 hidden_states=hidden_states,

@@ -172,15 +176,10 @@ class Qwen3_5MultiTokenPredictor(nn.Module):

176 {"hidden_states": hidden_states, "residual": residual}

177 )

178

175 - if mtp_layer.use_attn_reduce_scatter_for_moe:

176 - hidden_states, residual = _all_gather_hidden_and_residual(

177 - hidden_states,

178 - residual,

179 - positions.shape[-1],

180 - self.config.hidden_size,

181 - )

182 -

179 hidden_states, _ = self.norm(hidden_states, residual)

180 + if mtp_layer.use_attn_reduce_scatter_for_moe:

181 + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

182 + hidden_states = hidden_states[: positions.shape[-1]]

183 return hidden_states

184

185 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

vllm/model_executor/models/qwen3_next.py

+35 -57

@@ -78,6 +78,18 @@ logger = init_logger(__name__)

78 KVCache = tuple[torch.Tensor, torch.Tensor]

81 +def _should_use_sequence_parallel(vllm_config: VllmConfig) -> bool:

82 + config = vllm_config.model_config.hf_text_config

83 + parallel_config = vllm_config.parallel_config

84 + return (

85 + parallel_config.use_sequence_parallel_moe

86 + and parallel_config.pipeline_parallel_size == 1

87 + and getattr(config, "num_experts", 0) > 0

88 + and not getattr(config, "mlp_only_layers", [])

89 + and getattr(config, "decoder_sparse_step", 1) == 1

90 + )

91 +

92 +

93 def _is_shared_expert_fse_compatible(quant_config) -> bool:

94 """Check if shared expert can be fused with routed experts.

@@ -397,7 +409,6 @@ class Qwen3NextDecoderLayer(nn.Module):

409 model_config = vllm_config.model_config

410 cache_config = vllm_config.cache_config

411 quant_config = vllm_config.quant_config

400 - parallel_config = vllm_config.parallel_config

412

413 self.layer_type = layer_type

414 self.layer_idx = extract_layer_index(prefix)

@@ -409,10 +420,8 @@ class Qwen3NextDecoderLayer(nn.Module):

420 config.num_experts > 0

421 and (self.layer_idx + 1) % config.decoder_sparse_step == 0

422 )

412 - self.use_attn_reduce_scatter_for_moe = (

413 - parallel_config.use_sequence_parallel_moe

414 - and parallel_config.pipeline_parallel_size == 1

415 - and is_moe_layer

423 + self.use_attn_reduce_scatter_for_moe = _should_use_sequence_parallel(

424 + vllm_config

425 )

426

427 if self.layer_type == "linear_attention":

@@ -481,11 +490,6 @@ class Qwen3NextDecoderLayer(nn.Module):

490 **kwargs: object,

491 ):

492 full_num_tokens = positions.shape[-1]

484 - input_is_sequence_parallel = (

485 - self.use_attn_reduce_scatter_for_moe

486 - and residual is not None

487 - and hidden_states.shape[0] != full_num_tokens

488 - )

493

494 if residual is None:

495 residual = hidden_states

@@ -493,7 +497,7 @@ class Qwen3NextDecoderLayer(nn.Module):

497 else:

498 hidden_states, residual = self.input_layernorm(hidden_states, residual)

499

496 - if input_is_sequence_parallel:

500 + if self.use_attn_reduce_scatter_for_moe:

501 hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

502 hidden_states = hidden_states[:full_num_tokens]

503

@@ -524,8 +528,6 @@ class Qwen3NextDecoderLayer(nn.Module):

528 # pad if not divisible by world size

529 hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, sp_pad))

530 hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0)

527 - if not input_is_sequence_parallel:

528 - residual = sequence_parallel_chunk(residual)

531

532 # Fully Connected

533 hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)

@@ -554,24 +556,6 @@ class Qwen3NextDecoderLayer(nn.Module):

556 return hidden_states, residual

557

558

557 -def _all_gather_hidden_and_residual(

558 - hidden_states: torch.Tensor,

559 - residual: torch.Tensor | None,

560 - full_num_tokens: int,

561 - hidden_size: int,

562 -) -> tuple[torch.Tensor, torch.Tensor | None]:

563 - if residual is None:

564 - hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

565 - hidden_states = hidden_states[:full_num_tokens]

566 - return hidden_states, None

567 -

568 - combined_states = torch.cat([hidden_states, residual], dim=-1)

569 - combined_states = tensor_model_parallel_all_gather(combined_states, 0)

570 - combined_states = combined_states[:full_num_tokens]

571 - hidden_states, residual = combined_states.split([hidden_size, hidden_size], dim=-1)

572 - return hidden_states, residual

573 -

574 -

559 @support_torch_compile

560 class Qwen3NextModel(nn.Module, EagleModelMixin):

561 hf_to_vllm_mapper = WeightsMapper(

@@ -629,6 +613,10 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):

613 def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

614 return self.embed_tokens(input_ids)

615

616 + @property

617 + def use_sequence_parallel(self) -> bool:

618 + return self.layers[self.start_layer].use_attn_reduce_scatter_for_moe

619 +

620 def forward(

621 self,

622 input_ids: torch.Tensor | None,

@@ -648,35 +636,20 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):

636 residual = intermediate_tensors["residual"]

637

638 full_num_tokens = positions.shape[-1]

639 + if self.use_sequence_parallel:

640 + hidden_states = sequence_parallel_chunk(hidden_states)

641 + assert residual is None

642 +

643 aux_hidden_states = self._maybe_add_hidden_state([], 0, hidden_states, residual)

644 for layer_idx, layer in enumerate(

645 islice(self.layers, self.start_layer, self.end_layer),

646 start=self.start_layer,

647 ):

656 - if (

657 - hidden_states.shape[0] != full_num_tokens

658 - and not layer.use_attn_reduce_scatter_for_moe

659 - ):

660 - hidden_states, residual = _all_gather_hidden_and_residual(

661 - hidden_states,

662 - residual,

663 - full_num_tokens,

664 - self.config.hidden_size,

665 - )

648 hidden_states, residual = layer(

649 positions=positions,

650 hidden_states=hidden_states,

651 residual=residual,

652 )

671 - if (layer_idx + 1) in self.aux_hidden_state_layers and hidden_states.shape[

672 - 0

673 - ] != full_num_tokens:

674 - hidden_states, residual = _all_gather_hidden_and_residual(

675 - hidden_states,

676 - residual,

677 - full_num_tokens,

678 - self.config.hidden_size,

679 - )

653 self._maybe_add_hidden_state(

654 aux_hidden_states, layer_idx + 1, hidden_states, residual

655 )

@@ -685,14 +658,19 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):

658 return IntermediateTensors(

659 {"hidden_states": hidden_states, "residual": residual}

660 )

688 - if hidden_states.shape[0] != full_num_tokens:

689 - hidden_states, residual = _all_gather_hidden_and_residual(

690 - hidden_states,

691 - residual,

692 - full_num_tokens,

693 - self.config.hidden_size,

694 - )

661 hidden_states, _ = self.norm(hidden_states, residual)

662 + if self.use_sequence_parallel:

663 + if aux_hidden_states:

664 + hidden_size = hidden_states.shape[-1]

665 + hidden_states = torch.cat([hidden_states, *aux_hidden_states], dim=-1)

666 + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

667 + hidden_states = hidden_states[:full_num_tokens]

668 + hidden_states, *aux_hidden_states = hidden_states.split(

669 + hidden_size, dim=-1

670 + )

671 + else:

672 + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

673 + hidden_states = hidden_states[:full_num_tokens]

674 if aux_hidden_states:

675 return hidden_states, aux_hidden_states

676 return hidden_states

vllm/model_executor/models/qwen3_next_mtp.py

+9 -9

@@ -9,7 +9,7 @@ from torch import nn

10 from vllm.compilation.decorators import support_torch_compile

11 from vllm.config import VllmConfig

12 -from vllm.distributed.parallel_state import get_pp_group

12 +from vllm.distributed import get_pp_group, tensor_model_parallel_all_gather

13 from vllm.logger import init_logger

14 from vllm.model_executor.layers.linear import ColumnParallelLinear

15 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -22,8 +22,8 @@ from vllm.model_executor.models.qwen3_next import (

22 Qwen3NextModel,

23 Qwen3NextRMSNorm,

24 QwenNextMixtureOfExperts,

25 - _all_gather_hidden_and_residual,

25 )

26 +from vllm.model_executor.models.utils import sequence_parallel_chunk

27 from vllm.sequence import IntermediateTensors

28 from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig

@@ -131,6 +131,10 @@ class Qwen3NextMultiTokenPredictor(nn.Module):

131

132 current_step_idx = spec_step_idx % self.num_mtp_layers

133 mtp_layer = self.layers[current_step_idx]

134 + if mtp_layer.use_attn_reduce_scatter_for_moe:

135 + assert hidden_states.shape[0] == positions.shape[-1]

136 + hidden_states = sequence_parallel_chunk(hidden_states)

137 + assert residual is None

138 hidden_states, residual = mtp_layer(

139 positions=positions,

140 hidden_states=hidden_states,

@@ -142,14 +146,10 @@ class Qwen3NextMultiTokenPredictor(nn.Module):

146 {"hidden_states": hidden_states, "residual": residual}

147 )

148

145 - if mtp_layer.use_attn_reduce_scatter_for_moe:

146 - hidden_states, residual = _all_gather_hidden_and_residual(

147 - hidden_states,

148 - residual,

149 - positions.shape[-1],

150 - self.config.hidden_size,

151 - )

149 hidden_states, _ = self.norm(hidden_states, residual)

150 + if mtp_layer.use_attn_reduce_scatter_for_moe:

151 + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

152 + hidden_states = hidden_states[: positions.shape[-1]]

153 return hidden_states

154

155 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: