Training-Inference Consistent Segmented Execution for Long-Context LLMs

TL;DR

This paper introduces a training-inference consistent segmentation framework for long-context LLMs, improving scalability and efficiency while maintaining performance comparable to full-context attention.

Contribution

It proposes a novel segment-level generation method that aligns training and inference, reducing memory usage and enhancing scalability for long-context models.

Findings

Achieves performance comparable to full-context attention on benchmarks.

Reduces peak prefill memory by approximately 6x at 128K context length.

Offers competitive latency-memory trade-offs against inference-efficient baselines.

Abstract

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).