EarlySelectKV: Anticipatory KV Selection via Cross-Layer Query Proxies

AI summary

Overview

Research area: Natural Language Processing / efficient long-context LLM inference — specifically dynamic sparse key-value (KV) cache selection and its systems-level cost.

Technical level: Advanced. The paper assumes familiarity with KV caching, attention, page-level sparse attention, and inference serving stacks (TensorRT-LLM, CPU offloading).

Scope in one sentence: The paper proposes EarlySelectKV, a training-free wrapper for query-conditioned sparse KV selectors that generates KV indices from an earlier, proxy query instead of waiting for the exact target-layer query, and measures the resulting accuracy and serving-time effects.

What This Paper Is About

Long-context LLM inference is expensive because the KV cache grows with context length, batch size, and depth. "Query-conditioned" sparse methods reduce this cost by keeping a candidate cache and selecting a small subset of KV entries at each decoding step based on the current query — but the selector must wait for the target-layer query before it can score candidates and produce indices, putting routing directly on the critical path before attention. EarlySelectKV asks a different question than existing work: not which KV entries to keep, but when their indices can be determined, so that routing can happen earlier while the real query still performs the final attention.

Key Contributions

1. A training-free anticipatory routing framework. EarlySelectKV derives a proxy query from a preceding hidden state, precomputes the selected KV indices in an earlier computation window, and leaves the real target-layer query untouched for value aggregation. No retraining or selector redesign is required.

2. A selector-preserving interface with configurable anchors and projections. The original candidate construction, routing view, budget, and scoring rules are kept intact; only the routing query is replaced. The paper supports multiple anchor points (ES-In, ES-Mid, and an offline-calibrated ES-Mix) and native or low-rank query projections.

3. Plug-and-play validation across selector families. The transformation is instantiated on RocketKV (the primary setting), Quest, and Loki, plus an integration into the released InfLLM-V2-Long-Sparse-Base checkpoint without retraining.

4. Routing timing as an explicit optimization target. Beyond accuracy parity on LongBench, NIAH, and RULER, the paper reports deployment-level throughput and GPU-memory effects in a TensorRT-LLM setup with CPU-offloaded routing, arguing that "when indices are produced" is a first-class design axis alongside "which KV entries are selected."

Main Findings

• Near-parity on LongBench for the primary instantiation: ES-RocketKV differs from RocketKV by +0.018 points on average; at the 128-token budget the difference is 0.03 points, and across five backbones (7B–32B) every mean stays within 0.1 points.

• Near-parity on RULER: Across RULER-8K, 16K, and 32K, the mean ES-minus-RocketKV difference is approximately −0.01 points, and ES-RocketKV follows RocketKV's degradation as context length grows.

• Stable retrieval under low budgets on NIAH/Passkey: The mean difference across the representative settings in Table 1 is +0.25 percentage points. Results are identical on Llama and Mistral models; on Qwen2.5-7B-Instruct at budget 128, RocketKV scores 64.0 versus 66.0 for ES-RocketKV, and at budget 64 both score 48.0 against a TopK reference of 76.0.

• Generalization to other selectors is positive but bounded: Across eight Quest/Loki comparisons on two 7B backbones, the mean ES-minus-base difference is −0.07 points. Individual deltas range from −0.86 (Mistral-7B, Loki, budget 128) to +0.34 (Mistral-7B, Loki, budget 4096).

• Integration after sparse-attention training works without retraining: On InfLLM-V2-Long-Sparse-Base, exact W_Q changes the four-task average by +0.26 points, while a rank-256 approximation changes it by −0.28 points.

• Proxy fidelity is layer-dependent: Q_mid gives higher signed top-r channel recall and retained attention mass in early and middle layers. The retained-mass advantage of Q_mid over Q_in is 0.0220 before layer 20 and 0.0003 afterward, which is the motivation for ES-Mix.

• ES-Mix extends the preparation window without improving quality: Across 21 LongBench and 30 RULER settings, ES-Mix differs from ES-Mid by −0.006 points on average. Cross-domain calibration profiles are stable, with pairwise layer-gap Pearson correlations of at least 0.9724 and assignment agreement of at least 96.8%.

• Routing fidelity does not degrade along a fixed decode trajectory: On 20 MultiNews trajectories over decode positions 0 to 512, the retained-mass gap stays below 0.002 in magnitude, and page overlap and signed top-r recall trend intervals both include zero. On independent MATH-500 generation, ES-RocketKV stays within 1.0 percentage point of RocketKV on both tested models.

• Measured throughput gains grow with output length: On a 40K-token input with a 4096-token Stage-1 budget, a 128-token Stage-2 budget, and batch size 1, end-to-end output throughput improves from 1.3% at 128 output tokens to 12.4% at 8192 output tokens. Decode-only gains rise from 7.67% to 13.14%; with 8192 output tokens, end-to-end gains remain between 10.5% and 12.1% across batch sizes 1 to 8, while the short-output benefit approaches zero at batch 8.

• GPU memory reduction of about 5.3 GiB: This corresponds to 11.1% of the KV and K-summary pool, and reflects the CPU-offloaded deployment choice rather than an intrinsic property of the selector.

• Low-rank proxy projections trade cost for fidelity: On Mistral-7B-Instruct-v0.2 at budget 512, original RocketKV averages 41.96; rank 64 gives 41.32, rank 128 gives 41.82, rank 256 gives 42.09, and rank 512 gives 42.03 — saturating around rank 256. Rank 128 is used in the CPU-routing experiments.

Methodology in Plain English

The core observation is simple: generating sparse-attention indices requires a query, but it does not require the exact target-layer query. Routing only needs to preserve the relative preference among candidate KV pages, whereas final attention needs full query fidelity.

