From 3cb3f7c01814c90f3f4aacde83f6f2cfcd20ed35 Mon Sep 17 00:00:00 2001 From: xutizhou Date: Tue, 14 Apr 2026 17:33:27 +0800 Subject: [PATCH] fix: EPLB dispatch OOB when shared experts fusion enabled under DeepEP (#22525) --- python/sglang/srt/layers/moe/topk.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index ef7749c9a..c65942147 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -1017,9 +1017,21 @@ def _post_process_topk_ids( topk_ids=topk_ids, ) if _is_cuda: - topk_ids = _biased_grouped_topk_postprocess( - topk_ids, expert_location_dispatch_info, num_token_non_padded - ) + # When shared experts are fused (appended as extra columns in topk_ids), + # EPLB dispatch must only remap the routed expert columns. + # The shared expert column (value = n_routed_experts) would be out-of-bounds + # for the logical-to-physical dispatch table. + if num_fused_shared_experts > 0 and is_deepep_class_backend(): + shared_cols = topk_ids[:, -num_fused_shared_experts:] + routed_cols = topk_ids[:, :-num_fused_shared_experts] + routed_cols = _biased_grouped_topk_postprocess( + routed_cols, expert_location_dispatch_info, num_token_non_padded + ) + topk_ids = torch.cat([routed_cols, shared_cols], dim=-1) + else: + topk_ids = _biased_grouped_topk_postprocess( + topk_ids, expert_location_dispatch_info, num_token_non_padded + ) if num_fused_shared_experts > 0 and _use_aiter: M, N = router_logits.shape