DropKV: Decoupling Residual-Output Perturbation for Near-Optimal KV-Cache Eviction

Motivation

• At 128K tokens, Phi-3-mini’s FP16 KV cache requires approximately 48 GB, compared with roughly 7 GB for the model weights.The cache grows linearly with sequence length and can saturate GPU memory and increase memory I/O overhead.Source: §1 · pp. 1–2

Problem

• Value-aware eviction methods target output perturbation, but their NP-hard optimization has led to suboptimal heuristics and uncertain kernelization efficiency.Attention-only importance scores can miss value vectors whose magnitude or direction substantially shifts the output.Source: §1 · p. 2

Contributions

• DropKV decouples joint eviction into tractable per-token perturbation scoring and proves a constant-factor guarantee under a cone condition verified on three long-context LLMs.Its fused Triton implementation reduces scoring scratch from approximately 9 GB to 17 MB at n=128K.Source: §1 · p. 2
• On RULER and LongBench at 5% and 10% cache budgets, DropKV achieves the highest accuracy and lowest latency among query-aware and output-aware eviction baselines while remaining faster than dynamic cache.The contribution summary also reports matched-batch end-to-end prefill speedups.Source: §1 · p. 2

2.1 Causal Attention and KV Cache

• The causal attention mask blocks attention to future tokens; inference processes the prompt in a prefill pass and then generates tokens during decoding using cached keys and values.Prefill stores the prompt’s key and value projections, and each decoded token’s projections are appended to the cache.Source: §2.1 · pp. 2–3

2.2 KV-Cache Eviction

• Eviction methods range from position-, key-, and attention-based heuristics to output-aware approaches that incorporate value representations and perturbation.DropKV scores exact per-token output perturbation, unlike CriticalKV’s loose upper bound and AnDPro’s convex relaxation.Source: §2.2 · p. 3
• Dynamic sparse attention and compression or offload retain token positions, whereas eviction permanently removes tokens and shrinks the resident cache.These approaches target the same memory or bandwidth bottleneck but are orthogonal to eviction.Source: §2.2 · p. 3

3.1 Attention-Output Perturbation

• Evicting tokens shifts the attention output, with the perturbation depending on both their attention weights and their value residuals from the original output.For a single token, the supplied expression scales the residual a − v_j by p_j/(1 − p_j).Source: §3.1 · pp. 3–4

3.2 Combinatorial Objective and a Decoupled Proxy

• Because the joint k-token eviction objective is NP-hard, DropKV selects the k tokens with the smallest individual perturbation losses, which can be scored in parallel at O(nd) cost.The proxy omits interactions among evicted tokens and is near-optimal when attention weights are small and residual directions are approximately aligned.Source: §3.2 · p. 4

3.3 Multi-Query Aggregation

• DropKV aggregates squared per-token perturbation losses across an observation window and evicts the k tokens with the smallest scores.The squared form supports fused tiling and preserves the single-query proxy’s ranking when w=1; the multi-query score is also covered by an extension of the guarantee.Source: §3.3 · p. 4

3.4 Triton Kernel Implementation

• 19.8× scoring-step speedup at n=128K over the naïve PyTorch baseline was achieved by fused Triton kernels that stream over key/value tiles without materializing attention-weight or value-distance tensors.The fused path matched the PyTorch reference within BF16 round-off; the scoring-step speedup was 1.8× at n=4K and 19.8× at n=128K.Source: §3.4 · p. 5

4 Theoretical Guarantees

