Language Models Can Control Their Own Attention

Namgyu Ho, Huzama Ahmad, Woosung Koh, Se-Young Yun, Tal Schuster, Cicero Nogueira dos Santos

Declarative Attention lets LLMs explicitly declare their attention scope in chain-of-thought, cutting KV cache reads.

Can we reduce the computational cost of long-context LLM inference by allowing the model to dynamically choose which parts of the context to attend to?

Language models spend most of their decoding time reading the entire Key-Value (KV) cache, even though they only need a small fraction of the context to generate each token. The authors introduce Declarative Attention (DA), a protocol where the model explicitly declares its attention scope—global, focus, or local—within its chain-of-thought. An inference-time state machine parses these declarations to dynamically mask the KV cache, skipping irrelevant blocks. On long-context tasks, DA reduces attended tokens by up to 52% with only marginal accuracy drops, effectively trading extra decode steps for significantly lower per-step memory bandwidth usage.

Paper Primer

DA partitions generation into three modes: <global> for surveying the full context, <focus> for reasoning over specific addressable segments, and <local> for self-contained synthesis. The model emits these tags as parseable tokens, and the inference engine uses them to rewrite the KV-cache block table, ensuring only necessary blocks are read by the attention kernel.

DA significantly reduces KV cache memory access costs.

Across 15 long-context tasks, DA reduced average attended tokens by 52.0% on Gemma-4-31B and 31.1% on Qwen-3.6-27B. Up to 21 million tokens saved per response in the longest-context tasks.

The efficiency gain is driven by the dynamic mask, not the prompt format.

Comparing DA to a maskless baseline (DA-nm) shows that the mask cuts attended tokens by 71.1% on Gemma and 46.5% on Qwen relative to the maskless version. The mask converts a 66% token overhead (from extra decode steps) into a net saving.

Why is this approach better than existing sparse attention methods?

Existing methods approximate attention masks via lightweight scans over hidden activations, which still incur O(N) cost per step. DA eliminates this selection cost entirely by deriving the mask directly from the model's explicit reasoning trace.

What is the primary trade-off when using Declarative Attention?

DA increases the total number of decode steps by roughly 15–35% because the model must generate the protocol tags. The net efficiency gain depends on whether the reduction in per-step memory bandwidth outweighs the cost of these additional steps.

DA demonstrates that models can perform "system-2" attention management, turning selective attention into an auditable, steerable reasoning process that scales with model capability.

The Cost of Long Context

Standard LLMs read the entire KV cache each step, incurring an O(N) bottleneck.

Transformers compute attention over every preceding token at each decoding step, so the cost grows linearly with context length. In long‑context regimes the latency of reading the Key‑Value (KV) cache dominates decoding time, e.g., a 1 M‑token conversation forces a 15 GB KV read per step on Qwen‑3.5‑397B‑A17B.

In the baseline setting the model attends to the full KV cache at every generation step, never skipping tokens.

Instead of silently scanning the whole KV cache, the model declares which region it will attend to, allowing the inference engine to mask out irrelevant tokens.

The O(N) bottleneck of KV cache reading is the primary target of Declarative Attention.

How Declarative Attention Works

Declarative Attention lets the model select context segments via tags, cutting KV reads.

Even though attention weights concentrate on a few tokens, vanilla LLMs still read the full KV cache at every step, wasting compute.

How does Declarative Attention differ from simply prompting the model to “look at” a segment?

DA turns the model’s textual tag into a concrete block‑level mask that the inference engine enforces, whereas ordinary prompting relies on the model’s internal attention distribution, which still reads every KV block.

Model emits

Model extracts “2011” from Chunk 2 and stores it in its response.

Model computes 2011 − 2003 = 8 and outputs the answer.

DA makes the attention scope explicit, so the engine can skip reading irrelevant blocks, turning a textual plan into a concrete compute saving.

Magic chunks are fixed‑size, addressable slices of the long input that the model can refer to by number — like chapters in a book that can be opened on demand.

