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

moe: the shared-experts-fusion decision is a per-runner value the loader installs (#33889)

Browse Source
This commit is contained in:
Cheng Wan
2026-08-07 22:42:58 -07:00
committed by GitHub
parent eda0ddc260
commit b61a06921e
40 changed files with 1492 additions and 416 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/overrides.py
-12
View File
@@ -297,18 +297,6 @@ def mamba_extra_buffer_of(cfg: Any) -> bool:
)
def declare_load_time_override(source: str, declared: Dict[str, Any]) -> None:
"""Declare a load-time resolved field (model-file config overrides,
weight-resolved dtypes): validated against the resolvable whitelist, then
written to the config bags via ``get_context().override``; ``server_args``
stays the pristine startup record."""
from sglang.srt.runtime_context import get_context
context = get_context()
validate_declarations(context.server_args, [(source, dict(declared))])
context.override(source, **declared)
def collect_model_override_declarations(
architecture: str, server_args: Any, hf_config: Any
) -> List[Tuple[str, Dict[str, Any]]]:
python/sglang/srt/layers/moe/utils.py
+90
View File
@@ -21,6 +21,7 @@ if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs
from sglang.srt.runtime_context import get_server_args
from sglang.srt.utils.common import log_info_on_rank0
logger = logging.getLogger(__name__)
@@ -321,6 +322,12 @@ def initialize_moe_config(server_args: ServerArgs):
moe.tbo_token_distribution_threshold = server_args.tbo_token_distribution_threshold
moe.disable_fp4_allgather = server_args.disable_flashinfer_cutlass_moe_fp4_allgather
moe.quantization = server_args.quantization
# Seeded with the user's intent; each model's gate refines the ACTIVE
# value for its own build (install_shared_experts_fusion_decision).
moe.disable_shared_experts_fusion = server_args.disable_shared_experts_fusion
moe.speculative_disable_shared_experts_fusion = (
server_args.disable_shared_experts_fusion
)
def get_moe_a2a_backend() -> MoeA2ABackend:
@@ -357,6 +364,88 @@ def get_speculative_moe_a2a_backend() -> MoeA2ABackend:
return moe.speculative_a2a_backend
def is_shared_experts_fusion_disabled() -> bool:
"""The ACTIVE shared-experts-fusion decision for the model being built.
Written (both ways) by each MoE model's gate before its layers construct;
falls back to the config intent when no gate has run (models without an
auto-disable gate read the intent directly off the bag instead).
Construction-time only: a forward reads what its build baked in
(``num_fused_shared_experts`` on the layer). During a draft's build this
flag holds the DRAFT's decision, so a forward-time read would race the
build window — refuse it loudly."""
from sglang.srt.model_executor.forward_context import has_forward_context
if has_forward_context():
raise AssertionError(
"is_shared_experts_fusion_disabled() called inside a forward: the "
"fusion decision is construction-time state (it can hold the draft's "
"value while a draft builds). Read the value your build baked in, "
"e.g. the layer's num_fused_shared_experts."
)
moe = get_flags().moe
if moe.disable_shared_experts_fusion is None:
from sglang.srt.runtime_context import get_exec
return get_exec().moe.disable_shared_experts_fusion
return moe.disable_shared_experts_fusion
@contextmanager
def draft_model_build_scope():
"""Brackets a draft model's CONSTRUCTION: the gates it runs record their
fusion decision on the speculative leaf as well, and the target's ACTIVE
value returns on exit.
Deliberately does not touch ``runner_backend`` — swapping that is
``speculative_moe_backend_context``'s job and has to bracket the draft's
whole lifecycle (build + capture + forward), which not every worker does.
"""
moe = get_flags().moe
original_fusion = moe.disable_shared_experts_fusion
original_scope = moe.in_speculative_scope
try:
moe.in_speculative_scope = True
yield
finally:
moe.in_speculative_scope = original_scope
moe.disable_shared_experts_fusion = original_fusion
def install_shared_experts_fusion_decision(
model_class, hf_config, quant_config
) -> None:
"""Decide whether this runner's model fuses its shared experts, and install
the answer for the model it is about to build.
Called from the loader's single model-instantiation point, so the decision
is made once per runner — before any layer exists — and the model classes
are pure readers (``is_shared_experts_fusion_disabled``). A model family
that can auto-disable exposes the conditions as
``shared_experts_fusion_disable_reason(hf_config, quant_config)``; families
without one follow the user's intent.
Inside ``draft_model_build_scope`` the answer also lands on the speculative
leaf, so a flags dump afterwards shows both runners' decisions.
"""
from sglang.srt.runtime_context import get_exec
disabled = get_exec().moe.disable_shared_experts_fusion
if not disabled:
gate = getattr(model_class, "shared_experts_fusion_disable_reason", None)
reason = gate(hf_config, quant_config) if gate is not None else None
if reason:
log_info_on_rank0(
logger, f"{reason} Shared experts fusion optimization is disabled."
)
disabled = True
moe = get_flags().moe
moe.disable_shared_experts_fusion = disabled
if moe.in_speculative_scope:
moe.speculative_disable_shared_experts_fusion = disabled
def get_deepep_mode() -> DeepEPMode:
moe = get_flags().moe
if moe.deepep_mode is None:
@@ -526,6 +615,7 @@ def speculative_moe_backend_context():
"""
Context manager to temporarily use the speculative MoE backend for draft model operations.
This ensures that draft models in speculative decoding use the configured speculative backend.
"""
moe = get_flags().moe
original_backend = moe.runner_backend
python/sglang/srt/model_executor/model_runner.py
+8 -5
View File
@@ -2043,9 +2043,12 @@ class ModelRunner:
load_config: LoadConfig,
) -> None:
self.model = new_model
get_context().override(
"model_runner.update_model_fields",
model_path=model_path,
load_format=load_format,
)
# The record says what model this PROCESS serves; a draft's weight
# update is not that (its own state is on the runner).
if not self.is_draft_worker:
get_context().override(
"model_runner.update_model_fields",
model_path=model_path,
load_format=load_format,
)
self.load_config = load_config
python/sglang/srt/model_executor/model_runner_components/load_model_utils.py
+4 -4
View File
@@ -68,11 +68,11 @@ def maybe_downgrade_dtype_for_legacy_gpu(
logger.info(
"Compute capability below sm80. Use float16 due to lack of bfloat16 support."
)
from sglang.srt.arg_groups.overrides import declare_load_time_override
from sglang.srt.runtime_context import get_context
declare_load_time_override(
"ModelRunner._sm80_dtype_fallback", {"dtype": "float16"}
)
# Device-driven, so every runner in the process resolves the same way;
# the per-runner truth is model_config.dtype, this is the record.
get_context().override("ModelRunner._sm80_dtype_fallback", dtype="float16")
model_config.dtype = torch.float16
if torch.cuda.get_device_capability()[1] < 5:
raise RuntimeError("SGLang only supports sm75 and above.")
python/sglang/srt/model_loader/loader.py
+10
View File
@@ -83,6 +83,9 @@ from sglang.srt.distributed import (
model_parallel_is_initialized,
)
from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES
from sglang.srt.layers.moe.utils import (
install_shared_experts_fusion_decision,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
trigger_transferring_weights_request,
@@ -314,6 +317,13 @@ def _initialize_model(
) -> nn.Module:
"""Initialize a model with the given configurations."""
model_class, _ = get_model_architecture(model_config)
# Decide the shared-experts-fusion question here, once per runner, before any
# layer exists: this is the only place a model class is instantiated, and it
# is the last point that still knows both the checkpoint's quantization and
# (through the build scope) whether this runner is a draft.
install_shared_experts_fusion_decision(
model_class, model_config.hf_config, quant_config
)
kwargs = {
"config": model_config.hf_config,
"quant_config": quant_config,
python/sglang/srt/models/deepseek_nextn.py
+3 -1
View File
@@ -309,6 +309,8 @@ class DeepseekModelNextN(nn.Module):
class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
# The draft checkpoint reports the NextN architecture name.
fused_shared_experts_architecture = "DeepseekV3ForCausalLMNextN"
# Support amd/DeepSeek-R1-0528-MXFP4 renaming: model.layers.61*.
# Ref: HF config.json for amd/DeepSeek-R1-0528-MXFP4
@@ -343,7 +345,7 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
self.quant_config = quant_config
# if not set, model load will be broken in DeepseekV3ForCausalLM load_weights()
self.pp_group = get_pp_group()
self.determine_num_fused_shared_experts("DeepseekV3ForCausalLMNextN")
self.determine_num_fused_shared_experts()
self.use_dsa = is_deepseek_dsa(config)
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
self.mla_enable_prefill_cp = is_mla_prefill_cp_enabled() and not self.use_dsa
python/sglang/srt/models/deepseek_ocr.py
+21 -5
View File
@@ -1422,7 +1422,27 @@ def build_qwen2_decoder_as_encoder(
return decoder_as_encoder
def _is_ocr2(config: DeepseekVLV2Config) -> bool:
return (
str(getattr(config.vision_config, "model_name", "")).lower() == "deepencoderv2"
or getattr(config.projector_config, "input_dim", None) == 896
)
class DeepseekOCRForCausalLM(nn.Module):
@staticmethod
def shared_experts_fusion_disable_reason(hf_config, quant_config):
text_config = hf_config.text_config
if _is_ocr2(hf_config) or not (
text_config.topk_method == "noaux_tc" or text_config.use_mla
):
# Those branches build the dense DeepseekForCausalLM, which has no
# shared experts to fuse.
return None
return DeepseekV2ForCausalLM.shared_experts_fusion_disable_reason(
text_config, quant_config
)
def __init__(
self,
*,
@@ -1437,11 +1457,7 @@ class DeepseekOCRForCausalLM(nn.Module):
self.vision_config = config.vision_config
self.projector_config = config.projector_config
self.text_config = config.text_config
self.is_ocr2 = (
str(getattr(self.vision_config, "model_name", "")).lower()
== "deepencoderv2"
or getattr(self.projector_config, "input_dim", None) == 896
)
self.is_ocr2 = _is_ocr2(config)
n_embed = getattr(self.projector_config, "n_embed", 1280)
self.tile_tag = config.tile_tag
python/sglang/srt/models/deepseek_v2.py
+37 -44
View File
@@ -117,6 +117,7 @@ from sglang.srt.layers.moe.utils import (
has_per_rank_fused_shared_slots,
is_deepep_class_backend,
is_sbo_enabled,
is_shared_experts_fusion_disabled,
is_tbo_enabled,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -203,7 +204,6 @@ from sglang.srt.utils import (
LazyValue,
add_prefix,
is_non_idle_and_non_empty,
log_info_on_rank0,
make_layers,
use_intel_amx_backend,
)
@@ -571,7 +571,7 @@ class DeepseekV2MoE(nn.Module):
n_shared_experts = (
0 if config.n_shared_experts is None else int(config.n_shared_experts)
)
_fusion_disabled = get_exec().moe.disable_shared_experts_fusion
_fusion_disabled = is_shared_experts_fusion_disabled()
# num_fused_shared_experts drives weight remapping in deepseek_weight_loader:
# mlp.shared_experts → mlp.experts.256 when > 0.
@@ -2966,23 +2966,27 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
def routed_experts_weights_of_layer(self):
return self._routed_experts_weights_of_layer.value
def determine_num_fused_shared_experts(
self, architecture: str = "DeepseekV3ForCausalLM"
):
self.num_fused_shared_experts = 0
# The architecture this class fuses shared experts for; a subclass whose
# checkpoint reports a different name (the NextN drafts, GLM's DSA variant)
# overrides it.
fused_shared_experts_architecture = "DeepseekV3ForCausalLM"
if get_exec().moe.disable_shared_experts_fusion:
return
@classmethod
def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
"""Why this checkpoint cannot fuse its shared expert, or None.
disable_reason = None
Evaluated by the loader once per runner, before any layer is built (see
``install_shared_experts_fusion_decision``), so it takes the config and
quantization it is asked about rather than reading an instance.
"""
if get_exec().moe.enforce_shared_experts_fusion:
pass
elif is_sbo_enabled() or is_tbo_enabled():
disable_reason = "SBO/TBO enabled: incompatible with fusing shared expert into MoE kernel."
elif is_deepep_class_backend():
disable_reason = "DeepEP: fusion off by default (use --enforce-shared-experts-fusion to enable)."
elif (
self.config.architectures[0] != architecture
return None
if is_sbo_enabled() or is_tbo_enabled():
return "SBO/TBO enabled: incompatible with fusing shared expert into MoE kernel."
if is_deepep_class_backend():
return "DeepEP: fusion off by default (use --enforce-shared-experts-fusion to enable)."
if (
hf_config.architectures[0] != cls.fused_shared_experts_architecture
# Allow-list of n_routed_experts values that have been validated
# for shared-experts fusion under this code path. Currently:
# 256 -> DeepSeek-V3 / R1
@@ -2991,51 +2995,40 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
# moonshotai/Kimi-K2.5 (compressed-tensors) checkpoint
# stores the shared expert loose and is NOT pre-fused,
# so the fused path silently mis-loads it.
or self.config.n_routed_experts not in (256, 384)
or self.config.n_shared_experts != 1
or hf_config.n_routed_experts not in (256, 384)
or hf_config.n_shared_experts != 1
or (
self.config.n_routed_experts == 384
and (
self.quant_config is None or self.quant_config.get_name() != "quark"
)
hf_config.n_routed_experts == 384
and (quant_config is None or quant_config.get_name() != "quark")
)
):
disable_reason = "Config does not support fused shared expert(s)."
elif (
return "Config does not support fused shared expert(s)."
if (
(not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0))
and (not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4))
and (not _is_musa or torch.musa.get_device_capability("musa") < (3, 1))
):
disable_reason = (
return (
"Only Deepseek V3/R1 on NV-platform with capability >= 80 "
"or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization."
"or MT-platform with capability >= 31 can use shared experts fusion optimization."
)
elif get_parallel().moe_ep_size > 1 and (
if get_parallel().moe_ep_size > 1 and (
not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
):
disable_reason = (
return (
"Only Deepseek V3/R1 on AMD-platform with capability >= gfx942(MI30x) "
"can use shared experts fusion optimization under expert parallelism."
)
elif is_wint4afp8_or_wint4a16_config(self.quant_config):
disable_reason = "Deepseek V3/R1 W4AFP8/W4A16 model uses different quant method for routed experts and shared experts."
if is_wint4afp8_or_wint4a16_config(quant_config):
return "Deepseek V3/R1 W4AFP8/W4A16 model uses different quant method for routed experts and shared experts."
return None
if disable_reason is not None:
from sglang.srt.arg_groups.overrides import declare_load_time_override
declare_load_time_override(
"DeepseekV2ForCausalLM.determine_num_fused_shared_experts",
{"disable_shared_experts_fusion": True},
)
self.num_fused_shared_experts = 0
log_info_on_rank0(
logger,
f"{disable_reason} Shared experts fusion optimization is disabled.",
)
return
self.num_fused_shared_experts = self.config.n_shared_experts
def determine_num_fused_shared_experts(self):
# The decision was installed by the loader; this only reads it.
self.num_fused_shared_experts = (
0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
)
def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens
python/sglang/srt/models/deepseek_v4.py
+19 -28
View File
@@ -89,6 +89,7 @@ from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
from sglang.srt.layers.moe.utils import is_shared_experts_fusion_disabled
from sglang.srt.layers.quantization.fp8_utils import (
view_aiter_fused_rms_transposed_fp8_scale,
)
@@ -2649,35 +2650,25 @@ class DeepseekV4ForCausalLM(nn.Module):
self.capture_aux_hidden_states = True
self.model.dspark_layers_to_capture = list(layer_ids)
@classmethod
def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
"""V4 only fuses when explicitly asked to, and then the checkpoint must
carry exactly one shared expert. Asked by the loader before any layer is
built."""
if not get_exec().moe.enforce_shared_experts_fusion:
return "Config does not support fused shared expert(s)."
if hf_config.n_shared_experts != 1:
raise ValueError(
"DeepSeek V4 shared-experts fusion expects exactly one shared "
f"expert, but got n_shared_experts={hf_config.n_shared_experts}."
)
return None
def determine_num_fused_shared_experts(self):
self.num_fused_shared_experts = 0
if get_exec().moe.disable_shared_experts_fusion:
return
disable_reason = None
if get_exec().moe.enforce_shared_experts_fusion:
if self.config.n_shared_experts != 1:
raise ValueError(
"DeepSeek V4 shared-experts fusion expects exactly one shared "
f"expert, but got n_shared_experts={self.config.n_shared_experts}."
)
else:
disable_reason = "Config does not support fused shared expert(s)."
if disable_reason is not None:
from sglang.srt.arg_groups.overrides import declare_load_time_override
declare_load_time_override(
"DeepseekV4ForCausalLM.determine_num_fused_shared_experts",
{"disable_shared_experts_fusion": True},
)
log_info_on_rank0(
logger,
f"{disable_reason} Shared experts fusion optimization is disabled.",
)
return
self.num_fused_shared_experts = self.config.n_shared_experts
# The decision was installed by the loader; this only reads it.
self.num_fused_shared_experts = (
0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
)
@torch.no_grad()
def forward(
python/sglang/srt/models/deepseek_vl2.py
+10
View File
@@ -157,6 +157,16 @@ class DeepseekVL2MlpProjector(nn.Module):
class DeepseekVL2ForCausalLM(nn.Module):
@staticmethod
def shared_experts_fusion_disable_reason(hf_config, quant_config):
language_config = hf_config.language_config
if not language_config.use_mla:
return None
# The language model is built without a quantization config.
return DeepseekV2ForCausalLM.shared_experts_fusion_disable_reason(
language_config, None
)
def __init__(
self,
config: DeepseekVL2Config,
python/sglang/srt/models/dots_vlm.py
+8
View File
@@ -40,6 +40,14 @@ from .dots_vlm_vit import DotsVisionTransformer
class DotsVLMForCausalLM(nn.Module):
"""DotsVLM model for sglang inference"""
@staticmethod
def shared_experts_fusion_disable_reason(hf_config, quant_config):
if hf_config.encoder_only:
return None
return DeepseekV2ForCausalLM.shared_experts_fusion_disable_reason(
hf_config.language_config, quant_config
)
def __init__(
self, config: DotsVLMConfig, quant_config: Optional[QuantizationConfig] = None
) -> None:
python/sglang/srt/models/glm4_moe.py
+20 -35
View File
@@ -68,6 +68,7 @@ from sglang.srt.layers.moe.topk import TopK
from sglang.srt.layers.moe.utils import (
RoutingMethodType,
filter_moe_weight_param_global_expert,
is_shared_experts_fusion_disabled,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention
@@ -94,7 +95,6 @@ from sglang.srt.utils import (
is_hip,
is_non_idle_and_non_empty,
is_npu,
log_info_on_rank0,
make_layers,
)
from sglang.srt.utils.hf_transformers_utils import get_rope_config
@@ -400,9 +400,7 @@ class Glm4MoeSparseMoeBlock(nn.Module):
self.routed_scaling_factor = config.routed_scaling_factor
self.n_shared_experts = config.n_shared_experts
self.num_fused_shared_experts = (
0
if get_exec().moe.disable_shared_experts_fusion
else config.n_shared_experts
0 if is_shared_experts_fusion_disabled() else config.n_shared_experts
)
self.config = config
@@ -1172,44 +1170,32 @@ class Glm4MoeForCausalLM(nn.Module):
# For EAGLE3 support
self.capture_aux_hidden_states = False
def determine_num_fused_shared_experts(self):
if get_exec().moe.disable_shared_experts_fusion:
return
disable_reason = None
@classmethod
def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
"""Why this checkpoint cannot fuse its shared expert, or None. Asked by
the loader before any layer is built."""
if (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and (
not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
):
disable_reason = (
return (
"Only GLM-4.5 on NV-platform with capability >= 80 "
"or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization."
)
elif get_parallel().moe_ep_size > 1 and (
if get_parallel().moe_ep_size > 1 and (
not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4)
):
disable_reason = "Only GLM-4.5 on AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization under expert parallelism."
elif disable_reason is None and (
get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori()
):
disable_reason = "GLM-4.5 cannot use shared experts fusion optimization under deepep expert parallelism."
elif self.quant_config and self.quant_config.get_name() == "w4afp8":
disable_reason = "GLM-4.5 W4AFP8 model uses different quant method for routed experts and shared experts."
return "Only GLM-4.5 on AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization under expert parallelism."
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori():
return "GLM-4.5 cannot use shared experts fusion optimization under deepep expert parallelism."
if quant_config and quant_config.get_name() == "w4afp8":
return "GLM-4.5 W4AFP8 model uses different quant method for routed experts and shared experts."
return None
if disable_reason is not None:
from sglang.srt.arg_groups.overrides import declare_load_time_override
declare_load_time_override(
"Glm4MoeForCausalLM.determine_num_fused_shared_experts",
{"disable_shared_experts_fusion": True},
)
self.num_fused_shared_experts = 0
log_info_on_rank0(
logger,
f"{disable_reason} Shared experts fusion optimization is disabled.",
)
return
self.num_fused_shared_experts = self.config.n_shared_experts
def determine_num_fused_shared_experts(self):
# The decision was installed by the loader; this only reads it.
self.num_fused_shared_experts = (
0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
)
def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens
@@ -1459,8 +1445,7 @@ class Glm4MoeForCausalLM(nn.Module):
class GlmMoeDsaForCausalLM(DeepseekV2ForCausalLM):
def determine_num_fused_shared_experts(self):
super().determine_num_fused_shared_experts("GlmMoeDsaForCausalLM")
fused_shared_experts_architecture = "GlmMoeDsaForCausalLM"
class GlmMoeDsaForCausalLMNextN(DeepseekV3ForCausalLMNextN):
python/sglang/srt/models/glm4_moe_lite.py
+24 -33
View File
@@ -59,7 +59,10 @@ from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod
from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat
from sglang.srt.layers.moe.utils import filter_moe_weight_param_global_expert
from sglang.srt.layers.moe.utils import (
filter_moe_weight_param_global_expert,
is_shared_experts_fusion_disabled,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.layers.vocab_parallel_embedding import (
@@ -81,7 +84,6 @@ from sglang.srt.utils import (
add_prefix,
is_non_idle_and_non_empty,
is_npu,
log_info_on_rank0,
make_layers,
)
from sglang.srt.utils.hf_transformers_utils import get_rope_config
@@ -185,9 +187,7 @@ class Glm4MoeLiteSparseMoeBlock(nn.Module):
self.routed_scaling_factor = config.routed_scaling_factor
self.n_shared_experts = config.n_shared_experts
self.num_fused_shared_experts = (
0
if get_exec().moe.disable_shared_experts_fusion
else config.n_shared_experts
0 if is_shared_experts_fusion_disabled() else config.n_shared_experts
)
self.config = config
self.layer_id = layer_id
@@ -896,7 +896,7 @@ class Glm4MoeLiteForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
self.tp_size = get_parallel().tp_size
self.quant_config = quant_config
self.pp_group = get_pp_group()
self.determine_num_fused_shared_experts("Glm4MoeLiteForCausalLM")
self.determine_num_fused_shared_experts()
self.model = Glm4MoeLiteModel(
config, quant_config, prefix=add_prefix("model", prefix)
)
@@ -922,39 +922,30 @@ class Glm4MoeLiteForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
def routed_experts_weights_of_layer(self):
return self._routed_experts_weights_of_layer.value
def determine_num_fused_shared_experts(
self, architecture: str = "Glm4MoeLiteForCausalLM"
):
self.num_fused_shared_experts = 0
if get_exec().moe.disable_shared_experts_fusion:
return
# The architecture this class fuses shared experts for; the NextN draft
# reports its own name.
fused_shared_experts_architecture = "Glm4MoeLiteForCausalLM"
disable_reason = None
@classmethod
def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
"""Why this checkpoint cannot fuse its shared expert, or None. Asked by
the loader before any layer is built."""
if (
not _is_cuda
or torch.cuda.get_device_capability("cuda") < (8, 0)
or self.config.architectures[0] != architecture
or self.config.n_shared_experts != 1
or hf_config.architectures[0] != cls.fused_shared_experts_architecture
or hf_config.n_shared_experts != 1
):
disable_reason = "Only GLM-4.5 or GLM-4.6 on NV-platform with capability >= 80 can use shared experts fusion optimization."
elif get_parallel().moe_ep_size > 1:
disable_reason = "GLM-4.5 or GLM-4.6 cannot use shared experts fusion optimization under expert parallelism."
return "Only GLM-4.5 or GLM-4.6 on NV-platform with capability >= 80 can use shared experts fusion optimization."
if get_parallel().moe_ep_size > 1:
return "GLM-4.5 or GLM-4.6 cannot use shared experts fusion optimization under expert parallelism."
return None
if disable_reason is not None:
from sglang.srt.arg_groups.overrides import declare_load_time_override
declare_load_time_override(
"Glm4MoeLiteForCausalLM.determine_num_fused_shared_experts",
{"disable_shared_experts_fusion": True},
)
self.num_fused_shared_experts = 0
log_info_on_rank0(
logger,
f"{disable_reason} Shared experts fusion optimization is disabled.",
)
return
self.num_fused_shared_experts = self.config.n_shared_experts
def determine_num_fused_shared_experts(self):
# The decision was installed by the loader; this only reads it.
self.num_fused_shared_experts = (
0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
)
def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens
python/sglang/srt/models/glm4_moe_lite_nextn.py
+9 -5
View File
@@ -35,7 +35,7 @@ from sglang.srt.models.glm4_moe_lite import (
Glm4MoeLiteDecoderLayer,
Glm4MoeLiteForCausalLM,
)
from sglang.srt.runtime_context import get_exec, get_parallel, get_spec
from sglang.srt.runtime_context import get_parallel, get_spec
from sglang.srt.utils import BumpAllocator, add_prefix, is_npu
logger = logging.getLogger(__name__)
@@ -130,6 +130,9 @@ class Glm4MoeLiteModelNextN(nn.Module):
class Glm4MoeLiteForCausalLMNextN(Glm4MoeLiteForCausalLM):
# The draft checkpoint reports the NextN architecture name.
fused_shared_experts_architecture = "Glm4MoeLiteForCausalLMNextN"
def __init__(
self,
config: PretrainedConfig,
@@ -143,6 +146,11 @@ class Glm4MoeLiteForCausalLMNextN(Glm4MoeLiteForCausalLM):
quant_config = None
self.quant_config = quant_config
# The draft's own gate (its quantization can differ from the
# target's); the decoder below reads the ACTIVE decision as it builds,
# and num_fused_shared_experts drives the inherited loader's remap.
self.determine_num_fused_shared_experts()
self.model = Glm4MoeLiteModelNextN(
config, quant_config, prefix=add_prefix("model", prefix)
)
@@ -155,10 +163,6 @@ class Glm4MoeLiteForCausalLMNextN(Glm4MoeLiteForCausalLM):
)
self.logits_processor = LogitsProcessor(config)
self.num_fused_shared_experts = (
0 if get_exec().moe.disable_shared_experts_fusion else 1
)
@torch.no_grad()
def forward(
self,
python/sglang/srt/models/glm4_moe_nextn.py
+8 -5
View File
@@ -32,7 +32,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.models.glm4_moe import Glm4MoeDecoderLayer, Glm4MoeForCausalLM
from sglang.srt.runtime_context import get_exec, get_parallel, get_spec
from sglang.srt.runtime_context import get_parallel, get_spec
from sglang.srt.utils import add_prefix, is_npu
logger = logging.getLogger(__name__)
@@ -129,6 +129,13 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
quant_config = None
self.quant_config = quant_config
# The draft's own gate: its quantization can differ from the
# target's, and the decoder below reads the ACTIVE decision while it
# builds. Also sets num_fused_shared_experts, which drives the
# inherited loader's shared-expert remap.
self.num_fused_shared_experts = 0
self.determine_num_fused_shared_experts()
self.model = Glm4MoeModelNextN(
config, quant_config, prefix=add_prefix("model", prefix)
)
@@ -141,10 +148,6 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
)
self.logits_processor = LogitsProcessor(config)
self.num_fused_shared_experts = (
0 if get_exec().moe.disable_shared_experts_fusion else 1
)
@torch.no_grad()
def forward(
self,
python/sglang/srt/models/glm4v_moe.py
+20 -28
View File
@@ -11,6 +11,7 @@ from sglang.srt.layers.attention import vision_utils
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe import get_moe_a2a_backend
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.utils import is_shared_experts_fusion_disabled
from sglang.srt.layers.pooler import Pooler, PoolingType
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer
@@ -18,7 +19,7 @@ from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead
from sglang.srt.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.glm4_moe import Glm4MoeModel
from sglang.srt.models.glm4v import Glm4vForConditionalGeneration, Glm4vVisionModel
from sglang.srt.runtime_context import get_exec, get_mm, get_parallel
from sglang.srt.runtime_context import get_mm, get_parallel
from sglang.srt.utils import add_prefix, get_device_sm, is_cuda, log_info_on_rank0
from sglang.srt.utils.hf_transformers_utils import get_processor
@@ -82,35 +83,26 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
# For EAGLE3 support
self.capture_aux_hidden_states = False
@classmethod
def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
"""Why this checkpoint cannot fuse its shared expert, or None. Asked by
the loader before any layer is built."""
if not getattr(hf_config, "n_shared_experts", None):
return "No shared experts are defined in the config."
if not _is_cuda:
return "Shared experts fusion currently requires CUDA devices."
if _is_cuda and (_device_sm is not None) and (_device_sm < 80):
return "Shared experts fusion requires SM80 or newer GPUs."
if get_parallel().moe_ep_size > 1:
return "Shared experts fusion is not supported together with expert parallelism yet."
if get_moe_a2a_backend().is_deepep():
return "Shared experts fusion is not supported when Deepep MoE backend is enabled."
return None
def determine_num_fused_shared_experts(self):
if get_exec().moe.disable_shared_experts_fusion:
# The decision was installed by the loader; this only reads it.
if is_shared_experts_fusion_disabled():
return
disable_reason = None
if not getattr(self.config, "n_shared_experts", None):
disable_reason = "No shared experts are defined in the config."
elif not _is_cuda:
disable_reason = "Shared experts fusion currently requires CUDA devices."
elif _is_cuda and (_device_sm is not None) and (_device_sm < 80):
disable_reason = "Shared experts fusion requires SM80 or newer GPUs."
elif get_parallel().moe_ep_size > 1:
disable_reason = "Shared experts fusion is not supported together with expert parallelism yet."
elif get_moe_a2a_backend().is_deepep():
disable_reason = "Shared experts fusion is not supported when Deepep MoE backend is enabled."
if disable_reason is not None:
from sglang.srt.arg_groups.overrides import declare_load_time_override
declare_load_time_override(
"Glm4vMoeForConditionalGeneration.determine_num_fused_shared_experts",
{"disable_shared_experts_fusion": True},
)
log_info_on_rank0(
logger,
f"{disable_reason} Shared experts fusion optimization is disabled.",
)
return
self.num_fused_shared_experts = self.config.n_shared_experts
assert (
self.num_fused_shared_experts == 1
python/sglang/srt/models/glm_ocr_nextn.py
+3 -4
View File
@@ -25,6 +25,7 @@ from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_r
from sglang.srt.layers.dp_attention import is_dp_attention_enabled
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe.utils import is_shared_experts_fusion_disabled
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead,
@@ -33,7 +34,7 @@ from sglang.srt.layers.vocab_parallel_embedding import (
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.models.glm4 import Glm4DecoderLayer
from sglang.srt.models.glm_ocr import GlmOcrForConditionalGeneration
from sglang.srt.runtime_context import get_exec, get_parallel
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import add_prefix
logger = logging.getLogger(__name__)
@@ -138,9 +139,7 @@ class GlmOcrForConditionalGenerationNextN(GlmOcrForConditionalGeneration):
)
self.logits_processor = LogitsProcessor(config)
self.num_fused_shared_experts = (
0 if get_exec().moe.disable_shared_experts_fusion else 1
)
self.num_fused_shared_experts = 0 if is_shared_experts_fusion_disabled() else 1
@torch.no_grad()
def forward(
python/sglang/srt/models/kimi_k25.py
+8
View File
@@ -646,6 +646,14 @@ class KimiK25ForConditionalGeneration(nn.Module):
}
)
@staticmethod
def shared_experts_fusion_disable_reason(hf_config, quant_config):
if hf_config.encoder_only:
return None
return DeepseekV3ForCausalLM.shared_experts_fusion_disable_reason(
hf_config.text_config, quant_config
)
def __init__(
self,
config: KimiK25Config,
python/sglang/srt/models/kimi_vl.py
+15 -3
View File
@@ -114,7 +114,21 @@ class KimiVLMultiModalProjector(nn.Module):
return hidden_states
def _language_model_config(config: KimiVLConfig):
text_config = copy.deepcopy(config.text_config)
text_config.architectures = ["DeepseekV2ForCausalLM"]
return text_config
class KimiVLForConditionalGeneration(nn.Module):
@staticmethod
def shared_experts_fusion_disable_reason(hf_config, quant_config):
if hf_config.encoder_only:
return None
return DeepseekV2ForCausalLM.shared_experts_fusion_disable_reason(
_language_model_config(hf_config), quant_config
)
def __init__(
self,
config: KimiVLConfig,
@@ -138,10 +152,8 @@ class KimiVLForConditionalGeneration(nn.Module):
self.language_model = None
if not config.encoder_only:
text_config = copy.deepcopy(config.text_config)
text_config.architectures = ["DeepseekV2ForCausalLM"]
self.language_model = DeepseekV2ForCausalLM(
config=text_config,
config=_language_model_config(config),
quant_config=quant_config,
prefix=add_prefix("language_model", prefix),
)
python/sglang/srt/models/minicpmv.py
+8
View File
@@ -1640,6 +1640,14 @@ class MiniCPMV:
minicpmv: nn.Module
@staticmethod
def shared_experts_fusion_disable_reason(hf_config, quant_config):
# 4.6 nests a Qwen3.5 LLM under ``text_config``; every other version
# builds a dense LLM, for which the Qwen3.5 gate answers None.
return Qwen3_5ForCausalLM.shared_experts_fusion_disable_reason(
getattr(hf_config, "text_config", hf_config), quant_config
)
def __init__(
self,
config: PretrainedConfig,
python/sglang/srt/models/minimax_m3.py
+25 -36
View File
@@ -54,7 +54,10 @@ from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.topk import TopK
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.layers.moe.utils import (
get_moe_a2a_backend,
is_shared_experts_fusion_disabled,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.rotary_embedding import get_rope
@@ -287,9 +290,7 @@ class MiniMaxM3MoE(nn.Module):
self.tp_size = get_parallel().tp_size
self.n_shared_experts = getattr(config, "n_shared_experts", None)
self.num_fused_shared_experts = (
0
if get_exec().moe.disable_shared_experts_fusion
else config.n_shared_experts
0 if is_shared_experts_fusion_disabled() else config.n_shared_experts
)
if self.tp_size > config.num_local_experts:
@@ -1465,43 +1466,31 @@ class MiniMaxM3SparseForCausalLM(nn.Module):
def get_input_embeddings(self):
return self.model.get_input_embeddings()
def determine_num_fused_shared_experts(self):
if get_exec().moe.disable_shared_experts_fusion:
return
disable_reason = None
if not getattr(self.config, "n_shared_experts", None):
disable_reason = "No shared experts are defined in the config."
elif (
self.quant_config is not None
and self.quant_config.get_name() == "modelopt_mixed"
):
disable_reason = (
@classmethod
def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
"""Why this checkpoint cannot fuse its shared expert, or None. Asked by
the loader before any layer is built."""
if not getattr(hf_config, "n_shared_experts", None):
return "No shared experts are defined in the config."
if quant_config is not None and quant_config.get_name() == "modelopt_mixed":
return (
"Shared and routed experts may use different quantization formats "
"in ModelOpt mixed-precision checkpoints."
)
elif not _is_cuda:
disable_reason = "Shared experts fusion currently requires CUDA devices."
elif _is_cuda and (_device_sm is not None) and (_device_sm < 80):
disable_reason = "Shared experts fusion requires SM80 or newer GPUs."
elif get_parallel().moe_ep_size > 1:
disable_reason = "Shared experts fusion is not supported together with expert parallelism yet."
elif get_moe_a2a_backend().is_deepep():
disable_reason = "Shared experts fusion is not supported when Deepep MoE backend is enabled."
if not _is_cuda:
return "Shared experts fusion currently requires CUDA devices."
if _is_cuda and (_device_sm is not None) and (_device_sm < 80):
return "Shared experts fusion requires SM80 or newer GPUs."
if get_parallel().moe_ep_size > 1:
return "Shared experts fusion is not supported together with expert parallelism yet."
if get_moe_a2a_backend().is_deepep():
return "Shared experts fusion is not supported when Deepep MoE backend is enabled."
return None
if disable_reason is not None:
from sglang.srt.arg_groups.overrides import declare_load_time_override
declare_load_time_override(
"MiniMaxM3ForCausalLM.determine_num_fused_shared_experts",
{"disable_shared_experts_fusion": True},
)
log_info_on_rank0(
logger,
f"{disable_reason} Shared experts fusion optimization is disabled.",
)
def determine_num_fused_shared_experts(self):
# The decision was installed by the loader; this only reads it.
if is_shared_experts_fusion_disabled():
return
self.num_fused_shared_experts = self.config.n_shared_experts
assert (
self.num_fused_shared_experts == 1
python/sglang/srt/models/minimax_m3_vl.py
+27 -35
View File
@@ -10,7 +10,10 @@ from sglang.srt.distributed import (
get_pp_group,
)
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.moe.utils import get_moe_a2a_backend
from sglang.srt.layers.moe.utils import (
get_moe_a2a_backend,
is_shared_experts_fusion_disabled,
)
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.layers.utils.common import get_layer_id
@@ -43,7 +46,7 @@ from sglang.srt.models.minimax_vl_common import (
merge_vit_qkv_weights,
)
from sglang.srt.models.utils import WeightsMapper
from sglang.srt.runtime_context import get_exec, get_mm, get_parallel
from sglang.srt.runtime_context import get_mm, get_parallel
from sglang.srt.utils import add_prefix, get_device_sm, is_cuda, log_info_on_rank0
from sglang.srt.utils.hf_transformers_utils import get_rope_config
@@ -132,51 +135,40 @@ class MiniMaxM3SparseForConditionalGeneration(nn.Module):
self.logits_processor = LogitsProcessor(text_config)
def _determine_num_fused_shared_experts(self) -> None:
text_config = self.config.text_config
if get_exec().moe.disable_shared_experts_fusion:
return
disable_reason = None
@classmethod
def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
"""Why this checkpoint cannot fuse its shared expert, or None. Asked by
the loader before any layer is built; the experts live on the text
config."""
text_config = getattr(hf_config, "text_config", hf_config)
if not getattr(text_config, "n_shared_experts", None):
disable_reason = "No shared experts are defined in the config."
elif (
self.quant_config is not None
and self.quant_config.get_name() == "modelopt_mixed"
):
disable_reason = (
return "No shared experts are defined in the config."
if quant_config is not None and quant_config.get_name() == "modelopt_mixed":
return (
"Shared and routed experts may use different quantization formats "
"in ModelOpt mixed-precision checkpoints."
)
elif not _is_cuda:
disable_reason = "Shared experts fusion currently requires CUDA devices."
elif (_device_sm is not None) and (_device_sm < 80):
disable_reason = "Shared experts fusion requires SM80 or newer GPUs."
elif get_parallel().moe_ep_size > 1:
disable_reason = (
if not _is_cuda:
return "Shared experts fusion currently requires CUDA devices."
if (_device_sm is not None) and (_device_sm < 80):
return "Shared experts fusion requires SM80 or newer GPUs."
if get_parallel().moe_ep_size > 1:
return (
"Shared experts fusion is not supported together with expert "
"parallelism yet."
)
elif get_moe_a2a_backend().is_deepep():
disable_reason = (
if get_moe_a2a_backend().is_deepep():
return (
"Shared experts fusion is not supported when Deepep MoE backend "
"is enabled."
)
return None
if disable_reason is not None:
from sglang.srt.arg_groups.overrides import declare_load_time_override
declare_load_time_override(
"MiniMaxM3VLForCausalLM._determine_num_fused_shared_experts",
{"disable_shared_experts_fusion": True},
)
log_info_on_rank0(
logger,
f"{disable_reason} Shared experts fusion optimization is disabled.",
)
def _determine_num_fused_shared_experts(self) -> None:
# The decision was installed by the loader; this only reads it.
if is_shared_experts_fusion_disabled():
return
self.num_fused_shared_experts = text_config.n_shared_experts
self.num_fused_shared_experts = self.config.text_config.n_shared_experts
assert (
self.num_fused_shared_experts == 1
), "Only 1 fused shared expert is supported"
python/sglang/srt/models/pixtral.py
+10
View File
@@ -73,6 +73,16 @@ class VisionEncoderArgs:
class PixtralForConditionalGeneration(nn.Module):
merge_by_field_config = True
@staticmethod
def shared_experts_fusion_disable_reason(hf_config, quant_config):
text_config = hf_config.text_config
if getattr(text_config, "model_type", "") != "deepseek_v3":
# The GQA text config builds the dense Mistral backbone.
return None
return MistralLarge3ForCausalLM.shared_experts_fusion_disable_reason(
text_config, quant_config
)
@classmethod
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
if modality.startswith("image"):
python/sglang/srt/models/qwen3_5.py
+41 -25
View File
@@ -57,6 +57,9 @@ from sglang.srt.layers.linear import (
RowParallelLinear,
)
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.utils import (
is_shared_experts_fusion_disabled,
)
from sglang.srt.layers.parameter import (
BlockQuantScaleParameter,
PerTensorScaleParameter,
@@ -137,9 +140,10 @@ cached_get_processor = lru_cache(get_processor)
def _disable_shared_experts_fusion() -> bool:
# Resolved lazily: the global server args is not set at module import time
# (e.g. when this module is imported by unit tests).
return get_exec().moe.disable_shared_experts_fusion
# Resolved lazily: the flag is written by the owning model's gate before
# its layers build (per runner); models without a gate see the config
# intent through the accessor's fallback.
return is_shared_experts_fusion_disabled()
if _is_cuda:
@@ -1308,25 +1312,6 @@ class Qwen3_5ForCausalLM(nn.Module):
f"get_hidden_dim not implemented for {module_name}"
)
def _maybe_autodisable_shared_experts_fusion(self, config, quant_config):
# Auto-disable fusion when the checkpoint can't fuse (e.g. MXFP4 Qwen3.5)
# so the model still gets the #25885 multi-streaming path. ROCm-only.
if (
config.model_type == "qwen3_5_moe_text"
and not get_exec().moe.disable_shared_experts_fusion
and not can_fuse_shared_expert(config, quant_config)
):
from sglang.srt.arg_groups.overrides import declare_load_time_override
declare_load_time_override(
"Qwen3_5ForCausalLM._maybe_autodisable_shared_experts_fusion",
{"disable_shared_experts_fusion": True},
)
logger.info(
"Qwen3.5: shared-expert fusion not supported for this checkpoint; "
"auto-disabling (multi-streaming #25885 still applies)."
)
def __init__(
self,
config: Qwen3_5TextConfig,
@@ -1339,9 +1324,6 @@ class Qwen3_5ForCausalLM(nn.Module):
self.hidden_size = config.hidden_size
self.pp_group = get_pp_group()
if _is_hip:
self._maybe_autodisable_shared_experts_fusion(config, quant_config)
alt_stream = get_stream("alt") if _is_cuda or _hip_use_alt_stream else None
# Embedding layer
@@ -2316,4 +2298,38 @@ class Qwen3_5MoeForConditionalGeneration(Qwen3VLForConditionalGeneration):
)
def _qwen3_5_shared_experts_fusion_disable_reason(hf_config, quant_config):
"""Why this Qwen3.5 checkpoint cannot fuse its shared expert, or None.
ROCm-only: an MXFP4 checkpoint cannot fuse, and the model still wants the
#25885 multi-streaming path. Asked by the loader before any layer is built,
so it resolves the text config itself -- the loader hands over whichever
config the entry class takes.
"""
if not _is_hip:
return None
text_config = getattr(hf_config, "text_config", hf_config)
if getattr(text_config, "model_type", None) != "qwen3_5_moe_text":
return None
if can_fuse_shared_expert(text_config, quant_config):
return None
return (
"Qwen3.5: shared-expert fusion not supported for this checkpoint "
"(multi-streaming #25885 still applies)."
)
# Every class the loader may instantiate for a Qwen3.5 checkpoint answers the
# fusion question the same way.
for _entry_class in (
Qwen3_5ForCausalLM,
Qwen3_5MoeForCausalLM,
Qwen3_5ForConditionalGeneration,
Qwen3_5MoeForConditionalGeneration,
):
_entry_class.shared_experts_fusion_disable_reason = staticmethod(
_qwen3_5_shared_experts_fusion_disable_reason
)
EntryClass = [Qwen3_5MoeForConditionalGeneration, Qwen3_5ForConditionalGeneration]
python/sglang/srt/models/qwen3_5_mtp.py
+42 -25
View File
@@ -44,8 +44,49 @@ from sglang.srt.utils import add_prefix, is_npu
logger = logging.getLogger(__name__)
def _mtp_quant_config(quant_config):
"""The quantization the MTP module itself is built with.
The MTP module often ships unquantized even though the target checkpoint is
quantized; the loader's fusion gate has to see the same normalization the
constructor applies, or it would answer for the target's quantization.
"""
# Serialized Qwen3.5 ModelOpt checkpoints keep embedded MTP weights in
# BF16. Disable quantization for those checkpoints; non-serialized
# modelopt_fp4 still converts MoE expert weights on load.
if quant_config and (
quant_config.get_name() == "modelopt_mixed"
or (
quant_config.get_name() == "modelopt_fp4"
and quant_config.is_checkpoint_nvfp4_serialized
)
):
return None
if is_npu() and get_spec().speculative_draft_model_quantization is None:
return None
# Quark-quantized Qwen3.5 MXFP4 checkpoints ship the MTP module in bf16;
# every `mtp.*` layer appears under the quantization exclude list. Detect
# that and skip quantization here so linear/MoE weight loaders allocate
# bf16 shapes (see sgl-project/sglang#23113).
if quant_config and quant_config.get_name() == "quark":
exclude_layers = getattr(quant_config, "exclude_layers", [])
if any(
isinstance(layer, str) and layer.startswith("mtp.")
for layer in exclude_layers
):
return None
return quant_config
class Qwen3_5ForCausalLMMTP(nn.Module):
@staticmethod
def shared_experts_fusion_disable_reason(hf_config, quant_config):
return Qwen3_5ForCausalLM.shared_experts_fusion_disable_reason(
getattr(hf_config, "text_config", hf_config),
_mtp_quant_config(quant_config),
)
def __init__(
self,
config: PretrainedConfig,
@@ -61,31 +102,7 @@ class Qwen3_5ForCausalLMMTP(nn.Module):
# Deep-copy so MTP mutations below don't leak into the target's config.
config = copy.deepcopy(config)
# Serialized Qwen3.5 ModelOpt checkpoints keep embedded MTP weights in
# BF16. Disable quantization for those checkpoints; non-serialized
# modelopt_fp4 still converts MoE expert weights on load.
if quant_config and (
quant_config.get_name() == "modelopt_mixed"
or (
quant_config.get_name() == "modelopt_fp4"
and quant_config.is_checkpoint_nvfp4_serialized
)
):
quant_config = None
if is_npu() and get_spec().speculative_draft_model_quantization is None:
quant_config = None
# Quark-quantized Qwen3.5 MXFP4 checkpoints ship the MTP module in
# bf16; every `mtp.*` layer appears under the quantization exclude
# list. Detect that and skip quantization here so linear/MoE weight
# loaders allocate bf16 shapes (see sgl-project/sglang#23113).
if quant_config and quant_config.get_name() == "quark":
exclude_layers = getattr(quant_config, "exclude_layers", [])
if any(
isinstance(layer, str) and layer.startswith("mtp.")
for layer in exclude_layers
):
quant_config = None
quant_config = _mtp_quant_config(quant_config)
self.config = config
self.tp_size = get_parallel().tp_size
python/sglang/srt/models/qwen3_5_text.py
+7
View File
@@ -43,6 +43,13 @@ class Qwen3_5ForCausalLM(nn.Module):
packed_modules_mapping = qwen3_5.Qwen3_5ForCausalLM.packed_modules_mapping
supported_lora_modules = qwen3_5.Qwen3_5ForCausalLM.supported_lora_modules
@classmethod
def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
# The body decides; it is handed this config and quantization verbatim.
return cls.body_cls.shared_experts_fusion_disable_reason(
hf_config, quant_config
)
def __init__(
self,
config,
python/sglang/srt/runtime_context.py
+14
View File
@@ -376,6 +376,20 @@ class MoeFlags(_FlagGroupBase):
tbo_token_distribution_threshold: float | None = None
disable_fp4_allgather: bool | None = None
quantization: str | None = None
# The shared-experts-fusion decision, per runner — the runner_backend /
# speculative_runner_backend shape. Both leaves are seeded from the config
# intent by ``initialize_moe_config``; each MoE model's gate
# (determine_num_fused_shared_experts) refines the ACTIVE leaf, both ways,
# before its layers build and read it. ``speculative_moe_backend_context``
# brackets a draft's build: on exit the draft's effective decision is
# persisted onto the speculative leaf (inspectable afterwards) and the
# target's ACTIVE value returns.
disable_shared_experts_fusion: bool | None = None
speculative_disable_shared_experts_fusion: bool | None = None
# Lifecycle marker (the capture.disable_dispose_tensor family): set while
# speculative_moe_backend_context is active, so a draft gate's write also
# lands on the speculative leaf.
in_speculative_scope: bool = False
@dataclasses.dataclass
python/sglang/srt/speculative/draft_worker_common.py
+18 -10
View File
@@ -80,16 +80,24 @@ def build_draft_tp_worker(
server_args=server_args, algo_label=algo_label
)
)
draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
ps=ps,
nccl_port=nccl_port,
is_draft_worker=True,
# The draft runs at absolute target positions.
context_length=target_model_config.context_len,
draft_attention_backend=draft_backend,
)
from sglang.srt.layers.moe.utils import draft_model_build_scope
# The draft's model construction runs its own MoE gates; the scope routes
# their fusion decision to the speculative leaf and gives the target its
# ACTIVE value back. It deliberately does not swap runner_backend: these
# workers run the draft outside speculative_moe_backend_context, so a
# construction-only swap would build and execute under different backends.
with draft_model_build_scope():
draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
ps=ps,
nccl_port=nccl_port,
is_draft_worker=True,
# The draft runs at absolute target positions.
context_length=target_model_config.context_len,
draft_attention_backend=draft_backend,
)
draft_model_runner = draft_worker.model_runner
draft_worker.draft_runner = draft_model_runner
python/sglang/srt/speculative/eagle_worker_v2.py
+2 -1
View File
@@ -26,6 +26,7 @@ from sglang.srt.layers.attention.trtllm_mla_backend import (
TRTLLMMLABackend,
)
from sglang.srt.layers.moe.utils import (
draft_model_build_scope,
speculative_moe_a2a_backend_context,
speculative_moe_backend_context,
)
@@ -162,7 +163,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
ctx = empty_context()
with (
ctx
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), draft_model_build_scope():
self.draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py
+2 -1
View File
@@ -30,6 +30,7 @@ import torch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.moe.utils import (
draft_model_build_scope,
speculative_moe_a2a_backend_context,
speculative_moe_backend_context,
)
@@ -127,7 +128,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
with (
empty_context()
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), draft_model_build_scope():
# Both base classes own initialization, so initialize TpModelWorker
# explicitly after EagleDraftWorkerBase above.
TpModelWorker.__init__(
python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py
+5 -2
View File
@@ -26,7 +26,10 @@ from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.graph_runner.multi_layer_eagle_draft_extend_npu_graph_runner import (
MultiLayerEagleMultiStepDraftExtendNpuGraphRunner,
)
from sglang.srt.layers.moe.utils import speculative_moe_backend_context
from sglang.srt.layers.moe.utils import (
draft_model_build_scope,
speculative_moe_backend_context,
)
from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker
@@ -150,7 +153,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
)
# Load draft model weights only.
with empty_context(), speculative_moe_backend_context():
with empty_context(), speculative_moe_backend_context(), draft_model_build_scope():
self.draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
python/sglang/srt/speculative/standalone_worker_v2.py
+9 -3
View File
@@ -5,7 +5,10 @@ from typing import Optional
import torch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.moe.utils import speculative_moe_backend_context
from sglang.srt.layers.moe.utils import (
draft_model_build_scope,
speculative_moe_backend_context,
)
from sglang.srt.managers.tp_worker import TpModelWorker
from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.adaptive_runtime_state import (
@@ -66,8 +69,11 @@ class StandaloneDraftWorker(EagleDraftWorker):
self.speculative_num_steps * self.topk, self.speculative_num_draft_tokens
)
# Load draft model weights only.
with empty_context():
# Load draft model weights only. The standalone draft is a real model
# whose MoE gates run during construction; the scope routes their
# fusion decision to the speculative leaf (it does not swap
# runner_backend — the draft's forwards run outside that context).
with empty_context(), draft_model_build_scope():
self.draft_worker = TpModelWorker(
server_args=server_args,
gpu_id=gpu_id,
python/sglang/srt/state_capturer/routed_experts.py
+3 -6
View File
@@ -41,12 +41,9 @@ class RoutedExpertsCapturer(BaseTopkCapturer):
) -> Optional["RoutedExpertsCapturer"]:
if not get_exec().features.enable_return_routed_experts:
return None
if not get_exec().moe.disable_shared_experts_fusion and hasattr(
model, "num_fused_shared_experts"
):
num_fused_shared_experts = model.num_fused_shared_experts
else:
num_fused_shared_experts = 0
# The model's own attribute is the baked decision (0 when its gate
# disabled fusion); the ACTIVE flag can be holding another runner's.
num_fused_shared_experts = getattr(model, "num_fused_shared_experts", 0)
return RoutedExpertsCapturer(
model_config,
num_tokens=num_tokens,