KV selection stops waiting for the target-layer query: EarlySelectKV precomputes indices from proxy queries, lifting long-output throughput by up to 12.4%
Synopsis
Query-conditioned sparse KV selection must wait for the target-layer query before scoring candidates and generating indices, so routing sits on the immediate pre-attention path. EarlySelectKV replaces the routing query with a proxy built from an earlier hidden state, precomputing the KV indices the next layer will access while final attention still uses the real query. On RocketKV it differs by +0.018 points on LongBench, about −0.01 points on RULER and +0.25 percentage points on representative NIAH/Passkey settings; in a TensorRT-LLM deployment with CPU-offloaded routing, end-to-end output throughput rises by up to 12.4% and GPU-resident memory falls by 11.1% of the combined KV/K-summary pool.
Interpretation
In sparse long-context decoding, the routing step that decides which KV entries to read moves off the immediate pre-attention path, so selected indices can be prepared in an earlier computation window. Conventional query-conditioned selectors had to wait for the target-layer query before scoring candidates, generating indices and preparing the corresponding KV entries, leaving that serial dependency exposed ahead of sparse attention. On RocketKV, ES-RocketKV differs from native RocketKV by +0.018 points averaged over LongBench, about −0.01 points on RULER and +0.25 percentage points on representative NIAH/Passkey settings, indicating no consistent quality degradation from early routing.
The design splits dynamic KV selection into a main flow and a side flow: the main flow gathers KV entries using precomputed page indices and runs sparse attention with the real target-layer query. The side flow runs before attention reaches the target layer, builds a proxy query from a preceding hidden state through the query projection, and invokes the original selector to score and top-k over the K-summary cache, producing next-layer page indices. Two proxy anchors are studied: ES-In takes the current block input and offers a longer preparation window, while ES-Mid takes the post-attention intermediate state and gives a stronger proxy; ES-Mix assigns one per layer by offline calibration and averages 0.006 points below ES-Mid across 21 LongBench and 30 RULER settings.
Making indices available early also changes deployment: proxy-query construction, routing and K-summary state can move to the CPU and finish before target attention. On a single RTX 6000D with the TensorRT-LLM backend, 40K-token inputs and a 128-token Stage-2 budget, output throughput over RocketKV grows as outputs get longer. End-to-end output throughput gains rise from 1.3% at 128 output tokens to 12.4% at 8192 output tokens; with CPU-offloaded routing, GPU-resident memory drops by about 5.3 GiB on an allocation basis, 11.1% of the KV and K-summary pool.
Perspective
The demonstrated scope is quality evaluation on LongBench, NIAH/Passkey and RULER across four 7B–8B instruction models, plus runtime measurement on a single RTX 6000D with the TensorRT-LLM PyTorch backend, CPU-offloaded routing and GPU-resident KV; gains depend on how much proxy construction, routing and metadata preparation fit inside the available preparation window.
Runtime gains may be smaller for short outputs, large batches and resource-contended execution, falling to 0.38% at batch 8 with 128 output tokens; ES-Mix uses a fixed offline layer profile that may need re-profiling under substantial domain shifts; the InfLLM-V2 integration covers one checkpoint without runtime measurements, and long-generation evaluation is limited to the tested models and horizons.