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

Weaknesses

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.

• We agree, and this is indeed a known limitation of our current implementation. As noted, very large codebooks such as $b=16$ can offer better accuracy but exceed the capacity of on-chip memory (e.g., GPU shared memory), making them impractical in our current design. However, as shown in Table 3, recent calibration techniques such as PV-Tuning have proven highly effective at mitigating accuracy degradation even under smaller codebooks like $b=8$, which previously showed substantial loss. We are optimistic that as more advanced tuning and quantization algorithms are developed, the remaining accuracy gap will narrow further, without requiring excessively large codebooks.

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.

• Thank you for pointing this out. While Section 3 provides a high-level discussion of the complexity, cache behavior, and architectural constraints of the Psumbook design, we agree that a more rigorous and detailed analysis could enhance the clarity and completeness of the paper. Our current version emphasizes conceptual clarity and accessibility to help readers understand the motivation behind CodeGEMM. That said, your suggestion is well taken, and we plan to incorporate a more formal treatment in the final version.

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.

• Thank you for pointing this out. We agree that vector graphics would improve readability and scalability. While we are unable to modify the manuscript during the rebuttal phase, we will update these figures with vectorized versions in the camera-ready version if accepted.

If possible, add error bars or a variance table to the main throughput/accuracy results.

• We appreciate this suggestion. Although we cannot revise the main figures at this stage, we have annotated all numerical results used during rebuttal with ±2σ error margins to ensure transparency and reliability. We will include proper error bars and/or variance tables in the camera-ready version.

We sincerely appreciate the insightful suggestions provided. Throughout the rebuttal period, we have made our best effort to address the reviewers' feedback and strengthen the paper accordingly. That said, we recognize that there may still be areas for improvement. If there are additional suggestions or concerns, we would be more than happy to incorporate them. Thank you again for your thoughtful and constructive comments.

Weaknesses

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.

Can you provide more benchmark results for the kernel for different problem sizes (M, N, K)?

• Thank you for the suggestion, and we apologize for not including a more comprehensive set of benchmarks earlier. We have now conducted additional experiments across a wide range of (M, N, K) configurations. (Values in parentheses denote ±2σ error margins of 128 samples.)
• Unlike cuBLAS, which utilizes Tensor Cores and maintains relatively stable latency across batch sizes, recently proposed quantized kernels, including ours, run on CUDA cores and tend to exhibit increased latency as the batch size (M) grows. Among these state-of-the-art kernels, CodeGEMM consistently demonstrates strong performance on large-scale matrix multiplications, where its compute and memory efficiency are best utilized.

Implementation details such as choosing $t_w = 32, t_h = 2048$ are mentioned but not justified.

• Thank you for pointing this out. Initially, we chose the tile dimensions $t_w = 32$ and $t_h = 2048$ based on heuristic tuning. However, prompted by your feedback, we conducted a more systematic analysis of these implementation parameters.
• We found that $t_h = 2048$ consistently provided the best performance across a wide range of workloads, supporting our original choice. Interestingly, for $t_w$, smaller values such as $t_w = 32$ worked well for relatively small matrix sizes, whereas larger values like $t_w = 64$ yielded better performance on large-scale matrix multiplications. We believe this is because larger matrices benefit from coarser tiling, which reduces launch overhead and improves partial sum reduction at scale.
• We will include a brief discussion of these observations in the implementation section of the final version.

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.

• We agree, and this is indeed a known limitation of our current implementation. As noted, very large codebooks such as b=16 can offer better accuracy but exceed the capacity of on-chip memory (e.g., GPU shared memory), making them impractical in our current design. However, as shown in Table 3, recent calibration techniques such as PV-Tuning have proven highly effective at mitigating accuracy degradation even under smaller codebooks like b=8, which previously showed substantial loss. We are optimistic that as more advanced tuning and quantization algorithms are developed, the remaining accuracy gap will narrow further, without requiring excessively large codebooks.

Questions

Line 175 says "each input vector interacts with a limited set of centroids". Can you provide any evidence for this assumption?

• Thank you for pointing this out. The phrasing in Line 175 may have caused confusion. As illustrated in Figure 1 of the paper, each output activation is computed by combining the input activation with all $2^b$ centroids per codebook, meaning there is no restriction on which centroids are involved in the computation.
• However, when the number of output channels $N$ exceeds the number of centroid combinations $2^b$, which is common in practice, the pigeonhole principle implies that multiple output channels will share the same centroid indices. This leads to reuse opportunities in the partial sum computation. CodeGEMM exploits this by precomputing the Psumbook, which enables efficient reuse of these dot products across multiple output channels.

