🔥 今日值得一读
单令牌窥探就能锁定重要注意力头，PeekKV超半数场景准确率登顶！
PeekKV: The First Content Token Already Knows Which Heads Matter
👨💼 主理人解读 · 为什么值得关注
LLM长上下文显存吃紧时，可用首token信号动态选重要头压缩KV缓存，跨负载更稳定。
📖 AI 总结
该论文针对大语言模型推理中KV缓存随上下文长度线性增长导致的内存瓶颈问题，提出PeekKV方法。作者发现，现有按注意力头非均匀压缩KV缓存的方案存在精度波动，原因在于活跃注意力头在预填充与解码阶段之间会发生转移，但在解码路径中，过滤掉格式token后，活跃头组成在内容token间保持稳定。基于此，PeekKV在预填充后仅观察一个内容token，即“单token窥探”，同时获得头组得分与token重要性得分，据此按比例分配各头组的KV缓存预算，并在每组内仅保留得分最高的token。该方法无需离线剖析或训练，可直接应用于未见工作负载和新模型。在四个公开基准上，PeekKV在超过半数的按模型比较中取得最高精度，并在RULER和LongBench上最接近全缓存精度；在Llama-3.1-8B-Instruct上，16K上下文、10%内存预算时保留95.7%的全缓存精度，LongBench上1.5%预算时保留95.9%。
🔑 关键词速览
| KV cache | 键值缓存，LLM推理时存储注意力键和值以加速生成的机制，其大小随上下文长度线性增长。 |
| attention heads | 注意力头，Transformer中并行处理不同表示子空间的机制，不同头具有不同功能角色。 |
| prefill stage | 预填充阶段，LLM处理输入提示并生成初始KV缓存的阶段。 |
| decode stage | 解码阶段，LLM自回归生成输出令牌的阶段。 |
| PeekKV | 本文提出的方法，通过预填充后单令牌窥视来联合评估头组和令牌重要性，实现非均匀KV缓存压缩。 |
📱 每天一份 AI 前沿日报
关注公众号，每天 09:00 推送 · 不错过任何重磅