LOCKS: Page-Local Compact Key Summaries for Efficient Long-Context Decoding
2026-07-27 • Machine Learning
Machine LearningArtificial Intelligence
AI summaryⓘ
The authors address a speed and memory problem when using large language models on very long texts, caused by the need to look at all past data every time the model makes a new prediction. They propose LOCKS, a method that summarizes each page of text with a small, special summary and uses these summaries to decide which parts of the past to pay attention to, without reading everything. This approach performs almost as well as looking at all past data, but much faster and using less memory, especially on tasks requiring long reasoning. LOCKS can be added easily to existing systems and works well even on very long text inputs.
large language modelskey-value (KV) cachelong contextattention mechanismlow-rank approximationlog-sum-explong-document QAdecode latencyCUDA graphsbatched decoding
Authors
Junsung Hwang
Abstract
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: shared low-rank bases discard page-specific directions that a page's own compact basis retains. LOCKS gives every page its own spectral summary (resident, about a tenth the cache's size), 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), tracks the read-every-key oracle on retrieval-dense RULER down to the smallest budgets, and shows its largest margins on long-form reasoning (AIME26, MATH-500), where baseline selectors collapse. At its shipped $2048$-token budget LOCKS matches FullKV aggregate quality at $100$K$+$ context while attending about $2\%$ of the tokens, and halves per-token decode latency ($2.0\times$ at $1$M tokens) against dense attention. LOCKS ships as a drop-in plugin for unmodified vLLM, with batched decode running in full CUDA graphs.