If higher bits per weight were used, what would the speedup be like?

• Thank you for the suggestion. We also conducted latency measurements for higher-bit quantization settings using the kernel configuration ($g = 128, b = 8, t_w = 32, t_h = 2048$). For reference, FP16 cuBLAS latency is also included.
• As shown in the results, higher average bit precision generally leads to increased latency, especially when the number of codebooks (m) increases. This trend is more pronounced for larger matrix sizes, such as M = 8192. The measurements confirm that even at higher bit precisions, CodeGEMM maintains competitive performance compared to FP16 baselines, while offering a flexible trade-off between accuracy and efficiency depending on the (m, v) configuration.

We sincerely appreciate the insightful suggestions provided. If there are additional suggestions or concerns, we would be more than happy to incorporate them.

Weaknesses

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 .

• We agree, and this is indeed a clear limitation of our current work. We believe that model compression serves different goals depending on the deployment scenario. If the primary objective is to minimize memory footprint with minimal accuracy degradation, then using larger codebooks such as b = 16 is certainly effective.
• However, when practical speedup is also a goal, algorithm-hardware co-design becomes essential. To support this, we deliberately introduced constraints during the design of our compression algorithm to ensure compatibility with on-chip memory and efficient kernel execution.
• As shown in Table 3, recent calibration methods such as PV-Tuning offer promising ways to recover much of the accuracy lost under aggressive quantization. Notably, PV-Tuning was particularly effective in the b = 8 regime, where degradation was previously much higher compared to b = 16. We expect that as more advanced tuning and compression algorithms emerge, the remaining accuracy gap will continue to shrink even under shared memory constraints.

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 .

• We apologize for the omission. To support reproducibility, we provide the exact version information of the software environment used in our experiments: python==3.12.3, CUDA==12.6, torch==2.6.0, transformers==4.46.3, lm-eval==0.4.5
• We also fully agree that including error bars improves transparency. For this rebuttal, we have added 2-sigma error margins over 128 measurement samples to all reported numbers in our responses to the reviewers. While we are unable to modify the manuscript at this stage, we will include error bars in all relevant results in the camera-ready version.

Questions

The focus is on latency and perplexity; energy consumption and DRAM bandwidth savings which are critical in production are unreported.

Have you measured DRAM traffic reductions, and can you provide quantitative comparisons to dequantization-based kernels?

• Thank you for the suggestion. We agree on the importance of hardware utilization and conducted a set of experiments to compare the efficiency of different kernels.
• To evaluate DRAM traffic and power efficiency, we used nvidia-smi queries [1] sampled every 100 milliseconds for 10 seconds, then averaged the results. The results show that CodeGEMM achieves significantly better compute efficiency (FLOPS per watt) compared to dequantization-based kernels.
• The table below summarizes the results for various methods on a matrix multiplication workload with M = 1, N = 28672, and K = 8192. All values in parentheses indicate 2-sigma error margins over 128 measurement samples.

• These results highlight that CodeGEMM not only achieves lower latency but also delivers higher energy efficiency and better memory subsystem utilization, suggesting reduced and more structured DRAM access compared to dequantization-based approaches. We will include this analysis in the final version.

[1] https://docs.nvidia.com/deploy/nvidia-smi/index.html

Can you provide detailed microbenchmarks separating the time spent building versus reading the Psumbook, especially for varying $t_w$ and batch sizes?

• Thank you for the insightful suggestion. While it is challenging to measure precise cycle-level timing for each kernel phase due to concurrent execution across multiple SMs on the GPU, we performed controlled measurements by isolating execution to a single SM in order to estimate the relative cost of building and reading the Psumbook.
• The table below reports the percentage of execution cycles spent on each phase for various tile widths $t_w$ and batch sizes M. For a fixed $t_w$, increasing M shows that the ratio between building and reading remains largely stable, indicating that the Psumbook construction cost is effectively amortized across the batch. For a fixed M, increasing $t_w$ tends to increase the proportion of time spent on building in small matrices, but decrease it in large matrices.