• Under the stated assumptions, the proxy’s objective is at most 1/[c(1 − ε)(1 − kε)] times the optimal eviction objective.The bound applies to candidate-pool evictions of exactly k tokens.Source: §4 · pp. 5–6
• Proposition 1 extends Theorem 1 to the multi-query DropKV score under the same assumptions for each query.The guarantee analyzes eviction from a candidate pool L, such as tokens with the lowest attention scores, with exactly k tokens selected.Source: §4 · pp. 5–6
• Brute-force tests found median empirical approximation ratio ≤1.16 and 95th percentile ≤1.43 across all three models, below the worst-case bound.With both preconditions satisfied, the finite bound applied on 89.1%, 87.5%, and 73.4% of triples for Llama, Mistral, and Qwen2.5, respectively.Source: §4 · p. 6, Figure 1 · p. 6
• The cone condition requires removal residuals to align directionally rather than cancel, and its c⋆>0 criterion held on 91.2%, 91.1%, and 76.5% of triples for Llama, Mistral, and Qwen2.5.Prevalence stays high from Layer 1 onward on Llama and Mistral, with early-layer dips on Llama and Qwen2.5.Source: §4 · pp. 5–6, Figure 1 · p. 6
• When the cone condition does not strictly hold, a relaxed approximate cone condition bounds the weighted mass of outlier residuals by µ ≥ 0.The resulting additive bound recovers Theorem 1 when µ = 0.Source: §4 · p. 6

5 Experiments

• The evaluation covers three decoder-only LLMs, RULER and LongBench, and matched-budget comparisons on an A100-80GB GPU.The LongBench evaluation uses 16 English and code tasks; implementation details specify 5%, 10%, and 30% cache budgets.Source: §5 · p. 7

5.1 RULER

• At 5% cache budget, DropKV achieves the highest average RULER accuracy on all three model families, with gains of 2.35, 3.87, and 4.03 points over the best competing baseline.The results cover Llama-3.1-8B-Instruct, Qwen2.5-7B-1M, and Mistral-7B-512K; Table 1 reports the 128K results.Source: §5.1 · p. 7, Table 1 · p. 7

5.2 LongBench

• At 5% cache budget, DropKV achieves the highest LongBench average score on all three model families, improving over the strongest baseline by 0.87, 0.94, and 0.68 points.The reported averages are 45.51 for Llama-3.1-8B-Instruct, 42.66 for Qwen2.5-7B-1M, and 36.22 for Mistral-7B-512K.Source: §5.2 · p. 8, Table 2 · p. 8

5.3 Extended Results with Adaptive Head Budget

• AdaDropKV consistently outperforms competing adaptive baselines on RULER and LongBench across 5%, 10%, and 30% cache budgets.It combines DropKV token scoring with AdaKV's head-budget allocation strategy; Table 3 reports RULER at 64K and LongBench.Source: §5.3 · p. 9, Table 3 · p. 9

5.4 Memory and Latency

• At n=128K, the fused Triton scoring kernel is 19.8× faster than PyTorch, using ∼17 MB rather than ∼9 GB of scoring-kernel scratch memory.At matched batch sizes, prefill speedup is positive in eight of nine (n, model) cells, ranging from +1.9% to +9.9%; Qwen at n=32K is −0.1%.Source: §5.4 · p. 9, Figure 3 · p. 10
• At each method's maximum feasible batch, DropKV decoding throughput is 1.75–8.22× higher than the dynamic cache at 5% budget.The passage attributes this throughput gain to KV-bandwidth reduction shared by eviction methods; Triton's contribution is concentrated in prefill.Source: §5.4 · p. 9
• On Mistral-7B-512K at 32K sequence length and 5% budget, DropKV has the lowest full-generation latency among eviction methods at every batch size.The dynamic cache runs out of memory beyond B=4, while eviction methods scale to B=12; where dynamic is feasible, DropKV TTFT matches it within 0.6%.Source: §5.4 · pp. 9–10, Table 4 · p. 10

6 Discussion and Limitation

• DropKV decouples the joint perturbation problem, and the authors report that the cone condition holds at most layer, head, and query-position triples in the three target models.They describe the condition as alignment of value residuals and use it to justify the effectiveness of independent per-token scoring.Source: §6 · p. 10
• The current Triton implementation accelerates prefill scoring only; a Triton decode-time kernel is left for future work.Source: §6 · p. 10
• The paper frames DropKV as bridging theoretical optimality and practical scalability in KV-cache eviction.Source: §6 · p. 10

A.1 Extended RULER Results

