Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 6 additions & 1 deletion python/sglang/srt/layers/moe/token_dispatcher/moriep.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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=} "
Expand Down Expand Up @@ -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
Expand Down
28 changes: 28 additions & 0 deletions test/manual/ep/test_moriep_combine_dtype.py
Original file line number Diff line number Diff line change
@@ -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()
Loading