Title: CodeGEMM: A Codebook-Centric Approach to Efficient GEMM in Quantized LLMs
Venue: NeurIPS 2025
Accepted: yes
Keywords: Quantization

Matrix Multiplication

Large Language Models
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces CodeGEMM, an optimized matrix multiplication kernel for codebook-based quantized large language models. Traditional codebook-based quantization methods suffer from significant latency due to the need for repeated dequantization and frequent codebook access. \n\nCodeGEMM precompute and storing inner product results between codebook centroids and input data in the programmable cache, eliminating the need for on-the-fly dequantization. The kernel is designed to be flexible across various quantization hyperparameters and model configurations.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n1. The Psumbook-based approach is a practical and innovative solution to a key bottleneck in codebook-based quantization, significantly reducing both space and compute complexity.\n\n2. The kernel is adaptable to a wide range of quantization settings, supporting diverse codebook sizes, vector lengths, and group normalization strategies, making it practical for real deployment.\n\nWeaknesses\n1. CodeGEMM is constrained by on-chip cache size, making very large codebooks (e.g., b=16) infeasible. The workaround (smaller codebooks, fine-grained group normalization) is reasonable, but may limit certain applications.\n\n2. The performance improvements are well justified empirically, but the paper could benefit from a more formal treatment of complexity, cache usage, and the limits of the Psumbook design.\n\n3. Figure 1 - Figure 3 appears to be a non-vector image. It would be better to replace it with a vector graphic for improved clarity and scalability.\n\nQUESTIONS:\nsee weaknesses\n\nLIMITATIONS:\nProvide a small-scale analysis or discussion of possible failure cases where Psumbook overhead could become a new bottleneck (e.g., if codebook structure is highly non-uniform or input batch sizes are extreme).\n\nIf possible, add error bars or a variance table to the main throughput/accuracy results.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nCodeGEMM is a GPU GEMM kernel designed for large language models that eliminates the costly dequantization step by pre-computing all inner products between input activations and centroid vectors, storing the results in a small shared-memory “Psumbook.” At inference time, the kernel merely indexes this table with the quantized codes, reducing cache traffic and computational complexity while accommodating a broad range of quantization hyperparameters (bits per code, number of codebooks, vector length, group size) without kernel rewrites. Evaluations on 2-bit Llama-3 (8 B and 70 B) show up to 2.27× kernel-level speedups and 1.8× higher end-to-end throughput over state-of-the-art codebook methods, illustrating that a codebook-centric GEMM can translate the theoretical storage gains of extreme low-bit quantization into real-world LLM inference performance.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n- This paper is overall well written and easy to follow. The method and results are clearly presented.\n- Motivation and the design of the kernel are visualised with side-by-side diagrams (Figure 1) and a step-wise schematic (Figure 3). \n- The authors derive both computational and space complexity, showing CodeGEMM reduces the multiply-accumulate count by a factor ≈ m/v. \n- It provides up to 2.27× speed-up compared to other codebook-based quantization methods with comparable accuracy in the 2-bit configuration.\n\n\nWeaknesses: \n- It only provides benchmark results of the kernel under two different model sizes. Since this work is all about kernel optimization, more benchmark results under different problem sizes are supposed to be given. \n- Implementation details such as choosing t_w = 32, t_h = 2048 are mentioned but not justified. \n- Because the Psumbook must reside in shared memory, very large codebooks (b = 16) are out of scope; the authors themselves list this as a limitation.\n\nQUESTIONS:\n1. Can you provide more benchmark results for the kernel for different problem sizes (M, N, K)? \n2. Line 175 says \"each input vector interacts with a limited set of centroids\". Can you provide any evidence for this assumption? \n3. If higher bits per weight were used, what would the speedup be like?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper presents CodeGEMM, a new GPU kernel for weight-only, codebook-based quantization in large language models (LLMs). Instead of performing on-the-fly dequantization by loading full centroid vectors into cache, CodeGEMM precomputes all possible inner products between centroids and input subvectors—storing them in a compact Psumbook. During GEMM, it uses each code to index directly into this Psumbook, eliminating repeated dequantization and reducing both computational and space complexity. \nThe kernel is parameterized by vector length $v$, number of codebooks $m$, bits per code $b$, and tile width $t_w$, enabling exploration of the latency memory accuracy trade-off.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths:**\n- Precomputing and caching inner products removes the need to fetch full centroids for each code, directly addressing the shared-memory capacity bottleneck in existing kernels \n- Kernel-level benchmarks demonstrate up to 2.18× improvement over FP16 cuBLAS and 1.64× over AQLM on A100 GPUs, and throughput evaluations confirm real-world speedups in the HuggingFace Llama pipeline \n\n**Weaknesses:**\n- The reliance on fitting the Psumbook in shared memory forces $b\\le8$ and precludes exploration of larger codebook widths (e.g., $b=16$), limiting accuracy potential in extreme low-bit regimes .\n- The focus is on latency and perplexity; energy consumption and DRAM bandwidth savings which are critical in production are unreported.\n- It remains unclear how much time is spent constructing the Psumbook versus retrieving entries during GEMM. A breakdown would clarify whether build overheads offset retrieval gains at different batch sizes or tile shapes.\n- While NeurIPS checklist items are cited, the paper omits error bars (claimed negligible variance) and details on code release versioning, which may hinder exact reproduction .\n\nQUESTIONS:\n- Can you provide detailed microbenchmarks separating the time spent building versus reading the Psumbook, especially for varying $t_w$ and batch sizes?\n- How does CodeGEMM fare when coupled with codebooks generated by TurboQuant (https://arxiv.org/abs/2504.19874)?\n- Have you measured DRAM traffic reductions, and can you provide quantitative comparisons to dequantization-based kernels?\n- Do you envision streamed or hierarchical Psumbook designs that allow $b>8$ without cache overflow, and what challenges do these impose?\n\nLIMITATIONS:\nSee Weaknesses", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 1, "ORIGINALITY": 2, "DATE": "2025-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a codebook-based quantization method to represent weights using 2-bit configurations. Unlike previous approaches that stored the entire codebook in memory for dequantization, the proposed method stores the product-sum results, eliminating the need for dequantization. Hyperparameters for codebook selection such as group sizes, vector length, and number of codebooks are explored to find the best trade-off between performance accuracy and hardware complexity. The new approach achieves a 2.29× speed improvement compared to previous codebook-based quantization methods such as AQLM (2×8).\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\n1- The new approach is explored across various hardware constraints such as memory footprint, latency, and throughput, and it demonstrates the benefits of product-sum codebook-based quantization compared to previous approaches.\n\n2- The paper is well-written, and the experiments are elaborated in detail.\n\nWeaknesses:\n\n1- The accuracy degradation of the codebook-based approach is high compared to the FP16 approach in Table 3 (approximately 7%). The author is advised to elaborate on this accuracy degradation.\n\n2- It is suggested to include accuracy vs. latency and accuracy vs. throughput plots corresponding to Table 3 and Table 4.\n\n3- It is recommended that the comparison with these previous approaches [1, 2] be demonstrated in the paper.\n\n[1] Liu, Yifei, et al. \"Vptq: Extreme low-bit vector post-training quantization for large language models.\" arXiv preprint arXiv:2409.17066 (2024).\n \n[2] Liu, Zechun, et al. \"ParetoQ: Scaling Laws in Extremely Low-bit LLM Quantization.\" arXiv preprint arXiv:2502.02631 (2025).\n\nQUESTIONS:\nIs there any limitation on the size of the product-sum values stored in memory for large language models such as LLaMA 3.1 with 405B parameters?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-07"}