arXiv:2409.15355cs.LGcs.AI2024-09ICLR被引 32

通过分块注意力机制,大幅降低大上下文生成的延迟和计算开销。

Block-Attention for Efficient Prefilling

  • 将检索文本分块,仅对新块计算键值,复用旧块状态以提速
  • 32K长度输入下首令牌生成仅需45毫秒,延迟与算力消耗降低超98%
  • 支持无损切换全注意力与分块注意力,适合高吞吐场景

我们提出块注意力(Block-attention),一种用于缓解检索增强生成(RAG)中推理延迟与成本上升的新注意力机制。传统方法通常以自回归方式编码全部上下文,而块注意力将检索文档划分为离散块,除最后一块外,其余块独立计算键值(KV)状态。在RAG场景中,若每段文本作为一块,即可复用已见过段落的KV状态,显著减少推理时的延迟与计算开销。实现包括块分割、位置重编码及对大语言模型进行块注意力微调。在11个涵盖RAG、ICL及通用领域的基准测试中,经块微调后的模型性能与全注意力模型相当,并可无缝切换至全注意力模式且无性能损失。值得注意的是,块注意力将首次生成时间(TTFT)和浮点运算量(FLOPs)降至极低水平:对于总长度32K的输入序列,首令牌生成仅需45毫秒,相比全注意力模型分别降低98.7%和99.8%。附录A详细说明了其在游戏AI中的应用及其显著潜力,强烈建议游戏领域研究者关注。

原文摘要 · Abstract (English)

We introduce Block-attention, an attention mechanism designed to address the increased inference latency and cost in Retrieval-Augmented Generation (RAG) scenarios. Traditional approaches often encode the entire context in an auto-regressive manner. Instead, Block-attention divides retrieved documents into discrete blocks, with each block independently calculating key-value (KV) states except for the final block. In RAG scenarios, by defining each passage as a block, Block-attention enables us to reuse the KV states of passages that have been seen before, thereby significantly reducing the latency and the computation overhead during inference. The implementation of Block-attention involves block segmentation, position re-encoding, and fine-tuning the LLM to adapt to the Block-attention mechanism. Experiments on 11 diverse benchmarks, including RAG, ICL, and general domains, demonstrate that after block fine-tuning, the Block-attention model not only achieves performance comparable to that of full-attention models, but can also seamlessly switch between the block and full attention modes without any performance loss. Notably, Block-attention significantly reduces the time to first token (TTFT) and floating point operations (FLOPs) to a very low level. It only takes 45 ms to output the first token for an input sequence with a total length of 32K. Compared to the full-attention models, the TTFT and corresponding FLOPs are reduced by 98.7% and 99.8%, respectively. Additionally, in Appendix A, we elaborate on how Block-attention is applied in Game AI scenario and the substantial potential benefits it entails. We strongly suggest researchers in the gaming field not to overlook this section.

注意力机制推理优化RAG大模型

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。