From e4191708c9d6510f9c3e5178db3b261abbde7065 Mon Sep 17 00:00:00 2001 From: Alex Sun Date: Thu, 4 Jun 2026 14:04:39 +0800 Subject: [PATCH] =?UTF-8?q?[Qwen3.5][AMD]=20Fix=20shared-expert=20=C3=97ep?= =?UTF-8?q?=5Fsize=20over-count=20under=20allreduce-EP=20(#26845)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: HaiShaw --- python/sglang/srt/models/qwen2_moe.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py index 305155e13..2c1027fcd 100644 --- a/python/sglang/srt/models/qwen2_moe.py +++ b/python/sglang/srt/models/qwen2_moe.py @@ -71,6 +71,7 @@ from sglang.srt.layers.moe.topk import StandardTopKOutput, TopK, TopKOutputCheck from sglang.srt.layers.moe.utils import ( RoutingMethodType, filter_moe_weight_param_global_expert, + is_deepep_class_backend, ) from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention @@ -343,7 +344,19 @@ class Qwen2MoeSparseMoeBlock(nn.Module): return None shared_out = self.shared_expert_gate(hidden_states) shared_logits = shared_out[0] if isinstance(shared_out, tuple) else shared_out - return F.sigmoid(shared_logits) + w = F.sigmoid(shared_logits) + # This block runs only on the AMD AITER shared_expert_fusion path + # Allreduce-EP path: the fused shared expert occupies a single global + # slot loaded onto every EP rank (see FusedMoE.__init__: num_shared_slots + # == num_fused_shared_experts when not is_deepep_class_backend()). Every + # rank therefore computes the same full shared output, and the + # post-experts all_reduce sums it ep_size times. Pre-scale the per-token + # routing weight by 1/ep_size to cancel this, mirroring DeepSeek-V2's + # fused_shared_experts_scaling_factor pattern. + moe_ep_size = get_moe_expert_parallel_world_size() + if moe_ep_size > 1 and not is_deepep_class_backend(): + w = w / float(moe_ep_size) + return w def _append_shared_to_topk_output( self,