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

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

We extend vLLM ( Kwon et al., 2023 ) to support the DA state machine through hooks on its attention metadata builder, with no kernel modifications or scheduler changes. At each decode step, the hook rewrites the request’s KV-cache block table so that only the kept blocks remain visible to the attention kernel, which simply reads less. This design lets DA inherit the efficiency of state-of-the-art attention kernels such as FlashAttention ( Dao et al., 2022 ) , reducing the KV blocks read per step. We analyze the projected wall-clock effect in Section 5.4 and explain integration details in Section B .

#### Scope and tradeoffs.

DA applies to the global attention layers only. The efficient layers of modern architectures, such as sliding window attention (SWA) ( Beltagy et al., 2020 ) and Gated DeltaNet (GDN) ( Yang et al., 2025b ) , have per-step costs bounded by a fixed window or recurrent state rather than by context length, so there is little for a mask to save and DA leaves them untouched ( Section B ). Within the global attention layers, DA reduces the per-step access cost of attention by masking KV cache positions. The protocol trades more decode steps for lower per-step attention cost. Whether this nets out favorably depends on the deployment regime, which we examine in the next section.

## 3 When does DA pay off?

#### Overview.

DA trades more decode steps for lower per-step attention cost, and we expect the savings to outweigh the added steps in large-batch deployments where hardware is fully utilized. We frame this through roofline wall-time : the time each operation contributes to decode latency when charged at its own hardware ceiling ( Williams et al., 2009 ) . We compute it per response, summing each operation’s work over all decode steps. Roofline wall-time is defined per hardware target and does not depend on operating choices such as batch size, so the comparison rests on the hardware alone rather than on any particular serving configuration. We summarize the key argument below and defer the full background on inference costs and the full derivation to Section C .

#### FFN roofline wall-time.

In large-batch inference, FFN parameter loads amortize across the batch and FFN becomes compute-bound, with roofline wall-time T FFN = FLOPs / ( Peak FLOPS × MFU ) T_{\text{FFN}}=\text{FLOPs}\,/\,(\text{Peak FLOPS}\times\text{MFU}) where Model FLOPs Utilization (MFU) ( Chowdhery et al., 2023 ) is the achieved fraction of peak compute. At decode the per-step GEMMs are skinnier than at prefill, so MFU sits at the low end of the range that well-optimized serving stacks achieve. We state the value we adopt in Section 5.4 and justify it in Section C.5 .

#### Attention roofline wall-time.

Attention KV reads are per-sequence and remain memory-bound, with roofline wall-time T attn = KV bytes / ( Peak BW × MBU ) T_{\text{attn}}=\text{KV bytes}\,/\,(\text{Peak BW}\times\text{MBU}) where Model Bandwidth Utilization (MBU) ( Agarwal et al., 2023 ) is the achieved fraction of peak memory bandwidth. A well-optimized memory-bound attention kernel sits high in its utilization range at large-batch decode. Again, we adopt a specific value in Section 5.4 and justify it in Section C.5 .

#### Attention dominates FFN at scale.

In optimized large-batch deployments ( Yu et al., 2022 ; Kwon et al., 2023 ) , attention roofline wall-time grows with both context length (KV cache size) and the number of decode steps, while FFN roofline wall-time grows only with the number of decode steps, with its per-step cost fixed by the active-parameter FLOPs, independent of context length. The hybrid backbones we study carry a third context-independent cost: their efficient layers (GDN on Qwen, SWA on Gemma) add a fixed per-step memory read that the mask does not touch. DA’s per-step reduction therefore lands on the global-attention KV read, where decode work concentrates, and its benefit is bounded by the share of decode time spent there. The savings are therefore largest in large-batch, long-context serving.

## 4 Experimental settings

#### Models.

We evaluate DA on six models across two families: Gemma-4-{31B, 12B, E4B} ( Gemma Team et al., 2026 ) and Qwen-3.6-27B, Qwen-3.5-{9B, 4B} ( Qwen Team, 2026b ) . The main comparison ( Section 5.1 ) uses the two largest models, Gemma-4-31B and Qwen-3.6-27B, and the model-size analysis ( Section 5.2 ) uses all six models. All natively support 256K input tokens except Gemma-4-E4B (128K).

#### Evaluation datasets.

We evaluate on 15 long-context sources spanning two task categories where selective attention is particularly relevant: (1) single-span retrieval and reasoning, and (2) multi-span reasoning. The benchmark suite draws from RULER ( Hsieh et al., 2024 ) , LongBench v1 ( Bai et al., 2024 ) , LongBench v2 ( Bai et al., 2025 ) , LooGLE ( Li et al., 2024a ) , and ZeroSCROLLS ( Shaham et al., 2023 ) , abbreviated LBv1, LBv2, and ZS hereafter. Eleven sources use the original QA samples, and the remaining four use synthetic QA generated with Gemini-3-Flash to address annotation quality or extend coverage under our task taxonomy. We exclude samples whose context length exceeds 244K (or 116K for Gemma-4-E4B). From each source we then draw up to 128 examples with a fixed random seed; for a given model the same examples are used across all three methods (Vanilla, DA, DA nm {}^{\text{nm}} ), and across models the draw is identical except where a different tokenizer or context limit changes which examples pass the length filter. We explain details in Section D.1 .

