PM-KVQ: Progressive Mixed-Precision KV Cache Quantization for Long-CoT LLMs

发表时间: 2026-04 · arXiv:2505.18610 (ICLR 2026)

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

Tengxuan Liu, Shiyao Li, Jiayi Yang, Tianchen Zhao, Feng Zhou, Xiaohui Song, Guohao Dai, Shengen Yan, Huazhong Yang, Yu Wang
Tsinghua University, Infinigence-AI, Columbia University, OPPO AI Center, Beijing, China, Shanghai Jiaotong University

速读

一句话结论 本文提出了一种面向长思维链大语言模型的渐进式混合精度 KV Cache 量化方法 PM-KVQ，通过渐进式降比特、分块显存分配和带位置插值的校准策略，在相同显存预算下将推理基准测试准确率提升最高 8%，并实现相较 16-bit 模型 2.73 至 5.18 倍的吞吐量提升。

要解决什么问题 具备长思维链（long-CoT）推理能力的大语言模型在生成多达 128K 长度的回复时，会产生 10GB 到 100GB 的海量 KV Cache 显存开销。现有的训练后量化方法在短上下文场景表现良好，但直接用于长思维链模型会导致严重的性能崩塌，核心卡点有两处。第一是累积误差过大：现有方法在解码的每一步都直接将 KV Cache 量化到极低比特（如 2-bit），随着生成 token 数量的激增，量化误差不断累积，且在显存未满时白白浪费了可用预算。第二是短上下文校准失效：模型通过旋转位置编码（RoPE）将位置信息注入 Key Cache，导致不同通道呈现周期性变化。部分低频通道的周期长达 32K 个 token 以上，如果仅用 2K 长度的短序列进行校准，根本无法捕捉这些低频通道在长序列下的真实数据分布，进而引发巨大的量化误差。

怎么做的 PM-KVQ 的内核包含三个关键部件，分别从显存利用率和校准分布两方面绕开上述卡点。首先是渐进式量化（Progressive Quantization），其思路是“显存不满不降级”。在推理初期，KV Cache 以 16-bit 格式存储以规避早期累积误差；当显存预算耗尽时，再通过移位操作逐步将已缓存的张量位宽减半（16 降至 8，再降至 4 和 2）。为了在降比特时不引入反量化再量化的开销，作者设计了等效右移公式，仅通过整数加法和移位即可完成 $2b$ 比特到 $b$ 比特的转换，同时保持零点不变并自适应调整缩放因子：

$$ \mathbf{X}_b = \left( (2^{2b} - 2^b + 1) (\mathbf{X}_{2b} + 2^{b-1}) \right) >> 3b $$


其次是分块显存分配（Block-wise Memory Allocation）。由于 Transformer 较深层的网络通常对量化更敏感，采用统一位宽无法实现全局最优。该方法利用一阶泰勒展开评估各层 Key 和 Value 张量对量化扰动的敏感度 $s_{i,b}$，并将位宽分配建模为严格显存预算 $\mathcal{M}$ 下的整数规划问题，交由求解器在几秒内为敏感层分配更高的目标位宽：

$$ \underset{x_{i,b}}{\arg\min} \sum_i^N \sum_b x_{i,b} \cdot s_{i,b} $$
最后是带位置插值的校准（Calibration with Positional Interpolation）。为了在不增加注意力计算和显存开销的前提下捕捉长序列分布，该方法在短校准数据（如 2K 长度）的 RoPE 旋转矩阵中引入位置缩放因子 $s$。例如令 $s=4$，即可让 2K 的校准数据等效模拟出 8K 长度的位置信息分布，从而准确计算出用于平滑异常值的通道重参数化因子。

