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

[Perf][Qwen] Replace MOE all-reduce with reduce-scatter (#47006)

Signed-off-by: gcanlin <canlinguosdu@gmail.com> Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: yewentao256 <zhyanwentao@126.com>

Canlin Guo committed Jul 12, 2026 at 14:14 UTC a02984ed471488c0f0e8f73cab21be4325992d4c

3 files changed +113 -11

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

@@ -437,6 +437,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

437 vllm_config: VllmConfig,

438 prefix: str = "",

439 gqa_interleaved_layout=False,

440 + reduce_results: bool = True,

441 ) -> None:

442 super().__init__(config, vllm_config, prefix)

443

@@ -547,6 +548,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):

548 self.hidden_size,

549 bias=False,

550 input_is_parallel=True,

551 + reduce_results=reduce_results,

552 quant_config=self.quant_config,

553 prefix=f"{prefix}.out_proj",

554 )

vllm/model_executor/models/qwen3_5.py

@@ -121,10 +121,17 @@ class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):

121 config = vllm_config.model_config.hf_text_config

122 model_config = vllm_config.model_config

123 cache_config = vllm_config.cache_config

124 + parallel_config = vllm_config.parallel_config

125 quant_config = vllm_config.quant_config

126

127 self.layer_type = layer_type

128 self.layer_idx = extract_layer_index(prefix)

129 + is_moe_layer = config.model_type == "qwen3_5_moe_text"

130 + self.use_attn_reduce_scatter_for_moe = (

131 + parallel_config.use_sequence_parallel_moe

132 + and parallel_config.pipeline_parallel_size == 1

133 + and is_moe_layer

134 + )

135

136 if self.layer_type == "linear_attention":

137 self.linear_attn = QwenGatedDeltaNetAttention(

@@ -132,6 +139,7 @@ class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):

139 vllm_config=vllm_config,

140 prefix=f"{prefix}.linear_attn",

141 gqa_interleaved_layout=False,

142 + reduce_results=not self.use_attn_reduce_scatter_for_moe,

143 )

144 elif self.layer_type == "full_attention":

145 self.self_attn = Qwen3NextAttention(

@@ -140,6 +148,7 @@ class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):

148 cache_config=cache_config,

149 quant_config=quant_config,

150 prefix=f"{prefix}.self_attn",

151 + reduce_results=not self.use_attn_reduce_scatter_for_moe,

152 )

153 else:

154 raise ValueError(f"Invalid layer_type {self.layer_type}")

vllm/model_executor/models/qwen3_next.py

+102 -11

@@ -16,6 +16,7 @@ from vllm.distributed import (

16 get_pp_group,

17 get_tensor_model_parallel_world_size,

18 tensor_model_parallel_all_gather,

19 + tensor_model_parallel_reduce_scatter,

20 )

21 from vllm.logger import init_logger

22 from vllm.model_executor.layers.attention import Attention

@@ -192,13 +193,17 @@ class Qwen3NextSparseMoeBlock(nn.Module):

193 else None,

194 )

195

195 - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

196 + def forward(

197 + self,

198 + hidden_states: torch.Tensor,

199 + already_sequence_parallel: bool = False,

200 + ) -> torch.Tensor:

201 # NOTE: hidden_states can have either 1D or 2D shape.

202 orig_shape = hidden_states.shape

203 num_tokens, hidden_dim = hidden_states.shape

204 hidden_states = hidden_states.view(-1, hidden_dim)

205

201 - if self.is_sequence_parallel:

206 + if self.is_sequence_parallel and not already_sequence_parallel:

207 hidden_states = sequence_parallel_chunk(hidden_states)

208

209 if self.experts.is_internal_router:

@@ -213,7 +218,7 @@ class Qwen3NextSparseMoeBlock(nn.Module):

218 hidden_states=hidden_states, router_logits=router_logits

219 )

220

216 - if self.is_sequence_parallel:

221 + if self.is_sequence_parallel and not already_sequence_parallel:

222 final_hidden_states = tensor_model_parallel_all_gather(

223 final_hidden_states, 0

224 )

@@ -229,6 +234,7 @@ class Qwen3NextAttention(nn.Module):

234 model_config: ModelConfig | None = None,

235 cache_config: CacheConfig | None = None,

236 quant_config: QuantizationConfig | None = None,

237 + reduce_results: bool = True,

238 prefix: str = "",

239 ) -> None:

240 super().__init__()

@@ -271,6 +277,7 @@ class Qwen3NextAttention(nn.Module):

277 self.total_num_heads * self.head_dim,

278 config.hidden_size,

279 bias=False,

280 + reduce_results=reduce_results,

281 quant_config=quant_config,

282 prefix=f"{prefix}.o_proj",

283 )

@@ -405,16 +412,31 @@ class Qwen3NextDecoderLayer(nn.Module):

412 model_config = vllm_config.model_config

413 cache_config = vllm_config.cache_config

414 quant_config = vllm_config.quant_config

415 + parallel_config = vllm_config.parallel_config

416

417 self.layer_type = layer_type

418 self.layer_idx = extract_layer_index(prefix)

419