#### Baselines.

We compare DA against two baselines on the same tasks and models. All three arms share the same final-answer specification and <answer> convention, so the LLM judge scores them identically. See full prompts in Section F .

• Vanilla : serves the raw context inline with the question and a brief instruction to wrap the final answer in <answer> tags, with full causal attention. The gap between vanilla and DA nm {}^{\text{nm}} isolates the effect of the chunked prompt format.

• DA-no-mask (DA nm {}^{\text{nm}} ): uses the full DA prompt template but with full causal attention. The gap between DA nm {}^{\text{nm}} and DA isolates the effect of custom attention masking.

#### Evaluation.

Free-form generations make exact match unreliable, so we score with an LLM judge conditioned on the ground-truth answer ( Zheng et al., 2023 ) . For each question we generate a strict acceptance rubric with Gemini-3-Flash and apply it with thinking-enabled Qwen-3.5-4B as the judge, which achieves Pearson r = 0.99 r=0.99 with Gemini-3.1-Pro. See judge validation details in Section D.2 and full prompts in Section F .

#### Metrics.

We report two headline metrics for overall comparison: accuracy (fraction of responses judged as correct, with format failures counted as incorrect, averaged within each of the 15 sources and then equally across sources) and attended tokens (total KV positions attended during decoding, ∑ t attended ​ ( t ) \sum_{t}\text{attended}(t) , averaged per response within each source and then equally across the 15 sources, compared by percent reduction relative to vanilla). To estimate real-world cost in optimized large-batch deployments, we additionally report roofline wall-time as defined in Section 3 .

#### Inference settings.

We disable thinking mode for all models as preliminary experiments showed that they fail to follow the DA protocol within thinking tags. We discuss this limitation and its mitigation in Section 8 . For sampling, we use each model’s recommended parameters ( Section D.3 ). We use a max generation length of 8K tokens. All models are served on NVIDIA B200 GPUs via vLLM ( Kwon et al., 2023 ) , with our custom integration ( Section B ) used for DA.

## 5 Results

### 5.1 Main results

#### Accuracy drops modestly on average.

DA produces average accuracy drops of 1.27pp on Gemma-4-31B (87.01% to 85.74%) and 2.75pp on Qwen-3.6-27B (85.31% to 82.56%) across the 15 tasks. The aggregate change masks per-task variation. DA matches or improves on vanilla in 7 of 15 tasks on Gemma and 5 of 15 on Qwen, where many of the matches occur at the accuracy ceiling on niah retrieval tasks. The clearest single-task gains are longdep_qa on Gemma (+3.1pp) and code_repo on Qwen (+5.6pp). The losses are larger in multi-span reasoning than in single-span on both models (Gemma: 2.28pp vs 0.78pp category-average drop; Qwen: 3.59pp vs 2.34pp).

#### Attention cost falls below vanilla.

DA reduces the average attended-token count per response by 52.0% on Gemma-4-31B (13.43M to 6.45M) and 31.1% on Qwen-3.6-27B (22.54M to 15.52M). On Gemma the reduction holds across both categories (54.8% on single-span retrieval/reasoning, 48.5% on multi-span reasoning). On Qwen it holds as well (29.7% and 32.6%), though DA’s longer generations partly offset the per-step saving, and on five sources, most visibly qmsum and the two LBv2 QA sources, attended tokens rise above vanilla. In absolute terms, the largest per-task savings come from the longest-context tasks: code_repo (41.8M tokens saved on Gemma, 52.0M on Qwen) and dialogue_history (22.1M and 39.1M). Section 2 decomposes these effects into relative accuracy, decode steps, and attended tokens for all three arms on both headline models.

#### The attention mask is the source of the efficiency gain.

We isolate the contribution of the attention mask by comparing DA against DA-nm , which uses the same prompt template (chunked tool-use rendering and three-mode instructions) but with full causal attention. The chunked prompt format by itself is nearly lossless: DA-nm matches vanilla accuracy on Gemma (87.01% vs 87.01%) and stays within 0.69pp on Qwen (84.62% vs 85.31%). However, DA-nm ’s attended-token count is 66.2% higher than vanilla on Gemma and 28.8% higher on Qwen, reflecting the longer generations the DA protocol induces (DA and DA-nm both run roughly 15–35% more decode steps than vanilla; Section 2 ). The mask converts this overhead into a net saving: relative to DA-nm it cuts attended tokens by 71.1% on Gemma and 46.5% on Qwen, moving DA from 66.2% above vanilla to 52.0% below on Gemma, and from 28.8% above to 31.1% below on Qwen. The mask is therefore the source of both DA’s efficiency and most of its accuracy cost: it supplies essentially all of the token savings, and relative to DA-nm it accounts for the bulk of the accuracy drop ( − - 1.27pp on Gemma, − - 2.06pp on Qwen), since the chunked format alone is nearly free.

