diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index fb4d17797..9d95c4f23 100644 --- a/docs/references/environment_variables.md +++ b/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_FP4_DISP` | Use MXFP4 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_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` | diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index a2eb8f07b..3113c3cbe 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/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,23 +337,55 @@ class _MoriEPDispatcherImplBase: "SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK", 4096 ) - self.mori_op = init_mori_op( - self.group, - self.router_topk, - self.num_experts, - self.num_local_experts, - self.hidden_size, - self.params_dtype, - self.num_max_dispatch_tokens_per_rank, - self.deepep_mode, - self.instance_id, - ) + 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, + self.num_local_experts, + self.hidden_size, + self.params_dtype, + self.num_max_dispatch_tokens_per_rank, + self.deepep_mode, + self.instance_id, + self.fp8_dispatch, + self.fp4_dispatch, + ) + return self._mori_op + + def _apply_dispatch_dtype_override(self): + """Apply env var override to fp8_dispatch/fp4_dispatch flags.""" + if "SGLANG_MORI_DISPATCH_DTYPE" in os.environ: + 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, hidden_states: torch.Tensor, @@ -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 diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 234af35f5..7182e3d57 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/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 - # Block quant doesn't need to process weights after loading - if self.block_quant: + elif self.block_quant: + # Block quant doesn't need to process weights after loading 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 diff --git a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py b/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py index df6650be8..ae9743577 100644 --- a/python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py +++ b/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 ): diff --git a/test/registered/amd/test_moriep_small.py b/test/registered/amd/test_moriep_small.py index ffe56e926..3eca8ce27 100644 --- a/test/registered/amd/test_moriep_small.py +++ b/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