H2O 完全拆解：抓住 KV Cache 里的「重击者」，20% 预算无损生成（NeurIPS 2023）

观察一：注意力矩阵 95% 以上是稀疏的#

研究者拿预训练的 OPT 系列模型，在 WikiText-103 验证集上做零样本推理，统计每一层的注意力矩阵稀疏度。判定方法：对注意力概率矩阵（softmax 归一化后的 $QK^{\top}$）的每一行，取该行最大值的 1% 作为阈值，低于阈值的元素视为「不重要」。

结果出乎意料：模型虽然是稠密训练的，推理时几乎所有层的注意力矩阵稀疏度都超过 95%。这意味着生成下一个 token 时，真正贡献注意力的历史位置只有不到 5%，每一步实际需要用到的 KV 理论上可以砍到 1/20——这是 H2O 能成立的空间前提。

但稀疏性只是必要条件。注意力稀疏说明「每步只需少数 token」，却没说「该留哪少数 token」——而驱逐一旦选错，代价不是一次 cache miss，而是生成质量的连锁崩溃。95% 的稀疏率对应「任意选 5%」几乎必然灾难。

观察二：累计注意力服从幂律，少数重击者拿走大部分注意力#

第二个观察把「该留谁」的问题也回答了。把序列里每个 token 在注意力计算中累计收到的注意力概率（对所有 query 位置求和）画出来，分布是典型的幂律：一小部分 token 吸收了绝大部分注意力质量，其余 token 的得分长尾拖地。论文把这些极少数的关键 token 命名为重击者（Heavy Hitter，简写 $\text{H}_2$，这也是方法名 H2O 里「H2」的来源——注意不是水分子）。

图 2 的四个面板把这个观察钉死了：

• 面板 (a)：各层注意力稀疏度，几乎全部超过 95%，稀疏性在几乎每一层都成立；
• 面板 (b)：把累计注意力得分按词表里的词汇总（红色散点），和该词在语料中的共现次数（灰色曲线）叠在一起，两条曲线高度吻合。这说明重击者不是偶然涌现的——文本里频繁共现的词（主题词、高频功能词、反复出现的实体）天然是注意力的汇聚点，谁在语篇里「社交活跃」，谁就拿走最多的注意力。用论文里的话说，重击者的涌现是自然的，且与 token 在文本中的频繁共现强相关；
• 面板 (c)：把重击者对应的 token 从注意力中屏蔽掉（保留其余全部 KV），模型在多个下游任务上的准确率断崖式下跌——证明重击者不是「注意力集中的装饰品」，而是模型真正依赖的信息载体；
• 面板 (d)：预告了整个方法的可行性——在 20% 的 KV 预算下，用局部统计（只累加已发生步骤的注意力分数）决定驱逐的 H2O，和用全局统计（一个能看到未来所有注意力的 oracle）决定驱逐的 H2O，表现几乎一样好；而只保留最近 token 的 Local 策略大幅落后于全量基线。

面板 (d) 是全文最关键的图之一：它说明「过去被关注得多的 token，未来也大概率继续被关注」——过去的注意力分数是对未来重要性足够好的预测器。为什么？因为重击者的本质是语篇的主题与高频结构，主题在长文本中是缓慢变化的。自回归解码时我们无法预知未来每个 query 会注意谁，但用历史累计注意力做贪心近似，代价小到可以忽略。

问题形式化#

先把「带容量限制的生成过程」说清楚。设注意力 query 矩阵为 $Q \in \mathbb{R}^{n \times d}$，key 矩阵为 $K \in \mathbb{R}^{n \times d}$，第 $i$ 个 token 生成时，KV Cache 里实际保留的 token 集合记为 $S_i \subset [n]$（$[n]$ 表示全部 $n$ 个位置），预算为 $k$。驱逐策略 $g: S_{i-1} \to S_i$ 满足两条约束：

