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):