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

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>

test/srt/cpu/test_gemm.py

@@ -10,6 +10,7 @@ from utils import (
     per_token_quant_int8,
     precision,
     unpack_and_dequant_awq,
+    unpack_and_dequant_gptq,
 )
 
 from sglang.test.test_utils import CustomTestCase
@@ -44,6 +45,10 @@ class TestGemm(CustomTestCase):
     N_awq = [4096]
     K_awq = [4096]
 
+    M_gptq = [1, 32]
+    N_gptq = [4096]
+    K_gptq = [4096]
+
     def _bf16_gemm(self, M, N, K, has_bias):
 
         mat1 = torch.randn(M, K, dtype=torch.bfloat16)
@@ -250,7 +255,7 @@ class TestGemm(CustomTestCase):
 
         packed_weight, packed_zero, packed_scales = (
             torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
-                awq_weight, awq_zero, awq_scales
+                awq_weight, awq_zero, awq_scales, 0
             )
         )
         target_res = torch.ops.sgl_kernel.int4_scaled_mm_cpu(
@@ -277,6 +282,51 @@ class TestGemm(CustomTestCase):
             ):
                 self._int4_awq_gemm(*params)
 
+    def _int4_gptq_gemm(self, M, N, K, group_size, has_bias):
+        torch.manual_seed(127)
+        gptq_weight = torch.randint(-128, 128, (K // 8, N)).to(torch.int)
+        gptq_zero = torch.randint(0, 10, (K // group_size, N // 8)).to(torch.int)
+        gptq_scales = torch.rand(int(K // group_size), N).to(torch.bfloat16) // 10
+
+        bf16_weight = unpack_and_dequant_gptq(gptq_weight, gptq_zero, gptq_scales)
+        if has_bias:
+            bias = torch.rand(bf16_weight.shape[0]).to(torch.float)
+        else:
+            bias = None
+        x = torch.rand(M, bf16_weight.size(-1)).to(torch.bfloat16)
+        ref_res = torch.nn.functional.linear(
+            x, bf16_weight, bias=bias.to(torch.bfloat16) if has_bias else None
+        )
+
+        packed_weight, packed_zero, packed_scales = (
+            torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
+                gptq_weight, gptq_zero, gptq_scales, 1
+            )
+        )
+        target_res = torch.ops.sgl_kernel.int4_scaled_mm_cpu(
+            x,
+            packed_weight,
+            packed_zero,
+            packed_scales,
+            bias,
+        )
+
+        atol = rtol = precision[ref_res.dtype]
+        torch.testing.assert_close(ref_res, target_res, atol=atol, rtol=rtol)
+
+    def test_int4_gptq_gemm(self):
+        for params in itertools.product(
+            self.M_gptq, self.N_gptq, self.K_gptq, [128], self.has_bias
+        ):
+            with self.subTest(
+                M=params[0],
+                N=params[1],
+                K=params[2],
+                group_size=params[3],
+                has_bias=params[4],
+            ):
+                self._int4_gptq_gemm(*params)
+
 
 if __name__ == "__main__":
     unittest.main()

test/srt/cpu/test_moe.py

@@ -302,12 +302,12 @@ class TestFusedExperts(CustomTestCase):
         topk_weight, topk_ids = torch.topk(score, topk)
         awq_w13_weight_pack, awq_w13_zero_pack, awq_w13_scales_pack = (
             torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
-                awq_w13_weight, awq_w13_zero, awq_w13_scales
+                awq_w13_weight, awq_w13_zero, awq_w13_scales, 0
             )
         )
         awq_w2_weight_pack, awq_w2_zero_pack, awq_w2_scales_pack = (
             torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
-                awq_w2_weight, awq_w2_zero, awq_w2_scales
+                awq_w2_weight, awq_w2_zero, awq_w2_scales, 0
             )
         )
 

test/srt/cpu/utils.py

@@ -402,3 +402,39 @@ def unpack_and_dequant_awq(
     fp16_weight = qdq_weight_T.T
 
     return fp16_weight, zeros
+
+
+def unpack_4bit_to_32bit_signed(qweight, qzeros):
+    # Unpack 4-bit values and interpret them as signed integers
+    unpacked_weights = torch.zeros(
+        (qweight.shape[0] * 8, qweight.shape[1]),
+        dtype=torch.int8,
+        device=qweight.device,
+        requires_grad=False,
+    )
+    unpacked_zeros = torch.zeros(
+        (qzeros.shape[0], qzeros.shape[1] * 8),
+        dtype=torch.int8,
+        device=qzeros.device,
+        requires_grad=False,
+    )
+
+    for row in range(unpacked_weights.shape[0]):
+        i = row % 8
+        unpacked_weights[row, :] = (qweight[row // 8, :] >> (4 * i)) & 0xF
+
+    for col in range(unpacked_zeros.shape[1]):
+        i = col % 8
+        unpacked_zeros[:, col] = (qzeros[:, col // 8] >> (4 * i)) & 0xF
+
+    return unpacked_weights, unpacked_zeros + 1
+
+
+def unpack_and_dequant_gptq(qweight, qzeros, scales):
+    unpacked_qweight, unpacked_qzeros = unpack_4bit_to_32bit_signed(qweight, qzeros)
+    group_size = unpacked_qweight.shape[0] // scales.shape[0]
+    scales = scales.repeat_interleave(group_size, dim=0)
+    unpacked_qzeros = unpacked_qzeros.repeat_interleave(group_size, dim=0)
+    unpacked_qweight = (unpacked_qweight - unpacked_qzeros) * scales
+
+    return unpacked_qweight.T