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

Co-authored-by: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>

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)