Research map / Architectures and efficiency / Quantization
CodeGEMM: A Codebook-Centric Approach to Efficient GEMM in Quantized LLMs
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes CodeGEMM, a codebook-centric GEMM kernel that precomputes centroid-activation partial sums in a Psumbook so code indices gather results instead of dequantizing weights, speeding up 2-bit codebook-quantized Llama-3 inference.
Problem
Codebook-based weight-only quantization achieves strong 2-bit accuracy, but existing kernels rely on dequantization that repeatedly fetches centroids from a codebook, causing latency and cache pressure.
Approach
Replace dequantize-then-multiply with a Psumbook of precomputed inner products between centroids and input activations; code indices directly gather and accumulate these partial sums, reducing computation by roughly m/v and cache footprint by v.
Key results
- CodeGEMM achieves up to 2.18x and 1.64x kernel speedups over FP16 cuBLAS and AQLM respectively at the same bit width on Llama-3 8B.
- On Llama-3.1-8B-Instruct at about 2.13 bits, CodeGEMM-m1v4g128 with PV-Tuning reaches 63.96 average accuracy vs 65.82 for AQLM-1x16+PV-Tuning, at 228.3 tok/s vs 49.0.
- Abstract reports 1.83x (8B) and 8.93x (70B) end-to-end speedups in the 2-bit setting over codebook-based quantization at comparable accuracy.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| use Psumbook precomputed inner products instead of dequantizing codebook weights | CodeGEMM versus dequantization-based codebook kernels · Llama-3 8B and 70B, 2-bit linear layers | latency: 332.45 -> 152.69 (8B); 1111.36 -> 293.82 (70B) | Helped |
| use CodeGEMM-m1v4g128 instead of AQLM-1x16 | Llama-3.1-70B · Llama-3.1-70B | throughput: 5.5 -> 51.2 tok/s | Helped |
| add PV-Tuning to CodeGEMM-m1v4g128 (ablation) | CodeGEMM-m1v4g128 · Llama-3.1-8B-Instruct | average accuracy: 53.93 -> 63.96 | Helped |
| add PV-Tuning to CodeGEMM-m2v8g128 (ablation) | CodeGEMM-m2v8g128 · Llama-3.1-8B-Instruct | average accuracy: 52.67 -> 63.76 | Helped |
| use m1v4 instead of m2v8 (ablation) | CodeGEMM · Llama-3.1-8B-Instruct | average accuracy: 52.67 -> 53.93 | Helped |
| use m1v4 instead of m2v8 (ablation) | CodeGEMM · Llama-3.1-8B-Instruct | throughput: 214.4 -> 228.3 tok/s | Helped |
| use fine-grained group normalization (g=32) instead of g=128 (ablation) | CodeGEMM-m1v4 · Llama-3.1-70B | average accuracy: 70.11 -> 73.15 | Helped |
| increase tile height from 2048 to 4096 (ablation) | CodeGEMM-m1v4, tw32 · Llama linear layers, M=1, N=4096, K=4096 | latency: 25.07 -> 27.60 μs | Hurt |
| increase tile width from 32 to 128 (ablation) | CodeGEMM-m1v4, th2048 · Llama linear layers, M=1, N=4096, K=4096 | latency: 25.07 -> 26.81 μs | Hurt |
| increase tile width from 32 to 64 (ablation) | CodeGEMM-m1v4, th2048 · Llama linear layers, M=1, N=8192, K=8192 | latency: 36.02 -> 35.33 μs | Helped |
| increase the number of codebooks while holding average bit precision fixed (ablation) | CodeGEMM with row-wise normalization (g=-1) · WikiText-2 | perplexity | Helped |
| use smaller group size for finer-grained normalization (ablation) | CodeGEMM · WikiText-2 | perplexity | Helped |
Surprising findings
- Psumbook construction and reading are both substantial: build takes 20-46% of cycles while read takes 54-80%, so the kernel is not dominated by one phase.
- CodeGEMM underperforms Tensor Core cuBLAS at large batch sizes (M>32) and its advantage shrinks with batch size, a limitation shared by CUDA-core quantized kernels.
Limitations
- Psumbook must reside in on-chip shared memory, so codebook width is limited (b=8 in practice); b=16 exceeds cache.
- Underperforms Tensor Core cuBLAS at large batch sizes (M>32) since it runs on CUDA cores.
Experimental setup
Datasets: WikiText-2
Benchmarks: WikiText-2 perplexity MMLU WinoGrande HellaSwag ARC-Easy ARC-Challenge
Metrics: kernel latency (us) tokens/s perplexity accuracy GFLOPS/W memory utilization
Baselines: AQLM QuIP# QTIP LUT-GEMM FlexRound VPTQ GPTQ cuBLAS FP16
Builds on: AQLM GPTQ LUT-GEMM PV-Tuning
Closest papers by meaning
- Q-Palette: Fractional-Bit Quantizers Toward Optimal Bit Allocation for Efficient LLM Deployment (NeurIPS 2025 · similarity 0.89)
- BlockDialect: Block-wise Fine-grained Mixed Format Quantization for Energy-Efficient LLM Inference (ICML 2025 · similarity 0.89)
- SpQR: A Sparse-Quantized Representation for Near-Lossless LLM Weight Compression (ICLR 2024 · similarity 0.88)
- QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs (NeurIPS 2024 · similarity 0.88)
- ScaleBITS: Scalable Bitwidth Search for Hardware-Aligned Mixed-Precision LLMs (NeurIPS 2026 · similarity 0.88)
- LUT-GEMM: Quantized Matrix Multiplication based on LUTs for Efficient Inference in Large-Scale Generative Language Models (ICLR 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.