HACK: Homomorphic Acceleration via Compression of the Key-Value Cache for Disaggregated LLM Inference
HACK: Homomorphic Acceleration via Compression of the Key-Value Cache for Disaggregated LLM Inference
发表时间: 2025-08 · arXiv:2502.03589 (SIGCOMM 2025)
原文: https://arxiv.org/abs/2502.03589
Zeyu Zhang (University of Virginia), Haiying Shen (University of Virginia), Shay Vargaftik (VMware Research), Ran Ben Basat (University College London), Michael Mitzenmacher (Harvard University), Minlan Yu (Harvard University)
速读
一句话结论
HACK提出了一种用于分离式大语言模型推理的KV Cache同态量化加速方法，通过直接在2-bit量化数据上进行注意力矩阵乘法，彻底消除了反量化开销，将端到端作业完成时间（JCT）最多降低了70.9%。
要解决什么问题
分离式LLM推理（Disaggregated LLM Inference）将计算密集的Prefill阶段和显存密集的Decode阶段部署在不同的GPU上，以避免资源争抢并提高整体利用率。但这引入了一个致命卡点：Prefill节点必须通过网络将生成的KV Cache传输给Decode节点。由于云上高性价比的Prefill实例（如A10G、T4）通常只配备10至50 Gbps的低速网络，传输庞大的KV数据会造成严重的通信瓶颈，其耗时最多可占总作业完成时间（JCT）的42.2%。为了缓解传输压力，现有的KV量化基线方法（如CacheGen和KVQuant）会在发送前将数据压缩至2-bit，从而大幅降低通信耗时。然而，这些方法在Decode阶段遇到了新的机制卡点：在每一次自回归解码迭代中，模型都必须先将所有历史Token的KV数据从2-bit反量化回FP16，然后才能进行注意力矩阵乘法。随着上下文长度的增加，这种逐迭代的全局反量化操作带来了极其沉重的计算开销，最多可占JCT的37.9%。此外，由于最终计算仍在使用FP16，这些方法也无法加速实际的矩阵乘法过程。
怎么做的
HACK的核心思路是引入同态量化（Homomorphic Quantization），让注意力机制中的矩阵乘法直接在量化后的整数上执行，从而完全绕开反量化步骤，并利用GPU的INT8算力加速计算。具体而言，HACK采用非对称的2-bit随机量化，将矩阵按固定大小（如64）进行分块。对于注意力计算中的矩阵乘法 $C = AB$，假设 $A$ 和 $B$ 的元素分别被量化为 $a'_{iz}$ 和 $b'_{zj}$，且带有各自的缩放系数 $s$ 和最小值 $m$，HACK将原始的FP16乘法改写为如下近似更新式：
在这个公式中，核心项 $\sum_z a'_{iz}b'_{zj}$ 完全由量化后的整数构成，可以直接调用高效的INT8 Tensor Core进行加速。其余三项则是为了修正量化误差而引入的近似补偿项。为了让这个公式在Decode阶段真正高效，HACK设计了两个关键部件来消除边缘开销。第一是求和消除（Summation Elimination）。公式中的 $\sum_z b'_{zj}$ 涉及对历史KV Cache的遍历求和，如果每次迭代都重算，其计算量将达到 $2 d_h L_{KV}$，当序列长度大于30时，这甚至比直接反量化还要慢。HACK的解法是在生成KV时，提前计算好该块的整数和，并用极少量的额外显存（如INT16格式）将其缓存下来。每次迭代直接复用该求和结果，将近似计算的代价大幅降至 $10(d_h + L_{KV})$。第二是重量化消除（Requantization Elimination）。在Decode阶段，新生成的Value向量会不断追加到序列末尾。如果最后一个分块未满，新加入的值可能会打破该块原有的最大最小值范围，导致整个块需要重新量化。HACK通过维护一个极小的FP16缓冲区，专门存放这个未满的Value尾块，并对其使用标准的FP16乘法。只有当该块凑满64个Token时，才会触发一次性量化并移入主Cache，从而彻底避免了反复重量化的开销与精度折损。
效果如何
作者在vLLM框架上集成了修改后的FlashAttention-2内核来搭建实验。硬件采用A100作为Decode节点，搭配A10G、V100、T4或L4作为Prefill节点。评测模型覆盖了Mistral-v0.3 7B、Phi-3 14B、Yi 34B、Llama-3.1 70B以及Falcon 180B。对比基线包括：代表传统分离式推理路线的系统基线（未压缩的vLLM分离部署）、代表数据分布压缩路线的CacheGen，以及代表低比特量化路线的KVQuant。在Llama-3.1 70B模型和长文本检索数据集Cocktail（输入长度最高达28.8K）的设置下，HACK展现了最显著的优势。相比于必须进行反量化的CacheGen和KVQuant，HACK将JCT分别降低了41.5%和45.1%；而相比于完全不压缩的系统基线，HACK更是将JCT大幅降低了61.6%。在网络带宽极低（10 Gbps）的V100 Prefill节点上，HACK对基线的JCT降幅最高达到了70.9%。在代价方面，HACK的2-bit量化会带来轻微的精度损失，在各项任务中的准确率下降幅度在0.76%至1.56%之间，但整体精度仍优于CacheGen和KVQuant。作者也承认了当前工程实现上的一个局限：由于Triton编译器目前最低只支持INT8计算，HACK在底层执行时需要先在GPU本地显存中将2-bit数据转换为INT8再做乘法，未能完全释放INT4硬件的理论极限性能。
主要贡献
分离式大语言模型（LLM）推理通过将计算密集型的预填充（Prefill）阶段和内存密集型的解码（Decode）阶段分离，避免了预填充-解码干扰并提高了资源利用率，因而广受欢迎。然而，在这两个阶段之间传输键值（KV）数据可能会成为网络瓶颈，特别是对于长提示词而言。现有的KV量化方法虽然可以缓解传输瓶颈并降低内存需求，但它们引入了显著的反量化开销，加剧了计算时间的消耗。此外，低精度浮点格式（如FP4/6/8）缺乏高压缩率，且需要特定硬件支持。
为了解决这一问题，本文提出了用于分离式LLM推理的KV缓存同态加速压缩系统（HACK）。HACK消除了繁重的KV反量化步骤，直接在量化的KV数据上执行计算，以近似并降低昂贵的矩阵乘法步骤的成本。本文的主要贡献如下：
1. 提出了一种用于矩阵乘法的同态量化方法。它在不需要对量化矩阵进行反量化的情况下，对量化矩阵进行与KV相关的矩阵乘法，然后应用近似方法将量化输出转换为真实输出的近似值。该方法在避免昂贵的KV反量化开销的同时，降低了KV传输延迟、计算时间、内存访问延迟和内存需求。
2. 通过使用少量内存存储数据来消除冗余计算和更新量化数据的需求，进一步降低了同态量化的开销。
3. 将HACK集成到FlashAttention-2中，并在vLLM上构建了该系统。广泛的实验表明，与分离式LLM推理基线相比，HACK将作业完成时间（JCT）降低了高达$70.9\%$，与最先进的KV量化方法相比降低了高达$52.3\%$。该代码已开源。
表1：不同方法与我们的系统HACK的特征对比。
| 减少通信 | 减少计算 | 减少内存访问 | 避免反量化 | 实现高压缩率 | |
|---|---|---|---|---|---|
| Baseline | × | × | × | √ | × |
| KV quantization | √ | × | √ | × | √ |
| FP4/6/8 | √ | 需要硬件支持 | √ | 需要硬件支持 | × |
| HACK | √ | √ | √ | √ | √ |
背景知识与设计动机
分离式LLM推理的基础与瓶颈
LLM推理包含预填充和解码两个阶段。在注意力机制中，输入token被转换为查询矩阵$Q^h$、键矩阵$K^h$和值矩阵$V^h$。自注意力计算输出为$O^h = softmax(\frac{Q^h (K^h)^T}{\sqrt{d_h}}) V^h = P^h V^h$。分离式架构将预填充和解码分配给不同的GPU实例，预填充阶段输出第一个token及KV数据并传输给解码实例。然而，由于云服务商提供的廉价GPU实例通常缺乏高速网络，KV传输成为显著瓶颈。测试表明，在10-50 Gbps网络的实例上，通信时间占比高达$19.1\%-23.5\%$。对于长序列数据集（如arXiv和Cocktail），KV通信时间比短序列数据集高出$15.5-43.1$倍，计算时间高出$9.8-19.2$倍。此外，解码阶段的GPU内存使用率高达$93.7\%$，KV数据的内存访问延迟占JCT的高达$33.1\%$。流水线通信技术在通信时间远超预填充时间，或解码实例显存不足导致KV需暂存CPU时，会失去效用。
现有KV量化方法的反作用
CacheGen和KVQuant等量化方法通过混合精度和2位量化实现了约$86\%$的KV压缩率。虽然它们能将平均通信时间占比降低高达$34.1\%$，但必须在每次解码迭代中对检索到的所有token的KV值进行反量化。测试表明，这些方法引入的反量化时间占比高达$26.4\%-37.9\%$，且长序列数据集的反量化时间是短序列的$12.4-24.9$倍。它们虽然减少了内存访问时间，但完全无法减少计算时间。
方法细节
HACK系统工作流
为了在量化KV值的同时消除反量化时间开销，HACK对与KV相关的矩阵乘法提供同态量化。系统采用2位量化来处理KV数据以保证$86\%$的压缩率。首先，预填充实例从提示词token生成$Q, K, V$，并将它们量化为$Q', K', V'$。因为$Q$在计算后即被丢弃，无需极致压缩，故采用8位量化以提高精度。接着，$Q'$和$K'$之间的首次矩阵乘法使用同态量化执行，输出注意力分数$S$，此过程无需反量化$K'$，并由GPU的INT8计算能力加速。然后，$S$通过softmax转换为注意力概率$P$，并使用INT8量化为$P'$。随后，$P'$和$V'$再次使用同态量化进行乘法加速，最终输出第一个token。如果解码实例显存不足，预填充实例会将量化的KV数据交换到CPU内存中。当解码实例可用时，预填充实例将第一个token、$K', V'$以及量化元数据（最小值$m$和缩放值$s$）传输给解码实例，并存入KV缓存。在解码阶段，解码实例从第一个token生成$Q, K, V$并量化。新token的$K'$和$V'$在token维度上与所有先前token的$K'$和$V'$合并。最后，解码阶段的同态量化过程与预填充阶段相同，生成下一个token并进入下一次迭代。
用于矩阵乘法的非对称同态量化设计
为了利用GPU的INT4或INT8计算能力，HACK提出了一种同态量化方法，对于矩阵乘法$C = AB$，先将矩阵量化为$A'$和$B'$，执行量化乘法$C' = A'B'$，然后以最小开销将$C'$近似还原为$C$。系统采用非对称2位随机量化【索引1，Stochastic Quantization+1983】来减少量化误差。矩阵的行或列被划分为大小为$\Pi$的多个分区。在每个分区$i$中，提取最小值$min_i$和最大值$max_i$，计算缩放比例$scale = \frac{max_i - min_i}{2^2 - 1}$。原始值$x$通过随机舍入函数量化为整数$x' = round(\frac{x - min_i}{scale})$。
同态量化矩阵乘法的数学近似推导
设$a_{iz}$和$b_{zj}$为原矩阵元素，由于$a_{iz} \approx s_{a_i} a'_{iz} + m_{a_i}$且$b_{zj} \approx s_{b_j} b'_{zj} + m_{b_j}$，矩阵乘法$(AB)_{ij} = \sum_z a_{iz} b_{zj}$可以展开近似为：
其中，核心项$\sum_z a'_{iz} b'_{zj}$是可以通过INT8计算加速的量化矩阵乘法，其余项用于将量化结果近似还原为真实结果。该近似计算的额外成本为$9MN + MZ + NZ$。
自注意力机制中的分块与同态量化应用
在自注意力计算中，$Q$和$K$的内部维度是大小为$d_h$的注意力头维度。HACK将该维度划分为多个分区，分区大小$\Pi$必须是16的倍数以适配GPU硬件。$Q$和$K$经过同态量化输出注意力分数$S$，处理后得到概率$P$。$P$和$V$的内部维度是序列维度$L_{KV}$，同样根据$\Pi$进行分区。在解码阶段，新生成的token的$K'$和$V'$会沿着序列维度$L_{KV}$追加到先前的$K'$和$V'$中，并重复上述同态量化过程。
求和消除优化（Summation elimination）
在每次解码迭代中，近似公式中的求和项会带来额外开销。具体而言，$Q$和$K$的近似开销为$10(d_h + L_{KV}) + 2d_h L_{KV}$。如果不使用同态量化而直接反量化，成本为$4d_h L_{KV}$。为了消除$2d_h L_{KV}$的求和计算成本（即计算$\sum_z b'_{zj}$），HACK在解码期间将$K$和$V$的整数和存储起来，并在每次迭代中复用以避免重新计算。对于采用$b$位量化、大小为$\Pi$的分区，整数和最多需要$b + \lceil \log_2 \Pi \rceil$位来存储（例如2位量化、$\Pi=64$时仅需8位）。这仅占用极少的额外显存（约$2.7\%$）。优化后，每次解码迭代的近似成本降至$10(d_h + L_{KV})$。当序列长度$L_{KV} > 30$时，同态量化的近似成本将比传统KV反量化开销低一个数量级。
V矩阵最后分块的重量化消除（Requantization elimination）
每次解码迭代后，新token的KV值被追加。由于$K$的分区是沿固定的头维度排列的，新token的$K$自身形成独立分区，不会改变先前分区的$[min, max]$。然而，$V$的分区是沿序列维度排列的。如果$V$的最后一个分块中的token数量小于分区大小$\Pi$，新token的$V$元素会被添加到现有分区中。这可能导致新元素超出先前的$[min_j, max_j]$范围，从而迫使系统更新范围，并对该列上的所有其他token值进行重新量化（Requantization）。这不仅增加量化误差，还引入了额外开销。为了解决这个问题，当$V$最后一组的token数未达到$\Pi$时，HACK不存储其量化值，而是在一个独立的缓存中保留其原始的FP16值。此时，$P$的最后分块与$V$的最后分块的矩阵乘法直接以FP16格式进行。当token数量达到$\Pi$时，才对其进行量化并移至量化KV缓存。由于FP16乘法仅限于最后一个分块，其计算时间不会随序列长度扩展而增加。
补充细节
系统实现细节
HACK集成了内存高效的注意力后端FlashAttention-2，通过OpenAI Triton实现，并构建在vLLM之上。由于Triton目前支持的最低计算精度为INT8，HACK首先在本地GPU内存（而非全局内存）中将量化数据从2位转换为INT8，然后再执行矩阵乘法。系统实现了两个Triton内核：attn_prefill将QKV生成、量化和带有同态量化的自注意力融合到一个内核中；attn_decode除了上述步骤外，还将新token的量化KV数据与先前数据的拼接过程融入内核，并将$V$的最后一个分块分离到缓冲区以支持FP16计算。在数据管理方面，修改了vLLM的KV缓存结构，以FP16存储$m$和$s$。对于$\Pi=128$的2位量化，求和值需要9位，这会导致内存对齐问题。因此，HACK采用INT16来存储这种情况下的求和值，其占用的内存仅为量化KV数据的约$5\%$。跨实例通信通过NCCL实现。
实验环境
- 硬件配置：预填充实例采用AWS的g5.12xlarge (4xA10G)、p3.8xlarge (4xV100)、g4dn.12xlarge (4xT4)、g6.12xlarge (4xL4) 或 p4de.24xlarge (8xA100)。解码实例默认采用两台 p4de.24xlarge。网络带宽从10Gbps到400Gbps不等。
- 模型架构：Mistral-v0.3 7B (M)、Phi-3 14B (P)、Yi 34B (Y)、Llama-3.1 70B (L) 和 Falcon 180B (F)。配置了不同的张量并行 (TP) 和流水线并行 (PP) 策略。
- 数据集：IMDb分类（短序列，平均输入315，输出37）、arXiv摘要（长序列，最大输入24329）、Cocktail（长序列，最大输入28.8K）和HumanEval（代码生成，平均输入204，输出139）。
- 软件配置：基于vLLM构建，集成FlashAttention-2和Triton，修改了DistServe和SplitWise以支持以太网数据传输。
实验结果
端到端作业完成时间（JCT）性能
在使用Llama-3.1 70B和A10预填充实例的测试中，HACK在IMDb、HumanEval、arXiv和Cocktail数据集上，相比基线分别将平均JCT降低了$38.6\%$、$40.1\%$、$55.3\%$和$61.6\%$。相比CacheGen和KVQuant，HACK将JCT降低了$19.2\%-45.1\%$。HACK不仅加速了预填充和解码计算，还通过消除反量化开销（其仅占JCT的$1.53\%-3.18\%$的近似开销，而反量化占$17.2\%-30.4\%$）获得了显著提升。对于长序列数据集（arXiv和Cocktail），解码时间改善尤为明显（$32.1\%-33.7\%$）。
在峰值GPU内存使用方面，HACK相比基线降低了$25.0\%-33.6\%$（长序列），仅比其他量化方法高出$0.6\%-2.9\%$（由于存储求和值和FP16数据）。
跨模型与跨硬件性能
在A10G实例上测试不同模型处理Cocktail数据集时，HACK相比基线将平均JCT降低了$53.3\%-61.6\%$。在不同预填充实例（A10G, V100, T4, L4, A100）上，HACK相比基线将JCT降低了$59.3\%-70.9\%$。其中，在带宽最低的V100上，HACK对基线的提升最大（$70.9\%$），证明了其在低带宽环境下的有效性。
准确率性能
测试了不同分区大小（$\Pi=128, 64, 32$）。相比基线，HACK($\Pi=64$)的准确率损失仅为$0.76\%-1.56\%$，由于采用了细粒度的分区，其准确率优于CacheGen（损失$1.44\%-2.08\%$）和KVQuant（损失$1.46\%-2.33\%$）。因此，$\Pi=64$被选为默认配置。
消融实验与敏感性测试
移除求和消除（HACK/SE）导致长序列数据集的平均JCT增加了$22.1\%-25.9\%$，证明了存储求和值的必要性（仅占$2.2\%-2.7\%$显存）。移除最后分块重量化消除（HACK/RQE）导致短序列数据集JCT增加$17.8\%-21.7\%$，并引发额外的准确率下降（$0.08\%-0.29\%$）。敏感性测试表明，$\Pi=32$虽能提高准确率，但会使JCT增加高达$28\%$，$\Pi=64$是最佳折中。
可扩展性测试
当预填充模型副本数与解码模型副本数的比例（$p$）从1增加到8时，基线的平均JCT增加了$127\%$，而HACK仅增加了$31-43\%$，证明HACK在大规模部署中能有效缓解网络瓶颈。
结论
HACK是一种用于分离式LLM推理的新型量化方法。它在不引入昂贵KV反量化开销的情况下，有效降低了KV传输开销、计算时间和KV的内存访问延迟。通过集成到FlashAttention-2和vLLM中并进行广泛实验，证明HACK能够将JCT降低高达$70.9\%$。未来的工作计划探索能进一步改善准确率与JCT折中的新型量化方案，并直接在CUDA中实现HACK以支持INT4计算并减少Triton带来的执行开销。此外，也将探索如何将KV驱逐（KV eviction）技术与同态量化相结合。
参考文献索引汇总
- 【索引1，Stochastic Quantization + 1983 + Recent Developments in High-Energy Physics】在“用于矩阵乘法的非对称同态量化设计”段落中引用。原文描述为采用非对称2位随机量化（stochastic quantization）来减少量化误差，通过对矩阵元素进行分区并应用随机舍入机制来实现。
💬 评论讨论
欢迎在这里分享您的想法和见解！