通过块级结构化淘汰策略，提升大模型推理时的显存效率
PagedEviction: Structured Block-wise KV Cache Pruning for Efficient Large Language Model Inference
- 基于分页内存布局设计块级淘汰机制，精准控制缓存占用
- 在LongBench上实现更低显存使用与更高长文本生成准确率
- 无需修改现有内核，可无缝集成到vLLM框架中
KV缓存通过存储已处理标记的注意力状态，显著提升大语言模型推理效率。然而，随着序列长度增加，KV缓存迅速成为主要的内存瓶颈。为此，我们提出PagedEviction，一种细粒度、结构化的KV缓存剪枝策略，用于增强vLLM中PagedAttention的内存效率。与依赖注意力重要性的现有方法或跨页淘汰不同，PagedEviction针对分页内存布局设计了高效的块级淘汰算法。该方法可无缝集成至PagedAttention，无需修改其CUDA注意力内核。我们在Llama-3.1-8B-Instruct、Llama-3.2-1B-Instruct和Llama-3.2-3B-Instruct模型上，于LongBench基准测试中验证了其效果，结果表明，在长上下文任务中，该方法相比基线实现了更优的内存利用率与更高的准确性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。