From acea43079fa569d67172bf4754380bff2c8190be Mon Sep 17 00:00:00 2001 From: Joe Date: Wed, 2 Sep 2026 17:16:19 -0400 Subject: [PATCH] Fix native MoE handling of noncontiguous top-k IDs (#36407) Co-authored-by: BBuf <1182563586@qq.com> --- .../sglang/srt/layers/moe/fused_moe_native.py | 2 +- .../unit/layers/moe/test_fused_moe_native.py | 55 +++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) create mode 100644 test/registered/unit/layers/moe/test_fused_moe_native.py diff --git a/python/sglang/srt/layers/moe/fused_moe_native.py b/python/sglang/srt/layers/moe/fused_moe_native.py index d72f0e9a3..359ebfa26 100644 --- a/python/sglang/srt/layers/moe/fused_moe_native.py +++ b/python/sglang/srt/layers/moe/fused_moe_native.py @@ -76,7 +76,7 @@ def moe_forward_native( cnts = topk_ids.new_zeros((topk_ids.shape[0], len_experts)) cnts.scatter_(1, topk_ids.to(torch.int64), 1) tokens_per_expert = cnts.sum(dim=0) - idxs = topk_ids.view(-1).argsort() + idxs = topk_ids.reshape(-1).argsort() sorted_tokens = x[idxs // topk_ids.shape[1]] tokens_per_expert = tokens_per_expert.cpu().numpy() diff --git a/test/registered/unit/layers/moe/test_fused_moe_native.py b/test/registered/unit/layers/moe/test_fused_moe_native.py new file mode 100644 index 000000000..4d52819b1 --- /dev/null +++ b/test/registered/unit/layers/moe/test_fused_moe_native.py @@ -0,0 +1,55 @@ +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.layers.moe.fused_moe_native import moe_forward_native +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=3, suite="base-a-test-cpu") + + +class TestFusedMoeNative(CustomTestCase): + def test_noncontiguous_topk_ids(self): + torch.manual_seed(0) + layer = SimpleNamespace( + num_experts=3, + w13_weight=torch.randn(3, 4, 2), + w2_weight=torch.randn(3, 2, 2), + ) + hidden_states = torch.randn(3, 2) + topk_weights = torch.tensor( + [[0.75, 0.25], [0.4, 0.6], [0.9, 0.1]], dtype=torch.float32 + ) + padded_topk_ids = torch.tensor( + [[0, -1, 2, -1], [1, -1, 0, -1], [2, -1, 1, -1]], + dtype=torch.int64, + ) + topk_ids = padded_topk_ids[:, ::2] + self.assertFalse(topk_ids.is_contiguous()) + + config = SimpleNamespace( + activation="silu", + apply_router_weight_on_input=False, + gemm1_alpha=None, + gemm1_clamp_limit=None, + ) + output = moe_forward_native( + layer, + hidden_states, + (topk_weights, topk_ids, None), + config, + ) + reference = moe_forward_native( + layer, + hidden_states, + (topk_weights, topk_ids.contiguous(), None), + config, + ) + + torch.testing.assert_close(output, reference) + + +if __name__ == "__main__": + unittest.main()