# Training-Inference Consistent Segmented Execution for Long-Context LLMs

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 33 references that are in the map.
- First posted on arXiv 12 May 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [RULER: What's the Real Context Size of Your Long-Context Language Models?](https://mlsanity.com/paper/2404.06654/) (2024)
- [Compressive Transformers for Long-Range Sequence Modelling](https://mlsanity.com/paper/1911.05507/) (2019)
- [FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU](https://mlsanity.com/paper/2303.06865/) (2023)
- [Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention](https://mlsanity.com/paper/2502.11089/) (2025)
- [MInference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention](https://mlsanity.com/paper/2407.02490/) (2024)
- [LongBench v2: Towards Deeper Understanding and Reasoning on Realistic Long-context Multitasks](https://mlsanity.com/paper/2412.15204/) (2024)
- [DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads](https://mlsanity.com/paper/2410.10819/) (2024)
- [Retrieval Head Mechanistically Explains Long-Context Factuality](https://mlsanity.com/paper/2404.15574/) (2024)
- [Memorizing Transformers](https://mlsanity.com/paper/2203.08913/) (2022)
- [Recurrent Memory Transformer](https://mlsanity.com/paper/2207.06881/) (2022)
- [SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills](https://mlsanity.com/paper/2308.16369/) (2023)
- [Data Engineering for Scaling Language Models to 128K Context](https://mlsanity.com/paper/2402.10171/) (2024)
- [On The Computational Complexity of Self-Attention](https://mlsanity.com/paper/2209.04881/) (2022)
- [LM-Infinite: Zero-Shot Extreme Length Generalization for Large Language Models](https://mlsanity.com/paper/2308.16137/) (2023)
- [What Matters in Transformers? Not All Attention is Needed](https://mlsanity.com/paper/2406.15786/) (2024)
- [ChunkKV: Semantic-Preserving KV Cache Compression for Efficient Long-Context LLM Inference](https://mlsanity.com/paper/2502.00299/) (2025)
- [InfLLM-V2: Dense-Sparse Switchable Attention for Seamless Short-to-Long Adaptation](https://mlsanity.com/paper/2509.24663/) (2025)
- [Sliding Window Attention Training for Efficient Large Language Models](https://mlsanity.com/paper/2502.18845/) (2025)
- [Core Context Aware Transformers for Long Context Language Modeling](https://mlsanity.com/paper/2412.12465/) (2024)
- [Shiftable Context: Addressing Training-Inference Context Mismatch in Simultaneous Speech Translation](https://mlsanity.com/paper/2307.01377/) (2023)

## Built on by (most cited first, 1 of 1)

- [Distance-KV: Exploiting Relative Distance for Efficient Long-Context Inference](https://mlsanity.com/paper/2609.32663/) (2026)

## Abstract (arXiv)

Transformer-based large language models face severe scalability challenges in long-context generation due to the computational and memory costs of full-context attention. Under practical computation and memory constraints, many inference-efficient long-context methods improve efficiency by adopting bounded-context or segment-level execution only during inference, while continuing to train models under full-context attention, resulting in a mismatch between training and inference execution and state-transition semantics. Based on this insight, we propose a training-inference consistent segment-level generation framework, in which training and inference follow the same segment-level forward execution semantics. During training, consistency with inference is enforced by restricting gradient propagation to KV states carried over from the immediately preceding segment, while permitting head-specific access to past KV states during the forward pass without involving them in gradient propagation. Across long-context benchmarks, our approach achieves performance comparable to full-context attention, while achieving competitive latency-memory trade-offs against strong inference-efficient baselines, and substantially improving scalability at very long context lengths (e.g., approximately 6x lower peak prefill memory at 128K compared to full-context attention with FlashAttention).