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

文章标题：CodeGEMM：一种以码本为中心的量化大语言模型高效GEMM方法
作者/机构：Gunho Park, Jeongin Bae, Byeongwook Kim, Baeseong park, Jiwon Ryu, Hoseung Kim, Se Jung Kwon, Dongsoo Lee / NAVER Cloud

A1 主要贡献

核心问题：随着大型语言模型（LLMs）规模的不断增大，其部署面临着巨大的内存和存储挑战。仅权重量化是一种有效的模型压缩技术，尤其是在极低比特（如2-bit）场景下，基于码本的量化方法能保持较高的准确率。然而，现有的基于码本的量化核依赖于“解量化”过程，即在计算时需要反复获取质心并重构权重，这带来了巨大的延迟和缓存压力，甚至可能抵消量化带来的内存优势。

研究目标：本文旨在设计一种高效的通用矩阵乘法（GEMM）核，专门用于支持基于码本的量化，以在极低比特量化下实现实际的推理加速，并克服现有解量化方法的开销。

创新点 (CodeGEMM)：
本文提出了CodeGEMM，一种以码本为中心的GEMM核，其核心思想是取代传统的解量化步骤。具体创新如下：
1. 引入Psumbook：CodeGEMM不将整个码本加载到缓存中，而是预先计算质心向量和输入激活值之间的所有内积，并将这些部分和（Partial Sums, Psums）结果存储在一个轻量级的“Psumbook”中。这种方法避免了在计算过程中通过代码索引查找质心的解量化操作。
2. 消除冗余计算和降低缓存占用：在推理时，CodeGEMM直接使用代码索引从Psumbook中查询预先计算好的部分和，从而消除了重复计算，并显著减少了片上缓存的空间占用，因为存储的是标量内积结果而非完整的质心向量。
3. 统一且灵活的内核：CodeGEMM内核支持广泛的码本超参数（如码本数量、向量长度、分组大小），使得用户可以在一个统一的实现下系统地探索延迟、内存和准确率之间的权衡，为不同用例提供优化。
4. 显著的性能提升：在Llama-3模型上的实验表明，与现有先进的基于码本的量化方法相比，CodeGEMM在2-bit配置下，在8B模型上实现了1.83倍的加速，在70B模型上实现了8.93倍的加速，同时保持了相当的准确率，并提升了计算效率和内存子系统利用率。

图 1: 基于码本的量化模型矩阵乘法核的比较。基于解量化的核执行即时解量化，需要将整个码本加载到缓存中。相比之下，CodeGEMM预计算部分和并将其存储在Psumbook中，从而消除了量化开销和冗余计算。

A3 背景知识

2.1 仅权重量化

模型量化的驱动力与挑战。随着生成式语言模型规模越来越大，内存需求也随之增加，模型量化已成为减少内存占用、最小化模型大小和提高推理效率的关键技术。量化通常有两种应用方式：（1）仅量化权重，或（2）同时量化权重和激活值。然而，由于激活值的动态范围和极端离群值（被称为“巨量激活”【索引23，Massive activations in large language models，2024，arXiv】，其数值超过典型值2000倍以上），激活量化面临巨大挑战。

仅权重量化方法的发展。为了应对这些挑战，许多研究集中在仅权重量化上。GPTQ【索引6，Gptq: Accurate post-training quantization for generative pre-trained transformers，2022，arXiv】利用近似二阶信息将权重压缩到4比特，且准确率损失极小。AWQ【索引13，Awq: Activation-aware weight quantization for on-device llm compression and acceleration，2024，Proceedings of Machine Learning and Systems】根据激活值的大小来缩放显著权重，以减少量化引起的误差。仅权重量化实现了紧凑的模型压缩，通过专门的内核【索引6，Gptq: Accurate post-training quantization for generative pre-trained transformers，2022，arXiv；索引13，Awq: Activation-aware weight quantization for on-device llm compression and acceleration，2024，Proceedings of Machine Learning and Systems；索引20，Lut-gemm: Quantized matrix multiplication based on luts for efficient inference in large-scale generative language models，2024，arXiv；索引11，Finequant: Unlocking efficiency with fine-grained weight-only quantization for llms，2023，arXiv】减少了内存占用并加速了推理。然而，在极低精度下，它仍然会遭受显著的准确率下降。例如，QuIP【索引3，Quip: 2-bit quantization of large language models with guarantees，2024，Advances in Neural Information Processing Systems】使用随机旋转矩阵来缓解这个问题，实现了可接受的准确率，但在极低比特级别仍然面临挑战。为了克服这些限制，非均匀量化方法被引入【索引5，Extreme compression of large language models via additive quantization，2024，arXiv；索引25，Quip#: Even better llm quantization with hadamard incoherence and lattice codebooks，2024，arXiv；索引14，Vptq: Extreme low-bit vector post-training quantization for large language models，2024，arXiv；索引27，Gptvq: The blessing of dimensionality for llm quantization，2024，arXiv】。与传统的均匀量化不同，非均匀方法在表示更广泛的值范围方面提供了更大的灵活性，从而在极低比特量化场景中提高了准确率【索引5，Extreme compression of large language models via additive quantization，2024，arXiv】。

