VQ-LLM: High-Performance Code Generation for Vector Quantization Augmented LLM Inference

发表时间: 2025-03 · arXiv:2503.02236 (HPCA 2025)

原文: https://arxiv.org/abs/2503.02236

Zihan Liu1,2, Xinhao Luo1,2, Junxian Guo1, Wentao Ni1, Yangjie Zhou1, Yue Guan1,2, Cong Guo3, Weihao Cui1,4, Yu Feng1, Minyi Guo1,2, Yuhao Zhu5, Minjia Zhang6, Chen Jin7, Jingwen Leng1,2,*
1Shanghai Jiao Tong University, 2Shanghai Qi Zhi Institute, 3Duke University, 4National University of Singapore, 5University of Rochester, 6University of Illinois Urbana-Champaign, 7Magik Compute

速读

一句话结论
本文提出了一个名为 VQ-LLM 的自动代码生成框架,通过协同优化向量量化中的码本缓存与计算数据流,将大语言模型向量量化推理的延迟降低了 64.36% 到 99.1%,使其在保持极低位宽和高精度的同时,运行速度比肩甚至超越主流的逐元素量化方法。

要解决什么问题
向量量化(Vector Quantization)将多个元素组成的向量作为一个整体压缩为单个索引,相比传统的逐元素量化,它能捕获跨维度信息,从而在 1 到 2 位的极低位宽下依然保持高精度。然而,现有的向量量化大语言模型在实际推理时,延迟往往比未量化的 FP16 版本还要高。这个卡点主要源于码本(存储聚类中心的查找表)的访存机制与计算数据流存在严重脱节。首先是码本访存效率极低。常规做法是将整个码本存入 GPU 的共享内存,这不仅占用了大量空间,导致单个流多处理器上能并发执行的线程块数量锐减,而且由于码本条目数远超共享内存的 bank 数量,加之反量化时的随机访问特性,会引发严重的 bank 冲突。其次是码本加载与后续计算不协调。在将向量量化集成到注意力机制或矩阵乘法时,不同线程块在处理不同词元时会重复从全局内存加载相同的码本,造成极大的冗余带宽消耗。同时,线程在寄存器中反量化出的数据排布,往往与后续计算(如矩阵乘加指令)所需的数据排布不一致,迫使数据必须先写回共享内存再重新读取,产生了毫无必要的片上内存往返开销。此外,向量量化算法配置繁多,手动为每种情况编写高效算子并不现实。

怎么做的
VQ-LLM 的核心思路是围绕码本的访存特征,设计一套自适应的层次化缓存与融合计算引擎,从而彻底绕开冗余访存和排布冲突的卡点。该框架主要由两个关键部件构成。第一个部件是码本缓存(Codebook Cache)。它摒弃了将码本一股脑塞进共享内存的做法,而是根据离线分析出的条目访问频率,将码本打散放置在 GPU 的多级存储中。极少访问的冷条目留在全局内存,中等频率的条目放入共享内存,而访问极度频繁的热点条目则直接锁定在线程的局部寄存器中,从而从根本上消除了最严重的 bank 冲突。为了不影响线程块的并发度,系统会利用计算内核的资源空隙来决定寄存器和共享内存的分配边界。第二个部件是基于码本的计算引擎,它包含两项核心设计。其一是“以码本为中心的数据流”,该设计改变了原有的任务切分方式,直接沿着码本切换的维度来对计算任务进行切分和并行化。这样每个线程块只需加载一次码本,消除了重复加载的冗余流量。为了平衡并行化带来的全局归约开销与码本加载开销,系统会自适应寻找一个切分因子,使得两者流量相等以达到最优:

$$Traffic_{Reduce} = Traffic_{Codebook}$$


其二是“以码本为中心的层次化融合”。针对反量化数据与后续计算数据排布不匹配的问题,该机制利用 GPU 的束内数据交换(warp shuffle)功能,直接在寄存器级别完成数据重排,彻底跳过了共享内存。其核心是通过异或指令在寄存器间交换数据:

$$register \gets shfl\_xor(register, offset)$$
系统会根据排布差异计算所需的交换次数,当交换次数较少(如小于 5 次)时触发寄存器级融合,否则回退到共享内存融合。整个框架通过模板和启发式策略,自动为不同的量化配置生成最优的算子代码。

效果如何
实验在 RTX 4090 和受限于带宽的 Tesla A40 GPU 上进行,测试了 Llama-7B 和 Llama-65B 模型。对比的基线既包括将码本存放在全局内存或共享内存的朴素向量量化实现,也包括代表逐元素量化路线的先进方法 AWQ(用于权重量化)和 QoQ(用于 KV 缓存量化,集成在 qServe 系统中)。在算子层面的量化结果显示,VQ-LLM 相比现有的开源向量量化实现,延迟降低了 64.36% 到 99.1%。在等效 4 位宽的设置下,VQ-LLM 的实际运行速度已经追平甚至超越了 AWQ 和 QoQ,例如在矩阵向量乘法和矩阵乘法中,其延迟仅为逐元素量化基线的 88% 和 96%。在端到端的大模型文本生成测试中(批大小 16,上下文 1024),VQ-LLM 取得了与 qServe 相当的 2.2 倍加速比,同时在 arc-challenge 任务上的准确率比 qServe 高出 2.5%。在带宽更受限的 Tesla A40 上,VQ-LLM 的加速优势进一步放大。作者也指出了该方法的局限性:在矩阵乘法(GeMM)算子中,由于未集成 cutlass 库复杂的底层分块策略,其绝对性能仍不及 FP16 基线,但由于该算子主要用于预填充阶段,对整体生成延迟影响有限。此外,当前框架暂未原生支持多 GPU 的张量并行分布式推理。

主要贡献

随着大型语言模型(LLMs)的巨大成功,神经网络对当前硬件(尤其是内存系统)施加了巨大压力。为了部署这些大模型,量化技术变得至关重要。传统的逐元素量化方法通常在压缩到4比特以下时会导致显著的精度损失。相比之下,向量量化(Vector Quantization, VQ)将多个元素组成的向量作为一个压缩单元,能够捕获跨元素的信息,从而在相同的低比特率下保持更高的模型精度。然而,将VQ算法带来的内存减少转化为实际的延迟改善面临着巨大挑战。

本文深入分析了将VQ集成到计算内核中的现有方法,发现其主要低效性源于VQ算法中码本(codebooks)的访问效率低下以及计算数据流的不协调。同时,VQ算法的多样性(如不同的向量大小和条目数量)以及LLM计算内核的多样性(如矩阵乘法和注意力计算)使得为每种特定情况手动编写高效的内核实现变得不切实际。

为了解决这些问题,本文设计并实现了VQ-LLM,这是一个高效的融合VQ内核生成框架。主要贡献如下:
1. 首次深入探讨了向量量化的性能问题,并使其在LLM推理中具有实际可行性。
2. 引入了一种名为“码本缓存(codebook cache)”的软件抽象,用于优化码本访问效率,自适应地将不同条目存储在GPU的内存层级中(包括全局内存、共享内存和寄存器)。
3. 围绕码本缓存,设计了一个高效的计算引擎,该引擎采用了“以码本为中心的数据流(codebook-centric dataflow)”和“分层融合(hierarchical fusion)”优化,以减少计算过程中的内存流量。
4. 提供了自适应启发式方法,以针对不同的VQ配置定制优化参数。
5. 实验表明,与现有的开源实现相比,VQ-LLM实现了64.36%至99.1%的延迟降低。与AWQ和QoQ等最先进的逐元素量化方法相比,VQ-LLM在同等比特宽度的延迟上表现接近甚至更好,并可能提供更高的精度。

典型的向量量化流程。
典型的向量量化流程。

背景知识/关键Observation/设计原则

