DropKV speeds up cache management in long context language models
DropKV: Decoupling Residual-Output Perturbation for Near-Optimal KV-Cache Eviction
Artificial Intelligence
Summary
Managing memory in large language models takes more time and space as the context gets longer, which slows things down. The authors created DropKV, a new way to decide which parts of the memory to keep or remove that avoids complicated computations. DropKV is faster and more reliable, improving model speed without losing accuracy. They also built an efficient version that runs much quicker on hardware, tested on several benchmarks and models.
What this means in practice
- •For ai infrastructure engineers: Optimize cache eviction in large language models to reduce inference latency and resource usage in long-context deployments.
- •For cloud service providers: Deploy DropKV to accelerate prefill performance and overall response times for LLM-powered services with long input sequences.
Authors
Aozhong Zhang, Selcuk Gurses, Yanxia Deng, Naigang Wang, Chi-Chun Liu, Davis Wertheimer, Derrick Liu, Xin Li, Zi Yang, Felix X. -F. Ye, Penghang Yin
Abstract
Key-value (KV) cache management has become a critical bottleneck for scaling Large Language Models (LLMs) to long-context scenarios due to linear memory growth. While recent eviction methods have moved beyond simple attention-score heuristics to incorporate value representations, they typically face a fundamental trade-off: the underlying optimization for minimizing output perturbation is an NP-hard combinatorial problem, leading to mathematically suboptimal heuristics or implementations that are difficult to kernelize. In this work, we propose DropKV, a simple principled eviction framework based on Decoupled Residual-Output Perturbation. By decoupling the joint eviction decision into independent per-token scoring, DropKV circumvents combinatorial intractability and admits a provable constant-factor approximation guarantee under a cone condition that we empirically verify on three long-context LLMs. To ensure practical viability, we provide a high-performance implementation using fused Triton kernels that avoid materializing the full attention matrix, delivering a 19.8x scoring-kernel speedup that translates to up to a 9.9% matched-batch end-to-end prefill speedup. Extensive evaluations on the RULER and LongBench benchmarks demonstrate that DropKV consistently outperforms state-of-the-art baselines at equivalent cache budgets, achieving the lowest inference latency among eviction methods while also remaining faster than the dynamic cache baseline.