From 39008955ffc52dc5da783640ca7ef5c148daa430 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Sat, 14 Mar 2026 12:12:42 -0700 Subject: [PATCH] Revert "[AMD][MORI] Fix MTP crash with FP4/FP8 dispatch and add NEXTN dispatch env vars." (#20602) --- 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, 24 insertions(+), 46 deletions(-) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 7a421953d..98c50c03d 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -76,8 +76,6 @@ 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 c33f87e89..d507dad12 100644 --- a/python/sglang/srt/layers/moe/ep_moe/layer.py +++ b/python/sglang/srt/layers/moe/ep_moe/layer.py @@ -653,17 +653,14 @@ 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 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. + if is_quark_w4a4: + # W4A4 model with FP8 dispatch: must dequant FP8->BF16 first, + # because the FP4 per_1x32 quantization path needs BF16 input dispatch_a1 = upscale( dispatch_a1, dispatch_scale, dispatch_recv_token_num, output_dtype ) @@ -674,9 +671,10 @@ 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 or is_bf16_weights: - # FP8/BF16 weights + FP4 dispatch: no kernel supports fp4 act + fp8/bf16 weight. - # Must dequant FP4->BF16 first. + 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. 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 409b723db..a54e9ea7d 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -116,7 +116,6 @@ 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 @@ -195,7 +194,6 @@ 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: @@ -271,7 +269,6 @@ 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 f53021cd6..12dd2ba6a 100644 --- a/python/sglang/srt/layers/moe/moe_runner/base.py +++ b/python/sglang/srt/layers/moe/moe_runner/base.py @@ -48,7 +48,6 @@ 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 b3efa6fa3..85fe8653e 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py @@ -1,7 +1,6 @@ from __future__ import annotations import logging -import os from dataclasses import dataclass from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple @@ -171,31 +170,22 @@ def get_ep_dispatch_configs(num_max_dispatch_tokens_per_rank: int = 4096): } -@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") - +@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( - f"Both {fp8_var} and {fp4_var} are set to True. " - f"Using {fp4_var} and ignoring {fp8_var}." + "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=4) +@lru_cache(maxsize=2) def init_mori_op( group, router_topk, @@ -205,8 +195,6 @@ def init_mori_op( params_dtype, num_max_dispatch_tokens_per_rank, deepep_mode, - fp8_dispatch, - fp4_dispatch, ): import mori @@ -241,6 +229,8 @@ 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: @@ -318,7 +308,6 @@ class _MoriEPDispatcherImplBase: hidden_size: int, params_dtype: torch.dtype, deepep_mode: DeepEPMode, - is_nextn: bool = False, ): try: import mori # noqa: F401 @@ -332,8 +321,6 @@ 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 @@ -348,8 +335,6 @@ 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 @@ -436,7 +421,9 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase): output_dtype = hidden_states.dtype scale = None - if self.fp8_dispatch: + fp8_dispatch, fp4_dispatch = _get_mori_dispatch_quant_flags() + + if 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. @@ -453,7 +440,7 @@ class _MoriEPDispatcherImplNormal(_MoriEPDispatcherImplBase): device=hidden_states.device, ) - elif self.fp4_dispatch: + elif fp4_dispatch: # FP4 quant if num_tokens > 0: hidden_states, scale = self.fp4_quant_func(hidden_states, shuffle=False) @@ -658,7 +645,9 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase): output_dtype = hidden_states.dtype scale = None - if self.fp8_dispatch: + fp8_dispatch, fp4_dispatch = _get_mori_dispatch_quant_flags() + + if 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. @@ -675,7 +664,7 @@ class _MoriEPDispatcherImplLowLatency(_MoriEPDispatcherImplBase): device=hidden_states.device, ) - elif self.fp4_dispatch: + elif fp4_dispatch: # FP4 quant if num_tokens > 0: hidden_states, scale = self.fp4_quant_func(hidden_states, shuffle=False) @@ -829,7 +818,6 @@ class MoriEPDispatcher(BaseDispatcher): deepep_mode: DeepEPMode = DeepEPMode.AUTO, async_finish: bool = False, return_recv_hook: bool = False, - is_nextn: bool = False, ): super().__init__() @@ -844,7 +832,6 @@ 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 d10ff54c0..8f0617142 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -403,7 +403,6 @@ class DeepseekV2MoE(nn.Module): config, "routing_method_type", RoutingMethodType.DeepSeekV3 ), prefix=add_prefix("experts", prefix), - is_nextn=is_nextn, ) self.topk = TopK(