# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # 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. # ============================================================================== # Adapted from # https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/qwen2_moe.py """Inference-only Qwen2MoE model compatible with HuggingFace weights.""" import logging from contextlib import nullcontext from typing import Any, Dict, Iterable, List, Optional, Tuple, Union import torch import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig from sglang.srt.batch_overlap.two_batch_overlap import model_forward_maybe_tbo from sglang.srt.distributed import ( get_pp_group, get_pp_indices, moe_expert_parallel_all_reduce, moe_tensor_model_parallel_all_reduce, 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.activation import SiluAndMul from sglang.srt.layers.communicator import ( LayerCommunicator, LayerScatterModes, ScatterMode, ) from sglang.srt.layers.cp.utils import is_cp_v2_active from sglang.srt.layers.dp_attention import ( is_dp_attention_enabled, ) from sglang.srt.layers.elementwise import fused_gate_sigmoid_mul_add from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( MergedColumnParallelLinear, 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 import FusedMoE from sglang.srt.layers.moe.topk import StandardTopKOutput, TopK, TopKOutputChecker from sglang.srt.layers.moe.utils import ( RoutingMethodType, filter_moe_weight_param_global_expert, is_deepep_class_backend, ) from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.rotary_embedding import get_rope from sglang.srt.layers.utils import PPMissingLayer, get_layer_id from sglang.srt.layers.utils.cp_utils import ( cp_all_gather_rerange_output, cp_split_and_rebuild_data, cp_split_and_rebuild_position, is_prefill_context_parallel_enabled, ) from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from sglang.srt.model_executor.cuda_graph_config import ( Backend, Phase, check_cuda_graph_backend, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.runtime_context import get_forward, get_parallel, get_server_args from sglang.srt.utils import ( add_prefix, cpu_has_amx_support, get_bool_env_var, is_cpu, is_cuda, is_hip, is_npu, make_layers, use_intel_amx_backend, ) if is_npu(): from sglang.srt.hardware_backend.npu.cmo import ( shared_expert_on_independent_stream, wait_share_stream, ) from sglang.srt.environ import envs from sglang.srt.runtime_context import get_stream from sglang.srt.utils.hf_transformers_utils import get_rope_config _SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get() logger = logging.getLogger(__name__) _is_cuda = is_cuda() _is_cpu = is_cpu() _is_cpu_amx_available = cpu_has_amx_support() _is_hip = is_hip() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip def get_num_shared_experts(config: PretrainedConfig) -> int: n_shared_experts = getattr(config, "n_shared_experts", None) if n_shared_experts is not None: return n_shared_experts if ( hasattr(config, "shared_expert_intermediate_size") and config.shared_expert_intermediate_size > 0 ): return 1 return 0 def can_fuse_shared_expert( config: PretrainedConfig, quant_config: Optional[QuantizationConfig], ) -> bool: """Whether the shared expert may be fused as an extra MoE expert. Caller must still gate on the model/backend support flag. """ if ( get_server_args().disable_shared_experts_fusion is True or getattr(config, "shared_expert_intermediate_size", 0) <= 0 or config.shared_expert_intermediate_size != config.moe_intermediate_size or get_moe_a2a_backend().is_deepep() ): return False if quant_config is not None: exclude_layers = getattr(quant_config, "exclude_layers", None) if exclude_layers is None: exclude_layers = getattr(quant_config, "ignored_layers", []) # Other backends than quark do not exclude the shared expert here, so they # intentionally fall through and remain fusable can_fuse_fn = getattr(quant_config, "can_fuse_shared_expert", None) if can_fuse_fn is not None: if not can_fuse_fn(): return False return True class Qwen2MoeMLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, hidden_act: str, quant_config: Optional[QuantizationConfig] = None, reduce_results: bool = True, prefix: str = "", tp_rank: Optional[int] = None, tp_size: Optional[int] = None, ) -> None: super().__init__() self.gate_up_proj = MergedColumnParallelLinear( hidden_size, [intermediate_size] * 2, bias=False, quant_config=quant_config, prefix=add_prefix("gate_up_proj", prefix), tp_rank=tp_rank, tp_size=tp_size, ) self.down_proj = RowParallelLinear( intermediate_size, hidden_size, bias=False, quant_config=quant_config, reduce_results=reduce_results, prefix=add_prefix("down_proj", prefix), tp_rank=tp_rank, tp_size=tp_size, ) if hidden_act != "silu": raise ValueError( f"Unsupported activation: {hidden_act}. Only silu is supported for now." ) self.act_fn = SiluAndMul() def forward( self, x, ): gate_up, _ = self.gate_up_proj(x) x = self.act_fn(gate_up) x, _ = self.down_proj(x) return x class Qwen2MoeSparseMoeBlock(nn.Module): def __init__( self, layer_id: int, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, alt_stream: Optional[torch.cuda.Stream] = None, prefix: str = "", is_nextn: bool = False, support_shared_expert_fusion: bool = False, enable_cuda_shared_expert_fusion: bool = False, ): super().__init__() self.tp_size = get_parallel().tp_size self.layer_id = layer_id self.alt_stream = alt_stream 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}." ) self.num_experts = config.num_experts self.num_shared_experts = get_num_shared_experts(config) self.num_fused_shared_experts = 0 self.enable_shared_expert_fusion = False # default to False if support_shared_expert_fusion and ( _use_aiter or (_is_cuda and enable_cuda_shared_expert_fusion) ): self.enable_shared_expert_fusion = ( self.num_shared_experts > 0 and can_fuse_shared_expert(config, quant_config) ) if self.enable_shared_expert_fusion: self.num_fused_shared_experts = self.num_shared_experts self.topk = TopK( top_k=config.num_experts_per_tok, renormalize=config.norm_topk_prob, layer_id=layer_id, ) # Disable inplace MoE when fused gate will need hidden_states after experts _needs_hidden_after_experts = ( config.shared_expert_intermediate_size > 0 and not self.enable_shared_expert_fusion ) self.experts = get_moe_impl_class(quant_config)( layer_id=self.layer_id, top_k=( config.num_experts_per_tok if not self.enable_shared_expert_fusion else config.num_experts_per_tok + self.num_fused_shared_experts ), num_experts=( config.num_experts + get_server_args().ep_num_redundant_experts if not self.enable_shared_expert_fusion else config.num_experts + get_server_args().ep_num_redundant_experts + self.num_fused_shared_experts ), hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, quant_config=quant_config, prefix=add_prefix("experts", prefix), routing_method_type=RoutingMethodType.RenormalizeNaive, num_fused_shared_experts=self.num_fused_shared_experts, inplace=not _needs_hidden_after_experts, ) self.gate = ReplicatedLinear( config.hidden_size, config.num_experts, bias=False, quant_config=None, prefix=add_prefix("gate", prefix), ) # When enable_shared_expert_fusion, the shared expert runs inside the MoE kernel # (via _append_shared_to_topk_output); a separate shared_expert MLP would # double-count. If fusion is off (num_fused_shared_experts == 0), keep shared_expert. if ( config.shared_expert_intermediate_size > 0 and not self.enable_shared_expert_fusion ): self.shared_expert = Qwen2MoeMLP( hidden_size=config.hidden_size, intermediate_size=config.shared_expert_intermediate_size, hidden_act=config.hidden_act, quant_config=quant_config, reduce_results=False, prefix=add_prefix("shared_expert", prefix), **( dict(tp_rank=0, tp_size=1) if ( get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_flashinfer() ) else {} ), ) else: self.shared_expert = None if _is_cpu and _is_cpu_amx_available: self.shared_expert_gate = ReplicatedLinear( config.hidden_size, 1, bias=False, quant_config=None, prefix=add_prefix("shared_expert_gate", prefix), ) else: self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) 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 self.is_nextn = is_nextn 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 _get_shared_expert_weights( self, hidden_states: torch.Tensor ) -> Optional[Tuple[torch.Tensor, float]]: """Return the shared_expert_gate weights and the 1/ep_size scale. On the AMD AITER path the sigmoid activation and the scale are applied (in fp32) inside the fused append kernel, so this returns the raw gate logits to avoid a standalone activation kernel + cast. On the CUDA path the legacy eager ``sigmoid(logits) * scale`` is returned unchanged. """ if not self.enable_shared_expert_fusion or self.shared_expert_gate is None: return None shared_out = self.shared_expert_gate(hidden_states) shared_logits = shared_out[0] if isinstance(shared_out, tuple) else shared_out # Allreduce-EP path: the fused shared expert occupies a single global # slot loaded onto every EP rank (see FusedMoE.__init__: num_shared_slots # == num_fused_shared_experts when not is_deepep_class_backend()). Every # rank therefore computes the same full shared output, and the # post-experts all_reduce sums it ep_size times. Pre-scale the per-token # routing weight by 1/ep_size to cancel this, mirroring DeepSeek-V2's # fused_shared_experts_scaling_factor pattern. scale = 1.0 moe_ep_size = get_parallel().moe_ep_size if moe_ep_size > 1 and not is_deepep_class_backend(): scale = 1.0 / float(moe_ep_size) # Only AITER fuses sigmoid + cast in-kernel; on CUDA keep the legacy # eager activation so the NVIDIA path behavior is unchanged. if not _use_aiter: return F.sigmoid(shared_logits) * scale, 1.0 return shared_logits, scale def _append_shared_to_topk_output( self, topk_output: StandardTopKOutput, hidden_states: torch.Tensor, ) -> StandardTopKOutput: """Append shared expert ids and weights to topk output before fused MoE.""" if not self.enable_shared_expert_fusion: return topk_output shared = self._get_shared_expert_weights(hidden_states) if shared is None: return topk_output shared_weights, shared_scale = shared from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe_triton_kernels import ( fused_append_shared_experts_with_weights, ) # AITER returns raw logits + scale for in-kernel sigmoid fusion; CUDA # returns pre-activated weights (scale already folded in) → no fusion. fused_topk_ids, fused_topk_weights = fused_append_shared_experts_with_weights( topk_output.topk_ids, topk_output.topk_weights, shared_weights, self.num_fused_shared_experts, N=self.num_experts, apply_sigmoid=_use_aiter, scale=shared_scale, ) return StandardTopKOutput( topk_weights=fused_topk_weights, topk_ids=fused_topk_ids, router_logits=topk_output.router_logits, ) def _forward_shared_experts( self, hidden_states: torch.Tensor, apply_gate: bool = True ): shared_output = None if self.shared_expert is not None: shared_output = self.shared_expert(hidden_states) if self.shared_expert_gate is not None and apply_gate: if use_intel_amx_backend(self.shared_expert_gate): shared_output = torch.ops.sgl_kernel.fused_linear_sigmoid_mul( hidden_states, self.shared_expert_gate.weight, self.shared_expert_gate.bias, True, shared_output, ) elif _is_hip: from sglang.jit_kernel.triton.sigmoid_gate_mul import ( sigmoid_gate_mul_broadcast, ) gate = self.shared_expert_gate(hidden_states) shared_output = sigmoid_gate_mul_broadcast(shared_output, gate) else: shared_output = ( F.sigmoid(self.shared_expert_gate(hidden_states)) * shared_output ) return shared_output def _forward_deepep(self, hidden_states: torch.Tensor, forward_batch: ForwardBatch): enable_dual_stream = ( is_npu() and envs.SGLANG_NPU_USE_MULTI_STREAM.get() and forward_batch.forward_mode.is_cuda_graph() ) shared_output = None if hidden_states.shape[0] > 0: # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) if enable_dual_stream: shared_output = shared_expert_on_independent_stream( hidden_states.clone(), self._forward_shared_experts ) else: shared_output = self._forward_shared_experts(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, ) if not self.is_nextn else None ), ) else: topk_output = self.topk.empty_topk_output(hidden_states.device) final_hidden_states = self.experts( hidden_states=hidden_states, topk_output=topk_output, ) if enable_dual_stream: wait_share_stream() if shared_output is not None: final_hidden_states.add_(shared_output) return final_hidden_states def _forward_router_experts(self, hidden_states: torch.Tensor): # router_logits: (num_tokens, n_experts) router_logits, _ = self.gate(hidden_states) topk_output = self.topk(hidden_states, router_logits) if self.enable_shared_expert_fusion and TopKOutputChecker.format_is_standard( topk_output ): topk_output = self._append_shared_to_topk_output(topk_output, hidden_states) return self.experts(hidden_states, topk_output) def forward_normal_dual_stream( self, hidden_states: torch.Tensor, use_fused_gate: bool = False, ) -> torch.Tensor: current_stream = torch.cuda.current_stream() self.alt_stream.wait_stream(current_stream) shared_output = ( self._forward_shared_experts( hidden_states.clone(), apply_gate=not use_fused_gate ) if self.shared_expert is not None else None ) # ===== TO BE REFACTORED ==== # Shared-add overlap (SGLANG_OPT_LORA_SHARED_ADD_OVERLAP): hand the add to the LoRA # MoE dispatch so it overlaps the down-LoRA shrink on the alt stream. staged = False if shared_output is not None and _SGLANG_EXPERIMENTAL_LORA_OPTI: from sglang.srt.lora.trtllm_lora_temp.shared_add_overlap import ( shared_add_overlap_enabled, stage_shared_expert_add, unstage_shared_expert_add, ) if shared_add_overlap_enabled(): stage_shared_expert_add(shared_output, current_stream) staged = True # ===== END TO BE REFACTORED ==== with torch.cuda.stream(self.alt_stream): router_output = self._forward_router_experts(hidden_states) current_stream.wait_stream(self.alt_stream) if staged and unstage_shared_expert_add() is None: # The dispatch consumed the staging (add already enqueued); skip the caller's add. shared_output = None return router_output, shared_output def forward( self, hidden_states: torch.Tensor, forward_batch: Optional[ForwardBatch] = None, ) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) if get_moe_a2a_backend().is_deepep(): return self._forward_deepep(hidden_states, forward_batch) use_fused_gate = ( self.shared_expert_gate is not None and not use_intel_amx_backend(self.shared_expert_gate) and not is_npu() ) if hidden_states.shape[0] == 0: # M=0 guard for idle DP ranks: skip shared_experts and gate # (which crash on empty tensors in FP4 GEMM), but still call # self.experts() to participate in alltoall collective. shared_output = None topk_output = self.topk.empty_topk_output(hidden_states.device) final_hidden_states = self.experts(hidden_states, topk_output) elif self.alt_stream is not None and get_is_capture_mode(): final_hidden_states, shared_output = self.forward_normal_dual_stream( hidden_states, use_fused_gate=use_fused_gate ) else: shared_output = self._forward_shared_experts( hidden_states, apply_gate=not use_fused_gate ) final_hidden_states = self._forward_router_experts(hidden_states) if shared_output is not None: if use_fused_gate: fused_gate_sigmoid_mul_add( hidden_states, self.shared_expert_gate.weight.squeeze(), shared_output, final_hidden_states, ) else: final_hidden_states += shared_output if ( self.tp_size > 1 and not should_skip_post_experts_all_reduce( is_tp_path=True, ) and not get_moe_a2a_backend().is_flashinfer() ): final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) # Debug removed - was causing issues during CUDA graph capture return final_hidden_states.view(num_tokens, hidden_dim) class Qwen2MoeAttention(nn.Module): def __init__( self, hidden_size: int, num_heads: int, num_kv_heads: int, layer_id: int = 0, rope_theta: float = 10000, rope_scaling: Optional[Dict[str, Any]] = None, max_position_embeddings: int = 8192, qkv_bias: int = True, quant_config: Optional[QuantizationConfig] = None, dual_chunk_attention_config: Optional[dict[str, Any]] = None, prefix: str = "", ) -> None: super().__init__() self.hidden_size = hidden_size attn_tp_rank = get_parallel().attn_tp_rank attn_tp_size = get_parallel().attn_tp_size 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 = 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.qkv_proj = QKVParallelLinear( hidden_size, self.head_dim, self.total_num_heads, self.total_num_kv_heads, bias=qkv_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=False, 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.attn = RadixAttention( self.num_heads, self.head_dim, self.scaling, num_kv_heads=self.num_kv_heads, layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("attn", prefix), ) def forward( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) q, k = self.rotary_emb(positions, q, k) attn_output = self.attn(q, k, v, forward_batch) output, _ = self.o_proj(attn_output) return output class Qwen2MoeDecoderLayer(nn.Module): def __init__( self, config: PretrainedConfig, 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 self.start_layer = start_layer rope_theta, rope_scaling = get_rope_config(config) max_position_embeddings = getattr(config, "max_position_embeddings", 8192) qkv_bias = getattr(config, "qkv_bias", True) dual_chunk_attention_config = getattr( config, "dual_chunk_attention_config", None ) self.self_attn = Qwen2MoeAttention( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, num_kv_heads=config.num_key_value_heads, layer_id=layer_id, rope_theta=rope_theta, rope_scaling=rope_scaling, max_position_embeddings=max_position_embeddings, quant_config=quant_config, dual_chunk_attention_config=dual_chunk_attention_config, qkv_bias=qkv_bias, prefix=add_prefix("self_attn", prefix), ) self.layer_id = layer_id self.attn_tp_size = get_parallel().attn_tp_size self.attn_tp_rank = get_parallel().attn_tp_rank # Qwen2MoE 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 = Qwen2MoeSparseMoeBlock( layer_id=layer_id, config=config, quant_config=quant_config, alt_stream=alt_stream, prefix=add_prefix("mlp", prefix), ) else: self.mlp = Qwen2MoeMLP( 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, ) 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 ) # 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(mlp_reduce_scatter=mlp_reduce_scatter): hidden_states = self.mlp(hidden_states, forward_batch) hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual, forward_batch ) return hidden_states, residual class Qwen2MoeModel(nn.Module): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", decoder_layer_type: type[nn.Module] = Qwen2MoeDecoderLayer, alt_stream: Optional[torch.cuda.Stream] = None, ) -> None: super().__init__() self.config = config self.vocab_size = config.vocab_size self.pp_group = get_pp_group() self.moe_dp_size = get_parallel().moe_dp_size self.attn_cp_size = get_parallel().attn_cp_size if self.pp_group.is_first_rank: self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, use_attn_tp_group=is_dp_attention_enabled(), quant_config=quant_config, prefix=add_prefix("embed_tokens", prefix), ) else: self.embed_tokens = PPMissingLayer() # Use the provided decoder layer type or default to Qwen2MoeDecoderLayer decoder_layer_type = decoder_layer_type or Qwen2MoeDecoderLayer pp_start_layer, _ = get_pp_indices( config.num_hidden_layers, self.pp_group.rank_in_group, self.pp_group.world_size, ) self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers, lambda idx, prefix: decoder_layer_type( layer_id=idx, start_layer=pp_start_layer, config=config, quant_config=quant_config, prefix=prefix, alt_stream=alt_stream, ), pp_rank=self.pp_group.rank_in_group, pp_size=self.pp_group.world_size, prefix=add_prefix("layers", prefix), ) if self.pp_group.is_last_rank: self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) else: self.norm = PPMissingLayer(return_tuple=True) # For EAGLE3 support self.layers_to_capture = [] def set_eagle3_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) def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, pp_proxy_tensors: Optional[PPProxyTensors] = None, ) -> Union[torch.Tensor, PPProxyTensors]: if self.pp_group.is_first_rank: if input_embeds is None: hidden_states = self.embed_tokens(input_ids) else: hidden_states = input_embeds residual = None else: assert pp_proxy_tensors is not None hidden_states = pp_proxy_tensors["hidden_states"] residual = pp_proxy_tensors["residual"] if ( is_prefill_context_parallel_enabled() and not is_cp_v2_active(forward_batch) and forward_batch.forward_mode.is_context_parallel_extend() and forward_batch.attn_cp_metadata is not None ): if self.pp_group.is_first_rank: hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) positions = cp_split_and_rebuild_position(forward_batch, positions) aux_hidden_states = [] if forward_batch.can_run_tbo: hidden_states, residual = model_forward_maybe_tbo( layers=self.layers, enable_tbo=True, input_data_scatter_mode=ScatterMode.model_input_output(), positions=positions, forward_batch=forward_batch, hidden_states=hidden_states, residual=residual, ) else: for i in range(self.start_layer, self.end_layer): ctx = ( nullcontext() if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE) else get_global_expert_distribution_recorder().with_current_layer(i) ) with ctx: layer = self.layers[i] hidden_states, residual = layer( positions, hidden_states, forward_batch, residual, captured_last_layer_outputs=( aux_hidden_states if getattr(layer, "_is_layer_to_capture", False) else None ), ) if not self.pp_group.is_last_rank: if ( hidden_states is not None and hasattr(hidden_states, "_sglang_needs_allreduce_fusion") and hidden_states._sglang_needs_allreduce_fusion ): if get_parallel().moe_ep_size > 1: hidden_states = moe_expert_parallel_all_reduce(hidden_states) if get_parallel().moe_tp_size > 1: hidden_states = moe_tensor_model_parallel_all_reduce(hidden_states) hidden_states._sglang_needs_allreduce_fusion = False return PPProxyTensors( { "hidden_states": hidden_states, "residual": residual, } ) else: if hidden_states.shape[0] != 0: if residual is None: hidden_states = self.norm(hidden_states) else: hidden_states, _ = self.norm(hidden_states, residual) if ( self.pp_group.is_last_rank and not is_cp_v2_active(forward_batch) and is_prefill_context_parallel_enabled() and forward_batch.forward_mode.is_context_parallel_extend() and forward_batch.attn_cp_metadata is not None ): hidden_states = cp_all_gather_rerange_output( hidden_states, self.attn_cp_size, forward_batch, torch.cuda.current_stream(), ) if len(aux_hidden_states) == 0: return hidden_states return hidden_states, aux_hidden_states class Qwen2MoeForCausalLM(nn.Module): fall_back_to_pt_during_load = False def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.pp_group = get_pp_group() self.config = config self.quant_config = quant_config alt_stream = get_stream("alt") if _is_cuda else None self.model = Qwen2MoeModel( config, quant_config, prefix=add_prefix("model", prefix), alt_stream=alt_stream, ) 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) # For EAGLE3 support self.capture_aux_hidden_states = False @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: 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: return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states ) 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 load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): 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, ) params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: 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: for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue name = name.replace(weight_name, param_name) 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: # 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") @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, ) 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]) EntryClass = Qwen2MoeForCausalLM