chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Kernel Slimming] Migrate AWQ marlin repack kernel to JIT (#18949)

Browse Source 
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
This commit is contained in:
Linyu Wu
2026-02-23 22:05:27 +08:00
committed by GitHub
parent e0e0cad6bc
commit 2cdde5d4ab
11 changed files with 1336 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

164
python/sglang/jit_kernel/tests/test_awq_dequantize.py Normal file
View File

@@ -0,0 +1,164 @@
import itertools
import pytest
import torch
from sglang.jit_kernel.awq_dequantize import awq_dequantize as jit_awq_dequantize
try:
from sgl_kernel import awq_dequantize as aot_awq_dequantize
AOT_AVAILABLE = True
except ImportError:
AOT_AVAILABLE = False
def reverse_awq_order(t: torch.Tensor):
bits = 4
AWQ_REVERSE_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
reverse_order_tensor = torch.arange(
t.shape[-1],
dtype=torch.int32,
device=t.device,
)
reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits)
reverse_order_tensor = reverse_order_tensor[:, AWQ_REVERSE_ORDER]
reverse_order_tensor = reverse_order_tensor.view(-1)
t = t[:, reverse_order_tensor] & 0xF
return t
# qweights - [R , C // 8], int32
# scales - [R // G, C ], float16
# zeros - [R // G, C // 8], int32
def awq_dequantize_torch(
qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor, group_size: int
) -> torch.Tensor:
if group_size == -1:
group_size = qweight.shape[0]
bits = 4
shifts = torch.arange(0, 32, bits, device=qzeros.device)
iweights = torch.bitwise_right_shift(qweight[:, :, None], shifts[None, None, :]).to(
torch.int8
)
iweights = iweights.view(iweights.shape[0], -1)
zeros = torch.bitwise_right_shift(qzeros[:, :, None], shifts[None, None, :]).to(
torch.int8
)
zeros = zeros.view(qzeros.shape[0], -1)
zeros = reverse_awq_order(zeros)
iweights = reverse_awq_order(iweights)
iweights = torch.bitwise_and(iweights, (2**bits) - 1)
zeros = torch.bitwise_and(zeros, (2**bits) - 1)
scales = scales.repeat_interleave(group_size, dim=0)
zeros = zeros.repeat_interleave(group_size, dim=0)
return (iweights - zeros) * scales
@pytest.mark.parametrize(
"qweight_row,qweight_col,is_bf16_act",
list(
itertools.product(
[128, 256, 512, 1024, 3584],
[16, 32, 64, 128, 448],
[True, False],
)
),
)
def test_awq_dequantize_jit_vs_torch(
qweight_row: int, qweight_col: int, is_bf16_act: bool
):
device = torch.device("cuda")
qweight = torch.randint(
0,
torch.iinfo(torch.int32).max,
(qweight_row, qweight_col),
dtype=torch.int32,
device=device,
)
group_size = qweight_row
scales_row = qweight_row // group_size
scales_col = qweight_col * 8
if is_bf16_act:
scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device)
else:
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
qzeros = torch.randint(
0,
torch.iinfo(torch.int32).max,
(scales_row, qweight_col),
dtype=torch.int32,
device=device,
)
# Run both implementations
torch_out = awq_dequantize_torch(qweight, scales, qzeros, group_size)
jit_out = jit_awq_dequantize(qweight, scales, qzeros)
# Compare results (approximate due to different computation paths)
torch.testing.assert_close(
torch_out.to(torch.float32), jit_out.to(torch.float32), rtol=1e-3, atol=1e-5
)
@pytest.mark.parametrize(
"qweight_row,qweight_col,is_bf16_act",
list(
itertools.product(
[128, 256, 512, 1024, 3584],
[16, 32, 64, 128, 448],
[True, False],
)
),
)
def test_awq_dequantize_jit_vs_aot(
qweight_row: int, qweight_col: int, is_bf16_act: bool
):
if not AOT_AVAILABLE:
pytest.skip("sgl_kernel AOT not available")
device = torch.device("cuda")
qweight = torch.randint(
0,
torch.iinfo(torch.int32).max,
(qweight_row, qweight_col),
dtype=torch.int32,
device=device,
)
group_size = qweight_row
scales_row = qweight_row // group_size
scales_col = qweight_col * 8
if is_bf16_act:
scales = torch.rand(scales_row, scales_col, dtype=torch.bfloat16, device=device)
else:
scales = torch.rand(scales_row, scales_col, dtype=torch.float16, device=device)
qzeros = torch.randint(
0,
torch.iinfo(torch.int32).max,
(scales_row, qweight_col),
dtype=torch.int32,
device=device,
)
# Run both implementations
aot_out = aot_awq_dequantize(qweight, scales, qzeros)
jit_out = jit_awq_dequantize(qweight, scales, qzeros)
# Bitwise equality
torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0)
if __name__ == "__main__":
pytest.main([__file__])

117
python/sglang/jit_kernel/tests/test_awq_marlin_moe_repack.py Normal file
View File

@@ -0,0 +1,117 @@
import numpy as np
import pytest
import torch
from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.awq_marlin_repack import (
awq_marlin_moe_repack as jit_awq_marlin_moe_repack,
)
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
try:
from sgl_kernel import awq_marlin_moe_repack as aot_awq_marlin_moe_repack
AOT_AVAILABLE = True
except ImportError:
AOT_AVAILABLE = False
def awq_pack(
q_w: torch.Tensor,
num_bits: int,
size_k: int,
size_n: int,
):
assert q_w.shape == (size_k, size_n)
if num_bits == 4:
interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
elif num_bits == 8:
interleave = np.array([0, 2, 1, 3])
else:
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel()
q_w = q_w.reshape((-1, size_n)).contiguous()
return pack_cols(q_w, num_bits, size_k, size_n)
@pytest.mark.parametrize("num_bits", [4])
@pytest.mark.parametrize("num_experts", [2, 4, 8])
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2), (4, 4)])
@pytest.mark.parametrize("group_size", [16, 32])
def test_awq_marlin_moe_repack_jit_vs_aot(
num_bits, num_experts, k_tiles, n_tiles, group_size
):
if not AOT_AVAILABLE:
pytest.skip("sgl_kernel AOT not available")
tile_k, tile_n = 16, 64
size_k = k_tiles * tile_k
size_n = n_tiles * tile_n
pack_factor = 32 // num_bits
# Create per-expert AWQ-packed weights
b_q_weight = torch.empty(
(num_experts, size_k, size_n // pack_factor),
dtype=torch.int32,
device="cuda",
)
for e in range(num_experts):
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
w_ref, q_w, s, zp = quantize_weights(
b_weight, scalar_types.uint4, group_size, zero_points=True
)
b_q_weight[e] = awq_pack(q_w, num_bits, size_k, size_n)
perm = torch.empty((num_experts, 0), dtype=torch.int32, device="cuda")
out_jit = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, size_n, num_bits)
out_aot = aot_awq_marlin_moe_repack(b_q_weight, perm, size_k, size_n, num_bits)
torch.cuda.synchronize()
# Bitwise equality
torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
@pytest.mark.parametrize("num_bits", [4])
@pytest.mark.parametrize("num_experts", [2, 4])
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2)])
@pytest.mark.parametrize("group_size", [16, 32])
def test_awq_marlin_moe_repack_shape(
num_bits, num_experts, k_tiles, n_tiles, group_size
):
tile_k, tile_n = 16, 64
size_k = k_tiles * tile_k
size_n = n_tiles * tile_n
pack_factor = 32 // num_bits
# Create per-expert AWQ-packed weights
b_q_weight = torch.empty(
(num_experts, size_k, size_n // pack_factor),
dtype=torch.int32,
device="cuda",
)
for e in range(num_experts):
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
w_ref, q_w, s, zp = quantize_weights(
b_weight, scalar_types.uint4, group_size, zero_points=True
)
b_q_weight[e] = awq_pack(q_w, num_bits, size_k, size_n)
perm = torch.empty((num_experts, 0), dtype=torch.int32, device="cuda")
out = jit_awq_marlin_moe_repack(b_q_weight, perm, size_k, size_n, num_bits)
torch.cuda.synchronize()
assert out.is_cuda and out.dtype == torch.int32
expected_shape = (num_experts, size_k // 16, size_n * (num_bits // 2))
assert list(out.shape) == list(expected_shape)
if __name__ == "__main__":
import subprocess
subprocess.call(["pytest", "--tb=short", str(__file__)])

103
python/sglang/jit_kernel/tests/test_awq_marlin_repack.py Normal file
View File

@@ -0,0 +1,103 @@
import numpy as np
import pytest
import torch
from sgl_kernel.scalar_type import scalar_types
from sglang.jit_kernel.awq_marlin_repack import (
awq_marlin_repack as jit_awq_marlin_repack,
)
from sglang.srt.layers.quantization.utils import pack_cols, quantize_weights
from sglang.test.test_marlin_utils import get_weight_perm, marlin_weights
try:
from sgl_kernel import awq_marlin_repack as aot_awq_marlin_repack
AOT_AVAILABLE = True
except ImportError:
AOT_AVAILABLE = False
def awq_pack(
q_w: torch.Tensor,
num_bits: int,
size_k: int,
size_n: int,
):
assert q_w.shape == (size_k, size_n)
if num_bits == 4:
interleave = np.array([0, 2, 4, 6, 1, 3, 5, 7])
elif num_bits == 8:
interleave = np.array([0, 2, 1, 3])
else:
raise Exception("num_bits must be 4 or 8, got {}".format(num_bits))
q_w = q_w.reshape((-1, len(interleave)))[:, interleave].ravel()
q_w = q_w.reshape((-1, size_n)).contiguous()
return pack_cols(q_w, num_bits, size_k, size_n)
@pytest.mark.parametrize("num_bits", [4, 8])
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2), (4, 4)])
@pytest.mark.parametrize("group_size", [16, 32])
def test_awq_marlin_repack_jit_vs_aot(num_bits, k_tiles, n_tiles, group_size):
if not AOT_AVAILABLE:
pytest.skip("sgl_kernel AOT not available")
tile_k, tile_n = 16, 64
size_k = k_tiles * tile_k
size_n = n_tiles * tile_n
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
w_ref, q_w, s, zp = quantize_weights(
b_weight, scalar_types.uint4, group_size, zero_points=True
)
q_w_awq = awq_pack(q_w, num_bits, size_k, size_n)
out_jit = jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
out_aot = aot_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
torch.cuda.synchronize()
# Bitwise equality
torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
@pytest.mark.parametrize("num_bits", [4, 8])
@pytest.mark.parametrize("k_tiles,n_tiles", [(1, 1), (2, 2)])
@pytest.mark.parametrize("group_size", [16, 32])
def test_awq_marlin_repack_correct(num_bits, k_tiles, n_tiles, group_size):
tile_k, tile_n = 16, 64
size_k = k_tiles * tile_k
size_n = n_tiles * tile_n
pack_factor = 32 // num_bits
b_weight = torch.randn((size_k, size_n), dtype=torch.float16, device="cuda")
w_ref, q_w, s, zp = quantize_weights(
b_weight, scalar_types.uint4, group_size, zero_points=True
)
q_w_awq = awq_pack(q_w, num_bits, size_k, size_n)
weight_perm = get_weight_perm(num_bits)
q_w_marlin = marlin_weights(q_w, size_k, size_n, num_bits, weight_perm)
out_gpu = jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
assert out_gpu.is_cuda and out_gpu.dtype == torch.int32
expected_cols = size_n * tile_k // pack_factor
assert list(out_gpu.shape) == [size_k // tile_k, expected_cols]
torch.cuda.synchronize()
torch.testing.assert_close(out_gpu, q_w_marlin)
if __name__ == "__main__":
import subprocess
subprocess.call(["pytest", "--tb=short", str(__file__)])