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."""

sgl-kernel/csrc/cpu/moe.cpp

@@ -861,10 +861,14 @@ at::Tensor fused_experts_cpu(
   constexpr int64_t BLOCK_N = block_size_n();
 
   const auto st = hidden_states.scalar_type();
+  // TODO: fused_topk_torch_native (CPU fallback for models like MiniMax)
+  // returns int64 topk_ids; fused_experts_cpu requires int32. Remove the typecast after topk kernel is provided
+  auto topk_ids_ = topk_ids.scalar_type() == at::kInt ? topk_ids : topk_ids.to(at::kInt);
+
   CHECK_INPUT(hidden_states);
   CHECK_INPUT(w1);
   CHECK_INPUT(w2);
-  CHECK_EQ(topk_weights.sizes(), topk_ids.sizes());
+  CHECK_EQ(topk_weights.sizes(), topk_ids_.sizes());
   CHECK_DIM(2, hidden_states);
   if (moe_comp_method == CPUQuantMethod::INT4_W4A8 && is_vnni) {
     CHECK_DIM(4, w1);
@@ -874,9 +878,8 @@ at::Tensor fused_experts_cpu(
     CHECK_DIM(3, w2);
   }
   CHECK_DIM(2, topk_weights);
-  CHECK_DIM(2, topk_ids);
-
-  CHECK_EQ(topk_ids.scalar_type(), at::kInt);
+  CHECK_DIM(2, topk_ids_);
+  CHECK_EQ(topk_ids_.scalar_type(), at::kInt);
 
   // TODO: support topk_weights to be bf16 or fp16 in the kernel.
   // The topk_weights of llama4 is computed via Llama4MoE:custom_routing_function and is bf16/fp16
@@ -924,7 +927,7 @@ at::Tensor fused_experts_cpu(
   int64_t max_num_blocks = div_up(max_num_tokens_padded, BLOCK_M);
   auto buffer = at::empty(
       {max_num_tokens_padded + max_num_blocks + (num_threads + 1) * E + (E + 1) + (max_num_blocks + 1)},
-      topk_ids.options());
+      topk_ids_.options());
 
   int32_t* __restrict__ sorted_ids = buffer.data_ptr<int32_t>();
   int32_t* __restrict__ expert_ids = sorted_ids + max_num_tokens_padded;
@@ -948,7 +951,7 @@ at::Tensor fused_experts_cpu(
 
   // align experts index
   int64_t num_tokens_post_pad = moe_align_block_size<BLOCK_M>(
-      sorted_ids, expert_ids, topk_ids.data_ptr<int32_t>(), total_cnts, cumsums, offsets, E, numel, num_threads);
+      sorted_ids, expert_ids, topk_ids_.data_ptr<int32_t>(), total_cnts, cumsums, offsets, E, numel, num_threads);
 
   // unlike triton kernel, we fuse silu with gemm1 so only need 2 intermediate_caches:
   //   1. intermediate_cache1 : [M * topk, N]