2.2 基于码本的量化

基于码本的量化研究现状。学术界已进行了广泛研究，利用基于码本的量化技术来压缩大型语言模型，这得益于其非均匀特性和表示广泛值范围的灵活性【索引5，Extreme compression of large language models via additive quantization，2024，arXiv；索引16，Pv-tuning: Beyond straight-through estimation for extreme llm compression，2024，arXiv；索引25，Quip#: Even better llm quantization with hadamard incoherence and lattice codebooks，2024，arXiv；索引26，Qtip: Quantization with trellises and incoherence processing，2024，Advances in Neural Information Processing Systems；索引27，Gptvq: The blessing of dimensionality for llm quantization，2024，arXiv】。具体来说，GPTVQ【索引27，Gptvq: The blessing of dimensionality for llm quantization，2024，arXiv】将GPTQ框架扩展到基于码本的量化，将一个或多个列的量化与剩余未量化权重的更新交错进行。同样，AQLM【索引5，Extreme compression of large language models via additive quantization，2024，arXiv】为大型语言模型引入了一种多码本量化的改进方法，并提出了一种使用校准数据集优化码本的方法。QuIP#【索引25，Quip#: Even better llm quantization with hadamard incoherence and lattice codebooks，2024，arXiv】采用了与QuIP【索引3，Quip: 2-bit quantization of large language models with guarantees，2024，Advances in Neural Information Processing Systems】类似的平滑方法，在将权重映射到结构化格码本之前，应用旋转矩阵将权重转换到一个能最小化最坏情况量化误差的空间。类似地，QTIP【索引26，Qtip: Quantization with trellises and incoherence processing，2024，Advances in Neural Information Processing Systems】在此基础上，将基于旋转的平滑与格编码量化相结合，以进一步提升性能。

本文采用的加性码本量化策略。在本文中，我们基于AQLM【索引5，Extreme compression of large language models via additive quantization，2024，arXiv】中采用的加性码本量化策略，该策略将权重表示为多个质心向量的和。与依赖于矩阵变换（如旋转）的平滑方法相比，加性码本量化通过避免运行时的变换开销，提供了一种更直观且推理效率更高的替代方案。这种形式化方法支持基于快速查找的推理，并通过仔细控制超参数（如码本数量和向量长度），在准确率、内存占用和延迟之间实现了灵活的权衡。

码本量化过程详解。基于码本的量化通过将原始权重表示为一组向量，并用相应的质心向量来近似每个单独的向量，从而减少内存占用。如图2所示，一个大小为($M \times K$)的权重矩阵使用几个关键超参数进行量化，生成码本和相应的代码，这些超参数包括向量长度($v$)、分组大小($g$)、每个代码的比特数($b$)和码本数量($m$)。第一步，参数$v$指定了将($M \times K$)权重矩阵划分为向量的粒度。然后，每个向量通过一个缩放因子进行归一化，以便于权重向量的高效聚类（步骤1）。归一化分组大小$g$（总是大于或等于$v$）定义了执行归一化的范围，确保所有向量元素的一致缩放。具体来说，鉴于每个向量的长度为$v$，总共有$g/v$个权重向量被视为一个组进行归一化。

图 2: 一个(4 × 32)权重矩阵的量化过程示意图，其中b = 2, m = 1, v = 8, g = 16。

