技术摘要：Tamarin (HSQ) – 可逆分层 KV-Cache 压缩

1. 问题背景

在自回归大语言模型（LLM）推理中，键值（KV）缓存是 GPU 显存的主要消耗者，尤其是在长上下文场景下。与可以一次性量化的模型权重不同，KV 缓存随上下文长度线性增长。例如，一个使用 fp16 精度的 Qwen3-8B 模型在 32K token 上下文时需要约 4.5 GiB 的显存，这往往会超过标准 GPU 的容量。

现有的解决方案面临着一种带有共同盲点的二分困境：

• KV 量化： 将所有条目压缩为低精度（2–4 bit）。虽然这种方法是安全的，但其压缩比受到限制（通常为 fp16 的约 4 倍），并且在较低位宽下可能会降低质量。
• Token 剔除/合并： 诸如 StreamingLLM、H2O 和 SnapKV 等方法，会在预填充（prefill）阶段永久丢弃或融合被认为不重要的 token。这些方法实现了极高的压缩比，但具有不可逆性。如果一个在早期被剔除的 token 在后续解码步骤中变得至关重要，其信息将会永久丢失。

先前工作的核心局限在于：token 的重要性是取决于未来的查询（query）的，而基于剔除的方法必须在压缩时就对这种缺乏未来上下文信息的预测进行决策。

2. 方法论：Tamarin (HSQ)

Tamarin（分层语义量化，简称 HSQ）将 KV 缓存压缩重新定义为一个索引问题，而非删除问题。它维护了三个同步的上下文表示，确保每个压缩后的 token 都是可寻址且可恢复的。

三层架构

1. L1（焦点层，GPU）： 一组较小的、未经过摘要处理的 token，以 4-bit 量化向量的形式存储。这包括注意力汇点（attention sinks）、可选的近期窗口，以及由学习到的读取器（reader）明确标记为重要的 token。
2. L2（摘要层，GPU）： 剩余的 token 被划分为若干组（每组大小 $N \in \{8, 16\}$）。每个组被压缩为两个学习到的代表向量（$\hat{k}, \hat{v}$），并以 3-bit 精度存储。至关重要的是，这些向量充当的是路由键（routing keys），而非内容的替代品。
3. L3（存档层，CPU RAM）： 每个压缩 token 的原始（量化后）键和值以 4-bit 精度存储在 CPU RAM 中。这作为内容的“事实来源”（source of truth）。

解码过程

在解码期间，该过程运行为一个可逆的索引查找过程：

1. 路由（Routing）： 每个头的路由器会对 L2 摘要键进行评分，计算其与当前查询（query）的相关性。
2. 获取（Fetching）： 根据路由预算（$\beta$），路由器选择得分最高的组。系统从 L3 CPU 存档中提取相应的原始 token，在 GPU 上进行反量化，并将其旋转到语义位置。
3. 注意力（Attention）： 注意力机制在 L1 token、L2 摘要以及已恢复的 L3 原始 token 的并集上运行。对于已提取原始 token 的组，L2 摘要会被掩码（mask）处理，以防止重复计数。

核心技术组件

• 读取器（Reader）： 一个小型 MLP，根据早期解码器的隐藏状态对 token 重要性进行评分。它负责将 token 指向 L1、L2（中等组）或 L2（低等组）。
• 压缩器（Compressor）： 一个 Perceiver 风格的交叉注意力池化器，用于生成 3-bit 路由键。它的训练目标是在模型的查询投影下有效地进行排序，而不是重建内容。
• 量化（TurboQuant）： 使用带有固定随机正交旋转的随机标量量化来解相关坐标。L1 和 L3 使用 4-bit 精度；L2 使用 3-bit 精度。
• 实现： 基于 HuggingFace Transformers 实现了一个约 9000 行的 Python 库，无需自定义 CUDA 内核或修改模型权重。

3. 主要贡献

