From 6a21dd20b10614d7812c5c24ce78ff34b76d65d7 Mon Sep 17 00:00:00 2001 From: xutizhou Date: Mon, 18 May 2026 14:26:04 +0800 Subject: [PATCH] Fix EPLB mapping for TopK paths (#25285) --- python/sglang/srt/layers/moe/topk.py | 4 ---- test/registered/cpu/test_topk.py | 34 ++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index 96031a9dc..f47758996 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -860,8 +860,6 @@ def biased_topk_impl( topk_weights *= routed_scaling_factor topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32) - topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) - _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) return topk_weights, topk_ids @@ -893,8 +891,6 @@ def biased_topk_jit_kernel_impl( apply_routed_scaling_factor_on_output=apply_routed_scaling_factor_on_output, ) topk_weights, topk_ids = topk_weights.to(torch.float32), topk_ids.to(torch.int32) - topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info) - _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) return topk_weights, topk_ids diff --git a/test/registered/cpu/test_topk.py b/test/registered/cpu/test_topk.py index 58c334832..1b79e061d 100644 --- a/test/registered/cpu/test_topk.py +++ b/test/registered/cpu/test_topk.py @@ -2,9 +2,11 @@ import unittest import torch +from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers.moe.topk import ( biased_grouped_topk_impl as native_biased_grouped_topk, ) +from sglang.srt.layers.moe.topk import biased_topk_impl as native_biased_topk from sglang.srt.layers.moe.topk import fused_topk_torch_native as native_fused_topk from sglang.srt.layers.moe.topk import grouped_topk_gpu as native_grouped_topk from sglang.srt.models.llama4 import Llama4MoE @@ -138,6 +140,38 @@ class TestBiasedGroupedTopK(CustomTestCase): ) +class TestBiasedTopK(CustomTestCase): + def test_biased_topk_returns_logical_ids_with_eplb_info(self): + hidden_states = torch.ones(1, 4) + gating_output = torch.tensor([[10.0, 9.0, 1.0, 0.0]]) + correction_bias = torch.zeros(4) + dispatch_info = ExpertLocationDispatchInfo( + ep_dispatch_algorithm="static", + partial_logical_to_rank_dispatch_physical_map=torch.tensor( + [2, 3, 0, 1], dtype=torch.int64 + ), + partial_logical_to_all_physical_map=torch.tensor( + [[2], [3], [0], [1]], dtype=torch.int64 + ), + partial_logical_to_all_physical_map_num_valid=torch.ones( + 4, dtype=torch.int64 + ), + num_physical_experts=4, + ) + + _, topk_ids = native_biased_topk( + hidden_states=hidden_states, + gating_output=gating_output, + correction_bias=correction_bias, + topk=2, + renormalize=False, + scoring_func="sqrtsoftplus", + expert_location_dispatch_info=dispatch_info, + ) + + torch.testing.assert_close(topk_ids, torch.tensor([[0, 1]], dtype=torch.int32)) + + class TestTopK(CustomTestCase): def _run_single_test(self, M, E, topk, renormalize, dtype): torch.manual_seed(1998)