质心生成与码本构建。归一化之后，向量被聚类并映射到质心值，这些质心值作为每个簇的代表值（步骤2）。一种广泛用于确定这些质心向量的方法是K-means聚类算法，它将数据划分为预定数量的簇【索引5，Extreme compression of large language models via additive quantization，2024，arXiv；索引27，Gptvq: The blessing of dimensionality for llm quantization，2024，arXiv；索引14，Vptq: Extreme low-bit vector post-training quantization for large language models，2024，arXiv】。该算法迭代地初始化簇质心，并根据选定的距离度量将每个数据点分配到最近的簇，重复此过程直到收敛。

多码本概念与量化误差。每个质心向量被分配一个唯一的索引，称为代码（code）。簇的数量设置为$2^b$，其中$b$表示每个代码的比特数，允许表示从0到$2^b - 1$的值。然后构建一个码本，将每个代码映射到其对应的质心向量（步骤3），随后可以通过使用该码本来解码代码来重构原始权重值。为了最小化量化误差，提出了多码本概念【索引5，Extreme compression of large language models via additive quantization，2024，arXiv；索引2，Additive quantization for extreme vector compression，2014，Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition】，该概念将量化过程执行$m$次，构建$m$个码本，并通过将它们的值相加来更准确地表示原始权重值。

平均每权重比特数的计算。由于应用基于码本的量化使得$v$个全精度值由$m$个不同$b$比特精度的代码来表示，因此每个权重的平均比特数由这些超参数决定。具体而言，每个权重的比特数是通过将总量化比特数除以权重数量来计算的。总比特大小以及因此产生的平均每权重比特数$\bar{q}$可以表示如下：

$$\begin{aligned} \begin{aligned} \bar{q} & =\frac{S_{codebook}+S_{code}+S_{norm}}{M \cdot K} \\ & =\frac{16 \cdot m \cdot 2^{b} \cdot v+b \cdot m \cdot M \cdot \frac{K}{v}+16 \cdot M \cdot \frac{K}{g}}{M \cdot K}, \end{aligned} \end{aligned}$$

其中，码本中每个向量元素的精度设置为FP16。方程1还揭示了多个超参数组合可以实现相同的平均每权重比特数。例如，如表1所示，包括($v, m, b, g$) = (4, 1, 8, -1)和($v, m, b, g$) = (16, 3, 8, 32)在内的多种组合，其平均每权重比特数几乎相同，都为2左右。由于我们的目标是在极低比特量化下保持模型准确率，因此在比特精度和准确率之间实现最佳权衡至关重要。然而，这些超参数对模型准确率和推理延迟的影响在很大程度上尚未被探索，需要通过全面的实验研究进行广泛调查。在本文中，我们分析了具有相似平均每权重比特数的不同超参数组合如何影响内存占用和性能之间的权衡，为优化量化模型提供了宝贵的见解。

表 1: 不同量化配置下每个权重的平均比特数。$q_{code}$、$q_{codebook}$和$q_{norm}$分别代表分配给代码、码本和组归一化因子的平均比特数。$g = -1$表示逐行组归一化。

2.3 量化LLM的内核

现有量化内核的进展与局限性。最近的仅权重量化技术旨在将比特宽度降低到4比特以下以提高内存效率，同时也专注于实现实际的加速。例如，GPTQ【索引6，Gptq: Accurate post-training quantization for generative pre-trained transformers，2022，arXiv】和AWQ【索引13，Awq: Activation-aware weight quantization for on-device llm compression and acceleration，2024，Proceedings of Machine Learning and Systems】分别为均匀量化提供了INT3-FP和INT4-FP内核，并附带其量化方法。虽然这些内核相对于FP16 cuBLAS基线实现了延迟降低，但它们未能提高计算效率。这一限制源于其优势主要来自量化带来的更高效的数据移动，而非计算性能本身。为了解决这个问题，LUT-GEMM【索引20，Lut-gemm: Quantized matrix multiplication based on luts for efficient inference in large-scale generative language models，2024，arXiv】引入了一种利用查找表来降低计算复杂度的内核，它利用BCQ格式来表示量化权重。

