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