diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index d9127143a..59fb89794 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/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, diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 14afc0d2f..b40024dc0 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/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.""" diff --git a/sgl-kernel/csrc/cpu/moe.cpp b/sgl-kernel/csrc/cpu/moe.cpp index 8202a08e0..d1740c981 100644 --- a/sgl-kernel/csrc/cpu/moe.cpp +++ b/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* __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( - sorted_ids, expert_ids, topk_ids.data_ptr(), total_cnts, cumsums, offsets, E, numel, num_threads); + sorted_ids, expert_ids, topk_ids_.data_ptr(), 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]