diff --git a/python/sglang/srt/hardware_backend/npu/moe/topk.py b/python/sglang/srt/hardware_backend/npu/moe/topk.py index 6447d9a67..3ceb0eaad 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/topk.py +++ b/python/sglang/srt/hardware_backend/npu/moe/topk.py @@ -26,6 +26,7 @@ def fused_topk_npu( renormalize = topk_config.renormalize correction_bias = topk_config.correction_bias + # Fast path: simple top-k without grouped routing and bias if not use_grouped_topk and correction_bias is None: topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k_softmax( router_logits, @@ -40,8 +41,8 @@ def fused_topk_npu( ) topk_weights = topk_weights.to(torch.float32) + # Grouped top-k with correction bias elif use_grouped_topk and correction_bias is not None: - # Force set routed_scaling_factor = 1 to optimize renormalize topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k( router_logits.to(torch.float32), k=topk_config.top_k, @@ -57,6 +58,26 @@ def fused_topk_npu( eps=float(1e-20), ) + # npu_moe_gating_top_k is not yet supported custom_routing_function + # torch native is not yet supported num_token_non_padded + elif ( + topk_config.custom_routing_function is None + and num_token_non_padded is not None + and correction_bias is not None + ): + topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k( + router_logits.to(torch.float32), + k=topk_config.top_k, + bias=correction_bias.to(torch.float32), + renorm=0, + norm_type=1, + routed_scaling_factor=( + 1 if renormalize else topk_config.routed_scaling_factor + ), + eps=float(1e-20), + ) + + # Fallback to torch native implementation else: topk_config.torch_native = True return select_experts( diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index e5ef2d75c..d0d2d6c76 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -30,7 +30,6 @@ 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, ) @@ -41,6 +40,12 @@ from sglang.srt.layers.communicator import ( LayerScatterModes, ScatterMode, ) +from sglang.srt.layers.dp_attention import ( + attn_tp_all_reduce, + get_attention_tp_rank, + get_attention_tp_size, + is_dp_attention_enabled, +) from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ( QKVParallelLinear, @@ -250,11 +255,11 @@ class MiniMaxM2RMSNormTP(nn.Module): 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() + self.attn_tp_size = get_attention_tp_size() + self.attn_tp_rank = get_attention_tp_rank() # Weight parameter is sharded across TP ranks - self.weight = nn.Parameter(torch.ones(int(hidden_size / self.tp_world))) + self.weight = nn.Parameter(torch.ones(int(hidden_size / self.attn_tp_size))) self.weight.weight_loader = self.weight_loader self.variance_epsilon = eps @@ -264,11 +269,11 @@ class MiniMaxM2RMSNormTP(nn.Module): 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() + attn_tp_size = get_attention_tp_size() + attn_tp_rank = get_attention_tp_rank() - shard_size = loaded_weight.shape[0] // tp_world - shard = slice(tp_rank * shard_size, (tp_rank + 1) * shard_size) + shard_size = loaded_weight.shape[0] // attn_tp_size + shard = slice(attn_tp_rank * shard_size, (attn_tp_rank + 1) * shard_size) param.data.copy_(loaded_weight[shard]) @torch.compile(dynamic=True, backend=get_compiler_backend()) @@ -286,9 +291,9 @@ class MiniMaxM2RMSNormTP(nn.Module): # 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: + if self.attn_tp_size > 1: # All-reduce variance across TP ranks to get global variance - variance = tensor_model_parallel_all_reduce(variance) / self.tp_world + variance = attn_tp_all_reduce(variance) / self.attn_tp_size # Normalize and apply local weight shard x = x * torch.rsqrt(variance + self.variance_epsilon) @@ -304,8 +309,8 @@ class MiniMaxM2RMSNormTP(nn.Module): 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) + if q_norm.attn_tp_size > 1: + sum_sq = attn_tp_all_reduce(sum_sq) q, k = rms_apply_serial( q, @@ -313,7 +318,7 @@ class MiniMaxM2RMSNormTP(nn.Module): q_norm.weight, k_norm.weight, sum_sq, - q_norm.tp_world, + q_norm.attn_tp_size, q_norm.variance_epsilon, ) @@ -387,14 +392,28 @@ class MiniMaxM2MoE(nn.Module): param.data.copy_(loaded_weight.to(torch.float32)) def forward( - self, hidden_states: torch.Tensor, forward_batch: ForwardBatch + self, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, ) -> torch.Tensor: - if get_moe_a2a_backend().is_deepep(): - return self.forward_deepep(hidden_states, forward_batch) + if ( + not get_moe_a2a_backend().is_deepep() + and not get_moe_a2a_backend().is_ascend_fuseep() + ): + return self.forward_normal( + hidden_states, should_allreduce_fusion, use_reduce_scatter + ) else: - return self.forward_normal(hidden_states) + return self.forward_deepep(hidden_states, forward_batch) - def forward_normal(self, hidden_states: torch.Tensor) -> torch.Tensor: + def forward_normal( + self, + hidden_states: torch.Tensor, + should_allreduce_fusion: bool = False, + use_reduce_scatter: bool = False, + ) -> torch.Tensor: num_tokens, hidden_dim = hidden_states.shape hidden_states = hidden_states.view(-1, hidden_dim) @@ -403,7 +422,7 @@ class MiniMaxM2MoE(nn.Module): topk_output = self.topk(hidden_states, router_logits) final_hidden_states = self.experts(hidden_states, topk_output) - if self.tp_size > 1: + if self.tp_size > 1 and not should_allreduce_fusion and not use_reduce_scatter: final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) return final_hidden_states.view(num_tokens, hidden_dim) @@ -543,23 +562,26 @@ class MiniMaxM2Attention(nn.Module): ) -> None: super().__init__() self.hidden_size = config.hidden_size - tp_size = get_tensor_model_parallel_world_size() + + # Use attention TP rank/size for dp-attention support + attn_tp_rank = get_attention_tp_rank() + attn_tp_size = get_attention_tp_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 + assert self.total_num_heads % attn_tp_size == 0 + self.num_heads = self.total_num_heads // attn_tp_size self.total_num_kv_heads = config.num_key_value_heads - if self.total_num_kv_heads >= tp_size: + 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 % tp_size == 0 + 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 tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + assert attn_tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) # Use head_dim from config if available, otherwise calculate self.head_dim = getattr( @@ -588,6 +610,8 @@ class MiniMaxM2Attention(nn.Module): self.total_num_kv_heads, bias=False, quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, prefix=add_prefix("qkv_proj", prefix), ) @@ -597,6 +621,8 @@ class MiniMaxM2Attention(nn.Module): bias=False, reduce_results=False, quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, prefix=add_prefix("o_proj", prefix), ) @@ -751,12 +777,12 @@ class MiniMaxM2DecoderLayer(nn.Module): 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, - ) + if not forward_batch.forward_mode.is_idle(): + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + ) # Fully Connected (MLP or MoE) @@ -764,12 +790,27 @@ class MiniMaxM2DecoderLayer(nn.Module): 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 + should_allreduce_fusion = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) ) + use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) + + hidden_states = self.block_sparse_moe( + hidden_states, forward_batch, should_allreduce_fusion, use_reduce_scatter + ) + + if should_allreduce_fusion: + 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 # TBO Operations for MiniMax Decoder Layer @@ -851,6 +892,7 @@ class MiniMaxM2Model(nn.Module): self.embed_tokens = VocabParallelEmbedding( config.vocab_size, config.hidden_size, + use_attn_tp_group=is_dp_attention_enabled(), ) def layer_fn(idx, prefix: str) -> nn.Module: @@ -932,10 +974,11 @@ class MiniMaxM2Model(nn.Module): {"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 hidden_states.shape[0] != 0: + 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