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

[CPU] [Quantization] Add GPTQ/AWQ 4bits quantization support for CPU (#22685)

Browse Source 
Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
jianan-gu
2026-04-22 13:34:02 -07:00
committed by GitHub
co-authored by Ma Mingfei
parent 0b77284587
commit ad0fc88810
14 changed files with 835 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/cpu/torch_extension_cpu.cpp
+7 -4
View File
@@ -200,8 +200,11 @@ at::Tensor int4_scaled_mm_cpu(
at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional<at::Tensor> bias);
// weight prepack for int4 weights
std::tuple<at::Tensor, at::Tensor, at::Tensor>
convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32
at::Tensor qzeros, // awq: (*, K / group_size, N / 8) || gptq: (*, K / group_size, N / 8) , int32
at::Tensor scales, // awq: (*, K / group_size, N) || gptq: (*, K / group_size, N) , bfloat16
int64_t quant_method_4bit);
// bmm
void bmm_cpu(at::Tensor& out, at::Tensor& mat1, at::Tensor& mat2, bool is_vnni, const std::optional<at::Tensor>& scale);
@@ -520,8 +523,8 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
// weight prepack for int4 weights
m.def(
"convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor scales) -> (Tensor, Tensor, "
"Tensor)");
"convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor scales, int quant_method_4bit) -> (Tensor, "
"Tensor, Tensor)");
m.impl("convert_weight_packed_scale_zp", torch::kCPU, &convert_weight_packed_scale_zp);
// bmm