420 + mlp_only_layers = (

421 + [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers

422 + )

423 + is_moe_layer = (self.layer_idx not in mlp_only_layers) and (

424 + config.num_experts > 0

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

426 + )

427 + self.use_attn_reduce_scatter_for_moe = (

428 + parallel_config.use_sequence_parallel_moe

429 + and parallel_config.pipeline_parallel_size == 1

430 + and is_moe_layer

431 + )

432 +

433 if self.layer_type == "linear_attention":

434 self.linear_attn = QwenGatedDeltaNetAttention(

435 config,

436 vllm_config=vllm_config,

437 prefix=f"{prefix}.linear_attn",

438 gqa_interleaved_layout=True,

439 + reduce_results=not self.use_attn_reduce_scatter_for_moe,

440 )

441 elif self.layer_type == "full_attention":

442 self.self_attn = Qwen3NextAttention(

@@ -422,18 +444,13 @@ class Qwen3NextDecoderLayer(nn.Module):

444 model_config=model_config,

445 cache_config=cache_config,

446 quant_config=quant_config,

447 + reduce_results=not self.use_attn_reduce_scatter_for_moe,

448 prefix=f"{prefix}.self_attn",

449 )

450 else:

451 raise ValueError(f"Invalid layer_type {self.layer_type}")

452

430 - mlp_only_layers = (

431 - [] if not hasattr(config, "mlp_only_layers") else config.mlp_only_layers

432 - )

433 - if (self.layer_idx not in mlp_only_layers) and (

434 - config.num_experts > 0

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

436 - ):

453 + if is_moe_layer:

454 self.mlp = Qwen3NextSparseMoeBlock(

455 vllm_config=vllm_config,

456 prefix=f"{prefix}.mlp",

@@ -478,12 +495,23 @@ class Qwen3NextDecoderLayer(nn.Module):

495 positions: torch.Tensor = None,

496 **kwargs: object,

497 ):

498 + full_num_tokens = positions.shape[-1]

499 + input_is_sequence_parallel = (

500 + self.use_attn_reduce_scatter_for_moe

501 + and residual is not None

502 + and hidden_states.shape[0] != full_num_tokens

503 + )

504 +

505 if residual is None:

506 residual = hidden_states

507 hidden_states = self.input_layernorm(hidden_states)

508 else:

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

510

511 + if input_is_sequence_parallel:

512 + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

513 + hidden_states = hidden_states[:full_num_tokens]

514 +

515 if self.layer_type == "linear_attention":

516 hidden_states = self.linear_attn(hidden_states=hidden_states)

517 elif self.layer_type == "full_attention":

@@ -504,9 +532,25 @@ class Qwen3NextDecoderLayer(nn.Module):

532 self.attn_layer_scale.to(hidden_states.dtype) + 1

533 )

534

535 + if self.use_attn_reduce_scatter_for_moe:

536 + tp_world_size = get_tensor_model_parallel_world_size()

537 + # small trick using minus, eg. -17 % 8 = 7

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

539 + # pad if not divisible by world size

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

541 + hidden_states = tensor_model_parallel_reduce_scatter(hidden_states, 0)

542 + if not input_is_sequence_parallel:

543 + residual = sequence_parallel_chunk(residual)

544 +

545 # Fully Connected

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

509 - hidden_states = self.mlp(hidden_states)

547 + if self.use_attn_reduce_scatter_for_moe:

548 + hidden_states = self.mlp(

549 + hidden_states,

550 + already_sequence_parallel=True,

551 + )

552 + else:

553 + hidden_states = self.mlp(hidden_states)

554

555 if self.layer_scale:

556 if len(hidden_states.shape) == 2:

@@ -525,6 +569,24 @@ class Qwen3NextDecoderLayer(nn.Module):

569 return hidden_states, residual

570

571

572 +def _all_gather_hidden_and_residual(

573 + hidden_states: torch.Tensor,

574 + residual: torch.Tensor | None,

575 + full_num_tokens: int,

576 + hidden_size: int,

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

578 + if residual is None:

579 + hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)

580 + hidden_states = hidden_states[:full_num_tokens]

581 + return hidden_states, None

582 +

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

584 + combined_states = tensor_model_parallel_all_gather(combined_states, 0)

585 + combined_states = combined_states[:full_num_tokens]

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

587 + return hidden_states, residual

588 +

589 +

590 @support_torch_compile

591 class Qwen3NextModel(nn.Module, EagleModelMixin):

592 hf_to_vllm_mapper = WeightsMapper(

@@ -577,6 +639,8 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):

639 else:

640 self.norm = PPMissingLayer()

641

642 + self.aux_hidden_state_layers: tuple[int, ...] = ()

643 +

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

645 return self.embed_tokens(input_ids)

646

@@ -598,16 +662,36 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):

662 hidden_states = intermediate_tensors["hidden_states"]

663 residual = intermediate_tensors["residual"]

664

665 + full_num_tokens = positions.shape[-1]

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

667 for layer_idx, layer in enumerate(

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

669 start=self.start_layer,

670 ):

671 + if (

672 + hidden_states.shape[0] != full_num_tokens

673 + and not layer.use_attn_reduce_scatter_for_moe

674 + ):

675 + hidden_states, residual = _all_gather_hidden_and_residual(

676 + hidden_states,

677 + residual,

678 + full_num_tokens,

679 + self.config.hidden_size,

680 + )

681 hidden_states, residual = layer(

682 positions=positions,

683 hidden_states=hidden_states,

684 residual=residual,

685 )

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

687 + 0

688 + ] != 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 + )

695 self._maybe_add_hidden_state(

696 aux_hidden_states, layer_idx + 1, hidden_states, residual

697 )

@@ -616,6 +700,13 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):

700 return IntermediateTensors(

701 {"hidden_states": hidden_states, "residual": residual}

702 )

703 + if hidden_states.shape[0] != full_num_tokens:

704 + hidden_states, residual = _all_gather_hidden_and_residual(

705 + hidden_states,

706 + residual,

707 + full_num_tokens,

708 + self.config.hidden_size,

709 + )

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

711 if aux_hidden_states:

712 return hidden_states, aux_hidden_states