Agent初创实习-大模型推理加速02
H2O 方法汇报:Heavy-Hitter Oracle 如何动态压缩 KV Cache
参考论文:H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models
本汇报回答三个问题:
- H2O 的 pipeline 是怎么实现的?
- 它为什么能推理加速?
- 它和 StreamLLM 的“attention sink + sliding window”有什么区别?
1. 先说结论
H2O 做的事情很直接:
在生成过程中,不保存所有历史 token 的 KV cache,只动态保留“最近 token”和“历史上最常被注意到的 token”。
其中,“历史上最常被注意到的 token”就是 Heavy Hitters,也就是 H2。
它不是固定保留开头几个 token,也不是固定保留每隔几个 token,而是每一步生成时根据 attention 分数更新 token 的重要性。谁在过去生成过程中反复被后续 token 注意到,谁就更可能留在 KV cache 里。
一句话类比:
StreamLLM 像固定保留“开头几个主持人 + 最近聊天内容”;H2O 像动态保留“最近聊天内容 + 整场对话里一直被大家反复引用的关键人物”。
2. 背景:为什么 KV cache 会成为瓶颈
自回归生成时,模型每生成一个新 token,都要看前面所有 token。为了避免每一步都重新计算历史 token 的 key 和 value,推理系统会保存历史 token 的 KV cache。
标准做法是:
第 1 步:保存 token 1 的 KV 第 2 步:保存 token 1,2 的 KV 第 3 步:保存 token 1,2,3 的 KV ... 第 n 步:保存 token 1...n 的 KV问题是 KV cache 的显存开销会随着:
- 序列长度
- batch size
- 层数
- hidden size
线性增长。
长文本生成和大 batch 推理时,KV cache 可能比你想象中大得多。论文里举例,30B 模型、batch size 128、sequence length 1024 时,KV cache 可以到 180GB。
所以 H2O 的目标是:
不保存全部 KV,只保留一小部分关键 KV,同时尽量不掉效果。3. H2O 的两个核心观察
3.1 Attention 很稀疏
虽然 Transformer 是 dense attention,但实际推理时,每个新 token 通常只强烈关注少数历史 token。
也就是说:
当前 token 生成时,并不是每个历史 token 都同等重要。论文观察到,LLM 推理阶段的 attention matrix 很稀疏,大部分位置的 attention 分数很低。
这说明:
保留全部 KV 可能是浪费的。3.2 少数 token 长期很重要,也就是 Heavy Hitters
论文进一步发现,历史 token 的累计 attention 分数呈现长尾分布。
也就是说,少数 token 会反复被后续 token 注意到,它们贡献了大部分注意力价值。这些 token 就叫 Heavy Hitters。
举个直观例子:
输入:Children laughed and played in the sunny park ...在后续生成中,模型可能经常回看:
Childrenplayedpark
而一些功能词可能很少被回看。
H2O 的直觉是:
如果 KV cache 空间有限,与其随机留,不如留“最近 token + 历史高注意力 token”。4. H2O Pipeline
下面是 H2O 的整体流程。
