论文

LOCKS:用于高效长上下文解码的页面本地紧凑密钥摘要

LOCKS: Page-Local Compact Key Summaries for Efficient Long-Context Decoding

模型推理KV Cache

摘要

在长上下文中提供 大语言模型 的瓶颈是键值 (KV) 缓存,该缓存在每个解码步骤中都会读取。我们发现注意力键在页面内的排名大约较低。跨页面共享的单个低秩投影可能会错过特定于页面的方向;为每个页面拟合基础可以更好地识别在可比较的存储选择器成本下最受关注的页面。 LOCKS 存储每页的rank-$r$ 谱摘要,重建其页内logits,并通过log-sum-exp 质量选择页面,而不读取候选键或值。它在 LongBench-v1 上保持在 FullKV 的大约一个点内,在 RULER 上跟踪读取每个关键的精确 LSE 预言机直至最小的预算,并在 AIME26 和 MATH-500 上的紧张预算下最大限度地保持质量。在 $2048$-词元预算 上,它与超出 $100$K 上下文的 FullKV 聚合质量相匹配,同时参与约 $2\%$ 的词元。在 $2$-$8$ 的等级中,摘要使用 $4$-$10\%$ 的全 KV 字节。在具有 GPU 驻留 KV 的 GH200 上,LOCKS 在 $512$K 上下文中将完整解码步骤时间减少了 $1.8\times$。将全部 KV 卸载到 Grace 内存后,通过服务更大的批次,速度更快的密集后端的总吞吐量达到 $3.82$-$4.22\times$,达到 $64$K-$256$K。