diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index 82f13643cac4..5bc48d8d97b3 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py @@ -156,6 +156,7 @@ class CombineDtype(Enum): bf16 = "bfloat16" fp8 = "float8_blockwise" fp8_direct_cast = "float8_direct_cast" + fp4 = "fp4_blockwise" # packed E2M1, blockwise-scaled; ~half the combine transport of fp8 @dataclass(frozen=True) @@ -297,6 +298,8 @@ def init_mori_op( combine_quant_type = "fp8_blockwise" elif combine_dtype == CombineDtype.fp8_direct_cast: combine_quant_type = "fp8_direct_cast" + elif combine_dtype == CombineDtype.fp4: + combine_quant_type = "fp4_blockwise" logger.info( f"[MORI init] {world_size=} {rank=} {hidden_size=} {params_dtype=} " @@ -471,12 +474,14 @@ def _apply_dispatch_dtype_override(self): self.combine_dtype = CombineDtype.bf16 elif combine_dtype == "fp8_direct_cast": self.combine_dtype = CombineDtype.fp8_direct_cast + elif combine_dtype == "fp4": + self.combine_dtype = CombineDtype.fp4 elif "SGLANG_MORI_FP8_COMB" in os.environ: # Deprecated: will be removed in a future release logger.warning_once( "SGLANG_MORI_FP8_COMB is deprecated " "and will be removed in a future release. " - "Use SGLANG_MORI_COMBINE_DTYPE=auto|bf16|fp8|fp8_direct_cast instead." + "Use SGLANG_MORI_COMBINE_DTYPE=auto|bf16|fp8|fp4|fp8_direct_cast instead." ) if get_bool_env_var("SGLANG_MORI_FP8_COMB", "False"): self.combine_dtype = CombineDtype.fp8 diff --git a/test/manual/ep/test_moriep_combine_dtype.py b/test/manual/ep/test_moriep_combine_dtype.py new file mode 100644 index 000000000000..36b71e35e6e0 --- /dev/null +++ b/test/manual/ep/test_moriep_combine_dtype.py @@ -0,0 +1,28 @@ +"""Lightweight (no-GPU) wiring guards for the MoRI blockwise-FP4 combine dtype. + +Protects the sglang-side plumbing that exposes ``SGLANG_MORI_COMBINE_DTYPE=fp4`` and maps it to +MoRI's ``fp4_blockwise`` combine. The end-to-end kernel behaviour is covered by the GPU/eval tests. +""" + +import unittest + +from sglang.srt.layers.moe.token_dispatcher.moriep import CombineDtype + + +class TestMoriCombineDtypeFp4(unittest.TestCase): + def test_fp4_enum_member_exists(self): + self.assertTrue(hasattr(CombineDtype, "fp4")) + self.assertEqual(CombineDtype.fp4.value, "fp4_blockwise") + + def test_fp4_value_roundtrip(self): + self.assertIs(CombineDtype("fp4_blockwise"), CombineDtype.fp4) + + def test_existing_combine_dtypes_unchanged(self): + # Guard against accidentally breaking the existing dtypes when adding fp4. + self.assertEqual(CombineDtype.bf16.value, "bfloat16") + self.assertEqual(CombineDtype.fp8.value, "float8_blockwise") + self.assertEqual(CombineDtype.fp8_direct_cast.value, "float8_direct_cast") + + +if __name__ == "__main__": + unittest.main()