论文
KV 缓存的 JoLT:通过联合秩—位宽分配进行近无损 KV 缓存压缩
A JoLT for the KV cache: Near-Lossless KV Cache Compression via Joint Rank-bit Allocation
摘要
键值(KV)缓存是长上下文语言模型推理中的主要内存瓶颈。现有的压缩方法独立地应用低秩分解或量化,而不在共享存储预算下联合分配秩和精度。我们引入了 JoLT,一种 无需训练 压缩器,它将分组预填充缓存视为四阶张量,并沿词元和特征模式(携带低秩结构的两个轴)应用部分 Tucker 分解,同时保持头部和层模式完好无损。旋转的低位量化器捕获截断残差,单个拉格朗日对偶在全局字节约束下分配每组塔克等级和残差位宽度。 FlashJoLT 用随机近似值替换了精确的词元模式 SVD,该近似值以压缩成本的一小部分与自由区内的 JoLT 相匹配,并且融合的 Triton 解码内核直接评估存储因子的注意力,而无需具体化密集的 KV 张量。在来自四个架构系列的五个模型中,涵盖多头注意力、分组查询注意力和专家混合架构,JoLT 实现了 2 - 3 倍的压缩,且困惑度下降低于 0.2%,无需重新训练。在使用 LLaMA-3.1-8B 的 64K 环境下的 RULER 上,检索精度在 3 倍时仍保持近乎无损,在 4 倍和 5 倍时分别仅下降 0.90 和 2.40pp。 JoLT 证明了张量感知低秩分解和量化残差在单一存储预算下统一,无需重新训练即可跨不同模型架构实现近乎无损的 KV 缓存压缩。