• $|S_i| = k$：缓存大小不随时间变化；
• $|S_i \setminus S_{i-1}| \le 1$，等价于 $|S_i \cap S_{i-1}| \ge k-1$：每步至多驱逐一个 token，保证驱逐本身的开销有上限。

生成过程中，第 $i$ 步的注意力输出是：

其中 $Q_{i,*}$ 是 $Q$ 的第 $i$ 行（当前 query），$K_{S_i,*}$ 是只取 $S_i$ 中那些行的 key 子矩阵，$D_i$ 是归一化常数。被驱逐的 token 相当于 key 置零并从归一化项里剔除——这就是驱逐后注意力自动重新归一化（softmax 分母只对剩余 token 求和）的数学表达。目标：找一个驱逐策略，让受限生成过程的输出尽量接近无限制的原始过程。

贪心驱逐算法#

直接找最优的 $S_i$ 是组合问题，但有了「重击者」这把尺子就有了贪心方向。定义 token 集合的分数函数为集合内所有 token 的累计注意力得分之和：

其中 $o_s$ 是 token $s$ 从进入缓存以来累计收到的注意力概率。那么每一步，在「现有缓存 ∪ 新 token」里逐出累计分数最低的那一个，就能让剩余集合的累计分数最大：

伪代码如下（对应论文 Algorithm 1，把注意力概率累加进一个 per-token 分数向量，满员后每步贪心驱逐分数最低者）：

1
def h2o_generate(Q, K, budget_k):
2
    S = []                       # 缓存中的 token 集合
3
    acc = {}                     # token -> 累计注意力分数
4
    for i in range(len(Q)):
5
        scores = softmax(Q[i] @ K[S].T)          # 当前 query 对缓存内 token 的注意力
6
        for s, a in zip(S, scores):
7
            acc[s] += a                           # 累加进历史总分
8
        S.append(i); acc[i] = 0.0                 # 新 token 入缓存
9
        if len(S) > budget_k:
10
            u = min(S, key=lambda s: acc[s])      # 驱逐累计注意力最低的 token
11
            S.remove(u); del acc[u]
12
    return S

论文图 3 用一个小例子把驱逐过程画了出来（预算为 3）：

为什么是「重击者 + 最近」双缓冲，而不是纯重击者？#

论文的驱逐策略实际保留的是两部分各占一半预算：K 个重击者 + K 个最近 token（总预算 2K，通常按 prompt 长度的百分比折算）。论文与多个基线的对比中反复出现一个现象：只保留重击者、或只保留最近 token，都无法维持全量 KV 的质量（掉 2.85% 到 22.75% 不等），两者合体才和全量基线打平。为什么最近区不可或缺？

第一个理由是局部性：语言建模里，紧挨着当前位置的若干 token 与当前生成的关联最强——语法衔接、指代、正在展开的短语都发生在最近窗口内。重击者管「这篇文档在讲什么」（长期主题），最近区管「现在正说到哪」（短期工作记忆），两个尺度缺一不可。

第二个理由更微妙：资历歧视。累计分数是按时间攒出来的，一个刚生成的 token 还没来得及被任何后续 query 注意过，累计分数天然是零；如果驱逐只看累计分，新 token 会在进入缓存的下一秒就被当垃圾逐出，永远活不过「考察期」。给最近 K 个 token 一个固定席位（先进先出滑动窗口），相当于保护期——它们不参与驱逐竞争，等它们滑出最近区、有了足够的被关注历史之后，才有资格作为「重击者候选人」与历史 token 竞争。

这个设计与站内拆解过的 QEvict 如出一辙——QEvict 保留了 5 个 attention sink token 和 32 个 recent token 作为两个「保护区」，只是把窗口换成了层级结构（详见 QEvict 完全拆解）。H2O 是这套「保护区」设计的源头。

为什么要按头独立驱逐？#

