Oaken: Fast and Efficient LLM Serving with Online-Offline Hybrid KV Cache Quantization

发表时间: 2025-03 · arXiv:2503.18599

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

作者/机构: Minsu Kim (KAIST), Seongmin Hong (HyperAccel), RyeoWook Ko (KAIST), Soongyu Choi (KAIST), Hunjong Lee (HyperAccel), Junsoo Kim (HyperAccel), Joo-Young Kim (HyperAccel), Jongse Park (KAIST)

---

速读

一句话结论 本文提出了一种软硬协同的在线-离线混合 KV 缓存量化加速方案 Oaken，通过定制量化算法与专用硬件模块，在仅损失极小精度的前提下，将大语言模型推理的吞吐量提升了 1.58 倍。

要解决什么问题 大语言模型（LLM）推理服务通常通过增大批处理大小来提升吞吐量，但这在自回归生成阶段会遇到严重的内存墙瓶颈。具体而言，预填充阶段可以高度并行，但生成阶段是逐个词元串行生成的。此时，注意力机制中的键值（KV）缓存无法在不同请求间共享，随着批处理大小和上下文长度的增加，KV 缓存的体积会急剧膨胀。这不仅极大地消耗了内存容量，迫使系统横向扩展，而且其庞大的读取需求也迅速耗尽了内存带宽，导致核心利用率低下。为了缓解这一卡点，现有方法通常采用 KV 缓存量化技术（例如对大部分正常值用 INT4，对异常大值用 FP16），但这种做法面临两个致命缺陷：一是在线检测异常值需要进行昂贵的排序操作，严重拖慢了推理速度；二是为异常值保留高精度会引入稀疏矩阵存储开销，每个异常值往往需要额外的索引位（高达 23 位），这不仅增加了硬件设计的复杂性，还抵消了量化带来的容量和带宽收益。

怎么做的 Oaken 的核心思路是摒弃高开销的在线异常值检测与混合精度存储，转而采用一种“离线确定阈值、在线移位量化、软硬协同编码”的机制，将所有 KV 缓存都压缩到极低位宽。该方案由三个关键算法部件和配套硬件构成。首先是基于阈值的在线-离线混合量化。作者观察到 KV 缓存的数值分布在不同输入下保持一致，因此在离线阶段通过少量样本跑出四个固定的异常值阈值（$S_{low}$、$S_{high}$、$T_{low}$、$T_{high}$）。在线推理时，直接用这四个阈值将 KV 值划分为内部组（极小异常值）、中间组（正常值）和外部组（极大异常值），彻底避开了在线排序的开销。其次是组移位量化。为了让跨度极大的异常值也能用低位宽存储，Oaken 直接将异常值减去阈值，将其平移到较小的数值范围内再做均匀量化。其核心量化函数定义为： $$Q_{Oaken}(x) = \begin{cases} Q(x - S_{high}), & \text{if } x > S_{high} \\ Q(x), & \text{if } S_{low} \le x \le S_{high} \\ Q(x - S_{low}), & \text{if } x < S_{low} \end{cases}$$ 其中 $Q(x)$ 为常规的均匀量化函数。通过这种移位，占 90% 的中间组被量化为 4 位，而内外部的异常值被量化为 5 位。第三是融合的密集-稀疏编码。中间组作为密集矩阵存储，而异常值原本所在的位置会被置零。Oaken 巧妙地将 5 位异常值中的 4 位塞进密集矩阵的这些置零空位中，剩下的 1 个符号位连同 6 位位置索引和 1 位组标志，凑成标准的 8 位稀疏格式存储。这一设计将异常值的存储开销从常规的 23 位暴降至 8 位，且完美对齐内存。最后，Oaken 在硬件加速器的直接内存访问单元中定制了量化与反量化引擎，并设计了专门的内存管理单元，以页为粒度统一管理密集和稀疏数据，从而在物理层面打满内存突发读取带宽。

效果如何 实验在 Llama2、OPT、Mistral 和 Mixtral 等多种规模（6.7B 到 70B）的模型上展开，评估数据集包括 Wikitext2 以及 PIQA 等零样本任务。硬件基座为配备 80GB 高带宽内存的 NVIDIA A100 GPU 以及基于 LPU 扩展的 Oaken 模拟器（分为配备高带宽内存和配备大容量 LPDDR 内存两个版本）。对比基线包括代表全精度路线的 vLLM，以及代表不同 KV 缓存量化路线的 KIVI 和 KVQuant（逐向量混合精度）、QServe（通道重排与缩放）和 Tender（间接索引与通道分组）。在批处理大小为 256、输入输出上下文各 1K 的设置下，配备 LPDDR 的 Oaken 吞吐量达到了 vLLM 的 1.79 倍，并且比最先进的量化基线 QServe 高出 1.58 倍。在准确率方面，Oaken 相比全精度基线的平均精度损失仅为 0.87%，在 Llama2 模型上比 QServe 的准确率高出 1.38%。在代价方面，Oaken 的定制硬件模块仅带来 8.21% 的芯片面积开销。作者也指出了该方案的局限与失效场景：在输出极短的真实对话轨迹中，由于计算密集型操作占主导，Oaken 的优势并不明显；此外，尽管配备高带宽内存的 Oaken 在小模型上速度极快，但受限于物理容量，在处理超过 16K 长上下文时依然会因内存溢出而失效，此时必须依赖大容量的 LPDDR 版本才能支撑高达 32K 的序列推理。

A1 主要贡献

大型语言模型（LLM）服务系统通过批处理多个请求以实现高吞吐量，但注意力操作的批处理极具挑战性，使得内存带宽成为关键瓶颈。当前，业界依赖配备多个高带宽内存（HBM）通道的高端GPU来缓解此问题。然而，HBM的高带宽通常以牺牲有限的内存容量为代价，迫使系统进行横向扩展，从而降低了核心利用率并增加了成本。此外，支持更长上下文的最新进展大幅增加了键值（KV）缓存的大小，进一步加剧了内存容量的压力。

为缓解这一压力，现有研究探索了KV缓存量化技术，通常对大多数值使用低位宽（如INT4），并为异常值选择性地使用更高位宽（如FP16）。尽管这种方法有助于同时实现高精度和低位宽，但其在线异常值检测的成本过高，抵消了量化的优势。

本文基于上述洞察，提出了Oaken，一个通过算法和硬件协同设计同时实现高精度和高性能的加速解决方案。Oaken的核心贡献如下：

1. 在线-离线混合KV缓存量化算法： 为了在KV缓存量化的精度-性能权衡空间中有效找到最佳点，Oaken采用了一种在线-离线混合方法。

• 离线确定阈值： 在离线阶段，Oaken确定与数据无关的异常值阈值。
• 在线设置量化尺度： 在在线推理阶段，利用离线确定的阈值来设置量化尺度。
• 量化损失缓解技术： Oaken引入了一种量化损失缓解技术，该技术通过将值向更小的范围移动，在量化前将异常值转换为正常值。
• 融合的密集-稀疏存储： Oaken使用密集张量存储正常值，并将稀疏的异常值融合到密集张量中，以优化存储。

2. 面向量化感知的LLM加速器硬件模块： 为了将算法优势转化为实际性能增益，Oaken设计了定制的量化/反量化引擎和内存管理单元，这些模块可以集成到任何现有的LLM加速器中。

• 硬件模块集成： 这些模块被放置在现代LLM加速器中常见的DMA单元内。
• 高效内存管理单元（MMU）： MMU的设计面临的挑战是在有效布局密集和稀疏矩阵的同时，实现接近物理极限的最大带宽。Oaken的MMU为密集和稀疏数据设计了两个管理表，以处理虚拟到物理地址的映射，并以页为粒度管理单一地址空间。这种设计最大化了内存带宽利用率，同时避免了碎片化和突发顺序问题。

通过上述算法与硬件的协同设计，Oaken在有效带宽和容量方面达到了前所未有的水平，从而获得了比其他方案更高的吞吐量，如图1所示。

实验结果表明，与使用最先进KV缓存量化技术的NVIDIA A100 GPU相比，在批处理大小为256时，Oaken实现了高达1.58倍的吞-吐量提升，而平均精度损失仅为0.54%。同时，Oaken模块仅带来8.21%的面积开销，证明了其在构建快速且经济高效的LLM推理服务系统方面的重要价值。

---

A3 背景知识与关键观察

2. 背景知识

2.1 LLM推理入门

LLM推理的两个阶段。如图2(a)所示，大型语言模型的推理主要分为预填充（prefill）阶段和生成（generation）阶段【【4, PaLM 2 Technical Report, 2023】, 【8, Language Models are Few-Shot Learners, 2020, arXiv】, 【25, Mistral 7B, 2023, arXiv】, 【66, Llama 2: Open Foundation and Fine-Tuned Chat Models, 2023, arXiv】, 【84, OPT: Open Pre-trained Transformer Language Models, 2022, arXiv】】。预填充阶段接收输入token并通过解码器层生成单个输出token。在生成阶段，前一迭代生成的token被用来产生下一个token。这个过程是自回归的，每次迭代生成一个token。

多头注意力机制中的操作。LLM通常由多个解码器层组成，其中多头注意力是关键操作之一。如图2(b)所示，多头注意力首先生成查询（query）、键（key）和值（value）激活。键和值被缓存在片上内存以供后续操作使用，并作为KV缓存存储在片外内存中以备将来迭代使用。然后，查询与转置的键相乘，在预填充阶段，键直接从片上内存获取；而在生成阶段，则从片外KV缓存加载。同样，后续使用值的计算也需要以与键相同的方式访问片上内存或KV缓存。

2.2 用于高吞吐量LLM推理的批处理

批处理的优势。批处理是增强LLM服务系统推理吞吐量的常用方法【【1, Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve, 2024, arXiv】, 【20, NeuPIMs: NPU-PIM Heterogeneous Acceleration for Batched LLM Inferencing, 2024, ASPLOS ’24】, 【32, Efficient Memory Management for Large Language Model Serving with PagedAttention, 2023, SOSP ’23】, 【64, Flexgen: High-throughput generative inference of large language models with a single gpu, 2023, ICML】, 【77, Orca: A distributed serving system for {Transformer-Based} generative models, 2022, OSDI 22】】。它通过同时处理多个请求，并通过片上数据重用将内存密集型操作转换为计算密集型操作，从而提升推理吞吐量。例如，在前馈网络中，从内存读取的权重可以在多个请求之间重用，减少了内存访问并提高了吞吐量。

批处理在注意力机制中的局限性。图2(b)描绘了构成多头注意力层的两类操作。在激活-权重操作中，包括生成QKV和投影（Projection），请求的激活共享相同的权重。这使得片上数据重用成为可能，减少了内存访问和执行时间。然而，在激活-激活操作中，即查询与转置的键相乘或得分与值相乘时，每个请求需要不同的键或值，这使得批处理具有挑战性。因此，在这种情况下，批处理无法减少内存访问，在延迟方面没有性能优势。

2.3 LLM量化

量化作为内存压力的解决方案。随着LLM推理对内存容量的需求日益增加，社区已经探索通过量化技术来减少权重或激活的位精度，从而缓解这一需求。许多研究专注于减轻模型权重带来的内存压力，这对于加速小批量推理尤其有效【【17, OPTQ: Accurate quantization for generative pre-trained transformers, 2022, ICLR】, 【30, Squeezellm: Dense-and-sparse quantization, 2023, arXiv】, 【40, Awq: Activation-aware weight quantization for llm compression and acceleration, 2023, arXiv】, 【42, Qllm: Accurate and efficient low-bitwidth quantization for large language models, 2023, arXiv】, 【60, Omniquant: Omnidirectionally calibrated quantization for large language models, 2023, arXiv】, 【69, Outlier suppression+: Accurate quantization of large language models by equivalent and optimal shifting and scaling, 2023, arXiv】, 【70, Outlier suppression: Pushing the limit of low-bit transformer language models, 2022, NeurIPS】, 【78, Rptq: Reorder-based post-training quantization for large language models, 2023, arXiv】】。然而，由于KV缓存的大小随批处理大小和序列长度扩展，这些方法对于更大的批处理和长序列长度的加速效果有限。同时量化权重和激活的方法也面临类似的限制【【15, LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale, 2022, arXiv】, 【34, Enhancing Computation Efficiency in Large Language Models through Weight and Activation Quantization, 2023, EMNLP】, 【64, Flexgen: High-throughput generative inference of large language models with a single gpu, 2023, ICML】, 【71, Smoothquant: Accurate and efficient post-training quantization for large language models, 2023, ICML】】。