• Tables 5–8 extend RULER results to 16K, 32K, and 64K sequence lengths at 5%, 10%, and 30% cache budgets, and to 10% and 30% budgets at 128K.Source: §A.1 · p. 16
• Table 5 reports RULER results at 128K sequence length for 10% and 30% KV-cache budgets.Source: Table 5 · p. 16
• Tables 6–8 report RULER results at 64K, 32K, and 16K sequence lengths, respectively, across 5%, 10%, and 30% budgets.Source: Table 6 · p. 17, Table 7 · p. 18, Table 8 · p. 19

A.2 Extended LongBench Results

• Tables 9 and 10 extend LongBench results to 10% and 30% cache budgets.Table 10 identifies the best and second-best average results with gray highlighting and underlining.Source: §A.2 · p. 20, Table 9 · p. 20, Table 10 · p. 20

A.3 Extended Results with adaptive head budget

• Tables 11 and 12 extend AdaDropKV evaluation to broader sequence lengths and 5%, 10%, and 30% cache budgets on RULER and LongBench.The paper describes these results as evidence of versatility and robustness for combining DropKV scoring with adaptive head-budget allocation across diverse long-context workloads.Source: §A.3 · p. 20, Table 11 · p. 21, Table 12 · p. 22

A.4 Serving Throughput vs. Dynamic Cache

• At 5% budget, DropKV's maximum batch ranges from 1.75× to 6.87× the dynamic-cache baseline across the reported model and context settings.The appendix reports a Triton throughput lead over PyTorch on Mistral at every n, peaking at B=3 versus B=2 and 70.8→86.5 tok/s at n=128K; Table 13 defines its throughput and batch measures.Source: §A.4 · p. 21, Table 13 · p. 22

B Scoring Functions of Prior Eviction Methods

• Baselines approximate the same NP-hard objective through cache-side heuristics or output-perturbation surrogates, including bounds and convex relaxations.The scoring table classifies these as heuristic retention, explicit upper-bound, relaxation, or exact-singleton-decoupling strategies.Source: Table 14 · p. 23
• DropKV uses the exact singleton perturbation as a per-token eviction cost and decouples the multi-token objective into a window-summed sum of squared token costs.The decoupling is justified under the cone condition; the score uses (pj/(1 − pj))^2 weighting.Source: B Scoring Functions of Prior Eviction… · pp. 23–24
• At n=128K and 5% budget, DropKV exceeds SnapKV by +2.43, +7.68, and +7.09 average RULER points on Llama-3.1-8B, Mistral-7B-512K, and Qwen2.5-7B-1M, respectively.The comparison uses the same observation window and pooling kernel; the scores differ in their value-residual factor.Source: B Scoring Functions of Prior Eviction… · p. 23
• At a 5% RULER budget, DropKV improves average accuracy over AnDPro by +2.35, +3.87, and +4.03 points on Llama-3.1-8B-Instruct, Qwen2.5-7B-1M, and MegaBeam-Mistral-7B-512K, respectively.When vj is near a, DropKV assigns near-zero eviction cost, whereas AnDPro’s score is large and retains the token.Source: B Scoring Functions of Prior Eviction… · p. 24

C.1 NP-hardness of Exact Eviction Minimization

• Even with small eviction probabilities, minimizing the exact perturbation is NP-hard because the norm of the summed residuals couples all selected tokens.A reduction from Balanced Partition establishes hardness while satisfying the small-weight assumption.Source: §C.1 · p. 24

• For a single evicted token, the attention-output perturbation is pj/(1 − pj) times the residual a − vj.The multi-token output change follows from re-normalizing the attention probabilities over the remaining tokens.Source: §C.2 · pp. 24–25
• For the observation window, Proposition 1 extends near-optimality to summed squared perturbations, with a √w factor in the bound.The cone directions may differ across queries; under GQA, the corresponding factor becomes √gw.Source: §C.2 · p. 26

C.3 Numerical Verification of Theorem Assumptions

