chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[CPU] Optimize small oc GEMM for Qwen3-next on CPU (#12446)

Browse Source 
Co-authored-by: Zheng, Beilei <beilei.zheng@intel.com>
This commit is contained in:
jianan-gu
2025-12-04 00:38:47 -08:00
committed by GitHub
co-authored by Zheng, Beilei <beilei.zheng@intel.com>
parent 894c0dc57c
commit 70d2587324
5 changed files with 357 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+13
View File
@@ -124,6 +124,14 @@ std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A);
at::Tensor
weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni);
// gemm fusion
at::Tensor fused_linear_sigmoid_mul(
at::Tensor& mat1,
at::Tensor& mat2,
const std::optional<at::Tensor>& bias,
bool is_vnni,
const at::Tensor& post_mul_mat);
// igemm
at::Tensor int8_scaled_mm_cpu(
at::Tensor& mat1,
@@ -355,6 +363,11 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.def("weight_packed_linear(Tensor mat1, Tensor mat2, Tensor? bias, bool is_vnni) -> Tensor");
m.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);
// gemm fusion
m.def(
"fused_linear_sigmoid_mul(Tensor mat1, Tensor mat2, Tensor? bias, bool is_vnni, Tensor post_mul_mat) -> Tensor");
m.impl("fused_linear_sigmoid_mul", torch::kCPU, &fused_linear_sigmoid_mul);
// igemm
m.def(
"int8_scaled_mm_cpu(Tensor mat1, Tensor mat2, Tensor scales1, Tensor scales2, Tensor? bias, ScalarType "