基于码本量化的现有内核瓶颈。特别是在基于码本的量化中，当前的内核【索引5，Extreme compression of large language models via additive quantization，2024，arXiv】依赖于解量化过程。此过程通过使用代码从码本中检索相应的质心来重构原始权重。这些代码作为查询码本的索引，相应的质心被检索出来并用作解量化后的权重。($M \times K/v$)的代码矩阵检索长度为$v$的质心向量，重构出一个与原始矩阵($M \times K$)形状相同的解量化矩阵。随后与输入矩阵的矩阵乘法操作与原始权重矩阵的方式相同。为了支持高效的解量化，当前内核将整个码本加载到可编程缓存中，以实现质心向量的快速检索。然而，这种方法在码本加载阶段导致了显著的内存写入开销。此外，随着码本大小的增加，将整个码本加载到可编程缓存中变得更具挑战性。例如，在AQLM (1×16) 内核中，码本需要$2^{16}$个长度为$v = 8$的质心，每个都用FP16表示。这需要$2^{16} \times 1 \times 8 \times 2$ 字节（=1MB）的共享内存，远远超过了A100（164KB）和H100（224KB）的容量。因此，码本无法驻留在共享内存中，必须反复从DRAM中获取，从而显著增加了延迟。此外，虽然这种解量化方法有效优化了数据移动，但它没有降低计算复杂度，其计算复杂度与原始矩阵乘法相同。

A2 方法细节

3 CodeGEMM: 基于码本的GEMM

动机。图1展示了基于码本的量化中的矩阵乘法过程，并突出了两个关键挑战。首先，将整个码本加载到可编程缓存中会产生巨大的开销，该开销随码本的大小而增加，并且每个线程块都会重复此操作，导致在规模化部署中性能下降。其次，计算通常是冗余的，因为每个输入向量只与有限的一组质心相互作用，导致在整个输出矩阵中存在重复计算。为了解决这些挑战，我们提出了一种方法，旨在最小化码本加载开销并避免冗余计算，从而能够更有效地利用有限的片上缓存并提高整体效率。

方法论。CodeGEMM为量化LLM中的高效矩阵乘法引入了一种以码本为中心的方法。CodeGEMM不加载整个码本，而是在可编程缓存中存储质心向量和输入激活值之间内积的预计算结果。这种方法降低了空间复杂度，并消除了从码本中为每个代码检索质心的解量化操作。取而代之的是，内核重复利用预计算的结果，从而显著降低了计算复杂度。

CodeGEMM计算过程。图3展示了CodeGEMM的计算过程，描述了一个维度为($t_h \times t_w$)的权重块与维度为($t_w \times 1$)的输入块之间的矩阵乘法运算。首先，输入块被划分为维度为($t_w/v \times v \times 1$)的输入，以便与存储在码本中的质心进行高效计算（步骤1）。这些分段的输入表示为$x_k^j = x_{v \times j + k}$，其中$j \in [0, 1, \dots, t_w/v - 1]$，$k \in [0, 1, \dots, v - 1]$。相应的代码矩阵同样被划分为维度为($t_h \times t_w/v$)的代码块，以与输入块对齐进行有效计算。在我们的实现中，我们将$t_w = 32$和$t_h = 2048$设置为与硬件友好的分块策略对齐，并最大化每个线程块内Psumbook的重用。接下来，每个分段的输入与质心计算内积以生成部分和（Psums），然后将其作为Psumbook存储在可编程缓存中（步骤2）。Psumbook中的每个条目$p_i^j$计算如下：

$$p_i^j = \sum_{k=0}^{v-1} c_k^i \times x_k^j, \quad i \in [0, 1, \dots, 2^b - 1]$$

通过将Psumbook存储在可编程缓存中而不是完整的码本，CodeGEMM用一种更高效的、以代码为键检索预计算Psums的方式，取代了为每个代码获取并计算质心的传统过程。这不仅降低了计算复杂度，还降低了空间复杂度，因为缓存的是标量内积结果而不是完整的质心向量。最后，通过基于代码的索引从Psumbook中检索与每个输入对应的部分和，然后累加生成线程块的最终输出激活（步骤3）。与现有内核将整个码本加载到可编程缓存中相比，我们存储预计算内积结果的方法为基于码本的量化LLM实现了显著更高效的计算。

图 3: CodeGEMM核在基于码本的量化模型上的操作概览。1) 输入数据被重塑为向量以与码本维度对齐。2) 码本与输入向量之间的预计算内积被存储在可编程缓存的Psumbook中，显著减少了计算开销。3) 在计算过程中，代码从Psumbook中查询相应的部分和，然后累加以高效生成输出，无需即时解量化。