Why not let the model attend to arbitrary token ranges instead of predefined magic chunks?

Arbitrary ranges would require masking at token granularity, which vLLM’s block‑wise kernels cannot skip efficiently; magic chunks align with block boundaries, enabling real compute savings.

The three modes dictate how much of the context the model can see while reasoning.

Do the three modes correspond to existing attention mechanisms like sliding‑window or global attention?

They are a

Efficiency and Deployment Trade-offs

DA saves compute in large‑batch, long‑context decoding by cutting attention work.

Large‑batch decoding spreads the FFN parameter load across many tokens, making the FFN compute‑bound, while attention remains memory‑bound and grows with context length. The resulting imbalance means that cutting attention work can outweigh the extra decode steps.

The model predicts per‑step latency by dividing the work of each component by the hardware ceiling that component can actually achieve.

Compute FFN wall‑time: $T_{\text{FFN}} = \frac{2\cdot10^{9}}{100\cdot10^{12}\times0.1}=2\cdot10^{-4}\,$s ≈ 0.2 ms.

Compute attention wall‑time: $T_{\text{attn}} = \frac{200\cdot2^{20}}{1\cdot2^{40}\times0.5}= \frac{200}{1024\times0.5}\,$ms ≈ 0.39 ms.

Attention already consumes roughly twice the latency of the FFN.

If DA skips reading half the KV cache (100 MiB), $T_{\text{attn}}$ halves to ≈ 0.20 ms, making the total per‑step latency comparable to the FFN alone.

When context is large, the memory‑bound attention term dominates; cutting KV reads yields a net speedup even after adding extra decode steps.

How does this Roofline model differ from the classic Roofline model used for training performance?

The classic model treats the whole workload as a single operation and plots performance versus operational intensity. Here we split the inference workload into two distinct components—compute‑bound FFN and memory‑bound attention—each evaluated against its own hardware ceiling (MFU vs MBU). This separation lets us see how a technique like DA reshapes the memory‑bound term without affecting the compute‑bound term.

In large‑batch settings the FFN cost is amortized across many tokens, but the KV read cost stays per token. DA cuts that per‑token memory traffic, so the relative gain grows with batch size, making it especially attractive when the hardware is already saturated on bandwidth.

Evaluation Setup

Experimental settings detail models, datasets, baselines, metrics, and inference configuration.

We evaluate Declarative Attention (DA) on six large language models from two families: Gemma‑4‑31B, Gemma‑4‑12B, Gemma‑4‑E4B, Qwen‑3.6‑27B, Qwen‑3.5‑9B, and Qwen‑3.5‑4B. All models natively handle up to 256 K input tokens, except Gemma‑4‑E4B, which is limited to 128 K.

The benchmark comprises fifteen long‑context sources organized into two task categories: (1) single‑span retrieval and reasoning, and (2) multi‑span reasoning. Sources are drawn from RULER, LongBench v1/v2, LooGLE, and ZeroSCROLLS, abbreviated LBv1, LBv2, and ZS.

Four of the sources use synthetic QA generated with Gemini‑3‑Flash to improve annotation quality; the remaining eleven use the original QA samples. For each source we sample up to 128 examples with a fixed random seed, applying the same example set across all three methods (Vanilla, DA‑no‑mask, DA).

We compare three inference arms: Vanilla (full causal attention), DA‑no‑mask (DAnm, full DA prompt without custom masking), and Declarative Attention (DA) with its bespoke attention mask. All arms share the same final‑answer specification and convention, ensuring identical scoring conditions.

Performance is measured by two headline metrics: accuracy (fraction of responses judged correct) and attended $KV$ token reduction (percent decrease relative to Vanilla). Evaluations use an LLM judge (Qwen‑3.5‑4B) calibrated against Gemini‑3‑Flash rubrics, achieving Pearson r = 0.99 with Gemini‑3.1‑Pro.

