minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CPU] Implement MXFP4 Gemm kernels for intel AMX to support GPT OSS series. (#14385)

Browse Source
This commit is contained in:
Ma Mingfei
2026-03-29 23:44:12 -07:00
committed by GitHub
parent ed01e1d5d6
commit af62bd9486
6 changed files with 555 additions and 92 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+15
View File
@@ -143,6 +143,9 @@ std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
// weight prepack
at::Tensor convert_weight_packed(at::Tensor& weight);
// scale prepack for mxfp4
at::Tensor convert_scale_packed(at::Tensor& scale);
// quant
std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A);
@@ -178,6 +181,10 @@ at::Tensor fp8_scaled_mm_cpu(
at::ScalarType out_dtype,
bool is_vnni);
// mxfp4 gemm
at::Tensor mxfp4_scaled_mm_cpu(
at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2, const std::optional<at::Tensor>& bias, bool is_vnni);
// quant + igemm
at::Tensor int8_scaled_mm_with_quant(
at::Tensor& mat1,
@@ -465,6 +472,10 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
m.def("convert_weight_packed(Tensor weight) -> Tensor");
m.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);
// scale prepack for mxfp4
m.def("convert_scale_packed(Tensor scale) -> Tensor");
m.impl("convert_scale_packed", torch::kCPU, &convert_scale_packed);
// quant
m.def("per_token_quant_int8_cpu(Tensor A) -> (Tensor, Tensor)");
m.impl("per_token_quant_int8_cpu", torch::kCPU, &per_token_quant_int8_cpu);
@@ -490,6 +501,10 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
"out_dtype, bool is_vnni) -> Tensor");
m.impl("fp8_scaled_mm_cpu", torch::kCPU, &fp8_scaled_mm_cpu);
// mxfp4 gemm
m.def("mxfp4_scaled_mm_cpu(Tensor mat1, Tensor mat2, Tensor scales2, Tensor? bias, bool is_vnni) -> Tensor");
m.impl("mxfp4_scaled_mm_cpu", torch::kCPU, &mxfp4_scaled_mm_cpu);
// quant + igemm
m.def(
"int8_scaled_mm_with_quant(Tensor mat1, Tensor mat2, Tensor scales2, Tensor? bias, ScalarType out_dtype, bool "