BeaconKV: Key-Value Cache Compression Guided by Beacon Queries for Efficient Large Reasoning Model Inference

AI summary

Overview

Research area: Efficient inference for Large Reasoning Models (LRMs) — specifically, key-value (KV) cache compression for long Chain-of-Thought generation.

Technical level: Intermediate. The paper assumes familiarity with transformer attention, autoregressive decoding, KV caching, and positional embeddings (RoPE), but the core idea is geometric and intuitive.

Scope: This paper identifies a previously overlooked attention phenomenon in long-horizon reasoning ("Thought Revisiting Tokens") and exploits its geometric structure to build a training-free KV cache eviction method that cuts memory by up to 5.8× with near-lossless accuracy.

What This Paper Is About

Large Reasoning Models solve hard problems by generating very long chains of thought, but every generated token must be stored in a KV cache, so memory grows linearly with the reasoning trace — often exceeding a GPU's capacity. Existing compression methods decide which cached tokens to discard by looking only at the most recent queries, assuming recent attention is a good predictor of future attention. This paper shows that assumption breaks in long reasoning, and proposes a way to anticipate which old tokens will be needed again.

Key Contributions

1. Identification of Thought Revisiting Tokens (TRT): The authors show that certain decoding steps produce "global queries" that abruptly re-attend to distant earlier context — such as plans or problem constraints stated near the beginning of the trace — rather than to nearby tokens. These TRTs occur sporadically across many layers and heads, making them unpredictable for recent-query-based methods.

2. Discovery of geometric structure in global queries: Using pre-RoPE query states, they show that global queries are dissimilar to their neighbors but highly similar to each other, and that they cluster into a small number of groups in embedding space. This makes them representable by a compact set of "beacon queries."

3. BeaconKV, a training-free compression method: Beacon queries are combined with recent queries to score KV pairs, with differential RoPE alignment (beacons shifted to the current decoding step, recent queries left at their original positions) and max-pooling aggregation to preserve sparse but high-magnitude TRT signals.

4. Continual Farthest Point Sampling (FPS): A bounded-memory online algorithm that maintains beacon queries throughout decoding by filling a small per-head buffer and periodically downsampling it to the most geometrically distinctive queries, avoiding storage of the full query history.

Main Findings

• Recent-query baselines systematically evict reusable context: Because TRTs are sparse and event-driven, methods like RPC and R-KV assign low importance to distant tokens that are not currently attended, permanently discarding them before they are revisited.

• Accuracy gains over existing methods: BeaconKV achieves up to 31.7 percentage points higher accuracy than prior compression methods, with the largest gains in the low-budget regime (e.g., Qwen3-14B on AIME24 with a 1024-token KV budget).

• Memory and throughput: On Qwen3-4B with 32K-token generation on an A100 80GB, peak memory drops from 77.0 GB (full KV) to 13.3 GB — a 5.8× reduction — while throughput rises from 82.3 to 356.4 tokens/s (4.3×), enabling much larger batch sizes.

• Matched-budget efficiency parity: At a 2K KV budget (batch size 192), BeaconKV matches RPC in throughput and memory while improving LiveCodeBench accuracy by 6.3 points; at 1K (batch size 320), the gain is 12.3 points.

• Beacon queries must be balanced with recent queries: An ablation on Qwen3-4B/AIME24 shows (16 recent, 16 beacon) is the best trade-off. Over-allocating to beacons hurts both accuracy and latency (e.g., 1 beacon-31 recent gives 63.5% accuracy but 10,871s latency, versus 64.6% and 4,356s for 16/16).

• Max aggregation beats mean aggregation: Mean-pooling dilutes the sparse, high-magnitude TRT signals; max-pooling preserves them. The gap widens under tight budgets (e.g., 23.3 vs. 18.8 at a 256-token budget).

• Gains are not just from keeping early context: An "Initial+Recent" baseline that always retains the first queries and the most recent ones performs substantially worse, showing that the benefit comes from dynamically tracking emerging global query patterns.

• Online selection rivals offline ideals: Continual FPS achieves accuracy comparable to K-Means centroids or naive full-history FPS while using far less peak GPU memory.

Methodology in Plain English

