KVarN: 方差归一化的KV缓存量化减轻推理任务中的误差累积

摘要

测试时缩放是一种提升大语言模型推理性能的有效方法，但在长序列解码过程中，由于KV缓存不断增长，内存会成为瓶颈。KV缓存量化有助于缓解这一问题，但现有方法通常在预填充式设置下进行评估，而误差在自回归解码中的表现有所不同。我们发现，在后一种情况下，量化误差会随时间步累积，主要源于不正确的token尺度。为此，我们提出KVarN——一种免校准的KV缓存量化器，它先对K和V矩阵进行哈达玛变换，再沿两个轴施加双尺度方差归一化。实验表明，这种组合能够修正异常的token尺度误差，并显著减少相较于现有基线的误差累积。KVarN在包括MATH500、AIME24和HumanEval在内的生成式基准测试中，以2位精度实现了KV缓存量化的新最先进水平。KVarN方法的vLLM实现可参见：https://github.com/huawei-csl/KVarN

English

Test-time scaling is a powerful approach to obtain better reasoning in large language models, but it becomes memory-bottlenecked during long-horizon decoding, as the KV-cache grows. KV-cache quantization can help improve this, but current methods are evaluated under prefill-like settings and errors behave differently under autoregressive decoding. We show that in the latter regime, quantization errors accumulate across timesteps, driven primarily by incorrect token scales. We introduce KVarN, a calibration-free KV-cache quantizer that applies a Hadamard rotation followed by a dual-scaling variance normalization across both axes of the K and V matrices. We find that this combination fixes outlying token-scale errors and substantially reduces error accumulation over existing baselines. KVarN establishes a new state-of-theart for KV-cache quantization on generative benchmarks, including MATH500, AIME24 and HumanEval, at 2-bit precision. A vLLM implementation of the KVarN method is available at https://github.com/huawei-csl/KVarN