论文

OCTOPUS:在最佳平方误差量化下通过八面体参数化为 Transformer 优化 KV 缓存

OCTOPUS: Optimized KV Cache for Transformers via Octahedral Parametrization Under optimal Squared error quantization

模型推理KV Cache

摘要

键值 (KV) 缓存在长上下文自回归推理中主导内存带宽和占用空间。最近的旋转预处理编解码器(TurboQuant、PolarQuant)表明,结构化随机旋转后跟与分析上易于处理的边缘匹配的每坐标标量量化器是 KV 压缩的近乎最佳配方。 OCTOPUS 通过旋转坐标三元组的联合量化推进了这一范式。每个三元组的方向通过八面体参数化映射到一个正方形,并且两个结果坐标和三元组范数根据实现匹配的边际进行 Lloyd-Max 量化。优化每三元组的平方误差会给出严格不均匀的位分配,仅取决于密钥的总维数。我们发现,在我们测试的每个实际解码器上,扫描的有限维质量最优是恒定的。该编解码器是数据无关的、在线的并且在给定种子的情况下具有确定性。在文本、视频和音频方面,OCTOPUS 在每个报告的位宽和指标上都匹配或击败了之前的每个旋转编解码器,并且随着极端压缩的位下降,领先优势不断增大。此外,融合的 Triton 实现会动态重建密钥,而不会具体化未压缩的密钥,因此编解码器不会在现有的反量化上增加解码时带宽或延迟。项目页面:https://octopus-quant.github.io/