Tailoring the Quantization Space for 1-Bit KV Cache Compression

📄 arXiv: 2610.03027v1 📥 PDF

作者: Minsoo Cheong, Donghyun Son, Sungjoo Yoo

分类: cs.LG, cs.AI, cs.CL

发布日期: 2026-10-02

💡 一句话要点

提出TaSQ以解决1位KV缓存压缩中的性能瓶颈问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 键值缓存 向量量化 模型压缩 长上下文推理 深度学习优化 内存管理 高效推理

📋 核心要点

  1. 现有的VQ方法在1位压缩下性能显著下降，难以有效利用有限的质心表示更大通道组。
  2. 本文提出TaSQ，通过查询引导的通道加权、跨头归一化和协方差感知的通道分组，优化VQ目标空间。
  3. TaSQ在长链推理和长上下文检索基准测试中表现优异，支持更大的批量大小和更高的吞吐量。

📝 摘要（中文）

随着长上下文大语言模型（LLM）推理的普及，键值（KV）缓存成为主要的内存瓶颈，给内存容量和带宽带来了巨大压力。为缓解这一瓶颈，向量量化（VQ）成为了一种有前景的KV缓存压缩方法。然而，现有的VQ方法在1位压缩下性能显著下降。为此，本文提出了TaSQ，通过结合查询引导的通道加权、跨头归一化和协方差感知的通道分组，定制VQ目标空间，以更好地反映缓存激活的误差敏感性和统计结构。实验表明，TaSQ在多个基准测试中表现优异，且在单个RTX 6000 Ada GPU上，其SGLang实现支持高达14倍的批量大小，并实现了1.87倍的峰值吞吐量提升。

🔬 方法详解

问题定义：本文旨在解决在1位KV缓存压缩中，现有向量量化方法性能下降的问题。现有方法在极端压缩下难以有效利用有限的质心，导致性能瓶颈。

核心思路：论文提出的TaSQ通过结合查询引导的通道加权、跨头归一化和协方差感知的通道分组，定制VQ目标空间，以更好地适应缓存激活的统计特性和误差敏感性。

技术框架：TaSQ的整体架构包括三个主要模块：查询引导的通道加权用于调整通道的重要性，跨头归一化用于统一不同头的输出，协方差感知的通道分组则用于优化质心的分配。

关键创新：TaSQ的核心创新在于其定制的VQ目标空间设计，能够更有效地反映缓存激活的统计结构，与传统VQ方法相比，显著提升了压缩效率和推理稳定性。

关键设计：在参数设置上，TaSQ采用了RoPE兼容的变换，能够轻松与投影权重和代码本合并，保持了传统VQ查找结构，并且增加的服务开销微乎其微。实验中，TaSQ在单个RTX 6000 Ada GPU上实现了高达14倍的批量大小和1.87倍的吞吐量提升。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

在多个基准测试中，TaSQ表现出色，尤其是在长链推理和长上下文检索任务中，显著超越了现有低位KV缓存VQ基线。在单个RTX 6000 Ada GPU上，TaSQ实现了高达14倍的批量大小和1.87倍的峰值吞吐量提升，展示了其优越的性能。

🎯 应用场景

该研究的潜在应用领域包括大规模语言模型的推理优化、实时数据处理和高效的内存管理。通过提高KV缓存的压缩效率，TaSQ能够在资源受限的环境中支持更复杂的推理任务，具有重要的实际价值和未来影响。

📄 摘要（原文）

The key-value (KV) cache becomes a major memory bottleneck in long-context LLM inference, placing substantial pressure on memory capacity and bandwidth. To mitigate this bottleneck, vector quantization (VQ) has emerged as a promising approach for aggressive KV cache compression. However, existing VQ methods degrade substantially in the 1-bit regime. At such extreme compression, each codebook must represent a larger group of channels with a limited set of centroids, making effective use of its capacity increasingly challenging. To address this, we introduce $\textbf{TaSQ}$, which tailors the VQ target space by combining query-guided channel weighting, cross-head normalization, and covariance-aware channel grouping to better reflect the error sensitivity and statistical structure of cached activations. Since these transforms are RoPE-compatible and can be easily merged into projection weights and codebooks, TaSQ preserves the conventional VQ lookup structure and adds negligible serving overhead. Across general, long-chain-of-thought reasoning, and long-context retrieval benchmarks, TaSQ consistently outperforms existing low-bit KV cache VQ baselines while preserving reasoning stability. On a single RTX 6000 Ada GPU, its SGLang implementation supports up to $14\times$ larger batch sizes and achieves $1.87\times$ higher peak throughput compared to the BF16 baseline.