效果如何 实验在 A100 和 4090 GPU 上进行，评估了 7B 到 70B 规模的 DeepSeek-R1-Distill 系列以及 QwQ-32B 模型，最大输出长度设为 32K。测试任务涵盖 AIME-2024、AIME-2025、CMIMC-2024 数学基准和 LiveCodeBench 代码生成。对比基线包括代表统一位宽路线的 RotateKV 和 KIVI（保留首尾 token 为高精度），以及代表混合精度路线的 MiKV。在极低比特（2-bit 或 4-bit）设置下，RotateKV 和 MiKV 在长思维链模型上经常完全失效（输出乱码或得分为零）。相比最强基线 KIVI，PM-KVQ 在 10B 到 32B 模型上将 pass@1 准确率提升了最高 15.00%，多数投票准确率提升最高 17.78%。在吞吐量方面，PM-KVQ 达到了原始 16-bit 模型的 2.73 至 5.18 倍；由于推理前期使用了更高位宽，其吞吐量比纯 2-bit 的 KIVI 略低 2.45% 到 16.18%，但换取了极大的准确率收益。该方法的局限在于，过激的位置插值（如 $s=16$）会导致性能回落，且目前尚未与多头潜在注意力（MLA）等非 GQA 架构以及底层推理引擎进行系统级深度集成。

主要贡献

最近，通过长思维链（Chain-of-Thought, CoT）技术开发具备推理能力的大型语言模型（LLMs）取得了显著进展。然而，这种长CoT推理过程由于庞大的键值（Key-Value, KV）缓存内存开销，带来了巨大的内存负担。训练后KV缓存量化作为一种极具前景的压缩技术，已在短上下文场景中得到广泛研究。但是，直接将现有方法应用于长CoT LLMs会导致显著的性能下降，原因主要有两个方面：
（1）累积误差大：现有方法未能充分利用可用内存，它们在每个解码步骤中直接对KV缓存进行量化，导致较大的累积量化误差。
（2）短上下文校准失效：由于旋转位置编码（Rotary Positional Embedding, RoPE）的存在，在校准过程中使用短上下文数据无法捕捉到Key缓存中低频通道的数据分布，从而导致性能损失。

为了解决上述问题，本文针对长CoT LLMs提出了渐进式混合精度KV缓存量化（Progressive Mixed-Precision KV Cache Quantization, PM-KVQ）方法，主要创新点包括：
1. 为了减少累积误差，设计了一种渐进式量化策略，在每个Transformer块中逐步降低KV缓存的位宽。同时，提出了分块内存分配（block-wise memory allocation）技术，为更敏感的Transformer块分配更高的位宽。
2. 为了在不增加额外开销的情况下增加校准长度，提出了一种结合位置插值（positional interpolation）的新校准策略，利用带有位置插值的短校准数据来近似长上下文数据的分布。

广泛的实验表明，在7B–70B的长CoT LLMs上，在相同的内存预算下，PM-KVQ在推理基准测试中的性能比SOTA基线提升了最高$8\%$，并且相比原始16位LLMs实现了$2.73 - 5.18\times$的吞吐量提升。

表1展示了长CoT LLMs的内存开销。批处理大小为16，上下文长度为32K。可以看出KV缓存的内存开销远超模型权重。

Model Weights (GB) KV Cache (GB)
DeepSeek-LLaMA-8B 16 64
DeepSeek-Qwen-32B 64 128
DeepSeek-LLaMA-70B 140 160

背景知识与设计原则

长CoT大型语言模型的内存瓶颈。长CoT LLMs旨在增强复杂任务（如数学证明、科学推理和多跳问答）的多步推理能力。虽然长CoT能显著提升模型性能，但它引入了极多的解码Token（例如每个请求超过32K个Token）以及巨大的GPU内存开销。尽管采用了高效的注意力设计（如MQA、GQA和MLA），长CoT LLMs中KV缓存的内存开销仍然极其庞大，通常超过模型权重本身。因此，减少KV缓存的内存开销对于支持大批处理大小和长上下文需求至关重要。

训练后KV缓存量化的现状与局限。为了缓解长推理上下文带来的巨大内存开销，训练后KV缓存量化成为一种有前景的高效推理技术。现有方法通常对KV缓存应用非对称均匀量化：

$$ \mathbf{X}_{\mathrm{asym}} = \left\lfloor \frac{\mathbf{X}_{\mathrm{BF16}} - Z}{S_{\mathrm{asym}}} \right\rceil, $$

