Shard - 实现10倍KV缓存压缩

# Shard —— 实现 KV 缓存 10 倍压缩 来源：https://krishgarg.com/shard ← 返回 (https://krishgarg.com/writing) Krish Garg 和 Kirrithan Sathananthan · 2026 年 5 月 · github (https://github.com/krish1905/shard) **TL;DR\.***Shard* 是一个即插即用的 HuggingFace `Cache`，能让 Llama-3.1-8B 在 8K 上下文下的 KV 内存缩小约 **10 倍**（32K 下约 11 倍），同时对 NIAH 或 LongBench 没有可测量的影响。起初我们是为了复现 Google 的 TurboQuant[\[1\]](https://krishgarg.com/shard#fn1)，在 4 倍左右停滞不前，后来发现 K 和 V 需要不同处理方式，最终形成了不同的设计：对 K 进行 PCA 加 int4 量化（一旦撤销 RoPE，矩阵实际上是低秩的），对 V 进行 Hadamard 旋转加向量量化。注意力直接在压缩后的 K 上运行，无需 fp16 重建。代码：`krish1905/shard` (https://github.com/krish1905/shard)。 ## 1\. 缘起 我们观察 Llama-3.1-8B 处理 8K 提示时，发现 KV 缓存占用了 **1 GB** 的 VRAM。每个请求如此。而整个模型在磁盘上才 16 GB。模型六分之一的内存只是为了记住它已经看到的内容。这个比例感觉不对劲。 计算很简单。对于仅解码器的 Transformer，序列长度为 $T$： $$ \|\text{KV}\| = 2 \cdot L \cdot H_{kv} \cdot d \cdot T \cdot \text{bytes}(fp16) $$ 对于 Llama-3.1-8B-Instruct，超参数为 $L=32$, $H_{kv}=8$, $d=128$。在 $T=8192$ tokens 时，结果为 $2 \times 32 \times 8 \times 128 \times 8192 \times 2 = 1.07\text{ GB}$。在 128K 上下文时，达到 16 GB，超过模型权重。 过去两年中，每一篇关于 LLM 推理效率的论文都对此有所探讨：量化 KV 缓存[\[kivi\]](https://krishgarg.com/shard#fn-kivi)、驱逐不重要的 token[\[h2o\]](https://krishgarg.com/shard#fn-h2o)、投影到低秩子空间[\[kvtc\]](https://krishgarg.com/shard#fn-kvtc)、完全重构注意力机制[\[mla\]](https://krishgarg.com/shard#fn-mla)。不同的问题、不同的答案、不同的权衡。我们只写下一个目标：**在真实 LLM 上实现 10 倍压缩，且不损失精度。** 我们不知道是否可达。 ## 2\. 我们最初尝试复现的论文 一位朋友发来了 **TurboQuant**[\[1\]](https://krishgarg.com/shard#fn1)（Zandieh, Daliri, Hadian, Mirrokni 于 Google Research + Google DeepMind + NYU；arXiv:2504.19874 (https://arxiv.org/abs/2504.19874)；ICLR 2026）。其核心宣称是在 Llama-3.1-8B-Instruct 上“以每通道 3.5 比特实现绝对质量中立”，以及“以每通道 2.5 比特实现可忽略的质量下降”。我们两天内读了四遍。 核心技巧很漂亮。取任意单位向量 $x \in \mathbb{S}^{d-1}$，乘以随机正交矩阵 $\Pi$。他们的 **引理 1** 指出 $\Pi x$ 的每个坐标服从缩放后的 Beta 分布： $$ f_{X_j}(x) = \frac{\Gamma(d/2)}{\sqrt{\pi}\,\Gamma((d-1)/2)} (1-x^2)^{(d-3)/2} $$ 对于中等 $d$，这收敛于 $\mathcal{N}(0, 1/d)$。由此引出两个事实，将问题从困难变得几乎微不足道： 1. **每个坐标分布相同**，与输入数据无关。你可以预先计算一个最优标量量化器（Lloyd-Max[\[lloyd\]](https://krishgarg.com/shard#fn-lloyd)），并永远应用于每个坐标。 2. **不同坐标在高维下近似独立**（形式上是球面上的测度集中论证）。逐坐标的标量量化在均方误差上接近完全向量量化的最优结果。 他们在每坐标 $b$ 比特下的均方误差上界（**定理 1**）： $$ D_{\text{mse}}(Q_{\text{mse}}) \leq \frac{\sqrt{3}\,\pi}{2} \cdot 4^{-b} \approx 2.72 \cdot 4^{-b} $$ 以及信息论下界（**定理 3**，通过 Shannon 和 Yao 的极小化极大）： $$ D_{\text{mse}}(Q) \geq 4^{-b} $$ 差距是一个小常数（$\sqrt{3}\pi/2 \approx 2.72$）。对于内积失真，通过他们的 QJL（量化 Johnson–Lindenstrauss）残差校正[\[qjl\]](https://krishgarg.com/shard#fn-qjl)有类似的对偶结果：$D_{\text{prod}} \leq \frac{\sqrt{3}\pi^2}{d} \cdot 4^{-b} \cdot \|y\|^2$。 漂亮的数学。但在第 4 节中，我们发现他们实际在 Llama-3.1-8B-Instruct 上的实验最高只有 **每坐标 2.5–3.5 比特**，这意味着： $$ \text{compression} = \frac{16}{3.5} \approx 4.57\times \quad \text{或} \quad \frac{16}{2.5} \approx 6.4\times $$ 不错，但不是 10 倍。我们反复重读论文，试图找到遗漏之处，最终不得不接受：*该量化器是针对最坏情况输入分布（球面上的均匀分布）最优的*。这让你可以预先计算码本并跳过校准，但也阻止了你利用数据实际具有的结构。而 KV 缓存，大概不是球面上的均匀点。 ## 3\. 观察 K 缓存 我们在 8K fineweb 提示上运行 Llama-3.1-8B-Instruct，导出第 12 层的 K 张量，并进行了 SVD。**K 矩阵实际上是低秩的**。从数学上讲并非如此（它是满秩的），但 1024 个奇异值中的 192 个捕获了超过 99.5% 的 Frobenius 范数。尾部急剧下降。 然后我们在没有先撤销 RoPE 的情况下做了同样的事情。奇异值衰减慢得多。RoPE 隐藏了结构。 **原因。** Llama 的 $W_K$ 权重矩阵具有有限的有效秩。这是已知的，像 ASVD[\[asvd\]](https://krishgarg.com/shard#fn-asvd) 这样的权重级低秩压缩直接利用了这一点。当激活通过 $W_K$ 时，输出落在低维子空间上。但 RoPE 以*不同*的角度旋转每个 token 的 K，这会将能量分散到所有维度，并在你堆叠 token 时掩盖子空间。如果你对每个 token 撤销 RoPE（由于你知道角度，可以精确做到），子空间就会重新浮现。 KVTC[\[kvtc\]](https://krishgarg.com/shard#fn-kvtc)（Staniszewski & Łańcucki，arXiv:2511.01815 (https://arxiv.org/abs/2511.01815)，ICLR 2026）明确指出了这一点：*“在进行 PCA 之前移除键中的旋转位置编码（它们会掩盖低秩结构）。”* 完全正确。 K 的 SVD（一层，8K tokens），归一化奇异值： - 应用 RoPE：████████████████████████████████...（缓慢衰减） - 撤销 RoPE：███████▓▓▒▒▒░░░░（在 ≈ 192/1024 处急剧拐点） 所以 K 具有结构。而 TurboQuant 的下界 $D_{\text{mse}} \geq 4^{-b}$ *不适用于结构化数据*；它针对的是球面上的均匀点。如果 K 位于 $d$ 维背景空间的秩 $r$ 子空间上，那么在第一次量化之前，你就能获得对数因子无关的压缩。这是第一个“啊哈”时刻。 ## 4\. 但 V 没有这种结构 同样的转储，同样的 SVD，这次针对 V。奇异值几乎是平坦的。V 看起来基本是随机的。没有可利用的低秩结构。这一点感觉很重要。 我们读过的每篇论文都对 K 和 V 应用*相同*的压缩技术。TurboQuant 对两者进行旋转和标量量化；KIVI[\[kivi\]](https://krishgarg.com/shard#fn-kivi) 对两者进行量化；PolarQuant[\[polar\]](https://krishgarg.com/shard#fn-polar) 对两者进行旋转。但 K 和 V 在**结构上是不同的**。 我们寻找了已经认识到这一点的先前工作。两个发现： - **AsymKV**[\[asym1\]](https://krishgarg.com/shard#fn-asym1)（Tao 等人，arXiv:2410.13212 (https://arxiv.org/abs/2410.13212)，2024 年 10 月）主张对 K 和 V 使用不同的*位宽*，但对两者应用相同的方法（标量量化）。 - **“Homogeneous Keys, Heterogeneous Values”**[\[asym2\]](https://krishgarg.com/shard#fn-asym2)（arXiv:2506.05410 (https://arxiv.org/abs/2506.05410)，NeurIPS 2025）观察到相邻的键获得相似的注意力（因此键可以合并），而值则不同，并提出每侧不同的合并方案。 两者都没有将*对 K 的低秩 PCA* 与*对 V 的向量量化* 结合在一条流水线中。这正是我们希望填补的空白。 ## 5\. 构建 K 路径 Shard 的 K 侧最终设计为： 1. **撤销 RoPE。** Llama 使用旋转一半的 RoPE。对于位置 $p$ 的每个 token，应用逆变换： $$ k_{\text{no-rope}} = k \odot \cos(\theta p) - \text{rotate\_half}(k) \odot \sin(\theta p) $$ 其中 $\text{rotate\_half}([a, b]) = [-b, a]$。 2. **逐层 SVD。** 将 K 展平为 $(B \cdot T, H_{kv} \cdot d)$，减去均值，运行 `torch.svd_lowrank(centered, q=192, niter=24)`。**niter=24 不是可选的。** 在 niter=6 时，基的质量在不同运行间波动；在 niter=24 时稳定下来。我们通过本地测试验证：10 个随机种子的 SVD 恢复特征值方差从 niter=6 时的 3% 降至 niter=24 时的 0.1%。 3. **DP 比特分配。** 并非所有 192 个分量都同等重要。分成每组 64 个，在总预算下对比特选项 $\{0, 2, 4, 6, 8\}$ 运行 DP。这类似于 KVTC 的策略[\[kvtc\]](https://krishgarg.com/shard#fn-kvtc)。不同之处：我们为**零比特选项添加了 4 倍惩罚**。稍后会详细说明。 4. **对基进行 int8 量化**，带逐列缩放。**对系数进行对称 int4 量化**（范围 $[-7, 7]$），带逐分量缩放。这些存储在 GPU 上。每 token K 存储： $$ \underbrace{\text{rank} \cdot 4}_{\text{int4 coeffs}} + \underbrace{\frac{\text{rank} \cdot 16}{T}}_{\text{fp16 scale, amortized}} + \underbrace{\frac{H_{kv} \cdot d \cdot 8}{T}}_{\text{int8 basis, amortized}} \text{ bits} $$ 对于 rank=192 且 8K tokens，这大约为**每 K 元素 0.75 比特**。 ### 4 倍丢弃惩罚 我们从 DP 中的朴素 MSE 目标开始。从添加它的那一刻起，每个基准都退化了。修正来自一篇名为“Quantization Dominates Rank Reduction for KV-Cache Compression”[\[qdrr\]](https://krishgarg.com/shard#fn-qdrr) 的论文，该论文形式上论证，在 softmax 注意力的 Fisher 度量下，**丢弃一个方向在相同比特成本下比标量量化它差得多**（二次方级别）。直觉：如果你杀死了决定某个 token 路由的方向，该 token 的 argmax 会跳到别的东西上，这是一个分类性的、不可恢复的错误。标量量化只添加噪声。他们推导出量化相对于删除的比率为 $3 \times 2^{2b}$（INT4 下为 768 倍）。我们不需要那么多；我们只添加了一个 4 倍乘数： ``` err = gv * DROP_PENALTY if bits == 0 else gv / (3.0 * (1 << (2 * bits))) ``` 一个常数。NIAH 从 0.92 恢复到 1.000。 ## 6\. 构建 V 路径 V 平坦的奇异值谱排除了低秩技巧。我们依次尝试了三种变体： 1. **对原始 V 进行逐通道 NF4 标量量化**。可接受但有损。 2. **Hadamard 旋转 + NF4 标量量化**。Hadamard 去相关了通道（每个通道在旋转后独立地看起来像高斯分布），因此 NF4 的高斯最优码本工作得更好。小幅改进。 3. **Hadamard + 对 4 通道组进行 K 均值向量量化，256 条目码本。** VQ 捕获了标量量化无法捕获的联合结构。256 条目效果良好：足够大以覆盖分布，足够小以使一次查找只占用一条缓存行。 选项 3 在重建余弦和下游质量上以明显优势胜出。每 token V 存储：$H_{kv} \cdot d / 4 \cdot 8 = 256$ 字节（整个层），或**每 V 元素 2 比特**。结合 K 的 0.75 比特/元素，我们目标平均约为 **1.5 比特/元素**，这与实际测量到的 8K–32K 上下文下的 10–11 倍压缩相符。 ``` def vq_encode(v, group_size, codebook_size): bh, seq, hd = v.shape ng = hd // group_size ch_max = v.abs().amax(dim=1).clamp(min=1e-8) flat = (v / ch_max.unsqueeze(1)).reshape(bh * seq, ng, group_size).reshape(-1, group_size) centroids, idx = kmeans(flat.float(), codebook_size, n_iter=30) bits = max(1, (codebook_size - 1).bit_length()) packed, n_orig = pack_nbits(idx.to(torch.uint8).long(), bits) return {"centroids": centroids.half(), "packed": packed, "ch_maxabs": ch_max.half(), "shape": (bh, seq, hd), "group_size": group_size, "bits": bits, "n_orig": n_orig} ``` ## 7\. 注意力汇聚 我们端到端运行了完整流水线，得到 NIAH 召回率为 0.3。不只是差，是灾难性的。经过一天的分段排查，发现两个堆叠的问题。 **注意力汇聚。** Xiao 等人的 StreamingLLM[\[sinks\]](https://krishgarg.com/shard#fn-sinks)（arXiv:2309.17453 (https://arxiv.org/abs/2309.17453)，ICLR 2024）观察到 LLM 将大量注意力倾注到前几个 token 上，即使这些 token 在语义上无关紧要。Softmax 必须总和为 1；如果没有当前 token 是“合适查看的地方”，注意力就会停靠在起始位置。如果你有损地压缩前 4 个 token，汇聚点会失真，每个后续注意力头都会产生垃圾。我们添加了 **4 个 FP16 汇聚 token**，精确保留。 **近因偏差。** 语言具有很强的“查看刚刚发生的事”先验。有损压缩最近的 token 会严重损害下一 token 预测。我们添加了一个 **64 token 的 FP16 残差窗口**，也精确保留。两者都到位后，下一次运行时 NIAH 从 0.3 跃升至 1.000。在 8K 上下文中，汇聚+窗口的存储开销可以忽略不计（8192 个 token 中有 68 个 fp16 token，不到 1% 开销）。 最终存储布局（每层）： \[ 4 sink fp16 \]\[ middle: PCA(K) + VQ(V) \]\[ 64 window fp16 \] ^^^^^^^^^^^^^^^^^^^^^^^^^^ 这就是 10 倍压缩的来源 ## 8\. 解码 token：漂移问题 10 倍数字针对的是预填充阶段。在解码阶段，模型一次生成一个新 token。这些 token 放哪里？ 第一个尝试很笨：通过相同的 PCA 基压缩每个新的解码 token。前 20 个 token 效果良好。在 60 个 token 时崩溃。原因很微妙。 我们的 PCA 基是在预填充 K 分布上拟合的。解码 token 是*不同分布*：它们是模型自身的输出，以预填充为条件。将它们投影到一个对于它们来说太小的基上会产生有偏重建。由于每个解码步骤都依赖于之前的解码 token，**误差会累积**。第 20 个 token 偏差 5%，第 60 个偏差 30%，第 150 个无法辨认。 这时 TurboQuant 回来了。它的全部意义在于*数据无关*量化：码本不依赖于数据，因此每个 token 独立量化，没有漂移。我们划分了流水线： | 段 | 方法 | 原因 | |---|---|---| | 预填充中间 | PCA(K) + VQ(V) | 数据相关，利用结构，10 倍 | | 汇聚 + 窗口 | FP16 | 极小，精确保留 | | 解码流 | Hadamard + Lloyd-Max（可选 QJL） | 数据无关，无漂移 | 对于解码流，我们通过封闭形式更新 $E[X \mid X \in [a, b]]$ 实现了针对 $\mathcal{N}(0, 1/d)$ 的 Lloyd-Max 质心： $$ c_k^{\text{new}} = \frac{\phi(a_k/\sigma) - \phi(b_k/\sigma)}{\Phi(b_k/\sigma) - \Phi(a_k/\sigma)} \cdot \sigma $$ 其中 $a_k, b_k$ 是 Voronoi 边界（相邻质心的中点），$\phi, \Phi$ 是标准正态 PDF/CDF。对于 $b \in \{2, 3, 4, 8\}$，约 20 次迭代收敛，质心在初始化时预先计算一次。 ``` def lloyd_max_codebook(d, bits, n_iter=100): n = 1 << bits sigma = 1.0 / math.sqrt(d) c = [_inv_cdf((i + 0.5) / n) * sigma for i in range(n)] phi = lambda z: math.exp(-0.5 * z * z) / math.sqrt(2 * math.pi) Phi = lambda z: 0.5 * (1 + math.erf(z / math.sqrt(2))) for _ in range(n_iter): bounds = [-1e9] + [(c[k-1] + c[k]) / 2 for k in range(1, n)] + [1e9] nc = [] for k in range(n): a, b = bounds[k], bounds[k+1] m = Phi(b/sigma) - Phi(a/sigma) nc.append(c[k] if m < 1e-12 else (phi(a/sigma) - phi(b/sigma)) * sigma / m) if max(abs(x - y) for x, y in zip(nc, c)) < 1e-9: break c = nc return torch.tensor(c, dtype=torch.float32) ``` 测试了四种流式配置，在 5 个领域的 150 token 解码上……（原文似乎被截断，但翻译应忠实于原文。此处原文为“Tested four streaming configurations on 150-token decodes across 5 dive”，可能未写完。按原样翻译。） 测试了四种流式配置，在 5 个领域的 150 token 解码上……