From 4ed2548427a0f01a969d6e518088bcb62a568f5d Mon Sep 17 00:00:00 2001 From: Zheng Li Date: Thu, 12 Feb 2026 13:38:26 +0800 Subject: [PATCH] [Qwen3_5] Refactor `Qwen3_5ForCausalLMMTP` class implementation (#18538) --- python/sglang/srt/models/qwen3_5.py | 24 +++- python/sglang/srt/models/qwen3_5_mtp.py | 156 +++++++----------------- 2 files changed, 62 insertions(+), 118 deletions(-) diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py index c6c71247c..eb6166fad 100644 --- a/python/sglang/srt/models/qwen3_5.py +++ b/python/sglang/srt/models/qwen3_5.py @@ -330,6 +330,9 @@ class Qwen3_5LinearDecoderLayer(nn.Module): alt_stream=alt_stream, prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), ) + is_layer_sparse = True + is_previous_layer_sparse = True + is_next_layer_sparse = True elif config.model_type == "qwen3_5_text": self.mlp = Qwen2MoeMLP( hidden_size=config.hidden_size, @@ -338,15 +341,18 @@ class Qwen3_5LinearDecoderLayer(nn.Module): quant_config=quant_config, prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), ) + is_layer_sparse = False + is_previous_layer_sparse = False + is_next_layer_sparse = False else: raise ValueError(f"Invalid model type: {config.model_type}") self.layer_scatter_modes = LayerScatterModes.init_new( layer_id=layer_id, num_layers=config.num_hidden_layers, - is_layer_sparse=False, - is_previous_layer_sparse=False, - is_next_layer_sparse=False, + is_layer_sparse=is_layer_sparse, + is_previous_layer_sparse=is_previous_layer_sparse, + is_next_layer_sparse=is_next_layer_sparse, ) self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -491,6 +497,9 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): quant_config=quant_config, prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), ) + is_layer_sparse = False + is_previous_layer_sparse = False + is_next_layer_sparse = False elif config.model_type == "qwen3_5_moe_text": self.mlp = Qwen2MoeSparseMoeBlock( layer_id=layer_id, @@ -499,15 +508,18 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): alt_stream=alt_stream, prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), ) + is_layer_sparse = True + is_previous_layer_sparse = True + is_next_layer_sparse = True else: raise ValueError(f"Invalid model type: {config.model_type}") self.layer_scatter_modes = LayerScatterModes.init_new( layer_id=layer_id, num_layers=config.num_hidden_layers, - is_layer_sparse=False, - is_previous_layer_sparse=False, - is_next_layer_sparse=False, + is_layer_sparse=is_layer_sparse, + is_previous_layer_sparse=is_previous_layer_sparse, + is_next_layer_sparse=is_next_layer_sparse, ) self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) diff --git a/python/sglang/srt/models/qwen3_5_mtp.py b/python/sglang/srt/models/qwen3_5_mtp.py index c8d21cb3e..c4b18a19a 100644 --- a/python/sglang/srt/models/qwen3_5_mtp.py +++ b/python/sglang/srt/models/qwen3_5_mtp.py @@ -24,114 +24,15 @@ from sglang.srt.distributed import get_pp_group, get_tensor_model_parallel_world from sglang.srt.layers.layernorm import GemmaRMSNorm from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE -from sglang.srt.layers.vocab_parallel_embedding import ( - ParallelLMHead, - VocabParallelEmbedding, -) +from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.models.qwen3_5 import Qwen3_5AttentionDecoderLayer +from sglang.srt.models.qwen3_5 import Qwen3_5ForCausalLM from sglang.srt.utils import add_prefix logger = logging.getLogger(__name__) -class Qwen3_5MultiTokenPredictor(nn.Module): - def __init__(self, config: PretrainedConfig, quant_config=None, prefix: str = ""): - super().__init__() - - self.config = config - - self.vocab_size = config.vocab_size - - self.mtp_start_layer_idx = config.num_hidden_layers - self.num_mtp_layers = getattr(config, "mtp_num_hidden_layers", 1) - - self.embed_tokens = VocabParallelEmbedding( - self.vocab_size, - config.hidden_size, - ) - - self.fc = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) - - config.full_attention_interval = 1 - self.layers = torch.nn.ModuleList( - [ - Qwen3_5AttentionDecoderLayer( - config, - idx, - quant_config, - prefix=add_prefix(f"layers.{idx}", prefix), - ) - for idx in range(self.num_mtp_layers) - ] - ) - - self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.pre_fc_norm_hidden = GemmaRMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - self.pre_fc_norm_embedding = GemmaRMSNorm( - config.hidden_size, eps=config.rms_norm_eps - ) - - def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.embed_tokens(input_ids) - - @torch.no_grad() - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - forward_batch: torch.Tensor, - input_embeds: Optional[torch.Tensor] = None, - **kwargs, - ): - # if get_pp_group().is_first_rank: - assert input_embeds is None - input_embeds = forward_batch.mm_input_embeds - if ( - forward_batch.forward_mode.is_extend() - and forward_batch.contains_mm_inputs() - and not forward_batch.forward_mode.is_draft_extend() - ): - assert input_embeds is not None - input_embeds = torch.cat( - [input_embeds[:-1], self.embed_tokens(input_ids[-1].unsqueeze(0))] - ) - - if input_embeds is None: - input_embeds = self.embed_tokens(input_ids) - - hidden_states = forward_batch.spec_info.hidden_states - - # Some idle batch has 0 batch size. GemmaRMSNorm.forward would fail due to bs=0. - 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 = torch.cat([input_embeds, hidden_states], dim=-1) - - hidden_states = self.fc(hidden_states) - residual = None - - if self.num_mtp_layers == 1: - hidden_states, residual = self.layers[0]( - positions=positions, - hidden_states=hidden_states, - residual=residual, - forward_batch=forward_batch, - ) - else: - raise ("not implementation for other mtp layers[self.num_mtp_layers > 1]") - - if not get_pp_group().is_last_rank: - # For pipeline parallel, return intermediate tensors - return hidden_states - - hidden_states, _ = self.norm(hidden_states, residual) - return hidden_states - - class Qwen3_5ForCausalLMMTP(nn.Module): def __init__( @@ -140,7 +41,7 @@ class Qwen3_5ForCausalLMMTP(nn.Module): quant_config=None, prefix: str = "", ) -> None: - super().__init__() + nn.Module.__init__(self) self.is_multimodal = hasattr(config, "text_config") if self.is_multimodal: @@ -151,8 +52,18 @@ class Qwen3_5ForCausalLMMTP(nn.Module): self.quant_config = quant_config self.pp_group = get_pp_group() - self.model = Qwen3_5MultiTokenPredictor( - config, quant_config, prefix=add_prefix("mtp", prefix) + self.fc = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) + RMSNorm_cls = GemmaRMSNorm + self.pre_fc_norm_embedding = RMSNorm_cls( + config.hidden_size, config.rms_norm_eps + ) + self.pre_fc_norm_hidden = RMSNorm_cls(config.hidden_size, config.rms_norm_eps) + config.num_hidden_layers = 1 + config.full_attention_interval = 1 + self.model = Qwen3_5ForCausalLM( + config, + quant_config, + prefix=add_prefix("model", prefix), ) if get_pp_group().is_last_rank: @@ -165,9 +76,6 @@ class Qwen3_5ForCausalLMMTP(nn.Module): quant_config=quant_config, prefix=add_prefix("lm_head", prefix), ) - else: - # For pipeline parallel, create a placeholder layer - self.lm_head = nn.Linear(1, 1, bias=False) self.logits_processor = LogitsProcessor(config) @@ -193,17 +101,37 @@ class Qwen3_5ForCausalLMMTP(nn.Module): input_embeds: Optional[torch.Tensor] = None, **kwargs, ): + assert input_embeds is None + input_embeds = forward_batch.mm_input_embeds + if ( + forward_batch.forward_mode.is_extend() + and forward_batch.contains_mm_inputs() + and not forward_batch.forward_mode.is_draft_extend() + ): + assert input_embeds is not None + input_embeds = torch.cat( + [input_embeds[:-1], self.model.embed_tokens(input_ids[-1].unsqueeze(0))] + ) + + 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 = torch.cat([input_embeds, hidden_states], dim=-1) + + hidden_states = self.fc(hidden_states) + hidden_states = self.model( input_ids, positions, forward_batch, - input_embeds, + hidden_states, ) - if not get_pp_group().is_last_rank: - # For pipeline parallel, return intermediate results - return hidden_states - return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch ) @@ -293,6 +221,10 @@ class Qwen3_5ForCausalLMMTP(nn.Module): # Remove the mtp. prefix for processing name = name.replace("mtp.", "model.") + name = name.replace("model.fc", "fc") + name = name.replace("model.norm", "norm") + name = name.replace("model.pre_fc", "pre_fc") + if ".self_attn." in name: name = name.replace(".self_attn", "")