🤖 machine learning

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

LOCKS is a drop-in plugin for vLLM that accelerates long-context decoding by assigning each memory page a compact, low-rank spectral summary to efficiently estimate attention mass and select only the most relevant pages, thereby significantly reducing latency and memory usage while maintaining near-full-attention accuracy.

Original authors: Junsung Hwang

Published 2026-07-28

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Junsung Hwang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to read a massive library of books to answer a single question. In the world of artificial intelligence, Large Language Models (LLMs) are like brilliant students who have read the entire internet, but when they try to answer a question based on a very long document, they face a tricky problem. To think, they need to keep a "memory" of everything they've read so far. This memory is called the KV cache (Key-Value cache). Think of it as a giant whiteboard where the model writes down every single word it has processed.

The problem is that as the story gets longer, this whiteboard gets huge. Every time the model wants to write the next word, it has to scan the entire whiteboard to decide which past words are important. If the story is 100,000 words long, the model has to look at 100,000 words every single time it types one new letter. This is slow and uses up a massive amount of computer memory, like trying to find a specific needle in a haystack by moving the whole haystack every time you blink. Scientists have been trying to figure out how to make the model ignore the boring parts of the story and only look at the exciting parts, but they've struggled to do this without losing the ability to find the right answer.

This is where a new method called LOCKS comes in. The researchers behind this paper discovered a clever trick to speed things up without losing the plot. They realized that while the whole story is complex, small chunks of it (called "pages") have their own simple, unique patterns. Instead of trying to summarize the whole library with one giant, messy map, LOCKS gives every single page its own tiny, high-quality "spectral summary."

Think of it like this: Imagine you are a detective solving a mystery in a 1,000-page novel. Instead of reading every word of every page to find the killer, you create a tiny, 10% size "cheat sheet" for each page. This cheat sheet doesn't just list the words; it captures the vibe and the most important directions of that specific page's content. When the detective (the AI) needs to know where to look next, it doesn't read the full pages. It just glances at these tiny cheat sheets to see which pages have the most "clues" (attention mass).

The paper shows that this method is incredibly effective. By using these page-specific cheat sheets, the model can skip reading about 98% of the text in a 100,000-token context, yet it still finds the right answer almost as well as if it had read everything. In fact, on difficult math and reasoning tests, other methods that try to guess which pages are important often fail completely, but LOCKS keeps the "carrier" pages—the ones that actually hold the answer—safe and sound.

The researchers proved that trying to use one single map for the whole book (a "shared" summary) doesn't work because different pages have different secrets that get lost in the mix. They also showed that their method is "training-free," meaning it works with existing AI models without needing to re-teach them anything. When they tested it on real hardware, they found that it cut the time it takes to generate each word in half for very long documents. It's like turning a slow, plodding walk through a library into a high-speed teleportation system that only stops at the shelves that actually matter.

In short, LOCKS solves the "long-context" bottleneck by realizing that every page of a story has its own unique fingerprint. By creating a compact, page-specific summary for each one, the AI can instantly know which pages to read and which to ignore, making it possible to chat with models about books that are hundreds of thousands of words long without the computer getting overwhelmed.

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

1. Problem Statement

Serving Large Language Models (LLMs) at long context lengths is increasingly bottlenecked by the Key-Value (KV) cache, specifically by the KV-read bandwidth rather than memory capacity. In production regimes with high concurrency, the model weights are amortized across a batch, but every request must read its own unique KV cache. For models like Llama-3.1-8B, a 128K context requires ~16 GiB of KV cache per request, exceeding the model's weight size. Consequently, the cache, not the weights, limits the number of requests an accelerator can serve.

While attention is globally high-rank, it is locally low-rank: a few pages (blocks of tokens) carry almost all of a query's attention mass at any decode step. Sparse decoding aims to attend only to these critical pages. However, existing selection methods face a trade-off:

Exact selection (reading every key) defeats the purpose of sparsity by incurring full bandwidth costs.
Resident summaries (compact representations of pages) are required to avoid reading full keys, but existing approaches often fail to preserve the specific geometric structure needed to identify "carrier" pages (those containing the answer) under aggressive budget constraints. Shared projections across pages are provably blind to page-specific content, leading to task collapse in retrieval and reasoning scenarios.

2. Methodology: LOCKS

LOCKS (Page-Local Compact Key Summaries) addresses this by maintaining the full KV cache resident but avoiding the read of most keys during decoding. It operates on a page-local spectral summary that is training-free and integrated as a drop-in plugin for vLLM.

Core Mechanism

