跳转至

LOCKS：用于高效长文本解码的页局紧凑键摘要

文章背景与核心概要

在长文本大语言模型（LLM）的服务部署中，每个解码步都需要完整读取键值（KV）缓存，这成为了主要的性能瓶颈。为了解决这一痛点，本文提出了名为 LOCKS（Page-Local Compact Key Summaries，即“页局紧凑键摘要”）的全新插件与优化框架。该方法的核心洞察在于注意力键具有“局部低秩、全局高秩”的特性：共享的全局低秩草图会丢失特定页面的信息，而页面局部的秩-\(r\)（rank-\(r\)）谱摘要则能更好地保留关键载体。

通过为每个页面分配独立的秩-\(r\)谱摘要（在 \(r=8\) 时仅占缓存的十分之一，在 \(r=2\) 时占二十五分之一），LOCKS 能够重建页内逻辑、通过 log-sum-exp 估计注意力权重，并在无需在选择阶段读取原始候选键或值的情况下，仅对顶层页面执行注意力计算。实验表明，LOCKS 在长文档问答（LongBench-v1）和检索任务（RULER）上能够匹配全缓存基线的准确率；在长文本推理任务（AIME26、MATH-500）中，它也优于现有的选择器和基于淘汰机制的压缩器。在单张 H200 NVL 上处理 1M 长度文本时，它实现了每步 KV 读取量 10 到 25 倍的减少，将每个 Token 的解码延迟减半（加速比达 \(2.0\times\)），同时在 \(100\text{K}+\) 的上下文长度下保持了与 FullKV 相当的质量。该框架可作为未经修改的 vLLM 的即插即用插件，完美支持在完整 CUDA 图（CUDA graphs）中运行的批处理解码。

Summary

LOCKS (Page-Local Compact Key Summaries) is a novel plugin and optimization framework designed to accelerate long-context decoding in Large Language Models (LLMs) by addressing the performance bottlenecks of Key-Value (KV) caches.

LOCKS (Page-Local Compact Key Summaries) is a novel plugin and optimization framework designed to accelerate long-context decoding in Large Language Models (LLMs) by addressing the performance bottlenecks of Key-Value (KV) caches.

Key Innovations & Findings:

  • The Low-Rank Dilemma: Attention keys are locally low-rank but globally high-rank. While shared global low-rank sketches miss page-specific information, page-local rank-\(r\) spectral summaries preserve crucial carriers.
  • Efficient Attention Routing: LOCKS reconstructs within-page logits, estimates attention mass via log-sum-exp, and selectively attends to top pages without reading raw candidate keys or values during the selection phase.
  • Performance & Accuracy:
  • Matches full-cache baseline accuracy on long-document QA (LongBench-v1) and retrieval tasks (RULER).
  • Outperforms existing selectors and eviction-based compressors on long-form reasoning tasks (AIME26, MATH-500).
  • Achieves a 10–\(25\times\) reduction in per-step KV reads and halves per-token decode latency (\(2.0\times\) speedup at \(1\text{M}\) tokens on a single H200 NVL) while matching FullKV quality at \(100\text{K}+\) contexts.
  • Seamless Integration: Ships as a drop-in plugin for unmodified vLLM, supporting batched decoding running entirely in CUDA graphs.

Key Innovations & Findings:

  • The Low-Rank Dilemma: Attention keys are locally low-rank but globally high-rank. While shared global low-rank sketches miss page-specific information, page-local rank-\(r\) spectral summaries preserve crucial carriers.
  • Efficient Attention Routing: LOCKS reconstructs within-page logits, estimates attention mass via log-sum-exp, and selectively attends to top pages without reading raw candidate keys or values during the selection phase.
  • Performance & Accuracy:
  • Matches full-cache baseline accuracy on long-document QA (LongBench-v1) and retrieval tasks (RULER).
  • Outperforms existing selectors and eviction-based compressors on long-form reasoning tasks (AIME26, MATH-500).
  • Achieves a 10–\(25\times\) reduction in per-step KV reads and halves per-token decode latency (\(2.0\times\) speedup at \(1\text{M}\) tokens on a single H200 NVL) while matching FullKV quality at \(100\text{K}+\) contexts.
  • Seamless Integration: Ships as a drop-in plugin for unmodified vLLM, supporting batched decoding running entirely in CUDA graphs.

Document Metadata

  • arXiv ID: arXiv:2607.24555 [cs.LG]
  • Authors: Junsung Hwang
  • Primary Subject: Machine Learning (cs.LG)
  • Secondary Subjects: Artificial Intelligence (cs.AI)
  • Submission History:
  • v1: July 27, 2026
  • v2: August 24, 2026 (Current Version)
  • License: Creative Commons Attribution 4.0 license icon

Document Metadata

  • arXiv ID: arXiv:2607.24555 [cs.LG]
  • Authors: Junsung Hwang
  • Primary Subject: Machine Learning (cs.LG)
  • Secondary Subjects: Artificial Intelligence (cs.AI)
  • Submission History:
  • v1: July 27, 2026
  • v2: August 24, 2026 (Current Version)
  • License: Creative Commons Attribution 4.0 license icon

Abstract

在长文本下服务大语言模型时，其瓶颈在于键值（KV）缓存，该缓存需要在每个解码步被完整读取。注意力键呈现局部低秩但全局高秩的特点：跨页面共享的固定低秩草图被证明无法感知特定页面的方向，而在相同的摘要大小下，页面自身的基底能更好地对页面进行排序并保留载体。

LOCKS 为每个页面赋予其专属的秩-\(r\)谱摘要（常驻内存，在 \(r=8\) 时为缓存的十分之一，在 \(r=2\) 时为二十五分之一），重建页内逻辑，通过 log-sum-exp 估计每个页面的注意力权重，并仅对顶层页面计算注意力；选择过程本身不读取任何候选键或值。仅凭此摘要进行选择，在长文档问答（LongBench-v1；Llama-3.1-8B）上与全缓存相差约一个百分点；在密集检索的 RULER 任务上，它甚至在最小预算下也能紧跟“读取每一个键”的精确 LSE 预言机（oracle）；在长文本推理任务（AIME26、MATH-500；Qwen3-4B）的小预算区间内，它能最大程度地保持模型质量，而此时选择器和基于淘汰机制的推理压缩器均已失效。

在 \(2048\) 个 Token 的预算下，LOCKS 在 \(100\text{K}+\) 上下文（GLM-4-9B-Chat-1M）下能够匹配 FullKV 的总体质量，同时仅需关注约 \(2\%\) 的 Token；由于每一步都会完整扫描摘要，在该秩范围内，每步的 KV 读取量下降了约 \(10\)--\(25\times\)，与密集注意力相比，这使每个 Token 的解码延迟减半（在单个 H200 NVL 上针对 \(1\)M Token 测得 \(2.0\times\) 的加速，于 \(r=8\) 下测量）。LOCKS 作为未修改版 vLLM 的即插即用插件发布，其批处理解码完全在 CUDA 图中运行。

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.

Access & Resources

Access & Resources