From 0eea80bc001c991c9323ba7063f480fd4ae94c21 Mon Sep 17 00:00:00 2001 From: Duyi-Wang Date: Sat, 14 Mar 2026 05:03:17 +0800 Subject: [PATCH] [AMD][MORI] Fix MTP crash with FP4/FP8 dispatch and add NEXTN dispatch env vars. (#20453) --- docs/references/environment_variables.md | 2 + python/sglang/srt/layers/moe/ep_moe/layer.py | 16 ++++--- .../srt/layers/moe/fused_moe_triton/layer.py | 3 ++ .../sglang/srt/layers/moe/moe_runner/base.py | 1 + .../srt/layers/moe/token_dispatcher/moriep.py | 47 ++++++++++++------- python/sglang/srt/models/deepseek_v2.py | 1 + 6 files changed, 46 insertions(+), 24 deletions(-) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 98c50c03d..7a421953d 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -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` | diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py index d507dad12..c33f87e89 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -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 ) diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index a54e9ea7d..409b723db 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -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 diff --git a/python/sglang/srt/layers/moe/moe_runner/base.py b/python/sglang/srt/layers/moe/moe_runner/base.py index 12dd2ba6a..f53021cd6 100644 --- a/python/sglang/srt/layers/moe/moe_runner/base.py +++ b/python/sglang/srt/layers/moe/moe_runner/base.py @@ -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 diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index 85fe8653e..b3efa6fa3 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,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(): diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 8f0617142..d10ff54c0 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -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(