融合分页注意力与弹性注意力，显著提升长序列推理效率。
Paged Attention Meets FlexAttention: Unlocking Long-Context Efficiency in Deployed Inference
- 将分页注意力与弹性注意力结合，高效处理分散的键值缓存。
- 序列长度从128增至2048时延迟仅线性增长约2倍，远优于指数增长。
- 适合需部署长上下文大模型的工程师，开源实现便于落地。
大语言模型在长上下文推理中因传统键值（KV）缓存处理方式导致严重内存低效。本文提出将分页注意力与PyTorch的弹性注意力融合，解决单块缓存分配带来的内部碎片问题。该融合注意力核在IBM基础模型栈（FMS）中实现，能高效聚合分散的KV数据。在NVIDIA L4 GPU（24GB）上的基准测试表明，使用全局KV缓存时，推理延迟随序列长度从128增至2048仅线性增长约2倍，而无缓存时呈指数上升。单步评估峰值内存基本不变（由模型权重和激活主导），分页注意力仅在序列长度超过2048时产生微小增量内存开销，源于其2的幂次缓存分配策略。本文开源完整实现，并讨论其对长上下文模型部署的深远影响。
原文摘要 · Abstract (English)
Large Language Models (LLMs) encounter severe memory inefficiencies during long-context inference due to conventional handling of key-value (KV) caches. In this work, we introduce a novel integration of PagedAttention with PyTorch's FlexAttention, addressing internal fragmentation and inefficiencies associated with monolithic KV cache allocations. Implemented within IBM's Foundation Model Stack (FMS), our fused attention kernel efficiently gathers scattered KV data. Our benchmarks on an NVIDIA L4 GPU (24GB) demonstrate significantly reduced inference latency, growing only linearly (~2x) with sequence length from 128 to 2048 tokens when utilizing a global KV cache, compared to exponential latency increases without caching. While peak memory usage remains largely unchanged for single-step evaluations (dominated by model weights and activations), paged attention causes minimal incremental memory usage, observable only at sequence lengths exceeding 2048 tokens due to its power-of-two cache allocations. We open-source the full implementation and discuss its implications for future long-context model deployment.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。