CodeGEMM: A Codebook-Centric Approach to Efficient GEMM in Quantized LLMs

AI summary

Overview

Research area: Efficient inference systems for large language models — specifically GPU kernel design for codebook-based (vector-quantized) weight-only quantization.

Technical level: Advanced. The paper assumes familiarity with GEMM tiling, shared/programmable cache hierarchies, quantization hyperparameters, and GPU profiling telemetry.

Scope: CodeGEMM is a single GEMM kernel that replaces on-the-fly dequantization in codebook-quantized LLMs with precomputed centroid–activation inner products stored in an on-chip "Psumbook," and it is evaluated on Llama-3.1 8B and 70B models at roughly 2 bits per weight.

What This Paper Is About

Extremely low-bit (e.g., 2-bit) weight-only quantization is attractive because LLM inference is memory-bound, and codebook-based quantization preserves accuracy better than uniform quantization at those bit widths. The problem is that existing codebook kernels must load an entire codebook into programmable cache and dequantize weights on the fly — an operation that repeatedly fetches centroids, wastes on-chip space, and does not reduce arithmetic work at all. This paper asks whether the codebook can be handled differently so that matrix multiplication becomes cheaper in both space and computation, not just in data movement.

Key Contributions

1. CodeGEMM, a codebook-centric GEMM kernel. Instead of loading the full codebook into programmable cache and dequantizing, CodeGEMM precomputes inner products between centroid vectors and input activations and stores those scalars in a Psumbook. At inference, code indices directly gather precomputed partial sums, eliminating per-element centroid lookups.

2. A unified kernel spanning many codebook hyperparameters. One implementation supports variations in the number of codebooks (m), vector length (v), group size (g), and bits per code (b), enabling systematic study of latency–memory–accuracy trade-offs rather than requiring a separate kernel per configuration.

3. Complexity analysis showing a genuine reduction in computation, not only data movement. Standard GEMM and dequantization-based kernels are O(MNK). CodeGEMM is approximately O(MNK · m/v), a reduction factor of m/v, with space complexity O(m · 2^b · t_w/v) versus O(m · 2^b · v) for dequantization-based kernels.

4. Empirical speedups at comparable accuracy. On Llama-3 models, CodeGEMM reports 1.83× (8B) and 8.93× (70B) speedups in the 2-bit configuration compared to state-of-the-art codebook-based quantization, plus improved compute efficiency and memory-subsystem utilization.

Main Findings

• Kernel latency advantage. On 2-bit quantized Llama3 models measured on an NVIDIA A100 80GB, CodeGEMM (m1v4g128) reaches 152.69 µs on 8B and 293.82 µs on 70B, versus 332.45 µs / 1111.36 µs for FP16 cuBLAS, 250.12 µs / 674.67 µs for AQLM (2x8), and 645.51 µs / 2285.5 µs for AQLM (1x16). The paper states CodeGEMM achieves up to 2.18× and 1.64× speed-ups over the FP16 baseline and AQLM, respectively, at the same average bit precision.

• The large codebook is the bottleneck for AQLM. The AQLM (1×16) kernel needs 2^16 centroids of length v=8 in FP16, i.e. 1MB of shared memory, which exceeds both A100 (164KB) and H100 (224KB) shared memory capacity. This forces repeated DRAM fetches and explains its high latency (2285.5 µs on 70B).

• Higher energy efficiency and better memory behavior. On a GEMV with (M, N, K) = (1, 28672, 8192), CodeGEMM-m1v4g128 reaches 6.12 TFLOPS and 19.36 GFLOPS/W at 316.38 W, versus AQLM-2x8 at 2.59 TFLOPS and 10.18 GFLOPS/W, and cuBLAS at 1.58 TFLOPS and 4.95 GFLOPS/W. Memory utilization is 49.80% for CodeGEMM-m1v4g128 versus 19.96% for AQLM-2x8 and 6.00% for AQLM-1x16. Metrics were sampled at a 100 ms cadence over a 10 s window, with two-sigma error margins over 128 samples.

• Throughput–accuracy trade-off, not just memory–accuracy. On Llama-3.1-8B-Instruct, CodeGEMM-m1v4g128 with PV-Tuning reaches 228.3 tok/s at 63.96 average accuracy, versus AQLM-2x8 with PV-Tuning at 124.5 tok/s and 62.74, and AQLM-1x16 with PV-Tuning at 49.0 tok/s and 65.82. FP16 achieves 103.8 tok/s at 71.26.

• Uniform quantization collapses at 2 bits. FlexRound-q2g128 reaches 205.3 tok/s but only 41.65 average accuracy on the 8B model, and 41.7 tok/s / 36.58 on 70B — the worst accuracy among compared methods.

• VPTQ comparison. The vector-quantization approach VPTQ achieves an average accuracy of 57.98, slightly higher than CodeGEMM without PV-Tuning, but its kernel is a straightforward dequantize-then-multiply pipeline without operator fusion, yielding lower throughput than FP16.

• Scaling to 70B widens the gap. At 70B, CodeGEMM matches AQLM (1×16) accuracy while delivering an 8.93× throughput advantage; CodeGEMM-m1v4g32 reaches 73.15 average accuracy at 49.1 tok/s versus AQLM-1x16 at 73.89 average accuracy and 5.5 tok/s. The FP16 70B baseline is marked OOM.

• Hyperparameter interactions matter. Under row-wise normalization (g = -1), increasing the number of codebooks m at fixed average bit precision improves accuracy. As g becomes smaller (more fine-grained), that gain from increasing m diminishes and models perform similarly for a given average bits per weight. Group normalization overhead is minimal for g ≥ 32 but rises sharply at per-vector normalization (g = v).

