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

[Refactor] Refactor DeepEP dispatcher (#22822)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Cheng Wan <54331508+ch-wan@users.noreply.github.com>
This commit is contained in:
Артем Савкин
2026-05-18 04:36:42 +03:00
committed by GitHub
co-authored by gemini-code-assist[bot] Cheng Wan
parent 5147de26e4
commit a080358cac
30 changed files with 302 additions and 182 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+2 -2
View File
@@ -346,7 +346,7 @@ class Envs:
# Delay all-gather after qlora for better performance for Deepseek v3.2
SGLANG_USE_AG_AFTER_QLORA = EnvBool(False)
# Quantize x to int8 in the dispatch operator
DEEP_NORMAL_MODE_USE_INT8_QUANT = EnvBool(False)
DEEP_NORMAL_MODE_USE_INT8_QUANT = EnvBool(False) # This argument is deprecated
SGLANG_NPU_FUSED_MOE_MODE = EnvInt(1)
# MTHREADS & MUSA
@@ -411,7 +411,7 @@ class Envs:
SGLANG_MAX_KV_CHUNK_CAPACITY = EnvInt(128 * 1024)
# DeepEP
SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False)
SGLANG_DEEPEP_BF16_DISPATCH = EnvBool(False) # This argument is deprecated
SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(128)
SGLANG_DEEPEP_LL_COMBINE_SEND_NUM_SMS = EnvInt(32)
SGLANG_BLACKWELL_OVERLAP_SHARED_EXPERTS_OUTSIDE_SBO = EnvBool(False)
python/sglang/srt/hardware_backend/npu/quantization/awq_kernels.py
+18
View File
@@ -154,3 +154,21 @@ class AWQAscendMoEKernel:
dispatch_output: "StandardDispatchOutput",
) -> torch.Tensor:
return self.kernel.apply(layer, dispatch_output)
def apply_without_routing_weights(
self,
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
):
return self.kernel.apply_without_routing_weights(
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
)
python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py
+57
View File
@@ -426,6 +426,12 @@ class NPUW4A4Int4DynamicMoEMethod(_NPUFusedMoEMethodBase):
requires_grad=False,
)
# Quantizes in int4 separately from the dispatcher
# since deep_ep does not support quantization in int4
# dispatching works in bf16
if hasattr(layer, "dispatcher"):
layer.dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"})
def _pack_to_int32(self, weight: torch.Tensor):
# pack 8 int4 to int32, we use a int32 to represent a int4
assert (
@@ -460,6 +466,48 @@ class NPUW4A4Int4DynamicMoEMethod(_NPUFusedMoEMethodBase):
)
return StandardCombineInput(hidden_states=output)
def apply_without_routing_weights(
self,
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
):
hidden_states, hidden_states_scale = torch.ops.npu.npu_dynamic_quant(
hidden_states, dst_type=torch.quint4x2
)
# gmm1: up_gate_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w13_weight],
scale=[layer.w13_weight_scale],
per_token_scale=[hidden_states_scale],
split_item=2,
group_list_type=group_list_type,
group_type=0,
group_list=group_list,
output_dtype=output_dtype,
)[0]
# act_fn: swiglu
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(hidden_states)
# gmm2: down_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w2_weight],
scale=[layer.w2_weight_scale.to(output_dtype)],
per_token_scale=[pertoken_scale],
split_item=2,
group_list_type=group_list_type,
group_type=0,
group_list=group_list,
output_dtype=output_dtype,
)[0]
return hidden_states
class NPUW8A8Int8DynamicMoEMethod(_NPUFusedMoEMethodBase):
@@ -490,6 +538,9 @@ class NPUW8A8Int8DynamicMoEMethod(_NPUFusedMoEMethodBase):
requires_grad=False,
)
if hasattr(layer, "dispatcher"):
layer.dispatcher.set_quant_config({"dispatcher_output_dtype": "int8"})
def apply(
self,
layer,
@@ -656,6 +707,9 @@ class NPUW4A8Int8DynamicMoEMethod(_NPUFusedMoEMethodBase):
layer.w13_weight.data = self._pack_to_int32(layer.w13_weight.data)
layer.w2_weight.data = self._pack_to_int32(layer.w2_weight.data)
if hasattr(layer, "dispatcher"):
layer.dispatcher.set_quant_config({"dispatcher_output_dtype": "int8"})
def _process_weights_without_clip(
self, layer: torch.nn.Module, is_per_channel_weight
) -> None:
@@ -960,6 +1014,9 @@ class NPUW4A16Int4DynamicMoEMethod(_NPUFusedMoEMethodBase):
layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
if hasattr(layer, "dispatcher"):
layer.dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"})
def apply(
self,
layer,
python/sglang/srt/layers/moe/ep_moe/layer.py
+6 -21
View File
@@ -25,12 +25,6 @@ from sglang.srt.layers.moe.token_dispatcher.deepep import (
)
from sglang.srt.layers.moe.topk import TopKOutput, TopKOutputChecker
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
CompressedTensorsFusedMoEMethod,
)
from sglang.srt.layers.quantization.compressed_tensors.schemes import (
NPUCompressedTensorsW4A16Int4DynamicMoE,
)
from sglang.srt.layers.quantization.fp8 import Fp8Config
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config, W4AFp8MoEMethod
@@ -48,9 +42,6 @@ _is_npu = is_npu()
_is_fp8_fnuz = is_fp8_fnuz()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
if _is_npu:
import torch_npu
logger = logging.getLogger(__name__)
@@ -133,8 +124,13 @@ class DeepEPMoE(FusedMoE):
self.deepep_mode = get_deepep_mode()
# TODO: move this logic to process_weigths_after_loading, like:
# def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
# if hasattr(layer, "dispatcher"):
# layer.dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"})
if quant_config is None and hasattr(self.dispatcher, "set_quant_config"):
self.dispatcher.set_quant_config({"bf16_dispatch": True})
self.dispatcher.set_quant_config({"dispatcher_output_dtype": "bf16"})
if (
self.deepep_mode.enable_low_latency()
@@ -338,17 +334,6 @@ class DeepEPMoE(FusedMoE):
self, hidden_states, group_list_type, group_list, output_dtype
)
else:
input_quant = get_bool_env_var("DEEP_NORMAL_MODE_USE_INT8_QUANT")
if not input_quant and not isinstance(
self.quant_method,
(
NPUCompressedTensorsW4A16Int4DynamicMoE,
CompressedTensorsFusedMoEMethod,
),
):
hidden_states, hidden_states_scale = torch_npu.npu_dynamic_quant(
hidden_states
)
hidden_states = self.quant_method.apply_without_routing_weights(
self,
hidden_states,
python/sglang/srt/layers/moe/token_dispatcher/deepep.py
+77 -31
View File
@@ -1,6 +1,7 @@
from __future__ import annotations
import logging
import os
from contextlib import nullcontext
from dataclasses import dataclass
from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, Union
@@ -22,8 +23,9 @@ from sglang.srt.layers.moe.token_dispatcher.base import (
from sglang.srt.layers.moe.topk import TopKOutput
from sglang.srt.layers.moe.utils import (
DeepEPMode,
DeepEPOutputDtype,
get_deepep_config,
get_moe_runner_backend,
get_deepep_output_dtype,
is_tbo_enabled,
)
from sglang.srt.utils import (
@@ -344,6 +346,8 @@ class _DeepEPDispatcherImplBase:
self.overlap_args: Optional[CombineOverlapArgs] = None
self.meta_overlap_args: Optional[dict] = None
self.set_deepep_dispatcher_dtype()
def dispatch_a(
self,
hidden_states: torch.Tensor,
@@ -370,6 +374,74 @@ class _DeepEPDispatcherImplBase:
def set_quant_config(self, quant_config: dict) -> None:
self.quant_config = quant_config
self.set_deepep_dispatcher_dtype()
def set_deepep_dispatcher_dtype(self) -> None:
self.deepep_output_dtype = get_deepep_output_dtype(self)
# Configuration mapping for each dtype
config_map = {
DeepEPOutputDtype.BF16: {
"use_fp8": False,
"use_nvfp4": False,
},
DeepEPOutputDtype.FP8: {
"use_fp8": True,
"use_nvfp4": False,
},
# Needed for Ascend A2/A3 NPU case,
# despite the use_fp8 flag,
# quantization will be performed in int8
DeepEPOutputDtype.INT8: {
"use_fp8": True,
"use_nvfp4": False,
},
DeepEPOutputDtype.NVFP4: {
"use_fp8": False,
"use_nvfp4": True,
},
}
# Validate and apply hardware-specific adjustments
self._validate_and_adjust_dtype()
# Apply configuration
config = config_map[self.deepep_output_dtype]
self.use_fp8 = config["use_fp8"]
self.use_nvfp4 = config["use_nvfp4"]
# Handle environment variables
if _is_npu:
self._update_int8_quant_env()
def _validate_and_adjust_dtype(self) -> None:
"""Validate dtype against hardware and adjust if necessary."""
if _is_npu:
if self.deepep_output_dtype == DeepEPOutputDtype.FP8:
logger.warning_once(
"Ascend A2/A3 NPU does not support fp8 "
"deepep_dispatcher_output_dtype, switching to int8..."
)
self.deepep_output_dtype = DeepEPOutputDtype.INT8
elif self.deepep_output_dtype == DeepEPOutputDtype.NVFP4:
raise RuntimeError(
"Ascend A2/A3 NPU does not support nvfp4 deepep_dispatcher_output_dtype."
)
else:
if self.deepep_output_dtype == DeepEPOutputDtype.INT8:
logger.warning_once(
"GPU does not support int8 "
"deepep_dispatcher_output_dtype, switching to fp8..."
)
self.deepep_output_dtype = DeepEPOutputDtype.FP8
# NVFP4 is supported on GPU, no adjustment needed
def _update_int8_quant_env(self) -> None:
"""Update the DEEP_NORMAL_MODE_USE_INT8_QUANT environment variable."""
if self.use_fp8:
os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "1"
else:
os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "0"
def set_overlap_args(
self, combine_overlap_args: CombineOverlapArgs, meta_overlap_args: dict
@@ -397,14 +469,7 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
):
topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
topk_ids = topk_ids.to(torch.int64)
backend = get_moe_runner_backend()
# BF16 dispatch is needed when:
# - cutlass backend (uses different kernel)
# - deep_gemm backend with SGLANG_DEEPEP_BF16_DISPATCH enabled
need_bf16_dispatch = backend.is_cutlass() or (
backend.is_deep_gemm() and envs.SGLANG_DEEPEP_BF16_DISPATCH.get()
)
if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and not need_bf16_dispatch:
if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM and self.use_fp8:
# TODO hard code 128 block quant,use fp8 communication
hidden_states = sglang_per_token_group_quant_fp8(
hidden_states,
@@ -623,26 +688,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
hidden_states: torch.Tensor,
topk_ids: torch.Tensor,
):
use_nvfp4 = use_fp8 = False
input_global_scale = self.quant_config.get("input_global_scale", None)
bf16_dispatch = self.quant_config.get("bf16_dispatch", False)
if input_global_scale is not None:
use_nvfp4 = True
else:
backend = get_moe_runner_backend()
# BF16 dispatch is needed when:
# - quant_config requests BF16 dispatch explicitly
# - flashinfer_cutedsl: kernel quantizes to NVFP4 internally
# - NPU with SGLANG_DEEPEP_BF16_DISPATCH: INT8 input + BF16 weight GMM not supported
# - deep_gemm with SGLANG_DEEPEP_BF16_DISPATCH: user requests BF16 dispatch
need_bf16_dispatch = (
bf16_dispatch
or backend.is_flashinfer_cutedsl()
or (_is_npu and envs.SGLANG_DEEPEP_BF16_DISPATCH.get())
or (backend.is_deep_gemm() and envs.SGLANG_DEEPEP_BF16_DISPATCH.get())
)
if not need_bf16_dispatch:
use_fp8 = True
# round_scale / use_ue8m0 are FP8-DeepGEMM specific; they cause DeepEP
# to return int32-packed UE8M0 scales that don't feed the flashinfer
@@ -654,7 +700,7 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
use_ue8m0=deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
and deep_gemm_wrapper.DEEPGEMM_BLACKWELL,
)
if use_fp8
if self.use_fp8
else dict()
)
@@ -666,8 +712,8 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
topk_ids,
self.num_max_dispatch_tokens_per_rank,
self.num_experts,
use_fp8=use_fp8,
**(dict(use_nvfp4=True) if use_nvfp4 else dict()),
use_fp8=self.use_fp8,
**(dict(use_nvfp4=True) if self.use_nvfp4 else dict()),
**(
dict(x_global_scale=input_global_scale)
if input_global_scale is not None
python/sglang/srt/layers/moe/utils.py
+76
View File
@@ -9,14 +9,20 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.distributed.parallel_state import get_moe_expert_parallel_world_size
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
get_attention_dp_size,
is_dp_attention_enabled,
)
from sglang.srt.utils import is_npu
_is_npu = is_npu()
if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs
from sglang.srt.server_args import get_global_server_args
logger = logging.getLogger(__name__)
@@ -161,6 +167,76 @@ class DeepEPMode(Enum):
return self == DeepEPMode.AUTO
class DeepEPOutputDtype(Enum):
"""
Describes the dispatch output data type for DeepEP.
- BF16: dispatch hidden states in bf16
- FP8: dispatch hidden states in fp8
- INT8: dispatch hidden states in int8
- NVFP4: dispatch hidden states in nvfp4
"""
BF16 = "bf16"
FP8 = "fp8"
INT8 = "int8"
NVFP4 = "nvfp4"
def get_deepep_output_dtype(self) -> DeepEPOutputDtype:
"""
Automatically choose the dispatch output dtype for DeepEP.
The decision follows several checks in priority order:
0. Parse server argument.
1. Parse deprecated environment variables.
2. If quant_config contains input_global_scale → NVFP4 path.
3. Parse quant config
4. If flashinfer_cutedsl or is_cutlass backend is active → BF16 (it quantizes hidden_states internally).
5. Otherwise default for NPU → BF16 (the default for NPU).
6. Otherwise → FP8 (the default for most models like DeepSeek-V3).
"""
# 0. Parse server argument.
server_args = get_global_server_args()
if server_args and server_args.deepep_dispatcher_output_dtype != "auto":
return DeepEPOutputDtype(server_args.deepep_dispatcher_output_dtype)
# 1. Parse deprecated environment variables.
if envs.SGLANG_DEEPEP_BF16_DISPATCH.get():
logger.warning_once(
"Warning: The env variable SGLANG_DEEPEP_BF16_DISPATCH deprecated "
"and will be removed in future releases. Please use a new "
"`--deepep-dispatcher-output-dtype bf16` argument instead."
)
return DeepEPOutputDtype.BF16
# 2. NVFP4 is detected inside dispatch_a / _dispatch_core via quant_config; no need to infer here.
if self.quant_config is not None:
input_global_scale = self.quant_config.get("input_global_scale", None)
if input_global_scale is not None:
return DeepEPOutputDtype.NVFP4
# 3. Parse quant config to determine the output dtype of dispatcher
dispatcher_output_dtype = self.quant_config.get("dispatcher_output_dtype", None)
if dispatcher_output_dtype is not None:
return DeepEPOutputDtype(dispatcher_output_dtype)
# 4. flashinfer_cutedsl and is_cutlass expects BF16 dispatch
if (
get_moe_runner_backend().is_flashinfer_cutedsl()
or get_moe_runner_backend().is_cutlass()
):
return DeepEPOutputDtype.BF16
# 5. Default on NPU → BF16
if _is_npu:
return DeepEPOutputDtype.BF16
# 6. Default → FP8
return DeepEPOutputDtype.FP8
MOE_A2A_BACKEND: Optional[MoeA2ABackend] = None
MOE_RUNNER_BACKEND: Optional[MoeRunnerBackend] = None
SPECULATIVE_MOE_RUNNER_BACKEND: Optional[MoeRunnerBackend] = None
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py
+12
View File
@@ -606,6 +606,18 @@ class NPUCompressedTensorsW4A16Int4DynamicMoE(CompressedTensorsMoEScheme):
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
w13_weight_shape = torch.nn.Parameter(
torch.empty(num_experts, 2), requires_grad=False
)
layer.register_parameter("w13_weight_shape", w13_weight_shape)
set_weight_attrs(w13_weight_shape, extra_weight_attrs)
w2_weight_shape = torch.nn.Parameter(
torch.empty(num_experts, 2), requires_grad=False
)
layer.register_parameter("w2_weight_shape", w2_weight_shape)
set_weight_attrs(w2_weight_shape, extra_weight_attrs)
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
self.kernel.process_weights_after_loading(layer)
python/sglang/srt/layers/quantization/modelslim/schemes/modelslim_w4a4_int4_moe.py
+11 -3
View File
@@ -129,7 +129,15 @@ class ModelSlimW4A4Int4MoE(ModelSlimMoEScheme):
group_list,
output_dtype,
):
# FIXME W4A4 MoE does not work with DeepEP
raise NotImplementedError(
f"DeepEP currently does not support quantization in int4, please disable --moe-a2a-backend deepep"
logger.warning_once(
"Warning: Performance may be reduced, because DeepEP Dispatcher does not support 4-bit quantization, "
"switching to the bf16 dispatcher, quantization will be performed separately..."
)
return self.kernel.apply_without_routing_weights(
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
)
python/sglang/srt/models/deepseek_nextn.py
-6
View File
@@ -161,9 +161,6 @@ class DeepseekModelNextN(nn.Module):
forward_batch: ForwardBatch,
input_embeds: torch.Tensor = None,
) -> torch.Tensor:
if _is_npu and self.quant_config is None:
os.environ["SGLANG_DEEPEP_BF16_DISPATCH"] = "1"
os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "0"
zero_allocator = BumpAllocator(
buffer_size=2,
dtype=torch.float32,
@@ -224,9 +221,6 @@ class DeepseekModelNextN(nn.Module):
torch.cuda.current_stream(),
)
if _is_npu and self.quant_config is None:
os.environ["SGLANG_DEEPEP_BF16_DISPATCH"] = "0"
os.environ["DEEP_NORMAL_MODE_USE_INT8_QUANT"] = "1"
return hidden_states
python/sglang/srt/models/glm4_moe_nextn.py
+3 -1
View File
@@ -160,9 +160,11 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
if self.needs_quant_draft:
cxt = contextlib.nullcontext()
else:
# SGLANG_DEEPEP_BF16_DISPATCH is deprecated, will need
# to be removed in the future and moved to a new
# --deepep-dispatcher-output-dtype server argument.
unquant_patch = {
"SGLANG_DEEPEP_BF16_DISPATCH": "1",
"DEEP_NORMAL_MODE_USE_INT8_QUANT": "0",
}
cxt = temp_set_env(allow_sglang=True, **unquant_patch)
python/sglang/srt/models/qwen3_5_mtp.py
-15
View File
@@ -15,7 +15,6 @@
"""Inference-only Qwen3_5 MTP model."""
import logging
from contextlib import ExitStack
from typing import Iterable, Optional, Tuple
import torch
@@ -23,7 +22,6 @@ from torch import nn
from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
from sglang.srt.layers.layernorm import GemmaRMSNorm
@@ -138,17 +136,6 @@ class Qwen3_5ForCausalLMMTP(nn.Module):
input_embeds: Optional[torch.Tensor] = None,
**kwargs,
):
exit_stack = ExitStack()
if (
is_npu()
and self.quant_config is None
and get_global_server_args().quantization is not None
):
# ascend mtp unquant
exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True))
exit_stack.enter_context(
envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False)
)
assert input_embeds is None
input_embeds = forward_batch.mm_input_embeds
@@ -182,8 +169,6 @@ class Qwen3_5ForCausalLMMTP(nn.Module):
hidden_states,
)
exit_stack.close()
return self.logits_processor(
input_ids, hidden_states, self.lm_head, forward_batch
)
python/sglang/srt/models/qwen3_next_mtp.py
-15
View File
@@ -15,7 +15,6 @@
"""Inference-only Qwen3Next MTP Speculative Decoding."""
import logging
from contextlib import ExitStack
from typing import Iterable, Optional, Tuple
import torch
@@ -23,7 +22,6 @@ from torch import nn
from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size
from sglang.srt.environ import envs
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.layers.layernorm import GemmaRMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor
@@ -93,17 +91,6 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM):
input_embeds: Optional[torch.Tensor] = None,
**kwargs,
):
exit_stack = ExitStack()
if (
is_npu()
and self.quant_config is None
and get_global_server_args().quantization is not None
):
# ascend mtp unquant
exit_stack.enter_context(envs.SGLANG_DEEPEP_BF16_DISPATCH.override(True))
exit_stack.enter_context(
envs.DEEP_NORMAL_MODE_USE_INT8_QUANT.override(False)
)
if input_embeds is None:
input_embeds = self.model.embed_tokens(input_ids)
@@ -123,8 +110,6 @@ class Qwen3NextForCausalLMMTP(Qwen3NextForCausalLM):
hidden_states,
)
exit_stack.close()
return self.logits_processor(
input_ids, hidden_states, self.lm_head, forward_batch
)
python/sglang/srt/server_args.py
+10
View File
@@ -595,6 +595,9 @@ class ServerArgs:
enforce_disable_flashinfer_allreduce_fusion: bool = False
enable_aiter_allreduce_fusion: bool = False
deepep_mode: Literal["auto", "normal", "low_latency"] = "auto"
deepep_dispatcher_output_dtype: Literal["auto", "bf16", "fp8", "int8", "nvfp4"] = (
"auto"
)
ep_num_redundant_experts: int = 0
ep_dispatch_algorithm: Optional[Literal["static", "dynamic", "fake"]] = None
init_expert_location: str = "trivial"
@@ -5634,6 +5637,13 @@ class ServerArgs:
default="auto",
help="Select the mode when enable DeepEP or MoriEP MoE, could be `normal`, `low_latency` or `auto`. Default is `auto`, which means `low_latency` for decode batch and `normal` for prefill batch.",
)
parser.add_argument(
"--deepep-dispatcher-output-dtype",
type=str,
choices=["auto", "bf16", "fp8", "int8", "nvfp4"],
default="auto",
help="Select DeepEP dispatcher output dtype",
)
parser.add_argument(
"--ep-num-redundant-experts",
type=int,