Inference runs with thinking mode disabled, model‑specific sampling parameters (see Appendix D.3), and a maximum generation length of 8 K tokens. All models are served on NVIDIA B200 GPUs via vLLM with our custom DA integration (Appendix B).

Main Performance Results

Results quantify DA’s modest accuracy loss and large attention‑token savings across models.

Standard LLMs read the entire KV cache at every decode step; Declarative Attention (DA) lets the model skip irrelevant tokens, cutting compute while keeping accuracy near vanilla.

Declarative Attention reduces estimated decode wall‑clock time to 0.71× on Gemma‑4‑31B and 0.77× on Qwen‑3.6‑27B compared with vanilla inference.

Estimated via the roofline model (Section 3) on a single B200 accelerator in bf16, accounting for global‑memory reads, matmul costs, and local‑memory reads.

Across the 15 long‑context tasks, DA’s average accuracy drops are modest—1.27 pp on Gemma and 2.75 pp on Qwen—while it matches or exceeds vanilla on several tasks, notably gaining +3.1 pp on `longdep_qa` (Gemma) and +5.6 pp on `code_repo` (Qwen).

**Table 2.** DA reduces attention cost with a 1–3pp average accuracy drop. Across 15 long-context tasks, DA reduces total average attended tokens by 52.0% on Gemma-4-31B and 31.1% on Qwen-3.6-27B, with average accuracy drops of 1.27pp and 2.75pp respectively. DA^nm represents the DA ablation without custom attention masking. Averages are shown in bold. Attended-token ranks are highlighted by background color: best, middle, and worst. Columns are grouped as Accuracy (%) and Attended Tokens (M / sample); within each metric the three arms are Vanilla / DA^nm / DA, reported for Gemma-4-31B (G) and Qwen-3.6-27B (Q). Attended tokens are reported as logged; a small fraction of DA responses (under 1% on both models here, but up to ~6% on Gemma-4-12B in Figure 3) fail to terminate within the 8K generation budget and inflate the attended-token sum, a decode-length effect rather than a property of the mask (Section 5.2). Figure 3b marks each model's value with these responses excluded.

**Figure 2.** DA's efficiency comes from attention masking rather than shorter generation. Accuracy, decode steps, and attended tokens for Vanilla, DA, and DA^nm on the two headline models, each normalized to that model's vanilla. Even though DA is elicited zero-shot and generates more tokens than vanilla, it attends far fewer, so its efficiency comes entirely from the attention mask. (a) Accuracy is nearly matched across all three arms (a 1-3pp drop for DA). (b) DA and DA^nm both run more decode steps than vanilla (about 15-35% more), a byproduct of the zero-shot protocol. (c) Despite the extra steps, DA attends far fewer tokens than vanilla because of per-step masking, whereas DA^nm attends more, since its extra steps run at full attention.

**Figure 4.** DA holds accuracy near vanilla while its absolute token saving grows with context length. All 15 sources are pooled and binned by context length on Gemma-4-31B. (a) Relative accuracy (normalized to vanilla): DA stays within ~1pp of vanilla through 32K and then declines modestly at longer contexts, a decline absent from the maskless DA-nm line. (b) Absolute attended-token delta versus vanilla (millions per response): DA's saving grows from about -1M at short context to about -21M in the longest bin, while DA-nm's overhead grows positive. The constant-fraction (relative) attention view and the Qwen-3.6-27B counterpart are in Appendix E.

Overall, the results confirm that the attention mask is the primary source of DA’s efficiency, and that the approach scales favorably with longer contexts without sacrificing much accuracy.

Efficiency and Mode Usage

We quantify how each DA mode and protocol adherence affect efficiency and accuracy.

We now dissect how each Declarative Attention (DA) mode contributes to overall efficiency and how faithfully models follow the DA protocol.

Removing the mode eliminates up to 99 % of per‑token attention savings.

Local attends only ~6 % of tokens, delivering the largest savings.

Removing the mode cuts per‑token savings to roughly 76 %.

