diff --git a/python/sglang/jit_kernel/awq_dequantize.py b/python/sglang/jit_kernel/awq_dequantize.py new file mode 100644 index 000000000..4a188c02e --- /dev/null +++ b/python/sglang/jit_kernel/awq_dequantize.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_awq_dequantize_module(dtype: torch.dtype) -> Module: + args = make_cpp_args(dtype) + return load_jit( + "awq_dequantize", + *args, + cuda_files=["gemm/awq_dequantize.cuh"], + cuda_wrappers=[("awq_dequantize", f"awq_dequantize<{args}>")], + ) + + +def awq_dequantize( + qweight: torch.Tensor, + scales: torch.Tensor, + qzeros: torch.Tensor, +) -> torch.Tensor: + qweight_rows = qweight.shape[0] + qweight_cols = qweight.shape[1] + output = torch.empty( + (qweight_rows, qweight_cols * 8), + dtype=scales.dtype, + device=scales.device, + ) + module = _jit_awq_dequantize_module(scales.dtype) + module.awq_dequantize(output, qweight, scales, qzeros) + return output diff --git a/python/sglang/jit_kernel/awq_marlin_repack.py b/python/sglang/jit_kernel/awq_marlin_repack.py new file mode 100644 index 000000000..3b06144cb --- /dev/null +++ b/python/sglang/jit_kernel/awq_marlin_repack.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import cache_once, load_jit + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_awq_marlin_repack_module() -> Module: + return load_jit( + "awq_marlin_repack", + cuda_files=["gemm/marlin/awq_marlin_repack.cuh"], + cuda_wrappers=[("awq_marlin_repack", "awq_marlin_repack")], + ) + + +def awq_marlin_repack( + b_q_weight: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, +) -> torch.Tensor: + tile_size = 16 + pack_factor = 32 // num_bits + out = torch.empty( + (size_k // tile_size, size_n * tile_size // pack_factor), + dtype=b_q_weight.dtype, + device=b_q_weight.device, + ) + module = _jit_awq_marlin_repack_module() + module.awq_marlin_repack(out, b_q_weight, size_k, size_n, num_bits) + return out + + +def awq_marlin_moe_repack( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, +) -> torch.Tensor: + num_experts = b_q_weight.shape[0] + assert size_k % 16 == 0 + output = torch.empty( + (num_experts, size_k // 16, size_n * (num_bits // 2)), + device=b_q_weight.device, + dtype=b_q_weight.dtype, + ) + for e in range(num_experts): + output[e] = awq_marlin_repack(b_q_weight[e], size_k, size_n, num_bits) + return output diff --git a/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py b/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py new file mode 100644 index 000000000..09b6ccb3f --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_awq_dequantize.py @@ -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) diff --git a/python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py new file mode 100644 index 000000000..120c177d5 --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_moe_repack.py @@ -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) diff --git a/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py new file mode 100644 index 000000000..51403363d --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_awq_marlin_repack.py @@ -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) diff --git a/python/sglang/jit_kernel/csrc/gemm/awq_dequantize.cuh b/python/sglang/jit_kernel/csrc/gemm/awq_dequantize.cuh new file mode 100644 index 000000000..ac6b9a5ff --- /dev/null +++ b/python/sglang/jit_kernel/csrc/gemm/awq_dequantize.cuh @@ -0,0 +1,227 @@ +// Adapted from +// https://github.com/vllm-project/vllm/blob/eb59b5a6cba6727d3727c0372258db9002f687c1/csrc/quantization/awq/gemm_kernels.cu#L350 +#pragma once + +#include + +#include + +namespace device::awq { + +template +__device__ inline int lop3(int a, int b, int c) { + int res; + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" : "=r"(res) : "r"(a), "r"(b), "r"(c), "n"(lut)); + return res; +} + +__device__ uint4 dequantize_s4_to_fp16x2(uint32_t const& source) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750 + uint4 result; + + uint32_t* h = reinterpret_cast(&result); + uint32_t const i4s = reinterpret_cast(source); + + // First, we extract the i4s and construct an intermediate fp16 number. + static constexpr uint32_t immLut = (0xf0 & 0xcc) | 0xaa; + static constexpr uint32_t BOTTOM_MASK = 0x000f000f; + static constexpr uint32_t TOP_MASK = 0x00f000f0; + static constexpr uint32_t I4s_TO_F16s_MAGIC_NUM = 0x64006400; + + // Shift right by 8 to now consider elt_45 and elt_67. Issue first to hide RAW + // dependency if we issue immediately before required. + const uint32_t top_i4s = i4s >> 8; + // Extract elt_01 - (i4s & 0x000f000f) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[0]) + : "r"(i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_23 (i4s & 0x00f000f0) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[1]) + : "r"(i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_45 (top_i4s & 0x000f000f) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[2]) + : "r"(top_i4s), "n"(BOTTOM_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + // Extract elt_67 (top_i4s & 0x00f000f0) | 0x64006400 + asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n" + : "=r"(h[3]) + : "r"(top_i4s), "n"(TOP_MASK), "n"(I4s_TO_F16s_MAGIC_NUM), "n"(immLut)); + + // This is the half2 {1024, 1024} represented as an integer. + static constexpr uint32_t FP16_TOP_MAGIC_NUM = 0x64006400; + // This is the half2 {1 / 16, 1 / 16} represented as an integer. + static constexpr uint32_t ONE_SIXTEENTH = 0x2c002c00; + // This is the half2 {-64, -64} represented as an integer. + static constexpr uint32_t NEG_64 = 0xd400d400; + + // Finally, we construct the output numbers. + // Convert elt_01 + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[0]) : "r"(h[0]), "r"(FP16_TOP_MAGIC_NUM)); + // Convert elt_23 + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[1]) : "r"(h[1]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); + // Convert elt_45 + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(h[2]) : "r"(h[2]), "r"(FP16_TOP_MAGIC_NUM)); + // Convert elt_67 + asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(h[3]) : "r"(h[3]), "r"(ONE_SIXTEENTH), "r"(NEG_64)); + + return result; +#else + assert(false); + return {}; +#endif +} + +__device__ uint4 dequantize_s4_to_bf16x2(uint32_t const& source) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + uint4 result; + uint32_t* h = reinterpret_cast(&result); + uint32_t const i4s = source; + + // Define masks and constants + static constexpr uint32_t MASK = 0x000f000f; + static constexpr uint32_t EX = 0x43004300; + static constexpr uint32_t MUL = 0x3F803F80; + static constexpr uint32_t ADD = 0xC300C300; + + int lo0 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s, MASK, EX); + int hi0 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 4, MASK, EX); + int lo1 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 8, MASK, EX); + int hi1 = lop3<(0xf0 & 0xcc) | 0xaa>(i4s >> 12, MASK, EX); + + nv_bfloat162* res = reinterpret_cast(h); + res[0] = __hfma2( + *reinterpret_cast(&lo0), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + res[1] = __hfma2( + *reinterpret_cast(&hi0), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + res[2] = __hfma2( + *reinterpret_cast(&lo1), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + res[3] = __hfma2( + *reinterpret_cast(&hi1), + *reinterpret_cast(&MUL), + *reinterpret_cast(&ADD)); + + return result; +#else + assert(false); + return {}; +#endif +} + +template +__global__ void __launch_bounds__(256) dequantize_weights( + int* __restrict__ qweight, + OutputT* __restrict__ scales, + int* __restrict__ qzeros, + OutputT* __restrict__ output, + int group_size, + int qweight_cols, + int qweight_rows) { + int col = blockIdx.x * blockDim.x + threadIdx.x; + int row = blockIdx.y * blockDim.y + threadIdx.y; + if (col >= qweight_cols || row >= qweight_rows) return; + + int group_idx = row / group_size; + int scale_offset = 8 * col + group_idx * qweight_cols * 8; + uint4 loaded_scale = *(uint4*)(scales + scale_offset); + + // Handle different data types + if constexpr (std::is_same::value) { + // FP16 path + uint4 zeros = dequantize_s4_to_fp16x2(qzeros[col + group_idx * qweight_cols]); + uint4 weight_fp16 = dequantize_s4_to_fp16x2(qweight[col + row * qweight_cols]); + + // Use PTX assembly for FP16 operations + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.x) : "r"(weight_fp16.x), "r"(zeros.x)); + asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.x) : "r"(weight_fp16.x), "r"(loaded_scale.x)); + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.y) : "r"(weight_fp16.y), "r"(zeros.y)); + asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.y) : "r"(weight_fp16.y), "r"(loaded_scale.y)); + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.z) : "r"(weight_fp16.z), "r"(zeros.z)); + asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.z) : "r"(weight_fp16.z), "r"(loaded_scale.z)); + asm volatile("sub.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.w) : "r"(weight_fp16.w), "r"(zeros.w)); + asm volatile("mul.rn.f16x2 %0, %1, %2;\n" : "=r"(weight_fp16.w) : "r"(weight_fp16.w), "r"(loaded_scale.w)); + + OutputT* output_ptr = output + 8 * col + 8 * row * qweight_cols; + *(uint4*)output_ptr = weight_fp16; + } else if constexpr (std::is_same::value) { + uint4 weight_raw = dequantize_s4_to_bf16x2(qweight[col + row * qweight_cols]); + uint4 zero_raw = dequantize_s4_to_bf16x2(qzeros[col + group_idx * qweight_cols]); + uint4 scale_raw = *reinterpret_cast(scales + scale_offset); + + // Vectorized processing (each uint4 contains 4 nv_bfloat162) + nv_bfloat162* weight_vec = reinterpret_cast(&weight_raw); + nv_bfloat162* zero_vec = reinterpret_cast(&zero_raw); + nv_bfloat162* scale_vec = reinterpret_cast(&scale_raw); + +// Single instruction dual-channel operation +#pragma unroll + for (int i = 0; i < 4; ++i) { // uint4 = 4 * nv_bfloat162 + weight_vec[i] = __hmul2(__hsub2(weight_vec[i], zero_vec[i]), scale_vec[i]); + } + + // Directly store to OutputT array (guaranteed contiguous memory) + OutputT* output_ptr = output + 8 * col + row * qweight_cols * 8; + static_assert(sizeof(uint4) == 8 * sizeof(OutputT), "Memory layout mismatch"); + *reinterpret_cast(output_ptr) = weight_raw; + } +} + +} // namespace device::awq + +// Host wrapper +template +void awq_dequantize( + tvm::ffi::TensorView output, + tvm::ffi::TensorView qweight, + tvm::ffi::TensorView scales, + tvm::ffi::TensorView qzeros) { + using namespace host; + + int64_t qweight_rows = qweight.size(0); + int64_t qweight_cols = qweight.size(1); + int64_t scales_rows = scales.size(0); + + // Validate tensors + SymbolicDevice cuda_device; + cuda_device.set_options(); + + TensorMatcher({qweight_rows, qweight_cols}).with_dtype().with_device(cuda_device).verify(qweight); + TensorMatcher({scales_rows, qweight_cols * 8}).with_dtype().with_device(cuda_device).verify(scales); + TensorMatcher({scales_rows, qweight_cols}).with_dtype().with_device(cuda_device).verify(qzeros); + TensorMatcher({qweight_rows, qweight_cols * 8}).with_dtype().with_device(cuda_device).verify(output); + + // Get device and stream + auto device = cuda_device.unwrap(); + auto stream = LaunchKernel::resolve_device(device); + + int group_size = static_cast(qweight_rows / scales_rows); + int x_num_threads = 16; + int y_num_threads = 16; + int x_blocks = (static_cast(qweight_cols) + x_num_threads - 1) / x_num_threads; + int y_blocks = (static_cast(qweight_rows) + y_num_threads - 1) / y_num_threads; + + dim3 num_blocks(x_blocks, y_blocks); + dim3 threads_per_block(x_num_threads, y_num_threads); + + // Get pointers + auto* qweight_ptr = reinterpret_cast(qweight.data_ptr()); + auto* scales_ptr = reinterpret_cast(scales.data_ptr()); + auto* qzeros_ptr = reinterpret_cast(qzeros.data_ptr()); + auto* output_ptr = reinterpret_cast(output.data_ptr()); + + LaunchKernel(num_blocks, threads_per_block, stream)( + device::awq::dequantize_weights, + qweight_ptr, + scales_ptr, + qzeros_ptr, + output_ptr, + group_size, + static_cast(qweight_cols), + static_cast(qweight_rows)); +} diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin/awq_marlin_repack.cuh b/python/sglang/jit_kernel/csrc/gemm/marlin/awq_marlin_repack.cuh new file mode 100644 index 000000000..7f1735433 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/gemm/marlin/awq_marlin_repack.cuh @@ -0,0 +1,251 @@ +#pragma once + +#include + +#include + +#include "marlin.cuh" + +namespace device::marlin { + +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 +template +__global__ void awq_marlin_repack_kernel( + uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) { + return; +} +#else + +template +__global__ void awq_marlin_repack_kernel( + uint32_t const* __restrict__ b_q_weight_ptr, uint32_t* __restrict__ out_ptr, int size_k, int size_n) { + constexpr int pack_factor = 32 / num_bits; + + int k_tiles = size_k / tile_k_size; + int n_tiles = size_n / tile_n_size; + int block_k_tiles = div_ceil(k_tiles, (int)gridDim.x); + + auto start_k_tile = blockIdx.x * block_k_tiles; + if (start_k_tile >= k_tiles) { + return; + } + + int finish_k_tile = min(start_k_tile + block_k_tiles, k_tiles); + + // Wait until the next thread tile has been loaded to shared memory. + auto wait_for_stage = [&]() { + // We only have `stages - 2` active fetches since we are double buffering + // and can only issue the next fetch when it is guaranteed that the previous + // shared memory load is fully complete (as it may otherwise be + // overwritten). + cp_async_wait(); + __syncthreads(); + }; + + extern __shared__ int4 sh[]; + + constexpr int tile_n_ints = tile_n_size / pack_factor; + + constexpr int stage_n_threads = tile_n_ints / 4; + constexpr int stage_k_threads = tile_k_size; + constexpr int stage_size = stage_k_threads * stage_n_threads; + + auto fetch_to_shared = [&](int pipe, int k_tile_id, int n_tile_id) { + if (n_tile_id >= n_tiles) { + cp_async_fence(); + return; + } + + int first_n = n_tile_id * tile_n_size; + int first_n_packed = first_n / pack_factor; + + int4* sh_ptr = sh + stage_size * pipe; + + if (threadIdx.x < stage_size) { + auto k_id = threadIdx.x / stage_n_threads; + auto n_id = threadIdx.x % stage_n_threads; + + int first_k = k_tile_id * tile_k_size; + + cp_async4( + &sh_ptr[k_id * stage_n_threads + n_id], + reinterpret_cast( + &(b_q_weight_ptr[(first_k + k_id) * (size_n / pack_factor) + first_n_packed + (n_id * 4)]))); + } + + cp_async_fence(); + }; + + auto repack_tile = [&](int pipe, int k_tile_id, int n_tile_id) { + if (n_tile_id >= n_tiles) { + return; + } + + auto warp_id = threadIdx.x / 32; + auto th_id = threadIdx.x % 32; + + if (warp_id >= 4) { + return; + } + + int tc_col = th_id / 4; + int tc_row = (th_id % 4) * 2; + + constexpr int tc_offsets[4] = {0, 1, 8, 9}; + + int cur_n = warp_id * 16 + tc_col; + int cur_n_packed = cur_n / pack_factor; + int cur_n_pos = cur_n % pack_factor; + + constexpr int sh_stride = tile_n_ints; + constexpr uint32_t mask = (1 << num_bits) - 1; + + int4* sh_stage_ptr = sh + stage_size * pipe; + uint32_t* sh_stage_int_ptr = reinterpret_cast(sh_stage_ptr); + + // Undo interleaving + int cur_n_pos_unpacked; + if constexpr (num_bits == 4) { + constexpr int undo_pack[8] = {0, 4, 1, 5, 2, 6, 3, 7}; + cur_n_pos_unpacked = undo_pack[cur_n_pos]; + } else { + constexpr int undo_pack[4] = {0, 2, 1, 3}; + cur_n_pos_unpacked = undo_pack[cur_n_pos]; + } + + uint32_t vals[8]; +#pragma unroll + for (int i = 0; i < 4; i++) { + int cur_elem = tc_row + tc_offsets[i]; + + int packed_src_0 = sh_stage_int_ptr[cur_n_packed + sh_stride * cur_elem]; + int packed_src_1 = sh_stage_int_ptr[cur_n_packed + (8 / pack_factor) + sh_stride * cur_elem]; + + vals[i] = (packed_src_0 >> (cur_n_pos_unpacked * num_bits)) & mask; + vals[4 + i] = (packed_src_1 >> (cur_n_pos_unpacked * num_bits)) & mask; + } + + constexpr int tile_size_val = tile_k_size * tile_n_size / pack_factor; + int out_offset = (k_tile_id * n_tiles + n_tile_id) * tile_size_val; + + // Result of: + // https://github.com/NVIDIA/FasterTransformer/blob/main/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h + if constexpr (num_bits == 4) { + constexpr int pack_idx[8] = {0, 2, 4, 6, 1, 3, 5, 7}; + + uint32_t res = 0; +#pragma unroll + for (int i = 0; i < 8; i++) { + res |= vals[pack_idx[i]] << (i * 4); + } + + out_ptr[out_offset + th_id * 4 + warp_id] = res; + + } else { + constexpr int pack_idx[4] = {0, 2, 1, 3}; + + uint32_t res1 = 0; + uint32_t res2 = 0; +#pragma unroll + for (int i = 0; i < 4; i++) { + res1 |= vals[pack_idx[i]] << (i * 8); + res2 |= vals[4 + pack_idx[i]] << (i * 8); + } + + out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 0] = res1; + out_ptr[out_offset + th_id * 8 + (warp_id * 2) + 1] = res2; + } + }; + + auto start_pipes = [&](int k_tile_id, int n_tile_id) { +#pragma unroll + for (int pipe = 0; pipe < repack_stages - 1; pipe++) { + fetch_to_shared(pipe, k_tile_id, n_tile_id + pipe); + } + + wait_for_stage(); + }; +#pragma unroll + for (int k_tile_id = start_k_tile; k_tile_id < finish_k_tile; k_tile_id++) { + int n_tile_id = 0; + + start_pipes(k_tile_id, n_tile_id); + + while (n_tile_id < n_tiles) { +#pragma unroll + for (int pipe = 0; pipe < repack_stages; pipe++) { + fetch_to_shared((pipe + repack_stages - 1) % repack_stages, k_tile_id, n_tile_id + pipe + repack_stages - 1); + repack_tile(pipe, k_tile_id, n_tile_id + pipe); + wait_for_stage(); + } + n_tile_id += repack_stages; + } + } +} +#endif + +} // namespace device::marlin + +// Host wrapper +void awq_marlin_repack( + tvm::ffi::TensorView out, tvm::ffi::TensorView b_q_weight, int64_t size_k, int64_t size_n, int64_t num_bits) { + using namespace host; + using namespace device::marlin; + + // Validate alignment + RuntimeCheck(size_k % tile_k_size == 0, "size_k = ", size_k, " is not divisible by tile_k_size = ", tile_k_size); + RuntimeCheck(size_n % tile_n_size == 0, "size_n = ", size_n, " is not divisible by tile_n_size = ", tile_n_size); + RuntimeCheck(num_bits == 4 || num_bits == 8, "num_bits must be 4 or 8. Got = ", num_bits); + + int const pack_factor = 32 / num_bits; + + // Validate tensors + SymbolicDevice cuda_device; + cuda_device.set_options(); + + TensorMatcher({size_k, size_n / pack_factor}).with_dtype().with_device(cuda_device).verify(b_q_weight); + + TensorMatcher({size_k / tile_size, size_n * tile_size / pack_factor}) + .with_dtype() + .with_device(cuda_device) + .verify(out); + + // Get device and stream + auto device = cuda_device.unwrap(); + auto stream = LaunchKernel::resolve_device(device); + + // Get pointers + auto* b_q_weight_ptr = reinterpret_cast(b_q_weight.data_ptr()); + auto* out_ptr = reinterpret_cast(out.data_ptr()); + + // Get device attributes + int blocks = 0; + cudaDeviceGetAttribute(&blocks, cudaDevAttrMultiProcessorCount, device.device_id); + + int max_shared_mem = 0; + cudaDeviceGetAttribute(&max_shared_mem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device.device_id); + RuntimeCheck(max_shared_mem > 0, "max_shared_mem must be > 0"); + + // Dispatch based on num_bits + if (num_bits == 4) { + cudaFuncSetAttribute( + awq_marlin_repack_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); + LaunchKernel(blocks, repack_threads, stream, max_shared_mem)( + awq_marlin_repack_kernel, + b_q_weight_ptr, + out_ptr, + static_cast(size_k), + static_cast(size_n)); + } else if (num_bits == 8) { + cudaFuncSetAttribute( + awq_marlin_repack_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); + LaunchKernel(blocks, repack_threads, stream, max_shared_mem)( + awq_marlin_repack_kernel, + b_q_weight_ptr, + out_ptr, + static_cast(size_k), + static_cast(size_n)); + } else { + RuntimeCheck(false, "Unsupported repack config: num_bits = ", num_bits); + } +} diff --git a/python/sglang/jit_kernel/tests/test_awq_dequantize.py b/python/sglang/jit_kernel/tests/test_awq_dequantize.py new file mode 100644 index 000000000..e29475843 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_awq_dequantize.py @@ -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__]) diff --git a/python/sglang/jit_kernel/tests/test_awq_marlin_moe_repack.py b/python/sglang/jit_kernel/tests/test_awq_marlin_moe_repack.py new file mode 100644 index 000000000..217dfc0a6 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_awq_marlin_moe_repack.py @@ -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__)]) diff --git a/python/sglang/jit_kernel/tests/test_awq_marlin_repack.py b/python/sglang/jit_kernel/tests/test_awq_marlin_repack.py new file mode 100644 index 000000000..819fcf276 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_awq_marlin_repack.py @@ -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__)]) diff --git a/python/sglang/srt/layers/quantization/awq.py b/python/sglang/srt/layers/quantization/awq.py index 173bc2bb2..42d3a6c4d 100644 --- a/python/sglang/srt/layers/quantization/awq.py +++ b/python/sglang/srt/layers/quantization/awq.py @@ -60,7 +60,11 @@ if _is_npu: import torch_npu if _is_cuda: - from sgl_kernel import awq_dequantize, awq_marlin_moe_repack, awq_marlin_repack + from sglang.jit_kernel.awq_dequantize import awq_dequantize + from sglang.jit_kernel.awq_marlin_repack import ( + awq_marlin_moe_repack, + awq_marlin_repack, + ) elif _is_hip: