[CPU][INT4] Add INT4 kernels for CPU (#8226)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

test/srt/cpu/test_gemm.py

@@ -9,6 +9,7 @@ from utils import (
     native_w8a8_per_token_matmul,
     per_token_quant_int8,
     precision,
+    unpack_and_dequant_awq,
 )
 
 from sglang.test.test_utils import CustomTestCase
@@ -39,6 +40,10 @@ class TestGemm(CustomTestCase):
     N_fp8 = [128, 224]
     K_fp8 = [512, 576]
 
+    M_awq = [1, 32]
+    N_awq = [4096]
+    K_awq = [4096]
+
     def _bf16_gemm(self, M, N, K, has_bias):
 
         mat1 = torch.randn(M, K, dtype=torch.bfloat16)
@@ -227,6 +232,51 @@ class TestGemm(CustomTestCase):
             ):
                 self._fp8_gemm(*params)
 
+    def _int4_awq_gemm(self, M, N, K, group_size, has_bias):
+        awq_weight = torch.randint(-128, 128, (K, N // 8)).to(torch.int)
+        awq_zero = torch.randint(0, 10, (K // group_size, N // 8)).to(torch.int)
+        awq_scales = torch.rand(int(K // group_size), N).to(torch.bfloat16)
+        bf16_weight, _ = unpack_and_dequant_awq(
+            awq_weight, awq_zero, awq_scales, 4, 128
+        )
+        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(
+                awq_weight, awq_zero, awq_scales
+            )
+        )
+        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_awq_gemm(self):
+        for params in itertools.product(
+            self.M_awq, self.N_awq, self.K_awq, [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_awq_gemm(*params)
+
 
 if __name__ == "__main__":
     unittest.main()