Fix MiniMax-M2.7 on CPU (#25061)

python/sglang/srt/layers/moe/topk.py

@@ -610,6 +610,20 @@ def fused_topk_cpu(
     correction_bias: torch.Tensor = None,
     scoring_func: str = "softmax",
 ):
+    # TODO: add c++ kernel for cpu
+    # The topk_softmax_cpu kernel only handles vanilla softmax scoring with no
+    # correction bias. Fall back to the torch-native impl for the rest
+    # (e.g. MiniMax sets both correction_bias and scoring_func).
+    if correction_bias is not None or scoring_func != "softmax":
+        return fused_topk_torch_native(
+            hidden_states,
+            gating_output,
+            topk,
+            renormalize,
+            correction_bias=correction_bias,
+            scoring_func=scoring_func,
+        )
+
     topk_weights, topk_ids = torch.ops.sgl_kernel.topk_softmax_cpu(
         hidden_states=hidden_states,
         gating_output=gating_output,

python/sglang/srt/models/minimax_m2.py

@@ -78,6 +78,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTe
 from sglang.srt.model_loader.weight_utils import (
     default_weight_loader,
     maybe_remap_kv_scale_name,
+    narrow_padded_param_and_loaded_weight,
 )
 from sglang.srt.server_args import get_global_server_args
 
@@ -89,8 +90,10 @@ from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import (
     BumpAllocator,
     add_prefix,
+    cpu_has_amx_support,
     get_bool_env_var,
     get_compiler_backend,
+    is_cpu,
     is_cuda,
     is_non_idle_and_non_empty,
     is_npu,
@@ -100,6 +103,8 @@ from sglang.srt.utils.custom_op import register_custom_op
 from sglang.srt.utils.hf_transformers_utils import get_rope_config
 
 logger = logging.getLogger(__name__)
+_is_cpu = is_cpu()
+_is_amx_available = cpu_has_amx_support()
 _is_cuda = is_cuda()
 _is_npu = is_npu()
 
@@ -311,6 +316,20 @@ class MiniMaxM2RMSNormTP(nn.Module):
         """Custom weight loader that handles TP sharding."""
         shard_id = self.attn_tp_rank // self.num_head_replicas
         shard_size = param.data.shape[0]
+
+        if _is_cpu and _is_amx_available:
+            # Handle uneven TP sharding on CPU
+            param_data, loaded_weight = narrow_padded_param_and_loaded_weight(
+                param.data,
+                loaded_weight,
+                0,  # param_data_start
+                shard_id * shard_size,  # weight_start
+                0,  # shard_axis
+                shard_size,
+            )
+            param_data.copy_(loaded_weight)
+            return
+
         shard_end = (shard_id + 1) * shard_size
         assert shard_end <= loaded_weight.shape[0], (
             f"Weight shard out of bounds: shard [{shard_id * shard_size}:{shard_end}] "
@@ -394,6 +413,8 @@ class MiniMaxM2QKRMSNorm:
             if counter is not None:
                 self._counter = counter
                 self._forward_impl = self._forward_fused
+        elif _is_cpu and _is_amx_available:
+            self._forward_impl = self._forward_cpu
 
     @lru_cache
     @staticmethod
@@ -455,6 +476,12 @@ class MiniMaxM2QKRMSNorm:
         )
         return q, k
 
+    def _forward_cpu(self, q: torch.Tensor, k: torch.Tensor):
+        # TODO: add c++ kernel for cpu
+        q = self._q_norm(q.contiguous())
+        k = self._k_norm(k.contiguous())
+        return q, k
+
 
 class MiniMaxM2MoE(nn.Module):
     """MiniMax MoE implementation using DeepEP for Expert Parallel support."""