PagedEviction: Structured Block-wise KV Cache Pruning for Efficient Large Language Model Inference
This addresses memory efficiency issues for users of large language models during inference, though it is incremental as it builds on existing PagedAttention methods.
The paper tackles the memory bottleneck in Large Language Model inference caused by KV caching by proposing PagedEviction, a structured block-wise pruning strategy that improves memory efficiency while maintaining accuracy on long context tasks, as demonstrated on models like Llama-3.1-8B-Instruct with better performance than baselines.
KV caching significantly improves the efficiency of Large Language Model (LLM) inference by storing attention states from previously processed tokens, enabling faster generation of subsequent tokens. However, as sequence length increases, the KV cache quickly becomes a major memory bottleneck. To address this, we propose PagedEviction, a novel fine-grained, structured KV cache pruning strategy that enhances the memory efficiency of vLLM's PagedAttention. Unlike existing approaches that rely on attention-based token importance or evict tokens across different vLLM pages, PagedEviction introduces an efficient block-wise eviction algorithm tailored for paged memory layouts. Our method integrates seamlessly with PagedAttention without requiring any modifications to its CUDA attention kernels. We evaluate PagedEviction across Llama-3.1-8B-Instruct, Llama-3.2-1B-Instruct, and Llama-3.2-3B-Instruct models on the LongBench benchmark suite, demonstrating improved memory usage with better accuracy than baselines on long context tasks.