feat(moriep): add fp4 combine dtype (SGLANG_MORI_COMBINE_DTYPE=fp4) (#30706)

docker/rocm.Dockerfile

@@ -110,7 +110,7 @@ ARG ENABLE_MORI=0
 ARG NIC_BACKEND=none
 
 ARG MORI_REPO="https://github.com/ROCm/mori.git"
-ARG MORI_COMMIT="e31d426a13e96e1cbff96a1c904d291aefe8c46a"
+ARG MORI_COMMIT="6f072775a73518440b29be60e85dda70db63ca43"
 
 # NIXL (upstream ai-dynamo/nixl) — KV transfer backend for prefill/decode disaggregation.
 # Built from source for ROCm; needs UCX built --with-rocm (built here from openucx).

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 @@ class _MoriEPDispatcherImplBase:
                     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

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()