Focus attends ~12 % of tokens, providing substantial but smaller savings than .

**Figure 5.** DA spends most tokens in cheap modes, with larger savings at longer contexts. Results on Gemma-4-31B. (a) Token share by mode across context-length buckets; &lt;focus&gt; and &lt;local&gt; account for about 73% of generated tokens, and &lt;global&gt;'s share rises at the longest contexts. (b) Per-token attention saved relative to vanilla: &lt;global&gt; saves 0% by construction, while &lt;focus&gt; and &lt;local&gt; save 76–99% per token, with savings growing as context lengthens. The Qwen-3.6-27B counterpart is in Appendix E.

**Figure 3.** DA accuracy scales favorably with backbone size. Metrics are normalized to each model's vanilla baseline and the x-axis is total parameters. (a) Relative accuracy rises toward parity as the backbone grows, in both families. (b) Relative attended tokens: DA attends roughly half of vanilla on five of six models. Black lines mark each bar's value when non-terminating responses (those exhausting the 8K generation budget) are excluded from both arms. Gemma-4-12B sits above vanilla only because about 6% of its DA responses fail to terminate, a decode-length artifact rather than a property of the mask: excluding them brings it just below vanilla (98%). Qwen-3.5-4B shows a milder version of the same effect (50% to 38%), and non-termination occurs in its vanilla arm as well, a small-model trait rather than a DA effect.

Limitations and Future Directions

LLMs read the KV cache each step, and Declarative Attention picks needed segments, cutting compute with accuracy loss.

DA assumes a model can reason over long context while reading only a small part at each step. Our zero‑shot experiments show that off‑the‑shelf models already follow this pattern, achieving up to 99 % reduction in per‑step attention reads beyond 128 K tokens.

On the models we evaluate, vanilla global attention dominates decode cost (73 %–86 % of roofline‑estimated time). This raises the question whether the same dominance holds for newer architectures that already shrink KV reads.

Even when retrieval supplies new material, it does not replace DA: retrieval decides what enters the context, while DA decides which cached tokens to attend to. As tool results accumulate, they remain in context for many steps, making DA’s selective reading especially valuable.

Current zero‑shot prompting leaves several limitations: DA incurs ~⅓ more decode steps, relies on artificial “Magic Chunks”, and operates only in non‑thinking mode. Future work could close these gaps via fine‑tuning, more natural segment boundaries, and integration with thinking traces.

Lightweight scan‑based sparse attention can reduce the cost of the remaining global steps, and speculative decoding can offset DA’s extra steps. Because DA’s mask is applied by rewriting the KV‑cache table, it composes cleanly with both techniques.

Viewing DA as system‑2 sparse attention highlights its interpretability: the model’s textual declarations directly expose which tokens drive KV reads. This opens avenues such as reinforcement‑learning rewards for both accuracy and attention efficiency, and KV‑cache offloading that prefetches announced segments.

Supplementary Experimental Details

Supplementary details, related work, and implementation notes are collected here.

Appendix A surveys prior work on dynamic sparse attention, KV‑cache eviction, complementary cache axes, model‑controlled inference, and self‑declared attention.

Appendix B details how Declarative Attention (DA) is integrated into vLLM, describing the mask‑generation hooks and per‑kernel patches.

Appendix C provides background on inference costs, including the roofline model, with subsections on prefill vs decode, cost components, large‑batch behavior, wall‑time definition, utilization values, operating regimes, and the resulting attention‑dominance insight.

Appendix D enumerates experimental details such as evaluation datasets, validation of the local judge, sampling parameters, and discusses current limitations and future potential.

Appendix E presents additional experimental results not covered in the main text.

Appendix F describes the prompt construction used throughout the experiments.

Inference Cost Analysis

Appendix details the decode‑time breakdown and KV‑memory characteristics across models.

The decode step cost is fixed by the FLOPs budget once weight loads are amortized, so at large batch and context the attention term dominates the FFN term.

