[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