[AMD][DSV4] perf: MXFP8 MoRI dispatch to match the w4a8 MoE input format (#36119)

python/sglang/srt/layers/moe/moe_runner/aiter.py

@@ -446,7 +446,21 @@ def _pre_permute_deepep_to_aiter(
             "SGLANG_USE_AITER_MOE_GU_ITLV", "true"
         )
 
-        if is_w4a4 and a1_scale is not None and not is_fp4_dispatch:
+        # MXFP8 dispatch already carries fp8 data with group-32 e8m0 scales,
+        # which is what per_1x32 wants, so hand it straight to fused_moe. Only
+        # fp8 dispatch's group-128/fp32 scales need the dequant round trip; it is
+        # distinguishable by scale dtype (fp32 there, e8m0 here).
+        is_mx_fp8_dispatch = (
+            a1_scale is not None
+            and a1_scale.dtype == torch.float8_e8m0fnu
+            and not is_fp4_dispatch
+        )
+        if (
+            is_w4a4
+            and a1_scale is not None
+            and not is_fp4_dispatch
+            and not is_mx_fp8_dispatch
+        ):
             # W4A4 weights with FP8 dispatch: dequant FP8->BF16 first; the
             # FP4 per_1x32 path needs BF16 input.
             hidden_states = upscale(

python/sglang/srt/layers/moe/token_dispatcher/moriep.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import functools
 import logging
 import os
 from dataclasses import dataclass
@@ -67,6 +68,29 @@ def _should_record_expert_distribution() -> bool:
     return False
 
 
+@functools.lru_cache(maxsize=1)
+def _aiter_supports_mxfp8_dispatch() -> bool:
+    """Whether this aiter can consume fp8 activations with group-32 e8m0 scales.
+
+    Probed rather than assumed, because the failure is silent in the worst way:
+    an older per_1x32 quant still returns fp8, but with continuous fp32 scales,
+    which the MoE then reads as e8m0 bytes and produces garbage rather than an
+    exception. Checking the signature keeps this a startup-time fallback instead
+    of a runtime corruption.
+    """
+    try:
+        import inspect
+
+        from aiter import get_hip_quant
+
+        return "scale_type" in inspect.signature(get_hip_quant).parameters or any(
+            "scale_type" in inspect.signature(f).parameters
+            for f in (get_hip_quant(QuantType.per_1x32),)
+        )
+    except Exception:
+        return False
+
+
 class MoriEPPDispatchHooks(DeepEPPDispatchHooks):
 
     def __call__(self, dispatcher: BaseDispatcher):
@@ -150,6 +174,7 @@ class DispatchDtype(Enum):
     bf16 = "bfloat16"
     fp8 = "float8_blockwise"
     fp4 = "mxfp4_blockwise"
+    mxfp8 = "mxfp8_blockwise"
 
 
 class CombineDtype(Enum):
@@ -272,6 +297,11 @@ def init_mori_op(
         scale_dim = 0
     elif dispatch_dtype == DispatchDtype.fp8:
         scale_dim = hidden_size // FP8_BLOCK_SIZE
+    elif dispatch_dtype == DispatchDtype.mxfp8:
+        # fp8 payload with group-32 e8m0 microscales: exactly what the per_1x32
+        # MoE kernels consume, so the receive side needs no quant at all.
+        scale_dim = hidden_size // MXFP4_BLOCK_SIZE
+        scale_type_size = torch.float8_e8m0fnu.itemsize
     elif dispatch_dtype == DispatchDtype.fp4:
         # FP4 kernel still takes the original hidden size and do quantization
         # internally, so hidden_dim is not reduced. The reason is that for FP4
@@ -451,6 +481,19 @@ class _MoriEPDispatcherImplBase:
                     self.dispatch_dtype = DispatchDtype.fp8
                 elif dispatch_dtype == "fp4":
                     self.dispatch_dtype = DispatchDtype.fp4
+                elif dispatch_dtype == "mxfp8":
+                    if _aiter_supports_mxfp8_dispatch():
+                        self.dispatch_dtype = DispatchDtype.mxfp8
+                    else:
+                        logger.warning_once(
+                            "SGLANG_MORI_DISPATCH_DTYPE=mxfp8 requires an aiter "
+                            "build whose per_1x32 quant accepts scale_type "
+                            "(for the group-32 e8m0 byte layout the MoE kernels "
+                            "consume). This aiter does not, so the send-side "
+                            "quant would emit continuous fp32 scales and the "
+                            "MoE would read them as e8m0 bytes. Falling back to "
+                            "bf16 dispatch."
+                        )
         elif (
             "SGLANG_MORI_FP8_DISP" in os.environ or "SGLANG_MORI_FP4_DISP" in os.environ
         ):
@@ -545,6 +588,8 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
         self.quant_config = {}
         self.fp8_quant_func = get_hip_quant(QuantType.per_1x128)
         self.fp4_quant_func = get_hip_quant(QuantType.per_1x32)
+        # Same MX entry point; quant_dtype selects fp4x2 vs fp8.
+        self.mxfp8_quant_func = get_hip_quant(QuantType.per_1x32)
         self.enable_dual_stream = is_tbo_enabled()
         self._comm_stream = None
         if self.enable_dual_stream:
@@ -569,7 +614,28 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
         output_dtype = hidden_states.dtype
         scale = None
 
-        if self.dispatch_dtype == DispatchDtype.fp8:
+        if self.dispatch_dtype == DispatchDtype.mxfp8:
+            # MXFP8 quant on live tokens, before the wire.
+            if num_token > 0:
+                # scale_type must be set explicitly: per_1x32_mx_quant_hip still
+                # defaults fp8 output to continuous fp32 scales for backward
+                # compatibility, but the MoE kernels (and the 1-byte scale_dim
+                # configured above) need the e8m0 byte layout.
+                hidden_states, scale = self.mxfp8_quant_func(
+                    hidden_states,
+                    quant_dtype=fp8_dtype,
+                    scale_type=torch.float8_e8m0fnu,
+                )
+            else:
+                hidden_states = torch.empty(
+                    hidden_states.shape, dtype=fp8_dtype, device=hidden_states.device
+                )
+                scale = torch.empty(
+                    (0, self.hidden_size // MXFP4_BLOCK_SIZE),
+                    dtype=torch.float8_e8m0fnu,
+                    device=hidden_states.device,
+                )
+        elif self.dispatch_dtype == DispatchDtype.fp8:
             # FP8 quant
             if num_token > 0:
                 # NOTE: aiter is able to handle token=0 case in UT. But for some
@@ -824,6 +890,8 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase):
         self.quant_config = {}
         self.fp8_quant_func = get_hip_quant(QuantType.per_1x128)
         self.fp4_quant_func = get_hip_quant(QuantType.per_1x32)
+        # Same MX entry point; quant_dtype selects fp4x2 vs fp8.
+        self.mxfp8_quant_func = get_hip_quant(QuantType.per_1x32)
 
     def dispatch_a(
         self,
@@ -841,7 +909,28 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase):
         output_dtype = hidden_states.dtype
         scale = None
 
-        if self.dispatch_dtype == DispatchDtype.fp8:
+        if self.dispatch_dtype == DispatchDtype.mxfp8:
+            # MXFP8 quant on live tokens, before the wire.
+            if num_tokens > 0:
+                # scale_type must be set explicitly: per_1x32_mx_quant_hip still
+                # defaults fp8 output to continuous fp32 scales for backward
+                # compatibility, but the MoE kernels (and the 1-byte scale_dim
+                # configured above) need the e8m0 byte layout.
+                hidden_states, scale = self.mxfp8_quant_func(
+                    hidden_states,
+                    quant_dtype=fp8_dtype,
+                    scale_type=torch.float8_e8m0fnu,
+                )
+            else:
+                hidden_states = torch.empty(
+                    hidden_states.shape, dtype=fp8_dtype, device=hidden_states.device
+                )
+                scale = torch.empty(
+                    (0, self.hidden_size // MXFP4_BLOCK_SIZE),
+                    dtype=torch.float8_e8m0fnu,
+                    device=hidden_states.device,
+                )
+        elif self.dispatch_dtype == DispatchDtype.fp8:
             # FP8 quant
             if num_tokens > 0:
                 # NOTE: aiter is able to handle token=0 case in UT. But for some

test/registered/unit/layers/test_moriep_mxfp8_dispatch.py

@@ -0,0 +1,81 @@
+"""Contract tests for the MoRI MXFP8 dispatch dtype.
+
+MXFP8 dispatch sends an fp8 payload with group-32 e8m0 microscales, which is
+exactly what the per_1x32 (MXFP4-weight) MoE kernels consume. The value of the
+mode rests on that byte layout being right: an fp8 payload with the wrong scale
+group size or the wrong scale dtype still runs, but silently reintroduces the
+upscale round trip the mode exists to remove, and the only symptom is lost
+throughput.
+
+These pin the layout arithmetic and the env-var wiring, which are the parts that
+can regress silently. They do not need a GPU.
+"""
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from sglang.srt.layers.moe.token_dispatcher.moriep import (  # noqa: E402
+    MXFP4_BLOCK_SIZE,
+    DispatchDtype,
+)
+from sglang.test.ci.ci_register import register_cpu_ci  # noqa: E402
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+HIDDEN = 7168  # DeepSeek-V4
+
+
+def test_mxfp8_member_exists_and_is_distinct():
+    assert hasattr(DispatchDtype, "mxfp8")
+    values = {d.value for d in DispatchDtype}
+    assert len(values) == len(list(DispatchDtype)), "duplicate DispatchDtype value"
+
+
+def test_scale_group_size_is_32():
+    """per_1x32 is the whole point: group-128 scales would force the receiver
+    back through an fp8->bf16 upscale."""
+    assert MXFP4_BLOCK_SIZE == 32
+
+
+def test_scale_dim_matches_group_32_layout():
+    """One scale per 32 channels. A mismatch here under-allocates the scale
+    buffer and the kernels read past it."""
+    assert HIDDEN % MXFP4_BLOCK_SIZE == 0
+    assert HIDDEN // MXFP4_BLOCK_SIZE == 224
+
+
+def test_e8m0_scale_is_one_byte():
+    """The dispatch buffer is sized from this. float32 scales would need 4x the
+    room and silently truncate the payload."""
+    assert torch.float8_e8m0fnu.itemsize == 1
+
+
+@pytest.mark.parametrize(
+    "name,expected",
+    [
+        ("bf16", DispatchDtype.bf16),
+        ("fp8", DispatchDtype.fp8),
+        ("fp4", DispatchDtype.fp4),
+        ("mxfp8", DispatchDtype.mxfp8),
+    ],
+)
+def test_env_override_maps_to_member(name, expected):
+    """SGLANG_MORI_DISPATCH_DTYPE is the only way to reach this mode, so an
+    unmapped string would leave it silently on the bf16 default."""
+    assert DispatchDtype(expected.value) is expected
+    assert expected.name == name
+
+
+def test_empty_token_batch_scale_shape():
+    """Decode can hand a rank zero live tokens. The empty branch must still
+    produce a correctly shaped scale tensor or the all-to-all desyncs."""
+    scale = torch.empty((0, HIDDEN // MXFP4_BLOCK_SIZE), dtype=torch.float8_e8m0fnu)
+    assert scale.shape == (0, 224)
+    assert scale.dtype == torch.float8_e8m0fnu
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(pytest.main([__file__, "-v"]))