diff --git a/python/sglang/srt/layers/moe/ep_moe/kernels.py b/python/sglang/srt/layers/moe/ep_moe/kernels.py index 39f46ccef..086fa0db2 100644 --- a/python/sglang/srt/layers/moe/ep_moe/kernels.py +++ b/python/sglang/srt/layers/moe/ep_moe/kernels.py @@ -1307,7 +1307,9 @@ def zero_experts_compute_triton( zero_expert_scales[zero_expert_mask] = 0.0 normal_expert_mask = expert_indices >= num_experts - expert_indices[normal_expert_mask] = -1 + # Keep a valid routed-expert id for MoE kernels that do not accept negative + # ids. The zero scale below still removes the routed-expert contribution. + expert_indices[normal_expert_mask] = 0 expert_scales[normal_expert_mask] = 0.0 output = torch.zeros_like(hidden_states).to(hidden_states.device) diff --git a/test/registered/moe/test_zero_experts.py b/test/registered/moe/test_zero_experts.py new file mode 100644 index 000000000..b487bb021 --- /dev/null +++ b/test/registered/moe/test_zero_experts.py @@ -0,0 +1,66 @@ +import unittest + +import torch + +from sglang.srt.layers.moe.ep_moe.kernels import zero_experts_compute_triton +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small") + + +@unittest.skipIf(not torch.cuda.is_available(), "CUDA is required") +class TestZeroExpertsComputeTriton(CustomTestCase): + def test_zero_expert_routes_keep_valid_ids(self): + num_experts = 4 + hidden_states = torch.arange( + 2 * 512, dtype=torch.float32, device="cuda" + ).reshape(2, 512) + expert_indices = torch.tensor( + [[0, 4, 1], [5, 2, 6]], dtype=torch.int64, device="cuda" + ) + expert_scales = torch.tensor( + [[0.1, 0.3, 0.2], [0.4, 0.5, 0.6]], + dtype=torch.float32, + device="cuda", + ) + original_indices = expert_indices.clone() + original_scales = expert_scales.clone() + + output = zero_experts_compute_triton( + expert_indices=expert_indices, + expert_scales=expert_scales, + num_experts=num_experts, + zero_expert_type="identity", + hidden_states=hidden_states, + ) + torch.cuda.synchronize() + + zero_expert_mask = original_indices >= num_experts + normal_expert_mask = ~zero_expert_mask + self.assertTrue(torch.all(expert_indices >= 0).item()) + torch.testing.assert_close( + expert_indices[normal_expert_mask], + original_indices[normal_expert_mask], + ) + torch.testing.assert_close( + expert_indices[zero_expert_mask], + torch.zeros_like(expert_indices[zero_expert_mask]), + ) + torch.testing.assert_close( + expert_scales[normal_expert_mask], + original_scales[normal_expert_mask], + ) + torch.testing.assert_close( + expert_scales[zero_expert_mask], + torch.zeros_like(expert_scales[zero_expert_mask]), + ) + + expected_scale = ( + original_scales * zero_expert_mask.to(original_scales.dtype) + ).sum(dim=-1, keepdim=True) + torch.testing.assert_close(output, hidden_states * expected_scale) + + +if __name__ == "__main__": + unittest.main()