$$ S_{\mathrm{asym}} = \frac{\max(\mathbf{X}_{\mathrm{BF16}}) - Z}{2^b - 1}, $$
其中$\mathbf{X}_{\mathrm{BF16}}$表示16位BF16格式的Key或Value张量，$\mathbf{X}_{\mathrm{asym}}$表示整数张量，$S_{\mathrm{asym}}$和$Z = \min(\mathbf{X}_{\mathrm{BF16}})$分别表示缩放因子和零点，$b$表示量化位宽，$\lfloor \cdot \rceil$表示舍入函数。现有研究（如MKLV、WKVQuant、SKVQ、MiKV、KIVI等）发现Key和Value的敏感度不同，且Key张量中的异常值通常集中在特定通道中。本文采用了先前工作中的有效策略，例如将第一个Token存储为INT16格式，并对最近的Token使用滑动窗口。为了进一步减少量化误差，本文提出了渐进式量化和分块内存分配机制。

方法概览。 (a) 渐进式量化：我们逐步缩小 KV 缓存的位宽以充分利用内存预算。 (b) 块级内存分配：我们为具有更高敏感度的 Transformer 块分配更高的位宽。 (c) 使用位置插值进行校准，以使用短上下文数据近似长上下文数据的分布。
方法概览。 (a) 渐进式量化：我们逐步缩小 KV 缓存的位宽以充分利用内存预算。 (b) 块级内存分配：我们为具有更高敏感度的 Transformer 块分配更高的位宽。 (c) 使用位置插值进行校准，以使用短上下文数据近似长上下文数据的分布。

方法细节

现有量化方法对内存预算的利用不足。现有的训练后KV缓存量化方法在每个解码步骤都进行量化，这会导致巨大的累积误差。虽然带有高精度缓存的滑动窗口可以缓解这一问题，但在长CoT任务中，极低的位宽（例如2-bit）仍然会导致严重的精度损失。作者指出，现有的KV缓存量化方法未能充分利用目标硬件的内存预算，从而错失了减少累积误差的机会。如Fig 1(a)左侧所示，SOTA方法在每个解码步骤都存储2-bit的KV缓存，当内存预算未被完全占用时，这造成了大量的内存浪费。

渐进式量化策略（Progressive Quantization）。为了解决上述内存浪费问题，作者提出了一种渐进式量化策略，通过逐步缩小KV缓存的位宽来充分利用内存资源，从而显著减少累积量化误差。对于每个Transformer块，作者使用“Fbit”来表示渐进式量化过程的最终位宽。在这种情况下，可以根据目标长CoT LLM的最大上下文长度轻松计算出每个Transformer块的内存预算。如Fig 1(a)右侧所示，此示例中的Fbit为2-bit，最大上下文长度为32K。在生成过程中，首先将KV缓存以16位格式存储，以缓解较大的累积量化误差。一旦内存预算被完全占用，便应用位宽缩小策略（bit-width shrinking strategy），通过逐步降低现有KV缓存的位宽来容纳更多的Token。具体而言，量化位宽采用2的幂次方，按照16位、8位、4位和2位的顺序逐步递减。

等效右移策略（Equivalent Right Shift）。在位宽缩小策略中，作者设计了一种“等效右移”策略，该策略在数学上等效于先将$2b$位的KV缓存反量化，然后再将其量化为$b$位。这里的$b$可以是8、4或2，分别对应将KV缓存从16位缩小到8位、8位缩小到4位，以及4位缩小到2位。具体来说，通过使用整数加法和移位操作来公式化位宽缩小策略：

$$ \mathbf{X}_b = \left( (2^{2b} - 2^b + 1)(\mathbf{X}_{2b} + 2^{b-1}) \right) >> 3b, $$


其中$\mathbf{X}_b$和$\mathbf{X}_{2b}$分别表示$b$位和$2b$位张量。在转换过程中，保持零点不变（$Z_b = Z_{2b}$），并将缩放因子增加至$S_b = (2^b + 1)S_{2b}$，以保留数据分布的动态范围。

统一位宽分配的局限性。现有的KV缓存量化方法通常在所有Transformer块中应用统一的位宽，这可能无法充分利用目标硬件的内存资源。如Fig 1(b)左侧所示，在该示例中，目标硬件有足够的内存将KV缓存统一存储为2-Fbit格式，但这留下了一部分被浪费的内存。然而，如果切换到统一的4-Fbit格式，可能会超出内存限制并触发内存溢出（OOM）错误。因此，在具有不同内存资源的不同场景中，使用统一位宽可能无法充分利用可用内存。

