[CPU] Optimize Qwen3-next model on CPU (#12525)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com> Co-authored-by: Fan Yin <1106310035@qq.com>
2026-01-29 22:03:58 -08:00
parent 71e4d3b6bc
commit 336dc4579e
13 changed files with 366 additions and 41 deletions
@@ -525,6 +525,12 @@ topk_softmax_cpu(at::Tensor& hidden_states, at::Tensor& gating_output, int64_t t
      case 256:
        LAUNCH_TOPK_SOFTMAX_KERNEL(256);
        break;
+      case 384:
+        LAUNCH_TOPK_SOFTMAX_KERNEL(384);
+        break;
+      case 512:
+        LAUNCH_TOPK_SOFTMAX_KERNEL(512);
+        break;
      default:
        TORCH_CHECK(false, "Unexpected num_experts: ", num_experts);
    }