Support MXFP8 and deferred route weighting in DeepEP v2 (#40030)

Co-authored-by: metamergebot <324680979+metamergebot@users.noreply.github.com>
Co-authored-by: Xingyu Liu <38244988+charlotte12l@users.noreply.github.com>
Co-authored-by: pranjalssh <14260275+pranjalssh@users.noreply.github.com>

test/registered/kernels/ops/moe/test_minimax_quant_scatter.py

@@ -35,6 +35,35 @@ register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b
 dev = "cuda"
 
 
+def test_sm120_mxfp8_dispatch_preserves_activation_scale_recipe(monkeypatch):
+    """SM120 group-128 activations must not use the MXFP8 weight-scale recipe."""
+    from sglang.srt.layers import deep_gemm_wrapper
+    from sglang.srt.layers.moe.moe_runner import deep_gemm_sm120
+
+    monkeypatch.setattr(deep_gemm_sm120, "_is_sm120", True)
+    monkeypatch.setattr(deep_gemm_wrapper, "DEEPGEMM_SCALE_UE8M0", True)
+    config = MoeRunnerConfig(
+        num_experts=2, num_local_experts=2, top_k=1, hidden_size=512
+    )
+    quant = DeepGemmMoeQuantInfo(
+        torch.empty(1, dtype=torch.float8_e4m3fn),
+        None,
+        True,
+        block_shape=[1, 32],
+        use_mxfp8=True,
+    )
+    x = torch.randn(1024, 512, device=dev, dtype=torch.bfloat16)
+    ids = (torch.arange(1024, device=dev, dtype=torch.int32) % 2).view(-1, 1)
+    weights = torch.ones(1024, 1, device=dev)
+    result = deep_gemm_sm120.maybe_pre_permute(x, ids, weights, quant, config, {})
+
+    assert quant.scale_recipes(
+        activation_block_size=result.activation_scale_block_size,
+        hidden_size=result.hidden_states.shape[-1],
+        activation_scale_width=result.hidden_states_scale.shape[-1],
+    ) == ((1, 128), (1, 32))
+
+
 @pytest.mark.parametrize("num_tokens", [1, 7, 64, 256])
 @pytest.mark.parametrize("topk", [4, 5, 8])
 @pytest.mark.parametrize("hidden,group", [(6144, 32), (2048, 32), (4096, 128)])

test/registered/kernels/ops/quantization/test_per_token_group_quant.py

@@ -160,7 +160,9 @@ def test_ue8m0_bitexact(dtype, num_tokens, hidden):
     assert torch.equal(exp, exp_ref), "exponent bytes differ"
 
 
-@pytest.mark.parametrize("group_size", get_ci_test_range([16, 32, 64, 128], [16, 64]))
+@pytest.mark.parametrize(
+    "group_size", get_ci_test_range([16, 32, 64, 128], [16, 32, 64])
+)
 def test_ue8m0_group_sizes(group_size):
     """Group size is a template axis (v2 dispatched a runtime switch). Each size
     maps a group onto a different subwarp lane count; codes/exponents must stay
@@ -311,6 +313,104 @@ def _ref_silu_mul(x, hidden):
     return torch.nn.functional.silu(gate.float()).to(x.dtype) * up
 
 
+@pytest.mark.parametrize("group_size,hidden", [(32, 1792), (32, 6144), (128, 1024)])
+@pytest.mark.parametrize("swiglu_limit", [None, 10.0])
+def test_fp32_silu_post_quant(group_size, hidden, swiglu_limit):
+    """The post-quant kernels keep SiLU and the multiply in FP32 until FP8.
+
+    Unlike the generic fused quantizer below, there is no intermediate BF16
+    round. Reuse the independent UE8M0 oracle, and compare both layouts only
+    on active rows; an empty expert and a partial slab exercise masked counts.
+    """
+    from sglang.kernels.ops.attention.dsv4 import (
+        silu_and_mul_contig_post_quant,
+        silu_and_mul_masked_post_quant,
+    )
+
+    torch.manual_seed(123 + hidden)
+    experts, capacity = 3, 32
+    x = (
+        torch.randn(experts, capacity, hidden * 2, device="cuda", dtype=torch.bfloat16)
+        * 5
+    )
+    x[2, 0].zero_()
+    counts = torch.tensor([0, 17, 9], device="cuda", dtype=torch.int32)
+    gate, up = x.float().chunk(2, dim=-1)
+    if swiglu_limit is not None:
+        gate = gate.clamp_max(swiglu_limit)
+        up = up.clamp(-swiglu_limit, swiglu_limit)
+    activation = gate * torch.sigmoid(gate) * up
+    q_ref, exp_ref = ref_fp8_ue8m0(activation, group_size)
+
+    flat = x.flatten(0, 1)
+    q = torch.empty(experts * capacity, hidden, device="cuda", dtype=fp8_dtype)
+    scale = create_per_token_group_quant_fp8_output_scale(
+        x_shape=q.shape,
+        device="cuda",
+        group_size=group_size,
+        column_major_scales=True,
+        scale_tma_aligned=True,
+        scale_ue8m0=True,
+    )
+    silu_and_mul_contig_post_quant(
+        flat,
+        q,
+        scale,
+        group_size,
+        scale_ue8m0=True,
+        transposed=True,
+        swiglu_limit=swiglu_limit,
+    )
+    masked_q = torch.empty_like(q).view(experts, capacity, hidden)
+    masked_scale = torch.empty(
+        experts,
+        hidden // group_size // 4,
+        capacity,
+        device="cuda",
+        dtype=torch.int32,
+    )
+    silu_and_mul_masked_post_quant(
+        x,
+        masked_q,
+        masked_scale,
+        group_size,
+        counts,
+        scale_ue8m0=True,
+        transposed=True,
+        swiglu_limit=swiglu_limit,
+    )
+    exp = _decode_packed_exp(scale, hidden // group_size).view_as(exp_ref)
+    masked_exp = _decode_packed_exp(masked_scale.transpose(1, 2), hidden // group_size)
+    q = q.view_as(q_ref)
+    for expert, count in enumerate(counts.tolist()):
+        assert torch.equal(exp[expert, :count], exp_ref[expert, :count])
+        assert torch.equal(masked_exp[expert, :count], exp[expert, :count])
+        assert torch.equal(
+            masked_q[expert, :count].view(torch.uint8),
+            q[expert, :count].view(torch.uint8),
+        )
+        # Fast sigmoid may differ from torch by an FP32 ULP at an FP8
+        # rounding boundary; bound the resulting error, not arbitrary bytes.
+        torch.testing.assert_close(
+            q[expert, :count].float(),
+            q_ref[expert, :count].float(),
+            rtol=0.125,
+            atol=2**-9,
+        )
+        if count:
+            mismatch = (
+                (
+                    q[expert, :count].view(torch.uint8)
+                    != q_ref[expert, :count].view(torch.uint8)
+                )
+                .float()
+                .mean()
+            )
+            # A rare fast-math boundary flip is allowed, but systematic BF16
+            # intermediate rounding (the other fused path) must fail this gate.
+            assert mismatch.item() < 1e-4
+
+
 @pytest.mark.parametrize("column_major", [True, False])
 @pytest.mark.parametrize("scale_ue8m0", [True, False])
 def test_fused_silu(scale_ue8m0, column_major):

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

@@ -144,6 +144,7 @@ class TestDeepEPv2BufferLifecycle(CustomTestCase):
         impl = object.__new__(deepep_v2._DeepEPv2Impl)
         impl.num_max_dispatch_tokens_per_rank = 4
         impl.hidden_size = 128
+        impl.activation_scale_block_size = 128
         impl.router_topk = 2
         impl._validate_common(torch.empty(4, 128), torch.zeros(4, 2))
         with self.assertRaisesRegex(ValueError, "per-rank buffer capacity"):

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

@@ -1,6 +1,7 @@
 """Tests for the DeepEP v2 expanded/masked repack kernels."""
 
 import unittest
+from types import SimpleNamespace
 
 import torch
 
@@ -113,6 +114,80 @@ class TestDeepEPv2MaskedSlab(CustomTestCase):
     def test_empty_experts(self):
         self._check_expand_roundtrip([0, 0, 0, 0], torch.bfloat16, with_scale=False)
 
+    def test_runner_defers_expanded_route_weighting(self):
+        from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
+        from sglang.srt.layers.moe.moe_runner.deep_gemm import (
+            DeepGemmRunnerOutput,
+            post_permute_deep_gemm_to_deepep_v2,
+        )
+        from sglang.srt.layers.moe.token_dispatcher.base import RoutewiseLayout
+
+        counts = [3, 0, 2]
+        recv_x, _, psum, starts, total = _build_layout(
+            counts, self.ALIGN, self.HIDDEN, torch.bfloat16
+        )
+        masked_x, _, _ = expand_to_masked_slab(
+            recv_x, None, psum, len(counts), self.MAX_M, self.ALIGN
+        )
+        weights = torch.full((total,), 0.25, device=DEVICE)
+        state = {
+            "deepep_v2_expanded": True,
+            "deepep_v2_masked": True,
+            "deepep_v2_psum": psum,
+            "deepep_v2_total_expanded": total,
+            "deepep_v2_expert_alignment": self.ALIGN,
+            "topk_weights": weights,
+        }
+        rows = _real_rows(starts, counts)
+        for no_combine in (False, True):
+            with self.subTest(no_combine=no_combine):
+                output = post_permute_deep_gemm_to_deepep_v2(
+                    DeepGemmRunnerOutput(masked_x),
+                    None,
+                    MoeRunnerConfig(no_combine=no_combine),
+                    state,
+                )
+                expected = recv_x[rows] if no_combine else recv_x[rows] * 0.25
+                self.assertTrue(torch.equal(output.hidden_states[rows], expected))
+                self.assertEqual(
+                    output.routewise_layout,
+                    RoutewiseLayout.EXPANDED if no_combine else None,
+                )
+                if no_combine:
+                    self.assertTrue(torch.equal(output.topk_weights, weights))
+
+    def test_runner_restores_token_topk_routes_and_masks_nonlocal_slots(self):
+        from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
+        from sglang.srt.layers.moe.moe_runner.deep_gemm import (
+            DeepGemmRunnerOutput,
+            post_permute_deep_gemm_to_deepep_v2,
+        )
+        from sglang.srt.layers.moe.token_dispatcher.base import RoutewiseLayout
+
+        hidden = torch.tensor([[3.0], [5.0], [7.0]], device=DEVICE)
+        weights = torch.tensor([[0.25, 0.0], [0.5, 0.75]], device=DEVICE)
+        state = {
+            "topk_ids": torch.tensor([[0, -1], [1, 0]], device=DEVICE),
+            "topk_weights": weights,
+            "output_index": torch.tensor([[1, -1], [0, 2]], device=DEVICE),
+        }
+        for empty in (False, True):
+            with self.subTest(empty=empty):
+                if empty:
+                    state["output_index"] = torch.full((2, 2), -1, device=DEVICE)
+                output = post_permute_deep_gemm_to_deepep_v2(
+                    DeepGemmRunnerOutput(hidden[:0] if empty else hidden),
+                    None,
+                    MoeRunnerConfig(no_combine=True),
+                    state,
+                )
+                expected = torch.tensor([[[5.0], [0.0]], [[3.0], [7.0]]], device=DEVICE)
+                if empty:
+                    expected.zero_()
+                self.assertTrue(torch.equal(output.hidden_states, expected))
+                self.assertIs(output.topk_weights, weights)
+                self.assertEqual(output.routewise_layout, RoutewiseLayout.TOKEN_TOPK)
+
     def test_single_hot_expert(self):
         self._check_expand_roundtrip(
             [0, self.MAX_M, 0, 0], torch.bfloat16, with_scale=False, topk=True
@@ -133,7 +208,7 @@ class TestDeepEPv2MaskedSlab(CustomTestCase):
                 recv_x, None, psum, len(counts), self.MAX_M, self.ALIGN
             )
 
-    def _production_packed_ue8m0_layout(self, counts):
+    def _production_packed_ue8m0_layout(self, counts, group_size):
         """Build expanded rows with the production packed UE8M0 quantizer."""
         from sglang.kernels.ops.quantization.fp8_kernel import (
             sglang_per_token_group_quant_fp8,
@@ -146,7 +221,7 @@ class TestDeepEPv2MaskedSlab(CustomTestCase):
         )
         recv_x, recv_x_scale = sglang_per_token_group_quant_fp8(
             raw,
-            128,
+            group_size,
             column_major_scales=True,
             scale_tma_aligned=True,
             scale_ue8m0=True,
@@ -157,9 +232,14 @@ class TestDeepEPv2MaskedSlab(CustomTestCase):
         return recv_x, recv_x_scale, psum, starts, total, hidden
 
     def test_fp8_packed_ue8m0_scale_from_production_quantizer(self):
+        for group_size in (32, 128):
+            with self.subTest(group_size=group_size):
+                self._check_packed_scale(group_size)
+
+    def _check_packed_scale(self, group_size):
         counts = [3, 1, 6, 2]
         recv_x, recv_x_scale, psum, starts, _, hidden = (
-            self._production_packed_ue8m0_layout(counts)
+            self._production_packed_ue8m0_layout(counts, group_size)
         )
         E = len(counts)
         masked_x, masked_x_scale, masked_m = expand_to_masked_slab(
@@ -175,10 +255,15 @@ class TestDeepEPv2MaskedSlab(CustomTestCase):
                 torch.testing.assert_close(masked_x_scale[e, j], recv_x_scale[s + j])
 
     def test_expand_under_cuda_graph_capture(self):
+        for group_size in (32, 128):
+            with self.subTest(group_size=group_size):
+                self._check_graph_capture(group_size)
+
+    def _check_graph_capture(self, group_size):
         # Exercise replay with the production packed scale layout.
         counts = [3, 1, 6, 2]
         recv_x, recv_x_scale, psum, starts, _, _ = self._production_packed_ue8m0_layout(
-            counts
+            counts, group_size
         )
         E = len(counts)
         warm = torch.cuda.Stream()
@@ -235,10 +320,34 @@ class TestDeepEPv2HandleLifecycle(CustomTestCase):
 
         impl._get_buffer = _boom
         with self.assertRaisesRegex(RuntimeError, "boom"):
-            impl.combine(None)
+            impl.combine(SimpleNamespace(routewise_layout=None))
         self.assertIsNone(impl._handle)
         self.assertFalse(impl._pad_empty_combine)
 
+    def test_unfinalized_routes_are_rejected_and_release_handle(self):
+        from sglang.srt.layers.moe.token_dispatcher.base import (
+            CombineInputChecker,
+            RoutewiseLayout,
+        )
+        from sglang.srt.layers.moe.token_dispatcher.deepep_v2 import (
+            DeepEPv2CombineInput,
+        )
+
+        impl = self._bare_impl()
+        impl._handle = object()
+        output = DeepEPv2CombineInput(
+            torch.empty(2, 8), torch.ones(2), RoutewiseLayout.EXPANDED
+        )
+        self.assertTrue(CombineInputChecker.needs_model_route_finalization(output))
+        with self.assertRaisesRegex(ValueError, "model route finalization"):
+            impl.combine(output)
+        self.assertIsNone(impl._handle)
+        self.assertFalse(
+            CombineInputChecker.needs_model_route_finalization(
+                DeepEPv2CombineInput(torch.empty(2, 8), None)
+            )
+        )
+
 
 if __name__ == "__main__":
     unittest.main()

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

@@ -52,6 +52,13 @@ class _FakeBuffer:
         num_recv = (x[0] if isinstance(x, tuple) else x).shape[0]
         topk_idx = kwargs["topk_idx"]
         topk_weights = kwargs["topk_weights"]
+        if kwargs["do_expand"]:
+            if isinstance(x, tuple):
+                x = tuple(t.repeat_interleave(TOPK, dim=0) for t in x)
+            else:
+                x = x.repeat_interleave(TOPK, dim=0)
+            topk_idx = None
+            topk_weights = topk_weights.flatten()
         event = SimpleNamespace(event=None, current_stream_wait=lambda: None)
         return x, topk_idx, topk_weights, _FakeHandle(num_recv), event
 
@@ -93,7 +100,9 @@ class _DeepEPv2WireDtypeBase(CustomTestCase):
         for item in reversed(self._patches):
             item.stop()
 
-    def _dispatch(self, use_fp8_dispatch, num_tokens=8, is_extend_in_batch=True):
+    def _dispatch(
+        self, use_fp8_dispatch, num_tokens=8, is_extend_in_batch=True, group_size=128
+    ):
         dispatcher = deepep_v2.DeepEPv2Dispatcher(
             group=_FakeGroup(),
             router_topk=TOPK,
@@ -102,6 +111,7 @@ class _DeepEPv2WireDtypeBase(CustomTestCase):
             hidden_size=HIDDEN,
             params_dtype=torch.bfloat16,
             use_fp8_dispatch=use_fp8_dispatch,
+            activation_scale_block_size=group_size,
         )
         dispatcher._impl.num_max_dispatch_tokens_per_rank = NUM_MAX_TOKENS
         hidden_states = torch.randn((num_tokens, HIDDEN), dtype=torch.bfloat16)
@@ -116,6 +126,18 @@ class _DeepEPv2WireDtypeBase(CustomTestCase):
 
 
 class TestDeepEPv2WireDtype(_DeepEPv2WireDtypeBase):
+    def test_mxfp8_scale_group_survives_both_dispatch_layouts(self):
+        for is_extend in (True, False):
+            with self.subTest(is_extend=is_extend):
+                _, out = self._dispatch(
+                    use_fp8_dispatch=True,
+                    is_extend_in_batch=is_extend,
+                    group_size=32,
+                )
+                self.assertEqual(out.activation_scale_block_size, 32)
+                self.assertEqual(out.hidden_states_scale.shape[-1], HIDDEN // 32)
+                self.assertEqual(out.is_expanded, not is_extend)
+
     def test_bf16_dispatch_sends_unquantized_activations(self):
         hidden_states, out = self._dispatch(use_fp8_dispatch=False)
         self.assertIs(_FakeBuffer.last.dispatch_x, hidden_states)

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

@@ -76,13 +76,18 @@ def _fp8_method(**overrides):
     return method
 
 
-def test_deepep_v2_quant_contract_accepts_blockwise_fp8(_moe_flags):
+@pytest.mark.parametrize("use_mxfp8,block_size", [(False, [128, 128]), (True, [1, 32])])
+def test_deepep_v2_quant_contract_accepts_blockwise_fp8(
+    _moe_flags, use_mxfp8, block_size
+):
     from sglang.srt.layers.moe.fused_moe_triton.layer import (
         _validate_deepep_v2_quant_method,
     )
 
     _moe_flags.a2a_backend = MoeA2ABackend.DEEPEP_V2
-    _validate_deepep_v2_quant_method(_fp8_method(weight_block_size=[128, 128]))
+    _validate_deepep_v2_quant_method(
+        _fp8_method(weight_block_size=block_size, use_mxfp8=use_mxfp8)
+    )
 
 
 @pytest.mark.parametrize(
@@ -90,7 +95,7 @@ def test_deepep_v2_quant_contract_accepts_blockwise_fp8(_moe_flags):
     [
         ({"activation_scheme": "static"}, "activation_scheme"),
         ({"weight_block_size": None}, "weight_block_size"),
-        ({"weight_block_size": (1, 32), "use_mxfp8": True}, "MXFP8"),
+        ({"weight_block_size": (128, 128), "use_mxfp8": True}, "MXFP8"),
         ({"is_fp4_expert": True}, "FP4 experts"),
     ],
 )
@@ -149,5 +154,46 @@ def test_deepep_v2_runner_backstop(_moe_flags):
     assert MoeRunner(MoeRunnerBackend.DEEP_GEMM, MoeRunnerConfig()).runner_core
 
 
+def test_deepep_v2_registration_uses_primary_architecture_and_rejects_conflicts():
+    from sglang.srt.configs.moe_model_registry import (
+        model_requires_fp32_silu_mul,
+        model_supports_deepep_v2,
+        register_deepep_v2_model,
+    )
+
+    register_deepep_v2_model("TestRoutewiseMoe", silu_mul_keep_fp32=True)
+    config = SimpleNamespace(architectures=["TestRoutewiseMoe"])
+    assert model_supports_deepep_v2(config)
+    assert model_requires_fp32_silu_mul(config)
+    config.architectures = ["UnsupportedMoe", "TestRoutewiseMoe"]
+    assert not model_supports_deepep_v2(config)
+    assert not model_requires_fp32_silu_mul(config)
+    with pytest.raises(ValueError, match="Conflicting"):
+        register_deepep_v2_model("TestRoutewiseMoe", silu_mul_keep_fp32=False)
+
+
+@pytest.mark.parametrize("block_size,width", [(32, 16), (128, 4)])
+def test_mxfp8_recipes_keep_activation_and_weight_groups_separate(block_size, width):
+    # The packed format is a layout contract, independent of the host GPU.
+    from unittest.mock import patch
+
+    from sglang.srt.layers import deep_gemm_wrapper
+    from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmMoeQuantInfo
+
+    with patch.object(deep_gemm_wrapper, "DEEPGEMM_SCALE_UE8M0", True):
+        quant = DeepGemmMoeQuantInfo(
+            None, None, True, block_shape=[1, 32], use_mxfp8=True
+        )
+    assert quant.scale_recipes(
+        activation_block_size=block_size, hidden_size=2048, activation_scale_width=width
+    ) == ((1, block_size), (1, 32))
+    with pytest.raises(AssertionError, match="activation scale mismatch"):
+        quant.scale_recipes(
+            activation_block_size=block_size,
+            hidden_size=2048,
+            activation_scale_width=width + 1,
+        )
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))