[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/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

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

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)

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)

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)

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 <sgl_kernel/tensor.h>
+
+#include <sgl_kernel/utils.cuh>
+
+namespace device::awq {
+
+template <int lut>
+__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<uint32_t*>(&result);
+  uint32_t const i4s = reinterpret_cast<uint32_t const&>(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<uint32_t*>(&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<nv_bfloat162*>(h);
+  res[0] = __hfma2(
+      *reinterpret_cast<nv_bfloat162*>(&lo0),
+      *reinterpret_cast<const nv_bfloat162*>(&MUL),
+      *reinterpret_cast<const nv_bfloat162*>(&ADD));
+  res[1] = __hfma2(
+      *reinterpret_cast<nv_bfloat162*>(&hi0),
+      *reinterpret_cast<const nv_bfloat162*>(&MUL),
+      *reinterpret_cast<const nv_bfloat162*>(&ADD));
+  res[2] = __hfma2(
+      *reinterpret_cast<nv_bfloat162*>(&lo1),
+      *reinterpret_cast<const nv_bfloat162*>(&MUL),
+      *reinterpret_cast<const nv_bfloat162*>(&ADD));
+  res[3] = __hfma2(
+      *reinterpret_cast<nv_bfloat162*>(&hi1),
+      *reinterpret_cast<const nv_bfloat162*>(&MUL),
+      *reinterpret_cast<const nv_bfloat162*>(&ADD));
+
+  return result;
+#else
+  assert(false);
+  return {};
+#endif
+}
+
+template <typename OutputT>
+__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<OutputT, half>::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<OutputT, __nv_bfloat16>::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<uint4*>(scales + scale_offset);
+
+    // Vectorized processing (each uint4 contains 4 nv_bfloat162)
+    nv_bfloat162* weight_vec = reinterpret_cast<nv_bfloat162*>(&weight_raw);
+    nv_bfloat162* zero_vec = reinterpret_cast<nv_bfloat162*>(&zero_raw);
+    nv_bfloat162* scale_vec = reinterpret_cast<nv_bfloat162*>(&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<uint4*>(output_ptr) = weight_raw;
+  }
+}
+
+}  // namespace device::awq
+
+// Host wrapper
+template <typename OutputT>
+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<kDLCUDA>();
+
+  TensorMatcher({qweight_rows, qweight_cols}).with_dtype<int32_t>().with_device(cuda_device).verify(qweight);
+  TensorMatcher({scales_rows, qweight_cols * 8}).with_dtype<OutputT>().with_device(cuda_device).verify(scales);
+  TensorMatcher({scales_rows, qweight_cols}).with_dtype<int32_t>().with_device(cuda_device).verify(qzeros);
+  TensorMatcher({qweight_rows, qweight_cols * 8}).with_dtype<OutputT>().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<int>(qweight_rows / scales_rows);
+  int x_num_threads = 16;
+  int y_num_threads = 16;
+  int x_blocks = (static_cast<int>(qweight_cols) + x_num_threads - 1) / x_num_threads;
+  int y_blocks = (static_cast<int>(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<int*>(qweight.data_ptr());
+  auto* scales_ptr = reinterpret_cast<OutputT*>(scales.data_ptr());
+  auto* qzeros_ptr = reinterpret_cast<int*>(qzeros.data_ptr());
+  auto* output_ptr = reinterpret_cast<OutputT*>(output.data_ptr());
+
+  LaunchKernel(num_blocks, threads_per_block, stream)(
+      device::awq::dequantize_weights<OutputT>,
+      qweight_ptr,
+      scales_ptr,
+      qzeros_ptr,
+      output_ptr,
+      group_size,
+      static_cast<int>(qweight_cols),
+      static_cast<int>(qweight_rows));
+}

python/sglang/jit_kernel/csrc/gemm/marlin/awq_marlin_repack.cuh

@@ -0,0 +1,251 @@
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include "marlin.cuh"
+
+namespace device::marlin {
+
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
+template <int const num_threads, int const num_bits>
+__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 <int const num_threads, int const num_bits>
+__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<repack_stages - 2>();
+    __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<int4 const*>(
+              &(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<uint32_t*>(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<kDLCUDA>();
+
+  TensorMatcher({size_k, size_n / pack_factor}).with_dtype<int32_t>().with_device(cuda_device).verify(b_q_weight);
+
+  TensorMatcher({size_k / tile_size, size_n * tile_size / pack_factor})
+      .with_dtype<int32_t>()
+      .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<uint32_t const*>(b_q_weight.data_ptr());
+  auto* out_ptr = reinterpret_cast<uint32_t*>(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<repack_threads, 4>, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);
+    LaunchKernel(blocks, repack_threads, stream, max_shared_mem)(
+        awq_marlin_repack_kernel<repack_threads, 4>,
+        b_q_weight_ptr,
+        out_ptr,
+        static_cast<int>(size_k),
+        static_cast<int>(size_n));
+  } else if (num_bits == 8) {
+    cudaFuncSetAttribute(
+        awq_marlin_repack_kernel<repack_threads, 8>, cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem);
+    LaunchKernel(blocks, repack_threads, stream, max_shared_mem)(
+        awq_marlin_repack_kernel<repack_threads, 8>,
+        b_q_weight_ptr,
+        out_ptr,
+        static_cast<int>(size_k),
+        static_cast<int>(size_n));
+  } else {
+    RuntimeCheck(false, "Unsupported repack config: num_bits = ", num_bits);
+  }
+}

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__])

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__)])

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__)])

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: