🤖 AI Summary
This work addresses the significant computational overhead of key-value (KV) cache access and processing in long-context decoding, where existing pruning methods—often relying solely on key information—struggle to balance efficiency and accuracy. To overcome this limitation, the authors propose ART, a lightweight runtime mechanism that dynamically monitors the magnitude and directional stability of accumulated attention outputs during decoding. ART terminates KV cache access early once subsequent KV blocks contribute negligibly, requiring no additional overhead and remaining compatible with both dense and sparse attention variants. The method also provides theoretical guarantees on truncation error. Experiments demonstrate that ART achieves up to a 20% improvement in generation throughput over state-of-the-art KV cache optimization techniques on LongBench and RULER Needle-in-a-Haystack benchmarks, without compromising output quality.
📝 Abstract
Long-context decoding in Large Language Models (LLMs) is severely constrained by the memory bandwidth required to fetch the extensive Key-Value (KV) cache. Most existing KV management methods rely on key-only pruning before decoding, despite the evidence that attention outputs depend jointly on keys and values, as incorporating values in their methods incurs prohibitive additional overhead. In this paper, we propose Attention Run-time Termination (ART), a lightweight run-time mechanism that tracks accumulated attention outputs during kernel execution and terminates subsequent KV block accesses once further contributions become negligible. This design makes ART orthogonal to existing key-based KV cache management methods, enabling seamless integration with them. Experiments on LongBench benchmarks show that ART achieves 20% higher generation throughput in large batch size than state-of-the-art baseline while maintaining comparable accuracy.