通过提前计算减少内存占用，让大模型推理更快更省显存
ScoutAttention: Efficient KV Cache Offloading via Layer-Ahead CPU Pre-computation for LLM Inference
- CPU提前一层预计算注意力，与GPU协同处理
- 相比现有方法提速2.1倍，精度损失小于2.4%
- 适合长文本生成场景，尤其在显存受限时
大型语言模型在长上下文推理中面临关键的GPU显存瓶颈，KV缓存占用严重限制解码批处理大小。现有将KV缓存卸载到DRAM的方法或需频繁的GPU-CPU数据传输，或对CPU计算要求过高，导致系统等待I/O或计算完成，GPU利用率低下。我们提出ScoutAttention，一种新型的KV缓存卸载框架，通过GPU-CPU协同注意力计算加速大模型推理。为避免CPU计算成为瓶颈，ScoutAttention引入块级稀疏注意力机制，显著降低CPU负载。不同于传统并行计算，其创新性地采用层前预计算算法，使CPU可提前一层开始注意力计算，并结合异步周期性召回机制，维持极低的CPU计算负担。实验表明，ScoutAttention在精度仅比基线低2.4%的情况下，相较现有卸载方法实现2.1倍速度提升。
原文摘要 · Abstract (English)
Large language models encounter critical GPU memory capacity constraints during long-context inference, where KV cache memory consumption severely limits decode batch sizes. While existing research has explored offloading KV cache to DRAM, these approaches either demand frequent GPU-CPU data transfers or impose extensive CPU computation requirements, resulting in poor GPU utilization as the system waits for I/O operations or CPU processing to complete. We propose ScoutAttention, a novel KV cache offloading framework that accelerates LLM inference through collaborative GPU-CPU attention computation. To prevent CPU computation from bottlenecking the system, ScoutAttention introduces GPU-CPU collaborative block-wise sparse attention that significantly reduces CPU load. Unlike conventional parallel computing approaches, our framework features a novel layer-ahead CPU pre-computation algorithm, enabling the CPU to initiate attention computation one layer in advance, complemented by asynchronous periodic recall mechanisms to maintain minimal CPU compute load. Experimental results demonstrate that ScoutAttention maintains accuracy within 2.4% of baseline while achieving 2.1x speedup compared to existing offloading methods.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。