diff --git a/python/sglang/srt/layers/moe/moe_runner/aiter.py b/python/sglang/srt/layers/moe/moe_runner/aiter.py index 07020e62b..8737b93c2 100644 --- a/python/sglang/srt/layers/moe/moe_runner/aiter.py +++ b/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( diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index 2f2e04d8a..2616c148f 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/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 diff --git a/test/registered/unit/layers/test_moriep_mxfp8_dispatch.py b/test/registered/unit/layers/test_moriep_mxfp8_dispatch.py new file mode 100644 index 000000000..c8d89d9d1 --- /dev/null +++ b/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"]))