Block-Attention for Efficient Prefilling

1 INTRODUCTION

• RAG prompts commonly include 5–30 retrieved passages, increasing sequence length and TTFT; context-dependent KV states then require re-encoding for an unseen query.The KV states for the same passage vary across contexts, according to the paper.Source: §1 · pp. 1–2
• Block-attention divides the input into blocks whose KV states are calculated independently, while the final block can attend to preceding blocks.In RAG, retrieved passages can each form a block, allowing their KV states to be cached and reused.Source: §1 · p. 2
• To adapt LLMs to Block-attention, fine-tuning took approximately 500–1000 steps and achieved full recovery of original accuracy across all scenarios, increasing accuracy from 49.9% to 70.0%.Source: §1 · p. 2
• After fine-tuning, Block-attention’s average accuracy remained comparable to Llama-3.1-Tulu-3-8B-SFT across 11 RAG, ICL, and general-domain benchmarks.The evaluation includes RAG, in-context learning, and general-domain tasks.Source: §1 · p. 2
• At a 32K input length, Block-attention’s TTFT and FLOPs-TFT are 1.3% and 0.2% of full-attention’s, respectively.These are the reported fractions for Block-attention relative to full-attention.Source: §1 · p. 3

2.1 MAIN IDEA

• Independent blocks compute their own KV states, while the final block integrates preceding blocks; updating a block requires recomputing that block and the final block.This avoids re-encoding subsequent blocks whose KV states would change under autoregressive attention.Source: §2.1 · p. 3

2.2 BLOCK SEGMENTATION

• In RAG, each retrieved passage forms a block and the user query is the final block; semantically independent units can also define blocks in other tasks.Examples include functions in code generation, dialogue turns, and ICL demonstrations; the paper focuses primarily on RAG.Source: §2.2 · p. 3

2.3 POSITION RE-ENCODING

• A reused passage may occupy a different prompt position, so its positional information must be re-encoded before its KV states are reused.For RoPE, the procedure resets the positional encoding to zero and rotates it to the updated position; the paper excludes KV-cache cost from its accounting.Source: §2.3 · pp. 3–4

2.4 BLOCK FINE-TUNE

• Block fine-tuning modifies the standard SFT attention mask so non-final blocks attend only within their own block, aligning training with inference.The paper reports that direct inference-time switching from full attention may cause a substantial performance decrease.Source: §2.4 · p. 4

2.5 INFERENCE

• Inference retrieves cached KV states for the first k −1 blocks, re-encodes their positions, then computes the final block's KV states and model output.In RAG, the final block is the user query.Source: §2.5 · pp. 4–5

3.1 DATASETS

• Training draws on Tulu3-SFT samples split into blocks and 20,000 sampled RAG instances; retrieved passages are concatenated in ascending retrieval-score order.The Tulu3-SFT splitting rules make 23% of that dataset usable for block fine-tuning.Source: §3.1 · p. 5, §3.2 · p. 6

3.3 BASE MODEL & BASELINES

• The experiments implement Block-attention on Tulu3-SFT, compare four baselines and three ablations, and evaluate models using OpenCompass with DeepSpeed and Flash-Attention.Training uses eight NVIDIA H20 GPUs, a learning rate of 2 × 10^-5, batch size 64, one epoch, and 20 warmup steps.Source: §3.3 · pp. 6–7, §3.4 · p. 7

3.5 MAIN RESULTS

• Without block fine-tuning, Tulu3-block-w/o-ft loses an average of 16.2% across four RAG benchmarks; promptCache and Superposition degrade further and show degeneration in case studies.The authors report that direct switching from full attention to parallel prompt encoding leads to a sharp accuracy drop.Source: §3.5 · p. 7
• After block fine-tuning, Tulu3-block-ft is within 1% of Tulu3-RAG on four RAG benchmarks and matches or slightly exceeds full-attention models overall in ICL.It also outperforms Tulu3-SFT on 2Wiki and HQA.Source: §3.5 · pp. 7–8, Table 1 · p. 7, Table 2 · p. 7
• Removing position re-encoding decreases average accuracy by 2% across RAG datasets, and the model occasionally degenerates.The ablation evaluates the position re-encoding operation.Source: §3.5 · p. 8
• After about 800 fine-tuning steps, Tulu3-block-ft can switch between block and full attention without performance loss on four RAG benchmarks.In full-attention mode, it is comparable to or slightly better than strong full-attention baselines in RAG and single-block zero-shot tasks.Source: §3.5 · p. 8, Figure 4 · p. 8