还有一个常被忽略的实现细节：驱逐是按注意力头独立进行的——每一层里，每个 head 各自维护自己的重击者列表，各自逐出自己的最低分 token。论文在对比 SpAtten（一个用累计注意力选 token 的早期工作）时专门强调了这一点：SpAtten 把注意力分数跨头跨层累加在一起做统一决策，而 H2O 允许每个 token 在每个 head 上独立地被保留或驱逐。

直觉上这很重要：不同注意力头在预训练模型里分化出了不同职能——有的头追踪句法依赖，有的头做指代消解，有的头只对位置敏感。一个 token 对句法头是重击者、对位置头可能毫无意义；如果跨头汇总打分，小众但关键的功能就会被多数头稀释掉。按头驱逐等于给每个「专家」独立的预算，保住的是各个功能维度上的重要 token。官方实现（见下文）里「one heavy hitter will be evicted for each head」说的就是这个机制。

理论的角色：子模函数与贪心近似#

H2O 把「限制容量下选最优 token 集」建模成一个动态子模最大化（dynamic submodular maximization）问题，并给出了理论保证。这里把直觉讲清楚，不抄证明。

子模函数（submodular function）刻画的是收益递减：集合越大，新增一个元素的边际收益越小。形式化地说，对任意 $X \subseteq Y$ 和不在 $Y$ 里的新元素 $x$：

直观理解：给一个只有 3 个 token 的集合补充一个新 token，新增的信息量通常大于给一个已有 30 个 token 的集合补同样的 token——新概念可能已经被大集合覆盖过了。文档摘要（coverage）类函数是典型例子。论文论证注意力的「信息覆盖」天然具有这种递减回报性质，因此是一个（近似）子模函数。

子模函数最大化的经典结论是：在基数约束（最多选 $k$ 个）下，贪心算法——每步选边际收益最大的元素——能达到最优值的 $1 - 1/e \approx 63.2\%$，这是理论上能保证的最好常数。H2O 的驱逐正是「反向贪心」：保留能使剩余集合分数最大化的 $k$ 个 token。论文把「每步序列都在变长、都要新加一个 token 再驱逐一个」的版本形式化为 dynamic submodular 框架，给出（非正式的）定理：

即：在温和假设下，贪心得到的缓存集合 $\widetilde{S}_i$ 能达到最优集合函数值的 $(1-\alpha)(1-1/e)$ 倍，再减去一个可容忍的误差项 $\beta$。这套理论的实际影响力不如两个经验观察大（假设的抽象程度较高，后续工作引用它更多是为了说明「驱逐是子模问题」这一视角），但它给了贪心驱逐一个自洽的解释框架：如果注意力函数确实是子模的，那么每步驱逐「当前最没用」的 token 在理论上就是接近最优的贪心。

FlexGen 改造：预分配 + 原地覆写#

论文的高吞吐实现基于 FlexGen（一个可插拔权重/缓存 offloading 的白盒 OPT 推理引擎）。要做的事情其实很小：FlexGen 是 OPT 的显式逐层实现，KV Cache 的读写都经过显式的缓存管理函数，H2O 只在这些函数里做「手术」。

关键工程决策有三个：

1. 预分配内存，驱逐不搬数据。 缓存张量一次性分配好，容量固定为「K 个重击者槽位 + K 个最近槽位」。被驱逐的 token 不是从内存里挪走，而是把它占的槽位直接覆写成新 token 的 KV——一个 cache_replace 就完成驱逐+插入，避免任何数据搬移，保证驱逐不引入 I/O 开销。
2. 最近区用环形队列。 最后 K 个槽位循环使用，新 token 每步写入，最老的最近 token 被顶掉，头尾指针 $O(1)$ 推进。
3. 累计分数向量随注意力计算顺带更新。 注意力概率在每步的 kernel 里本来就要求出来，把它们累加进一个 per-head 的分数向量是几乎免费的额外操作；驱逐决策（找最低分）在分数向量上做线性扫描或维护小顶堆即可。

