Training–Inference Consistent Segmented Execution for Long-Context LLMs

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-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 $6\times$ lower peak prefill memory at 128K compared to full-context attention with FlashAttention).

Large language models are now used to work with long documents, conversations, and reports. But the longer the input becomes, the more memory and computation the model needs, which can make long-context systems costly and hard to deploy. Many current efficiency methods change how a model reads long text only when the model is being used. During training, however, the model may have learned under different conditions. This mismatch can make its behavior less reliable on long inputs. We address this problem by training and using the model in the same piece-by-piece manner: the model reads a long input in segments, carries recent information forward, and looks back at older information only when needed. In our experiments, this approach keeps strong long-context performance while using less memory and achieving a better balance between speed and memory use. This could make long-context AI systems more practical for tasks such as document understanding, long conversations, and long-form summarization.