ParisKV: Fast and Drift-Robust KV-Cache Retrieval for Long-Context LLMs

AI summary

Overview

Research area: Efficient LLM inference — specifically KV-cache retrieval and sparse attention for long-context decoding.

Technical level: Advanced. The paper assumes familiarity with transformer attention, KV caching, approximate nearest-neighbor search, quantization, and GPU systems programming.

Scope: A GPU-native KV-cache retrieval framework that uses data-independent analytic centroids on a rotated unit hypersphere to remain stable under decoding-time distribution drift while enabling million-token contexts via CPU offloading.

What This Paper Is About

Long-context LLM decoding is bottlenecked by the KV cache, which grows linearly with context and dominates both memory and latency. Sparse attention methods try to retrieve only the most relevant past tokens, but existing retrieval approaches fail in two ways: the clustering centroids they learn during prefill become stale as generation proceeds (distribution drift), and they rely on CPU-side search plus CPU-to-GPU data movement that adds substantial latency. ParisKV solves both by making retrieval decisions in a stable, data-independent representation space and keeping the entire retrieval path on the GPU.

Key Contributions

1. Drift-robust analytic centroids. A data-independent transformation — L2 normalization onto the unit hypersphere plus a shared random orthogonal rotation (SRHT) — followed by a fixed set of sign-pattern centroids. Because the centroids are uniformly distributed and never learned from data, they cannot become stale as decoding continues, so recall does not collapse over long generations.

2. GPU-native coarse-to-fine retrieval pipeline. A coarse stage uses multi-subspace collision counting (integer vote accumulation) to prune candidates to roughly 5–10% of the cache; a fine stage reranks those candidates using a calibrated inner-product estimator built from 4-bit quantized key codes plus a precomputed per-key scaling factor. Ranking is done entirely on compressed codes — no full-precision keys are needed on the GPU during selection.

3. Scalable UVA-based offloading system. The full-precision KV cache lives in CPU memory and is fetched on demand via Unified Virtual Addressing, letting GPU kernels read only the selected top-k pairs without explicit memcpy calls or CPU-side scheduling. Four custom CUDA kernels (bucket_topk, parallel collision, fused reranking, UVA fetch) eliminate sorting overhead and minimize kernel launches.

4. Comprehensive evaluation across three model families showing accuracy that matches or exceeds full attention in 7 of 9 settings, and decode-latency reductions of up to 17× (vs. MagicPIG) and 44× (vs. PQCache) at million-token scale.

Main Findings

• Accuracy matches or beats full attention. On Qwen-3-4B, ParisKV scores 72.22 on GPQA-Diamond and 92.80 on MATH500, exceeding PQCache by +33.84 and +34.00 points respectively. Across all settings it outperforms both retrieval baselines and matches or exceeds full attention in 7/9 configurations.

• Especially strong on long-generation reasoning. On AIME25 pass@8, ParisKV improves over PQCache/MagicPIG by +40.0 to +76.7 points. The paper argues this is because drift compounds over tens of thousands of decoding steps, making long-generation tasks the hardest test for retrieval stability.

• Drift ablation confirms the mechanism works. Normalization + rotation + theoretical spherical centroids raises coarse-stage Recall@100 from 6% to 16.1%, and end-to-end Recall@100 after exact reranking from 36.5% to 64.3%.

• Better batch scaling than full attention. Across 64K/128K/256K contexts, ParisKV delivers 2.1×–2.8× higher peak decoding throughput than full attention, and sustains larger runnable batch sizes (e.g., at 128K, full attention OOMs at batch ≥ 4 while ParisKV scales to batch 8).

• Massive latency advantage at extreme context. At 1024K tokens on Llama3.1-8B, ParisKV achieves 49 ms/step versus 2179 ms/step for PQCache (44.4× faster) and 830 ms/step for MagicPIG (16.9× faster). Full attention cannot run at all at 384K even at batch size 1.

• Long-input tasks improve too, though less dramatically. On LongBench-V2 with DeepSeek-R1-Llama-8B, ParisKV reaches 28.43 overall (vs. 19.90 PQCache, 13.92 MagicPIG), with the largest gains in Medium/Long buckets — because fewer decode tokens means less drift opportunity.

Methodology in Plain English

Standard retrieval-based sparse attention works by clustering keys during prefill into groups with representative "centroids," then at each decode step finding the nearest centroids to the query and scoring keys in those buckets. The problem: as generation proceeds, newly produced keys drift away from centroids that were fit only on prefill keys, so recall degrades.