向量量化(VQ)基础。相较于传统量化,VQ将多个元素的向量作为一个单元,并使用组织成码本的训练量化点将向量量化为单个元素。典型的VQ流程包括:首先将原始向量分割为子向量,然后进行$k$-means聚类将子向量分组为多个簇,并用最近的簇中心(码本条目)索引替换原始子向量。接着,收集原始子向量与其簇中心之间的残差,并进行新一轮量化。在重建原始数据(解量化)时,利用量化数据查找对应的码本条目,并在不同残差之间进行逐元素累加,最后拼接所有子空间的结果。整个过程中的向量大小(Vector size)、条目数(#Entry)和残差次数(Residual)均是可配置的。

大型语言模型(LLMs)基础。LLMs采用Transformer架构,其核心是多头注意力(MHA)。MHA的数学描述为:$MultiHead(Q,K,V) = Concat(head_1, \dots, head_h)W^O$,其中$head_i = Attention(Q=HW_i^Q, K=HW_i^K, V=HW_i^V)$,且$Attention(Q,K,V) = softmax(QK^T/\sqrt{d_k})V$。在文本生成中,LLMs通常包含预填充(prefill)阶段和解码(decode)阶段。为了在解码阶段高效重用先前计算的词元表示,通常利用键值(KV)缓存机制。

用于LLM加速的VQ。由于LLMs受限于内存带宽,权重和KV缓存占用了超过95%的内存,VQ凭借其优越的压缩比和重建质量成为加速焦点。通过跨维度信息,VQ能更好地捕捉数据分布特征,从而在相同等效比特宽下优于SOTA逐元素量化基线。然而,现有的VQ算法提供的内核延迟很高,使得其实际应用受限。在LLMs上下文中,解量化是主要瓶颈,因为每次计算前都需要解量化。
(上)VQ与逐元素量化的准确率对比,左侧为权重,右侧为KV缓存量化。(下)VQ(右)比逐元素量化(左)能更好地捕捉数据分布,利用了跨维度信息。

微基准测试研究设置。本文在一个RTX 4090 GPU上评估了具有32个头、头维度为128的Llama-7B注意力内核。研究了三种遵循CQ-2算法的VQ KV缓存实现:FP16-attn(基线)、VQ-attn-GC(码本存储在全局内存)和VQ-attn-SC(码本存储在共享内存)。
所研究的VQ内核的工作流。

低效性分析:码本访问效率低下。尽管VQ将KV缓存压缩至1/8,但VQ-attn-GC和VQ-attn-SC的性能均低于FP16基线。虽然VQ-attn-SC优于VQ-attn-GC,但其计算(SM)利用率下降了30%以上,这是因为VQ显著增加了共享内存占用,减少了每个SM上并发运行的线程块数量。此外,码本条目数量远超共享内存bank数(例如256个条目对32个bank),且访问是随机的,导致严重的bank冲突。结论是:将码本存储在快速片上缓冲区(如共享内存)是必要的,但并非易事。
(左)VQ-attn-GC和VQ-attn-SC相对于FP16-attn的延迟。(右)VQ-attn-SC的相对性能计数器。

低效性分析:码本加载与计算不协调。VQ版本的全局内存到共享内存流量高于FP16版本,这是因为将VQ集成到原始计算数据流中会导致码本的非协调和重复加载。例如,处理不同词元的线程块在处理所有通道时会访问和加载相同的码本。此外,共享内存到寄存器的流量也显著增加,原因是解量化数据的布局与后续计算所需的布局不匹配(解量化是按行进行的,而后续计算要求按列累加),迫使解量化数据必须经过共享内存进行重新排列。结论是:将VQ算法集成并融合到LLM内核中需要仔细协调码本加载与融合内核的计算数据流。
FP16-attn(内框)和VQ-attn-SC(外框)的数据流。
在注意力(解码)计算中,KV缓存的解量化数据布局与后续计算所需布局的对比。

VQ多样性带来的额外复杂性。最先进的VQ方法在向量大小、条目数、残差次数以及训练码本所用的张量部分上存在显著差异(如QuiP#-4、AQLM-3、GPTVQ-2、CQ-4等)。同时,与VQ算法关联的计算也多种多样(如GeMM、GeMV、注意力计算)。结论是:必须采用自适应解决方案,才能在各种VQ算法及其后续计算中实现最佳性能。

方法细节

VQ-LLM设计概览。基于上述挑战,本文设计了VQ-LLM。该框架引入了“码本缓存”软件抽象,通过离线分析码本条目的访问特征(冷、中、热),自适应地将条目存储在GPU的全局内存、共享内存和寄存器中。围绕码本缓存,设计了包含“以码本为中心的数据流”和“以码本为中心的分层融合”两个核心技术的计算引擎,并结合自适应启发式方法,根据算法和目标GPU配置自动选择最佳参数。
VQ-LLM设计概览。

码本缓存:设计直觉。单纯将整个码本放在共享内存中会导致共享内存使用量增加和严重的bank冲突。因此,我们提出根据访问频率将不同条目存储在不同的内存层级中:极少访问的条目存储在全局内存中以节省共享内存,最常访问的条目存储在线程局部寄存器中以消除bank冲突。分析表明,码本中超过一半的条目访问频率低于平均值,而少数“热”条目(访问频率高于$\mu + 3\sigma$)更容易引发bank冲突,这构成了码本缓存设计的基础。
VQ-GeMM内核中一个线程块的码本条目访问频率,配置为<12, 12, 2> (AQLM-3)。

码本缓存:实现细节。在实现中,我们采用了一种极度轻量级且可配置的基于重排序的静态映射机制。首先在基于分析的离线阶段,按访问频率降序对码本条目进行排序和重排序,确保最频繁访问的条目索引为0。接着,建立两个边界:$n_{reg}$和$n_{shared}$。我们将前$n_{reg}$个条目分配到寄存器,后续直至$n_{shared}$的条目分配到共享内存,剩余条目存储在全局内存。在运行时解量化中,通过简单的索引比较即可寻址码本条目。此外,这种基于频率的重排序是在张量级别全局进行的,因为分析表明这些热条目在张量的不同部分之间是一致被访问的。
张量不同部分的条目冷热情况。

码本缓存:自适应资源分配。共享内存和寄存器资源是有限的,过度使用会降低线程块的占用率。我们采用了一种自适应启发式方法来决定$n_{reg}$和$n_{shared}$。首先,识别两种资源使用中的“空闲区(slack)”,即在不影响并发性和GPU利用率的情况下可以占用的资源空间。不同的计算表现出不同的空闲区,可以通过离线分析得出。最后,通过将可用的空闲区除以单个码本条目的大小来确定$n_{reg}$和$n_{shared}$。
计算内核资源消耗及相应的硬件占用率。蓝色区域是可以在不影响性能的情况下使用的资源空闲区。

码本缓存:用户接口。我们提供了三个API供后续计算使用。第一个是Load,它接受存储在全局内存中的码本和内存空闲区,将码本加载到内存层级中,并返回缓存的码本及两个访问边界。第二个是Access,允许用户在解量化过程中访问特定条目,它利用边界来确定在何处定位条目。第三个是Switch,当算法针对张量的不同部分训练不同的码本时,该API便于根据用户处理的特定张量部分切换到新码本。

计算引擎:以码本为中心的数据流设计直觉。为了充分利用GPU的并行计算资源,通常采用分块(tiling)技术。但在VQ场景下,由于码本切换轴与任务规约(reduction)轴之间的冲突,朴素的并行化会引入过多的流量。为此,我们设计了以码本为中心的数据流,通过沿着码本切换轴(即每四个通道,对应一个码本)对任务进行划分和并行化,确保每个线程块只需加载一个码本,从而消除了重复的码本加载或切换。同时,对于那些原本进行时间累加但现在被并行化的轴,我们执行显式的全局规约以确保结果准确。
遵循CQ配置的注意力(解码)计算的以码本为中心的数据流示例。

计算引擎:以码本为中心的数据流实现与自适应。在实现中,我们首先识别发生规约的轴和需要切换码本的轴,然后沿着码本切换轴分割并并行化计算,最后对相交的轴进行显式全局规约。为了平衡全局规约的开销,我们利用一个分割因子(split factor)来控制并行化程度。较大的分割因子会减少重复码本流量,但需要更多的全局规约。我们通过令$Traffic_{Reduce} = Split~Factor \times Output~Size$与$Traffic_{Codebook} = Original~Codebook~Traffic / Split~Factor$相等(基于中值定理),自适应地确定最佳分割因子以最小化总开销。

计算引擎:以码本为中心的分层融合设计直觉。基线方法采用共享内存级别的融合,导致共享内存与寄存器之间产生过多流量。我们利用现代GPU支持寄存器级别数据交换的特性(即shfl_xor API),通过在寄存器中直接交换数据来绕过共享内存。例如,当解量化数据的布局为8,而mma计算指令所需布局为2时,我们通过特定的线程映射将数据交换限制在一个包含4个线程的mini-warp内。在这个mini-warp内执行三次shuffle操作后,每个线程寄存器持有的数据就能精确对齐计算指令的要求。
基于shuffle API的warp内数据交换示例,每个线程一次解量化8个元素,而后续计算要求一个线程只持有2个元素(mma指令)。

计算引擎:线程映射与分层融合实现。由于朴素的顺序线程映射会导致复杂的交换路径并占用额外寄存器,我们在离线阶段预先确定线程映射。如Alg. 1所示,首先根据解量化和计算找到每个元素的关联,然后为每个线程识别需要其解量化数据的所有线程,并将这些线程分组为一个mini-warp。最后按mini-warp重新映射所有线程。在自适应层面,由于共享内存访问延迟几乎是寄存器访问加shuffle组合延迟的五倍,因此对于需要少于五次shuffle操作的量化张量,我们实施寄存器级别融合;对于其他张量,保留传统的共享内存级别融合。

# Algorithm 1 Intra-warp data exchange based on shuffling
1: function THREAD_MAPPING(data, layout_dequant, compute)
2:     for item in data do
3:         item.tid_compute, dequant <- GetTid(item, layout_compute, dequant)
4:     mini_warps <- []
5:     for dequant_thread in warp do
6:         mw <- [data.tid_compute for data.tid_dequant=dequant_thread]
7:         if mw not in mini_warps then
8:             mini_warps[mw] <- []
9:         mini_warps[mw].append(dequant_thread)
10:    for mw in mini_warps do
11:        mini_warps[mw][i] mw[i] // Thread mapping we need

12: function REG_FUSION(data, iter)
13:    for off in [1, iters) do // intersected 0 no shuffle needed
14:        data[tid ^ off] <- shfl_xor(data[tid ^ off], off)
15:    return reg

计算引擎:整体工作流。我们的计算引擎采用基于模板的设计(Alg. 2)来生成最终的融合内核。首先在离线阶段,根据VQ配置和目标计算,确定共享内存/寄存器预算、分割因子、所需的shuffle次数以及相应的线程映射。随后,通过Parallel_For函数启动以码本为中心的数据流计算。在每个并行任务内,先将码本加载到缓存中,接着使用映射后的线程进行解量化。解量化后,使用Reg_FusionShared_Fusion执行分层融合使数据达到正确布局,随后进行计算。最后,如果有必要,执行全局规约。

# Algorithm 2 Complete VQ-aware computation template
1: function KERNEL_TEMPLATE
2:     All, Reduce <- compute_op.all_axes, reduce_axes
3:     layout_src, dst <- codebook.vector_size, compute_op.required_size
4:     Budget <- Free shared and reg to preserve occupancy
5:     factor <- Value to make TrafficReduce = TrafficCodebook
6:     n_shuffle <- layout_src / layout_dst
7:     if n_shuffle <= thres_shuffle(= 5) then
8:         Thread_Mapping(compute_op.warp_tile, layout_src, dst)
9:     Parallel_For(codebook.switch_axes, factor)
10:        if required by algorithm then
11:            CB <- Switch(New codebook ptr)
12:        CB_cached, boundry <- Load(CB, Budget)
13:        for id in quantized_data do
14:            data <- Access(CB_cached, boundry, CB, id)
15:        if n_shuffle <= thres_shuffle then
16:            data <- Reg_Fusion(data, n_shuffle)
17:        Else
18:            data <- Shared_Fusion(data, layout_src, dst)
19:        for temporal_iteration on All - codebook.switch_axes do
20:            partial <- compute_op(data, temporal_iteration)
21:        output <- Reduce(partial, Reduce ∩ codebook.switch_axes)
22:        Return output

实验环境

  • 数据集/任务:使用arc-challenge任务评估端到端模型的准确率,测试条件为Batch Size 16,序列长度1024,生成256个tokens。
  • 模型架构:Llama-7B 和 Llama-65B。
  • 硬件配置:单张 NVIDIA RTX 4090 24GB GPU;在端到端评估中加入了一张 Tesla A40 GPU 以探索带宽受限环境下的性能。
  • 软件/代码库配置:基于 CUDA 模板实现,依赖 cutlass 和 flash-attn 库,集成了 LMEval 框架用于精度评估,对比基线包括开源的 qServe 框架。评估的内核包括 GeMM, GeMV 和 FlashDecoding。VQ配置包括 QuiP#-4, AQLM-3, GPTVQ-2(权重压缩)和 CQ-2, CQ-4(KV缓存压缩)。

实验结果

整体加速比。在各类VQ配置下,VQ-LLM相比未优化的全局内存版本(GC)平均降低了46.13%的延迟(最高达53.73%),对应$1.9\times$到$2.2\times$的加速比。实验发现,计算密集型内核(GeMM)从优化中获得的加速比相对高于内存密集型内核(GeMV)。对于注意力内核,随着批处理大小(Batch Size)的增加,加速比也随之增加。Llama-65B展现出与Llama-7B几乎一致的加速比,证明了优化的良好可扩展性。
对于各种VQ配置,性能最佳版本相对于未优化版本的整体延迟降低。对于注意力(解码),1k和4k分别表示序列长度为1024和4096。

加速比分解分析(GeMM和GeMV)。对于码本较小的QuiP#-4,SC(共享内存缓存)和O1(中等频率条目缓存)效果相同。AQLM-3和GPTVQ-2因码本较大,O1优化提升显著。O2(寄存器缓存热条目)对AQLM-3提升最大,因其有大量高频条目。O3(以码本为中心的数据流)对GeMM产生负面影响(因输出尺寸大导致规约开销大),但对GeMV有利。O4(分层融合)显著提升了GeMM性能(因为mma指令只需极少的shuffle),但由于需要大量的shuffle操作,导致GeMV在QuiP#-4和AQLM-3上性能下降。
GeMM(上)和GeMV(下)的优化分解。

加速比分解分析(注意力解码)。对于CQ-2,由于码本较大,SC会显著降低性能,因此必须使用O1。O2提升较小,因为极高频条目较少。O3通过消除重复流量显著提升了性能。O4提供了轻微的改善。CQ-4相对于CQ-2也实现了类似的加速效果。
(左)注意力(解码)中CQ-2的优化分解。(右)CQ-4相对于CQ-2的延迟。

与FP16及逐元素量化的比较。在4-bit等效宽度下,VQ-LLM实现了与SOTA逐元素量化方法(AWQ和QoQ/qServe)相当(注意力解码为$1.01\times$)甚至更低(GeMV/GeMM为$0.88\times / 0.96\times$)的延迟。同时,开源的QuiP#和AQLM实现表现出$2.83\times$到$114.4\times$的延迟,无法投入实际应用,这证明了本文工作成功将算法改进转化为了实际应用。
与逐元素量化工作的延迟对比。

端到端评估结果。在等效4-bit设置下,VQ-LLM实现了与qServe相当的端到端推理加速(相较于FP16基线提升约$2.2\times$)。在arc-challenge任务上,VQ-LLM的准确率超过qServe约2.5%。在带宽受限的Tesla A40 GPU上,VQ-LLM展现出比RTX 4090更大的加速比,证明其在带宽受限环境下的有效性。内存使用方面,FP16基线消耗超过22GB,而VQ-LLM-4仅使用不到6GB。
(左)相对于FP16的整体加速比,以及(右)SOTA逐元素量化(qServe)与VQ-LLM在arc-challenge上的准确率。

补充细节

不同类型的注意力机制。除了Flash Decoding基线外,VQ-LLM在对比Flash Attention、Paged Flash Attention和Paged Flash Decoding时均超越了这些基线。在Batch Size为8、序列长度为4096的条件下,相比最佳的FP16基线,VQ-LLM实现了66.4%的延迟降低和75%的内存占用减少。
各种注意力基线相对于我们性能最佳的CQ-4实现的相对延迟。

量化开销。权重压缩不引入运行时量化开销。对于KV缓存压缩,在解码阶段为新词元即时量化新键值的运行时开销可忽略不计($<1 \mu s$)。在预填充阶段,量化所有提示词元的键值相比线性投影引入的开销不到10%,且后续计算不会立即需要量化后的KV缓存,因此这些开销可以忽略。

结论

本文提出了VQ-LLM,这是一个针对向量量化优化的代码生成框架,包含码本缓存和基于码本的计算引擎。通过该框架,相较于未优化版本平均降低了46.13%的延迟,相较于开源实现最高降低了99%的延迟。码本缓存提出了分层放置策略以保持硬件利用率并减少bank冲突;计算引擎提出了以码本为中心的数据流和融合方案,以减少过多的片外和片上流量。所有优化均通过启发式方法自适应配置。最终证明了VQ-LLM相较于未优化实现和逐元素量化方法在实际应用中的有效性和可行性。未来的工作将探索结合NCCL等通信库的多GPU分布式推理场景。

参考文献汇总

  • [10] 描述了FlashDecoding的数据流,并行化不同token的计算并在全局内存中计算局部softmax:Dao et al., "Flash-decoding for long-context inference", 2023.
  • [69] 描述了CQ-2及CQ-4算法,其中每四个通道需要切换不同的码本:Zhang et al., "KV cache is 1 bit per channel: Efficient large language model inference with coupled quantization", 2024.
  • [56] 描述了QuiP#算法,利用基于格的码本,使用整个权重张量训练一个码本避免重复流量,但向量大小为8:Tseng et al., "Quip#: Even better LLM quantization with hadamard incoherence and lattice codebooks", 2024.
  • [12] 描述了AQLM算法,具有未对齐的12-bit存储格式:Egiazarian et al., "Extreme compression of large language models via additive quantization", 2024.
  • [57] 描述了GPTVQ-2算法,权重矩阵每(256, 256)的块共享一个码本:van Baalen et al., "GPTVQ: the blessing of dimensionality for LLM quantization", 2024.
  • [45] 描述了现代GPU支持的warp内数据交换API (shfl_xor):NVIDIA, "Nvidia cuda warp shuffle functions", 2024.
  • [30] 描述了SOTA级别的权重逐元素量化方法AWQ:Lin et al., "AWQ: activation-aware weight quantization for on-device LLM compression and acceleration", MLSys 2024.
  • [31] 描述了SOTA级别的KV缓存逐元素量化方法QoQ及qServe框架:Lin et al., "Qserve: W4A8KV4 quantization and system co-design for efficient LLM serving", 2024.