[AWQ] Evaluate fused vs unfused GEMM on actual shape

vllm/_custom_ops.py

@@ -13,6 +13,7 @@
     flashinfer_quant_nvfp4_8x4_sf_layout,
 )
 from vllm.utils.math_utils import cdiv
+from vllm.utils.torch_utils import direct_register_custom_op
 
 logger = init_logger(__name__)
 
@@ -499,20 +500,46 @@ def awq_dequantize(
     return torch.ops._C.awq_dequantize(qweight, scales, zeros, split_k_iters, thx, thy)
 
 
-def awq_gemm(
+def _awq_gemm(
     input: torch.Tensor,
     qweight: torch.Tensor,
     scales: torch.Tensor,
     qzeros: torch.Tensor,
     split_k_iters: int,
 ) -> torch.Tensor:
+    # num_tokens >= threshold
+    FP16_MATMUL_HEURISTIC_CONDITION = input.shape[0] >= 256
+
+    if FP16_MATMUL_HEURISTIC_CONDITION:
+        out = awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
+        return torch.matmul(input, out)
+
     if envs.VLLM_USE_TRITON_AWQ:
         from vllm.model_executor.layers.quantization.awq_triton import awq_gemm_triton
 
         return awq_gemm_triton(input, qweight, scales, qzeros, split_k_iters)
     return torch.ops._C.awq_gemm(input, qweight, scales, qzeros, split_k_iters)
 
 
+def _awq_gemm_fake_impl(
+    input: torch.Tensor,
+    qweight: torch.Tensor,
+    scales: torch.Tensor,
+    qzeros: torch.Tensor,
+    split_k_iters: int,
+) -> torch.Tensor:
+    M, N = input.shape[0], qweight.shape[1] * 8
+    return torch.empty((M, N), dtype=input.dtype, device=input.device)
+
+
+direct_register_custom_op(
+    op_name="awq_gemm",
+    op_func=_awq_gemm,
+    fake_impl=_awq_gemm_fake_impl,
+)
+awq_gemm = torch.ops.vllm.awq_gemm
+
+
 # gptq
 def gptq_gemm(
     a: torch.Tensor,

vllm/model_executor/layers/quantization/awq.py

@@ -264,14 +264,7 @@ def apply(
         out_shape = x.shape[:-1] + (qweight.shape[-1] * pack_factor,)
         reshaped_x = x.reshape(-1, x.shape[-1])
 
-        # num_tokens >= threshold
-        FP16_MATMUL_HEURISTIC_CONDITION = x.shape[:-1].numel() >= 256
-
-        if FP16_MATMUL_HEURISTIC_CONDITION:
-            out = ops.awq_dequantize(qweight, scales, qzeros, 0, 0, 0)
-            out = torch.matmul(reshaped_x, out)
-        else:
-            out = ops.awq_gemm(reshaped_x, qweight, scales, qzeros, pack_factor)
+        out = ops.awq_gemm(reshaped_x, qweight, scales, qzeros, pack_factor)
         if bias is not None:
             out.add_(bias)
         return out.reshape(out_shape)