生成循环本身（权重/激活的预取与计算重叠、offloading 调度）完全不用动。论文给出了实现骨架：attention_forward 在 decode 分支调用 compute_attention()，它会返回新算出的 KV 以及本次要驱逐的槽位索引（evict_ids），store_cache() 拿到 evict_ids 后执行覆写。驱逐策略被封装在 compute_attention 这一个函数里，想换任何别的驱逐策略都只需要改这里——这也是为什么 H2O 的代码可以作为后续各种驱逐研究的实验平台。

两套代码：模拟版与真实丢弃版#

官方仓库 FMInference/H2O 提供两套实现：h2o_flexgen（基于 FlexGen，追求端到端吞吐，报告吞吐数字用它）和 h2o_hf（基于 Hugging Face Transformers，跑各种质量基准）。h2o_hf 里又分两种模式：

• 模拟版：不真正丢 KV，而是通过修改 attention mask 把被驱逐 token 的位置屏蔽掉，模拟驱逐的效果。优点是与标准 Transformers 前向完全兼容、实现简单、跑 benchmark 快；
• 真实丢弃版（utils_real_drop）：真正改写 attention 前向，把被驱逐 token 的 KV 从缓存里物理删除。它验证了一个重要事实：模拟版的结果是可信的——驱逐带来的质量变化纯粹来自「注意力看不到这些 token」，而不是什么 mask 伪影。

每步开销分析#

驱逐策略本身几乎不增加解码延迟：累计分数的累加复用注意力计算的结果，驱逐选择只在缓存满员后每步淘汰一个最低分 token。内存上，额外只需维护每个（层, 头）一个长度不超过预算的分数向量，和 KV 本身的体量相比可以忽略。算上这点开销，H2O 端到端（含 prefill 与解码、含驱逐决策本身的耗时）依然拿到了论文报告的全部加速比。

精度：预算压到 20%，性能追平全量#

论文在 OPT（6.7B–66B，并覆盖至更大尺寸）、LLaMA（7B/13B/30B）、GPT-NeoX-20B 三类模型上，用 HELM 与 lm-eval-harness 的八个任务（COPA、MathQA、OpenBookQA、PiQA、RTE、Winogrande、XSUM、CNN/Daily Mail）做评测，缓存预算从 4% 扫到 100%。图 4 汇总了各任务、各模型上的「KV 预算-准确率」曲线：

论文特别指出两个「失效点」：在 {LLaMA-13B, XSUM} 和 {LLaMA-7B, CNN/Daily Mail} 两组上，Local 策略在 60% 预算时就崩了，而 H2O 在 20% 预算下依然匹配全量性能——长文本摘要这类强依赖远距离信息的任务最能区分两者。总体看，多数任务上 H2O 与全量基线的差距在 1 个百分点以内，部分任务反超，论文称之为「20% 预算实现 5 倍内存压缩而几乎无损」。

一个反直觉的现象是：有些任务上 H2O 甚至超过全量基线。比如 OPT-66B 的 RTE 高 0.73%、OPT-30B 的 MathQA 高 0.64%、GPT-NeoX-20B 的 XSUM 高 0.18。论文把这种现象解释为驱逐带来的正则化效应——丢掉的冗余 KV 反而滤掉了一些干扰信息。表 1 给出了一组代表性数字（5-shot，20% 预算）：

注意最后一行的 Local 策略：同样是 20% 预算、同样的总 token 数，只因为保留的是「最近的 20%」而不是「重击者+最近」，性能直接崩到接近随机。表里还有一层值得玩味的对照：5-shot 的完整 prompt 很长，而 H2O 用 20% 预算大约只相当于保留 1.2 个样本的上下文——它却稳定胜过用全量 KV 的 0-shot 和 1-shot 模型。这说明被驱逐的主要是 prompt 里的冗余，而重击者机制保住了真正承载任务信息的示范样本。