How does CodeGEMM fare when coupled with codebooks generated by TurboQuant (https://arxiv.org/abs/2504.19874)?

• Thank you for the valuable suggestions. TurboQuant is an online vector quantization method primarily designed for KV‑cache compression. This aligns well with CodeGEMM, which dynamically operates on precomputed codebooks and codes, replacing on-the-fly dequantization with psum lookups. CodeGEMM is agnostic to how the codebooks are trained; it only requires the centroids, codes, and the associated hyperparameters (m, b, v) to construct a Psumbook for the current input tile. As such, codebooks generated by TurboQuant can be directly used with CodeGEMM, provided that the per-codebook bitwidth and vectorization choices fit within on-chip memory constraints.
• While TurboQuant is designed for online KV‑cache quantization, CodeGEMM naturally extends to this use case. In attention matmuls with quantized activations, the Psumbook is constructed between the current query tile and the key or value centroids, followed by a lookup using the stored codes. This maintains the same memory and compute characteristics, although the Psumbook must be rebuilt for each new query tile.

Do you envision streamed or hierarchical Psumbook designs that allow $b>8$ without cache overflow, and what challenges do these impose?

• We agree that fitting the entire Psumbook in shared memory limits the maximum codebook size $b$. We are exploring two complementary extensions: streamed or hierarchical Psumbook designs. First, a streamed Psumbook that partitions the table along the codebook dimension and prefetches chunks from global memory with double buffering, overlapping data transfer and lookup. This avoids cache overflow but introduces additional bandwidth demand and synchronization overhead. Second, a hierarchical Psumbook that decomposes a wide codebook into coarse and residual sub-codebooks. Partial sums are built separately for each level and then aggregated, which reduces the $2^b$ memory blow-up at the cost of extra lookups and combination logic. The main challenges include hiding global memory latency, preventing bank conflicts under larger tiles, and managing register and shared memory usage to sustain occupancy. Additionally, if the computational overhead of Psumbook building grows beyond a certain point, the net performance gain of CodeGEMM may diminish.

We sincerely appreciate the insightful suggestions provided. Throughout the rebuttal period, we have made our best effort to address the reviewers' feedback and strengthen the paper accordingly. That said, we recognize that there may still be areas for improvement. If there are additional suggestions or concerns, we would be more than happy to incorporate them. Thank you again for your thoughtful and constructive comments.

Weaknesses

1- 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.

• Thank you for the insightful feedback. We agree with your observation. In the context of LLM post-training quantization, 2-bit quantization still leads to notable accuracy degradation, especially under extreme compression. While advanced techniques such as PV-Tuning help mitigate this issue to some extent, a performance gap remains. We believe that with sufficient resources, quantization-aware training could be a promising direction to further narrow this gap.
• Our primary goal with CodeGEMM is to serve as a practical kernel-level bridge between model compression and real-world speedup. We are optimistic that as better quantization algorithms continue to emerge, CodeGEMM will enable them to realize their full latency-throughput benefits on modern hardware.

2- It is suggested to include accuracy vs. latency and accuracy vs. throughput plots corresponding to Table 3 and Table 4.

• We agree that figures are better suited than tables for illustrating trade-offs. Although we are unable to modify the manuscript during the rebuttal phase, we will include these plots in the camera-ready version, if accepted. We believe the visual representation will clearly highlight the accuracy-efficiency Pareto front achieved by CodeGEMM under different quantization settings.

3- It is recommended that the comparison with these previous approaches [1, 2] be demonstrated in the paper.

• Thank you for the helpful suggestion. While the reviewer recommended evaluations on both ParetoQ and VPTQ, we were only able to conduct experiments with VPTQ due to the lack of publicly available LLaMA 3 checkpoints for ParetoQ. VPTQ, which is also a vector-quantization-based method, shows comparable accuracy. However, similar to AQLM, it relies on dequantization-based kernels and therefore exhibits lower throughput than the FP16 baseline. We appreciate the suggestion, which helped enrich our study, and we plan to include this result in the camera-ready version if accepted.

Questions

Is 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?

• Our approach is independent of the total model size and scales well to extremely large models such as LLaMA 3.1–405B. The only memory constraint lies in the size of on-chip memory used to store the Psumbook. The size of Psumbook depends only on codebook parameters (e.g., number of codebooks and vector length), not on the number of parameters. Importantly, the performance of CodeGEMM improves with larger problem sizes. As shown in the table below, CodeGEMM achieves greater speedup over cuBLAS as the matrix dimensions grow. Each matrix shape corresponds to configurations used in LLaMA 3 models of varying sizes, including 1B, 8B, 70B, and 405B. (Values in parentheses indicate ±2σ error over 128 measurement samples.)

We sincerely appreciate the insightful suggestions provided. Throughout the rebuttal period, we have made our best effort to address the reviewers' feedback and strengthen the paper accordingly. That said, we recognize that there may still be areas for improvement. If there are additional suggestions or concerns, we would be more than happy to incorporate them. Thank you again for your thoughtful and constructive comments.