Language Models Can Control Their Own Attention — part 1/6

Cleaned-text code points [0, 15634). Bibliography omitted here; source.md includes it.

\reportnumber

# Language Models Can Control Their Own Attention

###### Abstract

Abstract: Language models spend most of their attention on a small fraction of context, yet they read the entire KV cache to find the few tokens that matter. If the user asks about a previous detail in a 1M-token conversation, global attention layers must scan the full context to generate each token of the reply. A prominent approach mitigates this cost by pre-selecting relevant tokens via lightweight proxy scores, but this extrinsic scoring still incurs O ⁡ ( N ) O(N) per step. We take an intrinsic approach motivated by the simple question: wouldn’t the model already know which parts of the context are relevant? To this end, we introduce Declarative Attention (DA), a protocol that elicits the model to declare where it needs to attend within its chain-of-thought, partitioning generation into three modes: <global> (full context), <focus> (a specific region), and <local> (recent output only). The inference engine parses these declarations like tool calls and skips most of the KV cache read. Under zero-shot evaluation across 15 long-context tasks, DA on off-the-shelf models (Gemma-4-31B, Qwen-3.6-27B) significantly reduces total attended tokens during decoding (52.0%, 31.1%) with modest accuracy drops (1.27pp, 2.75pp) that shrink with model scale. DA unlocks a new axis of sparse attention, with further potential under training-based methods that future work can explore.

## 1 Introduction

Transformers compute attention over every preceding token at each decoding step, rendering them computationally expensive for long-context tasks ( Deng et al., 2024 ) . Specifically, Key-Value (KV) cache memory access latency heavily dominates decoding time in long-context regimes. For instance, in Qwen-3.5-397B-A17B ( Qwen Team, 2026b ) , roughly 15 GB of KV cache must be loaded per sequence at every decoding step for a 1M-token context; a memory bandwidth requirement comparable to loading the model’s 17B active parameters ( Dao et al., 2022 ) .

However, this exhaustive mechanism diverges fundamentally from human cognition: when comprehending a lengthy document, humans do not continuously re-read every prior word to synthesize information and answer a question. Empirical literature corroborates this intuition, demonstrating that attention scores naturally concentrate on a small subset of context tokens ( Child et al., 2019 ; Zhang et al., 2023 ; Tang et al., 2024 ) .

While sparsifying context attention emerges as a natural solution, a fundamental challenge remains: true attention scores are unknown a priori . Because these scores only become available after computing the full attention matrix, dynamically identifying which tokens to attend to is prohibitively expensive. Prior works attempt to bypass this by predicting attention-heavy tokens and applying a mask that excludes irrelevant context from computation and memory access. Early approaches rely on static heuristics such as recency or historical attention magnitude ( Zhang et al., 2023 ; Xiao et al., 2024b ) . However, these fixed rules struggle to anticipate the specific tokens required by future queries, leading to degraded long-context performance ( Li et al., 2025 ; Moschella et al., 2026 ) . Conversely, more recent methods acknowledge the query-dependent nature of attention, approximating the mask via a lightweight scan over the KV cache at each decoding step ( Tang et al., 2024 ; Yang et al., 2025a ; DeepSeek-AI, 2025b ) . While this reduces the constant factor of the cost, the complexity per step remains O ⁡ ( N ) O(N) . We survey this and the adjacent lines of work in full in Section A .

#### Contribution.

We propose an orthogonal direction: eliciting the model to explicitly declare where it will attend, step by step. Recent studies demonstrate that language models encode information about future tokens within their hidden states ( Pal et al., 2023 ; Wu et al., 2024 ) , and chain-of-thought (CoT) prompting surfaces this latent computation as interpretable text ( Wei et al., 2022 ; Korbak et al., 2025 ) . We extend this principle from dictating what to think to where to attend. Encouragingly, a prior work demonstrates that this behavior can be trained into a model, fine-tuning separately per task on hand-designed context partitions and annotation syntax, within 2K-token contexts ( Jin et al., 2024 ) . We show that it can now be elicited through a single task-agnostic protocol: modern off-the-shelf models declare their attention zero-shot under one fixed prompt, across a wide range of long-context tasks with contexts on the order of 100K tokens. By deriving the attention mask directly from the model’s generated reasoning trace, rather than approximating it via hidden activations, our approach eliminates the selection cost entirely. O ⁡ ( N ) O(N) context reads remain only in the model’s declared global phases, not as a per-step overhead.

