FastKV: Decoupling of Context Reduction and KV Cache Compression for Prefill-Decoding Acceleration

Abstract: While LLMs excel at handling long-context sequences, they require substantial prefill computation and key-value (KV) cache, which can heavily burden computational efficiency and memory usage in both prefill and decoding stages. Recent works that compress KV caches with prefill acceleration reduce this cost but inadvertently tie the prefill compute reduction to the decoding KV budget. This coupling arises from overlooking the layer-dependent variation of critical context, often leading to accuracy degradation. To address this issue, we introduce FastKV, a KV cache compression framework designed to reduce latency in both prefill and decoding by leveraging the stabilization of token importance in later layers. FastKV performs full-context computation until a Token-Selective Propagation (TSP) layer, which forwards only the most informative tokens to subsequent layers. From these propagated tokens, FastKV independently selects salient KV entries for caching, thereby decoupling KV budget from the prefill compute reduction based on the TSP decision. This independent control of the TSP rate and KV retention rate enables flexible optimization of efficiency and accuracy. Experimental results show that FastKV achieves speedups of up to 1.82$\times$ in prefill and 2.87$\times$ in decoding compared to the full-context baseline, while matching the accuracy of the baselines that only accelerate the decoding stage. Our code is available at https://github.com/dongwonjo/FastKV.

• The paper introduces FastKV, which uses Token-Selective Propagation after an empirically selected stabilization layer and independently compresses each layer’s KV cache, separating prefill compute from decoding memory.
• FastKV retains roughly full-context accuracy, with LongBench drops of about 1.5–1.7% and RULER scores of 75.6 at 10% KV retention, while preserving performance at contexts up to 128K tokens.
• The method delivers up to 1.82× prefill and 2.87× decoding speedups, exceeding 2× end-to-end acceleration at 128K context length on an A100 GPU, though its generalization beyond 8B GQA models remains unverified.

FastKV is a KV cache compression framework for long-context LLM inference that decouples prefill compute reduction from decoding KV budget, addressing a coupling that limits prior prefill-aware compression methods. The paper's central claim is that token importance stabilizes across later layers, so full-context processing can be confined to early layers while both context propagation and KV retention are controlled independently thereafter. The reported results include up to 1.82× prefill and 2.87× decoding speedup over the full-context baseline on LLaMA-3.1-8B-Instruct, with LongBench accuracy drops within 1% (2502.01068).

Motivation: layer-dependent context dynamics

The authors analyze how critical tokens evolve across depth by feeding 128K-token inputs to LLaMA-3.1-8B-Instruct and measuring the overlap of top-512 attention-mass tokens between layers. In early layers (≤15), overlap decays sharply with layer distance; beyond layer 15 it decays slowly, indicating a stable subset of consistently important tokens. A second measurement of top-$K$ attention recall shows that even $K=512$ (0.38% of tokens) captures the majority of attention mass at every layer during decoding.

These two observations jointly motivate the design: pruning before stabilization irreversibly removes tokens later layers would need, whereas after stabilization aggressive pruning is safe; and since decoding attends sparsely, each layer need not cache everything it computed during prefill. Existing methods violate one or both conditions — GemFilter fixes a single filter-layer selection across all layers and restarts prefill on the reduced set, discarding information entirely, while PyramidInfer prunes gradually from the earliest layers under a cosine schedule and ties its KV retention rate to its prefill compute rate.

Method

FastKV introduces two mechanisms.

Token-Selective Propagation (TSP). Prefill proceeds in two stages. Up to a designated TSP layer (layer 15 for LLaMA-3.1-8B-Instruct, layer 17 for Ministral-8B-Instruct), all layers process the full context. At the TSP layer, per-token saliency scores are computed by averaging, over heads, the attention weights received from a recent observation window (window size 8, pooling kernel size 7). Top-ranked tokens up to the TSP rate (20% in evaluation), merged with all window tokens, are propagated as hidden states to subsequent layers, which then process only this compressed context. Crucially, discarded tokens still contribute through early-layer attention, so their information is integrated into the surviving representations rather than lost as in GemFilter's re-prefill scheme.

