Learning When to Attend: Conditional Memory Access for Long-Context LLMs

Language models struggle to generalize beyond pretraining context lengths, limiting long-horizon reasoning and retrieval. Continued long-context training is effective but expensive due to the quadratic cost of Attention. We observe that most tokens do not require (Global) Attention over the entire sequence and can rely on local context. Based on this, we propose L2A (Learning To Attend), a layer that enables token-wise conditional long-range memory access by learning \textit{when} to invoke Global Attention. Applied to Qwen2.5 and Qwen3 models, L2A extends the effective context length from 32K to 128K tokens, matching standard long-context training within 3\% while skipping Global Attention for $\sim$80\% of tokens and outperforming prior baselines. We further design custom Triton kernels to efficiently implement L2A on GPUs, achieving $\sim$2× improvements in training throughput and time-to-first-token over FlashAttention. Additionally, L2A enables post-training pruning of sparse Global Attention layers, reducing KV cache memory by up to 50\% with negligible performance loss.