chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[AMD][MORI] Fix MTP crash with FP4/FP8 dispatch and add NEXTN dispatch env vars. (#20453)

Browse Source
This commit is contained in:
Duyi-Wang
2026-03-14 05:03:17 +08:00
committed by GitHub
parent c37ef7f18b
commit 0eea80bc00
6 changed files with 46 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/references/environment_variables.md
View File

@@ -76,6 +76,8 @@ SGLang supports various environment variables that can be used to configure its
| --- | --- | --- |
| `SGLANG_MORI_FP8_DISP` | Use FP8 for dispatch | `"false"` |
| `SGLANG_MORI_FP4_DISP` | Use MXFP4 for dispatch | `"false"` |
| `SGLANG_MORI_NEXTN_FP8_DISP` | Use FP8 for MTP (NextN) layer dispatch. When neither `SGLANG_MORI_NEXTN_FP8_DISP` nor `SGLANG_MORI_NEXTN_FP4_DISP` is set, falls back to `SGLANG_MORI_FP8_DISP`/`SGLANG_MORI_FP4_DISP`. | `"false"` |
| `SGLANG_MORI_NEXTN_FP4_DISP` | Use MXFP4 for MTP (NextN) layer dispatch. When neither `SGLANG_MORI_NEXTN_FP8_DISP` nor `SGLANG_MORI_NEXTN_FP4_DISP` is set, falls back to `SGLANG_MORI_FP8_DISP`/`SGLANG_MORI_FP4_DISP`. | `"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` |

16
python/sglang/srt/layers/moe/ep_moe/layer.py
View File

@@ -653,14 +653,17 @@ class MoriEPMoE(DeepEPMoE):
quant_type = QuantType.No
is_bf16_weights = self.w13_weight.dtype == torch.bfloat16
if (
not is_fp8_quant
and dispatch_scale is not None
and dispatch_a1.dtype != torch.float4_e2m1fn_x2
):
if is_quark_w4a4:
# W4A4 model with FP8 dispatch: must dequant FP8->BF16 first,
# because the FP4 per_1x32 quantization path needs BF16 input
if is_quark_w4a4 or is_bf16_weights:
# W4A4 or BF16 weights with FP8 dispatch: must dequant FP8->BF16 first.
# W4A4 needs BF16 input for the FP4 per_1x32 path;
# BF16 weights don't support quantized activations.
dispatch_a1 = upscale(
dispatch_a1, dispatch_scale, dispatch_recv_token_num, output_dtype
)
@@ -671,10 +674,9 @@ class MoriEPMoE(DeepEPMoE):
quant_type = QuantType.per_128x128
if dispatch_a1.dtype == torch.float4_e2m1fn_x2 and dispatch_scale is not None:
if is_fp8_quant:
# FP8 weights + FP4 dispatch is not supported by fused_moe kernels
# (no kernel for q_dtype_a=fp4x2, q_dtype_w=fp8).
# Must dequant FP4->BF16 first; fused_moe will re-quant to FP8 internally.
if is_fp8_quant or is_bf16_weights:
# FP8/BF16 weights + FP4 dispatch: no kernel supports fp4 act + fp8/bf16 weight.
# Must dequant FP4->BF16 first.
dispatch_a1 = upscale_mxfp4(
dispatch_a1, dispatch_scale, dispatch_recv_token_num, output_dtype
)

3
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
View File

@@ -116,6 +116,7 @@ def create_moe_dispatcher(moe_runner_config: MoeRunnerConfig) -> BaseDispatcher:
deepep_mode=get_deepep_mode(),
async_finish=True,
return_recv_hook=True,
is_nextn=moe_runner_config.is_nextn,
)
elif a2a_backend.is_ascend_fuseep():
from sglang.srt.layers.moe.token_dispatcher import NpuFuseEPDispatcher
@@ -194,6 +195,7 @@ class FusedMoE(torch.nn.Module):
with_bias=False,
routing_method_type: Optional[RoutingMethodType] = None,
is_gated: bool = True,
is_nextn: bool = False,
):
super().__init__()
if params_dtype is None:
@@ -269,6 +271,7 @@ class FusedMoE(torch.nn.Module):
gemm1_clamp_limit=gemm1_clamp_limit,
is_gated=is_gated,
routing_method_type=routing_method_type,
is_nextn=is_nextn,
)
self.quant_method: Optional[FusedMoEMethodBase] = None

1
python/sglang/srt/layers/moe/moe_runner/base.py
View File

@@ -48,6 +48,7 @@ class MoeRunnerConfig:
routed_scaling_factor: Optional[float] = None
gemm1_alpha: Optional[float] = None
gemm1_clamp_limit: Optional[float] = None
is_nextn: bool = False
@dataclass

