Block-Attention for Efficient Prefilling
Authors: Dongyang Ma
Yan Wang
Tian Lan
Venue: ICLR 2025 Poster (2025)
Keywords: LLM
RAG
parallel context encoding
efficient language model
Primary area: foundation or frontier models, including LLMs
TL;DR: This paper enables cross-prompt passage KV-state reuse in RAG, drastically reducing TTFT and computational costs while maintaining full-attention performance, and releases an LLM capable of seamless switching between block and full attention modes.

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.

Published as a conference paper at ICLR 2025
BLOCK-ATTENTION FOR EFFICIENT PREFILLING
Dongyang Ma*
Yan Wang"i
Tencent
Tencent
dongyangma@tencent.ccom
yanwang.brandeen branden@gma i com
Tian Lan
lant lant sentianggttbey@miil.....
ABSTRACT
We introduce Block-attention, an attention mechanism designed to address the increased inference latency and cost in Retrieval-Augmented Generation (RAG) scenarios. 'rraditional approaches often encode the entire context in an auto- regressive manner. Instead, Block-attention divides retrieved documents into dis- crete blocks, with each block independently calculating key-value (KV) states except for the final block. In RAG scenarios, by defining each passage a block, Block-attentin been seen before, thereby significantly reducing the latency and the computa- tion overhead during inference. The implementation of Block-attention involves block segmentation, position re-encoding, and fine-tuning the LLM to adapt to the Block- mechanism. Experiments 11 diverse benchmarks, includ- ing RAG, ICL, and general domains, demonstrate that after block fine-tuning, the Block-attention model not only achieves performance comparable that of full- attention models, but can also seamlessly switch between the block and full atten- tion 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 output the first token for an input sequence with total length of 32K. Compared to the full-attention models, the TTFT and corresponding FLOPs are reduced by 98.7% and 99.8%, respectively. Addition- ally, in Appendix A we elaborate how Block-attention is applied in Game Al scenario and the substantial potential benefits entails. We strongly suggest researchers in the gaming field not overlook this section.
enables us to reuse the KV states of passages that have
INTRODUCTION
Cetrieval-Augmentie Generation (RAG) al. 2022 Lan et al. 2023) is crucial technology for mitigating knowledge hallucination in Large Language Models (LLMs). By utilizing retrieval technology, LLMs can seamlessly access passages stored in external databases, grounding their responses in the content of these passages. To the best of our understanding, RAG has emerged as the most effective method for infusing specific domain knowledge into L.MMs in real-world scenarios.
However, everything has two sides, and RAG is exception. Generally, for each user query, in order to ensure that a passage with the "correct knowledge" is successfully recalled, it is common practice to retrieve multiple passages-typiche
between 5 to 30 in most scenarios (Kwiatkowski et al., 2019 Joshi et al., 2017). These passages are then incorporated into the input prompt for the LLM. As result, the inference efficiency decreases significantly due to the increased sequence length of this extended input prompt. Specifically, the inference latency, measured as the time to first token (TTFT), is considerably higher for RAG-LLLM compared non-RAG LLM (Li al.
2023 Zhu et al. 2024).
Equal Contribution
1 Corresponding Author
'Codes, datasets and model weights have been publicly available httpa://github.com Veemporaryy..oA///Acc----- atteetttittttttinnnnooooo

---

