Repository: sgl-project/sglang
Pull request #4104: Add awq dequantize kernel to sgl with 1x to 3x speedup
Author: login: zcnrex
Labels: name: high priority
Created: 2025-03-05T19:27:39Z | Merged: 2025-03-12T07:10:02Z

<!-- Thank you for your contribution! We appreciate it. The following guidelines will help improve your pull request and facilitate feedback. If anything is unclear, don't hesitate to submit your pull request and ask the maintainers for assistance. -->

## Motivation
#2965 


Adapted from [vllm awq_dequant](https://github.com/vllm-project/vllm/blob/eb59b5a6cba6727d3727c0372258db9002f687c1/csrc/quantization/awq/gemm_kernels.cu#L413)

Achieved 1x - 4x speed up compared to vllm awq_dequantize cuda implementation.

The performance improvement comes from removing the copy & for loop [in the end](https://github.com/vllm-project/vllm/blob/eb59b5a6cba6727d3727c0372258db9002f687c1/csrc/quantization/awq/gemm_kernels.cu#L403) when assigning the output.

```
✅ All implementations match
awq-dequantize-performance:
    qweight_row  qweight_col         VLLM  SGL Kernel
0        3584.0        448.0    80.863997   27.519999
1        3584.0        576.0    96.832000   32.736000
2        3584.0       4736.0   669.727981  182.528004
3        3584.0         16.0    15.584000   13.296000
4        3584.0         32.0    17.983999   13.888000
5        3584.0         64.0    20.800000   14.784000
6        3584.0        128.0    29.472001   17.759999
7       18944.0        448.0   356.128007   99.232003
8       18944.0        576.0   441.632003  123.007998
9       18944.0       4736.0  3465.471983  897.119999
10      18944.0         16.0    24.992000   16.640000
11      18944.0         32.0    33.567999   19.040000
12      18944.0         64.0    49.791999   23.871999
13      18944.0        128.0    82.176000   36.511999
14        128.0        448.0    15.823999   13.584000
15        128.0        576.0    16.783999   13.824000
16        128.0       4736.0    39.584000   19.616000
17        128.0         16.0    12.992000   12.800000
18        128.0         32.0    12.896000   12.704000
19        128.0         64.0    12.992000   12.896000
20        128.0        128.0    13.504000   13.200000
21        256.0        448.0    19.040000   14.240000
22        256.0        576.0    20.352000   14.640000
23        256.0       4736.0    63.391998   24.480000
24        256.0         16.0    12.832000   12.720000
25        256.0         32.0    12.912000   12.832000
26        256.0         64.0    13.376000   13.056000
27        256.0        128.0    14.144000   13.312000
28        512.0        448.0    23.871999   15.328000
29        512.0        576.0    28.031999   16.832000
30        512.0       4736.0   110.048003   37.023999
31        512.0         16.0    12.896000   12.704000
32        512.0         32.0    13.264000   12.992000
33        512.0         64.0    14.144000   13.312000
34        512.0        128.0    15.552000   13.328000
35       1024.0        448.0    35.007998   18.224001
36       1024.0        576.0    39.391998   19.328000
37       1024.0       4736.0   203.455999   62.015999
38       1024.0         16.0    13.280000   13.024000
39       1024.0         32.0    14.464000   13.280000
40       1024.0         64.0    15.488000   13.216000
41       1024.0        128.0    18.271999   14.368000
```


There's also a [triton implementation](https://github.com/vllm-project/vllm/blob/6eaf93020dc467ccd98f107fbab99d958e0b442f/vllm/model_executor/layers/quantization/awq_triton.py#L11) that we could port later.

Slightly better performance compared to triton implementation on H100
```
$ export VLLM_USE_TRITON_AWQ=1

✅ All implementations match
awq-dequantize-performance:
    qweight_row  qweight_col        VLLM  SGL Kernel
0        3584.0        448.0   29.200001   28.160000
1        3584.0        576.0   33.696000   33.376001
2        3584.0       4736.0  183.824003  183.295995
3        3584.0         16.0   14.816000   13.920000
4        3584.0         32.0   15.424000   14.848000
5        3584.0         64.0   16.319999   15.520000
6        3584.0        128.0   19.743999   18.880000
7       18944.0        448.0   99.583998   99.808000
8       18944.0        576.0  123.680003  123.712003
9       18944.0       4736.0  898.576021  897.840023
10      18944.0         16.0   19.007999   17.664000
11      18944.0         32.0   20.608000   19.776000
12      18944.0         64.0   26.303999   24.607999
13      18944.0        128.0   37.344001   37.280001
14        128.0        448.0   15.040000   14.144000
15        128.0        576.0   15.200000   14.304000
16        128.0       4736.0   21.183999   20.288000
17        128.0         16.0   14.272000   13.472000
18        128.0         32.0   14.432000   13.488000
19        128.0         64.0   14.592000   13.664000
20        128.0        128.0   14.784000   13.728000
21        256.0        448.0   15.840000   14.976000
22        256.0        576.0   16.031999   15.168000
23        256.0       4736.0   26.880000   25.119999
24        256.0         16.0   14.368000   13.472000
25        256.0         32.0   14.560000   13.600000
26        256.0         64.0   14.752000   13.664000
27        256.0        128.0   14.976000   13.920000
28        512.0        448.0   16.480001   15.807999
29        512.0        576.0   18.656000   17.759999
30        512.0       4736.0   37.888002   37.664000
31        512.0         16.0   14.464000   13.568000
32        512.0         32.0   14.752000   13.664000
33        512.0         64.0   14.976000   13.888000
34        512.0        128.0   15.008000   14.016000
35       1024.0        448.0   19.904001   19.168001
36       1024.0        576.0   20.864001   19.936001
37       1024.0       4736.0   63.167997   62.816001
38       1024.0         16.0   14.560000   13.664000
39       1024.0         32.0   14.976000   13.920000
40       1024.0         64.0   14.848000   13.952000
41       1024.0        128.0   15.680000   14.944000
```

<!-- Explain the purpose of this PR and the goals it aims to achieve. -->

## Modifications

<!-- Describe the changes made in this PR. -->

## Checklist

- [ ] Format your code according to the [Code Formatting with Pre-Commit](https://docs.sglang.ai/references/contribution_guide.html#code-formatting-with-pre-commit).
- [ ] Add unit tests as outlined in the [Running Unit Tests](https://docs.sglang.ai/references/contribution_guide.html#running-unit-tests-adding-to-ci).
- [ ] Update documentation / docstrings / example tutorials as needed, according to [Writing Documentation](https://docs.sglang.ai/references/contribution_guide.html#writing-documentation-running-docs-ci).
- [ ] Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to [Benchmark and Profiling](https://docs.sglang.ai/references/benchmark_and_profiling.html) and [Accuracy Results](https://docs.sglang.ai/references/accuracy_evaluation.html).
- [ ] For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
- [ ] Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.



Diff:
diff --git a/sgl-kernel/benchmark/bench_awq_dequant.py b/sgl-kernel/benchmark/bench_awq_dequant.py
new file mode 100644
index 000000000000..22280c250fe5
--- /dev/null
+++ b/sgl-kernel/benchmark/bench_awq_dequant.py
@@ -0,0 +1,118 @@
+import itertools
+from typing import List, Tuple
+
+import torch
+import triton
+import triton.testing
+from sgl_kernel import awq_dequantize
+from vllm import _custom_ops as ops
+
+
+def vllm_awq_dequantize(
+    qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    return ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
+
+
+def sglang_awq_dequantize(
+    qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+    return awq_dequantize(qweight, scales, qzeros)
+
+
+def calculate_diff(qweight_row: int, qweight_col: int):
+    """Calculate difference between VLLM and SGLang implementations."""
+    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,
+    )
+
+    vllm_out = vllm_awq_dequantize(qweight, scales, qzeros)
+    sglang_out = sglang_awq_dequantize(qweight, scales, qzeros)
+
+    output_diff = torch.abs(vllm_out.float() - sglang_out.float()).mean().item()
+
+    if torch.allclose(
+        vllm_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5
+    ):
+        print("✅ All implementations match")
+    else:
+        print("❌ Implementations differ")
+
+
+qweight_row_range = [3584, 18944, 128, 256, 512, 1024]
+qweight_cols_range = [448, 576, 4736, 16, 32, 64, 128]
+
+configs = list(itertools.product(qweight_row_range, qweight_cols_range))
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["qweight_row", "qweight_col"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=["vllm", "sglang"],
+        line_names=["VLLM", "SGL Kernel"],
+        styles=[("blue", "-"), ("green", "-")],
+        ylabel="us",
+        plot_name="awq-dequantize-performance",
+        args={},
+    )
+)
+def benchmark(qweight_row, qweight_col, provider):
+    dtype = torch.float16
+    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 == "vllm":
+        fn = lambda: vllm_awq_dequantize(
+            qweight.clone(), scales.clone(), qzeros.clone()
+        )
+    elif provider == "sglang":
+        fn = lambda: sglang_awq_dequantize(
+            qweight.clone(), scales.clone(), qzeros.clone()
+        )
+
+    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    calculate_diff(qweight_row=3584, qweight_col=448)
+    benchmark.run(print_data=True)
diff --git a/sgl-kernel/csrc/gemm/awq_kernel.cu b/sgl-kernel/csrc/gemm/awq_kernel.cu
new file mode 100644
index 000000000000..2b697cae4016
--- /dev/null
+++ b/sgl-kernel/csrc/gemm/awq_kernel.cu
@@ -0,0 +1,127 @@
+// Adapted from
+// https://github.com/vllm-project/vllm/blob/eb59b5a6cba6727d3727c0372258db9002f687c1/csrc/quantization/awq/gemm_kernels.cu#L350
+#include <c10/cuda/CUDAGuard.h>
+#include <cuda_fp16.h>
+#include <torch/all.h>
+
+__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;
+
+  // Note that the entire sequence only requires 1 shift instruction. This is
+  // thanks to the register packing format and the fact that we force our
+  // integers to be unsigned, and account for this in the fp16 subtractions. In
+  // addition, I exploit the fact that sub and fma have the same throughput in
+  // order to convert elt_23 and elt_67 to fp16 without having to shift them to
+  // the bottom bits before hand.
+
+  // 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
+}
+
+__global__ void __launch_bounds__(256) dequantize_weights(
+    int* __restrict__ qweight,
+    half* __restrict__ scales,
+    int* __restrict__ qzeros,
+    half* __restrict__ output,
+    int group_size,
+    int qweight_cols) {
+  int col = blockIdx.x * blockDim.x + threadIdx.x;
+  int row = blockIdx.y * blockDim.y + threadIdx.y;
+
+  uint4 zeros = dequantize_s4_to_fp16x2(qzeros[col + (row / group_size) * qweight_cols]);
+  uint4 loaded_scale = *(uint4*)(scales + 8 * col + (row / group_size) * qweight_cols * 8);
+
+  uint4 weight_fp16 = dequantize_s4_to_fp16x2(qweight[col + row * qweight_cols]);
+
+  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));
+
+  half* output_ptr = output + 8 * col + 8 * row * qweight_cols;
+  *(uint4*)output_ptr = weight_fp16;
+}
+
+torch::Tensor awq_dequantize(torch::Tensor qweight, torch::Tensor scales, torch::Tensor qzeros) {
+  int qweight_rows = qweight.size(0);
+  int qweight_cols = qweight.size(1);
+  int group_size = qweight_rows / scales.size(0);
+
+  int x_num_threads = 16;
+  int y_num_threads = 16;
+  int x_blocks = qweight_cols / x_num_threads;
+  int y_blocks = qweight_rows / y_num_threads;
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(qweight));
+
+  auto output_tensor_options = torch::TensorOptions().dtype(scales.dtype()).device(scales.device());
+  at::Tensor output = torch::empty({qweight_rows, qweight_cols * 8}, output_tensor_options);
+
+  auto _qweight = reinterpret_cast<int*>(qweight.data_ptr<int>());
+  auto _scales = reinterpret_cast<half*>(scales.data_ptr<at::Half>());
+  auto _zeros = reinterpret_cast<int*>(qzeros.data_ptr<int>());
+  auto _output = reinterpret_cast<half*>(output.data_ptr<at::Half>());
+
+  dim3 num_blocks(x_blocks, y_blocks);
+  dim3 threads_per_block(x_num_threads, y_num_threads);
+
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  dequantize_weights<<<num_blocks, threads_per_block, 0, stream>>>(
+      _qweight, _scales, _zeros, _output, group_size, qweight_cols);
+
+  return output;
+}
diff --git a/sgl-kernel/csrc/torch_extension.cc b/sgl-kernel/csrc/torch_extension.cc
index d8bd89917cac..a5a115e3adfc 100644
--- a/sgl-kernel/csrc/torch_extension.cc
+++ b/sgl-kernel/csrc/torch_extension.cc
@@ -75,6 +75,9 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) {
   /*
    * From csrc/gemm
    */
+  m.def("awq_dequantize(Tensor qweight, Tensor scales, Tensor qzeros) -> Tensor");
+  m.impl("awq_dequantize", torch::kCUDA, &awq_dequantize);
+
   m.def(
       "int8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? "
       "bias) -> Tensor");
diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h
index 34ce443a2012..934478eec1fb 100644
--- a/sgl-kernel/include/sgl_kernel_ops.h
+++ b/sgl-kernel/include/sgl_kernel_ops.h
@@ -112,6 +112,7 @@ void apply_rope_pos_ids_cos_sin_cache(
 /*
  * From csrc/gemm
  */
+torch::Tensor awq_dequantize(torch::Tensor qweight, torch::Tensor scales, torch::Tensor qzeros);
 torch::Tensor int8_scaled_mm(
     const torch::Tensor& mat_a,
     const torch::Tensor& mat_b,
diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py
index c8cb0443d686..da2959269332 100644
--- a/sgl-kernel/python/sgl_kernel/__init__.py
+++ b/sgl-kernel/python/sgl_kernel/__init__.py
@@ -23,6 +23,7 @@
     silu_and_mul,
 )
 from sgl_kernel.gemm import (
+    awq_dequantize,
     bmm_fp8,
     cublas_grouped_gemm,
     fp8_blockwise_scaled_mm,
diff --git a/sgl-kernel/python/sgl_kernel/gemm.py b/sgl-kernel/python/sgl_kernel/gemm.py
index e5936da5677d..d68d0a437a7d 100644
--- a/sgl-kernel/python/sgl_kernel/gemm.py
+++ b/sgl-kernel/python/sgl_kernel/gemm.py
@@ -4,6 +4,12 @@
 from sgl_kernel.utils import _get_cache_buf, get_cuda_stream
 
 
+def awq_dequantize(
+    qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor
+) -> torch.ByteTensor:
+    return torch.ops.sgl_kernels.awq_dequantize(qweight, scales, qzeros)
+
+
 def int8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):
     return torch.ops.sgl_kernel.int8_scaled_mm(
         mat_a,
diff --git a/sgl-kernel/setup.py b/sgl-kernel/setup.py
index 0cf88ff06be3..885e654894a9 100644
--- a/sgl-kernel/setup.py
+++ b/sgl-kernel/setup.py
@@ -150,6 +150,7 @@ def copy_deepgemm_to_build_lib(self):
     "csrc/elementwise/rope.cu",
     "csrc/gemm/bmm_fp8.cu",
     "csrc/gemm/cublas_grouped_gemm.cu",
+    "csrc/gemm/awq_kernel.cu",
     "csrc/gemm/fp8_gemm_kernel.cu",
     "csrc/gemm/fp8_blockwise_gemm_kernel.cu",
     "csrc/gemm/int8_gemm_kernel.cu",
diff --git a/sgl-kernel/tests/test_awq_dequant.py b/sgl-kernel/tests/test_awq_dequant.py
new file mode 100644
index 000000000000..c2a2ee84d527
--- /dev/null
+++ b/sgl-kernel/tests/test_awq_dequant.py
@@ -0,0 +1,67 @@
+import itertools
+from typing import Optional, Tuple
+
+import pytest
+import torch
+from sgl_kernel import awq_dequantize
+from vllm import _custom_ops as ops
+
+
+def vllm_awq_dequantize(
+    qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor
+) -> torch.Tensor:
+    return ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
+
+
+def sglang_awq_dequantize(
+    qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor
+) -> torch.Tensor:
+    return awq_dequantize(qweight, scales, qzeros)
+
+
+@pytest.mark.parametrize(
+    "qweight_row,qweight_col",
+    list(
+        itertools.product(
+            [3584, 18944, 128, 256, 512, 1024], [448, 576, 4736, 16, 32, 64, 128]
+        )
+    ),
+)
+def test_awq_dequant_compare_implementations(
+    qweight_row: int,
+    qweight_col: int,
+):
+    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,
+    )
+
+    # Run both implementations
+    vllm_out = vllm_awq_dequantize(qweight, scales, qzeros)
+    sglang_out = sglang_awq_dequantize(qweight, scales, qzeros)
+
+    # Compare results
+    torch.testing.assert_close(
+        vllm_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5
+    )
+
+
+if __name__ == "__main__":
+    # Run the specific test function directly
+    pytest.main([__file__])


Reviews:
user: login: hebiao064
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-08T05:16:41Z
user: login: zcnrex
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-08T06:01:47Z
user: login: BBuf
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-09T11:11:38Z
user: login: zcnrex
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-09T20:01:45Z
user: login: BBuf
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-10T01:43:08Z

Review comments:
diff_hunk: @@ -0,0 +1,115 @@
+import itertools
+from typing import List, Tuple
+
+import torch
+import triton
+import triton.testing
+from sgl_kernel import awq_dequantize
+from vllm import _custom_ops as ops
+
+
+def vllm_awq_dequantize(
+    qweight: torch.Tensor,
+    scales: torch.Tensor,
+    qzeros: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    return ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
+
+
+def sglang_awq_dequantize(
+    qweight: torch.Tensor,
+    scales: torch.Tensor,
+    qzeros: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+    return awq_dequantize(qweight, scales, qzeros)
+
+
+def calculate_diff(qweight_row: int, qweight_col: int):
+    """Calculate difference between VLLM and SGLang implementations."""
+    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)
+
+    vllm_out = vllm_awq_dequantize(qweight, scales, qzeros)
+    sglang_out = sglang_awq_dequantize(qweight, scales, qzeros)
+
+    output_diff = torch.abs(vllm_out.float() - sglang_out.float()).mean().item()
+
+    if torch.allclose(
+        vllm_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5
+    ):
+        print("✅ All implementations match")
+    else:
+        print("❌ Implementations differ")
+
+
+qweight_row_range = [3584, 18944, 128, 256, 512, 1024]
+qweight_cols_range = [448, 576, 4736, 16, 32, 64, 128]
+
+configs = list(itertools.product(qweight_row_range, qweight_cols_range))
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["qweight_row", "qweight_col"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=["vllm", "sglang"],
+        line_names=["VLLM", "SGL Kernel"],
+        styles=[("blue", "-"), ("green", "-")],
+        ylabel="us",
+        plot_name="awq-dequantize-performance",
+        args={},
+    )
+)
+def benchmark(qweight_row, qweight_col, provider):
+    dtype = torch.float16
+    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 == "vllm":
+        fn = lambda: vllm_awq_dequantize(qweight.clone(), scales.clone(), qzeros.clone())
+    elif provider == "sglang":
+        fn = lambda: sglang_awq_dequantize(qweight.clone(), scales.clone(), qzeros.clone())
+
+    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    calculate_diff(qweight_row=3584, qweight_col=448)
path: sgl-kernel/benchmark/bench_awq_dequant.py
user: login: hebiao064
body: can you add a test for this kernel?
created_at: 2025-03-08T05:16:40Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,115 @@
+import itertools
+from typing import List, Tuple
+
+import torch
+import triton
+import triton.testing
+from sgl_kernel import awq_dequantize
+from vllm import _custom_ops as ops
+
+
+def vllm_awq_dequantize(
+    qweight: torch.Tensor,
+    scales: torch.Tensor,
+    qzeros: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    return ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
+
+
+def sglang_awq_dequantize(
+    qweight: torch.Tensor,
+    scales: torch.Tensor,
+    qzeros: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+
+    return awq_dequantize(qweight, scales, qzeros)
+
+
+def calculate_diff(qweight_row: int, qweight_col: int):
+    """Calculate difference between VLLM and SGLang implementations."""
+    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)
+
+    vllm_out = vllm_awq_dequantize(qweight, scales, qzeros)
+    sglang_out = sglang_awq_dequantize(qweight, scales, qzeros)
+
+    output_diff = torch.abs(vllm_out.float() - sglang_out.float()).mean().item()
+
+    if torch.allclose(
+        vllm_out.to(torch.float32), sglang_out.to(torch.float32), rtol=1e-3, atol=1e-5
+    ):
+        print("✅ All implementations match")
+    else:
+        print("❌ Implementations differ")
+
+
+qweight_row_range = [3584, 18944, 128, 256, 512, 1024]
+qweight_cols_range = [448, 576, 4736, 16, 32, 64, 128]
+
+configs = list(itertools.product(qweight_row_range, qweight_cols_range))
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["qweight_row", "qweight_col"],
+        x_vals=configs,
+        line_arg="provider",
+        line_vals=["vllm", "sglang"],
+        line_names=["VLLM", "SGL Kernel"],
+        styles=[("blue", "-"), ("green", "-")],
+        ylabel="us",
+        plot_name="awq-dequantize-performance",
+        args={},
+    )
+)
+def benchmark(qweight_row, qweight_col, provider):
+    dtype = torch.float16
+    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 == "vllm":
+        fn = lambda: vllm_awq_dequantize(qweight.clone(), scales.clone(), qzeros.clone())
+    elif provider == "sglang":
+        fn = lambda: sglang_awq_dequantize(qweight.clone(), scales.clone(), qzeros.clone())
+
+    ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles)
+
+    return 1000 * ms, 1000 * max_ms, 1000 * min_ms
+
+
+if __name__ == "__main__":
+    calculate_diff(qweight_row=3584, qweight_col=448)
path: sgl-kernel/benchmark/bench_awq_dequant.py
user: login: zcnrex
body: added
created_at: 2025-03-08T06:01:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,133 @@
+#include <c10/cuda/CUDAGuard.h>
path: sgl-kernel/src/sgl-kernel/csrc/gemm/awq_kernel.cu
user: login: BBuf
body: `# Adapted from https://github.com/vllm-project/vllm/blob/eb59b5a6cba6727d3727c0372258db9002f687c1/csrc/quantization/awq/gemm_kernels.cu#L350`
created_at: 2025-03-09T11:11:38Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,133 @@
+#include <c10/cuda/CUDAGuard.h>
path: sgl-kernel/src/sgl-kernel/csrc/gemm/awq_kernel.cu
user: login: zcnrex
body: Added thanks
created_at: 2025-03-09T20:01:45Z
author_association: CONTRIBUTOR

PR comments:
user: login: BBuf
created_at: 2025-03-07T01:46:16Z
body: Rebase main please. @zcnrex 
author_association: COLLABORATOR
user: login: zcnrex
created_at: 2025-03-07T06:07:00Z
body: @BBuf Rebased main, ready for review. Thanks.
author_association: CONTRIBUTOR
user: login: BBuf
created_at: 2025-03-07T09:14:28Z
body: > @BBuf Rebased main, ready for review. Thanks.

Thanks, I'll review it when I am free.
author_association: COLLABORATOR