Concretely, we introduce Declarative Attention (DA; Section 1 ), a protocol that partitions generation into three distinct attention modes: (1) global , where the model surveys the full context, (2) focus , where it commits to a specific contextual region, and (3) local , where it attends exclusively to its own recent output without seeing the full context ( Section 2 ). Mode transitions are emitted as parseable tokens within the chain-of-thought, which the inference engine reads to dynamically construct the attention mask at each decoding step.

We demonstrate DA’s value empirically ( Section 4 , Section 5 ): 1. Zero-shot efficacy as a lower bound. DA works zero-shot on off-the-shelf models without parameter updates. Consequently, our results represent a lower bound, with significant headroom expected if models are post-trained for the protocol itself ( Section 8 ).

2. Favorable cost-accuracy trade-offs. Across 15 long-context tasks, DA reduces average decoding attention cost by 52.0% on Gemma-4-31B and 31.1% on Qwen-3.6-27B , incurring only marginal accuracy drops (1.27pp and 2.75pp, respectively).

3. Positive scaling and cost savings. DA benefits directly from model capability, with the accuracy gap steadily closing as scale increases from 4B to 31B. Furthermore, absolute token savings grow sharply as context lengthens (saving up to 21M tokens per response), and ablations confirm the dynamic mask itself (cutting attended tokens by up to 71.1% relative to the maskless ablation) drives the savings, not the prompting format.

4. Efficient vLLM implementation. We integrate DA into vLLM ( Kwon et al., 2023 ) with block-aligned, in-place KV cache masking compatible with FlashAttention ( Dao et al., 2022 ) . A roofline-based wall-time analysis projects that DA’s attention savings would reduce decode wall-clock cost to 0.71 × \times of vanilla on Gemma-4-31B and 0.77 × \times on Qwen-3.6-27B on a well-optimized serving stack ( Section 5.4 ).

## 2 Declarative Attention

In a vanilla LLM forward pass, the bulk of attention weight is concentrated on a small fraction of the context, and this sparsity pattern shifts step to step in ways that cannot be identified ahead of time ( Child et al., 2019 ; Zhang et al., 2023 ; Tang et al., 2024 ) . Every decode step must therefore read the entire KV cache from HBM, even when most of it has negligible influence on the output.

The DA protocol elicits the model to restructure its chain-of-thought to make its attention plan legible. It requires the model to (i) organize its reasoning into contiguous spans where the attention scope stays stable, and (ii) declare that scope using a predefined tag syntax. To realize the corresponding mask, we introduce a DA state machine that runs alongside the inference engine: it watches for tag transitions in the model’s output and updates the attention mask at each decode step. DA defines three modes, <global> , <focus> , and <local> , each serving a distinct reasoning purpose.

A response to the question “How long after its founding did Acme Corp go public?” could look like:

<global> I need the founding year and the IPO year. The company history in Magic Chunk 2 should state the founding. </global> <focus magic_chunks="2"> "Acme Corp was founded in 2003 in San Jose." </focus> <global> The IPO year is still missing. Magic Chunk 7 covers Acme’s financial milestones. </global> <focus magic_chunks="7"> "Acme Corp went public on the NYSE in 2011." </focus> <local> 2011 - 2003 = 8 years. </local> <answer> 8 years </answer>

The magic chunks named in the focus declarations are the addressable segments of the long input, which we prepare by splitting the context into ∼ {\sim} 2K-token units ( Section 2.1 ). In every mode, the model keeps attending to the question, the instruction, and its own response so far. The modes differ only in how much of the context , the segmented long input region of the prompt ( Section 2.1 ), stays visible: • <global> attends to all context segments. It is the mode for navigation, surveying the full context to locate relevant information. Our zero-shot prompt instructs the model to use it to identify the next segment to focus on, briefly noting why that segment is relevant.

• <focus> attends only to the context segments named in the tag. It is the mode for reasoning over a specific contextual region without the cost of attending to the rest. Our prompt instructs the model to extract the needed values verbatim from the named segments.

• <local> attends to none of the context segments. It is the mode for self-contained reasoning over information already in the response. Our prompt instructs the model to plan over the question and to synthesize the final answer from previously extracted values.

DA itself does not restrict how, when, or how many times each mode is used. Because we elicit DA zero-shot from off-the-shelf models, our prompt provides per-mode guidance to scaffold their DA reasoning ( Section F ).

The remainder of this section describes (1) the structure of the prompt, (2) how the context is delivered as addressable segments for focus mode, and (3) the DA state machine’s decode-time interventions on the inference engine.