Decoupled KV retention. Independently of the TSP rate, each layer compresses its own KV cache to a configured retention rate (10–20%) using group-wise saliency scores averaged within GQA key-value groups. This yields two independent hyperparameters: TSP rate controls prefill compute; retention rate controls decoding memory and bandwidth.

The TSP layer itself is selected by minimizing the normalized L2 distance between final logits (or hidden states) under TSP and the full-context baseline over calibration inputs, subject to an upper bound $L_{\max}$ to avoid late placement with negligible speedup. Empirically, placing TSP too early causes large logit deviation; too late yields little latency benefit. Layer 15 saturates LongBench accuracy while providing meaningful savings.

Accuracy results

On LongBench (English tasks) across LLaMA-3.1-8B-Instruct and Ministral-8B-Instruct, FastKV at 60% prefill compute and 10% KV retention achieves average scores of 48.47 and 50.40 versus full-context baselines of 50.19 and 51.87 — drops of roughly 1.7% and 1.5%, comparable to SnapKV (48.73 / 50.68), which does not accelerate prefill at all. By contrast, GemFilter at similar budgets drops to 38.61 on LLaMA (an 11.58-point degradation, concentrated in code completion, which falls from 60.01 to 30.75), and PyramidInfer reaches only 39.32 despite retaining 60% of the KV cache. On RULER, FastKV averages 75.6 versus SnapKV's 73.6 and GemFilter's 69.6 at 10% retention, maintaining performance up to 128K contexts. Notably, on Needle-in-a-Haystack FastKV scores 99.9, slightly exceeding the full-context baseline's 99.0, which the authors attribute to TSP simplifying the input context around globally salient tokens.

The implication is that the accuracy cost of prefill-aware compression stems primarily from when context is reduced rather than how much: FastKV matches decoding-only methods' accuracy while adding prefill acceleration they lack.

Latency results

End-to-end benchmarks on a single A100 SXM GPU (256 generated tokens, varying context length) show that decoding-only baselines provide diminishing returns as prefill dominates at long contexts. H2O and PyramidInfer cannot use FlashAttention-2 and run out of memory beyond 8K tokens; H2O is slower than full-context even at 8K. FastKV achieves more than 2× end-to-end speedup at 128K, with peak stage-level speedups of 1.82× (prefill) and 2.87× (decoding). GemFilter is marginally faster on LLaMA because its filter layer (13) precedes FastKV's TSP layer (15); on Ministral, where both sit at layer 17, latencies are identical — but GemFilter pays for this with substantially lower accuracy.

Ablations confirm the design choices: accuracy degrades when the TSP rate falls below ~20%, and early TSP placement causes sharp accuracy loss consistent with the instability analysis.

Limitations and open questions

The paper concedes several constraints. FastKV's advantage materializes mainly in long-context regimes where KV cache handling dominates latency; for short contexts, quantization or kernel-level optimizations may be more effective, though the authors argue FastKV is orthogonal to such techniques. The method was evaluated only on 8B-scale models with GQA, and the TSP layer index, window size, and pooling kernel were fixed empirically per model — whether the calibration-based layer selection generalizes across architectures and scales is not established. The saliency criterion relies on window-token queries, inheriting assumptions from attention-sink literature about where informative queries concentrate; behavior under prompts whose relevant content lies far from the recent window is not separately analyzed. Finally, the claimed robustness rests on the observed stabilization threshold near layer 15; the paper does not characterize how this threshold varies with model family or training data distribution.

Conclusion

FastKV demonstrates that separating the decision of how much context to propagate during prefill from how much KV cache to retain for decoding resolves the accuracy–efficiency trade-off in prefill-aware KV compression. By confining full-context computation to early layers where attention focus is unstable and applying selective propagation plus independent KV retention afterward, it attains decoding-focused methods' accuracy alongside prefill-aware methods' speedup, with sub-2% LongBench degradation and over 2× end-to-end acceleration at 128K contexts.