# Towards Generalization of Block Attention via Automatic Segmentation and Block Distillation

arXiv: https://arxiv.org/abs/2605.15913 · HTML: https://mlsanity.com/paper/2605.15913/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Multimodal RAG, one of 15 directions in Retrieval-augmented generation.
- Builds on 32 of its 34 references that are in the map.
- First posted on arXiv 15 May 2026.

## Placed in

- Retrieval-augmented generation / [Multimodal RAG](https://mlsanity.com/rag/multimodal-rag/)
- Efficient LLM inference and serving / [KV cache optimization](https://mlsanity.com/inference/kv-cache-optimization/)

## Builds on (most cited first, 30 of 32)

- [ReAct: Synergizing Reasoning and Acting in Language Models](https://mlsanity.com/paper/2210.03629/) (2022)
- [HotpotQA: A Dataset for Diverse, Explainable Multi-hop Question Answering](https://mlsanity.com/paper/1809.09600/) (2018)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [LongBench: A Bilingual, Multitask Benchmark for Long Context Understanding](https://mlsanity.com/paper/2308.14508/) (2023)
- [DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence](https://mlsanity.com/paper/2606.19348/) (2026)
- [MuSiQue: Multihop Questions via Single-hop Question Composition](https://mlsanity.com/paper/2108.00573/) (2021)
- [Evaluating Very Long-Term Conversational Memory of LLM Agents](https://mlsanity.com/paper/2402.17753/) (2024)
- [LongMemEval: Benchmarking Chat Assistants on Long-Term Interactive Memory](https://mlsanity.com/paper/2410.10813/) (2024)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Scaling up Test-Time Compute with Latent Reasoning: A Recurrent Depth Approach](https://mlsanity.com/paper/2502.05171/) (2025)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [LongBench v2: Towards Deeper Understanding and Reasoning on Realistic Long-context Multitasks](https://mlsanity.com/paper/2412.15204/) (2024)
- [The Stack: 3 TB of permissively licensed source code](https://mlsanity.com/paper/2211.15533/) (2022)
- [Prompt Cache: Modular Attention Reuse for Low-Latency Inference](https://mlsanity.com/paper/2311.04934/) (2023)
- [LongLoRA: Efficient Fine-tuning of Long-Context Large Language Models](https://mlsanity.com/paper/2309.12307/) (2023)
- [Reasoning with Latent Thoughts: On the Power of Looped Transformers](https://mlsanity.com/paper/2502.17416/) (2025)
- [CacheBlend: Fast Large Language Model Serving for RAG with Cached Knowledge Fusion](https://mlsanity.com/paper/2405.16444/) (2024)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [BookSum: A Collection of Datasets for Long-form Narrative Summarization](https://mlsanity.com/paper/2105.08209/) (2021)
- [Scaling Latent Reasoning via Looped Language Models](https://mlsanity.com/paper/2510.25741/) (2025)
- [OpenWebMath: An Open Dataset of High-Quality Mathematical Web Text](https://mlsanity.com/paper/2310.06786/) (2023)
- [Looped Transformers are Better at Learning Learning Algorithms](https://mlsanity.com/paper/2311.12424/) (2023)
- [Liger Kernel: Efficient Triton Kernels for LLM Training](https://mlsanity.com/paper/2410.10989/) (2024)
- [SlimPajama-DC: Understanding Data Combinations for LLM Training](https://mlsanity.com/paper/2309.10818/) (2023)
- [LoopFormer: Elastic-Depth Looped Transformers for Latent Reasoning via Shortcut Modulation](https://mlsanity.com/paper/2602.11451/) (2026)
- [ChatQA 2: Bridging the Gap to Proprietary LLMs in Long Context and RAG Capabilities](https://mlsanity.com/paper/2407.14482/) (2024)
- [Block-Attention for Efficient Prefilling](https://mlsanity.com/paper/2409.15355/) (2024)
- [Attention Entropy is a Key Factor: An Analysis of Parallel Context Encoding with Full-attention-based Pre-trained Language Models](https://mlsanity.com/paper/2412.16545/) (2024)
- [Language Models as Science Tutors](https://mlsanity.com/paper/2402.11111/) (2024)
- [Superposition Prompting: Improving and Accelerating Retrieval-Augmented Generation](https://mlsanity.com/paper/2404.06910/) (2024)

## Abstract (arXiv)

Block attention, which processes the input as separate blocks that cannot attend to one another, offers significant potential to improve KV cache reuse in long-context scenarios such as Retrieval-Augmented Generation (RAG). However, its broader application is hindered by two key challenges: the difficulty of segmenting input text into meaningful, self-contained blocks, and the inefficiency of existing block fine-tuning methods that risk degrading performance. To address these, we first construct SemanticSeg, a large and diverse semantic segmentation dataset containing over 30k instances across 16 categories-including books, code, web text, and conversations with text lengths ranging from 2k to 32k. Using this dataset, we train a lightweight segmenter to automatically partition text into human-instinct-aligned blocks with controllable granularity. Second, we propose block distillation, a training framework that is more efficient than block fine-tuning, which uses a frozen full-attention teacher model to guide the block-attention student. This framework integrates three novel components: block sink tokens to mitigate information loss at block boundaries, block dropout to leverage training signals from all blocks, and token-level loss weighting to focus learning on block-attention-sensitive tokens. Experiments across multiple models and benchmarks demonstrate that our segmenter outperforms heuristic and statistical baselines, and block distillation achieves near-full-attention performance under block attention, establishing a practical and scalable pathway for deploying block attention.