Page-Local Spectral Decomposition:
- Instead of using a shared global or sequence-level basis, LOCKS constructs a unique low-rank basis for each page of the KV cache.
- When a page is finalized (during prefill or a page-crossing refresh), it undergoes a small eigendecomposition of its key deviations (centered keys).
- The summary stores the top- $r$ eigenvectors (basis) and the corresponding coefficients for each token in the page.
- Quantization: The basis is stored in int4 and coefficients/centroids in int8. At rank $r=8$ with page size $B=16$ , the summary is approximately 10% of the page's original KV bytes.
Reconstruction and Selection:
- At every decode step, LOCKS reconstructs the within-page logits for the query using the compact summary:
  $\hat{s}_j(q) = \log \sum_{i \in P_j} \exp\left( \frac{q^\top \mu_j + (V_j^\top q)^\top c_i}{\sqrt{d}} \right)$
  where $\mu_j$ is the page centroid, $V_j$ is the page-local basis, and $c_i$ are the coefficients.
- It computes the Log-Sum-Exp (LSE) of these reconstructed logits to estimate the page's attention mass.
- Selection is performed by ranking pages based on these estimated masses. Crucially, no candidate keys or values are read during this selection phase.
Group-Shared Selection (GQA):
- For Grouped-Query Attention (GQA), LOCKS averages the estimated mass shares across the query heads in a group to form a single shared selection table. This ensures all heads in a group attend to the same set of pages, avoiding traffic fragmentation.
Implementation:
- Implemented as a plugin for unmodified vLLM.
- The entire decode step (scoring, selection, and sparse attention) runs inside full CUDA graphs, ensuring zero overhead from Python dispatch and enabling efficient kernel fusion.
- The summary construction is decoupled from the decode critical path, occurring only when pages are finalized.

3. Key Contributions

Theoretical Findings

Locality of Structure: The paper demonstrates that the geometric structure relevant for page selection is concentrated at the page scope, not the sequence or cache scope. A page's own rank-8 eigenbasis captures the majority of its key energy, whereas shared bases (even when matched in state size) fail to capture page-specific directions.
Impossibility of Shared Projections: The authors prove (Proposition 1) that any fixed, query-independent shared low-rank linear projection is content-blind to a $(d-\rho)$ -dimensional family of within-page differences. This means shared bases cannot distinguish between pages that differ only in specific content directions, leading to selection errors.
Retention Guarantees: The paper provides a composed theoretical bound (Theorem 1) showing that the error in attention output is directly controlled by the spectral tail energy of the page-local reconstruction. The page-local basis is proven to be residual-optimal among orthogonal rank- $r$ reconstructions.

Empirical Results

Retrieval and QA (LongBench-v1, RULER): LOCKS tracks the "read-every-key" oracle (exact LSE selection) down to the smallest budgets. On retrieval-dense tasks (RULER), it significantly outperforms state-of-the-art selectors (Quest, KVzip, ShadowKV, RocketKV) in carrier retention (keeping the pages that contain the answer), which is the primary driver of task success.
Long-Form Reasoning (AIME26, MATH-500): In reasoning tasks where the model generates thousands of tokens, baseline selectors often collapse at tight budgets. LOCKS maintains high accuracy, matching FullKV performance at 100K+ contexts while attending to only ~2% of tokens.
Efficiency:
- At 1M context length, LOCKS achieves 2.0× speedup in per-token latency compared to dense attention (FlashAttention-3/FlashInfer).
- It reduces KV bytes read per decode step by 9.8× at 1M context.
- Throughput speedup scales with batch size, reaching 1.80× at 256K context with large batches.

4. Significance and Claims

The paper claims that page-local spectral concentration is the missing link for efficient sparse decoding. By shifting the scope of representation from global/sequence to local (per-page), LOCKS preserves the "carrier" pages necessary for task completion without requiring training or modifying the model architecture.

Drop-in Integration: Unlike many prior works that require model forks or training gates, LOCKS works with unmodified vLLM and standard GQA models.
Theoretical-Practical Alignment: The method is one of the few that connects a theoretical bound on reconstruction error directly to end-to-end task quality and carrier survival, validating that the "chain" from spectral concentration to task performance holds.
Scalability: The approach scales to extreme contexts (1M+ tokens) where memory bandwidth is the primary bottleneck, offering a practical path to serving long-context models without offloading to CPU or tiered memory.

The authors remain modest regarding limitations: the deployed quantized summary's retention is measured empirically rather than strictly bounded per step, and the method currently assumes a fixed page size. However, the results establish that a training-free, page-local spectral summary is sufficient to track the oracle selection down to aggressive sparsity budgets, enabling practical long-context acceleration.