[Kernel Slimming] Migrate AWQ marlin repack kernel to JIT (#18949)
Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
This commit is contained in:
Linyu Wu
125
python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py
Normal file
125
python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py
Normal file
@@ -0,0 +1,125 @@
|
||||
import itertools
|
||||
import os
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
|
||||
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
|
||||
|
||||
IS_CI = (
|
||||
os.getenv("CI", "false").lower() == "true"
|
||||
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
|
||||
)
|
||||
|
||||
# CI environment uses simplified parameters
|
||||
if IS_CI:
|
||||
qweight_row_range = [128]
|
||||
qweight_cols_range = [16]
|
||||
else:
|
||||
qweight_row_range = [128, 256, 512, 1024, 3584]
|
||||
qweight_cols_range = [16, 32, 64, 128, 448]
|
||||
|
||||
configs = list(itertools.product(qweight_row_range, qweight_cols_range))
|
||||
|
||||
|
||||
def check_correctness():
|
||||
if not AOT_AVAILABLE:
|
||||
print("sgl_kernel AOT not available, skipping correctness check")
|
||||
return
|
||||
|
||||
qweight_row, qweight_col = 128, 16
|
||||
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
|
||||
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,
|
||||
)
|
||||
|
||||
jit_out = jit_awq_dequantize(qweight, scales, qzeros)
|
||||
aot_out = aot_awq_dequantize(qweight, scales, qzeros)
|
||||
torch.cuda.synchronize()
|
||||
torch.testing.assert_close(jit_out, aot_out, rtol=0, atol=0)
|
||||
print("Correctness check passed (JIT vs AOT)")
|
||||
|
||||
|
||||
if AOT_AVAILABLE:
|
||||
line_vals = ["jit", "aot"]
|
||||
line_names = ["JIT Kernel", "AOT Kernel"]
|
||||
styles = [("blue", "-"), ("green", "-")]
|
||||
else:
|
||||
line_vals = ["jit"]
|
||||
line_names = ["JIT Kernel"]
|
||||
styles = [("blue", "-")]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["qweight_row", "qweight_col"],
|
||||
x_vals=configs,
|
||||
line_arg="provider",
|
||||
line_vals=line_vals,
|
||||
line_names=line_names,
|
||||
styles=styles,
|
||||
ylabel="us",
|
||||
plot_name="awq-dequantize-jit-vs-aot",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(qweight_row, qweight_col, provider):
|
||||
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
|
||||
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,
|
||||
)
|
||||
|
||||
quantiles = [0.5, 0.2, 0.8]
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: jit_awq_dequantize(qweight, scales, qzeros)
|
||||
elif provider == "aot":
|
||||
fn = lambda: aot_awq_dequantize(qweight, scales, qzeros)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
|
||||
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
check_correctness()
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,133 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
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
|
||||
|
||||
IS_CI = (
|
||||
os.getenv("CI", "false").lower() == "true"
|
||||
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
|
||||
)
|
||||
|
||||
# Fixed parameters
|
||||
NUM_BITS = 4
|
||||
GROUP_SIZE = 128
|
||||
SIZE_N = 4096
|
||||
|
||||
|
||||
def awq_pack(q_w, num_bits, 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)
|
||||
|
||||
|
||||
def make_moe_weights(num_experts, size_k, size_n, num_bits, group_size):
|
||||
pack_factor = 32 // num_bits
|
||||
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, min(group_size, size_k), 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")
|
||||
return b_q_weight, perm
|
||||
|
||||
|
||||
def check_correctness():
|
||||
if not AOT_AVAILABLE:
|
||||
print("sgl_kernel AOT not available, skipping correctness check")
|
||||
return
|
||||
|
||||
num_experts = 4
|
||||
size_k = 1024
|
||||
b_q_weight, perm = make_moe_weights(
|
||||
num_experts, size_k, SIZE_N, NUM_BITS, GROUP_SIZE
|
||||
)
|
||||
|
||||
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()
|
||||
torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
|
||||
print("Correctness check passed (JIT vs AOT)")
|
||||
|
||||
|
||||
if IS_CI:
|
||||
expert_range = [2, 4]
|
||||
else:
|
||||
expert_range = [2, 4, 8, 16]
|
||||
|
||||
if AOT_AVAILABLE:
|
||||
line_vals = ["jit", "aot"]
|
||||
line_names = ["JIT Kernel", "AOT Kernel"]
|
||||
styles = [("blue", "-"), ("green", "-")]
|
||||
else:
|
||||
line_vals = ["jit"]
|
||||
line_names = ["JIT Kernel"]
|
||||
styles = [("blue", "-")]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["num_experts"],
|
||||
x_vals=expert_range,
|
||||
line_arg="provider",
|
||||
line_vals=line_vals,
|
||||
line_names=line_names,
|
||||
styles=styles,
|
||||
ylabel="us",
|
||||
plot_name="awq-marlin-moe-repack-performance",
|
||||
args={"size_k": 4096, "size_n": SIZE_N, "num_bits": NUM_BITS},
|
||||
)
|
||||
)
|
||||
def benchmark(num_experts, size_k, size_n, num_bits, provider):
|
||||
group_size = min(GROUP_SIZE, size_k)
|
||||
b_q_weight, perm = make_moe_weights(
|
||||
num_experts, size_k, size_n, num_bits, group_size
|
||||
)
|
||||
|
||||
quantiles = [0.5, 0.2, 0.8]
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: jit_awq_marlin_moe_repack(
|
||||
b_q_weight, perm, size_k, size_n, num_bits
|
||||
)
|
||||
elif provider == "aot":
|
||||
fn = lambda: aot_awq_marlin_moe_repack(
|
||||
b_q_weight, perm, size_k, size_n, num_bits
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
|
||||
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
check_correctness()
|
||||
benchmark.run(print_data=True)
|
||||
117
python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py
Normal file
117
python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py
Normal file
@@ -0,0 +1,117 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import triton
|
||||
import triton.testing
|
||||
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
|
||||
|
||||
try:
|
||||
from sgl_kernel import awq_marlin_repack as aot_awq_marlin_repack
|
||||
|
||||
AOT_AVAILABLE = True
|
||||
except ImportError:
|
||||
AOT_AVAILABLE = False
|
||||
|
||||
IS_CI = (
|
||||
os.getenv("CI", "false").lower() == "true"
|
||||
or os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
|
||||
)
|
||||
|
||||
# Fixed problem dimensions
|
||||
SIZE_K = 4096
|
||||
SIZE_N = 4096
|
||||
NUM_BITS = 4
|
||||
GROUP_SIZE = 128
|
||||
|
||||
|
||||
def awq_pack(q_w, num_bits, 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)
|
||||
|
||||
|
||||
# Quantize weights once
|
||||
_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)
|
||||
|
||||
|
||||
def check_correctness():
|
||||
if not AOT_AVAILABLE:
|
||||
print("sgl_kernel AOT not available, skipping correctness check")
|
||||
return
|
||||
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()
|
||||
torch.testing.assert_close(out_jit, out_aot, rtol=0, atol=0)
|
||||
print("Correctness check passed (JIT vs AOT)")
|
||||
|
||||
|
||||
if IS_CI:
|
||||
k_range = [1024, 4096]
|
||||
else:
|
||||
k_range = [512, 1024, 2048, 4096, 8192]
|
||||
|
||||
if AOT_AVAILABLE:
|
||||
line_vals = ["jit", "aot"]
|
||||
line_names = ["JIT Kernel", "AOT Kernel"]
|
||||
styles = [("blue", "-"), ("green", "-")]
|
||||
else:
|
||||
line_vals = ["jit"]
|
||||
line_names = ["JIT Kernel"]
|
||||
styles = [("blue", "-")]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["size_k"],
|
||||
x_vals=k_range,
|
||||
line_arg="provider",
|
||||
line_vals=line_vals,
|
||||
line_names=line_names,
|
||||
styles=styles,
|
||||
ylabel="us",
|
||||
plot_name="awq-marlin-repack-performance",
|
||||
args={"size_n": SIZE_N, "num_bits": NUM_BITS},
|
||||
)
|
||||
)
|
||||
def benchmark(size_k, size_n, num_bits, provider):
|
||||
group_size = min(GROUP_SIZE, size_k)
|
||||
|
||||
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)
|
||||
|
||||
quantiles = [0.5, 0.2, 0.8]
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: jit_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
|
||||
elif provider == "aot":
|
||||
fn = lambda: aot_awq_marlin_repack(q_w_awq, size_k, size_n, num_bits)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
|
||||
return 1000 * ms, 1000 * max_ms, 1000 * min_ms
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
check_correctness()
|
||||
benchmark.run(print_data=True)
|
||||
Reference in New Issue
Block a user