### 2.1 Prompt overview

We design the prompt around two regions: a scaffold that is always visible to the model, and a context that contains the long input the model must reason over. The scaffold remains attended in every mode and provides the model with persistent grounding for the DA protocol. The context is the variable region whose visibility the DA state machine controls, and it holds the vast majority of the prompt’s tokens in long-context tasks. Section 1 shows the prompt structure, consisting of three scaffold sections (marked in blue ) surrounding the context:

• System instruction : a short fixed preamble whose content fills the attention sink so context never enters it ( Xiao et al., 2024b ) .

• Context : the long input the model reasons over, delivered as addressable segments in a simulated tool-use transcript ( Section 2.2 ).

• Question : the user’s query.

• Instruction : defines the three modes and provides usage guidance derived from failure modes observed during prompt development. The full DA Instruction Prompt and its vanilla counterpart Vanilla Instruction Prompt are shown in Section F .

### 2.2 Context delivery

To use focus mode, the model must name the context regions it intends to attend to. We therefore divide the context into addressable segments . Addressing a segment requires the model to track where the segment begins and ends. The text LLMs train on is pervasively structured: pre-training corpora carry document and paragraph boundaries, and post-training dialogues carry user, assistant, and tool turn boundaries ( Schick et al., 2023 ; Qwen Team, 2026b ; Gemma Team et al., 2026 ) . We expect models to track these familiar boundaries far more reliably than arbitrary novel delimiters, as we explain below. We therefore align both parts of context delivery with boundaries seen during training: (1) the segmentation that decides where segments begin and end, and (2) the formatting that presents each segment to the model.

#### Context segmentation.

We want segment boundaries to align with semantic content boundaries as closely as possible. For simplicity, we approximate this with heuristics. Targeting a 2048-token segment size, the segmenter splits only units that exceed the cap, cutting at the coarsest boundary available: paragraph breaks (double newlines) → \rightarrow single newlines → \rightarrow sentence ends ( . ! ? followed by whitespace) → \rightarrow clause ends ( ; : , followed by whitespace) → \rightarrow word boundaries. A segment edge therefore never falls inside a word. We detail the segmenter in Section F . This segmenter lets us evaluate DA on existing static long-context benchmarks, where the context is a single unstructured text with no markers of section or document boundaries. Many deployment scenarios instead carry naturally occurring, semantically aligned boundaries, such as user and assistant turns or tool responses containing retrieved context, which could serve as segments directly ( Section 8.2 ).

#### Segment formatting.

We present each segment to the model under the name magic chunk , marking segments as arbitrary retrieval splits so the model does not conflate them with the document’s own sections or chapters. The context region is a simulated tool-use transcript that we construct while preparing the prompt: for each segment, an assistant turn appears to call a get_magic_chunk tool that we declare through the model’s native tool-declaration format, and a tool response returns the segment text, headed Magic Chunk N . The transcript reads as if the model had already retrieved the document one magic chunk at a time, but no tool is ever executed: every segment is in place before generation begins. This format places segment boundaries on the special tokens that delimit user, assistant, and tool messages, boundaries the model saw throughout post-training and tracks fluently. We discuss how post-training could sharpen segment tracking further in Section 8 .

### 2.3 Decode-time interventions

#### Parsing and mode transitions.

The DA state machine starts in the default global mode and reads the model’s output stream to parse mode transitions. It transitions on the closing “ > ” character of an opening tag ( <focus magic_chunks="K"> or <local> ) and reverts to global on the matching closing tag. The <global> tag is prompting structure only: global is already the default state between declared spans, so it produces no transition. The tag instead helps the model understand and track the protocol’s structure, keeping the reasoning organized into contiguous spans with a declared scope.

#### Block-aligned mask construction.

vLLM stores the KV cache in small fixed-size blocks (typically 16 to 32 tokens), and its attention kernels read whole blocks. Masking therefore saves time only if it skips entire blocks: dropping scattered individual tokens would leave the memory reads unchanged. The state machine thus applies the mask at block granularity, rounding the kept token spans outward to block boundaries so that no token the model declared is ever dropped. The cost is at most one extra block at each edge of a kept span, a few dozen tokens against 2048-token segments. The result is an ordinary block list, so existing kernels such as FlashAttention ( Dao et al., 2022 ) run unchanged, following the block-sparse principle of Native Sparse Attention ( Yuan et al., 2025 ) .

#### Custom integration with vLLM.