# Adapted from qwen2_moe.py # 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 Qwen3MoE model compatible with HuggingFace weights.""" import logging import math from typing import Any, Dict, Iterable, List, Optional, Tuple, TypeVar import torch from torch import nn from transformers import PretrainedConfig from sglang.srt.distributed import ( get_pp_group, moe_expert_parallel_all_reduce, moe_tensor_model_parallel_all_reduce, ) from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes from sglang.srt.layers.cp.utils import is_cp_v2_active from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( QKVParallelLinear, ReplicatedLinear, RowParallelLinear, ) from sglang.srt.layers.logits_processor import LogitsProcessor from sglang.srt.layers.moe import ( get_moe_a2a_backend, should_skip_post_experts_all_reduce, ) from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.topk import TopK from sglang.srt.layers.moe.utils import ( RoutingMethodType, filter_moe_weight_param_global_expert, ) from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import MRotaryEmbedding, get_rope from sglang.srt.layers.utils import get_layer_id from sglang.srt.layers.utils.cp_utils import ( can_cp_split, is_prefill_context_parallel_enabled, prepare_context_parallel_metadata, ) from sglang.srt.layers.vocab_parallel_embedding import ParallelLMHead from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.qwen2_moe import Qwen2MoeMLP as Qwen3MoeMLP from sglang.srt.models.qwen2_moe import Qwen2MoeModel from sglang.srt.models.utils import ( apply_qk_norm, create_fused_set_kv_buffer_arg, enable_fused_set_kv_buffer, ) from sglang.srt.runtime_context import ( get_forward, get_parallel, get_server_args, get_stream, ) from sglang.srt.utils import ( LazyValue, add_prefix, is_cuda, is_flashinfer_available, is_non_idle_and_non_empty, is_npu, ) from sglang.srt.utils.hf_transformers_utils import get_rope_config _is_cuda = is_cuda() if _is_cuda: from sglang.jit_kernel.fused_qknorm_rope import ( can_use_fused_qk_norm_rope, fused_qk_norm_rope, ) TConfig = TypeVar("TConfig", bound=PretrainedConfig) Qwen3MoeConfig = None _is_flashinfer_available = is_flashinfer_available() logger = logging.getLogger(__name__) _is_cuda = is_cuda() _is_npu = is_npu() if _is_npu: from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import split_qkv_rmsnorm_rope def compute_yarn_parameters( config: PretrainedConfig, ) -> tuple[float, float, float, float]: """ Refer to https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_rope_utils.py#L197C1-L288C1 Computes the inverse frequencies with NTK scaling. Please refer to the [original paper](https://huggingface.co/papers/2309.00071) Args: config ([`~transformers.PretrainedConfig`]): The model configuration. Returns: factor: float, the scaling factor for the RoPE embeddings low: float, the lower bound of the dimension range high: float, the upper bound of the dimension range attention_factor: float, the post-processing scaling factor applied to the computed cos/sin """ # The config does not contain rope_scaling, which means the model is not using yarn. # In transformers v5, rope_parameters is never None (even for default rope), so also # check rope_type to distinguish actual yarn configs from plain rotary embeddings. rope_scaling = getattr(config, "rope_parameters", None) if rope_scaling is None: rope_scaling = getattr(config, "rope_scaling", None) if rope_scaling is None: return 1.0, 0, 0, 1.0 rope_type = rope_scaling.get("rope_type") or rope_scaling.get("type") or "default" if rope_type == "default": return 1.0, 0, 0, 1.0 base = rope_scaling.get("rope_theta") or getattr(config, "rope_theta", 10000) partial_rotary_factor = ( config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 ) head_dim = getattr( config, "head_dim", config.hidden_size // config.num_attention_heads ) dim = int(head_dim * partial_rotary_factor) factor = rope_scaling.get("factor", 1.0) attention_factor = rope_scaling.get("attention_factor") mscale = rope_scaling.get("mscale") mscale_all_dim = rope_scaling.get("mscale_all_dim") if "original_max_position_embeddings" in rope_scaling: original_max_position_embeddings = rope_scaling[ "original_max_position_embeddings" ] factor = config.max_position_embeddings / original_max_position_embeddings else: original_max_position_embeddings = config.max_position_embeddings def get_mscale(scale, mscale=1): if scale <= 1: return 1.0 return 0.1 * mscale * math.log(scale) + 1.0 # Sets the attention factor as suggested in the paper if attention_factor is None: if mscale and mscale_all_dim: attention_factor = float( get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dim) ) else: attention_factor = get_mscale(factor) # Optional config options # beta_fast/beta_slow: as suggested in the paper, default to 32/1 (correspondingly) beta_fast = rope_scaling.get("beta_fast") or 32 beta_slow = rope_scaling.get("beta_slow") or 1 # Compute the inverse frequencies def find_correction_dim(num_rotations, dim, base, max_position_embeddings): """Inverse dimension formula to find the dimension based on the number of rotations""" return ( dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) ) / (2 * math.log(base)) def find_correction_range( low_rot, high_rot, dim, base, max_position_embeddings, truncate ): """Find dimension range bounds based on rotations""" low = find_correction_dim(low_rot, dim, base, max_position_embeddings) high = find_correction_dim(high_rot, dim, base, max_position_embeddings) if truncate: low = math.floor(low) high = math.ceil(high) return max(low, 0), min(high, dim - 1) truncate = rope_scaling.get("truncate", True) low, high = find_correction_range( beta_fast, beta_slow, dim, base, original_max_position_embeddings, truncate ) # These parts are implemented in the fusedQKNormRopeKernel.cu # # def linear_ramp_factor(min, max, dim): # # if min == max: # # max += 0.001 # Prevent singularity # # linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) # # ramp_func = torch.clamp(linear_func, 0, 1) # # return ramp_func # # Note on variable naming: "interpolation" comes from the original technique, where we interpolate the position IDs # # to expand the possible context length. In other words, interpolation = apply scaling factor. # # pos_freqs = base ** (torch.arange(0, dim, 2).to(device=device, dtype=torch.float) / dim) # # inv_freq_extrapolation = 1.0 / pos_freqs # # inv_freq_interpolation = 1.0 / (factor * pos_freqs) # # # Get n-dimensional rotational scaling corrected for extrapolation # # inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2).to(device=device, dtype=torch.float) # # inv_freq = ( # # inv_freq_interpolation * (1 - inv_freq_extrapolation_factor) # # + inv_freq_extrapolation * inv_freq_extrapolation_factor # # ) # # return inv_freq, attention_factor return factor, low, high, attention_factor class Qwen3MoeSparseMoeBlock(nn.Module): def __init__( self, layer_id: int, config: Qwen3MoeConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.tp_size = get_parallel().moe_tp_size self.ep_size = get_parallel().moe_ep_size self.layer_id = layer_id if self.tp_size > config.num_experts: raise ValueError( f"Tensor parallel size {self.tp_size} is greater than " f"the number of experts {config.num_experts}." ) from sglang.srt.layers.quantization.gguf import GGUFConfig norm_topk_prob = getattr(config, "norm_topk_prob", True) if isinstance(quant_config, GGUFConfig): norm_topk_prob = False self.topk = TopK( top_k=config.num_experts_per_tok, renormalize=norm_topk_prob, use_grouped_topk=False, layer_id=layer_id, ) self.experts = get_moe_impl_class(quant_config)( num_experts=config.num_experts + get_server_args().ep_num_redundant_experts, top_k=config.num_experts_per_tok, layer_id=layer_id, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, quant_config=quant_config, prefix=add_prefix("experts", prefix), routing_method_type=RoutingMethodType.Renormalize, ) self.gate = ReplicatedLinear( config.hidden_size, config.num_experts, bias=False, quant_config=None, prefix=add_prefix("gate", prefix), ) if get_moe_a2a_backend().is_deepep(): # TODO: we will support tp < ep in the future self.ep_size = get_parallel().moe_ep_size self.num_experts = ( config.num_experts + get_server_args().ep_num_redundant_experts ) self.top_k = config.num_experts_per_tok def forward( self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] = None, ) -> torch.Tensor: if ( not get_moe_a2a_backend().is_deepep() and not get_moe_a2a_backend().is_ascend_fuseep() ): return self.forward_normal(hidden_states) else: return self.forward_deepep(hidden_states, forward_batch) def get_moe_weights(self): return [ x.data for name, x in self.experts.named_parameters() if name not in ["correction_bias"] and filter_moe_weight_param_global_expert( name, x, self.experts.num_local_experts ) ] def forward_normal( self, hidden_states: torch.Tensor, ) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) topk_output = self.topk(hidden_states, router_logits) final_hidden_states = self.experts(hidden_states, topk_output) if self.ep_size > 1 and not should_skip_post_experts_all_reduce( is_tp_path=False ): final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states) if self.tp_size > 1 and not should_skip_post_experts_all_reduce( is_tp_path=True ): final_hidden_states = moe_tensor_model_parallel_all_reduce( final_hidden_states ) return final_hidden_states.view(num_tokens, hidden_dim) def forward_deepep( self, hidden_states: torch.Tensor, forward_batch: ForwardBatch ) -> torch.Tensor: if hidden_states.shape[0] > 0: # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) topk_output = self.topk( hidden_states, router_logits, num_token_non_padded=forward_batch.num_token_non_padded, expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( layer_id=self.layer_id, ), ) else: topk_output = self.topk.empty_topk_output(hidden_states.device) final_hidden_states = self.experts( hidden_states=hidden_states, topk_output=topk_output, ) return final_hidden_states def op_gate(self, state): if is_non_idle_and_non_empty( state.forward_batch.forward_mode, state.hidden_states_mlp_input ): # router_logits: (num_tokens, n_experts) state.router_logits, _ = self.gate(state.hidden_states_mlp_input) else: state.router_logits = None def op_select_experts(self, state): router_logits = state.pop("router_logits") hidden_states = state.hidden_states_mlp_input if router_logits is not None: with get_global_expert_distribution_recorder().with_current_layer( self.layer_id ): state.topk_output = self.topk( hidden_states=hidden_states, router_logits=router_logits, num_token_non_padded=state.forward_batch.num_token_non_padded, expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( layer_id=self.layer_id, ), ) else: state.topk_output = self.topk.empty_topk_output(hidden_states.device) def op_dispatch_a(self, state): if self.ep_size > 1: self.experts.dispatcher.dispatch_a( hidden_states=state.pop("hidden_states_mlp_input"), topk_output=state.pop("topk_output"), tbo_subbatch_index=state.get("tbo_subbatch_index"), ) def op_dispatch_b(self, state): if self.ep_size > 1: with get_global_expert_distribution_recorder().with_current_layer( self.layer_id ): state.dispatch_output = self.experts.dispatcher.dispatch_b( tbo_subbatch_index=state.get("tbo_subbatch_index"), ) def op_experts(self, state): state.combine_input = self.experts.run_moe_core( dispatch_output=state.dispatch_output, ) def op_combine_a(self, state): if self.ep_size > 1: self.experts.dispatcher.combine_a( combine_input=state.pop("combine_input"), tbo_subbatch_index=state.get("tbo_subbatch_index"), ) state.pop("dispatch_output") def op_combine_b(self, state): if self.ep_size > 1: state.hidden_states_after_combine = self.experts.dispatcher.combine_b( tbo_subbatch_index=state.get("tbo_subbatch_index"), ) def op_output(self, state): state.hidden_states_mlp_output = state.pop("hidden_states_after_combine") class Qwen3MoeAttention(nn.Module): def __init__( self, hidden_size: int, num_heads: int, num_kv_heads: int, layer_id: int = 0, start_layer: int = 0, rope_theta: float = 10000, rope_scaling: Optional[Dict[str, Any]] = None, max_position_embeddings: int = 8192, head_dim: Optional[int] = None, rms_norm_eps: float = 1e-06, attention_bias: bool = False, config: Optional[TConfig] = None, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", dual_chunk_attention_config: Optional[dict[str, Any]] = None, alt_stream: Optional[torch.cuda.Stream] = None, ) -> None: super().__init__() self.hidden_size = hidden_size self.start_layer = start_layer attn_tp_rank = get_parallel().attn_tp_rank attn_tp_size = get_parallel().attn_tp_size self.config = config self.total_num_heads = num_heads assert self.total_num_heads % attn_tp_size == 0 self.num_heads = self.total_num_heads // attn_tp_size self.total_num_kv_heads = num_kv_heads if self.total_num_kv_heads >= attn_tp_size: # Number of KV heads is greater than TP size, so we partition # the KV heads across multiple tensor parallel GPUs. assert self.total_num_kv_heads % attn_tp_size == 0 else: # Number of KV heads is less than TP size, so we replicate # the KV heads across multiple tensor parallel GPUs. assert attn_tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) self.head_dim = head_dim or hidden_size // self.total_num_heads self.q_size = self.num_heads * self.head_dim self.kv_size = self.num_kv_heads * self.head_dim self.scaling = self.head_dim**-0.5 self.rope_theta = rope_theta self.max_position_embeddings = max_position_embeddings self.tp_rank = get_parallel().tp_rank self.qkv_proj = QKVParallelLinear( hidden_size, self.head_dim, self.total_num_heads, self.total_num_kv_heads, bias=attention_bias, quant_config=quant_config, tp_rank=attn_tp_rank, tp_size=attn_tp_size, prefix=add_prefix("qkv_proj", prefix), ) self.o_proj = RowParallelLinear( self.total_num_heads * self.head_dim, hidden_size, bias=attention_bias, quant_config=quant_config, tp_rank=attn_tp_rank, tp_size=attn_tp_size, reduce_results=False, prefix=add_prefix("o_proj", prefix), ) self.rotary_emb = get_rope( self.head_dim, rotary_dim=self.head_dim, max_position=max_position_embeddings, base=rope_theta, rope_scaling=rope_scaling, dual_chunk_attention_config=dual_chunk_attention_config, ) self.compatible_with_fused_kv_buffer = ( False if isinstance(self.rotary_emb, MRotaryEmbedding) else True ) self.compatible_with_fused_qk_norm_rope = not isinstance( self.rotary_emb, MRotaryEmbedding ) and self.head_dim in (64, 128, 256) _yarn_factor, _, _, _ = compute_yarn_parameters(config) self.use_fused_qk_norm_rope = ( get_server_args().enable_fused_qk_norm_rope and self.compatible_with_fused_qk_norm_rope and _is_cuda and can_use_fused_qk_norm_rope( self.head_dim, self.rotary_emb.is_neox_style, torch.bfloat16, _yarn_factor != 1.0, ) ) self._used_fused_qk_norm_rope_last_call = False self.attn = RadixAttention( self.num_heads, self.head_dim, self.scaling, num_kv_heads=self.num_kv_heads, layer_id=layer_id, prefix=add_prefix("attn", prefix), ) self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) self.alt_stream = alt_stream def op_prepare(self, state): state.attn_intermediate_state = self.forward_prepare( positions=state.positions, hidden_states=state.pop("hidden_states_after_comm_pre_attn"), forward_batch=state.forward_batch, ) def op_core(self, state): state.hidden_states_after_attn = self.forward_core( state.pop("attn_intermediate_state") ) def forward_prepare_npu( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ): qkv, _ = self.qkv_proj(hidden_states) if self.attn.layer_id == self.start_layer: self.rotary_emb.get_cos_sin_with_position(positions) q, k, v = split_qkv_rmsnorm_rope( qkv, self.rotary_emb.position_sin, self.rotary_emb.position_cos, self.q_size, self.kv_size, self.head_dim, eps=self.q_norm.variance_epsilon, q_weight=self.q_norm.weight, k_weight=self.k_norm.weight, q_bias=getattr(self.q_norm, "bias", None), k_bias=getattr(self.k_norm, "bias", None), ) inner_state = q, k, v, forward_batch return None, forward_batch, inner_state def forward_prepare_native( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ): qkv, _ = self.qkv_proj(hidden_states) q, k, v = self.apply_qk_norm_rope(qkv, positions, forward_batch) inner_state = q, k, v, forward_batch return None, forward_batch, inner_state def apply_qk_norm_rope(self, qkv, positions, forward_batch): use_fused = self.use_fused_qk_norm_rope and qkv.dtype == torch.bfloat16 if use_fused: theta = self.rope_theta positions = ( positions.view(-1).to(dtype=torch.int32, device=qkv.device).contiguous() ) factor, low, high, attention_factor = compute_yarn_parameters(self.config) fused_qk_norm_rope( qkv, self.num_heads, self.num_kv_heads, self.num_kv_heads, self.head_dim, self.q_norm.variance_epsilon, self.q_norm.weight, self.k_norm.weight, theta, self.rotary_emb.is_neox_style, positions, factor, low, high, attention_factor, ) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) self._used_fused_qk_norm_rope_last_call = True else: # Fallback to non-fused QK Norm & RoPE implementation q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q, k = apply_qk_norm( q=q, k=k, q_norm=self.q_norm, k_norm=self.k_norm, head_dim=self.head_dim, alt_stream=self.alt_stream, ) q, k = self.rotary_emb( positions, q, k, fused_set_kv_buffer_arg=( create_fused_set_kv_buffer_arg( value=v, layer=self.attn, forward_batch=forward_batch, ) if enable_fused_set_kv_buffer(forward_batch) and self.compatible_with_fused_kv_buffer else None ), ) self._used_fused_qk_norm_rope_last_call = False return q, k, v def forward_prepare( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ): if hidden_states.shape[0] == 0: return hidden_states, forward_batch, None if not _is_npu: return self.forward_prepare_native( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, ) else: return self.forward_prepare_npu( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, ) def forward_core(self, intermediate_state): hidden_states, forward_batch, inner_state = intermediate_state if inner_state is None: return hidden_states q, k, v, fb = inner_state must_save_kv = self._used_fused_qk_norm_rope_last_call save_kv_cache = must_save_kv or not ( enable_fused_set_kv_buffer(forward_batch) and self.compatible_with_fused_kv_buffer ) attn_output = self.attn( q, k, v, fb, save_kv_cache=save_kv_cache, ) output, _ = self.o_proj(attn_output) return output def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: s = self.forward_prepare( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, ) return self.forward_core(s) class Qwen3MoeDecoderLayer(nn.Module): def __init__( self, config: Qwen3MoeConfig, layer_id: int, start_layer: int = 0, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", alt_stream: Optional[torch.cuda.Stream] = None, ) -> None: super().__init__() self.config = config self.hidden_size = config.hidden_size rope_theta, rope_scaling = get_rope_config(config) self.rope_theta = rope_theta max_position_embeddings = getattr(config, "max_position_embeddings", 8192) head_dim = getattr( config, "head_dim", config.hidden_size // config.num_attention_heads ) rms_norm_eps = config.rms_norm_eps attention_bias = config.attention_bias dual_chunk_attention_config = getattr( config, "dual_chunk_attention_config", None ) self.self_attn = Qwen3MoeAttention( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, num_kv_heads=config.num_key_value_heads, layer_id=layer_id, start_layer=start_layer, rope_theta=rope_theta, rope_scaling=rope_scaling, max_position_embeddings=max_position_embeddings, head_dim=head_dim, rms_norm_eps=rms_norm_eps, attention_bias=attention_bias, config=config, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), dual_chunk_attention_config=dual_chunk_attention_config, alt_stream=alt_stream, ) self.layer_id = layer_id self.attn_tp_size = get_parallel().attn_tp_size self.attn_tp_rank = get_parallel().attn_tp_rank # Qwen3MoE all layers are sparse and have no nextn now self.is_layer_sparse = True is_previous_layer_sparse = True is_next_layer_sparse = True self.layer_scatter_modes = LayerScatterModes.init_new( layer_id=layer_id, num_layers=config.num_hidden_layers, is_layer_sparse=self.is_layer_sparse, is_previous_layer_sparse=is_previous_layer_sparse, is_next_layer_sparse=is_next_layer_sparse, ) if self.is_layer_sparse: self.mlp = Qwen3MoeSparseMoeBlock( layer_id=self.layer_id, config=config, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) else: self.mlp = Qwen3MoeMLP( hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps ) self.layer_communicator = LayerCommunicator( layer_scatter_modes=self.layer_scatter_modes, input_layernorm=self.input_layernorm, post_attention_layernorm=self.post_attention_layernorm, allow_reduce_scatter=True, is_last_layer=(self.layer_id == self.config.num_hidden_layers - 1), ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], captured_last_layer_outputs: Optional[List[torch.Tensor]] = None, **kwargs, ) -> Tuple[torch.Tensor, torch.Tensor]: hidden_states, residual = ( self.layer_communicator.prepare_attn_and_capture_last_layer_outputs( hidden_states, residual, forward_batch, captured_last_layer_outputs=captured_last_layer_outputs, **kwargs, ) ) if hidden_states.shape[0] != 0: hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, ) hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual, forward_batch ) fuse_mlp_allreduce = ( self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( forward_batch ) ) # For DP with padding, reduce scatter can be used instead of all-reduce. mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( forward_batch ) with get_forward().scoped( fuse_mlp_allreduce=fuse_mlp_allreduce, mlp_reduce_scatter=mlp_reduce_scatter, ): hidden_states = self.mlp(hidden_states, forward_batch) if fuse_mlp_allreduce: hidden_states._sglang_needs_allreduce_fusion = True else: hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual, forward_batch ) return hidden_states, residual def op_comm_prepare_attn( self, state, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], tbo_subbatch_index: Optional[int] = None, ): state.hidden_states_after_comm_pre_attn, state.residual_after_input_ln = ( self.layer_communicator.prepare_attn(hidden_states, residual, forward_batch) ) state.update( dict( forward_batch=forward_batch, positions=positions, tbo_subbatch_index=tbo_subbatch_index, ) ) def op_comm_prepare_mlp(self, state): state.hidden_states_mlp_input, state.residual_after_comm_pre_mlp = ( self.layer_communicator.prepare_mlp( state.pop("hidden_states_after_attn"), state.pop("residual_after_input_ln"), state.forward_batch, ) ) def op_comm_postprocess_layer(self, state): hidden_states, residual = self.layer_communicator.postprocess_layer( state.pop("hidden_states_mlp_output"), state.pop("residual_after_comm_pre_mlp"), state.forward_batch, ) output = dict( positions=state.positions, hidden_states=hidden_states, residual=residual, forward_batch=state.forward_batch, tbo_subbatch_index=state.tbo_subbatch_index, ) state.clear( expect_keys={ "positions", "forward_batch", "tbo_subbatch_index", } ) return output class Qwen3MoeModel(Qwen2MoeModel): def __init__( self, config: Qwen3MoeConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", decoder_layer_type=Qwen3MoeDecoderLayer, ) -> None: alt_stream = get_stream("alt") if _is_cuda else None super().__init__( config=config, quant_config=quant_config, prefix=prefix, decoder_layer_type=decoder_layer_type, alt_stream=alt_stream, ) def set_dflash_layers_to_capture(self, layers_to_capture: List[int]): self.layers_to_capture = layers_to_capture for layer_id in self.layers_to_capture: setattr(self.layers[layer_id], "_is_layer_to_capture", True) class Qwen3MoeForCausalLM(nn.Module): fall_back_to_pt_during_load = False # Mapping from fused module names to their component weight names. # Required for quantization configs (e.g., ModelOpt FP4) to correctly identify # which layers should be skipped based on the exclude_modules/ignore list. packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], "gate_up_proj": ["gate_proj", "up_proj"], } def __init__( self, config: Qwen3MoeConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config 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_server_args().enable_dp_lm_head, ) self.logits_processor = LogitsProcessor(config) self.capture_aux_hidden_states = False self.attn_cp_size = get_parallel().attn_cp_size self.attn_cp_rank = get_parallel().attn_cp_rank self.moe_dp_size = get_parallel().moe_dp_size assert self.attn_cp_size % self.moe_dp_size == 0, ( f"attn_cp_size ({self.attn_cp_size}) must be divisible by " f"moe_dp_size ({self.moe_dp_size})" ) def get_input_embeddings(self) -> nn.Embedding: return self.model.embed_tokens @torch.no_grad() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> torch.Tensor: if is_prefill_context_parallel_enabled() and not is_cp_v2_active(forward_batch): if can_cp_split(len(input_ids), self.attn_cp_size, forward_batch): forward_batch.attn_cp_metadata = prepare_context_parallel_metadata( len(input_ids), self.attn_cp_rank, self.attn_cp_size, forward_batch.seq_lens_cpu.tolist(), extend_seqs_len=forward_batch.extend_seq_lens_cpu, ) hidden_states = self.model( input_ids, positions, forward_batch, input_embeds, pp_proxy_tensors=pp_proxy_tensors, ) aux_hidden_states = None if self.capture_aux_hidden_states: hidden_states, aux_hidden_states = hidden_states if self.pp_group.is_last_rank: logits_output = self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states ) return logits_output else: return hidden_states @torch.no_grad() def forward_split_prefill( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, split_interval: Tuple[int, int], # [start, end) 0-based input_embeds: torch.Tensor = None, ): start, end = split_interval # embed if start == 0: if input_embeds is None: forward_batch.hidden_states = self.model.embed_tokens(input_ids) else: forward_batch.hidden_states = input_embeds # decoder layer for i in range(start, end): with get_global_expert_distribution_recorder().with_current_layer(i): layer = self.model.layers[i] forward_batch.hidden_states, forward_batch.residual = layer( positions, forward_batch.hidden_states, forward_batch, forward_batch.residual, ) if end == self.model.config.num_hidden_layers: # norm hidden_states, _ = self.model.norm( forward_batch.hidden_states, forward_batch.residual ) forward_batch.hidden_states = hidden_states # logits process result = self.logits_processor( input_ids, forward_batch.hidden_states, self.lm_head, forward_batch ) else: result = None return result @property def start_layer(self): return self.model.start_layer @property def end_layer(self): return self.model.end_layer def get_embed_and_head(self): return self.model.embed_tokens.weight, self.lm_head.weight def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None): if not self.pp_group.is_last_rank: return self.capture_aux_hidden_states = True if layer_ids is None: num_layers = self.config.num_hidden_layers self.model.set_eagle3_layers_to_capture( [ 2, num_layers // 2, num_layers - 3, ] ) # Specific layers for EAGLE3 support else: self.model.set_eagle3_layers_to_capture([val + 1 for val in layer_ids]) def set_dflash_layers_to_capture(self, layer_ids: List[int]): if not self.pp_group.is_last_rank: return if layer_ids is None: raise ValueError( "DFLASH requires explicit layer_ids for aux hidden capture." ) 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]], is_mtp: bool = False ): stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("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), ] 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, ) # Pre-define `params_dict` to avoid repeated expensive traversal of model parameters. 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 and hasattr(self.model, "start_layer") and ( layer_id < self.model.start_layer or layer_id >= self.model.end_layer ) ): continue if "rotary_emb.inv_freq" in name: continue for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if weight_name not in name: continue # We have mlp.experts[0].gate_proj in the checkpoint. # Since we handle the experts below in expert_params_mapping, # we need to skip here BEFORE we update the name, otherwise # name will be updated to mlp.experts[0].gate_up_proj, which # will then be updated below in expert_params_mapping # for mlp.experts[0].gate_gate_up_proj, which breaks load. if "mlp.experts" in name: continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: # Track if this is an expert weight to enable early skipping is_expert_weight = False for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue # Mark as expert weight regardless of whether we can process it is_expert_weight = True name = name.replace(weight_name, param_name) if name not in params_dict: # Expert weight not on this rank, will be skipped below continue param = params_dict[name] weight_loader = param.weight_loader weight_loader( param, loaded_weight, name, shard_id=shard_id, expert_id=expert_id, ) break else: if is_expert_weight: # This is an expert weight but not mapped to this rank, skip all remaining processing continue # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue if name not in params_dict: continue if name in params_dict.keys(): param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) else: logger.warning(f"Parameter {name} not found in params_dict") if not hasattr(self, "routed_experts_weights_of_layer"): self.routed_experts_weights_of_layer = LazyValue( lambda: { layer_id: self.model.layers[layer_id].mlp.get_moe_weights() for layer_id in range(self.start_layer, self.end_layer) if isinstance( self.model.layers[layer_id].mlp, Qwen3MoeSparseMoeBlock ) } ) @classmethod def get_model_config_for_expert_location(cls, config): return ModelConfigForExpertLocation( num_layers=config.num_hidden_layers, num_logical_experts=config.num_experts, num_groups=None, ) EntryClass = Qwen3MoeForCausalLM