Support DSV4 shared expert fusion for DeepEP and MegaMOE (#27349)

test/registered/moe/test_fused_append_remap_per_rank_shared_slots.py

@@ -1,8 +1,9 @@
-"""Unit tests for the fused append + DeepEP-remap shared-experts Triton kernel.
+"""Unit tests for fused append + per-rank shared-slot remap.
 
 Covers ``fused_append_remap_shared_experts_deepep``, which collapses
-``fused_append_shared_experts()`` followed by ``_remap_topk_for_deepep()`` into a
-single Triton launch on the aiter/DeepEP-class path. The kernel is GPU-only
+``fused_append_shared_experts()`` followed by
+``remap_topk_for_per_rank_shared_slots()`` into a
+single Triton launch on the per-rank shared-slot path. The kernel is GPU-only
 (Triton), so these tests are skipped when no accelerator is present.
 """
 
@@ -14,7 +15,11 @@ from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels impo
     fused_append_remap_shared_experts_deepep,
     fused_append_shared_experts,
 )
-from sglang.srt.layers.moe.topk import TopKConfig, _remap_topk_for_deepep, _use_aiter
+from sglang.srt.layers.moe.topk import (
+    TopKConfig,
+    _use_aiter,
+    remap_topk_for_per_rank_shared_slots,
+)
 from sglang.srt.runtime_context import get_parallel
 from sglang.srt.utils import get_device
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
@@ -50,7 +55,7 @@ def _reference_append_remap(
 @unittest.skipUnless(
     torch.cuda.is_available(), "fused append+remap kernel requires a GPU"
 )
-class TestFusedAppendRemapDeepEP(CustomTestCase):
+class TestFusedAppendRemapPerRankSharedSlots(CustomTestCase):
     # (m, k, num_physical_routed, ep_size, ep_rank, num_fused_shared_experts).
     # k and num_fused_shared_experts are kept powers of two (tl.arange constraint).
     CASES = [
@@ -107,7 +112,7 @@ class TestFusedAppendRemapDeepEP(CustomTestCase):
                 self.assertTrue(torch.allclose(got_w, exp_w))
 
     def test_equivalence_with_eager_append_then_remap(self):
-        """Fused kernel == fused_append_shared_experts() + _remap_topk_for_deepep().
+        """Fused kernel == append shared experts + per-rank shared-slot remap.
 
         The eager remap overwrites the shared weight: 1.0 on the aiter/HIP path
         (routed_scaling_factor is pre-folded into the routed topk weights), else
@@ -140,7 +145,7 @@ class TestFusedAppendRemapDeepEP(CustomTestCase):
                         scale_factor,
                         npr,  # shared-expert base id (overwritten by the remap)
                     )
-                    eager_ids, eager_w = _remap_topk_for_deepep(
+                    eager_ids, eager_w = remap_topk_for_per_rank_shared_slots(
                         eager_ids,
                         eager_w,
                         s,

test/registered/moe/test_hash_topk.py

@@ -0,0 +1,145 @@
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
+from sglang.srt.layers.moe import hash_topk as hash_topk_module
+from sglang.srt.layers.moe.hash_topk import HashTopK
+from sglang.srt.layers.moe.topk import (
+    StandardTopKOutput,
+)
+from sglang.srt.models.deepseek_v2 import DeepseekV2MoE
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="base-b-test-cpu")
+
+
+@pytest.fixture(autouse=True)
+def _set_dummy_server_args():
+    set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
+
+def test_hash_topk_remaps_per_rank_fused_shared_slots(monkeypatch):
+    monkeypatch.setattr(
+        hash_topk_module, "has_per_rank_fused_shared_slots", lambda *_args: True
+    )
+    recorded = {}
+
+    class FakeRecorder:
+        def on_select_experts(self, *, topk_ids):
+            recorded["topk_ids"] = topk_ids.clone()
+
+    monkeypatch.setattr(
+        hash_topk_module,
+        "get_global_expert_distribution_recorder",
+        lambda: FakeRecorder(),
+    )
+
+    topk = HashTopK(
+        topk=3,
+        num_experts=256,
+        num_fused_shared_experts=1,
+        vocab_size=2,
+        scoring_func="sqrtsoftplus",
+        routed_scaling_factor=2.5,
+    )
+    with torch.no_grad():
+        topk.tid2eid.copy_(torch.tensor([[0, 65], [63, 127]], dtype=torch.int32))
+
+    info = ExpertLocationDispatchInfo(
+        ep_dispatch_algorithm="static",
+        partial_logical_to_rank_dispatch_physical_map=torch.arange(
+            256, dtype=torch.int32
+        ),
+        partial_logical_to_all_physical_map=torch.arange(256, dtype=torch.int32).view(
+            256, 1
+        ),
+        partial_logical_to_all_physical_map_num_valid=torch.ones(
+            256, dtype=torch.int32
+        ),
+        num_physical_experts=256,
+    )
+
+    with (
+        get_parallel().override(moe_ep_size=4, moe_ep_rank=2),
+        hash_topk_module.envs.SGLANG_OPT_USE_FUSED_HASH_TOPK.override(False),
+    ):
+        output = topk(
+            hidden_states=torch.empty(2, 4),
+            router_logits=torch.ones(2, 256),
+            input_ids=torch.tensor([0, 1], dtype=torch.int64),
+            expert_location_dispatch_info=info,
+        )
+
+    # Physical layout for EP=4 has 64 routed slots per rank plus one local
+    # shared slot: [0..63, shared, 64..127, shared, ...].
+    assert output.topk_ids.tolist() == [[0, 66, 194], [63, 128, 194]]
+    assert torch.allclose(output.topk_weights[:, -1], torch.full((2,), 0.4))
+    assert recorded["topk_ids"].tolist() == [[0, 65], [63, 127]]
+
+
+def test_hash_topk_empty_output_keeps_per_rank_shared_slot(monkeypatch):
+    monkeypatch.setattr(
+        hash_topk_module, "has_per_rank_fused_shared_slots", lambda *_args: True
+    )
+
+    topk = HashTopK(
+        topk=7,
+        num_experts=256,
+        num_fused_shared_experts=1,
+        vocab_size=2,
+        scoring_func="softmax",
+    )
+
+    output = topk.empty_topk_output(torch.device("cpu"))
+
+    assert output.topk_ids.shape == (0, 7)
+    assert output.topk_weights.shape == (0, 7)
+    assert output.router_logits.shape == (0, 6)
+
+
+def test_deepep_empty_forward_does_not_append_shared_slot_twice():
+    captured = {}
+
+    class FakeTopK:
+        def empty_topk_output(self, device, *, layer_id=None):
+            return StandardTopKOutput(
+                topk_weights=torch.empty((0, 9), dtype=torch.float32, device=device),
+                topk_ids=torch.empty((0, 9), dtype=torch.int32, device=device),
+                router_logits=torch.empty((0, 8), dtype=torch.float32, device=device),
+            )
+
+    class FakeExperts:
+        should_fuse_routed_scaling_factor_in_topk = True
+
+        def __call__(self, hidden_states, topk_output):
+            captured["topk_ids_shape"] = tuple(topk_output.topk_ids.shape)
+            captured["topk_weights_shape"] = tuple(topk_output.topk_weights.shape)
+            return hidden_states
+
+    moe = SimpleNamespace(
+        _fuse_shared_experts_inside_sbo=False,
+        is_nextn=False,
+        num_fused_shared_experts=1,
+        layer_id=0,
+        topk=FakeTopK(),
+        experts=FakeExperts(),
+        alt_stream=None,
+        routed_scaling_factor=1.0,
+    )
+
+    hidden_states = torch.empty((0, 4), dtype=torch.float32)
+    forward_batch = SimpleNamespace(num_token_non_padded=None)
+
+    DeepseekV2MoE.forward_deepep(moe, hidden_states, forward_batch)
+
+    assert captured["topk_ids_shape"] == (0, 9)
+    assert captured["topk_weights_shape"] == (0, 9)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/unit/eplb/test_deepep_waterfill_eplb.py

@@ -56,7 +56,7 @@ class TestDeepEPWaterfillEPLB(CustomTestCase):
         self.assertEqual(len(weights), 1)
         self.assertEqual(weights[0].shape, (experts.num_local_experts, 2))
 
-    def test_topk_recorder_ids_exclude_deepep_fused_shared_slots(self):
+    def test_topk_recorder_ids_exclude_per_rank_fused_shared_slots(self):
         topk_ids = torch.tensor([[0, 33, 263, 256]], dtype=torch.int32)
         topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
         topk_config = TopKConfig(
@@ -74,7 +74,9 @@ class TestDeepEPWaterfillEPLB(CustomTestCase):
         with (
             patch.object(topk_module, "_is_cuda", True),
             patch.object(topk_module, "_use_aiter", False),
-            patch.object(topk_module, "is_deepep_class_backend", return_value=True),
+            patch.object(
+                topk_module, "has_per_rank_fused_shared_slots", return_value=True
+            ),
             get_parallel().override(moe_ep_size=8, moe_ep_rank=7),
             patch.object(
                 topk_module,
@@ -94,7 +96,7 @@ class TestDeepEPWaterfillEPLB(CustomTestCase):
         self.assertTrue(torch.equal(processed_ids, torch.tensor([[0, 34, 270, 271]])))
         self.assertTrue(torch.equal(recorder_ids, torch.tensor([[0, 33, 263]])))
 
-    def test_topk_recorder_ids_match_dispatch_ids_for_non_deepep_fusion(self):
+    def test_topk_recorder_ids_match_dispatch_ids_without_per_rank_shared_slots(self):
         topk_ids = torch.tensor([[0, 33, 263, 256]], dtype=torch.int32)
         topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
         topk_config = TopKConfig(
@@ -112,7 +114,9 @@ class TestDeepEPWaterfillEPLB(CustomTestCase):
         with (
             patch.object(topk_module, "_is_cuda", True),
             patch.object(topk_module, "_use_aiter", False),
-            patch.object(topk_module, "is_deepep_class_backend", return_value=False),
+            patch.object(
+                topk_module, "has_per_rank_fused_shared_slots", return_value=False
+            ),
             patch.object(
                 topk_module,
                 "_biased_grouped_topk_postprocess",

test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py

@@ -1,10 +1,10 @@
-"""Unit tests for fused shared-expert weight scaling on the DeepEP layout.
+"""Unit tests for fused shared-expert weight scaling on per-rank shared slots.
 
-These tests pin the contract of ``_remap_topk_for_deepep`` for the fused shared
-expert's topk weight on the two paths this fix covers:
+These tests pin the contract of ``remap_topk_for_per_rank_shared_slots`` for
+the fused shared expert's topk weight on the two paths this fix covers:
 
   * aiter (HIP) path: routed_scaling_factor is folded into the routed weights and
-    forward_deepep skips the post-MoE multiply, so the shared weight must be 1.0
+    the post-MoE multiply is skipped, so the shared weight must be 1.0
     for a net 1.0x contribution.
   * post-MoE scaling path (default): the whole MoE output is multiplied by
     routed_scaling_factor afterward, so the shared weight must be 1/rsf.
@@ -56,7 +56,7 @@ class TestFusedSharedExpertScaling(CustomTestCase):
                 ),
             ),
         ):
-            _out_ids, out_weights = topk_module._remap_topk_for_deepep(
+            _out_ids, out_weights = topk_module.remap_topk_for_per_rank_shared_slots(
                 topk_ids.clone(),
                 topk_weights.clone(),
                 num_fused_shared_experts=1,
@@ -100,7 +100,7 @@ class TestFusedSharedExpertScaling(CustomTestCase):
                 ),
             ),
         ):
-            out_ids, _ = topk_module._remap_topk_for_deepep(
+            out_ids, _ = topk_module.remap_topk_for_per_rank_shared_slots(
                 topk_ids.clone(),
                 topk_weights.clone(),
                 num_fused_shared_experts=1,

test/registered/unit/layers/quantization/test_fp8_utils_mxfp4.py

@@ -0,0 +1,50 @@
+import unittest
+
+import torch
+
+from sglang.srt.layers.quantization.fp8_utils import (
+    quantize_block_fp8_weight_to_mxfp4,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=4, suite="base-a-test-cpu")
+
+
+class TestFp8UtilsMxfp4(unittest.TestCase):
+    def test_quantize_block_fp8_weight_to_mxfp4_shapes_and_dtype(self):
+        fp8_weight = (
+            torch.linspace(-2.0, 2.0, 32 * 32, dtype=torch.float32)
+            .reshape(32, 32)
+            .to(torch.float8_e4m3fn)
+        )
+        fp8_scale = torch.ones(1, 1, dtype=torch.float8_e8m0fnu)
+
+        fp4_weight, fp4_scale = quantize_block_fp8_weight_to_mxfp4(
+            fp8_weight, fp8_scale, [128, 128]
+        )
+
+        self.assertEqual(fp4_weight.dtype, torch.int8)
+        self.assertEqual(fp4_weight.shape, torch.Size([32, 16]))
+        self.assertEqual(fp4_scale.dtype, torch.float8_e8m0fnu)
+        self.assertEqual(fp4_scale.shape, torch.Size([32, 1]))
+
+    def test_quantize_block_fp8_weight_to_mxfp4_grouped_weight(self):
+        fp8_weight = (
+            torch.linspace(-2.0, 2.0, 2 * 32 * 32, dtype=torch.float32)
+            .reshape(2, 32, 32)
+            .to(torch.float8_e4m3fn)
+        )
+        fp8_scale = torch.ones(2, 1, 1, dtype=torch.float8_e8m0fnu)
+
+        fp4_weight, fp4_scale = quantize_block_fp8_weight_to_mxfp4(
+            fp8_weight, fp8_scale, [128, 128]
+        )
+
+        self.assertEqual(fp4_weight.dtype, torch.int8)
+        self.assertEqual(fp4_weight.shape, torch.Size([2, 32, 16]))
+        self.assertEqual(fp4_scale.dtype, torch.float8_e8m0fnu)
+        self.assertEqual(fp4_scale.shape, torch.Size([2, 32, 1]))
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py

@@ -0,0 +1,50 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from sglang.srt.models import deepseek_v4 as deepseek_v4_module
+from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=4, suite="base-a-test-cpu")
+
+
+class TestDeepseekV4SharedExpertFusionPolicy(unittest.TestCase):
+    def _make_model(self, n_shared_experts=1):
+        return SimpleNamespace(
+            config=SimpleNamespace(n_shared_experts=n_shared_experts)
+        )
+
+    def test_disables_shared_fusion_without_enforce(self):
+        server_args = SimpleNamespace(
+            disable_shared_experts_fusion=False,
+            enforce_shared_experts_fusion=False,
+        )
+        model = self._make_model()
+
+        with patch.object(
+            deepseek_v4_module, "get_global_server_args", return_value=server_args
+        ):
+            DeepseekV4ForCausalLM.determine_num_fused_shared_experts(model)
+
+        self.assertEqual(model.num_fused_shared_experts, 0)
+        self.assertTrue(server_args.disable_shared_experts_fusion)
+
+    def test_enables_shared_fusion_when_enforced(self):
+        server_args = SimpleNamespace(
+            disable_shared_experts_fusion=False,
+            enforce_shared_experts_fusion=True,
+        )
+        model = self._make_model()
+
+        with patch.object(
+            deepseek_v4_module, "get_global_server_args", return_value=server_args
+        ):
+            DeepseekV4ForCausalLM.determine_num_fused_shared_experts(model)
+
+        self.assertEqual(model.num_fused_shared_experts, 1)
+        self.assertFalse(server_args.disable_shared_experts_fusion)
+
+
+if __name__ == "__main__":
+    unittest.main()