Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Block-Attention for Efficient Prefilling

Dongyang Ma, Yan Wang, Lan Tian

TL;DR

Block-attention introduces a blockwise attention mechanism for Retrieval-Augmented Generation to dramatically reduce inference latency and compute by partitioning the input into independent blocks and reusing KV states across prompts. Only the final block attends across preceding blocks, while block-wise fine-tuning recovers full-attention-level accuracy across 11 benchmarks, including RAG and ICL tasks. The approach yields substantial efficiency gains, with TTFT and FLOPs reductions up to near-100% at long input lengths, and enables seamless switching between block and full attention during inference. The method demonstrates strong potential for real-time LLM applications, including Game AI, by enabling efficient, scalable, and flexible knowledge-driven generation.

Abstract

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.

Block-Attention for Efficient Prefilling

TL;DR

Block-attention introduces a blockwise attention mechanism for Retrieval-Augmented Generation to dramatically reduce inference latency and compute by partitioning the input into independent blocks and reusing KV states across prompts. Only the final block attends across preceding blocks, while block-wise fine-tuning recovers full-attention-level accuracy across 11 benchmarks, including RAG and ICL tasks. The approach yields substantial efficiency gains, with TTFT and FLOPs reductions up to near-100% at long input lengths, and enables seamless switching between block and full attention during inference. The method demonstrates strong potential for real-time LLM applications, including Game AI, by enabling efficient, scalable, and flexible knowledge-driven generation.

Abstract

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.
Paper Structure (26 sections, 6 equations, 5 figures, 3 tables)

This paper contains 26 sections, 6 equations, 5 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Block-attention
  3. Main Idea
  4. Block Segmentation
  5. Position Re-Encoding
  6. Block Fine-tune
  7. Inference
  8. Experiments
  9. Datasets
  10. Train Dataset
  11. Evaluation Dataset
  12. Input Format
  13. Base Model & Baselines
  14. Experimental Setup
  15. Training and Inference
  16. ...and 11 more sections

Figures (5)

  • Figure 1: The Block-attention Masks
  • Figure 2: The Inference Pipeline of Block-attention Model
  • Figure 3: The Inference Pipeline of Block-attention Model. The retrieved documents at the top have the highest relevance to the user query.
  • Figure 4: The accuracy of model checkpoint
  • Figure 5: A case of Texas hold'em AI's gamecore data