For a single decode step on Qwen‑3.5‑397B‑A17B at 1 M context, the 15 global attention layers hold roughly 15.4 GB of KV cache, while the 45 Gated DeltaNet layers add only a fixed, context‑independent state.

Using a B200‑class accelerator (bandwidth $8\,\text{TB/s}$, compute $4.5\,\text{PFLOPS}$ dense FP8), the attention roofline wall‑time is about $15.4\,\text{GB}/(8\,\text{TB/s}\times0.70)\approx2.7\,\text{ms}$, whereas the FFN roofline wall‑time is $34\,\text{GFLOPs}/(4.5\,\text{PFLOPS}\times0.40)\approx0.019\,\text{ms}$.

**Table 1.** Language Models Can Control Their Own Attention

Under expert parallelism the per‑expert GEMM aggregates enough tokens to become compute‑bound, but its skinny GEMMs and routing keep utilization low, so attention dominates by roughly two orders of magnitude (≈ 145× in the example).

**Table 5.** Trace inputs (macro-averaged over the 15 tasks). Decode-step count $D$ and attended-token total $A$, measured directly from the generated responses of each arm.

Three features emerge: (1) the global read accounts for 73 % of Gemma’s decode time and 86 % of Qwen’s, making it the primary lever for DA; (2) DA increases $D$ (by 31 % for Qwen, 35 % for Gemma) but still cuts total time to 0.71× (Gemma) and 0.77× (Qwen); (3) the local‑term floor limits the net gain—only 5 % of Qwen’s attention time is lost to the local term, versus 42 % for Gemma.

This decomposition is a ceiling estimate: it omits prefill, ignores tiny per‑step ops (normalization, rotary embeddings, residual adds), and assumes token‑granular attended‑token counts, so the reported savings are optimistic.

**Table 6.** KV bytes per context token, as of August 2026. Stored bytes are what a context token occupies across the global attention layers, and the O(N) read is what one decode step must read per context token. The two coincide under full attention, at 13.8 to 254.0 KB per token as served, and separate by a factor of 6.6 to 33 under indexer-based sparse attention, where the step scans only the indexer keys at O(N) and reads a fixed top-k set besides, at 0.35 to 2.77 KB per token. MiniMax-M3 is the exception at 20.7 KB, because three of its 60 layers keep full attention and read their whole cache every step. The bf16 columns make the architectures comparable at one precision, and the as-served columns charge the dtype each vendor's recipe specifies. They agree wherever a width is fixed by the architecture, which is why every indexer read is unchanged between them. "Forced" marks a layout the model's tech report fixes, for which no uniform-bf16 figure is published. `L_global` counts the layers whose per-step read grows with context, which excludes sliding-window and linear-attention layers. †Not part of the roofline set of Table 7.

The table presents performance metrics for various large language models, including parameters ($P$), KV cache size ($b_{kv}$), local cache size ($s_{local}$), global attention time ($T_{global}$), matrix multiplication time ($T_{matmul}$), local attention time ($T_{local}$), and the attention share percentage.

The attention share is computed as $T_{\text{global}}/(T_{\text{global}}+T_{\text{matmul}}+T_{\text{local}})$, where $T_{\text{global}}=b_{\text{kv}}N/B_{\text{Weff}}$, $T_{\text{matmul}}=2P/C$, and $T_{\text{local}}=s_{\text{local}}/B_{\text{Weff}}$, with $B_{\text{Weff}}=8\,\text{TB/s}\times0.70=5.6\,\text{TB/s}$ and $C=2.25\,\text{PFLOPS}\times0.40=0.9\,\text{PFLOPS}$ (BF16) or $C=4.5\,\text{PFLOPS}\times0.40=1.8\,\text{PFLOPS}$ (FP8).