基于一阶泰勒近似的敏感度评估。为了在不同场景下充分利用内存资源以获得更好的性能，作者提出了一种分块内存分配策略，为更敏感的块分配更高的位宽。受现有混合精度量化方法的启发，作者采用一阶泰勒近似来估计模型输出对Key缓存和Value缓存扰动的敏感度。以Key缓存为例：

$$ \mathcal{L}(Q_b(\mathbf{K}_i)) \approx \mathcal{L}(\mathbf{K}) + \mathbf{G}_{\mathbf{K}_i} \odot (\mathbf{K}_i - Q_b(\mathbf{K}_i)), $$


其中$\mathcal{L}$是损失函数，$i$表示第$i$个Transformer块，$\mathbf{K}_i$是Key缓存，$Q_b(\cdot)$是$b$位量化函数，$\mathbf{G}_{\mathbf{K}_i}$是损失函数关于$\mathbf{K}_i$的梯度，$\odot$是逐元素乘法运算符。Value缓存遵循类似的公式。

最小化量化敏感度目标。为了最小化每个Transformer块中KV缓存量化的影响，目标是最小化以下敏感度项：

$$ s_{i,b} = \|\mathbf{G}_{\mathbf{K}_i} \odot (\mathbf{K}_i - Q_b(\mathbf{K}_i))\|_1 + \|\mathbf{G}_{\mathbf{V}_i} \odot (\mathbf{V}_i - Q_b(\mathbf{V}_i))\|_1, $$


其中$s_{i,b}$表示第$i$个Transformer块中的KV缓存对$b$位量化的敏感度。

整数规划问题建模（Integer Programming）。考虑到所有Transformer块的敏感度，目标是在给定内存预算的情况下，为每个块分配适当的位宽，以最小化对损失函数的影响。为此，作者将分块位宽分配形式化为以下整数规划问题：

$$\begin{aligned} \begin{array}{rl} \underset{x_{i,b}}{\arg\min} & \sum_i^N \sum_b x_{i,b} \cdot s_{i,b}, \\ & \sum_b x_{i,b} = 1, \\ & \sum_i^N \sum_b x_{i,b} \cdot (Mem(Q_b(\mathbf{K}_i)) + Mem(Q_b(\mathbf{V}_i))) \leq \mathcal{M}, \\ & x_{i,b} \in \{0, 1\}, b \in B, \end{array} \end{aligned}$$


其中$N$是Transformer块的数量，$Mem(\cdot)$是计算量化后KV缓存内存使用量的函数，$\mathcal{M}$是所有Transformer块KV缓存的内存预算，$x_{i,b}$是指示第$i$个块选择位宽$b$的one-hot向量，$B$是可选的位宽集合。提出的整数规划问题可以在几秒钟内由CVXPY【索引编号，CVXPY: A Python-embedded modeling language for convex optimization+2016+Journal of Machine Learning Research】有效求解。

通道级重参数化与短序列校准的局限。先前的研究观察到LLMs的Key缓存中某些通道存在异常值，这显著增加了量化误差。现有的方法如QServe引入了通道级重参数化方法，将Key张量中的异常值转移到Query张量中：

$$ \mathbf{P} = (\mathbf{Q}\mathbf{\Lambda}) \cdot Q\left((\mathbf{K}\mathbf{\Lambda}^{-1})^T\right), \mathbf{\Lambda} = diag(\lambda_i), $$


其中$i$是通道索引，$\lambda_i$是第$i$个通道的重参数化因子。通常，$\lambda_i$使用典型长度为512个Token的短序列数据集进行校准，校准过程如下：

$$ \lambda_i = \left(\max_m |K_{m,i}|\right)^\alpha, $$
其中$m$是Token位置索引，$\alpha$是调节异常值转移强度的参数。然而，将上述技术应用于长CoT LLMs时，由于使用了RoPE将位置信息注入Key缓存，使得不同通道呈现周期性变化：
$$\begin{aligned} \left[ \begin{array}{c} \widetilde{K}_{m,i} \\ \widetilde{K}_{m,i+\frac{d}{2}} \end{array} \right] = \sqrt{K_{m,i}^2 + K_{m,i+\frac{d}{2}}^2} \left[ \sin(m\theta_i + \varphi) \right], \end{aligned}$$
由于$\theta_i = \theta^{-2i/d}$随$i$增加而减小，在索引接近$d/2$和$d$的通道中，正弦曲线的频率极低（例如在DeepSeek-R1-Distill-Qwen-7B中周期长达54,410个Token）。因此，使用短序列校准无法获得能完全反映这种类正弦数据分布的重参数化因子（如Fig 1(c)顶部所示）。

