One-sentence Summary

Researchers from KAIST AI and Google DeepMind propose Declarative Attention (DA), an intrinsic protocol that lets language models declare where to attend within chain-of-thought via , , and modes, enabling inference engines to skip most KV cache reads, and in zero-shot evaluation across 15 long-context tasks, DA reduces attended tokens by 52.0% for Gemma-4-31B and 31.1% for Qwen-3.6-27B with accuracy drops of 1.27pp and 2.75pp, respectively.

Key Contributions

• Declarative Attention is a zero-shot protocol that elicits language models to declare their attention targets (global, focus, local) within their chain-of-thought, enabling the inference engine to skip most KV cache reads without extrinsic scoring.
• Across 15 long-context tasks, off-the-shelf models using Declarative Attention reduce total attended tokens during decoding by 52.0% (Gemma-4-31B) and 31.1% (Qwen-3.6-27B) with only 1.27pp and 2.75pp accuracy drops, and the accuracy gap narrows as model scale increases.
• An efficient vLLM integration with block-aligned, in-place KV cache masking compatible with FlashAttention is provided, and roofline-based wall-time analysis projects decode cost reductions to 0.71× and 0.77× of vanilla on the respective models.

Introduction

Transformer inference for long contexts is dominated by the memory bandwidth cost of loading the key-value (KV) cache at every decoding step, making full attention prohibitively expensive. Prior work attempts to sparsify attention by predicting which tokens are relevant, but static heuristics fail to adapt to query-specific needs, while dynamic methods still incur an O(N) per-step scanning cost. The authors introduce Declarative Attention, a zero-shot prompting protocol that elicits the model to explicitly state where it will attend as part of its chain-of-thought. The inference engine parses these declared mode transitions (global, focus, local) to construct the attention mask directly from generated text, eliminating the per-step selection overhead and enabling efficient long-context decoding with only marginal accuracy loss.

Method

The authors propose Declarative Attention (DA), a protocol that elicits large language models to restructure their chain-of-thought reasoning, making their attention plans explicit and legible. Instead of relying on auxiliary scorers or predicting attention weights, DA requires the model to organize its reasoning into contiguous spans where the attention scope remains stable and to declare that scope using a predefined tag syntax.

To facilitate this, the authors design a specific prompt structure consisting of a persistent scaffold and a variable context region. The scaffold includes a system instruction, the user's question, and detailed mode instructions, ensuring persistent grounding for the protocol. The long input context is divided into addressable segments of approximately 2048 tokens, referred to as "magic chunks". These segments are presented to the model within a simulated tool-use transcript, where an assistant appears to call a retrieval tool for each chunk. This formatting aligns segment boundaries with the special tokens delimiting user, assistant, and tool messages, which the model tracks fluently due to its post-training.

The model freely alternates between three distinct reasoning modes, each serving a specific purpose and dictating how much of the context remains visible. As shown in the figure below:

In <global> mode, the model attends to all context segments, which is utilized for navigating and surveying the full context to locate relevant information. In <focus> mode, the model attends only to the specific context segments named in the tag, enabling targeted reasoning over a specific region without the computational cost of processing the entire context. Finally, in <local> mode, the model attends to none of the context segments, relying solely on the information already accumulated in its response for self-contained reasoning and answer synthesis. In all modes, the model continuously attends to the scaffold and its own generated tokens.

To realize the corresponding attention masks, the authors introduce a DA state machine that operates alongside the inference engine. The state machine starts in the default <global> mode and parses the output stream to detect mode transitions based on the closing characters of the opening tags. Because modern inference engines like vLLM store the Key-Value cache in fixed-size blocks, the state machine applies the mask at block granularity. It rounds the kept token spans outward to block boundaries to ensure no declared tokens are dropped, allowing existing kernels such as FlashAttention to run unchanged. The system integrates with vLLM through hooks on the attention metadata builder, rewriting the request's KV-cache block table at each decode step so that only the kept blocks remain visible to the attention kernel. This design significantly reduces the KV blocks read per step, trading a higher number of decode steps for a lower per-step attention cost.

Experiment

the paper evaluate Differential Attention (DA) on 15 long-context retrieval and reasoning tasks using six models from the Gemma and Qwen families, comparing against vanilla full-attention and a maskless variant. DA achieves substantial token savings (up to 52% fewer attended tokens) with modest accuracy drops, and the attention mask is the primary source of both the savings and the accuracy cost. Accuracy relative to vanilla improves with model scale, and token savings increase with context length, translating to theoretical decode-time speedups of 1.3-1.4x on optimized hardware. The efficiency stems from cheap focus and local modes, and protocol adherence is robust for larger models.

The benchmark suite covers 15 long-context sources ranging from short documents to million-token code repositories, with single-span retrieval tasks averaging 9K–11.5K tokens. In differential attention experiments, the cheap focus and local modes dominate generation and their savings grow with context length, while global attention accounts for a minority of tokens. Protocol adherence is near-perfect for large models, with focus success rates reaching 99%, but smaller models struggle, linking adherence failures to accuracy drops. Single-span retrieval tasks in the benchmark have mean context lengths from 9K to 11.5K tokens, with standard deviations up to 7.1K. On Gemma-4-31B, focus and local attention modes together generate 73% of tokens and achieve 76–99% per-token attention savings relative to vanilla steps. Focus success rate rises sharply with model scale, from 58% on Gemma-4-E4B to 99% on the largest models, making protocol adherence a bottleneck only at small scales.

Across long-context retrieval and reasoning tasks, the DA protocol reduces total attended tokens by 52% on Gemma-4-31B and 31% on Qwen-3.6-27B, with average accuracy drops of only 1.3 and 2.8 percentage points. The custom attention masking is critical: without it, token reductions shrink and attended tokens can even increase over vanilla. The trade-off is slightly more favorable on Gemma, where larger token savings come with a smaller accuracy penalty. DA consistently reduces attended tokens on both models, with Gemma achieving roughly a halving across all tasks. Removing custom masking (DA^nm) erodes the savings and, on Qwen, leads to higher attended tokens than vanilla for some single-span tasks. Accuracy under DA remains high, with only isolated drops of 1–4 percentage points on a few tasks, and no loss on others. Qwen shows a larger average accuracy decline (2.75 pp) and a smaller token reduction (31%) compared to Gemma (1.27 pp, 52%).

The DA protocol reduces estimated decode wall-clock time to 71% of vanilla for Gemma-4-31B and 77% for Qwen-3.6-27B on a single B200 accelerator. Savings come entirely from cutting the global-attention KV read, the dominant cost in vanilla, while matmul and local-memory costs rise slightly because DA generates more decode steps. The benefit is larger on Qwen where the local memory (GDN state) is tiny, whereas Gemma's large sliding-window attention local read limits the overall gain. DA slashes global-memory read time by roughly half on both models, but the total speedup is capped by fixed local-memory costs that the mask does not reduce. Under DA, Gemma's local SWA read accounts for 42% of attention time, while Qwen's GDN state is only 5%, explaining why Qwen sees a larger relative reduction in total decode time.

The evaluation uses a long-context retrieval benchmark with mean context lengths of 9K–11.5K tokens and models such as Gemma-4-31B and Qwen-3.6-27B. The differential attention protocol shifts most computation to cheap focus and local modes, yielding large reductions in attended tokens with only minor accuracy penalties, while smaller models exhibit lower protocol adherence that correlates with accuracy loss. Custom attention masking is essential for these savings, and although decode wall-clock time improves, overall speedups are limited by fixed local-memory costs that the protocol does not reduce.