KV缓存量化的兴起。为此，研究人员最近开始将注意力转向KV缓存量化。其中一些工作提出在每个向量的基础上量化注意力层的键和值【【22, KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization, 2024, arXiv】, 【41, QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving, 2024, arXiv】, 【43, KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache, 2024, arXiv】】，而其他工作则建议将KV缓存量化为多个具有相似量级的向量组，并应用通道重排技术【【5, QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs, 2024, arXiv】, 【35, Tender: Accelerating Large Language Models via Tensor Decomposition and Runtime Requantization, 2024, ISCA】, 【39, DuQuant: Distributing Outliers via Dual Transformation Makes Stronger Quantized LLMs, 2024, arXiv】, 【86, Atom: Low-bit Quantization for Efficient and Accurate LLM Serving, 2024, arXiv】】。

3. 设计动机

本文专注于有效利用批处理LLM推理以提高吞吐量。为此，本节将深入分析LLM推理的特性，特别是比较非批处理和批处理执行的差异。

3.1 批处理LLM推理的特性分析

单实例推理。我们通过检查单请求和批处理请求场景下加速器核心利用率来分析批处理推理的瓶颈。图3(a)显示了处理单个请求时，预填充阶段和生成阶段核心利用率的差异。在预填充阶段，多个核心被用来并行处理多个输入token。然而，生成阶段仅处理前一迭代的一个token，这是一个高度串行的过程，限制了并行化。因此，加速器只能利用少数核心，导致生成所有输出token的延迟较高。

批处理推理。图3(b)展示了批处理推理的核心利用率。与单请求场景类似，批处理请求的预填充阶段也可以在多个核心上并行化，从而实现高核心利用率。然而，处理批处理请求时的生成阶段仍然显示出利用率不足，由于带宽争用导致延迟更长。这是因为来自注意力层的键值无法在请求之间进行批处理和共享，并且加速器为每个请求利用的核心数量有限。图3(c)显示了在NVIDIA A100 GPU上运行Llama2-13B模型生成阶段的GPU核心利用率，表明利用率不足主要源于多头注意力操作。

批处理：双刃剑。一个直接的解决方案是增加批处理大小以充分利用所有核心，但这引入了一个新挑战。注意力操作所需的键值大小与批处理大小成正比。然而，随着批处理大小的增加，KV缓存导致的内存占用也随之增长，导致操作延迟更长，这受限于内存带宽。最近的技术，如分组查询注意力和多查询注意力，通过缩小KV缓存大小来减少内存开销，但不能完全解决这个问题【【3, GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, 2023, arXiv】, 【61, Fast Transformer Decoding: One Write-Head is All You Need, 2019, arXiv】】。

关键观察。我们的分析提出了两个关键观察：
(1) 虽然预填充阶段充分利用了计算资源，但生成阶段往往未能做到。这种不平衡为通过更大的批处理大小提高吞吐量提供了机会，但同时也增加了对内存容量的需求，从而推高了系统构建成本。
(2) 由于激进的批处理导致的KV缓存大小增加，也对内存带宽提出了高要求，以高效计算注意力操作，从而在批处理推理期间延迟了整个生成阶段。

3.2 带宽与容量的权衡

内存技术的权衡。虽然LLM服务既需要高带宽又需要大容量，但内存技术在这两种资源之间存在权衡。高带宽内存（HBM）牺牲了大部分容量以提供卓越的带宽，而低功耗双倍数据速率（LPDDR）DRAM则处于权衡空间的另一端。基于这些观察和见解，我们进行了初步研究，以更好地理解这些冲突资源对LLM推理吞吐量的性能影响。

初步研究结果。图4(a)和(b)展示了使用两种不同内存类型的加速器变体对Llama2-13B和OPT-30B模型的吞吐量比较结果。我们使用了一个现有的LLM定制加速器【48, HyperAccel Latency Processing Unit (LPUTM) Accelerating Hyperscale Models for Generative AI, 2023, IEEE Hot Chips 35 Symposium (HCS)】，并将输入和输出序列长度都设置为1K。图4(c)列出了两种被评估加速器的规格。对于较小的LLM模型，两种加速器都实现了随批处理大小增加而次线性扩展的吞吐量，其中基于HBM的加速器因其优越的带宽而性能最佳。然而，对于更大的LLM模型和批处理大小，基于HBM的加速器面临扩展挑战。相比之下，基于LPDDR的加速器可以容纳更大的批处理，展现出最佳性能。

结论。总结来说，虽然高内存带宽对于实现高吞吐量至关重要，但充足的内存容量对于高效服务LLM同样必不可少，尤其是在处理较大模型和批次时。最近的一些研究也强调了内存容量在LLM推理中的重要性，并建议在LLM加速器中引入大容量内存【【53, An LPDDR-based CXL-PNM Platform for TCOefficient Inference of Transformer-based Large Language Models, 2024, HPCA】, 【83, LLMCompass: Enabling Efficient Hardware Design for Large Language Model Inference, 2024, ISCA】】。尽管业界对大规模LLM推理的需求日益增长，但当前的内存技术难以同时满足容量和带宽要求，迫使人们必须做出选择。我们通过在各种批处理推理场景中评估我们配备LPDDR或HBM内存的解决方案来探索这一权衡。

