[CPU][INT4] Add INT4 kernels for CPU (#8226)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
jianan-gu
co-authored by
gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
parent
88f7759402
commit
c35aa0238c
@@ -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()
|
||||
|
||||
+116
-9
@@ -5,9 +5,11 @@ import unittest
|
||||
# TODO: use interface in cpu.py
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.amx_utils import CPUQuantMethod
|
||||
|
||||
kernel = torch.ops.sgl_kernel
|
||||
|
||||
torch.manual_seed(1234)
|
||||
torch.manual_seed(128)
|
||||
|
||||
from utils import (
|
||||
BLOCK_K,
|
||||
@@ -20,6 +22,7 @@ from utils import (
|
||||
scaled_weight,
|
||||
torch_naive_fused_moe,
|
||||
torch_w8a8_per_column_fused_moe,
|
||||
unpack_and_dequant_awq,
|
||||
)
|
||||
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
@@ -48,8 +51,7 @@ def fused_moe(a, w1, w2, score, topk, renormalize, prepack):
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
inplace,
|
||||
False,
|
||||
False,
|
||||
CPUQuantMethod.UNQUANT,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
@@ -79,6 +81,12 @@ class TestFusedExperts(CustomTestCase):
|
||||
E_fp8 = [8]
|
||||
topk_fp8 = [4]
|
||||
|
||||
M_int4 = [1, 6]
|
||||
N_int4 = [512]
|
||||
K_int4 = [256]
|
||||
E_int4 = [8]
|
||||
topk_int4 = [4]
|
||||
|
||||
def _bf16_moe(self, m, n, k, e, topk, renormalize):
|
||||
dtype = torch.bfloat16
|
||||
prepack = True
|
||||
@@ -156,8 +164,7 @@ class TestFusedExperts(CustomTestCase):
|
||||
topk_weight,
|
||||
topk_ids.to(torch.int32),
|
||||
inplace,
|
||||
True,
|
||||
False,
|
||||
CPUQuantMethod.INT8_W8A8,
|
||||
w1_s,
|
||||
w2_s,
|
||||
None,
|
||||
@@ -229,13 +236,12 @@ class TestFusedExperts(CustomTestCase):
|
||||
topk_weight,
|
||||
topk_ids.to(torch.int32),
|
||||
False,
|
||||
False,
|
||||
True,
|
||||
CPUQuantMethod.FP8_W8A16,
|
||||
w1s,
|
||||
w2s,
|
||||
None,
|
||||
None,
|
||||
[BLOCK_N, BLOCK_K],
|
||||
None,
|
||||
None,
|
||||
True,
|
||||
)
|
||||
|
||||
@@ -259,6 +265,107 @@ class TestFusedExperts(CustomTestCase):
|
||||
):
|
||||
self._fp8_moe(*params)
|
||||
|
||||
def _int4_moe(self, M, N, K, E, topk, group_size=128):
|
||||
dtype = torch.bfloat16
|
||||
|
||||
a = torch.rand(M, K, dtype=dtype) / math.sqrt(K)
|
||||
|
||||
awq_w13_weight = torch.randint(-127, 128, (E, K, 2 * N // 8)).to(torch.int)
|
||||
awq_w13_zero = torch.randint(0, 10, (E, K // group_size, 2 * N // 8)).to(
|
||||
torch.int
|
||||
)
|
||||
awq_w13_scales = torch.rand(E, int(K // group_size), 2 * N).to(torch.bfloat16)
|
||||
|
||||
awq_w2_weight = torch.randint(-127, 128, (E, N, K // 8)).to(torch.int)
|
||||
awq_w2_zero = torch.randint(0, 10, (E, N // group_size, K // 8)).to(torch.int)
|
||||
awq_w2_scales = torch.rand(E, int(N // group_size), K).to(torch.bfloat16)
|
||||
bf16_w13_weight = []
|
||||
bf16_w2_weight = []
|
||||
for i in range(E):
|
||||
bf16_w13_weight_i, _ = unpack_and_dequant_awq(
|
||||
awq_w13_weight[i], awq_w13_zero[i], awq_w13_scales[i], 4, 128
|
||||
)
|
||||
bf16_w2_weight_i, _ = unpack_and_dequant_awq(
|
||||
awq_w2_weight[i], awq_w2_zero[i], awq_w2_scales[i], 4, 128
|
||||
)
|
||||
bf16_w13_weight.append(bf16_w13_weight_i)
|
||||
bf16_w2_weight.append(bf16_w2_weight_i)
|
||||
bf16_w13_weight = torch.stack(bf16_w13_weight).detach()
|
||||
bf16_w2_weight = torch.stack(bf16_w2_weight).detach()
|
||||
|
||||
score = torch.rand((M, E), dtype=dtype)
|
||||
|
||||
ref_out = torch_naive_fused_moe(
|
||||
a, bf16_w13_weight, bf16_w2_weight, score, topk, False
|
||||
)
|
||||
score = torch.softmax(score, dim=-1, dtype=torch.float32)
|
||||
topk_weight, topk_ids = torch.topk(score, topk)
|
||||
awq_w13_weight_pack = []
|
||||
awq_w13_zero_pack = []
|
||||
awq_w13_scales_pack = []
|
||||
awq_w2_weight_pack = []
|
||||
awq_w2_zero_pack = []
|
||||
awq_w2_scales_pack = []
|
||||
for i in range(E):
|
||||
packed_weight_13_i, packed_zero_13_i, packed_scales_13_i = (
|
||||
torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
|
||||
awq_w13_weight[i], awq_w13_zero[i], awq_w13_scales[i]
|
||||
)
|
||||
)
|
||||
awq_w13_weight_pack.append(packed_weight_13_i)
|
||||
awq_w13_zero_pack.append(packed_zero_13_i)
|
||||
awq_w13_scales_pack.append(packed_scales_13_i)
|
||||
packed_weight_2_i, packed_zero_2_i, packed_scales_2_i = (
|
||||
torch.ops.sgl_kernel.convert_weight_packed_scale_zp(
|
||||
awq_w2_weight[i], awq_w2_zero[i], awq_w2_scales[i]
|
||||
)
|
||||
)
|
||||
awq_w2_weight_pack.append(packed_weight_2_i)
|
||||
awq_w2_zero_pack.append(packed_zero_2_i)
|
||||
awq_w2_scales_pack.append(packed_scales_2_i)
|
||||
awq_w13_weight_pack = torch.stack(awq_w13_weight_pack).detach()
|
||||
awq_w13_zero_pack = torch.stack(awq_w13_zero_pack).detach()
|
||||
awq_w13_scales_pack = torch.stack(awq_w13_scales_pack).detach()
|
||||
awq_w2_weight_pack = torch.stack(awq_w2_weight_pack).detach()
|
||||
awq_w2_zero_pack = torch.stack(awq_w2_zero_pack).detach()
|
||||
awq_w2_scales_pack = torch.stack(awq_w2_scales_pack).detach()
|
||||
|
||||
out = kernel.fused_experts_cpu(
|
||||
a,
|
||||
awq_w13_weight_pack,
|
||||
awq_w2_weight_pack,
|
||||
topk_weight,
|
||||
topk_ids.to(torch.int32),
|
||||
False,
|
||||
CPUQuantMethod.INT4_W4A8,
|
||||
awq_w13_scales_pack,
|
||||
awq_w2_scales_pack,
|
||||
awq_w13_zero_pack,
|
||||
awq_w2_zero_pack,
|
||||
None,
|
||||
True,
|
||||
)
|
||||
|
||||
atol = rtol = precision[dtype]
|
||||
torch.testing.assert_close(ref_out.bfloat16(), out, atol=atol, rtol=rtol)
|
||||
|
||||
def test_int4_moe(self):
|
||||
for params in itertools.product(
|
||||
self.M_int4,
|
||||
self.N_int4,
|
||||
self.K_int4,
|
||||
self.E_int4,
|
||||
self.topk_int4,
|
||||
):
|
||||
with self.subTest(
|
||||
M=params[0],
|
||||
N=params[1],
|
||||
K=params[2],
|
||||
E=params[3],
|
||||
topk=params[4],
|
||||
):
|
||||
self._int4_moe(*params)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -66,8 +66,6 @@ class TestSharedExpert(CustomTestCase):
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
False,
|
||||
)
|
||||
|
||||
@@ -124,8 +122,6 @@ class TestSharedExpert(CustomTestCase):
|
||||
w1_s,
|
||||
w2_s,
|
||||
None,
|
||||
None,
|
||||
None,
|
||||
False,
|
||||
)
|
||||
|
||||
@@ -196,8 +192,6 @@ class TestSharedExpert(CustomTestCase):
|
||||
w1s,
|
||||
w2s,
|
||||
[BLOCK_N, BLOCK_K],
|
||||
None,
|
||||
None,
|
||||
True,
|
||||
)
|
||||
|
||||
|
||||
@@ -286,3 +286,105 @@ def make_non_contiguous(x: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
last_dim = x.shape[-1]
|
||||
return x[..., : last_dim // 2] if x.is_contiguous() else x
|
||||
|
||||
|
||||
def awq_reverse_reorder_int_tensor(int_tensor, bits: int):
|
||||
assert bits == 4
|
||||
|
||||
int_tensor = int_tensor.T.contiguous()
|
||||
compress_ratio = 32 // bits
|
||||
assert int_tensor.shape[-1] % compress_ratio == 0
|
||||
|
||||
order_map = [0, 2, 4, 6, 1, 3, 5, 7]
|
||||
order_tensor = torch.tensor(
|
||||
order_map, dtype=torch.int32, device=int_tensor.device
|
||||
).reshape(1, -1)
|
||||
order_tensor = order_tensor.repeat(int_tensor.shape[1] // compress_ratio, 1)
|
||||
order_tensor = order_tensor + torch.arange(
|
||||
0,
|
||||
int_tensor.shape[1],
|
||||
compress_ratio,
|
||||
dtype=torch.int32,
|
||||
device=int_tensor.device,
|
||||
).reshape(-1, 1)
|
||||
order_tensor = order_tensor.reshape(-1)
|
||||
|
||||
reverse_order_tensor = torch.arange(order_tensor.shape[0])[order_tensor]
|
||||
reverse_order_tensor = reverse_order_tensor[order_tensor]
|
||||
int_tensor = int_tensor[:, reverse_order_tensor]
|
||||
return int_tensor
|
||||
|
||||
|
||||
def unpack_and_dequant_awq(
|
||||
awq_qweight: torch.Tensor,
|
||||
awq_qzeros: torch.Tensor,
|
||||
awq_scales: torch.Tensor,
|
||||
bits: int,
|
||||
group_size: int,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
awq_qweight (`torch.LongTensor`):
|
||||
Expected shape: (in_features, out_features // (32 // bits))
|
||||
awq_qzeros (`torch.LongTensor`):
|
||||
Expected shape: (in_features // group_size, out_features // (32 // bits))
|
||||
awq_scales (`torch.LongTensor`):
|
||||
Expected shape: (in_features // group_size, out_features)
|
||||
|
||||
Returns:
|
||||
fp16_weight (`torch.LongTensor`):
|
||||
With shape (in_features, out_features).
|
||||
zeros (`torch.LongTensor`):
|
||||
With shape (in_features // group_size, out_features).
|
||||
"""
|
||||
assert bits == 4
|
||||
|
||||
qzeros = awq_qzeros
|
||||
qweight = awq_qweight
|
||||
qweight = qweight.T.contiguous()
|
||||
|
||||
scales = awq_scales
|
||||
scales = scales.reshape(-1, 1, scales.shape[-1])
|
||||
|
||||
infeatures = awq_qweight.shape[0]
|
||||
|
||||
wf = torch.tensor(
|
||||
list(range(0, 32, bits)), dtype=torch.int32, device=qzeros.device
|
||||
).unsqueeze(0)
|
||||
zeros = torch.bitwise_right_shift(torch.unsqueeze(qzeros, 2), wf.unsqueeze(0)).to(
|
||||
torch.int16 if bits == 8 else torch.int8
|
||||
)
|
||||
|
||||
torch.bitwise_and(zeros, (2**bits) - 1, out=zeros)
|
||||
|
||||
zeros = zeros.reshape(-1, 1, zeros.shape[1] * zeros.shape[2])
|
||||
|
||||
weight = torch.bitwise_right_shift(
|
||||
torch.unsqueeze(qweight, 1), wf.unsqueeze(-1)
|
||||
).to(torch.int16 if bits == 8 else torch.int8)
|
||||
torch.bitwise_and(weight, (2**bits) - 1, out=weight)
|
||||
weight = weight.reshape(-1, group_size, weight.shape[2])
|
||||
|
||||
weight = weight.view(-1, weight.shape[-1])
|
||||
zeros = zeros.view(-1, zeros.shape[-1])
|
||||
|
||||
zeros = zeros.T.contiguous()
|
||||
zeros = awq_reverse_reorder_int_tensor(zeros, bits)
|
||||
weight = awq_reverse_reorder_int_tensor(weight, bits)
|
||||
|
||||
# Dequantize weights.
|
||||
scales = awq_scales
|
||||
zeros = zeros.contiguous()
|
||||
scale_zeros = zeros * scales
|
||||
|
||||
g_idx = torch.tensor(
|
||||
[i // group_size for i in range(infeatures)], dtype=torch.int32
|
||||
)
|
||||
scale_mat = scales[g_idx]
|
||||
scale_zeros_mat = scale_zeros[g_idx].to(torch.bfloat16)
|
||||
|
||||
qdq_weight_T = weight * scale_mat - scale_zeros_mat.to(torch.bfloat16)
|
||||
|
||||
fp16_weight = qdq_weight_T.T
|
||||
|
||||
return fp16_weight, zeros
|
||||
|
||||
Reference in New Issue
Block a user