结合位置插值的校准策略（Calibration with Positional Interpolation）。直接增加校准数据的长度会因自注意力算子的$O(N^2)$复杂度而显著增加延迟和内存成本。相反，作者提出利用位置插值【索引编号，Extending context window of large language models via positional interpolation+2023+arXiv】将长上下文位置信息嵌入到短校准数据中。具体而言，在RoPE的旋转矩阵中，将位置缩放因子$s$乘以位置索引$m$进行位置插值：

$$\begin{aligned} \left[ \begin{array}{c} \widetilde{K}_{m,i} \\ \widetilde{K}_{m,i+\frac{d}{2}} \end{array} \right] = \sqrt{K_{m,i}^2 + K_{m,i+\frac{d}{2}}^2} \left[ \cos(s \cdot m\theta_i + \varphi) \right]. \end{aligned}$$


如Fig 1(c)底部所示，通过应用位置插值，可以在不增加额外计算和内存开销的情况下，将最大位置索引增加$s$倍。

方法整体执行流程。PM-KVQ结合了上述三种技术：（1）在推理过程之前，首先基于校准数据集分析每个Transformer块的敏感度，并求解整数规划问题为每个块设置合适的Fbit。接着，使用带有位置插值的校准数据集应用通道级重参数化技术。（2）在推理过程中，对KV缓存应用渐进式量化，将其位宽从16位逐步降低到分配的Fbit。

实验环境

实验结果

小模型（小于10B）的推理性能。如原文Table 2所示，作者将PM-KVQ与RotateKV、MiKV和KIVI进行了比较。在2-bit的DeepSeek-R1-Distill-Qwen-7B上，应用RotateKV或MiKV会导致模型无法生成有意义的回复；SOTA方法KIVI也遭受了高达$9\%$的性能损失。而PM-KVQ在应用统一Fbit（Batch Size=40）时，比KIVI高出最高$8\%$。当Batch Size减小到32时，PM-KVQ利用分块内存分配策略更好地利用了额外的内存，获得了最高$0.84\%$的额外性能提升。对于4-bit的DeepSeek-R1-Distill-LLaMA-8B，PM-KVQ在AIME-2024上超越SOTA方法最高$6.5\%$，甚至在数学基准测试上取得了比原始LLM更好的性能。此外，PM-KVQ的平均投票准确率比KIVI高出最高$15.56\%$，证明了其更强的稳定性。

中等模型（10B到32B）的推理性能。由于MiKV和RotateKV在2-bit下失效，作者仅将PM-KVQ与KIVI进行了比较（见Table 2）。PM-KVQ同样展现出优越的性能，在不同LLMs上的平均pass@1和投票准确率分别提升了最高$15.00\%$和$17.78\%$。特别是在DeepSeek-R1-Distill-Qwen-14B上，KIVI在CMIMC-2024上导致了$21.87\%$的性能下降，而PM-KVQ在Batch Size为16和12时的性能下降幅度极小，仅为$1.87\%$和$2.91\%$。

大型模型（70B级别）的推理性能。作者在AIME-2024上评估了2-bit的DeepSeek-R1-Distill-LLaMA-70B模型。原始16位模型的pass@1为$69.14\%$。使用KIVI量化为2-bit时，pass@1显著下降至$51.88\%$。相比之下，PM-KVQ在Batch Size为12和16时，使2-bit模型达到了$64.79\%$的高pass@1，超越KIVI基线$12.91\%$。

补充细节

推理效率分析（Efficiency Analysis）。作者在A100-80G GPU上评估了7B和32B长CoT LLMs，比较了PM-KVQ（Fbit=2）、原始16位LLMs和2-bit KIVI基线的吞吐量。为了充分利用显存，量化模型采用了比原始模型更大的批处理大小。如原文Table 3所示，PM-KVQ相比原始16位LLMs实现了$2.73 - 5.18\times$的吞吐量提升。与KIVI相比，PM-KVQ的吞吐量略微降低了$2.45 - 16.18\%$（主要是因为在推理早期使用了更高的位宽），但获得了$10.57 - 23.48\%$的显著相对精度提升。此外，位宽缩小的开销可以忽略不计。预推理过程（分块内存分配和校准）均在一小时内完成。

