03. InfiniGen — KV 动态卸载 + 重要 token 选择

元数据

论文标题:InfiniGen: Efficient Generative Inference of Large Language Models with Dynamic KV Cache Management
出处:OSDI 2024
作者方向:Seoul National University 等
关键词:KV offload、speculative attention、importance prediction

一句话精髓

不是把所有 KV 都存起来,而是在线预测”哪些 token 这一步用得上”,只把那些拉回 HBM——把 KV 容量瓶颈转化为”选择问题”。

解决的问题

长上下文场景下 KV 太大装不下。直接卸载到 CPU 内存延迟太高。已有方案(H2O、StreamingLLM)用启发式淘汰一些 token,但信息损失累积。

InfiniGen 的视角:Attention 本质上每步只看少数关键 token,何不在卸载场景下做”按需精确加载”?

关键 idea

设计	内容
KV 全量驻留 CPU 内存	HBM 只装当前 active 的少量 KV
轻量预测器	用前若干层的 attention 分数,预测后续层的 top-k 重要 token
按需加载	只把预测到的 top-k 从 CPU 拉回 HBM
流水重叠	加载 layer N+1 的 KV 与 layer N 的计算并行

🧠 关键洞察:Attention 分数在层间有相关性——前几层”重要”的 token,在后几层往往也重要。这给了 speculative load 的可行基础。

关键架构图

   layer 1 ── attention ── 用 score 预测 layer 2 的 top-k
                              │
                              ▼
                       从 CPU 拉这 top-k KV → HBM
                              │
   layer 2 ── attention(只用拉回的) ── 预测 layer 3
                              │
                              ▼
                            ……

HBM:  ███░░░ (active top-k only)
DRAM: ████████████████████ (full KV)

关键数据(公开报告)

相比全量卸载基线,端到端推理速度提升约数倍
在长上下文 benchmark 上,精度退化可忽略(因为预测准了)
HBM 占用大幅降低,可服务的最大上下文长度显著拉长

局限

预测器训练需要校准数据,对未见过的领域可能 miss
Top-k 预测错了就会 stall(回退去全量找)
仅针对 inference,不考虑 KV 持久化(SSD)层
多用户并发时,不同 session 的”重要 token”集合不同——预测器需要 per-session 维护

对本项目的启示

🌟 核心方法论价值:把”哪些数据放哪一级” 从启发式 LRU 升级到模型驱动的预测式选择。

我们可以扩展:

多类型数据的 importance 预测:KV / 向量索引 / RAG chunk 各有自己的 importance signal,统一抽象成”下一步访问概率”
三级化预测:不仅是”加载到 HBM”,还要预测”该卸到 DRAM 还是 SSD”——选择问题升级到放置问题
预测器 + 成本模型耦合:重要程度 × 当前层级带宽 → 是否值得搬运
训练侧 transfer:InfiniGen 的预测器只在推理时用,我们能不能在训练时学到一个跨任务通用的 importance head?

横向对比

方法	思路	信息损失
全量驻留 HBM	不卸载	0,但容量爆
H2O / StreamingLLM	启发式淘汰	累积损失
InfiniGen	全量留 DRAM,按需精确加载	接近 0
Quest(ICML’24)	块级稀疏 attention	小
本项目目标	InfiniGen 思想 + 多类型 + 三级	—

待精读

预测器具体网络结构、训练数据如何获得?
多 batch 并行下,top-k 重叠率多少?能否共享加载?
论文里给的”层间 attention 相关性”实证图(关键论据)

03. InfiniGen — 自适应 KV 卸载与检索

搜索