ParisKV sidesteps this by never learning centroids from data. Instead, it:

1. Normalizes every key and query to unit length, so similarity depends only on direction, not magnitude.
2. Rotates them with a shared random orthogonal matrix, which spreads information evenly across dimensions and makes the distribution approximately isotropic.
3. Splits the resulting vectors into B subspaces and, within each subspace of dimension m, uses a fixed codebook of 2^m sign-pattern directions (corners of a hypercube projected onto the sphere). Because these centroids are uniformly spread over the sphere, any new key is always close to some centroid — no staleness possible.
4. At decode time, runs a coarse stage where each key earns votes for landing in the same bucket as the query across subspaces; high-vote keys become candidates. Then a fine stage reranks candidates by estimating inner products from 4-bit quantized direction codes plus a precomputed correction weight that compensates for quantization shrinkage.
5. Offloads the full-precision KV cache to CPU memory, keeping only compact metadata (centroid IDs, 4-bit codes, weights) resident on the GPU. Only the final top-k KV pairs are fetched back over UVA for the actual attention computation.

Why This Matters

Impact on research. The paper reframes KV-cache retrieval as a representation-stability problem rather than a learned-index problem. Its core insight — that a fixed rotation makes centroids analytically derivable and inherently drift-proof — is a general principle that could apply to other approximate retrieval settings in ML systems, not just KV caches. It also provides a clean demonstration that sparsity and accuracy need not trade off in long-context inference.

Real-world applications:

• Long-running reasoning agents that generate tens of thousands of tokens of chain-of-thought, where prior retrieval methods collapse.
• Document analysis and summarization over book-length or codebase-length inputs that exceed GPU memory.
• Multi-turn conversational assistants with persistent long histories where early context remains relevant.
• On-device or memory-constrained inference where CPU offloading is the only way to handle long contexts within hardware budgets.

Industry relevance. The efficiency numbers are directly actionable for serving infrastructure: 2.8× throughput gains within full attention's runnable range, and the ability to serve million-token requests at all — a regime where full attention simply OOMs. The UVA-based design also fits naturally into existing GPU serving stacks without requiring CPU-side search servers. The code is publicly released.

Future Directions

• Integration with dropping methods. ParisKV retains the full KV cache; combining it with eviction could reduce the CPU-side memory footprint further, at some accuracy risk.
• Training-aware retrieval. The current pipeline is purely inference-time. Training models with retrieval-friendly attention patterns (as MoBA and NSA do) might raise the ceiling on how aggressive pruning can be.
• Adaptive candidate budgets. The paper uses a fixed top-k = 100 and a fixed candidate ratio. Making these query-adaptive — larger for ambiguous queries, smaller for confident ones — could yield further efficiency gains.
• Generalization beyond the tested model families. All three evaluated models are in the 4B–8B range; behavior at 70B+ or with mixture-of-experts architectures remains an open question, particularly whether the isotropic assumption after rotation holds equally well at larger dimensions and different attention head counts.

Target Audience

Machine learning systems researchers and inference engineers working on long-context LLM serving. Practitioners who need to deploy models at context lengths beyond GPU memory will find the offloading design directly useful. Readers should already understand transformer attention mechanics, KV caching, quantization basics, and approximate nearest-neighbor search; the paper is not accessible to newcomers without that background.

Authors’ abstract

KV-cache retrieval is essential for long-context LLM inference, yet existing methods struggle with distribution drift and high latency at scale. We introduce ParisKV, a drift-robust, GPU-native KV-cache retrieval framework based on collision-based candidate selection, followed by a quantized inner-product reranking estimator. For million-token contexts, ParisKV supports CPU-offloaded KV caches via Unified Virtual Addressing (UVA), enabling on-demand top-$k$ fetching with minimal overhead. ParisKV matches or outperforms full attention quality on long-input and long-generation benchmarks. It achieves state-of-the-art long-context decoding efficiency: it matches or exceeds full attention speed even at batch size 1 for long contexts, delivers up to 2.8$\times$ higher throughput within full attention's runnable range, and scales to million-token contexts where full attention runs out of memory. At million-token scale, ParisKV reduces decode latency by 17$\times$ and 44$\times$ compared to MagicPIG and PQCache, respectively, two state-of-the-art KV-cache Top-$k$ retrieval baselines, code is available at https://github.com/amy-77/ParisKV/tree/main.