From cf66693b3530430b8d23aa6eefdfd922c8bfcc1f Mon Sep 17 00:00:00 2001 From: Byron Hsu Date: Fri, 29 May 2026 14:31:05 -0700 Subject: [PATCH] [Model] Add Qwen3-MoE MTP (#26468) Co-authored-by: Byron Hsu Co-authored-by: Cursor Co-authored-by: root --- python/sglang/srt/configs/model_config.py | 4 + python/sglang/srt/models/qwen3_moe.py | 19 ++- python/sglang/srt/models/qwen3_moe_mtp.py | 131 ++++++++++++++++++ python/sglang/srt/speculative/eagle_worker.py | 6 +- .../sglang/srt/speculative/eagle_worker_v2.py | 7 + 5 files changed, 163 insertions(+), 4 deletions(-) create mode 100644 python/sglang/srt/models/qwen3_moe_mtp.py diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index a94b6638a..e3a6035b0 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -474,6 +474,10 @@ class ModelConfig: self.hf_config.architectures[0] = "Qwen3NextForCausalLMMTP" self.hf_config.num_nextn_predict_layers = 1 + if is_draft_model and self.hf_config.architectures[0] == "Qwen3MoeForCausalLM": + self.hf_config.architectures[0] = "Qwen3MoeForCausalLMMTP" + self.hf_config.num_nextn_predict_layers = 1 + if is_draft_model and self.hf_config.architectures[0] in [ "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py index bbd95e6ce..5c6c76bf1 100644 --- a/python/sglang/srt/models/qwen3_moe.py +++ b/python/sglang/srt/models/qwen3_moe.py @@ -1107,7 +1107,9 @@ class Qwen3MoeForCausalLM(nn.Module): self.capture_aux_hidden_states = True self.model.set_dflash_layers_to_capture([val + 1 for val in layer_ids]) - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): + def load_weights( + self, weights: Iterable[Tuple[str, torch.Tensor]], is_mtp: bool = False + ): stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), @@ -1128,6 +1130,21 @@ class Qwen3MoeForCausalLM(nn.Module): params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: + if is_mtp: + if "mtp" not in name: + continue + + if name in [ + "mtp.fc.weight", + "mtp.pre_fc_norm_embedding.weight", + "mtp.pre_fc_norm_hidden.weight", + ]: + name = name.replace("mtp.", "") + else: + name = name.replace("mtp", "model") + elif "mtp" in name: + continue + layer_id = get_layer_id(name) if ( layer_id is not None diff --git a/python/sglang/srt/models/qwen3_moe_mtp.py b/python/sglang/srt/models/qwen3_moe_mtp.py new file mode 100644 index 000000000..e6f825eef --- /dev/null +++ b/python/sglang/srt/models/qwen3_moe_mtp.py @@ -0,0 +1,131 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Inference-only Qwen3-MoE MTP speculative decoding.""" + +import logging +from typing import Iterable, Optional, Tuple + +import torch +from torch import nn +from transformers import PretrainedConfig + +from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world_size +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.models.qwen3_moe import Qwen3MoeForCausalLM, Qwen3MoeModel +from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils import add_prefix + +logger = logging.getLogger(__name__) + + +class Qwen3MoeForCausalLMMTP(Qwen3MoeForCausalLM): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + self.config = config + config.num_hidden_layers = 1 + self.tp_size = get_tensor_model_parallel_world_size() + self.quant_config = quant_config + self.pp_group = get_pp_group() + + self.fc = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) + self.pre_fc_norm_embedding = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.pre_fc_norm_hidden = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.model = Qwen3MoeModel( + config, quant_config, prefix=add_prefix("model", prefix) + ) + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + use_attn_tp_group=get_global_server_args().enable_dp_lm_head, + ) + self.logits_processor = LogitsProcessor(config) + + # Required by Qwen3MoeForCausalLM.load_weights(), which we reuse below. + self.stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + self.expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts, + ) + self.capture_aux_hidden_states = False + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + input_embeds: Optional[torch.Tensor] = None, + **kwargs, + ): + if input_embeds is None: + input_embeds = self.model.embed_tokens(input_ids) + + hidden_states = forward_batch.spec_info.hidden_states + + if not forward_batch.forward_mode.is_idle(): + input_embeds = self.pre_fc_norm_embedding(input_embeds) + hidden_states = self.pre_fc_norm_hidden(hidden_states) + hidden_states = self.fc(torch.cat((input_embeds, hidden_states), dim=-1)) + + with get_global_expert_distribution_recorder().disable_this_region(): + hidden_states = self.model( + input_ids, + positions, + forward_batch, + hidden_states, + ) + + return self.logits_processor( + input_ids, hidden_states, self.lm_head, forward_batch + ) + + def load_weights( + self, weights: Iterable[Tuple[str, torch.Tensor]], is_mtp: bool = False + ): + return super().load_weights(weights, is_mtp=True) + + +EntryClass = [Qwen3MoeForCausalLMMTP] diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 1bf696791..ae84f8226 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -877,12 +877,12 @@ class EAGLEWorker(TpModelWorker): # Set inputs forward_batch.input_ids = input_ids - # This is a temporary fix for the case that the user is using standalone - # speculative decoding and the draft model architecture is gpt-oss. gpt-oss - # rope kernel needs cache_loc to be contiguous. + # Some draft model RoPE kernels need cache_loc to be contiguous. if ( self.server_args.speculative_algorithm == "STANDALONE" and self.model_config.hf_config.architectures[0] == "GptOssForCausalLM" + ) or self.model_config.hf_config.architectures[0] == ( + "Qwen3MoeForCausalLMMTP" ): out_cache_loc = out_cache_loc.contiguous() forward_batch.out_cache_loc = out_cache_loc[i] diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 186dd0a1e..8672913af 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -486,6 +486,13 @@ class EagleDraftWorker(BaseDraftWorker): # Set inputs forward_batch.input_ids = input_ids + # Qwen3-MoE MTP uses a fused RoPE + KV-store path whose cache_loc + # argument must be contiguous. + if ( + self.draft_runner.model_config.hf_config.architectures[0] + == "Qwen3MoeForCausalLMMTP" + ): + out_cache_loc = out_cache_loc.contiguous() forward_batch.out_cache_loc = out_cache_loc[i] spec_info.hidden_states = hidden_states