[NPU] Support dp-attention for MiniMax2.5 (#20919)

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(

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