### 5.2 Model capability scaling

We examine how DA behaves as the backbone model becomes larger. As shown in Section 3 , DA’s relative accuracy improves consistently within both the Gemma 4 and Qwen 3.5/3.6 families as the backbone scales, narrowing the gap to vanilla. DA also attends fewer tokens than vanilla in five of the six models, giving an additional axis of inference-time efficiency; the lone exception, Gemma-4-12B , is a generation-length artifact rather than a masking failure (discussed below).

#### Stronger models close the accuracy gap.

Section 3(a) reports DA’s accuracy relative to vanilla across six models from the Qwen and Gemma families. Relative accuracy increases monotonically with model size on both families: from 29% at Gemma-4-E4B to 99% at Gemma-4-31B , and from 64% at Qwen-3.5-4B to 97% at Qwen-3.6-27B . Within each family, smaller models retain a smaller fraction of vanilla accuracy under DA. We hypothesize that this reflects two compounding requirements DA places on the model: it must follow the protocol (parseable tags, valid chunk references, coherent mode orderings) and produce useful answers under significant information restriction. Both requirements plausibly benefit from general model capability. At the low end, Gemma-4-E4B retains only 29% of vanilla accuracy, much of which is an outright failure to emit a parseable answer under the protocol: its focus-parse success rate is only 58%, against 99% at the largest models ( Section 6 , Section 6 ). This suggests that DA’s gains depend on a minimum level of base capability that the smallest models do not yet meet.

#### Token savings are roughly scale-independent.

Relative attended tokens ( Section 3(b) ) are less regular than accuracy and show no strong size trend: Gemma’s models attend roughly 46–48% of vanilla and Qwen’s 50–69%. The per-step masking saving is itself roughly scale-independent, at a per-step attention ratio near 0.5 on five of the six models (0.65 on the smallest, Gemma-4-E4B ); the spread in the totals comes instead from generation length, which each model sets for itself under the zero-shot prompt through its own mix of global, focus, and local modes. The one model above vanilla, Gemma-4-12B (183%), reflects a decode-length artifact rather than the mask ( Section 3 ). The per-step attention ratio is therefore the cleaner measure of the masking mechanism, while the total reflects the elicitation ( Section 8 ).

### 5.3 Context-length scaling

#### Accuracy holds while token saving grows with context.

We pool all 15 sources and bin by context length ( Section 4 ). DA’s relative accuracy holds within about one point of vanilla up to 32K tokens and then declines modestly, to roughly 96% in the longest bin ( Section 4(a) ); the maskless DA-nm line does not show this decline, indicating that the long-context accuracy cost comes from the attention mask rather than the chunked prompt format. In absolute terms, DA’s token saving grows sharply with context ( Section 4(b) ): DA attends about 1M fewer tokens than vanilla at short context and about 21M fewer in the longest bin, while DA-nm ’s overhead grows the other way. Because DA’s per-step masking saves a roughly constant fraction of tokens ( ∼ \sim 50%; Section E ), the absolute saving scales with the raw context length, so the benefit is largest exactly where long-context decoding is most expensive.

### 5.4 Efficiency results

#### Theoretical wall-clock time on optimized inference setups.

The attended-token reductions above translate into roofline-estimated decode-time wall-clock savings on the kind of inference setup that long-context serving targets. We estimate this for both headline models on a single B200 accelerator in bf16, following the roofline wall-time method of Section 3 and Section C : each operation is charged at its own hardware ceiling, a Model FLOPs Utilization (MFU) of 40% for the compute-bound matrix multiplies and a Model Bandwidth Utilization (MBU) of 70% for the memory-bound reads. We divide each decode step into three costs and sum them over the full generation, averaged across the 15 tasks. The matmul cost covers every projection, feed-forward, and output-head matrix multiply, and scales with the active parameters and the number of decode steps, independent of context length. The global-memory cost is the global-attention KV read, whose per-step size grows with the number of attended tokens, and is the one term the DA mask reduces. The local-memory cost is the bounded per-step read of the model’s efficient layers, the SWA KV cache on Gemma-4-31B and the GDN recurrent state on Qwen-3.6-27B , which is context-independent and which DA cannot touch. Section 3 reports the work and resulting wall-clock time for each.