Haojun Xia, Xiaoxia Wu, Jisen Li, Robert Wu, Junxiong Wang, Jue Wang, Chenxi Li, Aman Singhal, Alay Dilipbhai Shah, Alpay Ariyak, Donglin Zhuang, Zhongzhu Zhou, Ben Athiwaratkun, Zhen Zheng, Shuaiwen Leon Song
本文旨在解决大型语言模型(LLM)在长上下文推理中,由KV缓存(KV cache)引起的巨大内存占用瓶颈。虽然4-bit KV量化能保持准确率,但2-bit量化通常会导致性能下降,尤其是在长上下文推理任务上。本文通过算法-系统协同设计,提出了一种名为 Kitty 的混合精度KV缓存方案,以弥合2-bit量化的准确率差距。
核心问题与研究目标:
- 核心问题: LLM在处理长序列时,KV缓存的内存占用远超模型权重,成为主要的系统瓶颈。例如,LLaMA3-70B模型在处理32个128K序列长度的请求时,KV缓存需要超过1.2TB,远超单个GPU的内存容量。
- 研究目标: 实现接近2-bit的KV缓存压缩率,同时将模型准确率的损失降至可忽略不计的水平,并设计一个高效的推理系统来支持该方案,从而在有限的内存预算下提升吞吐量和批处理大小。
创新点与主要贡献:
1. 揭示了2-bit KV量化的瓶颈: 本文通过实验发现,像KIVI这样的先进2-bit KV缓存量化方法会显著降低推理LLM的准确率,揭示了低比特KV量化中存在的关键表示瓶颈。
2. 提出动态信道级精度提升算法: 本文提出了一种新颖的2-bit KV缓存量化算法——动态信道级精度提升(Dynamic Channel-wise Precision Boost)。该算法受到信道级模式和量化敏感度的观察启发,通过对Key缓存的信道按敏感度排序,并仅将一小部分关键信道保持在较高的精度(如4-bit),而将其余信道积极压缩至2-bit。结合其他优化,本文提出了一个名为Kitty的系统性KV缓存量化方案。实验证明,Kitty在准确率上显著优于先前的工作。
3. 设计并构建了端到端的推理系统: 为支持新颖的2-bit KV量化算法,本文设计并构建了一个端到端的GPU推理系统。该系统解决了在保持页面布局合并和反量化统一的情况下,处理动态4-bit信道提升的挑战,避免了分散读取或硬编码掩码。该系统能够将KV内存减少近8倍,批处理大小增加8倍,并在相同内存预算下实现2.1倍至4.1倍的吞吐量提升,同时准确率损失可忽略不计。
低比特量化的前景与挑战 低比特量化是将全精度(如FP16)的键(Key)和值(Value)转换为紧凑表示(如8位、4位甚至2位)的有前途的方向,以减少KV缓存的内存占用。与静态的权重化不同,KV缓存量化必须在运行时动态执行,这使得问题更具挑战性。
逐token与逐信道量化 KV缓存可表示为形状为(B, H, L, D)的矩阵,其中B是批大小,H是KV头数量,L是序列长度,D是每个头的隐藏大小。因此,对于特定请求的每个头,KV缓存是一个(L, D)的二维矩阵。量化可以沿L或D维度进行。逐token量化为每个token的键或值向量应用独立的缩放因子(和可选的零点),因此每个token可以独立量化。而逐信道量化则为每个信道计算一个跨所有token共享的缩放因子,这种方法在Key缓存量化中被观察到能更好地保持准确性【索引17,Kivi: a tuning-free asymmetric 2-bit quantization for kv cache,2024,ICML'24】、【索引11,Kvquant: Towards 10 million context length llm inference with kv cache quantization,2024】、【索引35,Kv cache is 1 bit per channel: Efficient large language model inference with coupled quantization,2024a】。实践中,最近的工作大多对K缓存采用逐信道量化,对V缓存采用逐token量化。
混合精度量化 将整个KV缓存量化到极低精度(例如2位)通常会严重损害模型质量。为缓解此问题,混合方案会保留一部分元素为较高精度,同时量化其余部分。例如,KVQuant【索引11,Kvquant: Towards 10 million context length llm inference with kv cache quantization,2024】识别异常值并以FP16稀疏表示存储它们,但这种方法对硬件不友好,系统效率低,因为它引入了稀疏-密集乘法的额外运行时开销,这在GPU上可能很慢【索引31,Flash-llm: Enabling cost-effective and highly-efficient large generative model inference with unstructured sparsity,2023】、【索引9,Sparse gpu kernels for deep learning,2020】。KIVI【索引17,Kivi: a tuning-free asymmetric 2-bit quantization for kv cache,2024,ICML'24】和BitDecoding【索引6,Bitdecoding: Unlocking tensor cores for longcontext llms with low-bit kv cache,2025】保留最近的token为全精度以维持准确的局部上下文,但在低精度(如2位)下仍有严重的精度下降。MiniKV【索引25,Minikv: Pushing the limits of llm inference via 2-bit layer-discriminative kv cache,2024】为每层提出了层判别式的比特分配方案。KVTuner【索引14,Kvtuner: Sensitivity-aware layerwise mixed-precision kv cache quantization for efficient and nearly lossless llm inference,2025】提出为KV缓存调整逐层的混合精度比特宽度。QuaRot【索引2,Quarot: Outlier-free 4-bit inference in rotated llms,2024】对激活值和KV缓存应用正交旋转,使分布无异常值。
非KV量化 尽管本文重点是KV量化,但值得注意的是,一些权重和激活值量化技术是正交且可能互补的【索引34,Smoothquant: Accurate and efficient post-training quantization for large language models,2024a】、【索引15,Awq: Activation-aware weight quantization for on-device llm compression and acceleration,2024】。
图1. 我们的KV缓存量化方案Kitty的图示。该图显示了解码阶段,其中只有当前的查询向量(单个token)参与计算,而KV缓存中所有先前存储的键和值向量都被重用。Key缓存分为三部分:(i) Sink(初始token,保持FP16),(ii) Q-Buffer(量化缓冲区,在形成量化组之前临时存储FP16的KV),以及(iii) 量化信道(大多数信道量化到INT2以实现最大压缩,一小部分信道保留在INT4以保持准确性)。Value缓存采用带滑动窗口的逐token量化,其中Sink(初始token)和最近的Local(本地token)都保留在FP16。这里,L表示序列长度,D表示头大小;S是保留的sink token数量,G是量化组大小,R是本地窗口大小。默认配置为:S = 32, R = 128, G = 128,这在准确性和内存节省之间提供了良好的平衡。
现有方法的局限性 尽管有所进展,但现有解决方案在将KV缓存精度推至4位以下时,难以维持准确率。Liu等人【索引16,Quantization hurts reasoning? an empirical study on quantized reasoning models,2025】的研究表明,使用最先进的量化算法将KV缓存量化到3位已经会导致显著的准确率下降【索引11,Kvquant: Towards 10 million context length llm inference with kv cache quantization,2024】、【索引2,Quarot: Outlier-free 4-bit inference in rotated llms,2024】。因此,2位精度应更具挑战性。
2-bit量化导致显著精度下降的实证观察 我们凭经验观察到,虽然4-bit KV缓存量化可以与FP16紧密匹配,但将其进一步降低到2-bit(INT2)会在推理和生成基准测试中导致严重的准确率下降。
定量分析 如表1所示,将Qwen3-8B和LLaMA3-8B的KV缓存精度从16位降至4位几乎不会导致准确率下降。然而,进一步将缓存降至2位会使平均准确率分别下降-15.23和-10.15。这些结果表明,简单地将现有量化方法应用于KV缓存是不够的,并引发了一个问题:更精细的技术能否减轻这些损失,并为长上下文LLM推理启用实用的低比特KV缓存量化?
Table 1. 低比特KV缓存量化在Qwen3-8B (Yang et al., 2025) 和 LLaMA3-8B (Dubey et al., 2024) 上的准确率下降情况。虽然4-bit KIVI (Liu et al., 2024) 保持了准确率,但2-bit KIVI 显示出显著下降。注意:“MATH”和“GPQA”分别表示MATH-Algebra和GPQA-Diamond子集。
保留初始Token为全精度 序列中的初始token有时被称为注意力池(attention sinks),正如StreamingLLM【索引33,Efficient streaming language models with attention sinks,2024b】所指出的。这些token通常获得不成比例的高注意力权重,因此对最终输出有重大贡献。本文规定,以可忽略的开销将初始token保留为全精度有助于减轻准确率损失。我们实现了KIVI-K2V2,这是KIVI的一个算法变体,其中前32个token在KV缓存中额外以全精度保留。表2中的实证结果显示,KIVI-K2V2在各项基准测试中显著减轻了准确率下降。平均而言,KIVI-K2V2*在Qwen3-8B和LLaMA3-8B上分别比原始的KIVI-K2V2高出+8.28和+5.33。
提升Key缓存的精度 尽管保留sink-token条目的精度恢复了部分性能,但与FP16基线相比,该技术仍有待改进。为了进一步缩小差距并更好地理解Key和Value缓存的相对重要性,我们还评估了两个设计变体:KIVI-K2V4 和 KIVI-K4V2。表2的结果表明,提高Key缓存的精度比提高Value缓存的精度要有效得多。特别是,KIVI-K4V2 的性能优于KIVI-K2V4,并在多个基准测试中达到了接近FP16基线的准确率,这突显了提高Key缓存精度相对于Value缓存的重要性。
Table 2. 不同KV缓存方案在Qwen3-8B (Yang et al., 2025) 和 LLaMA3-8B (Dubey et al., 2024) 上的准确率对比。行为任务(从原始表格转置);列为方法。在每个模型块内,每行的最佳结果以粗体显示。注意:“MATH” = MATH-Algebra,“GPQA” = GPQA-Diamond。以星号(*)结尾的算法名称表示初始32个token以全精度保留。
探索选择性精度提升的可能性 在前一小节(§ 3.1)中,我们观察到提高Key缓存的精度可以有效减轻KV缓存量化带来的准确率下降。这引出了一个关键问题:与其提升整个Key缓存的精度,我们能否通过仅提高Key缓存一小部分的精度来获得相当的准确率增益?为了探索这一想法,我们首先对Key缓存张量进行视觉和统计分析,然后根据观察结果提出我们的关键创新。
观察1:Key缓存中的信道级模式 我们可视化了Qwen3-8B上多个层的Key缓存激活值的绝对值,发现Key缓存中不同信道的量级差异很大。如图2a所示,一部分信道始终表现出较高的量级。虽然该图仅展示了第10层的可视化结果,但在其他层也观察到了类似的模式。这些观察结果与先前的工作【索引17,Kivi: a tuning-free asymmetric 2-bit quantization for kv cache,2024,ICML'24】、【索引11,Kvquant: Towards 10 million context length llm inference with kv cache quantization,2024】一致,后者利用逐信道量化来减少量化误差。
信道级视角的启发 这些观察表明,Key缓存中的所有信道行为并不相同。一些信道似乎发挥着更强的影响力,而其他信道可能只有边际效应。这种不均衡的贡献激发了一种信道感知视角:通过有选择地为最具影响力的信道分配更高的精度,可能在保持模型准确率的同时,仍然减少整体KV缓存的内存占用。
观察2:信道级量化敏感度 为了解每个信道的重要性,我们分离了量化对每个单独信道的影响。具体来说,我们每次将单个信道量化到2位(INT2),同时保持其他信道不变,并计算其对注意力分数的影响:
量化敏感度的量化 我们计算原始注意力分数矩阵与其扰动后(量化一个信道后)的均方误差(MSE),并使用此MSE作为量化每个信道敏感度的指标,如图2b所示。由于这些模型中使用了分组查询注意力(GQA)【索引1,Gqa: Training generalized ´ multi-query transformer models from multi-head checkpoints,2023】,其中一个键头被多个查询头共享,因此我们用不同颜色分别报告了每个查询头的MSE。
敏感度分析结果 如图2b所示,不同信道对量化的敏感度差异巨大。一小部分信道在2位量化后总是给注意力分数带来更大的误差,这表明它们应该被保留在更高的精度,例如4位。同时,将其他信道量化到2位对注意力分数造成的误差较小,表明它们可以安全地被更积极地量化到2位。在其他层和模型上也观察到了类似的模式。这种差异突显了一个明确的机会:与其将所有信道统一量化到2位,有选择地保留一小部分敏感信道在更高精度,有望在保持推理准确率的同时减少KV缓存的大小。这一见解构成了以下关键创新的基础。
图2. Qwen3-8B第10层Key缓存的视觉和统计分析。(a) 第一个KV头的Key缓存量级可视化。分布不均,少数信道始终显示高量级,这启发了一种信道感知的量化方法。垂直轴表示激活量级,而水平轴跨越token和信道维度。(b) 在量化Key缓存的每个信道后,原始注意力分数矩阵与其扰动后的对应矩阵之间的均方误差(MSE)。由于分组查询注意力(Ainslie et al., 2023),共享相同Key缓存的不同Q头之间的模式是一致的。在其他层/模型上也观察到了类似的模式。
关键创新:提升关键信道的精度 受先前观察的启发,我们提出了信道级精度提升方法,即将关键信道保留在更高精度(INT4),而将其他信道积极量化到更低精度(INT2),以减少统一2位量化引入的失真并提高推理准确率。
信道重要性的轻量级启发式计算 为了应用我们的信道级精度提升方法,我们首先需要确定要提升的信道子集。在运行时直接计算每个信道的敏感度,如我们在图2b中所做的那样,成本高得令人望而却步,因为它需要为每个信道计算注意力分数矩阵并测量其与FP16基线的MSE。为了解决这个问题,我们使用可以在单次传递中计算的轻量级启发式方法来近似信道重要性。在这项工作中,我们使用基于量级的选择启发式来近似信道重要性。这种启发式方法基于两个直觉。首先,平均量级较大的信道更容易受到量化误差的影响。其次,这类信道往往对注意力分数有更大的影响。因此,我们将重要性分数定义为信道i在所有token上的平均量级:
$$s_{i}=\frac{1}{T}\sum_{t=1}^{T}|x_{i,t}|.$$信道选择与未来工作 需要注意的是,每个信道的分数需要在推理运行时计算。基于这些信道级分数,选择得分最高的K个信道,并以更高精度(例如INT4)存储到KV缓存中。我们只评估了一种启发式方法和一种基线用于信道选择。虽然这足以证明信道级精度提升的有效性以及启发式引导选择优于随机选择的必要性,但更具原则性或自适应性的策略可能会产生更强的鲁棒性准确率恢复。探索这些策略是一个自然的下一步,我们将其作为未来的工作。
Kitty算法的构成 我们的Kitty量化算法建立在KIVI【索引17,Kivi: a tuning-free asymmetric 2-bit quantization for kv cache,2024,ICML'24】的基础上,同时引入了两项关键增强功能,以在激进的低比特设置下进一步提高准确率。(1) 我们为Key和Value缓存都保留了初始token(Sink tokens)的全精度。(2) 更重要的是,我们将新颖的信道级精度提升机制应用于Kitty。图1展示了我们量化方案Kitty的整体设计,它将上述所有优化集成到一个统一的方案中。
准确率验证框架 为了评估该量化算法对准确率的影响,我们将Kitty量化方案集成到HuggingFace Transformers框架【索引12,Transformers library,2025】中,并构建了一个准确率模拟框架来验证Kitty对准确率的影响。
设计目标 为了支持我们提出的Kitty算法,我们为KV缓存设计了一种内存和计算高效的布局。鉴于KV分页【索引13,Efficient memory management for large language model serving with pagedattention,22023】对于可扩展的缓存管理和减少内存碎片至关重要,我们通过将每个量化组(例如,128个token)视为一个独立的页面来实现此功能。为了实现高效的运行时加载,每个Key缓存页面被分解为两个2位矩阵:一个密集矩阵和一个结构化稀疏矩阵,它们分开存储。这种设计简化了加载过程,因为两个矩阵现在都只包含2位元素。
图3. Kitty页面中心KV缓存布局的图示。
分页和量化的键值缓存 如图3所示,我们将KV缓存存储在GPU内存中的几个全精度缓冲区(即Sink、Q-Buffer、Local)和许多固定大小的内存页面中(灵感来自Paged Attention【索引13,Efficient memory management for large language model serving with pagedattention,22023】)。在图1所示的Kitty设计中,G个(例如128个)token在Key缓存中形成一个量化组,我们自然地选择G作为页面大小。因此,每个页面包含(D, G)个元素,其中D是每个头的大小。为了高效的内存管理,我们也将Value缓存分割成页面。
量化页面的动态密集-稀疏分解 在注意力计算期间,KV缓存的所有页面都必须从HBM中获取。然而,在我们的Kitty算法中,Key缓存中的不同信道以混合精度(2位或4位)进行量化,这使得页面加载过程变得复杂。为了应对这一挑战,受Wu等人【索引29,Zeroquant(4+2): Redefining llms quantization with a new fp6-centric strategy for diverse generative tasks,2023】和Xia等人【索引32,Quant-LLM: Accelerating the serving of large language models via FP6-Centric Algorithm-System Co-Design on modern GPUs,2024】的启发,我们将每个混合精度Key页面分解为两个具有统一2位精度的张量,如图3a所示。
密集张量的构成 我们将所有2位信道以及提升的4位信道的低两位存储在一个形状为(D, G)的密集张量Tensor_2bits中。
稀疏张量的构成 如图3a下部所示,提升信道的高两位形成一个2位稀疏张量。为提高内存效率,我们仅紧凑地存储该稀疏张量的非零信道,从而得到形状为(Dboost, G)的Tensor_High_2bits,其中Dboost表示提升到4位精度的信道数量。
重构索引 为了重构完整的张量,我们维护一个形状为(D,)的索引张量Boost_IDX_uint8,它定义了一个从1, 2, . . . , D到1, 2, . . . , Dboost的映射,将每个逻辑信道索引映射到其在GPU内存中的物理偏移量。对于未提升到INT4精度的信道,相应的索引被设置为一个哨兵值Dboost + 1。
挑战与解决方案 从我们的量化表示中重构一个FP16页面并非易事,特别是对于Key缓存,其中每个信道可能有不同的精度(2位或4位)。为了应对这一挑战,我们的Kitty系统利用了之前介绍的密集-稀疏分解和索引张量Boost_IDX_uint8,以在GPU上实现完全并行的页面重构。
Triton风格的伪代码 算法1展示了该反量化核的Triton风格伪代码,它直接在芯片上从量化的内存页面(图3a)重构FP16张量。在第8行,一个形状为(D, T)的uint8张量从HBM加载,其中每个字节编码了四个连续的2位值。因此,同一个字节在逻辑上被重复使用了四次,在片上内存中引入了有意的冗余。在第9行,通过位移和掩码操作提取每个打包字节的低两位。对于高两位,第10-11行根据提升掩码有条件地从HBM中获取额外数据。只有提升到INT4精度的信道才会从次要张量(Tensor_High_2bits)中读取,其中提升索引张量提供了密集张量和紧凑的提升子空间之间的直接地址映射。最后,该核结合低两位和高两位分量,应用逐信道缩放和零点校正,并在片上重构全精度的Kfp16页面。这种设计避免了发散的内存访问,并实现了所有信道的高效位级解包和并行重构。
# Algorithm 1 Dequantize KeyCache Page()
1: shifts = [0, 2, 4, 6, 0, ...]
2: Input: Quantized cache C, metadata M, boosted channel number Dboost
3: Output: Dequantized key page Kfp16 ∈ R^(D×T)
4: scale, zero_point ← LoadMeta(M)
5: boost_idx ← Load(C)
6: boost_mask ← (boost_idx ≤ Dboost)
7: # Xlow is an uint8 tensor with shape (D, T).
8: Xlow ← LoadLowBits(C)
9: Xlow ← (Xlow ≫ shifts) & 0x3
10: # Xhigh is an uint8 tensor with shape (D, T).
11: Xhigh ← LoadHighBits(C, boost_idx, boost_mask)
12: Xhigh ← (Xhigh ≫ shifts) & 0x3
13: # Combining and dequantization.
14: X ← Xlow | (Xhigh ≪ 2)
15: Kfp16 ← X ⊙ scale + zero_point
16: return Kfp16
流水线设计概述 在本小节中,我们介绍了我们新颖的注意力层执行流水线。FP16注意力的执行流水线很简单,即先更新KV缓存,然后计算注意力输出。然而,对于Kitty KV缓存,情况变得更加复杂,因为它需要跨多个异构内存组件进行协调管理,包括全精度缓冲区(Sink、Q-Buffer和Local)和存储在HBM中的量化页面。每个组件在平衡精度、内存效率和时间局部性方面都有不同的作用。
三阶段执行流水线 为了高效地同步这些异构组件,Kitty采用了一个轻量级的三阶段执行流水线。(1) 将新的KV向量插入静态缓冲区(例如Sink、Q-Buffer或Local);(2) 从我们的异构内存布局中加载KV缓存,用片上内存重构FP16 KV缓存并计算注意力输出;(3) (可选地)对来自全精度缓冲区的KV向量执行量化和打包。
步骤(1):以全精度插入新的KV向量 新生成的KV向量将在注意力计算之前插入到KV缓存中。根据每个缓冲区的当前状态,KV向量可能有多个目标。如果Sink未满,新的KV将直接插入其中。否则,新的Key张量将被插入到Q-Buffer中,由于步骤(3)的存在,该缓冲区保证不会满。对于Value缓存,我们也分配了一个Q-Buffer来存储待量化的张量(以较大粒度而不是每个解码步骤进行,以降低成本)。不同的是,如果Sink已满,新的Value张量将首先插入到Local缓冲区。然后,如果Local已满,Local中最旧的token将被驱逐并插入到Q-Buffer中。这是为了保证Value缓存有一个固定的局部窗口。
步骤(2):注意力计算 我们使用两个定制的Triton GPU核(qk核和sv核)和一个PyTorch操作符(softmax)实现了整体的注意力计算。首先用qk核计算注意力logits,它从GPU内存加载Query向量和(量化的)Key缓存,并进行矩阵乘法。然后,用PyTorch softmax操作符计算注意力分数(全精度)。最后,用sv核将注意力分数和(量化的)Value缓存相乘。我们将这些核的进一步优化(例如,将它们融合成一个【索引5,Flashat-´ tention: fast and memory-efficient exact attention with io-awareness,2022】)留作未来工作。
步骤(3):量化和打包 在注意力计算之后,一旦Q-Buffer已满,量化和打包过程将启动。在此过程中,Q-Buffer中的所有KV向量将使用我们定制的Triton核进行量化,并打包到一个页面中(如图3所述)。鉴于Q-Buffer的大小为G,并且每个解码步骤最多可以向Q-Buffer插入一个向量,因此该量化过程最多每G个解码步骤启动一次。通过这种方式,量化开销被有效摊销并变得可以忽略不计。
评测模型:
评测数据集:
硬件配置:
软件与评测框架:
lm-evaluation-harness (EleutherAI Team, 2025) 进行。实验内容:
本节旨在验证第三节提出的量化方案的有效性。实验比较了不同KV缓存量化方法在Qwen3和LLaMA3模型上的准确率。比较的方法包括:
实验结果与分析(见表3):
- KIVI-K2V2 在激进的2-bit量化下,模型准确率显著下降,尤其是在GSM8K和MATH-Algebra等推理任务上。
- KIVI-K2V2* 通过保留初始token,准确率得到显著恢复,验证了3.1节的发现。然而,与FP16基线相比仍有差距,特别是在复杂推理任务上。
- Kitty 和 Kitty-Pro 凭借信道级精度提升策略,有效弥补了剩余的准确率差距。Kitty在所有模型上都恢复了准确率,例如在Qwen3-8B上平均准确率比KIVI-K2V2高出5.17。Kitty-Pro通过增加提升信道的比例,在多个模型(如Qwen3-14B和LLaMA3.3-70B-Instruct)上达到了与FP16基线相当甚至略有超过的准确率。
- 结论:* 实验结果表明,信道级精度提升是恢复激进KV量化准确率损失的有效且可扩展的方法。Kitty-Pro在保持低比特KV缓存内存优势的同时,实现了与FP16相当的准确率。
Table 3. 不同KV缓存量化方法的基准测试结果。
注:K16V16表示KV缓存为16位精度,而KIVI-K2V2使用KIVI (Liu et al., 2024) 算法进行2位量化;星号(*)表示前32个token保持全精度。Kitty和Kitty-Pro分别保留12.5%和25%的信道为INT4。最大生成长度为4096个token。
实验内容:
为了检验方法在更长上下文推理下的鲁棒性,本节在AIME24和AIME25数据集上进行了评测,最大生成长度设置为32,768个token。
实验结果与分析(见表4):
- 在长上下文场景下,2-bit量化的准确率下降问题更加明显。KIVI-K2V2相较于FP16基线,平均准确率下降约13个点。
- KIVI-K2V2 缓解了部分下降,但仍落后FP16基线6-8个点。
- 相比之下,本文提出的Kitty实现了显著的准确率恢复,将平均差距缩小到仅3-4个点。
- 结论:* Kitty即使在32k token的上下文长度下也能保持稳定的准确率,显示出在长上下文推理设置下的强大鲁棒性。
Table 4. Qwen3-8B在AIME24和AIME25上的结果(最大生成长度:32k tokens)。这里的KIVI-KV2*表示KIVI的一个变体,其中初始token的KV未被量化。
实验内容:
本节通过改变提升到INT4的Key缓存信道比例(channel boost rate),在Qwen3-8B上对信道级精度提升策略进行消融研究,并与随机选择信道的基线进行比较。
实验结果与分析(见图4):
- 单调的准确率提升: 如图4所示,在GSM8K和MATH Algebra任务上,准确率几乎随着信道提升比例的增加而单调提高。当25%的信道被提升到4-bit精度时,准确率基本得到恢复,这与表3中Kitty-Pro达到与基线相当水平的结果一致。
- 选择启发式的必要性: 通过与随机选择基线(为每个信道分配随机重要性分数)的比较,验证了基于启发式(量级)的信道选择比随机选择带来了显著更大的收益。
- 结论: 提高信道提升率能持续改善准确率,并且使用基于量级的启发式方法来选择关键信道是至关重要的。
图4. 在Qwen3-8B上通过信道级精度提升实现的准确率恢复。在其他任务上也观察到类似趋势。
实验内容:
本节使用第四节介绍的原型推理引擎,对Kitty的端到端系统效率进行评估。实验在单个NVIDIA A100 GPU上对Qwen3-8B模型进行推理,生成长度为8192个token。比较对象包括:
- HF Dynamic FP16: Hugging Face默认的动态FP16 KV缓存。
- HF Static FP16: Hugging Face预分配的静态FP16 KV缓存。
- HF KIVI INT4: Hugging Face实现的INT4 KIVI量化。
- Kitty-Pro: 本文的系统实现。
实验通过增加批处理大小直至内存溢出(OOM),报告不同批大小下的峰值内存使用和token生成吞吐量。
实验结果与分析(见图5):
- 在相同的内存预算下,Kitty-Pro能够支持比FP16基线大8倍的批处理大小。
- 通过支持更大的批处理,Kitty-Pro实现了比FP16基线高2.1倍至4.1倍的推理吞吐量。
- 结论: Kitty的系统实现显著提升了推理效率。作者指出,由于GPU核是使用Triton实现的,如果使用更底层的语言(如CUDA)并进行更细粒度的优化,效率还有进一步提升的空间。
图5. Qwen3-8B生成8192个token时的内存使用和吞吐量比较。Kitty通过支持更大的批处理大小可以实现更高的吞吐量。
本文首先指出,尽管4-bit KV缓存量化基本能保持模型准确率,但现有的2-bit量化方法会在推理密集型任务上导致严重的性能下降。保留初始token为全精度(FP16)有助于缓解这一问题,但无法完全弥合与FP16基线的差距。
本文的核心创新——信道级精度提升方法,进一步缩小了这一差距。研究发现,仅将12.5%至25%的关键Key缓存信道提升至更高精度,就足以恢复大部分的准确率损失。基于此创新,本文提出了一个准确且内存高效的量化框架 Kitty。
此外,本文还提出了相应的系统设计,以支持使用 Kitty 进行端到端推理。端到端评估结果显示,本文的2-bit推理系统能够显著降低GPU内存消耗,使推理时的批处理大小增加8倍,从而在相同的内存预算下,相较于FP16基线,实现了2.1倍至4.1倍的推理吞吐量提升。