TokenButler：Token 重要性是可预测的
摘要
大型语言模型（LLM）依赖键值缓存（KV-Cache）来存储 token 历史，从而实现 token 的高效解码。随着 KV-Cache 的增长，它成为主要的内存和计算瓶颈。然而，有机会缓解这一瓶颈，先前的研究表明，只有一小部分 token 对每个解码步骤有实质性贡献。寻找这些关键 token 的一个关键挑战在于它们是动态的，且高度依赖于输入查询。现有方法要么因永久驱逐 token 而面临质量风险，要么保留完整的 KV-Cache 但依赖于检索 token 块，并且许多现有的 KV-Cache 稀疏化方法依赖于不准确的 token 重要性代理指标。为了解决这些局限性，我们引入了 TokenButler，这是一种高粒度、查询感知的预测器，能够学习识别这些关键 token。TokenButler 以固定的深度步长预测低维重要性查询，并将它们与学习到的真实 KV-Cache 键的投影相结合，以低成本对 token 进行评分，从而在固定预算下实现动态的逐 token 选择，同时保留完整的 KV-Cache。我们通过蒸馏模型的掩码因果注意力分布来训练 TokenButler，以最小的参数开销优化轻量级预测器。我们在一个新颖的合成小上下文共指检索任务上评估了 TokenButler，在现有方法失败的情况下展示了接近 oracle 的准确率。此外，TokenButler 在长上下文基准测试（RULER、LongBench）上取得了具有竞争力或更优的性能，使用我们提出的*带邻居获取的预测区间*在保持准确率在 1.1\% 范围内的同时，实现了高达 的 GPU 加速，并且与使用 CPU 卸载的 Dense Attention 相比，延迟降低了 7.6。代码可用：https://github.com/abdelfattah-lab/TokenButler
引用
@article{arxiv.2503.07518,
title = {TokenButler: Token Importance is Predictable},
author = {Yash Akhauri and Ahmed F AbouElhamayed and Yifei Gao and Chi-Chih Chang and Sameh Gobriel and Nilesh Jain and Mohamed S. Abdelfattah},
journal= {arXiv preprint arXiv:2503.07518},
year = {2026}
}