Workflow
Cache Me If You Can:陈丹琦团队如何「抓住」关键缓存,解放LLM内存?
机器之心·2025-06-24 22:07

| 机器之心报道 | | --- | | 机器之心编辑部 | 普林斯顿大学计算机科学系助理教授陈丹琦团队又有了新论文了。 近期,诸如「长思维链」等技术的兴起,带来了需要模型生成数万个 token 的全新工作负载。 大多数语言模型都基于 Transformer 架构,其在进行自回归解码(即逐字生成文本)时,需要将所有先前 token 的注意力状态存储在一个名为 KV 缓存的 内存区域中。 KV 缓存是模型进行快速推理的基石,但它的大小会随着输入文本的长度线性增长。例如,使用 Llama-3-70B 模型处理一个长度为 128K token 的提示 (这大约相当于 Llama 3 技术报告本身的长度),就需要分配高达 42GB 的内存专门用于存储 KV 缓存。 许多先前的工作意识到了这个问题,并提出了从内存中丢弃(驱逐)部分键值对的方法,以实现所谓的「稀疏注意力」。然而,在一个公平的环境下对它们 进行横向比较却异常困难。 生成过程 = 预填充(对输入进行前向传播并保存键值对)+ 后填充(一次解码一个输出词元)。 有些论文旨在加速预填充阶段;另一些则忽略该阶段,转而致力于最小化后填充阶段的内存开销。同样,有的研究侧 ...