The authors start by visualizing attention patterns during reasoning and notice two kinds of queries: "local" ones that look at nearby tokens, and "global" ones that jump far back to earlier parts of the trace. They confirm this is a real, recurring phenomenon across layers and heads, then measure how similar these queries are to one another. Global queries turn out to be outliers relative to their neighbors but close to each other in embedding space — so a handful of representative points can stand in for all of them.

Building on that, BeaconKV keeps two small sets of queries at eviction time: the most recent queries (for local fluency) and a set of "beacon" queries that summarize the global query clusters seen so far. To pick beacons cheaply, they use farthest point sampling with a bounded buffer that is periodically compressed. When the KV cache hits its limit, these queries are run against the cached keys, scores are taken with a max over queries and heads, and the top-scoring KV pairs are kept while the rest are discarded. Beacon queries have RoPE applied so they look like they are being asked now, which simulates whether a future revisiting token would need the cached content. The whole process needs no training or model modification.

Why This Matters

Impact on research: The paper challenges a foundational assumption in the KV cache compression literature — that recent queries proxy future attention. It provides a concrete counterexample in reasoning workloads and offers a new geometric lens (clustering of global queries) that other efficiency methods could build on, including sparse attention and memory-augmented architectures.

Real-world applications:

• On-premise and edge deployment of reasoning models: A 5.8× memory reduction can bring long-trace reasoning within reach of single-GPU or smaller-footprint servers rather than multi-GPU clusters.
• High-throughput reasoning services: Larger batch sizes and 4.3× throughput directly translate into lower cost per query for API providers serving math, coding, and science tasks.
• Agentic and multi-step pipelines: Agents that reason over long plans benefit when the model can still consult its early plan after thousands of tokens.
• Energy and cost efficiency: Longer traces on smaller hardware reduce the compute and energy footprint of large-scale inference.

Industry relevance: KV cache memory is one of the main cost drivers for reasoning-model serving. A training-free method that plugs into existing inference stacks without fine-tuning or architectural change is straightforward to adopt for teams already running vLLM-style servers, and the accuracy–efficiency trade-off is directly monetizable in capacity planning.

Future Directions

• Learned or hybrid beacon selection: FPS is training-free and effective, but a lightweight learned selector could adapt better across domains without the latency overhead seen in the most beacon-heavy configurations.
• Extension beyond reasoning workloads: The current evaluation focuses on math, coding, and science benchmarks; it is unclear whether the query-cluster structure appears in general long-context tasks like document QA or multi-turn dialogue.
• Combination with orthogonal techniques: Pairing BeaconKV with sparse attention kernels, adaptive reasoning-length control (early exit, summarization), or redundancy-based scoring like R-KV could compound memory savings.
• System-level optimization of beacon scoring: The paper notes a modest decoding-latency increase from beacon-query scoring; custom kernels or amortized scoring schedules could remove that overhead, and per-layer/per-head budget allocation remains unexplored.

Target Audience

This paper is most valuable to ML systems and inference engineers optimizing the serving of reasoning models, researchers working on KV cache compression, sparse attention, or long-context efficiency, and practitioners deploying open-source LRMs (DeepSeek-R1 distills, Qwen3) under tight GPU memory budgets. Readers with a basic understanding of transformer attention and autoregressive decoding will be able to follow the argument without difficulty; the geometric analysis is illustrated with clear visualizations.

Authors’ abstract

Large Reasoning Models (LRMs) achieve superior problem-solving through extended Chain-of-Thought (CoT) generation, but the resulting key-value (KV) cache grows linearly with sequence length and creates severe memory bottlenecks, often exceeding GPU capacity for long reasoning traces. Existing KV cache compression methods rely on recent queries to estimate future token importance, implicitly assuming these serve as reliable proxies for future attention patterns. We demonstrate that this assumption fails in long-horizon reasoning: certain decoding steps generate Thought Revisiting Tokens (TRT) that re-attend to distant previous context, such as task-solving plans formulated early in the trace. Through systematic analysis, we discover that queries corresponding to the TRT cluster into a small number of similarity groups in the embedding space. Based on this insight, we propose BeaconKV, a training-free KV cache compression method that maintains beacon queries, compact representatives for each global query cluster, to anticipate which KV pairs will be revisited without storing the entire query history. Across four open-source LRMs and diverse reasoning benchmarks, BeaconKV generally outperforms existing compression methods, achieving up to $5.8\times$ memory reduction while nearly preserving full cache accuracy and improving throughput by over $4.3\times$.