3.3 KV缓存量化

算法方法应对内存墙。利用提供足够带宽或容量的内存，并结合系统横向扩展，可以帮助解决批处理推理中的瓶颈。然而，这种方法会带来巨大的系统构建成本和严重的资源利用不足，既不是根本性的也不是可持续的解决方案。为了从根本上解决内存墙挑战，社区探索了旨在减少对两种冲突资源需求的算法方法。量化是这样一个方向，因其能同时缓解容量和带宽瓶颈而广受认可。

纯权重化量化的局限性。如图5(a)所示，随着批处理大小的增加，模型权重所需的内存容量保持不变，而KV缓存的大小则不断增长，最终占据整个设备内存。图5(b)比较了4位权重量化和4位KV缓存量化的性能。结果显示，权重量化带来的性能提升微乎其微，表明其在解决批处理LLM推理的内存墙问题上效果不佳。然而，KV缓存量化带来了更大的性能增益，证明了其在缓解内存压力方面的有效性。

现有KV缓存量化方案的局限性。最近，大量前期工作探索了KV缓存量化方法【【5, QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs, 2024, arXiv】, 【35, Tender: Accelerating Large Language Models via Tensor Decomposition and Runtime Requantization, 2024, ISCA】, 【41, QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving, 2024, arXiv】, 【43, KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache, 2024, arXiv】, 【86, Atom: Low-bit Quantization for Efficient and Accurate LLM Serving, 2024, arXiv】】。虽然这些工作开辟了一个新的研究方向，但它们受限于为实现量化而产生的大量运行时开销、为追求更快的量化而牺牲准确性，或两者兼而有之。

• QServe【41】和Atom【86】通过应用RPTQ【78】中采用的变换矩阵来重排键值通道，而QServe还通过应用SmoothQuant【71】引入的缩放矩阵来处理异常值。
• Tender【35】通过间接索引执行通道重排，并将具有相似量级的键值通道分组。这些解决方案的局限性在于，由于其量化粒度较低，会导致精度下降，并且通道重排会带来额外的运行时开销。
• KIVI【43】和KVQuant【22】基于每个KV通道的值量级呈现不同模式的洞察，提出了基于每个向量的混合精度量化。这些方法通过将异常通道与其余部分隔离来最小化精度损失，但它们因排序操作或混合精度计算而产生巨大开销，这在很大程度上抵消了量化的性能增益。

Oaken的动机。现有解决方案的这些局限性强调了对一种KV缓存量化技术的需求，该技术能在不牺牲准确性的前提下实现低位宽。此外，这种低位宽必须在硬件层面转化为切实的吞吐量提升。为了满足这些苛刻的要求，我们提出了一种算法-硬件协同设计的加速解决方案，用于高效的批处理LLM推理，即Oaken。

---

A2 方法细节

4. Oaken的KV量化算法

我们的量化算法由我们通过观察LLM推理期间KV缓存值分布获得的三个主要经验属性驱动。基于这三个观察到的属性，我们将三种算法技术融入Oaken的KV缓存量化中。下面，我们将首先讨论KV缓存分布的经验分析结果，然后逐一描述这三种技术。

4.1 KV分布的观察

现有工作的启示。现有的LLM量化方法通常表明，处理大值对模型精度有显著影响【【15, LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale, 2022, arXiv】, 【19, OliVe: Accelerating Large Language Models via Hardware-friendly Outlier-Victim Pair Quantization, 2023, ISCA ’23】, 【22, KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization, 2024, arXiv】, 【33, OWQ: Outlier-Aware Weight Quantization for Efficient Fine-Tuning and Inference of Large Language Models, 2024, AAAI】, 【71, Smoothquant: Accurate and efficient post-training quantization for large language models, 2023, ICML】, 【80, Mokey: Enabling narrow fixed-point inference for out-of-the-box floatingpoint transformer models, 2022, ISCA】, 【86, Atom: Low-bit Quantization for Efficient and Accurate LLM Serving, 2024, arXiv】】。此外，其他先前工作表明，接近零的小值可能因量化期间的下溢而消失，导致更大的误差【【2, Intriguing Properties of Quantization at Scale, 2023, NeurIPS】, 【13, Accurate Neural Training with 4-bit Matrix Multiplications at Standard Formats, 2024, arXiv】, 【27, F8Net: Fixed-Point 8-bit Only Multiplication for Network Quantization, 2022, ICLR】, 【34, Enhancing Computation Efficiency in Large Language Models through Weight and Activation Quantization, 2023, EMNLP】】。这些观察强调了分析量化目标分布和特性的重要性。我们检查了多个LLM和数据集上KV缓存的值分布，并得出了设计有效量化技术的关键见解。

观察1。图6(a)展示了使用Wikitext2数据集时，不同LLM模型中每个解码器层的KV缓存值的最小和最大范围。值得注意的是，键和值的量级在不同模型之间以及同一模型内部的不同解码器层之间存在差异。这些变化是每个模型和解码器层的独特属性，由其模型权重的差异驱动。从这一观察中，我们得到的见解是，量化因子应为每个模型及其各个解码器层单独确定。

观察2。图6(b)显示了使用Llama2-7B模型和Wikitext2、PIQA、Hellaswag数据集时KV值的最小和最大范围。我们看到，KV缓存值的范围在这些数据集中保持一致。这一观察意味着量化因子无需针对输入序列进行定制，而只需每层使用一个全局量化因子。

观察3。先前的工作已经识别出KV缓存内通道量级的模式，并提出了逐向量或逐向量组的量化技术来利用这一模式【【22, KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization, 2024, arXiv】, 【41, QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving, 2024, arXiv】, 【43, KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache, 2024, arXiv】, 【86, Atom: Low-bit Quantization for Efficient and Accurate LLM Serving, 2024, arXiv】】。图6(c)展示了Llama2-7B和OPT-6.7B模型第六层键的前4%值的分布。分布呈现出多条垂直线，表明高量级值集中在特定通道中。通道量级的这种模式与先前的观察一致，表明需要进行逐向量或逐向量组的量化。然而，这一结果也揭示了该模式的例外情况，表现为不连续的线和点，这在使用每个向量或向量组仅一个量化尺度时会导致精度下降。基于这一观察，我们引入通道内的多个量化组，根据每个元素的大小将键值向量分割成组，以保持精度。

总结。总之，我们为设计Oaken的量化技术得出了以下三个见解：

• Oaken应为每个模型和解码器层分别确定量化因子。
• Oaken可以采用一个通用的缩放因子，而与输入提示无关，显示出其对数据模式的不敏感性。
• Oaken应使用根据每个向量内值的量级分段的多个量化组。

4.2 算法概述

Oaken算法流程。我们提出了一种量化算法，旨在通过最大化KV缓存的压缩率同时最小化量化损失来提高性能，这基于第4.1节的观察和见解。图7展示了Oaken量化算法的总体流程。Oaken的量化技术由三个部分组成：(1) 基于阈值的在线-离线混合量化，用于分离和量化正常值和异常值；(2) 组移位量化，用于量化具有较大量级的值；(3) 融合的密集-稀疏编码，用于最小化稀疏矩阵带来的容量开销。以下各节将介绍每个组件的详细设计。

4.3 基于阈值的在线-离线混合量化

细粒度分组。Oaken通过隔离与其他典型正常值相比异常大或异常小的值来最小化量化损失。我们提出了一种基于阈值的在线-离线混合量化方法，以实现更细粒度的分组。Oaken将每个token的KV向量分为三个量化组：外部组（outer）、中间组（middle）和内部组（inner）（图7(a)）。中间组由大多数KV值所属的正常值组成，而外部组和内部组则由分别具有大和小量级的异常值组成。Oaken通过隔离外部组来防止量化尺度因大量级异常值而偏斜，并通过隔离内部组来确保小量级异常值在量化过程中不会消失【【2, Intriguing Properties of Quantization at Scale, 2023, NeurIPS】, 【13, Accurate Neural Training with 4-bit Matrix Multiplications at Standard Formats, 2024, arXiv】, 【27, F8Net: Fixed-Point 8-bit Only Multiplication for Network Quantization, 2022, ICLR】, 【34, Enhancing Computation Efficiency in Large Language Models through Weight and Activation Quantization, 2023, EMNLP】】。

离线异常值阈值分析。为了将异常值与正常值分开，通常使用topK操作来保持恒定的异常值比例【22】。虽然这种方法能最小化量化损失，但topK操作本质上是排序，时间复杂度为$O(n \log n)$，在推理期间执行会引入巨大开销，降低端到端性能。Oaken采用离线异常值阈值分析，利用第4.1节讨论的KV缓存分布的一致特性来避免昂贵的在线操作。

分组标准。将KV缓存分为三组的标准基于通过离线阈值分析确定的四个阈值：$S_{low}$、$S_{high}$、$T_{low}$和$T_{high}$。使用这些阈值，外部组、中间组和内部组的定义如下：

Oaken使用样本输入提示执行大约一百次离线推理，以收集每个解码器层KV缓存的分布信息。在分析过程中，使用topK操作从每次推理运行的KV缓存中提取四个组阈值，并为每个解码器层计算它们的平均值。这些统计数据随后用于建立KV缓存的组阈值。如第4.1节所述，Oaken的离线分析应为每个模型单独执行，而KV缓存分布和分析过程与分析数据集和未来输入均无关。它仅需少量推理即可有效最小化精度损失，确保开销最小且一次性。

均匀量化。Oaken采用均匀量化，其中缩放因子S仅使用简单的统计数据计算，以最小化硬件复杂性：

其中B是量化值的位宽，Max和Min分别代表待量化值的最大值和最小值。将值x转换为其量化值的均匀量化函数，使用Eq. 2中的缩放因子定义如下：

Oaken为三个量化组中的每一个找到最小值和最大值，并在线动态计算每个组的量化缩放因子。

在线KV缓存量化。如前所述，Oaken对KV缓存执行逐token量化，仅关注每个注意力层中新生成的键值向量。Oaken使用从离线分析获得的四个组阈值动态地将KV缓存值分为三组。然后，它检索每个量化组的最小值和最大值，在线计算缩放因子，并相应地量化每个组内的值。

4.4 组移位量化

面临的挑战。Oaken基于阈值的在线分组通过将值分为外部、中间和内部三组来有效减轻信息损失。然而，这种方法带来了新的挑战。当量化外部组时，其值具有较大的量级，直接应用均匀量化会导致信息损失。先前的工作通过为异常值使用混合精度（例如FP16），与正常值使用的精度（例如INT4）不同来解决这个问题【【22, KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization, 2024, arXiv】, 【30, Squeezellm: Dense-and-sparse quantization, 2023, arXiv】, 【33, OWQ: Outlier-Aware Weight Quantization for Efficient Fine-Tuning and Inference of Large Language Models, 2024, AAAI】, 【86, Atom: Low-bit Quantization for Efficient and Accurate LLM Serving, 2024, arXiv】】。然而，为异常值使用混合精度会引入稀疏性，每个条目需要23位的存储成本：16位用于值，6位用于索引，1位用于指示组。处理稀疏的混合精度异常值还需要额外的硬件模块，这增加了加速器的复杂性。虽然当异常值比例很小时这些开销可以忽略不计，但随着比例增加，开销会变得相当可观。此外，这种额外的复杂性使得通过调整组阈值来探索精度-性能权衡的最佳点变得具有挑战性。

