[AMD][MoRI] Auto-select dispatch quantization type from MoE weight dtype. (#21040)

docs/references/environment_variables.md

@@ -74,8 +74,7 @@ SGLang supports various environment variables that can be used to configure its
 
 | Environment Variable | Description | Default Value |
 | --- | --- | --- |
-| `SGLANG_MORI_FP8_DISP` | Use FP8 for dispatch | `"false"` |
+| `SGLANG_MORI_DISPATCH_DTYPE` | Override MoRI-EP dispatch quantization type. `auto` uses auto-detection from weight dtype; `bf16`/`fp8`/`fp4` forces the specified type for all layers | `"auto"` |
-| `SGLANG_MORI_FP4_DISP` | Use MXFP4 for dispatch | `"false"` |
 | `SGLANG_MORI_FP8_COMB` | Use FP8 for combine | `"false"` |
 | `SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK` | Maximum number of dispatch tokens per rank for MORI-EP buffer allocation | `4096` |
 | `SGLANG_MORI_DISPATCH_INTER_KERNEL_SWITCH_THRESHOLD` | Threshold for switching between `InterNodeV1` and `InterNodeV1LL` kernel types. `InterNodeV1LL` is used if `SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK` is less than or equal to this threshold; otherwise, `InterNodeV1` is used. | `256` |

python/sglang/srt/layers/moe/token_dispatcher/moriep.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple
 
@@ -170,22 +171,9 @@ def get_ep_dispatch_configs(num_max_dispatch_tokens_per_rank: int = 4096):
     }
 
 
-@lru_cache(maxsize=2)
-def _get_mori_dispatch_quant_flags():
-    fp8_dispatch = get_bool_env_var("SGLANG_MORI_FP8_DISP", "False")
-    fp4_dispatch = get_bool_env_var("SGLANG_MORI_FP4_DISP", "False")
-    if fp8_dispatch and fp4_dispatch:
-        logger.warning(
-            "Both SGLANG_MORI_FP8_DISP and SGLANG_MORI_FP4_DISP are set to True. "
-            "Using SGLANG_MORI_FP4_DISP and ignoring SGLANG_MORI_FP8_DISP."
-        )
-        fp8_dispatch = False
-    return fp8_dispatch, fp4_dispatch
-
-
 # init_mori_op only needs do once in model initial stage
 # use lru_cache to reuse the same mori_op instance to avoid the init overhead for mori
-@lru_cache(maxsize=2)
+@lru_cache(maxsize=4)
 def init_mori_op(
     group,
     router_topk,
@@ -196,6 +184,8 @@ def init_mori_op(
     num_max_dispatch_tokens_per_rank,
     deepep_mode,
     instance_id=0,
+    fp8_dispatch=False,
+    fp4_dispatch=False,
 ):
 
     import mori
@@ -228,12 +218,6 @@ def init_mori_op(
     if async_mode:
         mode = EpMode.LOW_LATENCY
 
-    logger.info(
-        f"[MORI init] {world_size=} {rank=} {hidden_size=} {params_dtype=} "
-        f"{num_max_dispatch_tokens_per_rank=} {num_local_experts=} "
-        f"{router_topk=} {mode=}"
-    )
-
     cfg = get_ep_dispatch_configs(num_max_dispatch_tokens_per_rank)[mode]
 
     kernel_type = cfg.kernel_type
@@ -246,8 +230,6 @@ def init_mori_op(
     data_type = fp8_dtype
     scale_type_size = torch.float32.itemsize
 
-    fp8_dispatch, fp4_dispatch = _get_mori_dispatch_quant_flags()
-
     if fp8_dispatch:
         scale_dim = hidden_size // 128
     elif fp4_dispatch:
@@ -272,6 +254,13 @@ def init_mori_op(
     if get_bool_env_var("SGLANG_MORI_FP8_COMB", "False"):
         combine_quant_type = "fp8_direct_cast"
 
+    logger.info(
+        f"[MORI init] {world_size=} {rank=} {hidden_size=} {params_dtype=} "
+        f"{num_max_dispatch_tokens_per_rank=} {num_local_experts=} "
+        f"{router_topk=} {mode=} {fp8_dispatch=} {fp4_dispatch=} "
+        f"{combine_quant_type=}"
+    )
+
     mori_config = mori.ops.EpDispatchCombineConfig(
         rank=rank,
         world_size=world_size,
@@ -348,7 +337,22 @@ class _MoriEPDispatcherImplBase:
             "SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK", 4096
         )
 
-        self.mori_op = init_mori_op(
+        self._mori_op = None
+        self.fp8_dispatch = False
+        self.fp4_dispatch = False
+
+        self.quant_config: Optional[dict] = None
+
+        self.overlap_args: Optional[CombineOverlapArgs] = None
+        self.meta_overlap_args: Optional[dict] = None
+
+    @property
+    def mori_op(self):
+        if self._mori_op is None:
+            # If set_quant_config was never called, apply env var override now
+            if self.quant_config is None:
+                self._apply_dispatch_dtype_override()
+            self._mori_op = init_mori_op(
                 self.group,
                 self.router_topk,
                 self.num_experts,
@@ -358,12 +362,29 @@ class _MoriEPDispatcherImplBase:
                 self.num_max_dispatch_tokens_per_rank,
                 self.deepep_mode,
                 self.instance_id,
+                self.fp8_dispatch,
+                self.fp4_dispatch,
             )
+        return self._mori_op
 
-        self.quant_config: Optional[dict] = None
+    def _apply_dispatch_dtype_override(self):
-
+        """Apply env var override to fp8_dispatch/fp4_dispatch flags."""
-        self.overlap_args: Optional[CombineOverlapArgs] = None
+        if "SGLANG_MORI_DISPATCH_DTYPE" in os.environ:
-        self.meta_overlap_args: Optional[dict] = None
+            dispatch_dtype = os.environ["SGLANG_MORI_DISPATCH_DTYPE"].lower()
+            if dispatch_dtype != "auto":
+                self.fp8_dispatch = dispatch_dtype == "fp8"
+                self.fp4_dispatch = dispatch_dtype == "fp4"
+        elif (
+            "SGLANG_MORI_FP8_DISP" in os.environ or "SGLANG_MORI_FP4_DISP" in os.environ
+        ):
+            # Deprecated: will be removed in a future release
+            logger.warning_once(
+                "SGLANG_MORI_FP8_DISP and SGLANG_MORI_FP4_DISP are deprecated "
+                "and will be removed in a future release. "
+                "Use SGLANG_MORI_DISPATCH_DTYPE=auto|bf16|fp8|fp4 instead."
+            )
+            self.fp8_dispatch = get_bool_env_var("SGLANG_MORI_FP8_DISP", "False")
+            self.fp4_dispatch = get_bool_env_var("SGLANG_MORI_FP4_DISP", "False")
 
     def dispatch_a(
         self,
@@ -388,6 +409,19 @@ class _MoriEPDispatcherImplBase:
 
     def set_quant_config(self, quant_config: dict) -> None:
         self.quant_config = quant_config
+        # Auto-detect dispatch quantization from weight dtype
+        weight_dtype = quant_config.get("weight_dtype", None)
+        if weight_dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz):
+            self.fp8_dispatch = True
+            self.fp4_dispatch = False
+        elif weight_dtype == torch.float4_e2m1fn_x2:
+            self.fp8_dispatch = False
+            self.fp4_dispatch = True
+        else:
+            self.fp8_dispatch = False
+            self.fp4_dispatch = False
+        # Apply env var override immediately so dispatch_a sees correct flags
+        self._apply_dispatch_dtype_override()
 
     def set_overlap_args(
         self, combine_overlap_args: CombineOverlapArgs, meta_overlap_args: dict
@@ -432,7 +466,7 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
         output_dtype = hidden_states.dtype
         scale = None
 
-        fp8_dispatch, fp4_dispatch = _get_mori_dispatch_quant_flags()
+        fp8_dispatch, fp4_dispatch = self.fp8_dispatch, self.fp4_dispatch
 
         if fp8_dispatch:
             # FP8 quant
@@ -652,7 +686,7 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
         return combined_hidden_states, done_event
 
     def set_quant_config(self, quant_config: dict):
-        self.quant_config = quant_config
+        super().set_quant_config(quant_config)
 
 
 class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase):
@@ -678,7 +712,7 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase):
         output_dtype = hidden_states.dtype
         scale = None
 
-        fp8_dispatch, fp4_dispatch = _get_mori_dispatch_quant_flags()
+        fp8_dispatch, fp4_dispatch = self.fp8_dispatch, self.fp4_dispatch
 
         if fp8_dispatch:
             # FP8 quant
@@ -828,7 +862,7 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase):
         return combined_hidden_states
 
     def set_quant_config(self, quant_config: dict):
-        self.quant_config = quant_config
+        super().set_quant_config(quant_config)
 
 
 @dataclass

python/sglang/srt/layers/quantization/fp8.py

@@ -1232,15 +1232,13 @@ class Fp8MoEMethod(FusedMoEMethodBase):
     def process_weights_after_loading(self, layer: Module) -> None:
         if _is_hip and _use_hip_int4:
             self.process_weights_hip_int4(layer)
-            return
 
+        elif self.block_quant:
             # Block quant doesn't need to process weights after loading
-        if self.block_quant:
             self.process_weights_after_loading_block_quant(layer)
-            return
 
         # If checkpoint is fp16 or bfloat16, quantize in place.
-        if not self.quant_config.is_checkpoint_fp8_serialized:
+        elif not self.quant_config.is_checkpoint_fp8_serialized:
             # If ROCm, fp8_dtype will be float8_e4m3fnuz (MI300x HW)
             w13_weight = torch.empty_like(layer.w13_weight.data, dtype=fp8_dtype)
             w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype)
@@ -1267,7 +1265,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):
 
             if _is_hip:
                 self.process_weights_hip_scale_padding(layer)
-            return
 
         # If checkpoint is fp8, we need to handle that the
         # MoE kernels require single activation scale and single weight
@@ -1358,7 +1355,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
                 )
 
                 align_fp8_moe_weights_for_flashinfer_trtllm(layer)