EarlySelectKV splits decoding into two flows. The main flow behaves like normal sparse attention: it reads already-computed page indices, gathers the selected KV entries, and runs attention with the real query. The side flow runs earlier: it builds a proxy query from a hidden state produced before the target layer, runs the original selector's scoring rule over the same routing view and budget, and stores the resulting indices so they are ready when the next layer needs them.

Two anchor points are studied. ES-In uses the input hidden state of the current block, which is available very early and therefore gives a long preparation window for routing, metadata lookup, and KV preparation, at the cost of a weaker proxy. ES-Mid uses the post-attention residual state of the current block, which already contains attention-level context, giving a better proxy but a shorter window; it is the conservative default. ES-Mix profiles layers offline: it measures how much real-query attention mass each anchor retains, computes the gap between the ES-Mid and ES-In scores per layer, and assigns ES-In wherever that gap is within a tolerance threshold (τ = 0.01), keeping ES-Mid elsewhere. The assignment is fixed at inference time.

For evaluation, the authors compare RocketKV against ES-RocketKV on four instruction-tuned models (Llama-3.1-8B-Instruct, Llama-3-8B-Instruct, Mistral-7B-Instruct-v0.2, Qwen2.5-7B-Instruct) using LongBench, NIAH, and RULER, with a default sweep of 128 to 4096 selected tokens and a budget of 64 included for NIAH and anchor diagnostics. All sparse methods keep layer 0 dense and apply sparse selection from layer 1 onward; ESelector reserves the local budget outside routing and keeps the most recent tail page. Proxy fidelity is diagnosed with signed top-r channel recall and retained attention mass rather than exact query reconstruction, since only the routing signal is being replaced. Runtime is measured on a single NVIDIA RTX 6000D GPU with the TensorRT-LLM PyTorch backend and CUDA Graphs, using FP8 K-summary caches and an Intel Xeon Platinum CPU handling proxy construction, routing, and K-summary state.

Why This Matters

The paper reframes sparse long-context decoding as a scheduling problem in addition to a selection problem. Most work in this space optimizes the selector's accuracy or its cost; EarlySelectKV points out that even a cheap router blocks attention if it sits on the immediate pre-attention path, and that routing's share of attention-side latency grows as the token budget shrinks. Moving index generation earlier turns an exposed serial dependency into something that can be overlapped — and makes CPU-offloaded routing practical, since the indices exist early enough to hide host-side work.

Real-world applications that benefit:

• Long-document understanding and summarization over inputs where KV cache size forces aggressive sparsification.
• Repository-level code completion, where context windows are long and per-token latency is user-visible.
• Retrieval-augmented generation over large corpora, where the active evidence set shifts across decoding steps.
• Long-context reasoning and math, where the paper reports ES-RocketKV staying within 1.0 percentage point of RocketKV on MATH-500.

Industry relevance: the measured gains are reported in a production-oriented stack (TensorRT-LLM, CUDA Graphs, FP8 summaries) with batch-size and output-length sweeps, and the memory result (about 5.3 GiB, 11.1% of the KV and K-summary pool) follows directly from being able to place routing state on the CPU. For serving operators, that is a concrete lever on both throughput and GPU residency — though the paper is explicit that the gain is a deployment-level effect, not an isolated proxy-query microbenchmark.

Future Directions

• Co-training the proxy with the model. The authors note that the InfLLM-V2 experiment tests integration after sparse-attention training, and suggest incorporating EarlySelectKV during pretraining or long-context adaptation so hidden states, query projections, and routing behavior can co-adapt, potentially improving proxy fidelity.

• Extending beyond a single-GPU, GPU-resident-KV deployment. The runtime study covers one TensorRT-LLM configuration; multi-GPU serving, KV-page prefetch or offloading, and larger-model runtime scaling are untested, so measured gains may not transfer.

• Adaptive rather than fixed layer profiles. ES-Mix uses a fixed offline layer profile and may need re-profiling under substantial domain shifts. Qwen2.5-7B-Instruct is a concrete open case: its low-gap layers do not form a contiguous late-layer region, so ES-Mix results are not reported for it.

• Routers where the window may not be enough. The limitations section flags that gains may shrink for short outputs, fine-grained or expensive routers, large batches, and resource-contended execution — leaving open how much proxy work can realistically be hidden for heavier selectors.

Target Audience

Inference systems engineers and serving-infrastructure teams working on long-context deployment; researchers in KV-cache compression, eviction, and dynamic sparse attention who want a selector-agnostic technique that does not require retraining; and graduate students already comfortable with attention mechanics who want a clear example of how a timing-level reframing can produce measurable systems gains at near-zero accuracy cost. Readers without background in sparse attention or LLM serving will find the systems portions harder to place in context than the accuracy comparisons.

Authors’ abstract

Query-conditioned sparse KV selection reduces long-context inference cost by attending to a subset of historical KV entries at each decoding step. Existing selectors must wait for the target-layer query before scoring candidates and generating indices, placing routing on the immediate pre-attention path. We introduce EarlySelectKV, a training-free transformation that uses proxy queries from earlier hidden states to precompute selected KV indices while retaining the real target-layer query for final attention. Tests with RocketKV, Quest, and Loki show no consistent pattern of quality degradation across the evaluated settings. In the primary RocketKV evaluation, ES-RocketKV differs by +0.018 points on LongBench, approximately -0.01 points on RULER, and +0.25 percentage points on representative NIAH/Passkey settings. In a TensorRT-LLM deployment with CPU-offloaded routing, ES-RocketKV improves end-to-end output throughput by up to 12.4% and reduces GPU-resident memory by 11.1% of the combined KV/K-summary pool. This makes routing timing an explicit optimization target alongside KV subset selection in sparse long-context decoding. Code is available at https://github.com/sea-with-sakura/EarlySelectKV.