• The sampled candidate pools are typically in the small-weight regime: median ε is 2.6–4.3 × 10^-6 and the 95th percentile is at most about 1.2 × 10^-4 across the three models.The maximum is higher in a few early-layer/head triples, so the paper reports finite-bound coverage rather than claiming the denominator is positive uniformly.Source: §C.3 · p. 26, Table 15 · p. 26
• Among sampled triples, the strict cone condition holds for 91.21%, 91.12%, and 76.45% for Llama, Mistral, and Qwen2.5, respectively.Source: §C.3 · p. 27, Table 17 · p. 27
• Across 1,200 records per model, DropKV’s realized proxy ratio has median ≤1.16 and 95th percentile ≤1.43 in every model, stratification, and k cell; the worst record is 1.82.The ratio is much tighter than the theorem’s median bound factor of 10–12; on rank-disagreement cases, DropKV’s median is 1.03 versus 1.09–1.15 for attention-only selection.Source: §C.3 · p. 27, Table 19 · p. 28

C.4 Further Robustness Analysis

• The numerical verification covers Assumptions 1 and 2, not the relaxed Assumption 3, because estimating its out-of-cone weighted-mass parameter μ was left outside the check.The robust analysis bounds the weighted mass of residuals outside the primary alignment cone; the strict cone condition is the special case μ = 0.Source: §C.3 · pp. 27–28, §C.4 · p. 28

D Triton Implementation Details

• The Triton scoring path uses a FlashAttention-style forward pass to compute outputs and log-sum-exp values, then accumulates scores over key tiles without materializing the attention matrix.When split-K is enabled, a merge kernel combines partial online-softmax states.Source: D Triton Implementation Details · p. 29

• At n=16K, the Triton scoring call is 15.7× faster than the matched PyTorch reference, with 41× fewer instructions and 26× less HBM traffic.The reference materializes dense W×L attention and repeated K/V tensors, whereas Triton streams K/V tiles and accumulates scores directly.Source: §D.1 · p. 30, Table 20 · p. 31
• At n=16K, Kernel A reaches 49.0% of peak HBM bandwidth versus 22.7% SM throughput, while Kernel B reaches 54.0% versus 39.6%.Kernel B sustains roughly 1.10 TB/s on A100 HBM.Source: §D.1 · p. 30, Table 21 · p. 31
• The split three-kernel design preserves tile-level parallelism for memory-bound scoring; a persistent fused A+B kernel would leave only eight CTAs for the stated Llama-3.1-8B setup.Kernel A’s register use and Kernel B’s shared-memory use constrain occupancy, while the split avoids dense attention and repeated-KV materialization.Source: §D.1 · pp. 30–31

E Details of RULER and LongBench Benchmarks

• RULER covers retrieval, variable tracking, word-frequency extraction, and single- and multi-hop question answering through its task suite.Its NIAH variants test retrieving one or multiple values from long contexts; other tasks require identifying common or frequent words or integrating evidence.Source: E Details of RULER and LongBench… · pp. 31–32
• LongBench averages 6711 in length and spans single- and multi-document QA, summarization, few-shot learning, synthetic tasks, and code generation.Source: E Details of RULER and LongBench… · p. 32
• LongBench includes diverse source datasets and tasks, from legal and academic documents and multi-document QA to summarization, few-shot learning, synthetic passage tests, and code completion.Its synthetic tasks test passage retrieval and counting, while code completion uses LCC and RepoBench-P.Source: E Details of RULER and LongBench… · p. 32

F Broader Impact

• DropKV’s lower inference memory footprint and HBM traffic are described as reducing energy use and carbon footprint, while its lower cost and faster execution may also enable disinformation or surveillance.Source: F Broader Impact · pp. 32–33

NeurIPS Paper Checklist

• The checklist identifies DropKV’s contributions as near-optimality guarantees and a Triton implementation, and says reproducibility details cover models, datasets, evaluation, hyperparameters, scoring, and optimization.The authors say they use open-source datasets and plan to release code after acceptance.Source: NeurIPS Paper Checklist · pp. 34–36
• The paper marks crowdsourcing, human-subjects research, and related IRB disclosures N/A because the work involved no crowdsourcing or human subjects.The passages provide no benchmark details.Source: NeurIPS Paper Checklist · p. 39
• LLM use was limited to writing assistance, including grammar correction and language polishing, and did not contribute to core methods or experiments.The checklist notes that N/A means LLMs were not an important, original, or non-standard part of core method development.Source: NeurIPS Paper Checklist · p. 40