diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py index b44c5f040982..b68f7e718bd2 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_5_moe.py @@ -780,13 +780,14 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: layer_type="moe", ) - # Single merge-point all_reduce for routed + shared partial sums. - # Both branches produce per-rank partial outputs under TP/EP sharding - # (routed: MoEShardableNode; shared: rowwise down_proj inside the MLP). - # One reduction on the sum lifts both to full; reducing before the add - # would mix a full routed contribution with an unreduced shared one. - expert_output = expert_output + shared_expert_output + # The shared expert is replicated (Qwen3_5MoeMLP intentionally omits + # ``layer_type`` and the yaml ``shard_layers`` whitelist excludes it), + # so its output is already the full value on every rank. All-reduce + # the sharded routed-expert partial first, then add the replicated + # shared output; adding before would scale the shared output by the + # TP world size. expert_output = torch.ops.auto_deploy.all_reduce(expert_output, layer_type="moe") + expert_output = expert_output + shared_expert_output expert_output = expert_output.reshape(batch_size, sequence_length, hidden_dim) return expert_output