minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Duyi-Wang
2026-03-24 22:53:57 -07:00
committed by GitHub
parent 86e2622097
commit 61a902ce88
5 changed files with 90 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/moe/token_dispatcher/moriep.py
+71 -37
View File
@@ -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
python/sglang/srt/layers/quantization/fp8.py
+6 -7
View File
@@ -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
python/sglang/srt/layers/quantization/quark/schemes/quark_w4a4_mxfp4_moe.py
+4
View File
@@ -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
):