组移位量化算法。解决此问题的一个直接方法是也对异常值进行量化。然而，如第4.1节所述，由于其宽广的量级范围，量化外部组具有挑战性。为了解决这个问题，我们提出了组移位量化，这是一种硬件高效的算法，可在将异常值压缩到更低位宽的同时最小化量化损失。Oaken组移位算法的核心思想是使用从离线分析获得的阈值来移动整个组，以缩小值的范围，从而使低位量化成为可能。例如，对于外部组，我们从大于$S_{high}$的值中减去$S_{high}$，从小于$S_{low}$的值中减去$S_{low}$。虽然中间组对应于正常值，我们的组移位算法也可以以相同的方式应用。图7(b)显示，应用上述方法会移动外部组和中间组的分布，将它们集中在更窄的范围内。因此，Oaken可以使用组移位将跨越宽范围的组量化到低位宽，从而最小化量化损失。请注意，组移位方法不需要超出第4.3节描述的离线分析所获得的额外信息。

Oaken的量化函数。总之，Oaken的量化函数$Q_{Oaken}(x)$，将值x转换为量化值，定义如下：

其中Q(x)是Eq. 3中定义的量化函数。Oaken将中间组量化为4位，内部和外部组量化为5位。

4.5 融合的密集-稀疏编码

存储策略。Oaken采用先前工作中提出的密集-稀疏编码策略，以高效存储密集的正常值和稀疏的异常值【【22, KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization, 2024, arXiv】, 【30, Squeezellm: Dense-and-sparse quantization, 2023, arXiv】】。在Oaken中，由正常值组成并占KV缓存大部分的中间组存储在密集矩阵中。由异常值组成的外部组和内部组使用稀疏矩阵格式坐标列表（COO）【【16, Spatula: A Hardware Accelerator for Sparse Matrix Factorization, 2023, MICRO ’23】, 【18, SparseP: Towards Efficient Sparse Matrix Vector Multiplication on Real Processing-In-Memory Architectures, 2022, Proc. ACM Meas. Anal. Comput. Syst.】, 【72, SpaceA: Sparse Matrix Vector Multiplication on Processing-inMemory Accelerator, 2021, HPCA】, 【85, SpArch: Efficient Architecture for Sparse Matrix Multiplication, 2020, HPCA】】进行存储，密集矩阵中相应的元素被置零。Oaken中使用的COO格式需要额外的6位来指示每个值的位置，以及1位来表示量化组，以及用于表示每个条目值的位数。

融合编码方法。为了进一步减少容量开销，我们建议利用密集矩阵中的置零元素。这些置零元素对应于KV缓存中最初被异常值占据的位置，在将KV缓存分离为密集和稀疏矩阵后保持未使用。我们引入了一种融合的密集-稀疏编码方法，该方法重新利用这些未使用的4位来存储部分异常值，如图7(c)所示。具体来说，量化的5位异常值中的四位被嵌入到密集矩阵的置零元素中，而剩余的6个索引位、1个组位和1个符号位则以稀疏COO格式存储。

优势。由于COO格式中的索引位已经指示了异常值在密集矩阵中的位置，因此不需要专门的标志来表示它们的存在。此外，由于稀疏矩阵中的每个条目固定为8位且内存对齐，内存管理单元可以高效地以页为基础处理密集和稀疏矩阵，其细节将在第5.2节中描述。通过将我们的融合密集-稀疏编码策略与组移位量化相结合，我们将每个异常值条目的位宽从23位减少到仅8位，从而在保持内存对齐的同时，提高了KV缓存的压缩率。

5. Oaken的加速器架构

5.1 架构概述

整体架构。图8展示了所提出架构的概览，主要由计算核心、内存控制器、主机接口和互连网络组成。计算核心旨在支持端到端的LLM推理操作。内存控制器处理设备内存以读取数据，包括模型参数、键和值，并负责将键和值写回内存。主机接口采用基于PCIe的连接与主机系统通信。该接口还管理传入请求的调度，并将它们分配到计算核心以进行高效处理。互连网络连接这些组件，并经过优化以在内存读取期间最大化带宽利用率，确保数据高效传输到计算核心。这种设计使得所有内存控制器可以并发地用于读取模型参数并将其分配到计算核心。另一方面，从计算核心进行的内存写入频率较低，涉及的数据量较小，从而减少了带宽消耗并简化了逻辑设计，而不会影响性能。

5.2 Oaken计算核心

总体设计。所提出加速器的主要模块是计算核心，它们改编自LPU【48】中引入的架构，以实现端到端的LLM推理。每个计算核心由一个矩阵处理单元（MPU）和一个向量处理单元（VPU）组成，旨在逐个token地执行LLM推理操作。这些处理单元被设计为在整个过程中保持高利用率，同时最小化低效逻辑并确保低延迟。MPU被设计为流式读取内存中的权重以执行高效的矩阵-向量乘法，而VPU处理矩阵-向量计算之间的逐元素操作。直接内存访问（DMA）单元通过从内存读取权重以供给处理单元，并将KV缓存写回内存来促进数据传输。该DMA单元还集成了量化/反量化引擎和内存管理单元（MMU），所有这些对于实现所提出的KV量化技术至关重要。

量化引擎。图9(a)显示了DMA单元中的量化引擎，该引擎旨在执行在线KV缓存量化。

• 首先，① 分解器模块根据离线确定的异常值阈值将传入的激活划分为三个量化组。然后，它通过从这些组的值中减去其流式传输的阈值来对外部组和中间组执行组移位。
• 最后，中间组被导向正常值量化路径，而无论是来自内部组还是外部组的异常值都被路由到异常值量化路径，并在相应的备用路径中插入零。
• ② 正常值和 ③ 异常值量化器模块都处理每个键和值向量的量化。量化缩放因子在运行时根据每个组的逐token最小值和最大值动态计算。计算完缩放因子后，这两个模块最终执行4/5位均匀量化。
• 量化的正常值和异常值通过一个或门合并，并发送到量化的密集矩阵。异常值模块为COO转换生成索引和组标志。它采用一个去零移位器（zero-remove shifter）【【57, Compressing DMA engine: Leveraging activation sparsity for training deep neural networks, 2018, HPCA】, 【67, Spatten: Efficient sparse attention architecture with cascade token and head pruning, 2021, HPCA】】来实现融合的密集-稀疏编码，优化内存使用。

反量化引擎。反量化引擎，如图9(b)所示，也集成在DMA单元中，用于反量化从内存中检索的KV缓存。

