QUICK: Quantization-aware Interleaving and Conflict-free Kernel for efficient LLM inference

TL;DR

The paper addresses inefficiencies in weight-only quantized LLM inference caused by dequantization overhead and shared-memory bank conflicts in mixed-precision GEMMs. It proposes QUICK, a set of optimized CUDA kernels that offline interleave quantized weights to align with Tensor Core data-loading patterns (ldmatrix/mma) and skip shared-memory write-backs, thereby removing bank conflicts. Empirical results show substantial gains in matrix-multiply throughput and end-to-end token generation across several models and GPUs, including notable improvements over AutoAWQ and positive benefits in vLLM integration. This approach enables more effective deployment of weight-only quantized LLMs, especially for larger batch sizes, by improving throughput without increasing memory bandwidth requirements.

Abstract

We introduce QUICK, a group of novel optimized CUDA kernels for the efficient inference of quantized Large Language Models (LLMs). QUICK addresses the shared memory bank-conflict problem of state-of-the-art mixed precision matrix multiplication kernels. Our method interleaves the quantized weight matrices of LLMs offline to skip the shared memory write-back after the dequantization. We demonstrate up to 1.91x speedup over existing kernels of AutoAWQ on larger batches and up to 1.94x throughput gain on representative LLM models on various NVIDIA GPU devices.

Figures (8)

• Figure 1: Data loading pattern of ldmatrix instruction for a single $8\times8$ matrix. Two half-precision elements are loaded to the destination register d0 per each thread lane in a warp.
• Figure 2: Computation overview of original kernel and QUICK. Compared to original kernel, QUICK bypasses 3) shared memory write-back and 4) ldmatrix operation of dequantized weights by using interleaving data pattern.
• Figure 3: Number of bank conflicts from benchmark result using NVIDIA Nsight Compute nsight. A matrix multiplication of shape $64 \times 8192 \times 8192 (M \times N \times K)$ was used as the workload.
• Figure 4: ldmatrix instruction-aware weight interleaving to avoid shared memory conflicts. With interleaved weight matrix, direct load from DRAM to registers for each thread without ldmatrix is possible. Note that the figure is illustrating a case of computing $8 \times 4 \times 8 (M \times N \times K)$ GEMM using 8 threads for simplicity. QUICK utilizes $16 \times 8 \times 16 (M \times N \times K)$ GEMM with 32 threads and its corresponding interleaving pattern.
• Figure 5: Parallel i4-f16 dequantization kernel-aware weight reordering.