位宽缩小策略的消融实验。作者比较了三种将KV缓存从$2b$位缩小到$b$位的策略：（1）直接右移（Direct Right Shift）：仅保留高$b$位，调整缩放因子以补偿幅度下降；（2）修改右移（Modified Right Shift）：调整零点和缩放因子以映射均值；（3）等效右移（Equivalent Right Shift，本文采用）：等效于反量化后再量化。如原文Table 4和Fig 2证实，直接右移和修改右移策略导致了显著的性能下降（pass@1分别下降$32.09\%$和$15.42\%$），而等效右移策略表现出显著的改进，维持了无损的投票准确率。
当位宽从 4 位缩小到 2 位时，不同的位宽缩小策略。

不同Transformer块的敏感度分析。对于参数量小于10B的模型，如Fig 3所示，更深的块往往对量化更敏感，并获得了更大的KV缓存内存预算。此外，在DeepSeek-R1-Distill-Qwen-7B模型中，第一个块比其他浅层块更敏感。提出的内存分配策略准确捕捉到了这一特征，并相应地为第一个块分配了更高的内存预算。
不同 Transformer 块中 KV 缓存量化的敏感度。不同的颜色代表不同的内存预算。

位置插值的消融实验。如原文Table 5所示，当校准序列长度设为2,048时，应用$s=4$的位置插值比不使用位置插值的pass@1提高了$1.66\%$，达到了与使用8,192个Token校准序列相当的精度。但当$s$增加到16时，位置插值导致了性能下降，表明过于激进的缩放会引起精度损失。

附录

不同Transformer块敏感度的进一步分析（C.1）。对于参数量超过10B的模型，如Fig 4所示，较深层块的KV缓存依然比较浅层块更敏感。同时，在基于Qwen的模型中，第一个块表现出极大的敏感度（在前15个块中最大），而这一现象在基于LLaMA的模型中未被观察到。
不同 Transformer 块中 KV 缓存量化的敏感度。不同的颜色代表不同的内存预算。

短生成上下文任务的性能（C.2）。为了验证PM-KVQ对短上下文任务的泛化能力，作者在IFEval基准上进行了评估。相比于推理基准（平均输出长度13,904个Token），IFEval的输出较短（平均1,182个Token）。结果表明，尽管PM-KVQ并非专为短输出场景设计，但它依然超越了KIVI，并达到了与原始16位模型相当的准确率。

与不同位宽KIVI的比较（C.3）。在相同的GPU和Batch Size约束下，如果KIVI采用更高的量化位宽，它将消耗更多内存并在达到最大输出长度前耗尽预算。作者评估了不同位宽的KIVI，当其耗尽内存时截断输出。结果表明，PM-KVQ依然比KIVI高出$0.83\% - 18.33\%$。

与混合精度量化基线的比较（C.4）。作者在相同的GPU和Batch Size设置下，将PM-KVQ与KVTuner进行了比较。结果显示，PM-KVQ的pass@1超越了KVTuner $2.71 - 6.04\%$。

跨不同硬件配置的性能（C.5）。当目标GPU显存容量增加时，PM-KVQ可以通过增加Batch Size来维持每个请求的内存分配，或者为每一层分配更高的Fbit。实验证实，在不同的Fbit设置下，PM-KVQ始终能达到与原始LLM相当的精度。

与稀疏注意力的集成（C.6）。渐进式量化和分块内存分配可以自然地应用于稀疏注意力机制。作者将PM-KVQ与QuestAttention结合，结果表明该方法不会降低QuestAttention的pass@1性能。

预推理过程的效率分析（C.7）。相比于QServe，PM-KVQ利用位置插值将校准序列长度从8,192减少到2,048个Token，使校准时间大幅减少了最高$77.21\%$。额外的分块内存分配过程占预推理时间的$22.50 - 23.53\%$。

渐进式量化在不同响应长度下的优势（C.8）。长CoT任务的响应长度差异巨大。对于较短的响应（如<16K Token），PM-KVQ在整个解码过程中保持比Fbit更高的位宽，从而减少了累积误差。对于较长的响应（如>16K Token），尽管KV缓存最终会缩小至Fbit，但前16K个Token是使用高精度缓存生成的，模型因此从更准确的早期生成中受益。