**Table 8.** Attention share against context length and against the utilization assumptions. The left block holds MFU and MBU fixed and varies N, and the right block holds N at 1M and takes the two corners of the reported utilization ranges of Appendix C.5, MFU 40 to 60% and MBU 60 to 85%. Context length moves the share far more than the utilization choice does, which is why the headline states 1M explicitly. Across the whole utilization box the 1M shares span 54.4 to 99.59%, so the conclusion does not rest on the particular values we adopt. See Table 7.

Sweeping MFU from 40 % to 60 % and MBU from 60 % to 85 % moves the 1 M‑token attention share of every full‑attention model by less than three percentage points, confirming that utilization choices have modest impact compared to context length.

Raw Performance Data

Additional experimental details, dataset statistics, and supplemental figures.

This appendix gathers the auxiliary material referenced throughout the paper: detailed dataset statistics, validation of the local judge, sampling settings, limitations, and the full prompt specifications used for Declarative Attention (DA). It also reproduces the performance table and two key figures that illustrate DA’s token‑saving behavior.

**Figure 9.** Qwen matches Gemma's per-token savings but shifts more into global mode at long context. Counterpart to Figure 5, on Qwen-3.6-27B. Per-token savings match Gemma's (focus and local save 74–99% per token on Qwen, against 76–99% on Gemma), but Qwen spends an increasing share of tokens in the unmasked global mode as context grows, reaching about 55% in the longest bucket versus roughly 45% on Gemma, which is why its total token savings are smaller at long context.

**Figure 10.** DA's per-step masking saves a roughly constant fraction of tokens across context length. The same Gemma-4-31B pooling as Figure 4b, normalized to vanilla rather than shown as an absolute delta. DA attends about 50-64% of vanilla in every context bin, while DA-nm stays above vanilla. Together with the absolute view in Figure 4b, this shows that DA's saving is a stable fraction of a growing total, so the absolute token saving scales up with context length.

**Table.** Context source, domain, task, QA source, and context length statistics for the evaluated datasets.

Section D.1 describes the 15 datasets used for long‑context QA, noting that 11 retain their original QA annotations while four are augmented with synthetic QA generated by Gemini‑3‑Flash to cover additional context types.

Section D.2 validates the Qwen‑3.5‑4B local judge by re‑evaluating a stratified sample of 2 993 responses with Gemini‑3.1‑Pro; the two judges agree on 98.53 % of verdicts ($\kappa$ = 0.940) and exhibit a Pearson correlation of 0.992 on per‑cell accuracies.

Section D.3 lists the sampling hyper‑parameters for each model family: Qwen‑3.5/6 uses temperature 0.7, top‑p 0.80, top‑k 20, presence penalty 1.5; Gemma‑4 employs temperature 1.0, top‑p 0.95, top‑k 64.

Section D.4 enumerates current limitations: DA relies solely on prompting (no fine‑tuning), which fails on sources that forbid intermediate reasoning or require global counts; six datasets (Table 10) exhibit large accuracy drops (‑30 pp to ‑21 pp) despite per‑step attention reductions of 39–67 %.

Section E reports additional experimental results, referencing Figure 4b (absolute DA token savings) and Figure 8/9 (Qwen counterparts), while Figure 10 re‑expresses the same data as a ratio to vanilla.

Section F details the prompt‑construction pipeline: a tokenizer‑aware semantic segmenter splits context into 2 048‑token segments (hard cap 2 560 tokens) using a hierarchy of delimiters, guaranteeing lossless reconstruction and handling over‑cap atomic strings.

The DA Instruction Prompt defines three reasoning modes (, , ) for extracting values from “magic chunks” and mandates at least one block before answering, ensuring the model does not hallucinate unseen content.

The Vanilla Instruction Prompt provides the baseline prompting format, requiring the model to answer solely from the supplied context without any magic‑chunk machinery.

The Rubric Generation Prompt specifies how to create a binary evaluation rubric that lists the exact correct answer and a concise set of decision rules for external judges.

The Judge Prompt instructs the evaluation LLM to apply the provided rubric to a target model’s response, ignoring superficial differences such as case or punctuation.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers