Declarative Attention Cuts Gemma-4-31B KV Reads 52% With a 1.27pp Accuracy Drop

2026-09-05

DA lets off-the-shelf models declare global/focus/local spans in CoT. Gemma-4-31B cuts attended tokens 52.0% at 1.27pp accuracy cost; the mask, not the prompt, supplies the save.

What problem this solves

At long-context decode, every step reloads the full KV cache from HBM. Attention mass sits on a small slice of tokens, but those scores exist only after the full matrix is computed, so there is no cheap way to skip ahead of time. Static rules (keep recent tokens, keep historically high-scoring ones) miss what a future query will need. Lightweight scanners that score a proxy over the cache cut the constant, and still pay O(N) per step.

Declarative Attention (DA) asks a different question: does the model already know which region matters? It has the model write that region into the chain-of-thought. The inference engine parses the tags the way it would parse a tool call, and skips KV blocks that were never named.

Method

DA is a protocol, not a new set of weights. Context is split into addressable segments of about 2048 tokens, cutting at paragraph breaks, then newlines, then sentence punctuation, never inside a word. Each segment is rendered as a simulated getmagicchunk tool response named Magic Chunk N, so the boundary sits on the same special tokens the model saw for tool messages in post-training. No tool actually runs; every segment is in the prompt before generation starts.

Generation splits into three modes. The model flips between them with tags in its chain-of-thought; the engine updates the mask on the closing > of an opening tag:

The system instruction, the question, and the DA usage note stay visible in every mode and fill the attention sink. Masking is block-aligned in vLLM, rounded outward so a declared token is never dropped; FlashAttention and similar kernels run unchanged. DA touches global-attention layers only. Gemma's sliding-window layers and Qwen's Gated DeltaNet already have per-step cost independent of context length, so the mask leaves them alone.

Two controls: Vanilla is ordinary full causal attention; DA-no-mask uses the same chunked prompt with full attention, which separates format from masking. Thinking mode is off. In pilots the models could not follow the protocol inside thinking tags.

Results

Fifteen long-context sources from RULER, LongBench v1/v2, LooGLE, and ZeroSCROLLS, from roughly 6K tokens up to million-token code repositories. Up to 128 examples per source; samples above 244K are dropped.

ModelAccuracy Vanilla → DAAttended tokens / sampleReduction
Gemma-4-31B87.01% → 85.74% (-1.27pp)13.43M → 6.45M52.0%
Qwen-3.6-27B85.31% → 82.56% (-2.75pp)22.54M → 15.52M31.1%

DA matches or beats Vanilla on 7 of 15 tasks on Gemma and 5 of 15 on Qwen. The drop is larger on multi-span reasoning than on single-span retrieval (2.28pp vs 0.78pp on Gemma). The biggest absolute saves sit on the longest tasks: 41.8M attended tokens on Gemma's coderepo, 52.0M on Qwen's.

DA-no-mask accuracy is essentially Vanilla (tied at 87.01% on Gemma, 0.69pp down on Qwen) but attended tokens rise 66.2% / 28.8% versus Vanilla, because the protocol writes about 15–35% more decode steps. Turning the mask on then cuts attended tokens 71.1% (Gemma) and 46.5% (Qwen) relative to DA-no-mask. The format is nearly free. The savings, and most of the accuracy cost, come from the mask.

By scale, relative accuracy rises with size: Gemma-4-E4B keeps 29% of Vanilla (58% focus-parse success) versus 99% at 31B; Qwen goes from 64% at 4B to 97% at 27B. Focus success at the largest models is 99%. By context length on Gemma, relative accuracy stays within about 1pp through 32K and falls to about 96% of Vanilla in the longest bin; the absolute save grows from about 1M tokens at short context to about 21M in that longest bin.

A roofline estimate on one B200 in bf16 (MFU 40%, MBU 70%) puts decode wall-clock at 0.71× of Vanilla on Gemma and 0.77× on Qwen. All of the save is the global-attention KV read; matmul and local-layer reads go up slightly because DA runs more steps. That is a theoretical ceiling, not measured latency, and it excludes prefill.

On Gemma-4-31B, about 73% of generated tokens run in focus or local, which attend to roughly 12% and 6% of a Vanilla step. The global share grows with context, to about 45% in the longest Gemma bucket and about 55% on Qwen, which is what caps the total save.

Why it matters

This is an axis orthogonal to "infer a sparse mask from activations": the selection is written in text, a person can read it, the engine can execute it, and it runs zero-shot on off-the-shelf models. For long-context serving, the global KV read already dominates decode roofline in this setup (73% on Gemma, 86% on Qwen). Agent traces may be an even better fit. Retrieval decides what enters context; DA decides, among what already entered, what to read at this step. Tool results linger for the rest of the episode, which is exactly the redundancy DA is built to skip.

What you can use today is the protocol plus a vLLM hook, with no kernel change. The numbers are a prompting floor. Post-training, exposing DA as ordinary tools, and building a segment index for global steps are all listed as future work.

Limitations

The authors are specific. The zero-shot policy is not optimal: DA runs about a third more decode steps than Vanilla, and the mode mix does not always match the task. Benchmarks force manufactured chunks; real dialogues and tool traces already have addressable boundaries, and this paper does not measure that distribution. Thinking mode cannot follow the protocol, so every number is non-thinking. Global steps still pay full price and account for more than 80% of DA's attended tokens. The wall-clock figures depend on MFU/MBU assumptions and were not measured end to end.

Qwen leans harder on global at long context, so both the save and the accuracy hold up less well than on Gemma. About 6% of Gemma-4-12B DA responses hit the 8K cap without finishing, which pushes attended-token totals above Vanilla; drop those traces and the number falls just below Vanilla. Scoring uses an LLM judge with a gold answer (Qwen-3.5-4B, Pearson r=0.99 versus Gemini-3.1-Pro), and four sources use questions synthesized by Gemini-3-Flash. Small models mostly cannot follow the protocol; do not expect the 4B class to work zero-shot.

Terms

Source

What people are saying

All paper explainers