A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 11, 12:30 UTC
[Bugfix] awq_gemm: fix argument order swap (#30364)
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Matthias Gehre committed
Dec 14, 2025 at 11:15 UTC
e9add129ad9daf7a9e00381da318db271646813a
2 files changed
+7
-7
tests/kernels/quantization/test_awq.py
+3
-3
@@ -41,9 +41,9 @@ def test_awq_gemm_opcheck(monkeypatch: pytest.MonkeyPatch):
41
qweight = torch.randint(
42
-2000000000, 2000000000, (8192, 256), device="cuda", dtype=torch.int32
43
)
44
- scales = torch.randint(
44
+ scales = torch.empty((64, 2048), device="cuda", dtype=torch.float16)
45
+ qzeros = torch.randint(
46
-2000000000, 2000000000, (64, 256), device="cuda", dtype=torch.int32
47
)
47
- qzeros = torch.empty((64, 2048), device="cuda", dtype=torch.float16)
48
split_k_iters = 8
49
- opcheck(torch.ops._C.awq_gemm, (input, qweight, qzeros, scales, split_k_iters))
49
+ opcheck(torch.ops._C.awq_gemm, (input, qweight, scales, qzeros, split_k_iters))
vllm/_custom_ops.py
+4
-4
@@ -498,15 +498,15 @@ def awq_dequantize(
498
def awq_gemm(
499
input: torch.Tensor,
500
qweight: torch.Tensor,
501
- qzeros: torch.Tensor,
501
scales: torch.Tensor,
502
+ qzeros: torch.Tensor,
503
split_k_iters: int,
504
) -> torch.Tensor:
505
if envs.VLLM_USE_TRITON_AWQ:
506
from vllm.model_executor.layers.quantization.awq_triton import awq_gemm_triton
507
508
- return awq_gemm_triton(input, qweight, qzeros, scales, split_k_iters)
509
- return torch.ops._C.awq_gemm(input, qweight, qzeros, scales, split_k_iters)
508
+ return awq_gemm_triton(input, qweight, scales, qzeros, split_k_iters)
509
+ return torch.ops._C.awq_gemm(input, qweight, scales, qzeros, split_k_iters)
510
511
512
# gptq
@@ -632,8 +632,8 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):
632
def _awq_gemm_fake(
633
input: torch.Tensor,
634
qweight: torch.Tensor,
635
- qzeros: torch.Tensor,
635
scales: torch.Tensor,
636
+ qzeros: torch.Tensor,
637
split_k_iters: torch.SymInt,
638
) -> torch.Tensor:
639
num_in_feats = input.size(0)