计算复杂度。在一个权重矩阵大小为($M \times K$)、输入矩阵大小为($K \times N$)的矩阵乘法运算中，标准GEMM内核的计算复杂度为$O(MNK)$。这个复杂度同样适用于基于解量化的内核，因为它们通过量化提高了数据移动效率，但并未降低总计算成本。

CodeGEMM的计算复杂度优势。相比之下，CodeGEMM通过预先计算质心向量和输入激活值之间的所有内积结果，将它们存储在Psumbook中，并用简单的检索操作替代重复计算，从而减少了计算工作量。这使得CodeGEMM能够实现比传统方法更低的计算复杂度。假设$M \gg 2^b$，CodeGEMM的计算复杂度表示如下：

$$\begin{aligned} \begin{aligned} C &=C_{build}+C_{read}=\mathcal{O}\left(m \cdot 2^{b} \cdot v \cdot \frac{K}{v} \cdot N\right)+\mathcal{O}\left(m \cdot M \cdot \frac{K}{v} \cdot N\right) \\ &=\mathcal{O}\left(m M N K\left(\frac{2^{b}}{M}+\frac{1}{v}\right)\right) \approx \mathcal{O}\left(M N K \cdot \frac{m}{v}\right), \end{aligned} \end{aligned}$$

其中，$C_{build}$和$C_{read}$分别表示构建Psumbook和从中检索值的计算复杂度。因此，与传统内核相比，CodeGEMM实现了$(m/v)$的计算缩减因子，其中$m$是码本数量，$v$是向量长度——这两者对于优化性能至关重要。这使得CodeGEMM能够像传统的基于解量化的内核一样通过量化提高数据移动效率，同时还降低了计算复杂度以实现更高的效率。

空间复杂度。基于解量化的GEMM内核将整个码本加载到可编程缓存中进行计算。在这种情况下，空间复杂度为$O(m \cdot 2^b \cdot v)$，这对应于码本的完整大小。相比之下，CodeGEMM的空间复杂度为$O(m \cdot 2^b \cdot t_w/v)$，其中$t_w$表示权重块的宽度。由于CodeGEMM存储的是预计算的内积结果而不是完整的质心向量，其空间复杂度与向量长度$v$成反比，从而导致更小的缓存占用。这种内存需求的减少使得CodeGEMM能够实现更高效的缓存利用，使其更适合于可编程缓存大小有限的加速器。

A4 实验环境与结果

实验环境

我们通过探索关键超参数之间的权衡来评估CodeGEMM，重点关注LLM压缩相关的三个主要指标：内存占用、延迟和准确率。

• 内存占用：使用每个权重的平均比特数$\bar{q}$来量化，并根据公式1计算。

• 延迟测量：

  • 内核级延迟：测量矩阵乘法的原始内核性能，假设线性层的形状与广泛采用的Llama-3【索引4，The llama 3 herd of models，2024，arXiv】架构相同。报告的延迟是单个Transformer解码器块中所有线性层内核执行时间的总和，不进行层融合。
  • 端到端吞吐量：使用HuggingFace【索引29，Transformers: State-of-the-art natural language processing，2020，Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations】库提供的Llama实现进行测量，并启用了层融合。

• 硬件配置：所有延迟测量均在NVIDIA A100 80GB GPU上进行。

• 准确率评估：使用lm-eval-harness【索引7，The language model evaluation harness，2024，URL】基准测试套件，在标准任务上进行零样本和5样本设置的评估。

实验结果

内存占用与延迟
* 实验内容：比较了多种量化矩阵乘法方法在2-bit量化的Llama3模型（8B和70B）上的内核级延迟。
* 实验结果：如表2所示，LUTGEMM【索引20，Lut-gemm: Quantized matrix multiplication based on luts for efficient inference in large-scale generative language models，2024，arXiv】在均匀量化上表现出色。QuIP#【索引25，Quip#: Even better llm quantization with hadamard incoherence and lattice codebooks，2024，arXiv】和QTIP【索引26，Qtip: Quantization with trellises and incoherence processing，2024，Advances in Neural Information Processing Systems】通过高度优化的融合内核展示了有竞争力的延迟。AQLM【索引5，Extreme compression of large language models via additive quantization，2024，arXiv】（2×8配置）有中等延迟降低，但AQLM-1×16由于其巨大的码本（$2^{16}$个条目）导致延迟显著增加。相比之下，CodeGEMM在相同平均比特精度下，分别比FP16基线和AQLM实现了高达2.18倍和1.64倍的加速。
* 分析结论：CodeGEMM的性能提升归功于使用预计算内积降低了计算复杂度和高效利用共享内存。图4(a)进一步展示了在Llama-3.1-8B模型单批次操作下，内存占用与延迟的关系。结果表明，分组大小g越小，内存占用和延迟越高，尤其是在g值较小（如g=v）时延迟急剧上升。