• 架构： 一种三层可逆层次结构，将路由（L2）与内容（L3）分离，从而实现了高压缩比且不造成永久性的信息丢失。
• 验证方法论： 使用严谨的统计框架进行评估，包括配对试验（每单元 $n=500$）、精确的 McNemar 检验、用于二进制检索等价性的 TOST 检验，以及针对困惑度（perplexity）的单侧非劣效性边界测试。
• 规模化发现： 研究表明，随着模型规模的增大，质量成本会降低；此外，该方法可以与 $\ge$ 4B 参数的模型进行 4-bit 权重量化的自由组合。
• 开源： 以 Apache 2.0 协议发布了 beta 版 HSQ 软件包，包括模型适配器和评估工具包。

4. 结果

在 Qwen3 系列（0.6B–14B）以及 fp16、int8 和 int4 权重下，针对高达 32K 的上下文进行了评估：

• 压缩比： Tamarin 将 GPU 驻留的 KV 缓存减少了 12–28×（例如，Qwen3-8B 在 32K 长度下，GPU 显存占用从 4.5 GiB 降至 172 MiB，同时 CPU 占用为 2.29 GiB）。
• 检索性能：
  • 在 4B 和 8B 模型上，检索性能在所有深度上与未压缩的基准线在统计上是等价的（TOST 等价性）。
  • 在 32K 上下文中，Qwen3-8B-int4 实现了 500/500 的针尖检索（needle retrieval）一致性，零不一致对。
• 困惑度（PPL）开销：
  • 对于 4B–14B 模型，PPL 开销极小（0.1–0.4%）。
  • 在 4B 模型上，多语言困惑度测量值甚至低于基准线（比例为 0.9972），这表明压缩成本低于测量误差范围。
  • 小型模型（0.6B）的成本较高（fp16 下为 4–6%，int4 下为 6–10%）。
• Token 一致性： 对于 4B–14B 模型，贪婪（greedy）层级 token 一致性与基准线保持在 96–97%。
• 速度：
  • 预填充（prefill）开销在长上下文时得到摊薄（在 32K 时为 1.2–1.3×）。
  • 解码吞吐量目前是瓶颈，在 8K 以内的上下文中运行速度为基准线的 40–73%，而在 32K 时由于从 CPU 获取和反量化的开销，下降至 17–18%。

5. 主张与意义

论文声称 Tamarin 成功地将 KV 缓存管理范式从删除转向了索引。通过利用路由器在查询时刻才决定哪些 token “重要”，该系统避免了基于剔除方法的不可逆性。

具体主张：

• 容量优先于延迟： Tamarin 被定位为一种在有限 GPU 显存中容纳更大上下文和更大 Batch Size 的解决方案，而非旨在降低延迟。
• 可逆性： 每个压缩后的 token 都是可寻址的；系统不会永久丢弃信息。
• 模型缩放： 该方法随着模型容量的增加而变得更加有效，4B 和 8B 模型表现出与基准线近乎完美的等价性。
• 组合性： 对于至少 4B 参数的模型，它可以与 4-bit 权重量化无缝结合，且不会产生可测量的额外成本。

明确的局限性与边界：

• 小型模型： 低于 4B 参数的模型（特别是 0.6B）会受到噪声叠加的影响，导致较高的 PPL 开销。
• 脆弱的检查点： Qwen3-14B-int4 检查点在特定深度的单针检索中显示出细微的“空洞”，这归因于量化的脆弱性而非方法本身。
• 模型家族： 该方法在 Qwen3 上得到了统计验证。在对 Mistral-7B 进行的初步测试中出现了性能下降的“家族效应”，据推测是由于缺乏每头键归一化（per-head key normalization）导致的。
• 吞吐量： 当前的实现是用解码速度换取内存容量，在极长上下文（32K）下会面临显著的吞吐量惩罚。

论文结论指出，Tamarin 为在消费级和数据中心 GPU 上实现长上下文推理提供了一条实用的、可逆的路径，前提是用户接受当前在解码吞吐量方面的权衡。