47
python/sglang/srt/layers/moe/token_dispatcher/moriep.py
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,31 @@ 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")
@lru_cache(maxsize=4)
def _get_mori_dispatch_quant_flags(is_nextn=False):
fp8_var, fp4_var = "SGLANG_MORI_FP8_DISP", "SGLANG_MORI_FP4_DISP"
if is_nextn and (
"SGLANG_MORI_NEXTN_FP8_DISP" in os.environ
or "SGLANG_MORI_NEXTN_FP4_DISP" in os.environ
):
fp8_var, fp4_var = "SGLANG_MORI_NEXTN_FP8_DISP", "SGLANG_MORI_NEXTN_FP4_DISP"
fp8_dispatch = get_bool_env_var(fp8_var, "False")
fp4_dispatch = get_bool_env_var(fp4_var, "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."
f"Both {fp8_var} and {fp4_var} are set to True. "
f"Using {fp4_var} and ignoring {fp8_var}."
)
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,
@@ -195,6 +205,8 @@ def init_mori_op(
params_dtype,
num_max_dispatch_tokens_per_rank,
deepep_mode,
fp8_dispatch,
fp4_dispatch,
):
import mori
@@ -229,8 +241,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:
@@ -308,6 +318,7 @@ class _MoriEPDispatcherImplBase:
hidden_size: int,
params_dtype: torch.dtype,
deepep_mode: DeepEPMode,
is_nextn: bool = False,
):
try:
import mori # noqa: F401
@@ -321,6 +332,8 @@ class _MoriEPDispatcherImplBase:
self.hidden_size = hidden_size
self.params_dtype = params_dtype
self.deepep_mode = deepep_mode
self.is_nextn = is_nextn
self.fp8_dispatch, self.fp4_dispatch = _get_mori_dispatch_quant_flags(is_nextn)
self.num_max_dispatch_tokens_per_rank = get_int_env_var(
"SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK", 4096
@@ -335,6 +348,8 @@ class _MoriEPDispatcherImplBase:
self.params_dtype,
self.num_max_dispatch_tokens_per_rank,
self.deepep_mode,
self.fp8_dispatch,
self.fp4_dispatch,
)
self.quant_config: Optional[dict] = None
@@ -421,9 +436,7 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
output_dtype = hidden_states.dtype
scale = None
fp8_dispatch, fp4_dispatch = _get_mori_dispatch_quant_flags()
if fp8_dispatch:
if self.fp8_dispatch:
# FP8 quant
if num_tokens > 0:
# NOTE: aiter is able to handle token=0 case in UT. But for some reason it failed at e2e case. Root cause TBD.
@@ -440,7 +453,7 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase):
device=hidden_states.device,
)
elif fp4_dispatch:
elif self.fp4_dispatch:
# FP4 quant
if num_tokens > 0:
hidden_states, scale = self.fp4_quant_func(hidden_states, shuffle=False)
@@ -645,9 +658,7 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase):
output_dtype = hidden_states.dtype
scale = None
fp8_dispatch, fp4_dispatch = _get_mori_dispatch_quant_flags()
if fp8_dispatch:
if self.fp8_dispatch:
# FP8 quant
if num_tokens > 0:
# NOTE: aiter is able to handle token=0 case in UT. But for some reason it failed at e2e case. Root cause TBD.
@@ -664,7 +675,7 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase):
device=hidden_states.device,
)
elif fp4_dispatch:
elif self.fp4_dispatch:
# FP4 quant
if num_tokens > 0:
hidden_states, scale = self.fp4_quant_func(hidden_states, shuffle=False)
@@ -818,6 +829,7 @@ class MoriEPDispatcher(BaseDispatcher):
deepep_mode: DeepEPMode = DeepEPMode.AUTO,
async_finish: bool = False,
return_recv_hook: bool = False,
is_nextn: bool = False,
):
super().__init__()
@@ -832,6 +844,7 @@ class MoriEPDispatcher(BaseDispatcher):
hidden_size=hidden_size,
params_dtype=params_dtype,
deepep_mode=deepep_mode,
is_nextn=is_nextn,
)
if self.deepep_mode.enable_low_latency():

1
python/sglang/srt/models/deepseek_v2.py
View File

@@ -403,6 +403,7 @@ class DeepseekV2MoE(nn.Module):
config, "routing_method_type", RoutingMethodType.DeepSeekV3
),
prefix=add_prefix("experts", prefix),
is_nextn=is_nextn,
)
self.topk = TopK(