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

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
+15 -9
View File
@@ -142,14 +142,20 @@ Never assign `server_args` fields from model code. Declare instead
callable receives *pristine* `server_args` + `hf_config` and must not write. callable receives *pristine* `server_args` + `hf_config` and must not write.
- Normalization that must see earlier declarations → a post-process pass invoked via - Normalization that must see earlier declarations → a post-process pass invoked via
`run_post_process_pass` at its slot (reads a view, returns a declaration dict). `run_post_process_pass` at its slot (reads a view, returns a declaration dict).
- Values only knowable at weight-load time → `declare_load_time_override(source, {...})` - Values only knowable at load time are **per-runner state**, not declarations:
— validates the whitelist, then routes through `get_context().override` (**bag-only**; there is no `declare_load_time_override` any more. A model-family decision that
the declaration lands on the published bags, not on any `ServerArgs` instance). its checkpoint drives (shared-experts fusion) is a question the *loader* asks
Scope caveat for draft models: only a draft build that publishes a private copy the model class — `shared_experts_fusion_disable_reason(hf_config,
under `preserve_config` discards its declarations with the scope. Draft loads quant_config)`, a classmethod answering without an instance — at the single
that skip publish share the process bags, so their declarations land model-instantiation point, and
process-wide — declares reachable from a draft load must be draft-safe (guard `install_shared_experts_fusion_decision` writes the answer to the ACTIVE moe
or same-value). flag before that model's layers build and read it
(`is_shared_experts_fusion_disabled`, config-intent fallback).
`draft_model_build_scope` brackets every draft build and routes the draft's
answer to the speculative leaf, so a draft's decision never overwrites the
target's. A process-level load-time fact (the sm80 dtype fallback —
device-driven, identical for every runner) records directly via
`get_context().override`.
Declarable fields form a whitelist: `Arg(..., resolvable=True)` in the `ServerArgs` Declarable fields form a whitelist: `Arg(..., resolvable=True)` in the `ServerArgs`
dataclass. A declaration against a non-whitelisted field fails at its slot. dataclass. A declaration against a non-whitelisted field fails at its slot.
@@ -307,7 +313,7 @@ Never module-skip a test "until the migration settles" — seed the context inst
Key source files: `python/sglang/srt/runtime_context.py` (the container, every tier, Key source files: `python/sglang/srt/runtime_context.py` (the container, every tier,
`publish`, `_ConfigBag`, `preserve_config`, `override_server_args`), `publish`, `_ConfigBag`, `preserve_config`, `override_server_args`),
`python/sglang/srt/arg_groups/overrides.py` (override registry, passes, `python/sglang/srt/arg_groups/overrides.py` (override registry, passes,
`declare_load_time_override`), `python/sglang/srt/server_args.py` (`NS` metadata, `declare_late_resolution`), `python/sglang/srt/server_args.py` (`NS` metadata,
`Arg(..., resolvable=True)`, `__setattr__` strict guard), and the guardrail tests under `Arg(..., resolvable=True)`, `__setattr__` strict guard), and the guardrail tests under
`test/registered/unit/` (`test_server_args_mutation_ratchet.py`, `test/registered/unit/` (`test_server_args_mutation_ratchet.py`,
`test_server_args_writer_ratchet.py`, `test_legacy_global_ratchet.py`, `test_server_args_writer_ratchet.py`, `test_legacy_global_ratchet.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( def collect_model_override_declarations(
architecture: str, server_args: Any, hf_config: Any architecture: str, server_args: Any, hf_config: Any
) -> List[Tuple[str, Dict[str, Any]]]: ) -> List[Tuple[str, Dict[str, Any]]]:
+90
View File
@@ -21,6 +21,7 @@ if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
from sglang.srt.runtime_context import get_server_args from sglang.srt.runtime_context import get_server_args
from sglang.srt.utils.common import log_info_on_rank0
logger = logging.getLogger(__name__) 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.tbo_token_distribution_threshold = server_args.tbo_token_distribution_threshold
moe.disable_fp4_allgather = server_args.disable_flashinfer_cutlass_moe_fp4_allgather moe.disable_fp4_allgather = server_args.disable_flashinfer_cutlass_moe_fp4_allgather
moe.quantization = server_args.quantization 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: def get_moe_a2a_backend() -> MoeA2ABackend:
@@ -357,6 +364,88 @@ def get_speculative_moe_a2a_backend() -> MoeA2ABackend:
return moe.speculative_a2a_backend 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: def get_deepep_mode() -> DeepEPMode:
moe = get_flags().moe moe = get_flags().moe
if moe.deepep_mode is None: 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. 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. This ensures that draft models in speculative decoding use the configured speculative backend.
""" """
moe = get_flags().moe moe = get_flags().moe
original_backend = moe.runner_backend original_backend = moe.runner_backend
@@ -2043,6 +2043,9 @@ class ModelRunner:
load_config: LoadConfig, load_config: LoadConfig,
) -> None: ) -> None:
self.model = new_model self.model = new_model
# 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( get_context().override(
"model_runner.update_model_fields", "model_runner.update_model_fields",
model_path=model_path, model_path=model_path,
@@ -68,11 +68,11 @@ def maybe_downgrade_dtype_for_legacy_gpu(
logger.info( logger.info(
"Compute capability below sm80. Use float16 due to lack of bfloat16 support." "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( # Device-driven, so every runner in the process resolves the same way;
"ModelRunner._sm80_dtype_fallback", {"dtype": "float16"} # 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 model_config.dtype = torch.float16
if torch.cuda.get_device_capability()[1] < 5: if torch.cuda.get_device_capability()[1] < 5:
raise RuntimeError("SGLang only supports sm75 and above.") raise RuntimeError("SGLang only supports sm75 and above.")
+10
View File
@@ -83,6 +83,9 @@ from sglang.srt.distributed import (
model_parallel_is_initialized, model_parallel_is_initialized,
) )
from sglang.srt.layers.modelopt_utils import QUANT_CFG_CHOICES 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.layers.quantization.base_config import QuantizationConfig
from sglang.srt.model_loader.remote_instance_weight_loader_utils import ( from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
trigger_transferring_weights_request, trigger_transferring_weights_request,
@@ -314,6 +317,13 @@ def _initialize_model(
) -> nn.Module: ) -> nn.Module:
"""Initialize a model with the given configurations.""" """Initialize a model with the given configurations."""
model_class, _ = get_model_architecture(model_config) 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 = { kwargs = {
"config": model_config.hf_config, "config": model_config.hf_config,
"quant_config": quant_config, "quant_config": quant_config,
+3 -1
View File
@@ -309,6 +309,8 @@ class DeepseekModelNextN(nn.Module):
class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): 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*. # Support amd/DeepSeek-R1-0528-MXFP4 renaming: model.layers.61*.
# Ref: HF config.json for amd/DeepSeek-R1-0528-MXFP4 # Ref: HF config.json for amd/DeepSeek-R1-0528-MXFP4
@@ -343,7 +345,7 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
self.quant_config = quant_config self.quant_config = quant_config
# if not set, model load will be broken in DeepseekV3ForCausalLM load_weights() # if not set, model load will be broken in DeepseekV3ForCausalLM load_weights()
self.pp_group = get_pp_group() 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.use_dsa = is_deepseek_dsa(config)
self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp() 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 self.mla_enable_prefill_cp = is_mla_prefill_cp_enabled() and not self.use_dsa
+21 -5
View File
@@ -1422,7 +1422,27 @@ def build_qwen2_decoder_as_encoder(
return 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): 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__( def __init__(
self, self,
*, *,
@@ -1437,11 +1457,7 @@ class DeepseekOCRForCausalLM(nn.Module):
self.vision_config = config.vision_config self.vision_config = config.vision_config
self.projector_config = config.projector_config self.projector_config = config.projector_config
self.text_config = config.text_config self.text_config = config.text_config
self.is_ocr2 = ( self.is_ocr2 = _is_ocr2(config)
str(getattr(self.vision_config, "model_name", "")).lower()
== "deepencoderv2"
or getattr(self.projector_config, "input_dim", None) == 896
)
n_embed = getattr(self.projector_config, "n_embed", 1280) n_embed = getattr(self.projector_config, "n_embed", 1280)
self.tile_tag = config.tile_tag self.tile_tag = config.tile_tag
+36 -43
View File
@@ -117,6 +117,7 @@ from sglang.srt.layers.moe.utils import (
has_per_rank_fused_shared_slots, has_per_rank_fused_shared_slots,
is_deepep_class_backend, is_deepep_class_backend,
is_sbo_enabled, is_sbo_enabled,
is_shared_experts_fusion_disabled,
is_tbo_enabled, is_tbo_enabled,
) )
from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig
@@ -203,7 +204,6 @@ from sglang.srt.utils import (
LazyValue, LazyValue,
add_prefix, add_prefix,
is_non_idle_and_non_empty, is_non_idle_and_non_empty,
log_info_on_rank0,
make_layers, make_layers,
use_intel_amx_backend, use_intel_amx_backend,
) )
@@ -571,7 +571,7 @@ class DeepseekV2MoE(nn.Module):
n_shared_experts = ( n_shared_experts = (
0 if config.n_shared_experts is None else int(config.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: # num_fused_shared_experts drives weight remapping in deepseek_weight_loader:
# mlp.shared_experts → mlp.experts.256 when > 0. # mlp.shared_experts → mlp.experts.256 when > 0.
@@ -2966,23 +2966,27 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
def routed_experts_weights_of_layer(self): def routed_experts_weights_of_layer(self):
return self._routed_experts_weights_of_layer.value return self._routed_experts_weights_of_layer.value
def determine_num_fused_shared_experts( # The architecture this class fuses shared experts for; a subclass whose
self, architecture: str = "DeepseekV3ForCausalLM" # checkpoint reports a different name (the NextN drafts, GLM's DSA variant)
): # overrides it.
self.num_fused_shared_experts = 0 fused_shared_experts_architecture = "DeepseekV3ForCausalLM"
if get_exec().moe.disable_shared_experts_fusion: @classmethod
return 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: if get_exec().moe.enforce_shared_experts_fusion:
pass return None
elif is_sbo_enabled() or is_tbo_enabled(): if is_sbo_enabled() or is_tbo_enabled():
disable_reason = "SBO/TBO enabled: incompatible with fusing shared expert into MoE kernel." return "SBO/TBO enabled: incompatible with fusing shared expert into MoE kernel."
elif is_deepep_class_backend(): if is_deepep_class_backend():
disable_reason = "DeepEP: fusion off by default (use --enforce-shared-experts-fusion to enable)." return "DeepEP: fusion off by default (use --enforce-shared-experts-fusion to enable)."
elif ( if (
self.config.architectures[0] != architecture hf_config.architectures[0] != cls.fused_shared_experts_architecture
# Allow-list of n_routed_experts values that have been validated # Allow-list of n_routed_experts values that have been validated
# for shared-experts fusion under this code path. Currently: # for shared-experts fusion under this code path. Currently:
# 256 -> DeepSeek-V3 / R1 # 256 -> DeepSeek-V3 / R1
@@ -2991,51 +2995,40 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
# moonshotai/Kimi-K2.5 (compressed-tensors) checkpoint # moonshotai/Kimi-K2.5 (compressed-tensors) checkpoint
# stores the shared expert loose and is NOT pre-fused, # stores the shared expert loose and is NOT pre-fused,
# so the fused path silently mis-loads it. # so the fused path silently mis-loads it.
or self.config.n_routed_experts not in (256, 384) or hf_config.n_routed_experts not in (256, 384)
or self.config.n_shared_experts != 1 or hf_config.n_shared_experts != 1
or ( or (
self.config.n_routed_experts == 384 hf_config.n_routed_experts == 384
and ( and (quant_config is None or quant_config.get_name() != "quark")
self.quant_config is None or self.quant_config.get_name() != "quark"
)
) )
): ):
disable_reason = "Config does not support fused shared expert(s)." return "Config does not support fused shared expert(s)."
elif ( if (
(not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) (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_hip or torch.cuda.get_device_capability("cuda") < (9, 4))
and (not _is_musa or torch.musa.get_device_capability("musa") < (3, 1)) 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 " "Only Deepseek V3/R1 on NV-platform with capability >= 80 "
"or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization." "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." "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) 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) " "Only Deepseek V3/R1 on AMD-platform with capability >= gfx942(MI30x) "
"can use shared experts fusion optimization under expert parallelism." "can use shared experts fusion optimization under expert parallelism."
) )
elif is_wint4afp8_or_wint4a16_config(self.quant_config): if is_wint4afp8_or_wint4a16_config(quant_config):
disable_reason = "Deepseek V3/R1 W4AFP8/W4A16 model uses different quant method for routed experts and shared experts." 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: def determine_num_fused_shared_experts(self):
from sglang.srt.arg_groups.overrides import declare_load_time_override # The decision was installed by the loader; this only reads it.
self.num_fused_shared_experts = (
declare_load_time_override( 0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
"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 get_input_embeddings(self) -> nn.Embedding: def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens return self.model.embed_tokens
+15 -24
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.logits_processor import LogitsProcessor
from sglang.srt.layers.moe import get_moe_a2a_backend, should_use_dp_reduce_scatterv 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.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 ( from sglang.srt.layers.quantization.fp8_utils import (
view_aiter_fused_rms_transposed_fp8_scale, view_aiter_fused_rms_transposed_fp8_scale,
) )
@@ -2649,35 +2650,25 @@ class DeepseekV4ForCausalLM(nn.Module):
self.capture_aux_hidden_states = True self.capture_aux_hidden_states = True
self.model.dspark_layers_to_capture = list(layer_ids) self.model.dspark_layers_to_capture = list(layer_ids)
def determine_num_fused_shared_experts(self): @classmethod
self.num_fused_shared_experts = 0 def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
if get_exec().moe.disable_shared_experts_fusion: """V4 only fuses when explicitly asked to, and then the checkpoint must
return carry exactly one shared expert. Asked by the loader before any layer is
built."""
disable_reason = None if not get_exec().moe.enforce_shared_experts_fusion:
if get_exec().moe.enforce_shared_experts_fusion: return "Config does not support fused shared expert(s)."
if self.config.n_shared_experts != 1: if hf_config.n_shared_experts != 1:
raise ValueError( raise ValueError(
"DeepSeek V4 shared-experts fusion expects exactly one shared " "DeepSeek V4 shared-experts fusion expects exactly one shared "
f"expert, but got n_shared_experts={self.config.n_shared_experts}." f"expert, but got n_shared_experts={hf_config.n_shared_experts}."
) )
else: return None
disable_reason = "Config does not support fused shared expert(s)."
if disable_reason is not None: def determine_num_fused_shared_experts(self):
from sglang.srt.arg_groups.overrides import declare_load_time_override # The decision was installed by the loader; this only reads it.
self.num_fused_shared_experts = (
declare_load_time_override( 0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
"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
@torch.no_grad() @torch.no_grad()
def forward( def forward(
+10
View File
@@ -157,6 +157,16 @@ class DeepseekVL2MlpProjector(nn.Module):
class DeepseekVL2ForCausalLM(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__( def __init__(
self, self,
config: DeepseekVL2Config, config: DeepseekVL2Config,
+8
View File
@@ -40,6 +40,14 @@ from .dots_vlm_vit import DotsVisionTransformer
class DotsVLMForCausalLM(nn.Module): class DotsVLMForCausalLM(nn.Module):
"""DotsVLM model for sglang inference""" """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__( def __init__(
self, config: DotsVLMConfig, quant_config: Optional[QuantizationConfig] = None self, config: DotsVLMConfig, quant_config: Optional[QuantizationConfig] = None
) -> None: ) -> None:
+19 -34
View File
@@ -68,6 +68,7 @@ from sglang.srt.layers.moe.topk import TopK
from sglang.srt.layers.moe.utils import ( from sglang.srt.layers.moe.utils import (
RoutingMethodType, RoutingMethodType,
filter_moe_weight_param_global_expert, filter_moe_weight_param_global_expert,
is_shared_experts_fusion_disabled,
) )
from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
@@ -94,7 +95,6 @@ from sglang.srt.utils import (
is_hip, is_hip,
is_non_idle_and_non_empty, is_non_idle_and_non_empty,
is_npu, is_npu,
log_info_on_rank0,
make_layers, make_layers,
) )
from sglang.srt.utils.hf_transformers_utils import get_rope_config 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.routed_scaling_factor = config.routed_scaling_factor
self.n_shared_experts = config.n_shared_experts self.n_shared_experts = config.n_shared_experts
self.num_fused_shared_experts = ( self.num_fused_shared_experts = (
0 0 if is_shared_experts_fusion_disabled() else config.n_shared_experts
if get_exec().moe.disable_shared_experts_fusion
else config.n_shared_experts
) )
self.config = config self.config = config
@@ -1172,44 +1170,32 @@ class Glm4MoeForCausalLM(nn.Module):
# For EAGLE3 support # For EAGLE3 support
self.capture_aux_hidden_states = False self.capture_aux_hidden_states = False
def determine_num_fused_shared_experts(self): @classmethod
if get_exec().moe.disable_shared_experts_fusion: def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
return """Why this checkpoint cannot fuse its shared expert, or None. Asked by
the loader before any layer is built."""
disable_reason = None
if (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and ( 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) 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 " "Only GLM-4.5 on NV-platform with capability >= 80 "
"or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization." "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) 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." return "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 ( if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mori():
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":
disable_reason = "GLM-4.5 cannot use shared experts fusion optimization under deepep expert parallelism." return "GLM-4.5 W4AFP8 model uses different quant method for routed experts and shared experts."
elif self.quant_config and self.quant_config.get_name() == "w4afp8": return None
disable_reason = "GLM-4.5 W4AFP8 model uses different quant method for routed experts and shared experts."
if disable_reason is not None: def determine_num_fused_shared_experts(self):
from sglang.srt.arg_groups.overrides import declare_load_time_override # The decision was installed by the loader; this only reads it.
self.num_fused_shared_experts = (
declare_load_time_override( 0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
"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 get_input_embeddings(self) -> nn.Embedding: def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens return self.model.embed_tokens
@@ -1459,8 +1445,7 @@ class Glm4MoeForCausalLM(nn.Module):
class GlmMoeDsaForCausalLM(DeepseekV2ForCausalLM): class GlmMoeDsaForCausalLM(DeepseekV2ForCausalLM):
def determine_num_fused_shared_experts(self): fused_shared_experts_architecture = "GlmMoeDsaForCausalLM"
super().determine_num_fused_shared_experts("GlmMoeDsaForCausalLM")
class GlmMoeDsaForCausalLMNextN(DeepseekV3ForCausalLMNextN): class GlmMoeDsaForCausalLMNextN(DeepseekV3ForCausalLMNextN):
+23 -32
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.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod
from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat 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.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.layers.vocab_parallel_embedding import (
@@ -81,7 +84,6 @@ from sglang.srt.utils import (
add_prefix, add_prefix,
is_non_idle_and_non_empty, is_non_idle_and_non_empty,
is_npu, is_npu,
log_info_on_rank0,
make_layers, make_layers,
) )
from sglang.srt.utils.hf_transformers_utils import get_rope_config 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.routed_scaling_factor = config.routed_scaling_factor
self.n_shared_experts = config.n_shared_experts self.n_shared_experts = config.n_shared_experts
self.num_fused_shared_experts = ( self.num_fused_shared_experts = (
0 0 if is_shared_experts_fusion_disabled() else config.n_shared_experts
if get_exec().moe.disable_shared_experts_fusion
else config.n_shared_experts
) )
self.config = config self.config = config
self.layer_id = layer_id self.layer_id = layer_id
@@ -896,7 +896,7 @@ class Glm4MoeLiteForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
self.quant_config = quant_config self.quant_config = quant_config
self.pp_group = get_pp_group() self.pp_group = get_pp_group()
self.determine_num_fused_shared_experts("Glm4MoeLiteForCausalLM") self.determine_num_fused_shared_experts()
self.model = Glm4MoeLiteModel( self.model = Glm4MoeLiteModel(
config, quant_config, prefix=add_prefix("model", prefix) config, quant_config, prefix=add_prefix("model", prefix)
) )
@@ -922,39 +922,30 @@ class Glm4MoeLiteForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
def routed_experts_weights_of_layer(self): def routed_experts_weights_of_layer(self):
return self._routed_experts_weights_of_layer.value return self._routed_experts_weights_of_layer.value
def determine_num_fused_shared_experts( # The architecture this class fuses shared experts for; the NextN draft
self, architecture: str = "Glm4MoeLiteForCausalLM" # reports its own name.
): fused_shared_experts_architecture = "Glm4MoeLiteForCausalLM"
self.num_fused_shared_experts = 0
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 ( if (
not _is_cuda not _is_cuda
or torch.cuda.get_device_capability("cuda") < (8, 0) or torch.cuda.get_device_capability("cuda") < (8, 0)
or self.config.architectures[0] != architecture or hf_config.architectures[0] != cls.fused_shared_experts_architecture
or self.config.n_shared_experts != 1 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." return "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: if 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 "GLM-4.5 or GLM-4.6 cannot use shared experts fusion optimization under expert parallelism."
return None
if disable_reason is not None: def determine_num_fused_shared_experts(self):
from sglang.srt.arg_groups.overrides import declare_load_time_override # The decision was installed by the loader; this only reads it.
self.num_fused_shared_experts = (
declare_load_time_override( 0 if is_shared_experts_fusion_disabled() else self.config.n_shared_experts
"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 get_input_embeddings(self) -> nn.Embedding: def get_input_embeddings(self) -> nn.Embedding:
return self.model.embed_tokens return self.model.embed_tokens
@@ -35,7 +35,7 @@ from sglang.srt.models.glm4_moe_lite import (
Glm4MoeLiteDecoderLayer, Glm4MoeLiteDecoderLayer,
Glm4MoeLiteForCausalLM, 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 from sglang.srt.utils import BumpAllocator, add_prefix, is_npu
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -130,6 +130,9 @@ class Glm4MoeLiteModelNextN(nn.Module):
class Glm4MoeLiteForCausalLMNextN(Glm4MoeLiteForCausalLM): class Glm4MoeLiteForCausalLMNextN(Glm4MoeLiteForCausalLM):
# The draft checkpoint reports the NextN architecture name.
fused_shared_experts_architecture = "Glm4MoeLiteForCausalLMNextN"
def __init__( def __init__(
self, self,
config: PretrainedConfig, config: PretrainedConfig,
@@ -143,6 +146,11 @@ class Glm4MoeLiteForCausalLMNextN(Glm4MoeLiteForCausalLM):
quant_config = None quant_config = None
self.quant_config = quant_config 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( self.model = Glm4MoeLiteModelNextN(
config, quant_config, prefix=add_prefix("model", prefix) config, quant_config, prefix=add_prefix("model", prefix)
) )
@@ -155,10 +163,6 @@ class Glm4MoeLiteForCausalLMNextN(Glm4MoeLiteForCausalLM):
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
self.num_fused_shared_experts = (
0 if get_exec().moe.disable_shared_experts_fusion else 1
)
@torch.no_grad() @torch.no_grad()
def forward( def forward(
self, self,
+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.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.models.glm4_moe import Glm4MoeDecoderLayer, Glm4MoeForCausalLM 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 from sglang.srt.utils import add_prefix, is_npu
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -129,6 +129,13 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
quant_config = None quant_config = None
self.quant_config = quant_config 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( self.model = Glm4MoeModelNextN(
config, quant_config, prefix=add_prefix("model", prefix) config, quant_config, prefix=add_prefix("model", prefix)
) )
@@ -141,10 +148,6 @@ class Glm4MoeForCausalLMNextN(Glm4MoeForCausalLM):
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
self.num_fused_shared_experts = (
0 if get_exec().moe.disable_shared_experts_fusion else 1
)
@torch.no_grad() @torch.no_grad()
def forward( def forward(
self, self,
+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.logits_processor import LogitsProcessor
from sglang.srt.layers.moe import get_moe_a2a_backend 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.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.pooler import Pooler, PoolingType
from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer 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.model_loader.weight_utils import default_weight_loader
from sglang.srt.models.glm4_moe import Glm4MoeModel from sglang.srt.models.glm4_moe import Glm4MoeModel
from sglang.srt.models.glm4v import Glm4vForConditionalGeneration, Glm4vVisionModel 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 import add_prefix, get_device_sm, is_cuda, log_info_on_rank0
from sglang.srt.utils.hf_transformers_utils import get_processor from sglang.srt.utils.hf_transformers_utils import get_processor
@@ -82,35 +83,26 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
# For EAGLE3 support # For EAGLE3 support
self.capture_aux_hidden_states = False 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): 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 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 self.num_fused_shared_experts = self.config.n_shared_experts
assert ( assert (
self.num_fused_shared_experts == 1 self.num_fused_shared_experts == 1
+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.dp_attention import is_dp_attention_enabled
from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.logits_processor import LogitsProcessor 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.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import ( from sglang.srt.layers.vocab_parallel_embedding import (
ParallelLMHead, 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.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.models.glm4 import Glm4DecoderLayer from sglang.srt.models.glm4 import Glm4DecoderLayer
from sglang.srt.models.glm_ocr import GlmOcrForConditionalGeneration 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 from sglang.srt.utils import add_prefix
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
@@ -138,9 +139,7 @@ class GlmOcrForConditionalGenerationNextN(GlmOcrForConditionalGeneration):
) )
self.logits_processor = LogitsProcessor(config) self.logits_processor = LogitsProcessor(config)
self.num_fused_shared_experts = ( self.num_fused_shared_experts = 0 if is_shared_experts_fusion_disabled() else 1
0 if get_exec().moe.disable_shared_experts_fusion else 1
)
@torch.no_grad() @torch.no_grad()
def forward( def forward(
+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__( def __init__(
self, self,
config: KimiK25Config, config: KimiK25Config,
+15 -3
View File
@@ -114,7 +114,21 @@ class KimiVLMultiModalProjector(nn.Module):
return hidden_states 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): 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__( def __init__(
self, self,
config: KimiVLConfig, config: KimiVLConfig,
@@ -138,10 +152,8 @@ class KimiVLForConditionalGeneration(nn.Module):
self.language_model = None self.language_model = None
if not config.encoder_only: if not config.encoder_only:
text_config = copy.deepcopy(config.text_config)
text_config.architectures = ["DeepseekV2ForCausalLM"]
self.language_model = DeepseekV2ForCausalLM( self.language_model = DeepseekV2ForCausalLM(
config=text_config, config=_language_model_config(config),
quant_config=quant_config, quant_config=quant_config,
prefix=add_prefix("language_model", prefix), prefix=add_prefix("language_model", prefix),
) )
+8
View File
@@ -1640,6 +1640,14 @@ class MiniCPMV:
minicpmv: nn.Module 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__( def __init__(
self, self,
config: PretrainedConfig, config: PretrainedConfig,
+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.ep_moe.layer import get_moe_impl_class
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
from sglang.srt.layers.moe.topk import TopK 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.quantization.base_config import QuantizationConfig
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.layers.rotary_embedding import get_rope 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.tp_size = get_parallel().tp_size
self.n_shared_experts = getattr(config, "n_shared_experts", None) self.n_shared_experts = getattr(config, "n_shared_experts", None)
self.num_fused_shared_experts = ( self.num_fused_shared_experts = (
0 0 if is_shared_experts_fusion_disabled() else config.n_shared_experts
if get_exec().moe.disable_shared_experts_fusion
else config.n_shared_experts
) )
if self.tp_size > config.num_local_experts: if self.tp_size > config.num_local_experts:
@@ -1465,43 +1466,31 @@ class MiniMaxM3SparseForCausalLM(nn.Module):
def get_input_embeddings(self): def get_input_embeddings(self):
return self.model.get_input_embeddings() return self.model.get_input_embeddings()
def determine_num_fused_shared_experts(self): @classmethod
if get_exec().moe.disable_shared_experts_fusion: def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
return """Why this checkpoint cannot fuse its shared expert, or None. Asked by
the loader before any layer is built."""
disable_reason = None if not getattr(hf_config, "n_shared_experts", None):
if not getattr(self.config, "n_shared_experts", None): return "No shared experts are defined in the config."
disable_reason = "No shared experts are defined in the config." if quant_config is not None and quant_config.get_name() == "modelopt_mixed":
elif ( return (
self.quant_config is not None
and self.quant_config.get_name() == "modelopt_mixed"
):
disable_reason = (
"Shared and routed experts may use different quantization formats " "Shared and routed experts may use different quantization formats "
"in ModelOpt mixed-precision checkpoints." "in ModelOpt mixed-precision checkpoints."
) )
elif not _is_cuda: if not _is_cuda:
disable_reason = "Shared experts fusion currently requires CUDA devices." return "Shared experts fusion currently requires CUDA devices."
elif _is_cuda and (_device_sm is not None) and (_device_sm < 80): if _is_cuda and (_device_sm is not None) and (_device_sm < 80):
disable_reason = "Shared experts fusion requires SM80 or newer GPUs." return "Shared experts fusion requires SM80 or newer GPUs."
elif get_parallel().moe_ep_size > 1: if get_parallel().moe_ep_size > 1:
disable_reason = "Shared experts fusion is not supported together with expert parallelism yet." return "Shared experts fusion is not supported together with expert parallelism yet."
elif get_moe_a2a_backend().is_deepep(): if get_moe_a2a_backend().is_deepep():
disable_reason = "Shared experts fusion is not supported when Deepep MoE backend is enabled." return "Shared experts fusion is not supported when Deepep MoE backend is enabled."
return None
if disable_reason is not None: def determine_num_fused_shared_experts(self):
from sglang.srt.arg_groups.overrides import declare_load_time_override # The decision was installed by the loader; this only reads it.
if is_shared_experts_fusion_disabled():
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.",
)
return return
self.num_fused_shared_experts = self.config.n_shared_experts self.num_fused_shared_experts = self.config.n_shared_experts
assert ( assert (
self.num_fused_shared_experts == 1 self.num_fused_shared_experts == 1
+27 -35
View File
@@ -10,7 +10,10 @@ from sglang.srt.distributed import (
get_pp_group, get_pp_group,
) )
from sglang.srt.layers.logits_processor import LogitsProcessor 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.quantization.base_config import QuantizationConfig
from sglang.srt.layers.utils import PPMissingLayer from sglang.srt.layers.utils import PPMissingLayer
from sglang.srt.layers.utils.common import get_layer_id 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, merge_vit_qkv_weights,
) )
from sglang.srt.models.utils import WeightsMapper 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 import add_prefix, get_device_sm, is_cuda, log_info_on_rank0
from sglang.srt.utils.hf_transformers_utils import get_rope_config 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) self.logits_processor = LogitsProcessor(text_config)
def _determine_num_fused_shared_experts(self) -> None: @classmethod
text_config = self.config.text_config def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
if get_exec().moe.disable_shared_experts_fusion: """Why this checkpoint cannot fuse its shared expert, or None. Asked by
return the loader before any layer is built; the experts live on the text
config."""
disable_reason = None text_config = getattr(hf_config, "text_config", hf_config)
if not getattr(text_config, "n_shared_experts", None): if not getattr(text_config, "n_shared_experts", None):
disable_reason = "No shared experts are defined in the config." return "No shared experts are defined in the config."
elif ( if quant_config is not None and quant_config.get_name() == "modelopt_mixed":
self.quant_config is not None return (
and self.quant_config.get_name() == "modelopt_mixed"
):
disable_reason = (
"Shared and routed experts may use different quantization formats " "Shared and routed experts may use different quantization formats "
"in ModelOpt mixed-precision checkpoints." "in ModelOpt mixed-precision checkpoints."
) )
elif not _is_cuda: if not _is_cuda:
disable_reason = "Shared experts fusion currently requires CUDA devices." return "Shared experts fusion currently requires CUDA devices."
elif (_device_sm is not None) and (_device_sm < 80): if (_device_sm is not None) and (_device_sm < 80):
disable_reason = "Shared experts fusion requires SM80 or newer GPUs." return "Shared experts fusion requires SM80 or newer GPUs."
elif get_parallel().moe_ep_size > 1: if get_parallel().moe_ep_size > 1:
disable_reason = ( return (
"Shared experts fusion is not supported together with expert " "Shared experts fusion is not supported together with expert "
"parallelism yet." "parallelism yet."
) )
elif get_moe_a2a_backend().is_deepep(): if get_moe_a2a_backend().is_deepep():
disable_reason = ( return (
"Shared experts fusion is not supported when Deepep MoE backend " "Shared experts fusion is not supported when Deepep MoE backend "
"is enabled." "is enabled."
) )
return None
if disable_reason is not None: def _determine_num_fused_shared_experts(self) -> None:
from sglang.srt.arg_groups.overrides import declare_load_time_override # The decision was installed by the loader; this only reads it.
if is_shared_experts_fusion_disabled():
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.",
)
return return
self.num_fused_shared_experts = self.config.text_config.n_shared_experts
self.num_fused_shared_experts = text_config.n_shared_experts
assert ( assert (
self.num_fused_shared_experts == 1 self.num_fused_shared_experts == 1
), "Only 1 fused shared expert is supported" ), "Only 1 fused shared expert is supported"
+10
View File
@@ -73,6 +73,16 @@ class VisionEncoderArgs:
class PixtralForConditionalGeneration(nn.Module): class PixtralForConditionalGeneration(nn.Module):
merge_by_field_config = True 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 @classmethod
def get_placeholder_str(cls, modality: str, i: int) -> str | None: def get_placeholder_str(cls, modality: str, i: int) -> str | None:
if modality.startswith("image"): if modality.startswith("image"):
+41 -25
View File
@@ -57,6 +57,9 @@ from sglang.srt.layers.linear import (
RowParallelLinear, RowParallelLinear,
) )
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE 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 ( from sglang.srt.layers.parameter import (
BlockQuantScaleParameter, BlockQuantScaleParameter,
PerTensorScaleParameter, PerTensorScaleParameter,
@@ -137,9 +140,10 @@ cached_get_processor = lru_cache(get_processor)
def _disable_shared_experts_fusion() -> bool: def _disable_shared_experts_fusion() -> bool:
# Resolved lazily: the global server args is not set at module import time # Resolved lazily: the flag is written by the owning model's gate before
# (e.g. when this module is imported by unit tests). # its layers build (per runner); models without a gate see the config
return get_exec().moe.disable_shared_experts_fusion # intent through the accessor's fallback.
return is_shared_experts_fusion_disabled()
if _is_cuda: if _is_cuda:
@@ -1308,25 +1312,6 @@ class Qwen3_5ForCausalLM(nn.Module):
f"get_hidden_dim not implemented for {module_name}" 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__( def __init__(
self, self,
config: Qwen3_5TextConfig, config: Qwen3_5TextConfig,
@@ -1339,9 +1324,6 @@ class Qwen3_5ForCausalLM(nn.Module):
self.hidden_size = config.hidden_size self.hidden_size = config.hidden_size
self.pp_group = get_pp_group() 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 alt_stream = get_stream("alt") if _is_cuda or _hip_use_alt_stream else None
# Embedding layer # 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] EntryClass = [Qwen3_5MoeForConditionalGeneration, Qwen3_5ForConditionalGeneration]
+42 -25
View File
@@ -44,8 +44,49 @@ from sglang.srt.utils import add_prefix, is_npu
logger = logging.getLogger(__name__) 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): 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__( def __init__(
self, self,
config: PretrainedConfig, 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. # Deep-copy so MTP mutations below don't leak into the target's config.
config = copy.deepcopy(config) config = copy.deepcopy(config)
# Serialized Qwen3.5 ModelOpt checkpoints keep embedded MTP weights in quant_config = _mtp_quant_config(quant_config)
# 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
self.config = config self.config = config
self.tp_size = get_parallel().tp_size self.tp_size = get_parallel().tp_size
+7
View File
@@ -43,6 +43,13 @@ class Qwen3_5ForCausalLM(nn.Module):
packed_modules_mapping = qwen3_5.Qwen3_5ForCausalLM.packed_modules_mapping packed_modules_mapping = qwen3_5.Qwen3_5ForCausalLM.packed_modules_mapping
supported_lora_modules = qwen3_5.Qwen3_5ForCausalLM.supported_lora_modules 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__( def __init__(
self, self,
config, config,
+14
View File
@@ -376,6 +376,20 @@ class MoeFlags(_FlagGroupBase):
tbo_token_distribution_threshold: float | None = None tbo_token_distribution_threshold: float | None = None
disable_fp4_allgather: bool | None = None disable_fp4_allgather: bool | None = None
quantization: str | 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 @dataclasses.dataclass
@@ -80,6 +80,14 @@ def build_draft_tp_worker(
server_args=server_args, algo_label=algo_label server_args=server_args, algo_label=algo_label
) )
) )
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( draft_worker = TpModelWorker(
server_args=server_args, server_args=server_args,
gpu_id=gpu_id, gpu_id=gpu_id,
@@ -26,6 +26,7 @@ from sglang.srt.layers.attention.trtllm_mla_backend import (
TRTLLMMLABackend, TRTLLMMLABackend,
) )
from sglang.srt.layers.moe.utils import ( from sglang.srt.layers.moe.utils import (
draft_model_build_scope,
speculative_moe_a2a_backend_context, speculative_moe_a2a_backend_context,
speculative_moe_backend_context, speculative_moe_backend_context,
) )
@@ -162,7 +163,7 @@ class EagleDraftWorker(EagleDraftWorkerBase):
ctx = empty_context() ctx = empty_context()
with ( with (
ctx 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( self.draft_worker = TpModelWorker(
server_args=server_args, server_args=server_args,
gpu_id=gpu_id, gpu_id=gpu_id,
@@ -30,6 +30,7 @@ import torch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.distributed.parallel_state_wrapper import ParallelState
from sglang.srt.layers.moe.utils import ( from sglang.srt.layers.moe.utils import (
draft_model_build_scope,
speculative_moe_a2a_backend_context, speculative_moe_a2a_backend_context,
speculative_moe_backend_context, speculative_moe_backend_context,
) )
@@ -127,7 +128,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker):
with ( with (
empty_context() 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 # Both base classes own initialization, so initialize TpModelWorker
# explicitly after EagleDraftWorkerBase above. # explicitly after EagleDraftWorkerBase above.
TpModelWorker.__init__( TpModelWorker.__init__(
@@ -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 ( from sglang.srt.hardware_backend.npu.graph_runner.multi_layer_eagle_draft_extend_npu_graph_runner import (
MultiLayerEagleMultiStepDraftExtendNpuGraphRunner, 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.schedule_batch import ScheduleBatch
from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.scheduler import GenerationBatchResult
from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.managers.tp_worker import TpModelWorker
@@ -150,7 +153,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase):
) )
# Load draft model weights only. # 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( self.draft_worker = TpModelWorker(
server_args=server_args, server_args=server_args,
gpu_id=gpu_id, gpu_id=gpu_id,
@@ -5,7 +5,10 @@ from typing import Optional
import torch import torch
from sglang.srt.distributed.parallel_state_wrapper import ParallelState 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.managers.tp_worker import TpModelWorker
from sglang.srt.server_args import ServerArgs from sglang.srt.server_args import ServerArgs
from sglang.srt.speculative.adaptive_runtime_state import ( 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 self.speculative_num_steps * self.topk, self.speculative_num_draft_tokens
) )
# Load draft model weights only. # Load draft model weights only. The standalone draft is a real model
with empty_context(): # 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( self.draft_worker = TpModelWorker(
server_args=server_args, server_args=server_args,
gpu_id=gpu_id, gpu_id=gpu_id,
@@ -41,12 +41,9 @@ class RoutedExpertsCapturer(BaseTopkCapturer):
) -> Optional["RoutedExpertsCapturer"]: ) -> Optional["RoutedExpertsCapturer"]:
if not get_exec().features.enable_return_routed_experts: if not get_exec().features.enable_return_routed_experts:
return None return None
if not get_exec().moe.disable_shared_experts_fusion and hasattr( # The model's own attribute is the baked decision (0 when its gate
model, "num_fused_shared_experts" # disabled fusion); the ACTIVE flag can be holding another runner's.
): num_fused_shared_experts = getattr(model, "num_fused_shared_experts", 0)
num_fused_shared_experts = model.num_fused_shared_experts
else:
num_fused_shared_experts = 0
return RoutedExpertsCapturer( return RoutedExpertsCapturer(
model_config, model_config,
num_tokens=num_tokens, num_tokens=num_tokens,
@@ -1,22 +1,24 @@
import unittest import unittest
from types import SimpleNamespace from types import SimpleNamespace
from sglang.srt.layers.moe.utils import (
install_shared_experts_fusion_decision,
is_shared_experts_fusion_disabled,
)
from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM
from sglang.srt.runtime_context import get_context, get_exec from sglang.srt.runtime_context import get_context, get_exec, get_flags
from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=4, suite="base-a-test-cpu") register_cpu_ci(est_time=4, suite="base-a-test-cpu")
class TestDeepseekV4SharedExpertFusionPolicy(unittest.TestCase): class TestDeepseekV4SharedExpertFusionPolicy(unittest.TestCase):
"""The disable decision is a load-time resolution: it lands on the """V4 fuses its shared expert only when explicitly asked to.
published config bag via declare_load_time_override (bag-only; the
ServerArgs instance stays pristine)."""
def _make_model(self, n_shared_experts=1): The gate is a question the loader asks the model class before any layer
return SimpleNamespace( exists (``shared_experts_fusion_disable_reason``); the answer is installed
config=SimpleNamespace(n_shared_experts=n_shared_experts) on the ACTIVE moe flag, and the config bag keeps the user's intent.
) """
def _publish(self, enforce): def _publish(self, enforce):
override = get_context().override_server_args( override = get_context().override_server_args(
@@ -24,25 +26,47 @@ class TestDeepseekV4SharedExpertFusionPolicy(unittest.TestCase):
) )
override.install() override.install()
self.addCleanup(override.restore) self.addCleanup(override.restore)
get_flags().moe.disable_shared_experts_fusion = None
self.addCleanup(
lambda: setattr(get_flags().moe, "disable_shared_experts_fusion", None)
)
def _install(self, n_shared_experts=1):
install_shared_experts_fusion_decision(
DeepseekV4ForCausalLM,
SimpleNamespace(n_shared_experts=n_shared_experts),
None,
)
def test_disables_shared_fusion_without_enforce(self): def test_disables_shared_fusion_without_enforce(self):
self._publish(enforce=False) self._publish(enforce=False)
model = self._make_model() self.assertEqual(
DeepseekV4ForCausalLM.shared_experts_fusion_disable_reason(
DeepseekV4ForCausalLM.determine_num_fused_shared_experts(model) SimpleNamespace(n_shared_experts=1), None
),
self.assertEqual(model.num_fused_shared_experts, 0) "Config does not support fused shared expert(s).",
# post-init declaration lands on the published config bag )
self.assertTrue(get_exec().moe.disable_shared_experts_fusion) self._install()
# The decision lands on the ACTIVE flag; the config intent is untouched.
self.assertTrue(is_shared_experts_fusion_disabled())
self.assertFalse(get_exec().moe.disable_shared_experts_fusion)
def test_enables_shared_fusion_when_enforced(self): def test_enables_shared_fusion_when_enforced(self):
self._publish(enforce=True) self._publish(enforce=True)
model = self._make_model() self.assertIsNone(
DeepseekV4ForCausalLM.shared_experts_fusion_disable_reason(
SimpleNamespace(n_shared_experts=1), None
)
)
self._install()
self.assertFalse(is_shared_experts_fusion_disabled())
DeepseekV4ForCausalLM.determine_num_fused_shared_experts(model) def test_enforcing_with_more_than_one_shared_expert_is_rejected(self):
self._publish(enforce=True)
self.assertEqual(model.num_fused_shared_experts, 1) with self.assertRaisesRegex(ValueError, "exactly one shared"):
self.assertFalse(get_exec().moe.disable_shared_experts_fusion) DeepseekV4ForCausalLM.shared_experts_fusion_disable_reason(
SimpleNamespace(n_shared_experts=2), None
)
if __name__ == "__main__": if __name__ == "__main__":
@@ -0,0 +1,181 @@
"""Every loader entry class that can reach a fusion-gated family answers for it.
The loader installs the shared-experts-fusion decision for the class it
instantiates (`install_shared_experts_fusion_decision`). A model whose layers
read `is_shared_experts_fusion_disabled()` therefore gets whatever answer that
*entry* class produced — and an entry class with no
`shared_experts_fusion_disable_reason` falls back to the user's intent, silently
skipping the family's auto-disable conditions.
That is easy to miss for a wrapper: `KimiVLForConditionalGeneration` is the
registered arch, but a DeepSeek body is built inside it, and the DeepSeek
conditions used to be evaluated during that nested construction. This case walks
the registry so a new wrapper (or a new MTP/nextn entry) cannot reintroduce the
gap.
"""
import ast
import importlib
import inspect
import os
import sys
import unittest
from sglang.srt.models.registry import ModelRegistry
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=90, suite="base-a-test-cpu")
GATE = "shared_experts_fusion_disable_reason"
FLAG_READERS = (
"is_shared_experts_fusion_disabled",
"determine_num_fused_shared_experts",
)
# Archs that read the fusion flag but deliberately have no gate: nothing in
# their lineage carries auto-disable conditions, so they follow the user's
# intent — the behavior they had before the decision moved to the loader.
GATELESS_BY_DESIGN = {
# The in-tree class has no ``determine_num_fused_shared_experts`` at all
# (the call is guarded by ``hasattr`` for a downstream variant).
"BailingMoeForCausalLMNextN",
# Its target family (Glm4v) is dense; there is no gate to inherit.
"GlmOcrForConditionalGenerationNextN",
# The vision tower registered on its own: it shares a module with
# PixtralForConditionalGeneration (which does answer) but builds no language
# model, so there is nothing for a gate to decide.
"PixtralVisionModel",
}
def _gated_classes(source: str) -> set:
"""Classes in this module that define or receive a fusion gate."""
names = set()
tree = ast.parse(source)
for node in ast.walk(tree):
if isinstance(node, ast.ClassDef):
for body in node.body:
if (
isinstance(body, (ast.FunctionDef, ast.AsyncFunctionDef))
and body.name == GATE
):
names.add(node.name)
if isinstance(body, ast.Assign) and any(
isinstance(t, ast.Name) and t.id == GATE for t in body.targets
):
names.add(node.name)
# ``for cls in (A, B): cls.<GATE> = ...``
if (
isinstance(node, ast.For)
and isinstance(node.iter, (ast.Tuple, ast.List))
and GATE in ast.dump(node)
):
names |= {e.id for e in node.iter.elts if isinstance(e, ast.Name)}
if isinstance(node, ast.Assign):
for target in node.targets:
if (
isinstance(target, ast.Attribute)
and target.attr == GATE
and isinstance(target.value, ast.Name)
):
names.add(target.value.id)
return names
def gated_class_names() -> set:
"""Every model class that *resolves* a gate, inherited ones included.
A subclass like `DeepseekV3ForCausalLM` inherits the gate without naming it,
so collecting names from class bodies alone would let a wrapper that builds
the subclass slip through.
"""
names = set()
for module_name, module in list(sys.modules.items()):
if not module_name.startswith("sglang.srt.models.") or module is None:
continue
for member in vars(module).values():
# transformers re-exports lazy placeholders that raise on any
# attribute access when their optional backend is missing.
try:
if (
inspect.isclass(member)
and (member.__module__ or "").startswith("sglang.srt.models.")
and hasattr(member, GATE)
):
names.add(member.__name__)
except Exception:
continue
return names
class TestFusionGateCoverage(CustomTestCase):
def test_every_entry_class_reaching_a_gated_family_has_a_gate(self):
models_dir = list(importlib.import_module("sglang.srt.models").__path__)[0]
gates_by_module = {}
for name in sorted(os.listdir(models_dir)):
if not name.endswith(".py"):
continue
with open(os.path.join(models_dir, name), encoding="utf-8") as f:
try:
gates_by_module[f"sglang.srt.models.{name[:-3]}"] = _gated_classes(
f.read()
)
except SyntaxError:
continue
missing = []
all_gated = None
for arch in sorted(ModelRegistry.get_supported_archs()):
try:
model_class, _ = ModelRegistry.resolve_model_cls(arch)
except Exception:
continue
if hasattr(model_class, GATE) or arch in GATELESS_BY_DESIGN:
continue
module = importlib.import_module(model_class.__module__)
try:
source = inspect.getsource(module)
except OSError:
continue
reasons = []
if any(reader in source for reader in FLAG_READERS):
reasons.append("reads the fusion flag")
try:
tree = ast.parse(source)
except SyntaxError:
tree = None
if tree is not None:
# Any *use* of a gated class counts, whatever the shape: a
# direct call (`DeepseekV2ForCausalLM(...)`), a module attribute
# (`qwen3_5.Qwen3_5MoeForCausalLM`), or a class attribute the
# constructor later calls (`body_cls = qwen3_5.Qwen3_5...`).
# Only matching calls would miss the last two.
if all_gated is None:
all_gated = gated_class_names()
used = set()
for node in ast.walk(tree):
if isinstance(node, ast.Name) and node.id in all_gated:
used.add(node.id)
elif isinstance(node, ast.Attribute) and node.attr in all_gated:
used.add(node.attr)
for name in sorted(used):
if name != model_class.__name__:
reasons.append(f"references {name}")
if reasons:
missing.append(
f"{arch} ({model_class.__module__}): {', '.join(reasons)}"
)
self.assertEqual(
[],
missing,
"these entry classes reach a fusion-gated family but resolve no "
f"{GATE}, so the loader falls back to the user's intent for them and "
"the family's auto-disable conditions never run:\n "
+ "\n ".join(missing),
)
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,514 @@
"""Every MoE family's fusion gate, asked the way the loader asks it.
`install_shared_experts_fusion_decision` calls
`<model class>.shared_experts_fusion_disable_reason(hf_config, quant_config)`
before the model is built, so the gate must answer from the config and
quantization it is handed — no instance, no layers. These cases pin each
family's branch table, which matters because most of these checkpoints cannot
be run on a single dev box: a wrong answer here is a silently wrong weight
remap (the loader remaps `mlp.shared_experts` into a fused slot the layers
never allocated), not a crash.
Conditions that depend on the device or the parallel topology are exercised
through `get_parallel().override(...)`; the ones that are pure config /
quantization are exercised directly.
"""
import unittest
import unittest.mock
from types import SimpleNamespace
from sglang.srt.runtime_context import get_context, get_parallel
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=6, suite="base-a-test-cpu")
def _quant(name: str):
return SimpleNamespace(get_name=lambda: name)
class _FusionGateCase(CustomTestCase):
def _seed(self, **fields):
override = get_context().override_server_args(**fields)
override.install()
self.addCleanup(override.restore)
def _reason(self, model_class, hf_config, quant_config=None, moe_ep_size=1):
# The gates consult the live EP size; without a group installed the
# canonical getter asserts, so every case states a topology.
with get_parallel().override(moe_ep_size=moe_ep_size):
return model_class.shared_experts_fusion_disable_reason(
hf_config, quant_config
)
class TestDeepseekV2Gate(_FusionGateCase):
def _config(self, **kw):
base = dict(
architectures=["DeepseekV3ForCausalLM"],
n_routed_experts=256,
n_shared_experts=1,
)
base.update(kw)
return SimpleNamespace(**base)
def test_a_foreign_architecture_cannot_fuse(self):
from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
self._seed()
self.assertIn(
"does not support",
self._reason(
DeepseekV2ForCausalLM,
self._config(architectures=["SomeOtherForCausalLM"]),
),
)
def test_an_unvalidated_expert_count_cannot_fuse(self):
from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
self._seed()
self.assertIn(
"does not support",
self._reason(DeepseekV2ForCausalLM, self._config(n_routed_experts=128)),
)
def test_the_384_expert_layout_needs_a_quark_checkpoint(self):
from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
self._seed()
config = self._config(n_routed_experts=384)
self.assertIn(
"does not support",
self._reason(DeepseekV2ForCausalLM, config, _quant("compressed-tensors")),
)
# With Quark the layout is pre-fused, so this branch stops objecting.
self.assertNotIn(
"does not support",
self._reason(DeepseekV2ForCausalLM, config, _quant("quark")) or "",
)
def test_the_nextn_draft_declares_its_own_architecture(self):
from sglang.srt.models.deepseek_nextn import DeepseekV3ForCausalLMNextN
from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
self.assertEqual(
DeepseekV3ForCausalLMNextN.fused_shared_experts_architecture,
"DeepseekV3ForCausalLMNextN",
)
self._seed()
draft_config = self._config(architectures=["DeepseekV3ForCausalLMNextN"])
# The draft's own class accepts it; the target's class does not.
self.assertNotIn(
"does not support",
self._reason(DeepseekV3ForCausalLMNextN, draft_config) or "",
)
self.assertIn(
"does not support", self._reason(DeepseekV2ForCausalLM, draft_config)
)
def test_expert_parallelism_blocks_fusion_off_rocm(self):
from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
self._seed()
self.assertTrue(
self._reason(DeepseekV2ForCausalLM, self._config(), moe_ep_size=2)
)
class TestGlmMoeLiteGate(_FusionGateCase):
def _config(self, **kw):
base = dict(architectures=["Glm4MoeLiteForCausalLM"], n_shared_experts=1)
base.update(kw)
return SimpleNamespace(**base)
def test_more_than_one_shared_expert_cannot_fuse(self):
from sglang.srt.models.glm4_moe_lite import Glm4MoeLiteForCausalLM
self._seed()
self.assertTrue(
self._reason(Glm4MoeLiteForCausalLM, self._config(n_shared_experts=2))
)
def test_expert_parallelism_blocks_fusion(self):
from sglang.srt.models.glm4_moe_lite import Glm4MoeLiteForCausalLM
self._seed()
config = self._config()
reason = self._reason(Glm4MoeLiteForCausalLM, config, moe_ep_size=2)
self.assertTrue(reason)
# This family checks the device capability before expert parallelism, so
# only ask *which* branch refused on a device that would otherwise fuse
# (a CPU runner never gets past the capability check).
if self._reason(Glm4MoeLiteForCausalLM, config) is None:
self.assertIn("expert parallelism", reason)
def test_the_nextn_draft_declares_its_own_architecture(self):
from sglang.srt.models.glm4_moe_lite_nextn import Glm4MoeLiteForCausalLMNextN
self.assertEqual(
Glm4MoeLiteForCausalLMNextN.fused_shared_experts_architecture,
"Glm4MoeLiteForCausalLMNextN",
)
class TestGlmMoeGate(_FusionGateCase):
def test_a_w4afp8_checkpoint_cannot_fuse(self):
from sglang.srt.models.glm4_moe import Glm4MoeForCausalLM
self._seed()
reason = self._reason(
Glm4MoeForCausalLM, SimpleNamespace(n_shared_experts=1), _quant("w4afp8")
)
self.assertTrue(reason)
def test_the_dsa_variant_declares_its_own_architecture(self):
from sglang.srt.models.glm4_moe import GlmMoeDsaForCausalLM
self.assertEqual(
GlmMoeDsaForCausalLM.fused_shared_experts_architecture,
"GlmMoeDsaForCausalLM",
)
class TestMiniMaxGates(_FusionGateCase):
def test_a_config_without_shared_experts_cannot_fuse(self):
from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM
self._seed()
self.assertIn(
"No shared experts",
self._reason(
MiniMaxM3SparseForCausalLM, SimpleNamespace(n_shared_experts=0)
),
)
def test_a_modelopt_mixed_checkpoint_cannot_fuse(self):
from sglang.srt.models.minimax_m3 import MiniMaxM3SparseForCausalLM
self._seed()
reason = self._reason(
MiniMaxM3SparseForCausalLM,
SimpleNamespace(n_shared_experts=1),
_quant("modelopt_mixed"),
)
self.assertIn("quantization formats", reason)
def test_the_vl_variant_reads_the_text_config(self):
from sglang.srt.models.minimax_m3_vl import (
MiniMaxM3SparseForConditionalGeneration,
)
self._seed()
wrapper = SimpleNamespace(text_config=SimpleNamespace(n_shared_experts=0))
self.assertIn(
"No shared experts",
self._reason(MiniMaxM3SparseForConditionalGeneration, wrapper),
)
class TestQwen3_5Gate(_FusionGateCase):
def test_every_entry_class_answers(self):
import sglang.srt.models.qwen3_5 as qwen3_5
for cls in (
qwen3_5.Qwen3_5ForCausalLM,
qwen3_5.Qwen3_5MoeForCausalLM,
qwen3_5.Qwen3_5ForConditionalGeneration,
qwen3_5.Qwen3_5MoeForConditionalGeneration,
):
self.assertTrue(
hasattr(cls, "shared_experts_fusion_disable_reason"),
f"{cls.__name__} would silently skip the ROCm auto-disable",
)
def test_the_auto_disable_is_rocm_only(self):
import sglang.srt.models.qwen3_5 as qwen3_5
self._seed()
# On a non-ROCm build the gate never objects, whatever the checkpoint is.
wrapper = SimpleNamespace(
text_config=SimpleNamespace(model_type="qwen3_5_moe_text")
)
if not qwen3_5._is_hip:
self.assertIsNone(
self._reason(qwen3_5.Qwen3_5MoeForConditionalGeneration, wrapper)
)
class TestWrapperEntryClassGates(_FusionGateCase):
"""A wrapper model answers with the config it hands its nested family.
The loader asks the class it instantiates, which for these models is the
wrapper — not the DeepSeek/Qwen3.5 body inside it. Each wrapper therefore
delegates to its family's gate with the config (and quantization) the
nested construction uses; these cases pin *what gets handed over*, because
handing over the top-level config instead would answer for the wrong
checkpoint (or raise on a config that has no expert counts at all).
"""
def _recording_gate(self, family_cls):
seen = {}
def recorder(hf_config, quant_config):
seen["config"] = hf_config
seen["quant"] = quant_config
return None
return seen, unittest.mock.patch.object(
family_cls,
"shared_experts_fusion_disable_reason",
staticmethod(recorder),
)
def test_kimi_vl_never_fuses_and_says_why(self):
from sglang.srt.models.kimi_vl import KimiVLForConditionalGeneration
self._seed()
config = SimpleNamespace(
encoder_only=False,
text_config=SimpleNamespace(
architectures=["Whatever"], n_routed_experts=256, n_shared_experts=1
),
)
# The construction rewrites the architecture to DeepseekV2ForCausalLM,
# which is not the architecture the fused path validated.
self.assertIn(
"does not support",
self._reason(KimiVLForConditionalGeneration, config),
)
self.assertIsNone(
self._reason(
KimiVLForConditionalGeneration,
SimpleNamespace(encoder_only=True, text_config=None),
),
"an encoder-only Kimi-VL builds no language model",
)
def test_kimi_k25_hands_over_its_text_config(self):
from sglang.srt.models.deepseek_v2 import DeepseekV3ForCausalLM
from sglang.srt.models.kimi_k25 import KimiK25ForConditionalGeneration
self._seed()
text_config = SimpleNamespace(
architectures=["DeepseekV3ForCausalLM"],
n_routed_experts=384,
n_shared_experts=1,
)
config = SimpleNamespace(encoder_only=False, text_config=text_config)
# The standard compressed-tensors Kimi-K2.5 checkpoint stores its shared
# expert loose, so this must refuse to fuse.
self.assertIn(
"does not support",
self._reason(
KimiK25ForConditionalGeneration,
config,
_quant("compressed-tensors"),
),
)
seen, patcher = self._recording_gate(DeepseekV3ForCausalLM)
with patcher:
self._reason(KimiK25ForConditionalGeneration, config, _quant("quark"))
self.assertIs(seen["config"], text_config)
self.assertIsNone(
self._reason(
KimiK25ForConditionalGeneration,
SimpleNamespace(encoder_only=True, text_config=None),
)
)
def test_pixtral_only_asks_for_its_mla_backbone(self):
from sglang.srt.models.mistral_large_3 import MistralLarge3ForCausalLM
from sglang.srt.models.pixtral import PixtralForConditionalGeneration
self._seed()
mla_text = SimpleNamespace(
model_type="deepseek_v3",
architectures=["DeepseekV3ForCausalLM"],
n_routed_experts=256,
n_shared_experts=1,
)
seen, patcher = self._recording_gate(MistralLarge3ForCausalLM)
with patcher:
self._reason(
PixtralForConditionalGeneration,
SimpleNamespace(text_config=mla_text),
)
self.assertIs(seen["config"], mla_text)
# A GQA text config builds the dense Mistral backbone instead.
self.assertIsNone(
self._reason(
PixtralForConditionalGeneration,
SimpleNamespace(text_config=SimpleNamespace(model_type="mistral")),
)
)
def test_dots_vlm_hands_over_the_language_config(self):
from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
from sglang.srt.models.dots_vlm import DotsVLMForCausalLM
language_config = SimpleNamespace(
architectures=["DeepseekV3ForCausalLM"],
n_routed_experts=256,
n_shared_experts=1,
)
config = SimpleNamespace(encoder_only=False, language_config=language_config)
seen, patcher = self._recording_gate(DeepseekV2ForCausalLM)
with patcher:
self._reason(DotsVLMForCausalLM, config, _quant("fp8"))
self.assertIs(seen["config"], language_config)
self.assertEqual(seen["quant"].get_name(), "fp8")
def test_deepseek_vl2_mirrors_its_unquantized_language_model(self):
from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
from sglang.srt.models.deepseek_vl2 import DeepseekVL2ForCausalLM
language_config = SimpleNamespace(
use_mla=True,
architectures=["DeepseekV3ForCausalLM"],
n_routed_experts=256,
n_shared_experts=1,
)
seen, patcher = self._recording_gate(DeepseekV2ForCausalLM)
with patcher:
self._reason(
DeepseekVL2ForCausalLM,
SimpleNamespace(language_config=language_config),
_quant("fp8"),
)
self.assertIs(seen["config"], language_config)
self.assertIsNone(
seen["quant"], "the language model is constructed without quantization"
)
# deepseek-vl2-tiny forbids MLA and builds the dense model instead.
self.assertIsNone(
self._reason(
DeepseekVL2ForCausalLM,
SimpleNamespace(language_config=SimpleNamespace(use_mla=False)),
)
)
def test_deepseek_ocr_only_asks_for_its_moe_branches(self):
from sglang.srt.models.deepseek_ocr import DeepseekOCRForCausalLM
from sglang.srt.models.deepseek_v2 import DeepseekV2ForCausalLM
text_config = SimpleNamespace(
topk_method="noaux_tc",
use_mla=True,
architectures=["DeepseekV3ForCausalLM"],
n_routed_experts=256,
n_shared_experts=1,
)
moe_config = SimpleNamespace(
vision_config=SimpleNamespace(model_name="deepencoder"),
projector_config=SimpleNamespace(input_dim=1280),
text_config=text_config,
)
seen, patcher = self._recording_gate(DeepseekV2ForCausalLM)
with patcher:
self._reason(DeepseekOCRForCausalLM, moe_config, _quant("fp8"))
self.assertIs(seen["config"], text_config)
# OCR2 (and any non-MLA, non-noaux_tc config) builds the dense model.
ocr2 = SimpleNamespace(
vision_config=SimpleNamespace(model_name="DeepEncoderV2"),
projector_config=SimpleNamespace(input_dim=896),
text_config=text_config,
)
self.assertIsNone(self._reason(DeepseekOCRForCausalLM, ocr2))
dense = SimpleNamespace(
vision_config=SimpleNamespace(model_name="deepencoder"),
projector_config=SimpleNamespace(input_dim=1280),
text_config=SimpleNamespace(topk_method="greedy", use_mla=False),
)
self.assertIsNone(self._reason(DeepseekOCRForCausalLM, dense))
def test_minicpmv_entries_delegate_to_the_qwen3_5_gate(self):
from sglang.srt.models.minicpmv import (
MiniCPMV,
MiniCPMV4_6ForConditionalGeneration,
)
from sglang.srt.models.qwen3_5 import Qwen3_5ForCausalLM
text_config = SimpleNamespace(model_type="qwen3_5_moe_text")
for cls in (MiniCPMV, MiniCPMV4_6ForConditionalGeneration):
seen, patcher = self._recording_gate(Qwen3_5ForCausalLM)
with patcher:
self._reason(cls, SimpleNamespace(text_config=text_config))
self.assertIs(seen["config"], text_config, cls.__name__)
def test_the_text_only_qwen3_5_entries_delegate_to_their_body(self):
import sglang.srt.models.qwen3_5 as qwen3_5
import sglang.srt.models.qwen3_5_text as qwen3_5_text
# A text-only Qwen3.5 checkpoint resolves to these classes, which shadow
# the multimodal ones by name — attaching the gate to the multimodal
# classes alone leaves the registry's text-only entries gate-less.
self.assertIs(
qwen3_5_text.Qwen3_5MoeForCausalLM.body_cls,
qwen3_5.Qwen3_5MoeForCausalLM,
)
text_config = SimpleNamespace(model_type="qwen3_5_moe_text")
seen, patcher = self._recording_gate(qwen3_5.Qwen3_5MoeForCausalLM)
with patcher:
self._reason(
qwen3_5_text.Qwen3_5MoeForCausalLM, text_config, _quant("quark")
)
self.assertIs(seen["config"], text_config)
self.assertEqual(seen["quant"].get_name(), "quark")
def test_the_qwen3_5_mtp_entry_normalizes_its_quantization(self):
from sglang.srt.models.qwen3_5 import Qwen3_5ForCausalLM
from sglang.srt.models.qwen3_5_mtp import (
Qwen3_5ForCausalLMMTP,
_mtp_quant_config,
)
# The normalization the constructor applies, shared with the gate.
self.assertIsNone(_mtp_quant_config(_quant("modelopt_mixed")))
serialized = SimpleNamespace(
get_name=lambda: "modelopt_fp4", is_checkpoint_nvfp4_serialized=True
)
self.assertIsNone(_mtp_quant_config(serialized))
# A non-serialized modelopt_fp4 checkpoint still converts on load, so
# the MTP module keeps the quantization.
online = SimpleNamespace(
get_name=lambda: "modelopt_fp4", is_checkpoint_nvfp4_serialized=False
)
self.assertIs(_mtp_quant_config(online), online)
quark_mtp = SimpleNamespace(
get_name=lambda: "quark", exclude_layers=["mtp.mlp.experts"]
)
self.assertIsNone(_mtp_quant_config(quark_mtp))
kept = _quant("fp8")
self.assertIs(_mtp_quant_config(kept), kept)
text_config = SimpleNamespace(model_type="qwen3_5_moe_text")
seen, patcher = self._recording_gate(Qwen3_5ForCausalLM)
with patcher:
self._reason(
Qwen3_5ForCausalLMMTP,
SimpleNamespace(text_config=text_config),
serialized,
)
self.assertIs(seen["config"], text_config)
self.assertIsNone(
seen["quant"], "the MTP module ships unquantized in that checkpoint"
)
class TestFamiliesWithoutAGate(_FusionGateCase):
def test_qwen2_moe_style_families_follow_the_intent(self):
"""A family with no gate must not grow one by accident: the installer
falls back to the user's intent for it."""
from sglang.srt.models.qwen2_moe import Qwen2MoeForCausalLM
self.assertFalse(
hasattr(Qwen2MoeForCausalLM, "shared_experts_fusion_disable_reason")
)
if __name__ == "__main__":
unittest.main()
@@ -0,0 +1,206 @@
"""A draft's construction decides for itself and leaves the process state alone.
The shared-experts-fusion decision is per checkpoint: each MoE model's gate
writes the ACTIVE moe flag (both ways) before its own layers build and read
it, and ``draft_model_build_scope`` — which brackets every draft
construction — records it on the speculative leaf and restores the target's
value on exit. The config bag keeps the
user's intent. A draft's weight update does not rewrite the
process's model_path record.
"""
import unittest
from types import SimpleNamespace
from sglang.srt.layers.moe.utils import (
draft_model_build_scope,
install_shared_experts_fusion_decision,
is_shared_experts_fusion_disabled,
speculative_moe_backend_context,
)
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.runtime_context import get_context, get_flags, get_model
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
class _AlwaysDisables:
"""A model class whose checkpoint can never fuse."""
@staticmethod
def shared_experts_fusion_disable_reason(hf_config, quant_config):
return "stand-in: this checkpoint cannot fuse."
class _NoGate:
"""A model family without an auto-disable gate: it follows the intent."""
def _install(model_class):
install_shared_experts_fusion_decision(model_class, SimpleNamespace(), None)
class TestFusionDecisionFlag(CustomTestCase):
def setUp(self):
super().setUp()
moe = get_flags().moe
self._saved = (
moe.disable_shared_experts_fusion,
moe.speculative_disable_shared_experts_fusion,
)
moe.disable_shared_experts_fusion = None
moe.speculative_disable_shared_experts_fusion = None
moe.in_speculative_scope = False
def tearDown(self):
moe = get_flags().moe
(
moe.disable_shared_experts_fusion,
moe.speculative_disable_shared_experts_fusion,
) = self._saved
super().tearDown()
def _seed(self, **fields):
override = get_context().override_server_args(**fields)
override.install()
self.addCleanup(override.restore)
def test_unset_flag_falls_back_to_the_config_intent(self):
self._seed(disable_shared_experts_fusion=True)
self.assertTrue(is_shared_experts_fusion_disabled())
self._seed(disable_shared_experts_fusion=False)
# A fresh install replaces the published config; the flag is still None.
self.assertFalse(is_shared_experts_fusion_disabled())
def test_the_installed_decision_wins_over_the_intent(self):
self._seed(disable_shared_experts_fusion=False)
_install(_AlwaysDisables)
self.assertTrue(is_shared_experts_fusion_disabled())
_install(_NoGate)
self.assertFalse(is_shared_experts_fusion_disabled())
def test_the_intent_short_circuits_the_gate(self):
# A user who passed --disable-shared-experts-fusion is not overruled,
# and the gate is not even asked.
self._seed(disable_shared_experts_fusion=True)
_install(_NoGate)
self.assertTrue(is_shared_experts_fusion_disabled())
def test_the_draft_build_scope_restores_the_targets_decision(self):
self._seed(disable_shared_experts_fusion=False)
_install(_NoGate) # the target's build
with draft_model_build_scope():
_install(_AlwaysDisables) # the draft's build
self.assertTrue(is_shared_experts_fusion_disabled())
self.assertFalse(is_shared_experts_fusion_disabled())
# The draft's decision stays inspectable on the twin leaf.
self.assertTrue(get_flags().moe.speculative_disable_shared_experts_fusion)
def test_a_gateless_draft_inherits_the_active_decision(self):
self._seed(disable_shared_experts_fusion=False)
_install(_AlwaysDisables) # the target's build
with draft_model_build_scope():
# A draft whose family has no gate follows the intent, which is what
# the target's own build already resolved to here.
self.assertTrue(is_shared_experts_fusion_disabled())
self.assertTrue(is_shared_experts_fusion_disabled())
def test_post_build_scopes_do_not_clobber_the_draft_leaf(self):
# init_attention_backends / cuda-graph capture / draft forwards enter
# scopes after construction; no gate runs there, so the persisted
# draft decision must survive.
self._seed(disable_shared_experts_fusion=False)
_install(_NoGate) # target's build
with draft_model_build_scope():
_install(_AlwaysDisables) # draft's build
for _ in range(3):
with draft_model_build_scope():
pass
with speculative_moe_backend_context():
pass
self.assertTrue(get_flags().moe.speculative_disable_shared_experts_fusion)
self.assertFalse(get_flags().moe.disable_shared_experts_fusion)
def test_the_build_scope_leaves_the_runner_backend_alone(self):
# Swapping runner_backend is speculative_moe_backend_context's job and
# must bracket the draft's whole lifecycle; dflash/dspark run their
# draft outside it, so a construction-only swap would build and
# execute the draft under different backends.
self._seed()
before = get_flags().moe.runner_backend
with draft_model_build_scope():
self.assertEqual(get_flags().moe.runner_backend, before)
self.assertEqual(get_flags().moe.runner_backend, before)
def test_a_record_outside_any_scope_is_target_only(self):
self._seed(disable_shared_experts_fusion=False)
get_flags().moe.speculative_disable_shared_experts_fusion = True
_install(_NoGate) # target's build
self.assertTrue(get_flags().moe.speculative_disable_shared_experts_fusion)
def test_initialize_moe_config_seeds_both_leaves(self):
from sglang.srt.layers.moe.utils import initialize_moe_config
from sglang.srt.server_args import ServerArgs
self._seed()
initialize_moe_config(
ServerArgs(model_path="dummy", disable_shared_experts_fusion=True)
)
moe = get_flags().moe
self.assertTrue(moe.disable_shared_experts_fusion)
self.assertTrue(moe.speculative_disable_shared_experts_fusion)
def test_a_forward_time_read_is_refused(self):
# The invariant behind the whole design: the decision is consumed at
# construction only. During a draft's build the flag holds the draft's
# value, so a forward reading it would race the build window.
from sglang.srt.model_executor.forward_context import (
ForwardContext,
forward_context,
)
self._seed()
with forward_context(ForwardContext(attn_backend=SimpleNamespace())):
with self.assertRaises(AssertionError):
is_shared_experts_fusion_disabled()
def test_the_intent_stays_on_the_bag(self):
self._seed(disable_shared_experts_fusion=False)
_install(_AlwaysDisables)
from sglang.srt.runtime_context import get_exec
self.assertFalse(get_exec().moe.disable_shared_experts_fusion)
class TestDraftWeightUpdateRecord(CustomTestCase):
def _seed(self, **fields):
override = get_context().override_server_args(**fields)
server_args = override.install()
self.addCleanup(override.restore)
return server_args
def _update(self, *, is_draft_worker: bool):
runner = ModelRunner.__new__(ModelRunner)
runner.is_draft_worker = is_draft_worker
runner.update_model_fields(
object(),
model_path="/new/checkpoint",
load_format="auto",
load_config=object(),
)
def test_a_target_update_is_recorded(self):
self._seed()
self._update(is_draft_worker=False)
self.assertEqual(get_model().model_path, "/new/checkpoint")
def test_a_draft_update_keeps_the_targets_record(self):
seeded = self._seed()
self._update(is_draft_worker=True)
self.assertEqual(get_model().model_path, seeded.model_path)
if __name__ == "__main__":
unittest.main()
+1 -30
View File
@@ -19,7 +19,6 @@ from sglang.srt.runtime_context import (
get_context, get_context,
get_flags, get_flags,
get_parallel, get_parallel,
get_schedule,
get_server_args, get_server_args,
reset_context, reset_context,
) )
@@ -403,6 +402,7 @@ class TestMoeFlagsGroup(_IsolatedServerArgs):
tbo_token_distribution_threshold=0.48, tbo_token_distribution_threshold=0.48,
disable_flashinfer_cutlass_moe_fp4_allgather=False, disable_flashinfer_cutlass_moe_fp4_allgather=False,
quantization=None, quantization=None,
disable_shared_experts_fusion=False,
) )
defaults.update(kw) defaults.update(kw)
initialize_moe_config(SimpleNamespace(**defaults)) initialize_moe_config(SimpleNamespace(**defaults))
@@ -964,35 +964,6 @@ class TestPublishLifecycle(_IsolatedServerArgs):
get_context().set_server_args(object()) get_context().set_server_args(object())
self.assertFalse(get_flags().capture.enable_torch_compile) self.assertFalse(get_flags().capture.enable_torch_compile)
def test_declare_load_time_override_writes_the_bag(self):
from sglang.srt.arg_groups.overrides import declare_load_time_override
args = self._publish(page_size=1)
declare_load_time_override("model.load_time", {"page_size": 64})
# The declaration lands on the config bag; the pristine startup record
# (server_args) is untouched.
self.assertEqual(get_schedule().page_size, 64)
self.assertEqual(args.page_size, 1)
def test_declare_load_time_override_validates_whitelist(self):
from sglang.srt.arg_groups.overrides import declare_load_time_override
args = self._publish(page_size=1)
with self.assertRaises(ValueError):
declare_load_time_override("bad", {"nope": 1})
self.assertEqual(args.page_size, 1)
def test_declare_load_time_override_records_provenance(self):
from sglang.srt.arg_groups.overrides import declare_load_time_override
self._publish(page_size=1)
declare_load_time_override("model.load_time", {"page_size": 64})
self.assertEqual(get_schedule().page_size, 64)
self.assertIn(
("model.load_time", {"page_size": 64}),
get_context().overrides_log(),
)
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()
@@ -69,7 +69,7 @@ class TestServerArgsMutationRatchet(CustomTestCase):
f"server_args mutations outside the resolution pipeline grew: " f"server_args mutations outside the resolution pipeline grew: "
f"{count} > baseline {_BASELINE}. Configuration is resolved in " f"{count} > baseline {_BASELINE}. Configuration is resolved in "
"ServerArgs.__post_init__; declare through the pipeline " "ServerArgs.__post_init__; declare through the pipeline "
"(passes / declare_load_time_override), change resolved config " "(passes / declare_late_resolution), change resolved config "
"with get_context().override(source, ...), or hand the value " "with get_context().override(source, ...), or hand the value "
"to its runner as a constructor argument — do not assign fields." "to its runner as a constructor argument — do not assign fields."
) )