• ④ 正常值反量化器模块缓冲传入的密集数据，以与由异常值反量化器模块处理的稀疏数据同步。
• ⑤ 异常值反量化器模块通过执行一个插零（zero-insert）【【57, Compressing DMA engine: Leveraging activation sparsity for training deep neural networks, 2018, HPCA】, 【67, Spatten: Efficient sparse attention architecture with cascade token and head pruning, 2021, HPCA】】操作来处理稀疏COO数据，以恢复原始数据对齐。它使用稀疏数据的索引和组信息识别融合异常值的原始位置，并相应地插入必要的零。
• 然后，两个反量化器模块都恢复数据，这些数据被缓冲以与对应模块的输出对齐。最后，来自两个反量化器模块的输出通过一个或门合并，并转发到处理单元。

由于反量化引擎不需要整个KV缓存来进行操作，我们设计它以流式方式工作。这种设计使反量化引擎能够保持低延迟，同时高效处理所有过去的KV缓存。

内存管理单元（MMU）。图10展示了内存管理单元的操作，它管理量化的KV缓存的读写。我们设计MMU单元以基于页面的方式处理密集和稀疏矩阵，优化带宽利用率。它支持突发模式下的多个内存访问和流线型操作，以隐藏内存和量化/反量化操作的延迟。由于Oaken的MMU单元共享一个公共地址空间，每个计算核心中的MMU独立运行，防止干扰。如果没有这个专门的MMU，处理可变大小的稀疏矩阵将需要额外的索引、重塑和后续操作的开销。MMU单元的设计存在两个主要挑战：

• (1) 地址和传输大小管理： 密集矩阵具有可预测的大小，并且在内存空间中良好对齐，而稀疏矩阵的大小各不相同。因此，需要为密集和稀疏数据都设置管理表以适应这种可变性。这些表包含KV缓存的虚拟到物理地址映射和传输大小，考虑了每个注意力头的最大序列长度。物理地址和传输大小在推理期间通过按需检查可用页面动态计算。
• (2) 读写粒度和顺序确定： 为了最大化内存带宽利用率，应尽可能利用突发访问以减少内存事务的总数。写入KV缓存涉及相对较小的大小，因为它只包含当前token的键值，而读取则需要检索所有先前token的KV缓存。为了解决这个问题，Oaken将当前token的KV缓存组织成一种便于后续操作中突发读取的布局。当前层生成的键值向量按注意力头划分并写入不同的页面，如第5.2节所述。当为下一个token生成KV缓存时，它会以类似的方式划分并顺序写入，紧跟在前一个token的KV缓存之后。这种顺序排列允许所有先前token的KV缓存可以以突发模式读取，使Oaken能够高效利用带宽。

5.3 Token级批处理调度

高效调度。高效的调度对于Oaken高效服务LLM推理至关重要。Oaken中的每个计算核心都经过优化，可以高效地处理单个token。在预填充阶段，来自每个请求的输入token被调度到多个核心上进行并行处理。然而，在生成阶段，每个核心处理来自一个请求的单个输出token，这会降低硬件利用率。对于更大的批次，Oaken通过并行处理多个请求来提高整体核心利用率。

延迟隐藏。尽管KV缓存量化和反量化的开销很小，但Oaken通过将它们与其他操作重叠来进一步最小化这一开销。在批处理推理中，KV缓存不能在核心之间共享，因为每个核心处理不同的请求，迫使每个核心独占内存带宽。Oaken采用一种调度策略，通过将KV量化和反量化与来自其他请求的DMA读取和注意力计算重叠来隐藏延迟。

---

A4 实验

实验环境

模型与数据集

• 模型：实验使用了Llama2-7B, 13B, 70B 【66】, OPT-6.7B, 13B, 30B 【84】, Mistral-7B 【25】 和 Mixtral-8x7B 【26】。其中Llama2、Mistral和Mixtral模型实现了分组查询注意力，Mistral和Mixtral还包含滑动窗口机制，Mixtral进一步集成了专家混合（MoE）层。
• 数据集：
  • 模型精度评估：使用Wikitext2, PIQA, Winogrande, 和 Hellaswag。Wikitext2用于评估困惑度（Perplexity，越低越好），其他三个数据集用于评估零样本准确率（zero-shot accuracy，越高越好）。
  • 真实世界基准测试：使用来自Azure LLM推理服务的两个开源生产轨迹：Conversation【47, 54】和BurstGPT【68】，模拟真实的推理服务场景。

硬件配置

• Oaken加速器：
  • 基于LPU【21, 53】的硬件模拟器进行扩展开发。
  • 提供两个版本：Oaken-LPDDR（配备LPDDR内存）和Oaken-HBM（配备与A100 GPU相同的HBM内存）。
  • HBM提供更高带宽但容量较小，LPDDR则相反。
• 基准平台：
  • NVIDIA A100 GPU，配备80 GB HBM内存【49】。
  • 对于小型模型（Llama2-7B/13B, Mistral-7B）使用单张A100。
  • 对于大型模型（OPT-30B, Mixtral-8x7B, Llama2-70B），使用两张A100，采用流水线并行。
• 硬件规格：详细规格如表1所示。

软件配置

• Oaken实现：使用SystemVerilog以RTL实现，并通过Synopsys VCS进行功能验证，使用Synopsys Design Compiler在TSMC 28nm工艺下进行综合，目标时钟频率1 GHz。
• 基准软件：
  • vLLM【32】作为FP16操作的GPU基准。
  • 量化基准包括KIVI【43】, QServe【41】, KVQuant【22】（均在A100上运行），以及Tender【35】（一个开源的LLM推理加速器模拟器，其内存和计算能力与A100对齐）。
  • 所有量化基准均采用4位仅KV缓存量化，以保证公平比较。

实验设置

