# 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 DeepSeek and Mixtral implementation """Inference-only MiniMax M2 model compatible with HuggingFace weights.""" import logging from contextlib import nullcontext from typing import Iterable, Optional, Set, Tuple, Union import torch import triton import triton.language as tl 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_moe_expert_parallel_world_size, get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_reduce, ) from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo from sglang.srt.layers.communicator import ( LayerCommunicator, LayerScatterModes, ScatterMode, ) 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.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 get_moe_a2a_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 from sglang.srt.layers.vocab_parallel_embedding import ( ParallelLMHead, VocabParallelEmbedding, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_loader.weight_utils import ( default_weight_loader, maybe_remap_kv_scale_name, ) from sglang.srt.server_args import get_global_server_args from sglang.srt.utils import ( BumpAllocator, add_prefix, get_compiler_backend, is_non_idle_and_non_empty, make_layers, ) logger = logging.getLogger(__name__) @triton.jit def rmsnorm_sumsq_kernel_serial( x1_ptr, # T* [B, D] x2_ptr, # T* [B, D] stride_x1, # int stride_x2, # int sum_sq_ptr, # float* [B] B, # int D1, # int D2, # int BLOCK_SIZE1: tl.constexpr, BLOCK_SIZE2: tl.constexpr, ): row_id = tl.program_id(0) x1_row = x1_ptr + row_id * stride_x1 x2_row = x2_ptr + row_id * stride_x2 offsets1 = tl.arange(0, BLOCK_SIZE1) mask1 = offsets1 < D1 offsets2 = tl.arange(0, BLOCK_SIZE2) mask2 = offsets2 < D2 x1 = tl.load(x1_row + offsets1, mask=mask1, other=0.0) x2 = tl.load(x2_row + offsets2, mask=mask2, other=0.0) x1_f32 = x1.to(tl.float32) sum_sq1 = tl.sum(x1_f32 * x1_f32, axis=0) x2_f32 = x2.to(tl.float32) sum_sq2 = tl.sum(x2_f32 * x2_f32, axis=0) tl.store(sum_sq_ptr + row_id, sum_sq1) tl.store(sum_sq_ptr + row_id + B, sum_sq2) @triton.jit def rmsnorm_apply_kernel_serial( x1_ptr, # T* [B, D] x2_ptr, # T* [B, D] w1_ptr, # T* [D] w2_ptr, # T* [D] sum_sq_ptr, # float* [B] out1_ptr, # T* [B, D] out2_ptr, # T* [B, D] B, # int D1, # int D2, # int stride_x1, # int stride_x2, # int tp_world, # int eps, # float BLOCK_SIZE1: tl.constexpr, BLOCK_SIZE2: tl.constexpr, ): row_id = tl.program_id(0) x1_row = x1_ptr + row_id * stride_x1 x2_row = x2_ptr + row_id * stride_x2 out1_row = out1_ptr + row_id * stride_x1 out2_row = out2_ptr + row_id * stride_x2 sum_sq1 = tl.load(sum_sq_ptr + row_id) sum_sq2 = tl.load(sum_sq_ptr + row_id + B) inv_rms1 = tl.rsqrt(sum_sq1 / D1 / tp_world + eps) inv_rms2 = tl.rsqrt(sum_sq2 / D2 / tp_world + eps) offsets1 = tl.arange(0, BLOCK_SIZE1) offsets2 = tl.arange(0, BLOCK_SIZE2) mask1 = offsets1 < D1 mask2 = offsets2 < D2 x1 = tl.load(x1_row + offsets1, mask=mask1, other=0.0) w1 = tl.load(w1_ptr + offsets1, mask=mask1, other=1.0) x2 = tl.load(x2_row + offsets2, mask=mask2, other=0.0) w2 = tl.load(w2_ptr + offsets2, mask=mask2, other=1.0) out1 = (x1.to(tl.float32) * inv_rms1 * w1.to(tl.float32)).to(x1.dtype) out2 = (x2.to(tl.float32) * inv_rms2 * w2.to(tl.float32)).to(x2.dtype) tl.store(out1_row + offsets1, out1, mask=mask1) tl.store(out2_row + offsets2, out2, mask=mask2) def rms_sumsq_serial(x1: torch.Tensor, x2: torch.Tensor) -> torch.Tensor: assert x1.is_cuda and x2.is_cuda B, D1 = x1.shape B2, D2 = x2.shape assert B == B2 stride_x1 = x1.stride(0) stride_x2 = x2.stride(0) # We found that custom all-reduce `sglang::cross_device_reduce_1stage` # is much faster than the nccl all-reduce in torch. # However, `should_custom_ar` checks if the reduced buffer is 16-byte aligned. # RMSNormTP reduces a [B, 2] fp32 tensor, so we pad the total element count to # satisfy the alignment requirement. B_padded = (B + B2 + 3) // 4 * 4 sum_sq = torch.empty(B_padded, device=x1.device, dtype=torch.float32) BLOCK_SIZE1 = triton.next_power_of_2(D1) BLOCK_SIZE2 = triton.next_power_of_2(D2) grid = (B,) rmsnorm_sumsq_kernel_serial[grid]( x1, x2, stride_x1, stride_x2, sum_sq, B, D1, D2, BLOCK_SIZE1, BLOCK_SIZE2, ) return sum_sq def rms_apply_serial( x1: torch.Tensor, x2: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, sum_sq: torch.Tensor, tp_world: int = 1, eps: float = 1e-5, ) -> torch.Tensor: assert x1.is_cuda and x2.is_cuda and w1.is_cuda and w2.is_cuda and sum_sq.is_cuda B, D1 = x1.shape B2, D2 = x2.shape assert B == B2 stride_x1 = x1.stride(0) stride_x2 = x2.stride(0) out1 = torch.empty(B, D1, device=x1.device, dtype=x1.dtype) out2 = torch.empty(B, D2, device=x2.device, dtype=x2.dtype) BLOCK_SIZE1 = triton.next_power_of_2(D1) BLOCK_SIZE2 = triton.next_power_of_2(D2) grid = (B,) rmsnorm_apply_kernel_serial[grid]( x1, x2, w1, w2, sum_sq, out1, out2, B, D1, D2, stride_x1, stride_x2, tp_world, eps, BLOCK_SIZE1, BLOCK_SIZE2, ) return out1, out2 class MiniMaxM2RMSNormTP(nn.Module): """RMSNorm with Tensor Parallel support for QK normalization.""" def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: super().__init__() self.tp_world = get_tensor_model_parallel_world_size() self.tp_rank = get_tensor_model_parallel_rank() # Weight parameter is sharded across TP ranks self.weight = nn.Parameter(torch.ones(int(hidden_size / self.tp_world))) self.weight.weight_loader = self.weight_loader self.variance_epsilon = eps @staticmethod def weight_loader( param: nn.Parameter, loaded_weight: torch.Tensor, ) -> None: """Custom weight loader that handles TP sharding.""" tp_world = get_tensor_model_parallel_world_size() tp_rank = get_tensor_model_parallel_rank() shard_size = loaded_weight.shape[0] // tp_world shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size) param.data.copy_(loaded_weight[shard]) @torch.compile(dynamic=True, backend=get_compiler_backend()) def forward( self, x: torch.Tensor, residual: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: """Forward pass with TP-aware variance computation.""" assert residual is None, "RMSNormTP does not support residual connection." orig_dtype = x.dtype x = x.to(torch.float32) # Compute variance across the full dimension (not just local shard) variance = x.pow(2).mean(dim=-1, keepdim=True, dtype=torch.float32) if self.tp_world > 1: # All-reduce variance across TP ranks to get global variance variance = tensor_model_parallel_all_reduce(variance) / self.tp_world # Normalize and apply local weight shard x = x * torch.rsqrt(variance + self.variance_epsilon) x = (x * self.weight).to(orig_dtype) return x @staticmethod def forward_qk( q_norm: "MiniMaxM2RMSNormTP", k_norm: "MiniMaxM2RMSNormTP", q: torch.Tensor, k: torch.Tensor, ) -> torch.Tensor: sum_sq = rms_sumsq_serial(q, k) if q_norm.tp_world > 1: sum_sq = tensor_model_parallel_all_reduce(sum_sq) q, k = rms_apply_serial( q, k, q_norm.weight, k_norm.weight, sum_sq, q_norm.tp_world, q_norm.variance_epsilon, ) return q, k class MiniMaxM2MoE(nn.Module): """MiniMax MoE implementation using DeepEP for Expert Parallel support.""" def __init__( self, config: PretrainedConfig, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ): super().__init__() self.tp_size = get_tensor_model_parallel_world_size() if self.tp_size > config.num_local_experts: raise ValueError( f"Tensor parallel size {self.tp_size} is greater than " f"the number of experts {config.num_local_experts}." ) self.use_routing_bias = getattr(config, "use_routing_bias", False) if self.use_routing_bias: self.e_score_correction_bias = nn.Parameter( torch.empty(config.num_local_experts, dtype=torch.float32) ) self.e_score_correction_bias.weight_loader = ( MiniMaxM2MoE.ebias_weight_loader ) else: self.e_score_correction_bias = None self.experts = get_moe_impl_class(quant_config)( num_experts=config.num_local_experts + get_global_server_args().ep_num_redundant_experts, top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.intermediate_size, layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("experts", prefix), ) self.topk = TopK( top_k=config.num_experts_per_tok, renormalize=True, scoring_func=config.scoring_func, correction_bias=self.e_score_correction_bias, routed_scaling_factor=1.0, ) self.gate = ReplicatedLinear( config.hidden_size, config.num_local_experts, bias=False, params_dtype=torch.float32, quant_config=None, prefix=add_prefix("gate", prefix), ) self.layer_id = layer_id if get_moe_a2a_backend().is_deepep(): self.ep_size = get_moe_expert_parallel_world_size() self.top_k = config.num_experts_per_tok @staticmethod def ebias_weight_loader(param: nn.Parameter, loaded_weight: torch.Tensor) -> None: assert param.size() == loaded_weight.size() param.data.copy_(loaded_weight.to(torch.float32)) def forward( self, hidden_states: torch.Tensor, forward_batch: ForwardBatch ) -> torch.Tensor: if get_moe_a2a_backend().is_deepep(): return self.forward_deepep(hidden_states, forward_batch) else: return self.forward_normal(hidden_states) 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.to(torch.float32)) topk_output = self.topk(hidden_states, router_logits) final_hidden_states = self.experts(hidden_states, topk_output) if self.tp_size > 1: final_hidden_states = 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.to(torch.float32)) 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(device=hidden_states.device) final_hidden_states = self.experts( hidden_states=hidden_states, topk_output=topk_output, ) return final_hidden_states # TBO Operations for MiniMax MoE def op_gate(self, state): """Gate operation for TBO - compute router logits""" if is_non_idle_and_non_empty( state.forward_batch.forward_mode, state.hidden_states_mlp_input ): # router_logits: (num_tokens, num_experts) state.router_logits, _ = self.gate(state.hidden_states_mlp_input) else: state.router_logits = None def op_select_experts(self, state): """Expert selection operation for TBO""" router_logits = state.pop("router_logits") hidden_states = state.hidden_states_mlp_input if router_logits is not None: ctx = ( nullcontext() if get_global_server_args().enable_piecewise_cuda_graph else get_global_expert_distribution_recorder().with_current_layer( self.layer_id ) ) with ctx: state.topk_weights_local, state.topk_idx_local, _ = 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_idx_local = torch.full( (0, self.top_k), -1, dtype=torch.int, device=hidden_states.device ) state.topk_weights_local = torch.empty( (0, self.top_k), dtype=torch.float32, device=hidden_states.device ) def op_dispatch_a(self, state): """Dispatch A operation for TBO - start async dispatch""" if self.ep_size > 1: self.experts.deepep_dispatcher.dispatch_a( hidden_states=state.pop("hidden_states_mlp_input"), topk_idx=state.pop("topk_idx_local"), topk_weights=state.pop("topk_weights_local"), forward_batch=state.forward_batch, tbo_subbatch_index=state.get("tbo_subbatch_index"), ) def op_dispatch_b(self, state): """Dispatch B operation for TBO - complete async dispatch""" if self.ep_size > 1: ctx = ( nullcontext() if get_global_server_args().enable_piecewise_cuda_graph else get_global_expert_distribution_recorder().with_current_layer( self.layer_id ) ) with ctx: state.dispatch_output = self.experts.deepep_dispatcher.dispatch_b( tbo_subbatch_index=state.get("tbo_subbatch_index"), ) def op_experts(self, state): """Expert computation for TBO""" state.hidden_states_experts_output = self.experts.moe_impl( dispatch_output=state.dispatch_output, ) def op_combine_a(self, state): """Combine A operation for TBO - start async combine""" if self.ep_size > 1: self.experts.deepep_dispatcher.combine_a( hidden_states=state.pop("hidden_states_experts_output"), topk_idx=state.dispatch_output.topk_idx, topk_weights=state.dispatch_output.topk_weights, forward_batch=state.forward_batch, tbo_subbatch_index=state.get("tbo_subbatch_index"), ) state.pop("dispatch_output") def op_combine_b(self, state): """Combine B operation for TBO - complete async combine""" if self.ep_size > 1: state.hidden_states_after_combine = ( self.experts.deepep_dispatcher.combine_b( tbo_subbatch_index=state.get("tbo_subbatch_index"), ) ) def op_output(self, state): """Output operation for TBO - final MLP output""" final_hidden_states = state.pop("hidden_states_after_combine") # MiniMax doesn't have shared experts like DeepSeek, so no need to add them state.hidden_states_mlp_output = final_hidden_states class MiniMaxM2Attention(nn.Module): """MiniMax Attention implementation with QK normalization and partial RoPE.""" def __init__( self, config: PretrainedConfig, layer_id: int = 0, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.hidden_size = config.hidden_size tp_size = get_tensor_model_parallel_world_size() # Get dimensions from config self.total_num_heads = config.num_attention_heads assert self.total_num_heads % tp_size == 0 self.num_heads = self.total_num_heads // tp_size self.total_num_kv_heads = config.num_key_value_heads if self.total_num_kv_heads >= 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 % 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 tp_size % self.total_num_kv_heads == 0 self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) # Use head_dim from config if available, otherwise calculate self.head_dim = getattr( config, "head_dim", self.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 # RoPE settings - support partial RoPE self.rope_theta = getattr(config, "rope_theta", 10000) self.max_position_embeddings = getattr(config, "max_position_embeddings", 8192) self.rotary_dim = getattr( config, "rotary_dim", self.head_dim ) # MiniMax uses rotary_dim=64 # QK Normalization settings self.use_qk_norm = getattr(config, "use_qk_norm", False) self.qk_norm_type = getattr(config, "qk_norm_type", "per_layer") self.qkv_proj = QKVParallelLinear( self.hidden_size, self.head_dim, self.total_num_heads, self.total_num_kv_heads, bias=False, quant_config=quant_config, prefix=add_prefix("qkv_proj", prefix), ) self.o_proj = RowParallelLinear( self.total_num_heads * self.head_dim, self.hidden_size, bias=False, reduce_results=False, quant_config=quant_config, prefix=add_prefix("o_proj", prefix), ) # Setup RoPE with partial rotary dimension rope_scaling = getattr(config, "rope_scaling", None) self.rotary_emb = get_rope( self.head_dim, rotary_dim=self.rotary_dim, # Use partial rotary dimension max_position=self.max_position_embeddings, base=self.rope_theta, rope_scaling=rope_scaling, ) # QK Normalization layers if self.use_qk_norm: if self.qk_norm_type == "per_layer": # Use RMSNormTP for proper tensor parallel support # Use total dimensions (before TP sharding) for correct normalization self.q_norm = MiniMaxM2RMSNormTP( self.total_num_heads * self.head_dim, eps=config.rms_norm_eps ) self.k_norm = MiniMaxM2RMSNormTP( self.total_num_kv_heads * self.head_dim, eps=config.rms_norm_eps ) else: raise ValueError(f"Unsupported qk_norm_type: {self.qk_norm_type}") 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_prepare( self, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, ): qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) if self.use_qk_norm: # q = self.q_norm(q.contiguous()) # k = self.k_norm(k.contiguous()) q, k = MiniMaxM2RMSNormTP.forward_qk( self.q_norm, self.k_norm, q.contiguous(), k.contiguous() ) else: q, k = q.contiguous(), k.contiguous() q, k = self.rotary_emb(positions, q, k) inner_state = q, k, v, forward_batch return None, forward_batch, inner_state def forward_core(self, intermediate_state): _, _, inner_state = intermediate_state attn_output = self.attn(*inner_state) 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) 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") ) class MiniMaxM2DecoderLayer(nn.Module): """MiniMax Decoder Layer implementation with MoE support.""" def __init__( self, config: PretrainedConfig, layer_id: int, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.hidden_size = config.hidden_size self.layer_id = layer_id # TBO support: All MiniMax layers are sparse (MoE) self.is_layer_sparse = True self.self_attn = MiniMaxM2Attention( config=config, layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("self_attn", prefix), ) self.block_sparse_moe = MiniMaxM2MoE( config=config, layer_id=layer_id, quant_config=quant_config, prefix=add_prefix("mlp", prefix), ) self.input_layernorm = RMSNorm( config.hidden_size, eps=getattr(config, "rms_norm_eps", 1e-6) ) self.post_attention_layernorm = RMSNorm( config.hidden_size, eps=getattr(config, "rms_norm_eps", 1e-6) ) 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, ) 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], ) -> torch.Tensor: # Self Attention hidden_states, residual = self.layer_communicator.prepare_attn( hidden_states, residual, forward_batch ) hidden_states = self.self_attn( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, ) # Fully Connected (MLP or MoE) hidden_states, residual = self.layer_communicator.prepare_mlp( hidden_states, residual, forward_batch ) hidden_states = self.block_sparse_moe(hidden_states, forward_batch) hidden_states, residual = self.layer_communicator.postprocess_layer( hidden_states, residual, forward_batch ) return hidden_states, residual # TBO Operations for MiniMax Decoder Layer def op_comm_prepare_attn( self, state, positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, residual: Optional[torch.Tensor], zero_allocator: BumpAllocator, tbo_subbatch_index: Optional[int] = None, ): """Communication prepare for attention - TBO operation""" 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, zero_allocator=zero_allocator, tbo_subbatch_index=tbo_subbatch_index, ) ) def op_comm_prepare_mlp(self, state): """Communication prepare for MLP - TBO operation""" 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_mlp(self, state): hidden_states = state.pop("hidden_states_mlp_input") state.hidden_states_mlp_output = self.block_sparse_moe( hidden_states, state.forward_batch ) def op_comm_postprocess_layer(self, state): """Communication postprocess for layer - TBO operation""" 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, zero_allocator=state.zero_allocator, tbo_subbatch_index=state.tbo_subbatch_index, ) return output class MiniMaxM2Model(nn.Module): """MiniMax Model implementation.""" fall_back_to_pt_during_load = False def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.padding_idx = getattr(config, "pad_token_id", 0) self.vocab_size = config.vocab_size self.pp_group = get_pp_group() self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, ) def layer_fn(idx, prefix: str) -> nn.Module: return MiniMaxM2DecoderLayer( config=config, layer_id=idx, quant_config=quant_config, prefix=prefix, ) self.layers, self.start_layer, self.end_layer = make_layers( config.num_hidden_layers, layer_fn, 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 get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) 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, Tuple[torch.Tensor, list[torch.Tensor]]]: if self.pp_group.is_first_rank: if input_embeds is None: hidden_states = self.get_input_embeddings(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"] 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 get_global_server_args().enable_piecewise_cuda_graph else get_global_expert_distribution_recorder().with_current_layer(i) ) with ctx: if i in self.layers_to_capture: aux_hidden_states.append(hidden_states + residual) layer = self.layers[i] hidden_states, residual = layer( positions=positions, forward_batch=forward_batch, hidden_states=hidden_states, residual=residual, ) if not self.pp_group.is_last_rank: return PPProxyTensors( {"hidden_states": hidden_states, "residual": residual} ) if residual is not None: hidden_states, _ = self.norm(hidden_states, residual) else: hidden_states = self.norm(hidden_states) if len(aux_hidden_states) == 0: return hidden_states return hidden_states, aux_hidden_states class MiniMaxM2ForCausalLM(nn.Module): """MiniMax M2 model for causal language modeling.""" def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, prefix: str = "", ) -> None: super().__init__() self.config = config self.quant_config = quant_config self.model = MiniMaxM2Model( config, quant_config, prefix=add_prefix("model", prefix) ) if get_pp_group().is_last_rank: self.lm_head = ParallelLMHead( config.vocab_size, config.hidden_size, quant_config=None, prefix=add_prefix("lm_head", prefix), ) else: self.lm_head = PPMissingLayer() self.logits_processor = LogitsProcessor(config) # For EAGLE3 self.capture_aux_hidden_states = False def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: return self.model.get_input_embeddings(input_ids) def set_eagle3_layers_to_capture(self, layer_ids: Optional[list[int]] = None): if not get_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.layers_to_capture = [ 2, num_layers // 2, num_layers - 3, ] # Specific layers for EAGLE3 support else: self.model.layers_to_capture = [val + 1 for val in layer_ids] def get_embed_and_head(self): return self.model.embed_tokens.weight, self.lm_head.weight @torch.no_grad() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, ) -> torch.Tensor: # _print_tensor_info(input_ids, "input_ids") hidden_states = self.model(input_ids, positions, forward_batch, input_embeds) aux_hidden_states = None if self.capture_aux_hidden_states: hidden_states, aux_hidden_states = hidden_states return self.logits_processor( input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states ) def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): """Load model weights with proper mapping for MiniMax architecture.""" 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), ] # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) expert_params_mapping = FusedMoE.make_expert_params_mapping( ckpt_gate_proj_name="w1", ckpt_down_proj_name="w2", ckpt_up_proj_name="w3", num_experts=self.config.num_local_experts, ) params_dict = dict(self.named_parameters()) loaded_params: Set[str] = set() for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue spec_layer = get_spec_layer_idx_from_weight_name(self.config, name) if spec_layer is not None: continue # skip spec decode layers for main model 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) and name not in params_dict: 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 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 # Remapping the name of FP8 kv-scale. name = maybe_remap_kv_scale_name(name, params_dict) if name is None: continue param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) loaded_params.add(name) return loaded_params @classmethod def get_model_config_for_expert_location(cls, config): from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation return ModelConfigForExpertLocation( num_layers=config.num_hidden_layers, num_logical_experts=config.num_local_experts, num_groups=None, ) def get_spec_layer_idx_from_weight_name( config: PretrainedConfig, weight_name: str ) -> Optional[int]: if hasattr(config, "num_mtp_modules") and (config.num_mtp_modules > 0): layer_idx = config.num_hidden_layers for i in range(config.num_mtp_modules): if weight_name.startswith(f"model.layers.{layer_idx + i}."): return layer_idx + i return None # Entry class for model registration EntryClass = MiniMaxM2ForCausalLM