FibQuant：面向随机访问KV缓存压缩的通用向量量化

摘要

长上下文推理日益成为一个内存流量问题。罪魁祸首是键值（KV）缓存：它随上下文长度、批大小、层数和头数增长，并且在每个解码步都要被读取。基于旋转的标量编解码器通过存储一个范数、施加共享随机旋转并逐坐标量化来满足这一系统约束。它们是通用且可随机访问的，但丢弃了归一化步骤所创造的几何结构。经Haar旋转后，$k$个连续坐标构成的块不再是乘积信源，而是单位球上的球面-Beta信源。我们提出FibQuant，一个通用定速率向量量化器，保持相同的归一化-旋转-存储接口，同时用与该典型信源匹配的共享径向-角向码本取代标量码表。该码本结合了Beta分位数半径、Fibonacci/Roberts–Kronecker准均匀方向以及多次重启的Lloyd–Max精调。我们证明所得的向量码在匹配速率下严格优于其标量乘积特例，其高速率增益可分解为单元整形因子与密度匹配因子。同一构造还提供了密集的速率轴，包括小数比特与低于一比特的工作点，且无需校准或变长地址。在GPT-2 small的KV缓存上，FibQuant描绘了一条内存-保真度前沿：注意力余弦相似度$0.99$时压缩$5\times$，$0.95$时达$34\times$。在TinyLlama-1.1B上端到端来看，它在$4\times$压缩时与fp16的困惑度差距在$0.10$以内，并在$b = 2$（$8\times$压缩）——即标量随机访问量化开始失效之处——困惑度比标量TurboQuant低$3.6\times$。

FibQuant：面向随机访问KV缓存压缩的通用向量量化：论文配图 — 图 1：FibQuant 编码器 - 解码器管道。缓存向量 xεℝdx\in\mathbb{R}^{d} 被分割为标量范数标头 ν=‖x‖2\nu=\|x\|_{2} 和单位方向 Πx/ν\Pi x/\nu，其中 Πεℝd×d\Pi\in\mathbb{R}^{d\times d} 是跨多个共享的单个 Haar 随机正交矩阵层、头、提示和标记。单位方向被划分为 kk 个连续坐标的 d/kd/k 个块，每个块落在单位球 𝔹k\mathbb{B}^{k} 上，边缘密度为 fd,kf_{d,k} （引理 1）。通过针对共享离线码本的最近码字查找，将每个块编码为 NN 索引之一 𝒞={cn}n=1N\mathcal{C}=\{c_{n}\}_{n=1}^{N} （等式（1））；解码器反转查找和旋转以恢复 x^\hat{x}。输出比特流是固定速率 (dlog2N)/k(dlog_{2}N)/k 位有效负载，每个缓存向量加上一个 fp16 范数，因此令牌槽位于仿射偏移处，并且任何过去的键/值都可以独立恢复 - 这是将可部署缓存编解码器与可变长度替代方案区分开来的随机访问要求。 TurboQuant 管道的唯一变化是用一个 kk 维向量量化器替换了每坐标标量表。