• Perplexity and throughput rankings can diverge. The m1v4 variant consistently beats m2v8 in both throughput and accuracy, consistent with kernel latency trends but only partially consistent with WikiText-2 perplexity trends.

Methodology in Plain English

Codebook quantization represents a weight vector not by individually quantized numbers but by pointing to a small set of prototype vectors (centroids). A weight matrix is split into vectors of length v, normalized over groups of size g, clustered into 2^b centroids per codebook, and stored as m sets of b-bit codes. Reconstruction requires reading the code and fetching the corresponding centroid — the dequantization step.

CodeGEMM flips this. Rather than fetching centroids during the matrix multiply, it first computes the inner product between every possible centroid and each segmented piece of the input activation, and stores those scalar results in a Psumbook that lives in programmable cache. During the actual multiply, the code is used only as an index into this small scalar table; the gathered values are accumulated to produce output activations. No centroids are touched at compute time, and each input–centroid product is computed once instead of once per output row.

The implementation uses a weight tile of t_h = 2048 and t_w = 32 to maximize Psumbook reuse within a thread block. Evaluation covers three axes: memory footprint (average bits per weight q̄, computed from the paper's Equation 1), latency (summed kernel execution times over all linear layers in one Transformer decoder block, without layer fusion, on an A100 80GB), and accuracy (lm-eval-harness across zero-shot and 5-shot settings on MMLU, WinoGrande, HellaSwag, ARC-Easy, and ARC-Challenge, plus WikiText-2 perplexity). End-to-end throughput was measured with the HuggingFace Llama implementation using layer fusion, and power/memory telemetry via nvidia-smi.

Why This Matters

Impact on research. The paper reframes codebook quantization as a problem where computation itself can be reduced, not merely data movement. Most quantized GEMM kernels inherit the full O(MNK) arithmetic cost of FP16; CodeGEMM argues that putting the table lookup at the level of precomputed partial sums (rather than dequantized weights, as in LUT-GEMM) changes the asymptotics. It also provides one configurable kernel for systematically mapping latency–memory–accuracy trade-offs that were previously only partly explored.

Real-world applications:

• Serving large quantized models on a single memory-constrained GPU or node, where shrinking weights to ~2 bits makes previously unloadable models feasible.
• Cost- and power-sensitive cloud inference, where the reported GFLOPS/W improvements (19.36 vs 10.18 and 4.95) directly affect serving economics.
• Edge and on-device deployment of LLMs, where on-chip memory is the binding constraint and a smaller cache footprint is decisive.
• Fixed-function ASIC or accelerator design: the authors argue CodeGEMM's generality — uniform and binary quantization can be expressed by defining centroids as c ∈ {0,1}^v and c ∈ {-1,1}^v respectively — makes it suitable as a unified hardware kernel.

Industry relevance. The work comes from NAVER Cloud, with an open implementation at github.com/naver-aics/codegemm, and targets deployment rather than a purely algorithmic advance. It addresses a practical pain point: state-of-the-art codebook methods like AQLM achieve the best 2-bit accuracy but their kernels are slow precisely because of codebook size, making them hard to justify in production.

Future Directions

1. Breaking the Psumbook size ceiling. The paper acknowledges that CodeGEMM requires the Psumbook to fit in on-chip shared memory, which rules out very large codebooks (b = 16, 2^16 entries). The authors fix b = 8 and recover accuracy through fine-grained group normalization. Whether hierarchical, multi-level, or partly off-chip Psumbook designs could restore access to large codebooks is left open.

2. Large-batch and prefill performance. CodeGEMM underperforms Tensor Core–based cuBLAS when batch size is large (e.g., M > 32). The paper attributes this to CUDA Core–based quantized GEMM generally and the current commercial GPU architecture, but a path toward Tensor Core integration or a hybrid kernel is not demonstrated.

3. Automated hyperparameter selection. The work shows that many combinations of (v, m, b, g) yield nearly the same average bits per weight but differ in accuracy and latency — for example (4, 1, 8, -1) and (16, 3, 8, 32) both land near 2 average bits. Turning this into an automatic search procedure per model and deployment target is a natural next step.

4. Portability of the Psumbook approach. The results are measured on an A100 80GB, and the design is motivated by shared-memory limits that differ across accelerators (A100 164KB vs H100 224KB). How the trade-offs shift on other hardware, and whether the same kernel structure transfers, is not reported.

Target Audience

Systems and ML researchers working on LLM inference efficiency, GPU kernel engineers, and quantization specialists who need sub-4-bit weight compression to be practically fast rather than merely accurate. It is also relevant to practitioners deploying Llama-scale models under hard memory budgets, and to hardware architects designing accelerators for non-uniform quantized inference. Readers without background in GEMM tiling, shared-memory hierarchies, or vector quantization will find the kernel-level sections demanding.

Authors’ abstract

Weight-only quantization is widely used to mitigate the memory-bound nature of LLM inference. Codebook-based methods extend this trend by achieving strong accuracy in the extremely low-bit regime (e.g., 2-bit). However, current kernels rely on dequantization, which repeatedly fetches centroids and reconstructs weights, incurring substantial latency and cache pressure. We present CodeGEMM, a codebook-centric GEMM kernel that replaces dequantization with precomputed inner products between centroids and activations stored in a lightweight Psumbook. At inference, code indices directly gather these partial sums, eliminating per-element lookups and reducing the on-chip footprint. The kernel supports the systematic exploration of latency-memory-accuracy trade-offs under a unified implementation. On Llama-3 models, CodeGEMM delivers 1.83x (8B) and 8.93x (70B) speedups in the 2-bit configuration compared to state-of-the-art codebook-based quantization at comparable accuracy and further improves computing efficiency and memory subsystem utilization.