On-Demand Attention: Language Models Know When to Recall

RESUMO

Reasoning and agentic workloads increasingly demand efficient long-context inference. Yet full-attention decoding reads the growing history at every step, regardless of its benefit to next prediction. We show that a pretrained model's decoding states already contain information predictive of this benefit, before the global read. Building on this finding, we introduce On-Demand Attention (ODA), a local-first decoding method that uses a lightweight recall head to selectively invoke global attention as its predicted benefit changes during generation. ODA trains only the recall head, leaving pretrained weights unchanged and keeping complete historical KV cache available for future recall. We further implement GPU-side conditional execution in vLLM, translating reduced global attention into practical decoding speedups over full attention at long context lengths. Experiments across Qwen and Gemma models, including hybrid-attention backbones, show selective recall recovers most performance lost under local attention while substantially reducing global reads. These findings support long-context inference in which pretrained models guide their own access patterns according to information they retain.