From 6e6009fb6be70fe0dab76330f346bfd71309fe93 Mon Sep 17 00:00:00 2001
From: jianan-gu <jianan.gu@intel.com>
Date: Fri, 7 Nov 2025 03:24:48 +0800
Subject: [PATCH] [CPU] Fix TP padding case with weight block size (#8243)

---
 python/sglang/srt/configs/update_config.py    | 28 +++++++++++++++----
 .../sglang/srt/layers/quantization/unquant.py |  5 +---
 sgl-kernel/csrc/cpu/topk.cpp                  |  3 ++
 3 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/python/sglang/srt/configs/update_config.py b/python/sglang/srt/configs/update_config.py
index abbd724fb..46d156d7b 100644
--- a/python/sglang/srt/configs/update_config.py
+++ b/python/sglang/srt/configs/update_config.py
@@ -41,10 +41,16 @@ def get_moe_padding_size(weight_block_size):
     return DEFAULT_MOE_PADDING_SIZE
 
 
-def get_num_heads_padding_size(tp_size, weight_block_size):
-    pad_size = (
-        tp_size * 2 if tp_size % 2 == 1 and weight_block_size is not None else tp_size
-    )
+def get_num_heads_padding_size(tp_size, weight_block_size, head_dim):
+    pad_size = tp_size
+
+    if weight_block_size is not None and head_dim % weight_block_size[0] != 0:
+        import math
+
+        pad_size = tp_size * (
+            math.lcm(head_dim, weight_block_size[0]) // weight_block_size[0]
+        )
+
     return pad_size
 
 
@@ -100,6 +106,13 @@ def adjust_config_with_unaligned_cpu_tp(
             model_config.hf_config.head_dim = (
                 model_config.hidden_size // model_config.num_attention_heads
             )
+        if hasattr(model_config.hf_config, "qk_nope_head_dim") and hasattr(
+            model_config.hf_config, "qk_rope_head_dim"
+        ):
+            model_config.hf_config.qk_head_dim = (
+                model_config.hf_config.qk_nope_head_dim
+                + model_config.hf_config.qk_rope_head_dim
+            )
 
         query_heads_per_kv = (
             model_config.num_attention_heads // model_config.get_total_num_kv_heads()
@@ -107,7 +120,12 @@ def adjust_config_with_unaligned_cpu_tp(
         total_kv_heads = model_config.get_total_num_kv_heads()
         from sglang.srt.layers.vocab_parallel_embedding import pad_vocab_size
 
-        pad_size = get_num_heads_padding_size(tp_size, weight_block_size)
+        head_dim = (
+            model_config.hf_config.qk_head_dim
+            if hasattr(model_config.hf_config, "qk_head_dim")
+            else model_config.hf_config.head_dim
+        )
+        pad_size = get_num_heads_padding_size(tp_size, weight_block_size, head_dim)
         num_key_value_heads = pad_vocab_size(total_kv_heads, pad_size)
 
         model_config.num_key_value_heads = num_key_value_heads
diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py
index 55c5b8e3f..26a83c560 100644
--- a/python/sglang/srt/layers/quantization/unquant.py
+++ b/python/sglang/srt/layers/quantization/unquant.py
@@ -316,10 +316,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
             moe_runner_config.activation == "silu"
         ), f"activation = {moe_runner_config.activation} is not supported."
 
-        if (
-            use_intel_amx_backend(layer)
-            and not moe_runner_config.apply_router_weight_on_input
-        ):
+        if use_intel_amx_backend(layer):
             from sglang.srt.layers.moe.topk import apply_topk_weights_cpu
 
             topk_weights, topk_ids, _ = topk_output
diff --git a/sgl-kernel/csrc/cpu/topk.cpp b/sgl-kernel/csrc/cpu/topk.cpp
index abc5a34fa..b4bdcd0b7 100644
--- a/sgl-kernel/csrc/cpu/topk.cpp
+++ b/sgl-kernel/csrc/cpu/topk.cpp
@@ -654,6 +654,9 @@ std::tuple<at::Tensor, at::Tensor> biased_grouped_topk_cpu(
       case 256:
         LAUNCH_BIASED_GROUPED_TOPK_KERNEL(256, 8);
         break;
+      case 384:
+        LAUNCH_BIASED_GROUPED_TOPK_KERNEL(384, 8);
+        break;
       default:
         TORCH_CHECK(false, "Unexpected num_experts: ", num_experts);
     }