Published as a conference paper at ICLR 2025
System: you are a. Below are some reference documents that may help you in answering... Vi
Blocko
Document [1] (Tièe: Manhattan Life Insurance Company) Manhatian Life Insurance ...In
Block,
Document, [2] (Title: 712 Fifth Avenue) 712 5th Avenue is a 650 skyscraper at 56 Street tn
Block;
Question: Which tower is tallec 712 Fifth Avenue or Manballan Life Insurance Building?
Black,
Figure 1: The Block-attention Masks
Given that the passages in the external databases might have been computed, it is natural to restore their KV cache for fast inference and avoid re-computing these passages. Nonetheless, for autore- gressive LLMs, the KV states are inherently context-dependent, which means the KV states for the same passage will vary in different contexts. As result, when encountering an unseen query, the model must undertake complete re-encoding of the KV states to ensure an accurate prediction.
In this paper. we propose the Block-attention method, which reduces the TTFT and computation FLOPs of RAG-LLMs to level nearly equivalent to that of non-RAG LLMs, while fully maintain- ing the same accuracy level. As shown in Figure [ the main idea of Block-attention is to divide the entire input sequence into several blocks. Each block independently calculates its KV states through full-attention, without any attention to other blocks. Only the final block able to attend other blocks, i.e., the user query is able to attend all the retrieved documents in the previous blocks. When utilizing Block-attention in RAG scenarios, we may achieve substantial benefits by defining each passage single block and caching its KV states in the memory for further reuse.
The implementation of Block-attention can be easily achieved through the following steps: 1) En- code all blocks except the last separately; 2) calculating the positional encoding for each based on its position the input prompt: 3) Integrating these blocks to compute the KV states for the final block. However, the primary challenge in utilizing Block-attention is that the LLMs have been exposed to such an attention mechanism during their training, leading to difficulties in accurately interpreting the input promp In preliminary experiments, attempted to directly implement Block- attentotnn in the LLMs without updating any parameters. Unfortunately. this approach resulted in a substantial decrease in performance, with the average accuracy of Llama-3. 1-Tull----BB---- on four RAG benchmarks falling from 66.1% to 49.9%.
To address this challenge, we implemented fine-tur tineituning process for the LLMs to adapt to the Block- attention mechanism. Our experiments demonstrated that, after approximately 1000 fine-tuning steps, the Block-attention model achieved a full recovery of its original accuracy across all scenar- ios, impressively increasing from 49.9% to 70.0%. This outcome underscores the Block-attention LLMs' capability to uphold inference accuracy while significantly enhancing inference efficiency in RAG scenarios.
We conduct comprehensive evaluations the Block-attention mechanism across 11 diverse bench- marks, including RAG, (CL(In-Contex Learning), and general domains. Experimental results demonstrate that after fine-tuning, the average accuracy of the Block-attention model on the bench- marks remain comparable to Llama- 3.1 --T-lu----B--FF In terms of efficiency, counted the TTFT and FLOPs to the first token (FLOPs-TFT) of the Block-attention model when the length of user input 50 and the total length of the input sequence gradually increases. We found that the longer the total length, the more obvious the improvement of Block- on efficiency. When
2 AAthhough some studies have proposed training-free frameworks Gim et al. 2024 Merth et al, 2024 utilizing parallel context encoding and KV caching to achieve passage-level KV cache reuse, they encounter significant degeneration issues practical applications. Our experimental results also demonstrate that the results of these methods are far from satisfactory.
We also proposed some training-fre free solutions mitigate this huge gap. someone is interested in switch- ing to Block-attention in an on-the-fly manner, please refer /Zhang al. (2024) for the detail.
2

---

Published as a conference paper at ICLR 2025
the length of the input sequence reaches 32K, the TTFT and FLOPs-TFT of the Block-attention model are 1.3% and 0.2% of that of the full-attention model, respectively" 4
BLOCK-ATTENTION
2.1 MAIN IDEA
Let S = {so, Sx) represents the input sequence, where each represents token. We denote the KV states associated with S as K (ko, ki.... kn} and V {up, Vi,..., Cn),, re- spectively. For an auto-regressive model OLLM. since the computation of the KV states is de- pendent on their preceding tokens, when a text block {si,... 8j) changes to a new text block (s'....,&&,, .,. , for the new sequence S' {so, ...., 8/; .... S': 8j+1..., Sa its KV states become K' = {Ko.... KGy.... "i.,...,,, Km. k/+t:..., kn and V = {vo, ".,,.,, U' + v. It is evident that although only one block {8,...., }j} has changed, due to the auto-regressive nature of LLMs, the KV states of all subsequent blocks must be re-encoded.
Our research focuses on exploring a novel attention mechanism named Block- This mech- anism is designed in such way that it only requires the computation of the text blocks that have been altered between two input sequences. As a result, it can achieve an outcome that is equivalent to that obtained from fully encoding the entire sequence.
As illustrated in Figure 1, the essence of Block-attention is to divide the input sequence S into sev- eral independent blocks. Each block autonomously calculates its KV states through self-attention, without considering other blocks. The final block, however, has the unique capability to integrate information from preceding blocks. A primary advantage of this method is the modular indepen- dence it provides: when a block 6, is updated to 6', re-encoding only the KV states of the affected block kig Ubg, and those of the final block kbg, Ubg is sufficient to obtain the updated KV states.
To develop a Block-attentic LLM capable of precise inference, we must tackle three challenges:
1) How do we segment blocks?
2) How should the positional encoding be calculated for each block?
3) How can the LLM be adapted the Block-attention mechanism?
These issues will be addressed in detail in Sections 2.2
2.3 and 2.4 respectively.
2.2 BLOCK SEGMENTATION
The primary principle of block division is segment semantically independent parts of the prompt into separate blocks. In RAG scenarios, since the retrieved passages are originally mutually inde- pendent, is natural divide them into different blocks. Therefore, let's go back to the left part of Figure | where we allocate each passage to single block and designate the user's query as the final block. This principle extends other scenarios as well. For example, in the context of code generation tasks, function may be treated one block; in multi-turn dialogues, each turn could be segmented into an individual block; while In-context Learning (ICL., each demonstration natu- rally forms self-contained block (we will validate the efficacy of Block-attention in ICL scenarios in our experiments). this paper, our primary focus is on the application of Block-attention in RAG, with the exploration of other scenarios reserved for future research.
2.3 POSITION RE- EE-CNNGG
The second problem is to re-encoding the positional information. Although the same passage may appear in multiple input prompts, its position generally varies. Therefore. when we attempt reuse the KV states a block, we need to re-encode its positional information. The process of re-encoding simple and straightforward: taking the rotary positional encoding (RoPE) an example, assume we wish to change the positional encoding of a block b = {si, ..., 8,} to b' = {Sia) 8ia J. then only need three steps:
'Given that the KV cache is already a mature and low-cost technology Qin et al. 2024 Lee et al. in this paper we do not take the cost of KV cache into account.
2021
3

---

Published as a conference paper at ICLR 2025
1) For the token 8,, its positional encoding vector f(s i) is calculated using the following formula:
cos i0o sin i0g 0
sin cos i0
Po
0
f(z;,
0
cos i sin at sin i0
4-1
Pt
(1)
0
COS i0
Pd-1
2) We rotating z; counterclockwise by i0 degrees, to re-set its positional encoding to zero:
cos i00
sin i0
0 0
- sin i00
cos i0
(s;;
cos i0
sin i0
f1r,;) (2)
sin
cos i0
3) Then, by performing clockwise rotation of (ia) ) degrees, we obtain the final positional encod- ing vector:
[ = f(vi, ia)
cos (i )00
sin (i
0
0
sin (ia)0
cos (is)0
0
0
(3)
0
cos (ia)0g
- sin (ia)
f(z;,
-1
0
0
sin (1a)0g
cos (is 10
For the remaining tokens within block b, namely $141,...,8p we can re-encode their positional information in similar manner. Al- though the formulas presented above may seem complex, the principle is quite straightforward: first set the positional encoding to zero, and then rotate it to the updated position. One might wonder why we do not simply rotate by (is -i) degrees directly? The reason is to mitigate the potential for errors in updating po- sitional encodings within practical applications: in the KV cache, the positional encoding of the initial token of each block standardized to zero, and with only the updated positional in- dex ia, we can readily determine their new po- sitional encoding vectors as per Equation
Outpu
Position Ro croodingg
Cached KV states
Cache
Block
Block
Block
Block.
Sequercee
Figure 2: The Inference Pipeline of Block- attention Model
2.4 BLOCK FINE-TUNE
Due to the LLM's reliance on full-attention during the training phase, a direct switch to Block- attention during inference might result in significant discrepancy between the training and infer- ence states. Our preliminary findings indicate that introducing Block-attention without subsequent fine-tuning could precipitate substantial decrease in performance, with the average accuracy drop- ping significantly from 67.9% to 48.0%. Adapting the LLM to Block-attention through fine-tuning, "block fine-tune," is quite straightforward The only difference from the stan- dard SFT process is the modification of the traditional lower- triangular attention mask matrix to the attention mask matrix depicted in the right part of Figure [ With this masking matrix, tokens in all blocks except the last are restricted to attending only to information within their own block, thereby ensuring consistency between training and inference.
which we refer
2.5 INFERENCE
In inference, the Block-attention model retrieves the KV states of blocks from the KV cache and concatenates them into the complete input KV states. The detailed process the inference stage is depicted in Figure Initially, query and extract the KV states of the first k 1 blocks from
4

---

Published as a conference paper at ICLR 2025
the cache. Then, based on the position of each block within the input sequence, calculate their positional encoding using Equation 3 Finally, using the KV states of the first i blocks, we compute the KV states of the last block well the model output. In the RAG scenarios. the last block the user query.
3 EXPERIMENTS
After the above analysis, there exist three concerns about the Block-attention method: 1) Can the Block-attention model achieve the same accuracy full-attention in multi-block scenarios such as RAG and ICL? 2) Can the Block-attention model still adapt to the full-attention mechanism? 3) How much can the Block-attention mechanism improve the efficiency? The following experimental results will reveal the answers to these three questions. In Sections 3.5 we explored the answers the first two questions and analyzed the accuracy of Block-attention models 11 diverse tasks. Meanwhile, in Section 3.6, we demonstrate the efficiency of Block- in RAG scenarios.
3.1 DATASETS
Train Dataset The first part of our training set comes from the SFT dataset of Tulu For the samples in this dataset, divide them into blocks according to three simple rules: 1) If it multi turn sample, then we divide each turn (a user message and an assistant message) into an independent block; 2) The system message and the user message are assigned two different blocks; 3) We directly use some newline characters, such as "\nin"" "_", "==", "\nkt", as block division labels. That is, when we encounter these characters, we divide the subsequent content into a new block. In this way. 23% of the Tulu3-SFT data can be used for block fine-tuning.
Another part of our training dataset is RAG samples. We randomly sample 20,000 instances from TriviaQA(TQA) (Joshi et al. 2017) and SWikiMuliiHopQQA(WWik (Ho et al. 2020) for fine-tuning models. Each training sample consists of (1) a question. (2) 10 passages retrieved from these two datasets using the Contriever toolki 5 which identifies the 10 most relevant passages, and (3) an answer generated by .lama3.3-77BB Instruutt based on the retrieved passages. The reason for using the Llama3 answer instead of the ground- -ruth answers is that the answer might not be present in our retrieved passages. This discrepancy could lead the model to overlook the content of the retrieved passages and generate outputs directly.
In the experiment, to maintain the full-attention ability of the model, we train the Block-attention model using both full-attention and Block- mechanism simultancously"" In other words, all samples in the training set will be trained in both ways. For the full-ateention baselines, only the full-attention method is used for training.
We evaluate the performance of our proposed Block-attention mechanism and baseline models on four widely-used RAG benchmarks: Natural Questions (NQ) (Kwiatkowski et al. 2019). TriviaQA (TQA) (Joshi et al. 2017), HotpotQA (HQA) (Yang et al. 2018). 2Wiki- MultiHopQA (2Wiki) (Ho et al.) 2020). and NarrtiveQA (NQA) (Kočiský al.J 2017). Following Kandpal et al. (2033) and Liu et al. (2024), we use accuracy our primary evaluation metric, judging whether any correct answers appear in the predicted output. To mitigate biases arising from output length, we set maximum token limit of 200 for the output sequences.
Evaluation Dataset
In addition, we also evaluated the performance of the Block-attention model and the full-attention model in seven benchmarks of the general domain: MMLU (Hendrycks et al. (2021a). BigBench-
Hard (BBH) (Suzgun et al. 2022), DROP (Dua et al. 2019). MATH (Hendrveks et al. 202
GSM8K (Cobbe et al.) 2021). HumanEval (Chen et al. 2021), and IFEval (Zhou et al. 2023
Please note that, among them, BBH, DROP, GSMSK, and MATH are ICL tasks with several in- dependent in-context examples. We will divide each example into separate block. For zero-shot
$ bhttss//hugggngffcccccm https///uuggiifiiccoo/ussttaaaat www.e//acccclllllliiiia..------a
The details about our processed datasets can be found httpa//////.///gttbbbcc//////rpp
Block- - Aet.ttteeiiiiin.... on
*Since do not have access to the OpenAl API, we are unable to conduct evaluations datasets that
require subjective assessment.
5

---

Published as a conference paper at ICLR 2025
Input Prompt
You intelligent Al assistant. Pleuse answer questions based the user's instructions. Below
are some reference documents that may help you in answering the user's question. 'n
Dec-
-Title: Munhattan Life Insurance Company in Manhattan Life Insurance Company, incorporated on May 29, 1850, life......In
Doo
Title: New York Life Building in
Life Insurance Huilding, York, located Madison Avenue, Manhattan, New York City, across......Iu
User Quay
Please write high-quality answer for the given question using only the provided search
documents (some which might be irrelevant).
Question: Which lower taller, 712 Fifth Avenue Munhuttan Life Insurance Building?
Output
The 712 Fifth Avenue tower
Figure 3: The Inference Pipeline of Block-attention Model. The retrieved documents at the top have the highest relevance to the user query.
datasets like MMLU, HumanEval, and IFEval, the block attention model will directly switch to the full-attention mode.
3.2
INPUT FORMAT
The format of input prompt for all datasets follows Liu et al. (2024). For retrieved passages, we concatenate them in ascending order of retrieval score. An example is shown in Figure
3.3 BASE MODEL & BASELINES
(denoted as Tulu3- SFT The reason why choose this model is that it is model that be reproduced (They released their SFT data). By comparing different fine-tuning methods with the same data, we can easily demonstrate the effectiveness of Block-attention in different scenarios. After applying block fine-tuning, the model is denoted as -ulu3-block-ft. For comparison, we also implemented four baselines and three ablated models.
We implement the Block-attention mechanism on Llama-3.1-Tulu----BB-
Tulu3-SFT: Our base model, which also serves as the performance ceiling for Block- model's general
attention models in general tasks. We aim to make the
performance approach this benchmark as closely as possible.
Tulu3-RAG: Since Tulu3-SFT was trained on RAG data, it demonstrates suboptimal performance in RAG tasks. To ensure fair comparison, we conducted supervised fine- tuning of Tul using our training data with full-attention mechanism. This model establishes the performance ceiling for Block-attentic models in RAG scenarios. We sim- ilarly aim to have the Block-attention model's RAG performance approach this upper bound as closely as possible.
A bascline model that applies the PromptCache (Gim et al., 2024 the Tulu3-RAG model. This enables training-free reuse of attention states across different prompts of large language models (LLMss..
Seluu-RAAG-romppe
Tulus-RAG-SSpeeppssttoo Serving as a baseline model, it uses the Superposition method the TAG model (Merth et al. 2024). This method allows the LLM process input documents along parallel prompt paths. This parallel processing mechanism not only en- hances the model's 's processing speed but also optimizes resource utilization by eliminating unnecessary computations.
TuluS-block-ft-full An ablated model that switches the Tulu3-block to the full-attention mode. We want observe through the performance of this ablated model whether the Block-attention and the full-attention can coexist in the same model.
hhtpss//tuggingffaccc.lllllllll wttp://c....o//lllaaaaaaa Jama-3.1 1--ee.3.--Tuu-------- S-TSST

---

Published as a conference paper at ICLR 2025
Models
2wiki
HQA
NQ
TQA
Tulu3-SFT
62.0
68.4
58.6
75.7
Tulu3-RAG
73.2
74.8
61.5
75.8
Tulu3-RAG-Sun Sa-pppprrsuperooiittnnnn
30.1
32.3
35,9
58.9
Tuuu3-RAG-promppcCo
32.4
31.6
44.4
61.8
Tulu3-block-ft
72.2
72.3
60.4
75.1
Tulu3-block-fffflll
73.6
75.2
62.2
76.2
Tubuu-block-ftww/orr
68.9
69.9
59.2
74.4
Tulu3-block-woofft
42.9
42.1
48.3
66.5
Table 1: Accuracy of different models four RAG benchmarks.
Task Type
General HumanEval
ICL
dataset
IFEval
MMLU
GMSSK
MATH
BBH
DROP
setup
0-shot
O-shot
O-shot
4-shot
4-shot
3-shot
3-shot
Tulu3-SFT
68.5
58.5
63.7
75.5
29.2
68.5
9.4
Tulu3-RAG
68.3
65.2
63.6
75.6
28.6
68.5
10.4
Tulu3-block-ft
70.0
59.1
63.0
75.7
28.8
65.3
14.4
Table 2: Accuracy of different models on seven general benchmarks. For the first three zero-shot benchmarks, the Block-attention will fall back to full-attention. For the subsequent four ICL datasets with few-shot examples, each sample will be divided into an independent block. Therefore, for k-shot sample, it will be divided into k+1 blocks.
Tulu3-block-wooft. An ablated model that transitions the attention mechanism of Tulu3- RAG to Block-attention mechanism without any block fine-tuning. The outcomes of this model represent the lower bounds for the Block-attention model's effectiveness, given that it has not undergone any adaptation to Block-attention during the training phase.
Another ablated model that is also fine-tuned to adapt the Block- attention mechanism, while no additional position re-encoding operations described in Sec- 2i.3 are conducted. This model will be used to evaluate the effectiveness of the proposed position re-encoding process.
Tulu3-block-w/o
3.4 EXPERIMENTAL SETUP
Training and Inference All experiments are conducted using 8 NVIDIA H20 GPUs with follow-
ing hyper-parameters: (1) learning rate a 2 x 10-5; (2) batch size b 64: (3) epochs n 1; and
(4) 20 warmup steps. The DeepSpeed 100 and Flash-Atter Aatsh-Atteniion (Dao et al.
2022 toolkits are utilized
accelerate our training procedure using bíloat16 format. Additionally, Flash Aess--Atention is utilized for efficient inference of our fine-tuned models and baselines.
Evaluation OpenCompass toolkit (Contributors 2023 is used to evaluate the performance of models under general benchmarks and RAG benchmarks.
3.5
MAIN RESULTS
From the results in Table we can draw four important conclusions:
1) It is not advisable to directly switch from full-attention to parallel prompt encoding methods including Block-attention, as it will lead sharp drop in accuracy. For instance, as can be seen from the experimental results of malu3-block-w/o-ff removing the Block fine-tune process causes the Tulu3-RAG model experience an average absolute performance decrease 16.2% across all four RAG benchmarks. Additionally, the performance degradation of Tulu3-prompiCache and Tulu3-Sune Suttt--uuppposttiion is even more significant, which also indicates that the current parallel prompt encoding methods are still far from satisfactory. Through case studies, found that these models have serious degeneration problems.

---

Published as a conference paper at ICLR 2025
2) However, if use the Block-attention mechanism in the fine-tuning stage, then the resultant model has comparable performance with its full-attention counterpart. In RAG scenarios, the per- formance gap between Tulu3-block-ft and T-la3-RAG on four benchmarks is less than equal to 1%. Meanwhile, Tulu3-block-ft significantly outperforms Tulu3-SFT on 2wiki and HQA. In the ICI. scenario, Tulu3-block-fi also has overall performance almost the same as than that of the two full-attention models. This conclusion indicates that in RAG and ICL scenarios, is completely feasible to replace full-attention with Block-attention, and there will be almost no performance loss.
slightly higher
3) Block and full attention can seamlessly transition. When we switch 3 the full- attention mode, we are pleasantly surprised to find that, whether in the RAG scenario (denoted as Tullu--blockfftfuulll or the naturally single-block zero-shot scenario (IFEval, HumanEval, and MMLU), the Block-attention model can achieve performance comparable to or even slightly higher than that of the two strong full-attention baselines. This conclusion, together with the previous one, indicates that using both Block-attention and full-attention during training can lead to the acquisition of a model that is capable of seamlessly transitioning between the two attention mechanisms.
4) The position re-encoding operations are essential for the Block-attention Model. Removing it leads to a certain degree of performance drop- an average 2% decrease in accuracy on all RAG datasets. Additionally. in this case, the model occasionally experiences the problem of degeneration.
Finally, one may still be interested in knowing exactly how many training steps needed for the model to adapt to the Block-attention mech- anism. Therefore, we counted the average ac- curacy of Tulu3-block-ft and Tulu3- tulu3-block-fi-full on four RAG benchmarks tuning steps and plotted it in Figure 4 It can be observed that the beginning stage of fine- tuning, there is a huge performance difference between the two attention modes. makes sense because the model needs more training steps to adapt to the Block-attention manner. After about 800 training steps, the model can switch seamlessly between the two attention modes without any performance loss.
Accuracy over Training Steps
70
65
different fine-
Accuracy
60
55
Full attention
50
Bloc k-attentioo
200
400
1200
Training Steps
Figure 4: The accuracy of model checkpoint
3.6 INFERENCE EFFICIENCY OF BLOCK-
In the previous section, already addressed first concern: After fine-tuning, the Block-attention model can achieve similar or even better performance than the full-attention model. In this section, focus on our third concern: How much the Block-attention mechanism reduce the TTFT and FLOPs-TFT?
To quantify the effects of the Block-attention mechanism the efficiency, we show in Table the respective TTFTs and FLOPs-TFT of the Llama3-block-ft and Llama3-vanilla-se when the KV states of all retrieved passages have been pre-computed and cached in memory. Obviously, the acceleration effect is gratifying: Once the length of the input sequence is 512 and the length of user input is 50, using Block-attention can reduce TTFT by 48% and reduce FLOPs-TFT by 90.1 %. As the total length increases, the TTFT and FLOPs-TTF of the Block-attention model maintain an essentially unchanged trend. When the length reaches 32K, the acceleration effect reaches an astonishing 98.77 78,7% and the consumption of FLOPs-TFT is even reduced by 99.8%. We may simply conclude the results as: with greater text comes greater necessity for ftock-attention.
3.7 DISCUSSION
From the experimental results, we can figure out the effects of Block- Block-attention on existing RAG ap- plications: Under the existing technical paradigm, developers need to deliberately balance the trade- off between accuracy and latency. Therefore, they have to limit the number retrieved passages to certain number, such as 3 to 10. With Block-attention, the impact of the number of retrieved
8

---

Published as a conference paper at ICLR 2025
Prompt Length
50
512
1K
2K
4K
8K
16K
32K
TTFT-vanilla
26
50
87
167
330
691
1515
3638
TTFT-block
26
26(48%)
26(71%)
26(84%))
27(91%)
29(95%)
34(97%)
45(98.7%)
FLOPs-TFT-vanill
7,5c+11
7.6c+12
1,5c+13
3.0c+13
6,1c+13
1.2c+14
2,45c+14
4,9c+14
FLOPs-TFT-block
7.5e+11
7.5++11
7.5e+11
7.5e+11
7.5e+11
7.5e+11
7.5e+11
7.5e+11
Reduction
90.1%
95.0%
97.5%
98.7%
99.3%
99.6%
99.8%
Table 3: The Time and FLOPs consumed by the first token of a user input with length of 50 tokens under different total length of the retrieved passages
passages on inference speed will be greatly reduced, which will empower the developers to freely choose the number of retrieved passages that gives the optimal effect without any hesitation.
While our experiments in this paper focus on the RAG and ICL scenario using publicly available datasets,
transformative potential extends far beyond these domains. Through validation across multiple internal tasks, we observed consistent efficiency gains. However, due confidentiality constraints preventing data/code disclosure, we strategically selected the RAG scenario for reproducibility and benchmarking purposes.
Notably, Appendix A details how Block-attention addresses critical latency challenges in Game AL. We believe that Block-attention's 's true disruptive potential lies in enabling real-time LLM agents- a vision technically unattainable through full-attention due to the constraints of inference costs.
4 RELATED WORKS
4.1 CONCURRENT WORK
We were surprised to find that another ICLR 2025 submission, TurboRAG (Lu & Tang. 2025), in- , dependently proposed similar methods and reached similar conclusions as ours. They also proposed independent attention (which is the block attention in this paper) and reordered positions (the posi- tion re-encoding in this paper). Unfortunately, their work was not accepted by ICLR 2025. We sin- cerely hope that the community will also recognize them as one of the pioneers of Block-attention, and we hope that this similarity will not affect the publication of their paper in the future. Another study, DecoupledRAG (Dong et al.) 2025). addresses the inefficiency issues in traditional Retrieval- Augmented Generation (RA methods by employing cross-attention mechanism to directly inject external knowledge into the LLM's reasoning process.
Recently, two highly discussed papers in the sparse attention domain-N5t (Native Sparse At- tention) from Deepseek (Yuan et al. 2025) and MoBA (Mixture of Block Attention) from Moon- shot (Lu al., (2025) independently proposed block-based attention mechanisms similar to our Block-attention While these works are not strictly parallel to ours (their preprints were re- leased five months after ours), we aim to clarify the distinctive contributions and orthogonality of our approach compared to these studies. After partitioning inputs into blocks, these methods em- ploy trainable block selection operation to filter out irrelevant blocks. The retained blocks are then concatenated and processed in an auto-regressive pre-filling manner. Our Block-attention, in the contrary, focuses on two novel aspects: 1) Parallel context encoding, and 2) Cross Cross-prompl block KV cache reuse. Neither NSA nor MoBA supports these capabilities.
In our prior work Zhang et al. (2024), we demonstrated that even simple block selection mech- mnisms-when applied to parallel context encoding can substantially improve model accuracy. This suggests the "Block-attention we proposed" addresses complementary optimization dimensions compared the "Block-attention proposed by DeepSeek and Moonshot". In their future models, they may take our work as another pathway to significantly reduce their inference cost further.
4.2 RETRIEVALLAAGGEENTED GENERATION (RAG)
RAG is a widely used technique to improve generations of language models by using retrieved nearest-neighbor documents passages as references, which typically involves two stages: retrieval

---

Published as a conference paper at ICLR 2025
and generation. Before generation, retrieval finds most similar passages with the user query the context, by using BM25 or dense retrieval model (Lee et al. 2021 Lan et al. 2023: 2024 Ma
et al. 2025 Che al. 2024). After collecting retrieved passages, there are numerous techniques
to incorporate the knowledge during generation. Earlier works includes concatenation (Izacard &
Grave 2021), cross-attention (Borgeaud et al. 2022] and distribution interpolation Khandelwa al.
(2020). Some studies have also begun to attempt to directly use the retriever as text generator, that
1s, text generation performed by selecting context-aware phrases from a collection of supporting
documents Lann et al., 2023 Cao et al., 2024).
Recently, LLMs becomes the most powerful paradigm for most NLP tasks, and simply concatenating all retrieved documents into the context of LLMs becomes the most simple and effective way for retrieval- -uggeentd generation (RAG). For example, Self-RAG (Asai et al. 2023) leverage a critic model to decide which content in retrieved passages should be used during generation. As specific
application of RAG, tool learning widely used call external APIs to retrieve related passages
from external database or tools to solve knowledge-intensive tasks (Schick et al. 2023
4.3 PARALLEI CONTEXT ENCODING
Some researches focus on individually and parallely process each documents, which is related to our work, such as SGLang (Zheng et al. 2024), FiD ([zacard & Grave, 2021), PCW (Ratner et al. 2023), PromptCache (Gim al. 2024 and CacheBlend Yao et al. 2024 Although SGLang could maintain the comparable generation quality as the original model, its reuse conditions are very strict, making it difficult to improve the inference efficiency. FiD is widely used for encode- decoder architecture, and not compatible for the decoder-only architecture of LLMs, since they need concatenate the hidden states for decoder during inference. PCW focuses on extending the context window rather than efficient inference. The acceleration effect of PromptCache is similar to Block- Attention; however, as shown our experiments, due to not properly handling positional encoding, its RAG performance is rather poor. CacheBlend introduces a trade-off between generation quality and KV cache reuse efficiency. Compared to SGLangg (Zheng et al.) (024). CacheBlend achieves higher reuse efficiency and lower TTFT but at the cost of slightly reduced generation quality. On the other hand, compared to PromptCache, CacheBlend incurs overhead due to cross- cross-attention recovery but delivers better generation quality.
From the above analysis, one may derive an intriguing conclusion: existing studies are forced to compromise either on generation quality or inference efficiency. Some works maintain high quality with low efficiency (SGLang), others prioritize high efficiency (PromptCache), while some attempt to balance both aspects (CacheBlend). In contrast, our proposed Block-Attention eliminates the need for such trade-offs through position re-encoding and block fine-tuning. It achieves both the same generation quality as the original model and the reuse efficiency comparable PrompiCache, with- out introducing additional overhead. More importantly, the resulting model even achieves seamless switching between Block-Attention and full-attention modes, significantly enhancing the flexibility of online services.
5 CONCLUSION
We introduce Block-attention to optimize the inference efficiency of LLM in RAG scenarios. Its essence lies in separately calculating the KV states of independent blocks in the input prompt. Block-attention enables us to pre-compute the KV states for all passages and cache them in memory, thus avoiding repeated computation of the KV states of the same passage during inference.
Our experimental results across various RAG benchmarks demonstrated the profound impact of Block-attention. We showed that Block-attention can maintain the original reasoning accuracy of LLM while significantly reducing its TTFT and FLOPs-TTF in RAG scenarios. The effectiveness of Block-attention becomes increasingly apparent as the number of passages increases or the fre- quency of retrievals increases. When considering the implementation of Block-attention in your applications, bear in mind this guiding principle: With greater text comes greater necessity for Block-attention.
10

---

Published as a conference paper at ICLR 2025
6 ACKNOWLEDGEMENTS
We sincerely express our gratitude to our friends, JCY, ZZL, HWY, and QXT, for their insightful suggestions and unwavering support during the inception of this idea. Unfortunately, due their af- filiation with confidential institution, we are unable to include them in the author list. We earnestly hope that they will eventually enjoy a more open research environment.
We also extend our heartfelt thanks to our colleagues at Tencent AI Lab- srin-Xintin Huang. Tian Liang, Jiahao Xu, and Zhaopeng Tu-for their valuable advice and resource support during both the rebuttal stage and the camera-ready phase of this work.
REFERENCES
Akari Asai, Zeqiu Wu, Yizhong Wang, Avirup Sil, and Hannaneh Hajishirzi. Self-rag: Learning to retrieve, generate, and critique through self-reflection, 2023. URL httts:://axxi......//
abs/2310.1152
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Mil- lican, George van den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego de Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irv- ing, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack W. Rae, Erich Elsen, and Lau- rent Sifre. Improving language models by retrieving from trillions of tokens, 2022. URL atttps://arxiae org/ abs 204422
Bowen Cao, Deng Cai, Leyang Cui, Xuxin Cheng, Wei Bi, Yuexian Zou, and Shuming Shi. Retrieval is accurate generation. In The Twelfth International Conference on Learning Representations, 2024. URLpttts:///peneeeee
net wtt.t//oorrmmmmdo22200..
Tian- Yi Che, Xian-Ling Mao, Tian Lan, and Heyan Huang. A hierarchical context augmentation method to improve retrieval-augmented Ilms scientific papers. In Pruceedings of the 30th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, KDD '24, pp. 243-254, New York, NY. USA, 2024. Association for Computing Machinery. ISBN 9798400704901. doi: 10.1144/3637528.3661
(RLLhttss://dbll wrg/10. 363145/3333288888 3671847
Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings. Matthias Plappert, Fo- tios Chantzis, Elizabeth Barnes, Ariel Herbert- Voss,, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob Me- Grew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba. Evaluating large language models trained on code. 2021.
Karl Cobbe, Vincet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reicchro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint arXiv:2 10, 14168, 2021.
OpenCompass Contributors. Opencompass: A universal evaluation platform for foundation models. 2023.
https://gtt wtwp.tcccoopeeeecoomoooooopooppp wwtp:///ppeeccccoppps///eeeoneee
Tri Dao, Daniel Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and mefory-efficient exact attention with io-awareness, 2022. URL
wwwpsiiii.rgaaaaao abs/
2205. 14135
Qian Dong, Qingyao Ai, Hongning Wang, Yiding Liu, Haitao Li, Weihang Su, Yiqun LIU, Tat-Seng Chua, and Shaoping Ma. Decoupling knowledge and context: An efficient and effective retrieval
11

---

Published as a conference paper at ICLR 2025
augmented generation framework via cross attention. In THE WEB CONFERENCE 2025, 2025.
URL httpps://ppenreviev..eee. http://wwww..........eeeoeeeee.e ceret//frummidde Law.ummm--d--EEEV.E.... Srt.:/aw.........oo.oo.o...EEaa.
Dheeru Dua, Yizhong Wang. Pradeep Dasigi. Gabriel Stanovsky, Sameer Singh, and Matt Gardner.
Drop: A reading comprehension benchmark requiring discrete reasoning over paragraphs. In North American Chupter of the Association for Compulational Linguisties, 2019. URL https:
//api. www.ntmattt.... wwsiticccchlaal. www.orr///eeeeeeeee CorpusTD: :TTTppsuTD:6888888881 67855846
In Gim, Guojun Chen, Seung seob Lee, Nikhil Sarda, Anurag Khandelwal, and Lin Zhong. Prompt cache: Modular attention reuse for low-latency inference, 2024. URL htt.bttsss//arii..rrt
abs/2311. 04934
Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, and Jacob Steinhardt. Measuring massive multitask language understanding. Proceedings of the Interna- fional Conference on Learning Representations (ICLR), 2021a.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang, Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. NeurIPS, 2021b.
Xanh Ho, Anh-Khoa Duong Nguyen, Saku Sugawara, and Akiko Aizawa. Constructing a multi- hop QA dataset for comprehensive evaluation of reasoning steps. In Donia Scott. Nuria Bel, and Chengqing Zong (eds.), Proceedings of the 28th International Conference on Computational Linguistics, pp. 6609-6625, Barcelona. Spain (Online), December 2020. International Com- mittee on Computational Linguistics. doi: 10. 580
http://00..0..ggmmaaa..... URL https:
//aclanthar s/aclanthology.o org. cowinollinnnooo
Gautier Izacard and Edouard Grave Leveraging passage retrieval with generative models for open domain question answering. In Paola Merlo, Jorg Tiedemann, and Reut Tsarfaty (eds.), Proceed- ings of the 16th Conference of the European Chapter of the Association for Computational Lin- guistics: Main Volume, pp. 874-880, Online, April 2021. Association for Computational Linguis- tics. doi: 10.18653/v1 eecl-mmii. 74. URL https: ac www.ennnnnoooot.oooooooooo org. Dacl-main.17
Mandar Joshi, Eunsol Choi, Daniel Weld, and Luke Zettlemoyer. TriviaQA: A large scale distantly supervised challenge dataset for reading comprehension. In Regina Barzilay and Min- Yen Kan (eds.), Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 1601-1611, Vancouver, Canada, July 2017. Association for Com- putational Linguistics. doi: 10011853/v1/P17-111 URL
org/
P17-1147
Nikhil Kandpal, Haikang Deng, Adam Roberts, Eric Wallace, and Colin Raffel. Large language models struggle learn long-tail knowledge, 2023. URL htttp:://erxiv.orr////
2211.08411
Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models, 2020. URL htttps://arxiv.
ht///bs/113 1911 00172
Tomáš Kočiský, Jonathan Schwarz, Phil Blunsom, Chris Dyer, Karl Moritz Hermann, Gábor Melis, and Edward Grefenstette. The narrativeqa reading comprehension challenge, 2017. https:
//arxiv.org//aae /1w.2 0711007040
Tom Kwiatkowski, Jennimaria Palomaki, Olivia Redfield, Michael Collins, Ankur Parikh, Chris Alberti, Danielle Epstein, Illia Polosukhin, Jacob Devlin, Kenton Lee, Kristina Toutanova, Llion Jones, Matthew Kelcey, Ming-Wai Wei Chang, Andrew M. Dai, Jakob Uszkoreit, Quoe Le, and Slav Petrov. Natural questions: A benchmark for question answering research. Transactions of the Association for Computational Linguistics, 7:452-466, 2019, doi: 1621/a/ 00276. URL htttss://aclawe wwwpnthhhooy..oor//////
Tian Lan, Deng Cai, Yan Wang, Heyan Huang, and Xian-Ling Mao. Copy is all you need, 2023.
URL htt.hhttp:/aaaae wtt.g///aa//....... 06962
12

---

Published as a conference paper at ICLR 2025
Tian Lan, Deng Cai, Yan Wang, Yixuan Su, Heyan Huang, and Xian-Ling Mao. Exploring dense retrieval for dialogue response selection, ACM Trans. Inf. Syst., 42(3), jan 2024. ISSN 1046-8188. doi: 10.1145/3632750.
Jinhyuk Lee, Mujeen Sung, Jaewoo Kang, and Danqi Chen. Learning dense representations of phrases scale, 2021. URL htt.htttss///xii...o.om
abs /212 12624
Huayang Li, Yixuan Su, Deng Cai, Yan Wang, and Lemao Liu. A survey on sttreval-augmentt text generation, 2022. URL httph::tts/attt httpss//aaxiv.ooo / abs/ wttp://wpww...//....ooooo 01110
Yucheng Li, Bo Dong, Frank Guerin, and Chenghua Lin. Compressing context to enhance in- ference efficiency large language models. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Pro- cessing, pp. 5342-6353, Singapore, December 2023, Association for Computational Linguis- tics. doi: 112116653/1/2222332 http:/222...mmmmmmmaa...33333. URL /t///ttpss///eeaaallooogg...... @mm paanlp-mminn 391
Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024. doi: 10. a a006338. URL
https: aaaaclle wew.annnhc...ooooooo conthology. ogy.ooy....o.................... org/ taw.-1......11......
Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong. Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang. Zhiqi Huang, Huan Yuan. Suting Xu, Xinran Xu, Guokun Lai, Yanru Chen, Huabin Zheng, Junjie Yan, Jianlin Su, Yuxin Wu, Neo Zhang, Zhilin Yang. Xinyu Zhou, Mingxing Zhang, and Jiezhong Qiu. Moba: Mixture of block attention for long-context llms, 2025. URL http://////aaiioooo/////////////
13189
Songshuo Lu and Yaohua Tang. TurboRAG: Accelerating retrooralllcvllaaamatotte generation with pre- computed KV caches for chunked text, 2025. URL UtL.
1d=x7NbaUERSU
Zi-Ao Ma, Tian Lan, Rong-Cheng Tu, Yong Hu, Yu-Shi Zhu, Tong Zhang, Heyan Huang, and Xian-Ling Mao. Multi-modal retrieval augmented multi-modal generation: Datasets, evaluation metrics and strong baselines, 2025. URL /tthttts///aaa/t wttp://www..rr///////111113.1..
Thomas Merth, Qichen Fu, Mohammad Rastegari, and Mahyar Najibi. Superposition prompting: Improving and accelerating retrieval-augmented generation, 2024. URL /tttps://arxiv.
erq/abs/2404 06910
Ruoyu Qin, Zheming Li, Weiran He, Mingxing Zhang, Yongwei Wu, Weimin Zheng, and Xin- ran Xu. Mooncake: kvcache-centric disaggregated architecture for 1l serving, 2024. URL
https://arxite wttp://wwwwwppp//...ooooo 00079
Nir Ratner, Yoav Levine, Yonatan Belinkov, Ori Ram, Inbal Magar, Omri Abend, Ehud Karpas, Brown, and Yoav Shoham. Parallel context windows for large Amnon Shashua, Kevin Leyton-Brown,
language models. In Anna Rogers, Jordan Boyd-Graber, and Naoaki Okazaki (eds.), Proceed- ings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume I: Long Papers), pp. 6383-6402, Toronto, Canada, July 2023, Association for Computational Lin- guisties. doi: 10.18653/1//023..coo
URL /ttp:::::/aacaaaa http://////aaaaahcccacgee wwtp://ooooyy.//2233...3/
acl-long. 352
Timo Schick, Jane Dwivedi- Yu, Roberto Dessi, Roberta Raileanu, Maria Lomeli, Eric Hambro, Luke Zettlemoyer, Nicola Cancedda, and Thomas Scialom. Toolformer: Language models can teach themselves to use tools. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL
Mirac Suzgun, Nathan Scales, Nathanael Schärli, Sebastian Gehrmann, Yi Tay, Hyung Won Chung, Aakanksha Chowdhery, Quoc V Le, Ed H Chi, Denny Zhou, , and Jason Wei. Challenging big- bench tasks and whether chain-of saiin-of-thought can solve them. arXiv preprint arXiv:2210.09261,, 2022.
13

---

Published as a conference paper at ICLR 2025
Oriol Vinyals, Igor Babuschkin, Wojciech M. Czarnecki, Michaël Mathieu, Max Jaderberg, Andrew Dudzik, Junyoung Chung, David H. Choi, Richard Powell, Timo Ewalds, et al. Grandmaster level in StarCraft II using multi-agent reinforcement learning. Nature, 575((7822)350- 354, 2019. doi: 10. (0008//41158-019--77722
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W. Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. Hotpotqa: A dataset for diverse, explainable multi-hop question answering, 2018. URL Btl.hhtss://arxtt htttpp://arx/vv.rrrr/a////////// wwt./ergg/bbbbeeeeeeeee
Jiayi Yao, Hanchen Li, Yuhan Liu, Siddhant Ray, Yihua Cheng, Qizheng Zhang, Kuntai Du, Shan Lu, and Junchen Jiang. Cacheblend: Fast large language model serving for rag with cached knowledge fusion, 2024. URL aa//aaxxvvv. org/ abs/ 20t.:/2005.14444
Deheng Ye, Zhao Liu, Mingfei Sun, Bei Shi, Peilin Zhao, Hao Wu, Hongsheng Yu, Shaojie Yang, Xipeng Wu, Qingwei Guo, Qiaobo Chen, Yinyuting Yin, Hao Zhang. Tengfei Shi, Liang Wang. Qiang Fu, Wei Yang, and Lanxiao Huang. Mastering complex control in moba games with deep reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence. 34(04); 6672-66792 Apr. 2020. doi: 10.1609/aaai, v34i04.6144. URL https: aaai. .org. index.
/ 6144
Jingyang Yuan, Huazuo Gao, Damai Dai, Junyu Luo, Liang Zhao, Zhengyan Zhang, Zhenda Xie, Y. X. Wei, Lean Wang, Zhiping Xiao, Yuqing Wang, Chong Ruan, Ming Zhang, Wenfeng Liang, and Wangding Zeng. Native sparse attention: Hardware aligned and natively trainable sparse attention, 2025. httptttss//xxxx..o/////aaaa
11089
Zhisong Zhang, Yan Wang. Xinting Huang, Tianqing Fang, Hongming Zhang, Chenlong Deng, Shuaiyi Li, and Dong Yu. Attention entropy key factor: An analysis parallel context encoding with full aalll--tteniion-aaae based praitee language models, 2024. URL httpss//a/// org/abs/ 2412. 16545
Lianmin Zheng. Liangsheng Yin, Zhiqiang Xie, Chuyue Sun, Jeff Huang, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E. Gonzalez, Clark Barrett, and Ying Sheng. Sglang: Efficient execution of structured language model programs, 2024. URL
org/abs/2312 07104
Jeffrey Zhou, Tianjian Lu, Swaroop Mishra, Siddhartha Brahma, Sujoy Basu, Yi Luan, Denny Zhou, and Le Hou. Instruction-followii evaluation for large language models. arXiv preprint arXiv:2311.07911, 2023.
Yun Zhu, Jia-Chen Gu, Caitlin Sikora, Ho Ko, Yinxiao Liu, Chu-Cheng Lin, Lei Shu, Liangchen Luo, Lei Meng, Bang Liu, and Jindong Chen. Accelerating inference of retrieval-augmennte gen- cration via sparse context selection, 2024. URL httpt::///// /tatpps///axiv...o wtw../...rgabb///25....o. 16178
A
IN GAME AI
Before detailing the application of Block-attention in Game AI, we first outline the characteristics of LLM tasks in gaming scenarios. In such contexts, the LLM input typically consists of the game's current state, known as gamecore data in the industry. This data is structured as JSON with lengths ranging from thousands to hundreds of thousands of tokens. Different tasks produce varied outputs, but the input generally remains the gamecore data the current frame. For example: for Al players like Al AlphaStar (Vinyals et al. 2019) and JueWu (Ye al. 2020), their problem can be formulated as:
P(Player Actionlgamecorv)
For AI NPCs, their problem then becomes:
PC Action gamecore).
As for tasks like Al commentary, they can be defined as:
P(Game Commentss
14

---

Published as a conference paper at ICLR 2025
ssate
Player 25s vien
"coatig ('plassenss l'pl's "pu", 17a
"chape's ('ad's Cbatts "romain's PIC), 1921 ('bects Lb, constats 2902, 1pats
"Grakfa feedle
1:077, "Ce"), "go"s 11877 1831 "D3"T (1012 1BE
"nt'l,
'cice's ('ataccl, "azudcla", "balledt, "deatat, "beet, "Clopa"),
"masage seelr engine", tengt p?". "costent your bet. 310 Raise 501
('ace': 'p2", "top" engien" content Raten to 5011)
suate
Player view
"coutia's C'plasorn's l'ptts "yo", "patt)),
"chipe's ('ll's C'at'' trematn's siC), 1p2"s ('bects 50, ronatats 2500, "pata
"seakte
"sssats 1.27,, "Ce'), "go": 1187, 'HIt I. "p3"r (1022 "BC03).
"closts ("etactly "aaadtte", "balladr, ideatst, "belt, (Clega1),
"maangee' seell engine i "eept po". consbene's yeur turn bet. 319
ccoon ('sce "pl", "tng" a engiee
"content" "0a22"3
Figure A case of Texas hold' em Al's gamecore data
As shown in the formulations above, most gaming tasks involve extremely long contexts (from several K to hundreds K) but relatively short output sequences (tens to hundreds of tokens). If the AI must prefill data from scratch for every decision, the resulting excessive inference latency would render LLMs impractical for real-time applications. This explains why current MOBA and FPS games still rely on small small-parann paraametr models for online deployment.
To illustrate how block-attention reduces TTFT in LLMs, consider Texas Hold' em example. Fig- ure compares gamecore data at two consecutive states: State N-1 (previous) and State (current). Notably, these states are nearly identical except for the value of statef "chip'' I['p2"] which changes from ('ber': 10, 'remain': 990) to {'bet': 50, 'remain': 950). This high inter-frame repetition (exceeding 99.5% based
analysis) is common in games. Furthermore, the structured JSON format of gamecore data allows rule-based partitioning into independent blocks. By encoding only new modified blocks during prefilling (instead of prefilling from scratch), achieve equivalent accuracy to full prefilling while significantly reducing latency.
In unreleased game (not yet available to the public), by dividing the game state into over 300 independent blocks, we achieved the same accuracy as full-attention mechanisms in these tasks, while cutting the average TTFT from 2,800ms 100ms and reducing query latency from 3,000ms to under 300ms. Our Block-attention method aims to help Al game developers effortlessly create more games powered by advanced AI systems. We envision that Block-attention will empower AI researchers in the gaming industry to seamlessly develop more Al-centric gaming experiences.
15


Reviews:
[{"review_id": "N2R8vrMIhK", "reviewer": "Reviewer_2J6u", "summary": "This paper introduces the Block-Attention method, an efficient approach tailored for RAG scenarios that enhances inference efficiency while preserving performance through fine-tuning. Block-Attention operates by dividing the input sequence into multiple independent blocks, each calculating its key-value (KV) states separately via self-attention. Only the final block can attend to previous blocks, allowing the user query to reference all prior retrieved documents. Evaluations on four RAG benchmarks reveal that Block-Attention, after fine-tuning, maintains comparable accuracy to traditional self-attention models (e.g., 68.4% vs. 67.9% on Llama3) and even slightly outperforms in some cases (e.g., 62.3% vs. 59.6% on Mistral). The efficiency benefits, measured in TTFT and FLOPs, grow with input length; at 32K input length, Block-Attention’s TTFT and FLOPs reach just 1.3% and 0.2%, respectively, of those in standard self-attention models.", "questions": "Only tested on relatively small models (7B-8B parameters) - would it work as well on larger models?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- Novel and practical idea\n- easy to follow\n- Simplicity: The solution is relatively simple to implement and can be integrated into existing LLM architectures.\n- Impressive performance: this method reduces Time to First Token (TTFT) by up to 98.7% and FLOPs by up to 99.8%.", "weaknesses": "### Needs for finetuning:  \nThe method requires additional fine-tuning, which might be resource-intensive for larger models. It would be better to explore a training-free approach with the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YMy0wSDsLp", "reviewer": "Reviewer_wUEZ", "summary": "The paper proposes Block Attention for retrieval augmented generation (RAG). The idea is to confine attention within each retrieved document (a.k.a. block) and only let the query attend to all documents. In this way, the KVs for each document can be cached and reused if the document re-appears in multiple queries. The method also proposes to explicitly correct RoPE positional encoding before reusing the cached KVs. After fine-tuning, this technique can retain the accuracy of the baseline while speeding up the inference due to smaller attention span and avoiding the KV recomputation thanks to the block caching.", "questions": "- Additional questions and comments:\n\n1) To my understanding, the main technical contribution of the method compared to the previous approaches is the proposal to explicitly re-adjust the positional encodings. What is the practical advantage of the proposed RoPE encoding adjustment? Does it lead to any improvement on the final RAG task quality metrics over PromptCache? From the current manuscript I did not find an ablation showing its importance to improve prior arts. \n\n2) What is the comparison between the proposed method and PromptCache in terms of efficiency? How much does the proposed method improve PromptCache in terms of wall-time?\n\n3) How does the method compare with respect to the Superposition Prompting paper above [1] both in terms of accuracy and efficiency?\n\n4) The paper mentions fine-tuning on RAG-specific datasets necessary for the proposed approach. Expectedly, fine-tuning can improve the accuracy on the specific tasks at hand, but it is also important to maintain the generalizability of the pre-trained models. This can be difficult especially when the training recipes are not available, or they are instruction fine-tuned on non-available data. How does the performance of the reported models change on common benchmarks such as MMLU and HumanEval after the suggested fine-tuning?\n\n5) The positional encoding re-adjustment is only discussed for RoPE. As a suggestion for further improvement, it would be nice to have a more thorough analysis on different positional encoding approaches.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The method is simple and easy to implement: It only involves limiting the attention computation within each document in RAG scenarios and caching the KVs for each block for document re-use as proposed in PromptCache [Gim et. al, 2024]. The newly introduced RoPE rotation correction is also straightforward to implement.\n- Local attention and prompt caching are sound techniques and suitable for RAG applications.\n- The paper is clear and easy to follow.", "weaknesses": "My primary concern is regarding the paper's contribution compared to the existing works and the need for a better experimental analysis to understand where this method stands among previous methods that share very similar ideas or how it improves them. The authors can find more detailed comments and suggestions below.\n\n- Contribution and limitations compared to existing works\n\nTo me the main contributions of the paper to improve the efficiency of RAG can be summarized as 1) limiting the attention within each document 2) caching the KVs of the documents to reuse them 3) Adjusting the rotation in the RoPE position encoding based on the position of the reappearing cached document. Points 1 and 2 have been already proposed in prior works. E.g. PromptCache [Gim et. al. 2024] proposed the same techniques to speed up the RAG applications. This makes contribution 3, the main new addition to the existing works. However, according to the experiments reported in the submission, the RoPE position correction does not lead to any improvement over PromptCache (lines 353-356). To my understanding this makes fine-tuning the only source of improvement over PromptCache, which is an expected gain which comes with known challenges of such a fine-tuning.\n\nI want to also point the authors to [1] below that is a closely relevant prior art which is currently missing from the draft. [1] also addresses the three aforementioned points above. Namely, it proposes limiting attention to attention blocks, caching the KVs for the documents to re-use them, and addresses the position encoding problems by overlaying the documents as parallel edges without requiring any fine-tuning. I would suggest the authors add a discussion on the advantages of the proposed method compared to [1] and also add controlled experimental comparisons to better highlight the contributions.\n\n[1] Superposition Prompting: Improving and Accelerating Retrieval-Augmented Generation, Merth et. al.\n\n* Experimental Results\n\nThe experimental section needs to be revised to clearly highlight the main sources of improvements and contributions between the proposed method and existing prior methods (under same settings, e.g. with or without fine-tuning). Based on the submission, comparisons with similar approaches are missing (see [1] above), additional speed up on top of PromptCache (if any) is not verified, and it is not clear to me if anything other than the fine-tuning is the source of accuracy improvement over PromptCache. The authors can find more detailed comments and questions regarding the experiments in the following section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ln5ECY0vh9", "reviewer": "Reviewer_1DTJ", "summary": "This paper tackles the efficiency challenges in RAG. The key observation is that different paragraphs retrieved in RAG are independent to each other. Therefore, it is possible to perform independent self-attention within each paragraph, and only allow the final question query to attend to all previous passages. This will dramatically improve TTFT in long-context RAG. The authors identify two important steps to recover accuracy after converting traditional self-attention to block attention (proposed in this paper): positional re-encoding and fine-tuning. Their final results show no loss of accuracy and significant speedup in long-context RAG applications.", "questions": "Please respond to my questions in \"Weaknesses\".\n\n===\n\nPost rebuttal: my concerns have been resolved and I will keep my score.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea is simple and reasonable: different retrieved passages are not necessarily related to each other so the attention mask can be sparsified.\n2. Positional re-encoding is an elegant way to solve inconsistencies in positional encodings.\n3. Finetuning results show no loss of accuracy and significant speedup over self-attention baselines.", "weaknesses": "1. The method requires fine-tuning, which limits its scalability.\n2. The evaluation seems to be a bit weak. \n- For example, the method trains on TQA/2Wiki's training set and evaluates on their validation sets. The results on these two benchmarks are not zero-shot and are not very representative. \n3. It's unclear what is the overhead associated with positional re-encoding.\n4. Can you construct some examples where different retrieved passages are related to each other? In this case, will the proposed method fail? For example, you might retrieve several chapters in a textbook, where later chapters are dependent on earlier chapters.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M2oDy6nZk8", "reviewer": "Reviewer_MNR9", "summary": "This paper presents the Block-Attention mechanism aimed at optimizing inference efficiency in retrieval-augmented generation (RAG) scenarios. The authors assert that by modifying the attention calculation to process tokens in independent blocks, they can significantly enhance both accuracy and computational efficiency compared to traditional self-attention. The proposed method is evaluated against well-known RAG benchmarks, and the authors claim notable improvements in inference speed and resource consumption.\n\nWhile the topic is pertinent, the contributions of the paper raise several concerns. The approach, although presented as innovative, relies on established principles of attention mechanisms without sufficiently distinguishing itself from previous work. Additionally, the experimental results, while promising, lack depth in exploring the nuances of the proposed method's effectiveness across varying conditions.", "questions": "1. How do the authors justify the choice of block size in the Block-Attention mechanism, and what implications does this have for different types of queries?\n2. Can the authors elaborate on how the performance of Block-Attention varies with input length and retrieval volume? Are there any thresholds beyond which the performance might degrade?\n3. Given the reliance on fine-tuning, how do the authors plan to address potential overfitting issues, particularly in diverse RAG settings?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is generally well-organized, with clear explanations of the Block-Attention mechanism and its proposed advantages.\n2. The investigation into improving inference efficiency in LLMs through Block-Attention is timely.\n3. The reported gains in inference speed and reduced computational load are compelling.", "weaknesses": "1. The concept of segmenting attention into blocks is not entirely novel and has been explored in various forms in the literature. The paper does not convincingly articulate how Block-Attention offers distinct advantages over similar approaches, such as sparse attention methods or hierarchical attention mechanisms.\n2. The paper presents empirical results without adequately establishing the theoretical underpinnings of the Block-Attention mechanism. For example, the claim that \"tokens in all blocks except the last are restricted to attending only to information within their own block\" (Section 3.4) needs further elaboration on how this restriction impacts the model's representational capacity compared to full self-attention.\n3.  The experimental section would benefit from a broader set of benchmarks and comparison against state-of-the-art methods beyond just self-attention. \n4. The authors should discuss how the model performs under different conditions, such as varying levels of noise in the retrieved passages or the diversity of user queries.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)