论文还做了「单独效果」消融（表 9），把预算分别只给重击者、只给最近区、两者各半，用 OPT-13B/30B 在四个任务上对比。以 OPT-30B 的 PiQA 为例：全量 78.51，只留最近区 55.82，只留重击者 67.25，两者合体 78.45——重击者单独贡献远大于最近区，但两者缺一不可（详见上文「双缓冲」一节）。另一组重要实验说明 H2O 的重击者区可以当外挂增强件：把 strided 或 fixed 的 Sparse Transformer 稀疏模式叠加重击者区之后，strided 模式几乎追平全量（表 2 中 COPA 从 50.00 回到 83.00，全量为 85.00），fixed 模式也大幅回升（COPA 从 61.00 回到 76.00）；对 Top-K 策略则有最多 2 个百分点的提升（表 8）。注意论文正文说两种模式叠加重击者后「达到与全量相似的表现」，严格看 fixed 模式在部分任务上仍差数个点——表里的数字比正文的表述更诚实。但这不影响核心结论：「保留累计注意力最高的历史 token」是一个与具体稀疏模式正交的通用增强，任何按位置/模式剪枝的方案，只要把省下来的预算部分让给重击者，都能显著止血。

吞吐与延迟：省下的显存换成更大的 batch#

精度之外，论文在系统层面报告了吞吐和延迟收益。方法是在 FlexGen 里实现驱逐，与 FlexGen 已有的 offloading、量化正交组合。收益逻辑很直接：KV Cache 从 100% 缩到 20%，省出的显存能换来两样东西——(1) 同一个 GPU 上塞下大得多的 batch；(2) 原本要依赖 CPU offloading 的配置可以完全不用 offload。两样都直接作用于解码吞吐。

T4（16GB）上的合成数据吞吐对比（token/s，括号里是有效 batch 大小与所需的最深 offloading 层级，C 表示 CPU）：

OPT-30B 的场景最能说明问题：模型权重加全量 KV 在 16GB 显存上根本放不下，三个基线都被迫深度 offload 到 CPU 内存，吞吐被 PCIe 带宽卡死在 1 token/s 上下；H2O 把 KV 砍到 20% 后，有效 batch 从 DeepSpeed/Accelerate 的 4–8（FlexGen 为 80）提升到 416，吞吐 18.8 token/s——比 DeepSpeed 与 Accelerate 高约 30 倍、比 FlexGen 高约 2.2 倍。论文标题级的结论用的是「最高」口径：相对 DeepSpeed ZeRO-Inference 与 HF Accelerate 最高 29×（出现在 XSUM/OPT-30B 的 offload 场景，6.70 vs 0.23 token/s），相对 FlexGen 最高 3×（出现在 OPT-6.7B 纯 GPU 场景，51.7 vs 16.8 token/s）。OPT-6.7B 纯 GPU 场景同样翻了三倍多：51.7 vs FlexGen 的 16.8，且 H2O 还顶着 4 倍于基线的 batch。XSUM 真实数据上结论一致。

A100（80GB）上的结果更直接，同一张卡、无 offloading，比的是同 batch 下的延迟与大 batch 下的吞吐上限：

同 batch 下延迟降低 1.1–1.9 倍（主要来自每步读取的 KV 变少、访存带宽压力下降），而 FlexGen 在 batch 64 直接 OOM 的配置，H2O 跑到 1161 token/s——相对 FlexGen 最好成绩 494.1 又提高了 2.35 倍。注意这里的模型（OPT）训练序列只有 2K，跑 10K 序列本来就超出设计范围，论文用它演示的是「KV 驱逐让更大 batch 成为可能」的系统学收益，真正面向长上下文的模型收益只会更大。

与量化的组合也有报告：H2O 与 4-bit 权重量化叠加后精度几乎不变（COPA/OpenBookQA/PiQA 上 84.00/43.20/78.80，vs 全量 85.00/43.20/78.51），而量化省下的显存让 OPT-6.7B 的 batch 从 4 提升到 44–70，同配置的吞吐从 35.1–52.1 token/s 涨到 50.5–72.5 token/s。驱逐和量化一个砍 token 数量、一个砍每个字节的位数，天然正交。