-            return
+
+        if hasattr(layer, "dispatcher"):
+            layer.dispatcher.set_quant_config({"weight_dtype": layer.w13_weight.dtype})
 
     def process_weights_hip_int4(self, layer: Module):
         # TODO: _use_aiter: add after triton kernel added

python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py

@@ -160,6 +160,10 @@ class QuarkW4A4MXFp4MoE(QuarkMoEScheme):
             layer.w13_weight.is_shuffled = True
             layer.w2_weight.is_shuffled = True
 
+        if hasattr(layer, "dispatcher"):
+            # Weights are stored as torch.uint8 but semantically MXFP4
+            layer.dispatcher.set_quant_config({"weight_dtype": torch.float4_e2m1fn_x2})
+
     def create_moe_runner(
         self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
     ):

test/registered/amd/test_moriep_small.py

@@ -69,7 +69,7 @@ class TestPureDP(CustomTestCase):
 
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
-        env["SGLANG_MORI_FP8_DISP"] = "False"
+        env["SGLANG_MORI_DISPATCH_DTYPE"] = "bf16"
         env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
 
@@ -113,7 +113,7 @@ class TestMTP(CustomTestCase):
 
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
-        env["SGLANG_MORI_FP8_DISP"] = "False"
+        env["SGLANG_MORI_DISPATCH_DTYPE"] = "bf16"
         env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
 
@@ -166,7 +166,7 @@ class TestNormal(CustomTestCase):
 
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
-        env["SGLANG_MORI_FP8_DISP"] = "False"
+        env["SGLANG_MORI_DISPATCH_DTYPE"] = "bf16"
         env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
 
@@ -213,7 +213,7 @@ class TestLowLatency(CustomTestCase):
 
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
-        env["SGLANG_MORI_FP8_DISP"] = "False"
+        env["SGLANG_MORI_DISPATCH_DTYPE"] = "bf16"
         env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
         # FIXME(billishyahao): enable p2p due to no rdma devices on CI machine
@@ -263,7 +263,7 @@ class TestTBOwithNormal(CustomTestCase):
 
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
-        env["SGLANG_MORI_FP8_DISP"] = "False"
+        env["SGLANG_MORI_DISPATCH_DTYPE"] = "bf16"
         env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
 
@@ -311,7 +311,7 @@ class TestTBOwithLowLatency(CustomTestCase):
 
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
-        env["SGLANG_MORI_FP8_DISP"] = "False"
+        env["SGLANG_MORI_DISPATCH_DTYPE"] = "bf16"
         env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
         # FIXME(billishyahao): enable p2p due to no rdma devices on CI machine
@@ -365,7 +365,7 @@ class TestMTPwithTBONormal(CustomTestCase):
 
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
-        env["SGLANG_MORI_FP8_DISP"] = "False"
+        env["SGLANG_MORI_DISPATCH_DTYPE"] = "bf16"
         env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
 
@@ -423,7 +423,7 @@ class TestMTPwithTBOLowLatency(CustomTestCase):
 
         env = dict(os.environ)
         env["SGLANG_USE_AITER"] = "1"
-        env["SGLANG_MORI_FP8_DISP"] = "False"
+        env["SGLANG_MORI_DISPATCH_DTYPE"] = "bf16"
         env["SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK"] = "4096"
         env["MORI_SHMEM_MODE"] = "ISOLATION"  # avoid out of symmetric heap memory
         # FIXME(billishyahao): enable p2p due to no rdma devices on CI machine