• 阈值设置：在所有评估中，外部组、中间组、内部组的比例固定为4%、90%和6%。这一全局配置适用于所有模型和数据集。
• 离线分析：Oaken的离线分析通过在推理过程中收集topK值来确定量化组的边界，并对收集到的值取平均。该过程使用Wikitext2数据集，对所有LLM模型进行，大约需要一百次推理，耗时约十分钟（即使对于Llama2-70B），开销可忽略不计。

实验结果

吞吐量

• 实验内容：比较了GPU基准（vLLM, KVQuant, KIVI, QServe）和ASIC加速器（Tender, Oaken）在六种LLM上的端到端吞吐量，批处理大小从16到256不等。
• 实验结果（图11）：在批处理大小为256时，Oaken-LPDDR相比vLLM和QServe平均分别取得了1.79倍和1.58倍的吞吐量提升。这是因为Oaken通过最小化对KV缓存的内存访问，减轻了带宽和容量瓶颈，从而减少了占推理时间大部分的注意力操作的执行时间。
• 分析结论：GPU基准和Tender在小批量和小模型上表现良好，但随着批量增大，因容量限制而性能饱和。Oaken-HBM在小模型上优于其他方案，但处理大模型或大批量时也受容量限制。Oaken-LPDDR则展现出良好的可扩展性，证明了其KV缓存量化技术的有效性。

准确率

• 实验内容：在Wikitext2, PIQA, Winogrande, 和 Hellaswag数据集上评估了各基准在八种LLM上的准确率，并报告了Llama2模型上的有效位宽。
• 实验结果（表2）：Oaken与原始FP16基准相比，平均准确率损失为0.87%。其准确率比KVQuant和KIVI分别低0.54%和0.32%，但比QServe高1.38%。
• 分析结论：KVQuant和KIVI通过使用稀疏布局和细粒度分组，以更大的有效位宽换取了更高的准确率，但其在线排序和混合精度操作的开销巨大。Tender, Atom和QServe的开销较小，但由于采用粗粒度的量化方案，准确率损失更大。Oaken在准确率和性能开销之间取得了良好的平衡。

准确率与压缩率的权衡

• 实验内容：通过调整量化组的比例，探索Oaken在Llama2-7B模型上的准确率与压缩率（性能）之间的权衡。同时，分析了使用不同分组数量（2, 3, 4, 5组）对准确率和有效位宽的影响。
• 实验结果（图12(a), 表3）：实验表明，4%（外部组）、90%（中间组）和6%（内部组）的比例是一个帕累托最优选择。增加分组数量可以提高准确率，但会增加位宽和硬件复杂性（如内存对齐问题）。
• 分析结论：三组量化方案（外部、中间、内部）在成本和准确率之间提供了最佳平衡。

延迟分解

• 实验内容：分解了LPU和Oaken-LPDDR在不同批处理大小下的端到端推理延迟，并测量了Oaken量化算法在GPU上的操作延迟。
• 实验结果（图12(b)）：Oaken通过KV量化缓解了带宽瓶颈，使得注意力操作的执行时间平均比LPU短55.0%。在批处理大小为64时，量化和反量化延迟仅占总延迟的1.29%和3.23%，且这些延迟被调度策略隐藏。相比之下，Oaken算法在GPU上的实现由于CUDA中的线程束分化（warp divergence）而延迟较长。
• 分析结论：Oaken的定制硬件能高效执行其量化算法，并将开销成功隐藏，从而实现端到端延迟的降低。

对序列长度的敏感性

• 实验内容：在Llama2-13B模型上，将总序列长度从1K增加到32K，测试吞吐量变化。
• 实验结果（图13）：在短序列（<8K）下，计算密集型操作占比较大，GPU基准（QServe, vLLM）因其更高的并行资源而表现更优。随着序列长度增加，Oaken-HBM凭借其高内存带宽和KV缓存量化超越了其他基准。然而，所有基于HBM的系统（包括QServe和Oaken-HBM）在处理超过16K的序列时因容量不足而失败。
• 分析结论：Oaken-LPDDR通过KV量化和大容量内存缓解了带宽和容量双重压力，能够处理长达32K的序列，展示了其在长上下文场景下的优势。

真实世界基准测试

• 实验内容：使用Conversation和BurstGPT两个真实世界LLM推理轨迹，在Llama2-13B和Mixtral-8x7B模型上评估生成吞吐量。
• 实验结果（图14）：Conversation轨迹的输出较短，Oaken的优势不明显。BurstGPT轨迹的输出较长，Oaken的KV缓存量化优势得以体现。对于Mixtral-8x7B模型，由于其本身使用了分组查询注意力减少了KV缓存，量化带来的性能增益相对较小，但随着批处理大小增加或在长生成场景下，Oaken-LPDDR的性能优势依然显现。
• 分析结论：在真实世界场景中，Oaken相比现有解决方案在Llama2-13B和Mixtral-8x7B模型上均表现出优势，尤其是在长序列生成和较大批处理大小的场景下。

面积与功耗

• 实验内容：分析了Oaken量化和反量化引擎在TSMC 28nm工艺下的面积开销，并评估了整个加速器的功耗。
• 实验结果（表4）：量化和反量化引擎分别仅占总计算核心面积的1.86%和6.35%。嵌入Oaken模块的整个加速器功耗为222.7W，比A100 GPU的400W TDP低44.3%。
• 分析结论：集成Oaken的量化和反量化模块只带来极小的硬件开销，同时相比GPU提升了性能并实现了更高的能效。

---

A5 结论

批处理LLM推理面临着由高内存带宽和容量需求带来的重大挑战，这一问题因现代LLM产生长序列输出导致KV缓存日益增大而加剧。本文通过提出一个加速解决方案Oaken来应对这一挑战，该方案联合利用了(1)一种离线-在线混合KV缓存压缩技术和(2)为该算法量身定制的、可与LLM加速器集成的自定义硬件模块。Oaken有效地释放了足够的带宽和容量，这是其他方式难以实现的，从而在仅有边际精度损失的情况下带来了显著的吞吐量提升。这些引人注目的优势表明，Oaken高效地解决了现代LLM服务的两大主要瓶颈。