表 2: 2-bit量化的Llama3模型（8B和70B）中量化矩阵乘法的内核级延迟（µs）。CodeGEMM始终比其他方法延迟更低。

效率与利用率
* 实验内容：使用nvidia-smi【索引18，Nvidia system management interface (nvidia-smi) documentation，2025，NVIDIA Corporation】遥测数据，在M=1, N=28672, K=8192的矩阵乘法工作负载上测量DRAM流量代理和功率效率。
* 实验结果：如表3所示，CodeGEMM的计算效率（GFLOPS/W）远高于基于解量化的内核。括号内的值为128个样本上的2-sigma误差范围。
* 分析结论：除了更低的延迟，CodeGEMM还显示出更高的能效和改进的内存子系统利用率，表明其相对于基于解量化的方法，DRAM访问更少且更有序。

表 3: 在GEMV（M, N, K）=（1, 28672, 8192）上的内核级性能评估。性能指标来自nvidia-smi遥测。CodeGEMM比基于解量化的码本内核提供更高的功率效率和改进的内存利用率。

内存占用与准确率
* 实验内容：在各种超参数配置下，评估量化模型的内存占用与准确率（以WikiText-2数据集上的困惑度衡量）之间的关系。
* 实验结果：如图4(b)所示，随着平均每权重比特数的增加，模型的困惑度降低。细粒度的组归一化虽然增加了内存占用，但能减少量化误差。在逐行组归一化（g = -1）下，增加码本数量m可以提高准确率。但当g变小时，增加m带来的准确率增益减小。
* 分析结论：CodeGEMM在块级码本优化【索引5，Extreme compression of large language models via additive quantization，2024，arXiv】的基础上，结合细粒度组归一化，实现了性能增强。

图 4: Llama-3.1-8B模型在不同配置下的延迟与准确率权衡。

吞吐量与准确率
* 实验内容：比较CodeGEMM与均匀量化（FlexRound【索引12，Flexround: Learnable rounding based on element-wise division for post-training quantization，2023，International Conference on Machine Learning】）和码本量化（AQLM【索引5，Extreme compression of large language models via additive quantization，2024，arXiv】）的准确率-吞吐量权衡。同时应用了PV-Tuning【索引16，Pv-tuning: Beyond straight-through estimation for extreme llm compression，2024，arXiv】后量化校准方法。
* 实验结果：如表4和图5(a)所示，FlexRound在2-bit下吞吐量高但准确率最差。AQLM(1×16)准确率最高但吞吐量极差。AQLM(2×8)吞吐量提升有限。CodeGEMM在保持有竞争力的准确率的同时，吞吐量显著提高。与PV-Tuning结合后，效果更佳。CodeGEMM m1v4配置在吞吐量和准确率上均优于m2v8。
* 分析结论：CodeGEMM在Llama-3.1-8B模型上，以相当的准确率实现了比基于解量化的内核高1.83倍的端到端加速，提供了更优的吞吐量-准确率权衡。

表 4: 不同量化方法在Llama-3.1-8B-Instruct模型上的准确率。Average列代表所有任务的平均准确率，包括MMLU（5-shot）和WinoGrande（WG）、HellaSwag（HS）、ARC-Easy（ARC-E）和ARC-Challenge（ARC-C）等0-shot任务。

扩展到更大模型
* 实验内容：在Llama-3.1-70B模型上评估CodeGEMM的可扩展性，并与GPTQ【索引6，Gptq: Accurate post-training quantization for generative pre-trained transformers，2022，arXiv】、FlexRound和AQLM进行比较。
* 实验结果：如表5和图5(b)所示，随着模型从8B增至70B，CodeGEMM相对于AQLM的延迟改进更为显著。AQLM(1×16)在70B模型上的吞吐量严重下降。CodeGEMM通过细粒度组归一化，在准确率上匹配AQLM(1×16)的同时，实现了8.93倍的吞吐量优势。
* 分析结论：CodeGEMM能有效扩展到大型模型，并保持有竞争力的性能。