3.6 INFERENCE EFFICIENCY OF BLOCK-ATTENTION

• At 32K input length, Block-attention reduces TTFT by 98.7% and FLOPs-TFT by 99.8% versus full attention, with retrieved-passage KV states precomputed and cached.The comparison uses a 50-token user input.Source: §3.6 · p. 8, Table 3 · p. 9

3.7 DISCUSSION

• The discussion argues that reduced speed sensitivity to passage count could let RAG developers choose retrieval counts for answer quality; experiments cover public RAG and ICL datasets, while internal-task data and code remain confidential.The authors selected RAG for reproducible benchmarking because confidentiality prevents disclosure of internal-task data and code.Source: §3.7 · pp. 8–9

4.1 CONCURRENT WORK

• Block-attention focuses on parallel context encoding and cross-prompt block KV reuse, capabilities NSA and MoBA do not support (Yuan et al., 2025; Lu et al., 2025).TurboRAG independently proposed independent attention and reordered positions similar to this paper’s block attention and position re-encoding.Source: §4.1 · p. 9
• Prior work found that simple block selection with parallel context encoding can improve accuracy, suggesting complementary optimization dimensions.The paper suggests this approach could complement block-selection methods from DeepSeek and Moonshot.Source: §4.1 · p. 9

4.2 RETRIEVAL-AUGMENTED GENERATION (RAG)

• RAG retrieves similar passages using BM25 or dense retrieval, then incorporates them through concatenation, cross-attention, distribution interpolation, or context-aware phrase selection.These approaches span retrieval and generation, including selecting supporting-document phrases during text generation.Source: §4.2 · pp. 9–10
• LLM-based RAG commonly concatenates retrieved documents; Self-RAG uses a critic to select passage content, while tool learning retrieves material for knowledge-intensive tasks.The cited examples are Self-RAG and tool learning (Schick et al., 2023).Source: §4.2 · p. 10

4.3 PARALLEL CONTEXT ENCODING

• Existing parallel-context systems trade off generation quality and efficiency: SGLang has strict reuse conditions, PromptCache has poor RAG performance, and CacheBlend incurs overhead or quality costs.The passage attributes PromptCache’s RAG performance to positional-encoding handling and describes CacheBlend’s trade-offs against SGLang and PromptCache.Source: §4.3 · p. 10
• Block-attention combines original-model generation quality with PromptCache-comparable reuse efficiency, without additional overhead, and switches between block and full-attention modes.The paper attributes this comparison to position re-encoding and block fine-tuning.Source: §4.3 · p. 10

• Block-attention separately calculates KV states for independent prompt blocks, allowing passage states to be precomputed and cached to avoid repeated computation.The method is presented as an inference-efficiency optimization for LLMs in RAG scenarios.Source: §5 · p. 10
• Across RAG benchmarks, Block-attention maintains original reasoning accuracy while reducing TTFT and FLOPs-TTF, with greater effects as passage count or retrieval frequency increases.The conclusion states that its impact becomes increasingly apparent as the number of passages or retrieval frequency rises.Source: §5 · p. 10

• Gamecore data is current-frame game state structured as JSON, with input lengths ranging from thousands to hundreds of thousands of tokens.The input generally remains the gamecore data of the current frame across different gaming tasks.Source: A BLOCK-ATTENTION IN GAME AI · p. 14
• Gaming tasks map gamecore data to different outputs, including player actions and game comments; AI NPCs are another described task type.The passage formulates player action and game comments as conditional on gamecore data.Source: A BLOCK-ATTENTION IN GAME AI · p. 14
• Prefilling long contexts from scratch for every decision would cause excessive latency for real-time applications; the paper links this to current MOBA and FPS games using small models.Source: A BLOCK-ATTENTION IN GAME AI · p. 15
• In the Texas Hold’em example, more than 99.5% of game-state data repeats across consecutive frames, and rule-based partitioning lets prefilling encode only new or modified blocks.The example’s chips field changes between states while the remainder is nearly identical; the passage reports equivalent accuracy to full prefilling with significantly reduced latency.Source: A BLOCK-ATTENTION IN GAME AI · p. 15, Figure 5 · p. 15
• 2,800ms to 100ms average TTFT and 3,000ms to under 300ms query latency, with the same accuracy as full-attention mechanisms, in an unreleased game.The game state was divided into over 300 independent blocks.Source: A BLOCK-ATTENTION IN GAME AI · p. 15