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