表 5: 不同量化方法在Llama-3.1-70B模型上的准确率。

图 5: Llama-3.1-8B和Llama-3.1-70B模型在不同配置下的延迟与准确率权衡。

A7 补充细节

5 相关工作

基于查找表的计算。先前的一些工作利用查找表（LUT）来提高计算效率，无论是在内核层面【索引20，Lut-gemm: Quantized matrix multiplication based on luts for efficient inference in large-scale generative language models，2024，arXiv；索引15，Look-up mai gemm: Increasing ai gemms performance by nearly 2.5 x via msgemm，2023，arXiv】还是在专门的硬件设计中【索引28，T-mac: Cpu renaissance via table lookup for low-bit llm deployment on edge，2025，Proceedings of the Twentieth European Conference on Computer Systems；索引17，Lut tensor core: Lookup table enables efficient low-bit llm inference acceleration，2024，arXiv；索引19，Figlut: An energy-efficient accelerator design for fp-int gemm using look-up tables，2025，2025 IEEE International Symposium on High Performance Computer Architecture (HPCA)】。然而，这些方法主要支持均匀或二进制编码的量化格式，并未扩展到基于码本的非均匀量化。相比之下，CodeGEMM被设计为自然支持基于码本的量化，同时通过适当的码本配置也能够推广到均匀和二进制格式。具体来说，通过将质心定义为$c \in \{0, 1\}^v$用于均匀量化和$c \in \{-1, 1\}^v$用于二进制量化，可以适应这两种格式。这种通用性使得CodeGEMM特别适合作为固定功能的ASIC内核部署，因为它能够在一个统一的硬件架构下支持多种量化格式。

基于码本的量化。除了加性码本，最近的工作探索了将旋转与码本本身集成的高度结构化设计。QuIP#【索引25，Quip#: Even better llm quantization with hadamard incoherence and lattice codebooks，2024，arXiv】和QTIP【索引26，Qtip: Quantization with trellises and incoherence processing，2024，Advances in Neural Information Processing Systems】都将其基于格和网格编码的码本与推理时的平滑旋转相结合，并且每项工作都提供了一个精心手工调优的内核，该内核将旋转与后续的查找操作融合在一起。相比之下，我们所针对的加性码本仅需要轻量级的表查找和累加，无需旋转，因此解量化过程保持简单和快速。这种设计使得CodeGEMM能够在提供有竞争力的2-bit准确率的同时，保持一个单一、可重用的内核，该内核可在广泛的量化设置中扩展。

A5 结论

总结
本文介绍了CodeGEMM，一种针对基于码本的量化模型的高效矩阵乘法方法。CodeGEMM不将完整的码本加载到可编程缓存中，而是预先计算结果并将其存储在Psumbook中，从而降低了空间和计算复杂度。实验表明，CodeGEMM在延迟方面优于当前最先进的内核，同时在极低比特量化下保持了高准确率。

局限性
1. 缓存大小限制：CodeGEMM要求Psumbook必须驻留在片上共享内存中，这限制了可用的码本大小。非常大的码本（例如，b=16，有$2^{16}$个条目）会超过当前GPU的缓存限制。在实践中，我们通过将每个码本的宽度固定为b=8，并通过细粒度组归一化来恢复准确率，这种方法增加的延迟可以忽略不计，并在现实硬件约束下产生了有利的准确率-效率权衡。
2. 大批量吞吐量：当批量较大时（例如，M > 32），CodeGEMM的性能不如基于Tensor Core的cuBLAS。这种行为在基于CUDA Core的量化GEMM内核中很常见，反映了当前商用GPU架构的限制，而非算法本身的低效。

总的来说，虽然CodeGEMM排除了极大的码本，并且没有针对大批量进行优化，但其硬件-算法协同设计解决了传统码本量化的关键效率瓶颈，降低了计算和空间复杂度。

A6 附录

A.1 Psumbook构建与读取的成本分解

实验设计。我们通过将执行隔离到单个SM（Streaming Multiprocessor）并将运行时间分为两个阶段来量化构建和使用Psumbook的相对成本：为每个tile构建Psumbook和在主累加循环中读取它。我们扫描了tile宽度$t_w$和批量大小M，同时固定了(N, K)，并报告了两种内核变体（m2v8和m1v4）中每个阶段所占用的执行周期百分比。

结果分析。在不同设置中，构建/读取的比例在固定$t_w$时对M保持稳定，这表明Psumbook的构建成本在整个批次中被分摊。在固定M时，较大的$t_w$会增加小矩阵的构建份额，而减少大矩阵的构建份额。典型的范围是：m2v8的构建周期约占28-46%，读取周期约占54-72%；m1v4的构建周期约占20-42%，读取周期约占58-80%。

表 6: 在不同$t_w$和M下，用于构建与读取Psumbook的周期份额（%）。

A.2 瓦片大小敏感性

实验设计。我们重新审视了我们对瓦片（tile）维度的启发式选择，并对$t_w \in \{32, 64, 128\}$和$t_h \in \{2048, 4096\}$在代表性矩阵形状上进行了系统性扫描。

结果分析。我们发现$t_h=2048$在广泛的工作负载中始终能产生最佳性能，这支持了我们最初的选择。对于水平瓦片，$t_w=32$这样较小的值对于相对较小的矩阵效果很好，而$t_w=64$则倾向于在大型矩阵上表现更佳。我们将此归因于更粗粒度的瓦片划分减少了内核启动开销，并提高了大规模下的部分和规约效率。表7总结了在M=1时，(N, K) ∈ {(4096, 4096), (8192, 8192)}的代表性结果。总的来说，这些观察结果证明了我们对中小型矩阵形状默认选择($t_w, t_h$) = (32, 2048)的合理性，而随着矩阵变大，首选$t_w=64$。

表 7: 瓦片维度对端到端延迟（µs）的影响。

A.3 更高比特精度的影响

实验设计。我们还使用内核配置(g=128, b=8, tw=32, th=2048)测量了更高平均比特精度下的延迟。作为参考，我们包含了FP16 cuBLAS的延迟。

结果分析。正如预期的那样，增加每个权重的有效比特数（通过增加码本数量m或减小向量长度v）通常会提高延迟，并且这种趋势对于更大的矩阵（例如，M=1, N=K=8192）更为明显。即使在更高的比特精度下，CodeGEMM与FP16基线相比仍然具有竞争力，同时通过(m, v)配置提供了灵活的准确率-效率权衡。

表 8: 在(g=128, b=8, tw=32, th=2048)下，更高有效比特精度的矩阵乘法延迟（µs）。FP16 cuBLAS作为参考。

A.4 批量大小敏感性与公平的cuBLAS核算

公平比较。为了进行公平比较，我们在报告FP16 cuBLAS延迟时包含了去量化阶段，因为在基于码本的流程中，去量化必须在GEMM之前进行。

结果分析。在这种核算方式下，即使在批量大小为8和16时，CodeGEMM仍然具有竞争力。在数据中心部署中，连续批处理可以聚合请求并增加解码期间的有效批量大小；然而，许多场景仍然在小批量下运行（例如，设备端推理）。在基于CUDA核心的量化内核中观察到的批量大小敏感性反映了架构限制，如占用率限制和共享内存带宽。这种限制是最近的方法（例如，QuIP#和QTIP）所共有的，并不表示算法效率低下。表9报告了Llama-3-8B上的线性层延迟。

表 9: Llama-3-8B解码器块内线性层的总延迟（µs）与批量大小的关系。

A.5 跨问题规模的额外基准测试

扩展评估。我们将评估扩展到了广泛的矩阵形状(M, N, K)扫描。延迟是在固定的硬件和软件堆栈上进行端到端测量的，所有内核都在相同的工具链下编译。

结果分析。表10报告了完整结果。cuBLAS在Tensor Cores上运行，随着批量大小M的增长，其延迟往往保持相对稳定，而最近提出的量化内核（包括我们的）在CUDA核心上执行，并且通常随着M的增大显示出延迟增加。在这个CUDA核心类别中，CodeGEMM在大型矩阵上始终表现良好，这些矩阵的算术强度和内存重用率很高。在实践中，该方法在各种(M, N, K)设置中保持竞争力，并且对于主导LLM推理的大规模矩阵乘法特别有效。

表 10: 在不同(M, N, K)配置下的内核延迟（µs）。