用局部压缩键值提升长文本生成效率,大幅降低计算开销。
LOCKS: Page-Local Compact Key Summaries for Efficient Long-Context Decoding
- 为每页独立构建低秩键摘要,避免全局压缩丢失页面特异性。
- 在2048令牌预算下仅读取2%的令牌,质量接近全量缓存,延迟减半。
- 可直接插入vLLM框架,适用于长文档问答、推理等场景。
长上下文语言模型服务受限于键值(KV)缓存的全量读取。尽管注意力键整体高秩,但局部呈低秩特性:共享的全局低秩草图会忽略页面特异性方向,而每个页面自建的低秩基能更好保留关键信息。LOCKS为每页生成独立的秩-$r$谱摘要($r=8$时缓存占用为原大小的十分之一,$r=2$时为二十五分之一),在页内重构逻辑,通过log-sum-exp估算各页注意力权重,仅关注前几页;选择过程无需读取候选键值。仅基于摘要选择即可在长文档问答(LongBench-v1;Llama-3.1-8B)上保持约1点误差,检索密集任务(RULER)中逼近读取所有键的精确-LSE基准,且在小预算下长期推理(AIME26, MATH-500;Qwen3-4B)表现最优。在2048令牌预算下,对100K+上下文(GLM-4-9B-Chat-1M)的聚合质量媲美全量缓存,仅需读取约2%的令牌;因摘要全步扫描,每步KV读取量减少10–25倍,解码延迟降低一半(单个H200 NVL上$1$M令牌时达$2.0\times$加速)。LOCKS可作为不修改vLLM的即插即用模块,支持批处理解码并运行在完整CUDA图中。
原文摘要 · Abstract (English)
Serving large language models at long context is bottlenecked by the key-value (KV) cache, which is read in full at every decode step. Attention keys are locally low-rank though globally high-rank: a fixed low-rank sketch shared across pages is provably blind to page-specific directions, while at the same summary size a page's own basis ranks pages and keeps carriers far better. LOCKS gives every page its own rank-$r$ spectral summary (resident, a tenth of the cache at $r{=}8$ and a twenty-fifth at $r{=}2$), reconstructs within-page logits, estimates each page's attention mass by log-sum-exp, and attends only the top pages; selection itself reads no candidate keys or values. Selecting on this summary alone stays within about a point of the full cache on long-document QA (LongBench-v1; Llama-3.1-8B), tracks the read-every-key exact-LSE oracle on retrieval-dense RULER down to the smallest budgets, and holds quality furthest into the small-budget regime on long-form reasoning (AIME26, MATH-500; Qwen3-4B), where selectors and eviction-based reasoning compressors both fall away. At a $2048$-token budget LOCKS matches FullKV aggregate quality at $100$K$+$ context (GLM-4-9B-Chat-1M) while attending about $2\%$ of the tokens; since the summary is scanned in full each step, the per-step KV read falls by about $10$--$25\times$ across that rank range, and this halves per-token decode latency ($2.0\times$ at $1$M tokens on one H200 NVL, measured at $r{=}8$) against dense attention. LOCKS ships as a drop-in plugin for unmodified vLLM, with batched decode running in full CUDA graphs.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。