等效右移的数学证明（D）。
定理与证明目标。定理D.1证明了等效右移公式：给定16位浮点张量$\mathbf{X}_{\mathrm{BF16}}$，其$2b$位和$b$位量化张量分别记为$\mathbf{X}_{2b}$和$\mathbf{X}_b$，则$\mathbf{X}_b = \left( (2^{2b} - 2^b + 1)(\mathbf{X}_{2b} + 2^{b-1}) \right) >> 3b$成立。
缩放因子与量化定义。设零点$Z_{2b} = Z_b = Z$。根据量化公式，缩放因子满足$S_b = (2^b + 1)S_{2b}$。定义中间变量$\widetilde{\mathbf{X}}_{2b} = \frac{\mathbf{X}_{\mathrm{BF16}} - Z}{S_{2b}}$和$\widetilde{\mathbf{X}}_b = \frac{\mathbf{X}_{\mathrm{BF16}} - Z}{S_b}$，则量化张量为$\mathbf{X}_{2b} = \left\lfloor \widetilde{\mathbf{X}}_{2b} \right\rceil$且$\mathbf{X}_b = \left\lfloor \widetilde{\mathbf{X}}_b \right\rceil$，且有$\widetilde{\mathbf{X}}_{2b} = (2^b + 1)\widetilde{\mathbf{X}}_b$。
舍入与不等式推导。根据舍入定义，得到不等式$\frac{\mathbf{X}_{2b} - 1/2}{2^b + 1} \leq \widetilde{\mathbf{X}}_b < \frac{\mathbf{X}_{2b} + 1/2}{2^b + 1}$。对$\mathbf{X}_{2b}$除以$2^b+1$进行带余除法，得到商$q$和余数$r$（$0 \leq q \leq 2^b - 1$，$0 \leq r \leq 2^b$）。
分情况讨论余数。
情况1：当$0 \leq r \leq 2^{b-1}$时，推导出$q - 1/2 < \widetilde{\mathbf{X}}_b < q + 1/2$，舍入得到$\mathbf{X}_b = q$。同时证明了右移公式的计算结果也等于$q$。
情况2：当$2^{b-1} + 1 \leq r \leq 2^b$时，推导出$q + 1/2 \leq \widetilde{\mathbf{X}}_b < q + 1$，舍入得到$\mathbf{X}_b = q + 1$。同时证明了右移公式的计算结果也等于$q + 1$。两种情况均证明了该位移操作等效于反量化并重新量化。

结论

本文提出了一种专为长CoT LLMs设计的训练后KV缓存量化方法——渐进式混合精度KV缓存量化（PM-KVQ）。为了减少统一位宽量化引起的大量累积误差，设计了渐进式量化和分块内存分配技术。为了在不产生额外开销的情况下增加有效校准长度，提出了一种带有位置插值的新校准策略。广泛的实验和消融研究证明了PM-KVQ及其各项技术的有效性。总体而言，PM-KVQ在推理相关的数学和编码基准测试中显著优于SOTA基线（最高提升$8\%$），并实现了原始16位LLMs $2.73 - 5.18\times$的吞吐量。未来工作将探索该方法与MLA等其他注意力机制、系统级优化技术和推理引擎的结合。

参考文献引用汇总

本文在方法设计和分析中引用了以下关键文献：
* 在描述分块位宽分配的整数规划问题求解时，引用了【Diamond & Boyd, 2016，CVXPY: A Python-embedded modeling language for convex optimization+2016+Journal of Machine Learning Research】，指出该问题可以通过CVXPY在几秒钟内有效求解。
* 在描述长上下文校准策略时，引用了【Chen et al., 2023，Extending context window of large language models via positional interpolation+2023+arXiv】，提出利用位置插值将长上下文位置信息嵌入到短校准数据中。
* 在讨论现有通道级重参数化技术时，引用了【Lin* et al., 2024，QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving+2024+arXiv】，说明了如何通过校准转移Key张量中的异常值。
* 在解释短序列校准失效原因时，引用了【Su et al., 2024，Roformer: Enhanced transformer with rotary position embedding+2024+Neurocomputing】，指出了RoPE将位置信息注入Key缓存后引入的周期性变化。