RAT+: Train Dense, Infer Sparse -- Recurrence Augmented Attention for Dilated Inference

Xiuying Wei; Caglar Gulcehre

cs.LGFeb 20, 2026

RAT+: Train Dense, Infer Sparse -- Recurrence Augmented Attention for Dilated Inference

Authors: Xiuying Wei, Caglar Gulcehre

Organizations: CLAIRE lab at EPFL, Lausanne, Switzerland.

Abstract

Structured dilated attention has an appealing inference-time efficiency knob: it reduces the FLOPs of attention and the KV cache size by a factor of the dilation size D, while preserving long-range connectivity. While prior work studies it by training each configuration from scratch, directly sparsifying a pretrained attention model into a dilated pattern leads to severe accuracy degradation, preventing flexible reuse across inference scenarios. We introduce RAT+, a dense-pretraining architecture that augments attention with full-sequence recurrence and active recurrence learning. A single RAT+ model is pretrained densely once and can then be flexibly switched at inference time to dilated attention (optionally with local windows) or hybrid layer/head compositions, requiring only a short 1B-token resolution adaptation rather than retraining separate sparse models. At 1.5B parameters trained on 100B tokens, RAT+ closely matches dense accuracy at D = 16, and drops by about 2-3 points at D = 64 on commonsense reasoning and LongBench tasks. We further scale to 2.6B and 7.6B parameters and observe even more promising performance (e.g., a 1-point average accuracy loss with a 64x reduction in attention FLOPs and KV cache size). Code is available at https://github.com/wimh966/rat-plus.

Figures & tables

Figure 1: (a) For architectural simplicity, we adopt an extreme overlapped setting, i.e., full-sequence recurrence with $\mathtt{L=T}$ . (b) Joint training to preserve dense attention capability while enforcing active recurrence learning with desired effective length $\mathtt{L^{*}}=64$ . (c) After pretraining, the resulting model can be efficiently adapted to various sparse inference patterns including effective results on dilated attention (optional local attention). It also has better performance on top-k block attention compared to attention.

Figure 2: Efficiency results of the temporal-mixing operator on a single GH200 GPU, covering both prefilling and decoding scenarios with hidden dimension

\mathtt{H}

. Prefilling latency is measured on sequences of 262K tokens. Decoding latency is measured for 256 or 128 batches of tokens for the two hidden dimensions, respectively; the baseline runs out of memory beyond 32K tokens. We use FlexAttention ( Dong et al., 2024 ) for prefilling (with

\mathtt{D=1}

reducing to FlashAttention), and FlashAttention ( Dao et al., 2022 ) for decoding. Recurrence is implemented by associative scan in PyTorch for prefilling and a simple step update for decoding. Additional results are reported in Subsec. B.4 .

Figure 3: Results of top-k block attention with block size

\mathtt{D}

and number of selected blocks

\mathtt{K}

on the hard NIAH-MK-2 and NIAH-MK-3 tasks from the RULER benchmark with

\mathtt{T=4096}

. RAT+ (no ARL in SFT) means disabling active recurrence learning during the SFT, which further demonstrates the benefit of the recurrence. More results include remaining tasks in Table 17 , MoBA-style top-k block attention in Table 18 .

Figure 4: Maximum decoding throughput (tokens/sec) of the full 1.5B and 7B models for decoding 1024 tokens, measured at context lengths of 4096 and 16384, corresponding to prefilling lengths of 3072 and 15360 tokens, respectively.

Figure 5: Comparison with GQA/MQA using different numbers of KV heads. Joint training is also applied to GQA/MQA (

\mathtt{D^{\dagger}=1}

\mathtt{W=64}

) to match training FLOPs. RAT+ achieves lower PPL and is more flexible: it requires only a single pretraining run and allows the local KV cache size to be configured at inference time.

Figure 6: Scaling-up experiments: we report validation loss on a held-out 0.5B-token subset to illustrate the even smaller loss gap between dense and sparse variants as model size increases. The starred points refer to attention models trained with

\mathtt{D^{\dagger}=1}

and

\mathtt{W=64}

, matched in training FLOPs at the same model scale, and are included as a reference for comparison with dense RAT+.

Figure 7: Choice of

\mathtt{L=64}

\mathtt{L^{*}=16}

means

\mathtt{D=1}

and

\mathtt{D=16}

joint training to enable 16 length capability for recurrence. Different performance on

\mathtt{D=1}

comes from different training FLOPs and different recurrence abilities in capturing the corresponding length. Grey points correspond to recurrence length generalization.

Figure 8: 1B-token adaptation on two pretrained models. It is evident that various dilated patterns quickly achieve stable loss values within a few hundred million tokens. We employed a simple optimization scheme with no warmup, which may explain the slight loss increase of

\mathtt{D=1}

at first, after which it recovers.

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 9: L2 norm values of recurrence outputs at different time steps. We observe that the outputs at early time steps differ significantly. The first row shows an initialized network using our simple recurrence at layers 0, 6, 18, and 23. The second row corresponds to the same initialized network but with a non-zero initial cell state provided to the recurrence. The third row shows the results of the pretrained network. We visualize the outputs of our simple recurrence here; a similar phenomenon has also been observed in standard recurrent networks, as reported in Cooijmans et al. (2016) .

SFT models	NIAH-S-1	NIAH-S-2	NIAH-S-3	NIAH-MK-1	NIAH-MK-2	NIAH-MK-3	NIAH-MV	NIAH-MQ
$\mathtt{D=1}$	100	100	100	100	100	99.6	99.85	100
$\mathtt{D=2}$	100	100	99.2	100	99.2	96	99.8	99.8
$\mathtt{D=4}$	100	100	97.8	99.8	97.4	88.8	99.75	99.5
$\mathtt{D=8}$	100	100	96	99.8	96.4	73.8	99.4	99.15
$\mathtt{D=16}$	100	100	89.2	99.6	95.2	54.6	99.3	98.15
$\mathtt{D=32}$	100	100	78.8	99.2	78	40.4	97.35	96.35

Appendix

Table 16: Retrieval ability : Accuracy performance with exact match scoring on the Needle-in-Haystack tasks with different configurations from the RULER benchmark ( Hsieh et al., 2024 ) . We use

\mathtt{T=4096}

and bold less favorable results. We found they mainly occur when the text query/key/values are the UUID, and may incorrectly copy parts of the answers. Also, in NIAH-MK-2, larger dilation sizes also perform worse.

Quest top-k	NIAH-S-1	NIAH-S-2	NIAH-S-3	NIAH-MK-1	NIAH-MK-2	NIAH-MK-3	NIAH-MV	NIAH-MQ
attention
$\mathtt{D^{\dagger}=1}$ , All	100	100	99.8	100	100	99.6	100	99.9
$\mathtt{D^{\dagger}=16}$ , $\mathtt{K=16}$	100	99.8	99.4	99.6	71.2	5.4	93.95	94.25
$\mathtt{D^{\dagger}=16}$ , $\mathtt{K=64}$	100	100	100	100	99.6	85.2	99.9	99.6
$\mathtt{D^{\dagger}=64}$ , $\mathtt{K=8}$	100	100	85.8	98.4	20.6	0.2	87.55	88.55
$\mathtt{D^{\dagger}=64}$ , $\mathtt{K=16}$	100	100	99	99.8	63.2	15.2	97.45	96.85

Appendix

Table 17: Retrieval ability of the top-k block attention for different models. On the following three pretrained-plus-SFT models, we apply Quest ( Tang et al., 2024 ) style top-k block attention, where the critical block is determined by the Min/Max statistics of attention keys within each block.

MoBA top-k	NIAH-S-1	NIAH-S-2	NIAH-S-3	NIAH-MK-1	NIAH-MK-2	NIAH-MK-3	NIAH-MV	NIAH-MQ
attention
$\mathtt{D^{\dagger}=1}$ , All	100	100	99.8	100	100	99.6	100	99.9
$\mathtt{D^{\dagger}=16}$ , $\mathtt{K=16}$	100	91	56	84	65.2	6.2	48.1	55.9
$\mathtt{D^{\dagger}=16}$ , $\mathtt{K=64}$	100	99.8	96.2	99.6	97.8	77	95.35	96.45
$\mathtt{D^{\dagger}=64}$ , $\mathtt{K=8}$	90.2	66.8	20	54.2	29.8	5.8	30.1	34.55
$\mathtt{D^{\dagger}=64}$ , $\mathtt{K=16}$	99.2	95.8	79	91.8	66.2	22.6	70.7	76.35

Appendix

Table 18: Retrieval ability of top-k block attention across different models. For the following three pretrained-plus-SFT models, we apply MoBA-style top-k block attention ( Lu et al., 2025 ) , where the critical block is selected based on the mean-pooled attention keys within each block. Notably, we apply this mechanism to all layers and to both the prefilling and decoding stages, without any further fine-tuning. In contrast, the original MoBA paper uses continued pretraining and reports applying the method only during the prefilling stage to avoid noticeable accuracy degradation. Our goal here is primarily to demonstrate the effectiveness of RAT+ over the attention backbone.

LLM Inference Acceleration

LLM inference acceleration develops algorithms and systems that reduce language-model latency, memory use, and serving cost while preserving output quality.

See all 1,084 papers

May 27, 2026

Augmenting Attention with Exponentially Decaying Memory Improves Query-Aware KV Sparsity

Efficient inference is critical for long-context language models, where attention computation and KV-cache access dominate the cost. Recent work RAT+, introduces a recurrence-augmented attention backbone that enables flexible dilated attention at inference time. In this paper, we investigate whether this exponentially decaying memory can also improve existing query-aware sparse inference methods. Using representative methods including Quest, MoBA, and SnapKV, we show that RAT+ consistently improves accuracy over standard attention across sparse budgets on eight needle-in-a-haystack tasks. We validate these gains both on the released checkpoints from the RAT+ paper and on OLMo2-7B, which we continue pretraining with the added memory module for 10B tokens. Finally, we propose two hypotheses explaining why this memory module benefits query-aware sparse inference and design targeted experiments to support them.

Xiuying Wei, Caglar Gulcehre

EPFL, CLAIRE Lausanne, Switzerland

KV Caching Memory-Augmented Language Models Self-Attention Sparse Attention Long-Context Language Model Inference LLM Inference Acceleration+1

Ask PDF

Sep 17, 2026

RBS-Attention: Radius-Bounded Sparse Prefill for Long-Context Large Language Models

Long-context large language model inference is increasingly limited by prefill, where dense self-attention processes the entire prompt before generation begins. Sparse block selection can reduce this cost, but a block centroid may hide a highly relevant token among many irrelevant ones. We call this failure mode mean dilution and propose RBS-Attention, a training-free sparse-prefill method with two complementary selection branches. A centroid base branch captures average relevance, while a rescue branch uses the maximum key-block radius and its prompt-, layer-, and head-dependent distribution to identify blocks at risk of underestimation. Independently thresholding the two branches and combining their masks controls the contribution of rescue blocks while preserving regular block-sparse FlashAttention execution. On H100 GPUs, RBS-Attention achieves 20.65

\times

standalone prefill-attention speedup, 11.92

\times

vLLM prefill-attention speedup, and 5.97

\times

end-to-end time-to-first-token speedup at 128K on Qwen3-30B-A3B-Instruct-2507-FP8. On the dense Qwen3-32B model, it obtains 88.65 overall RULER accuracy versus 89.52 for dense attention; LongBench-v2, InfiniteBench, and Video-MME provide additional quality evaluation. Supporting experiments measure actual retention, compare selectors at matched density, and characterize block-size, threshold, and memory behavior. Together, these results support radius-adaptive dual-branch selection as an effective approach to long-context prefill.

Chuxu Song, Jiuqi Wei, Zhencan Peng

Department of Computer Science, Rutgers University · OceanBase, Ant Group

Long-Context Language Model Inference LLM Inference Acceleration Sparse Attention Block-Sparse Attention

Ask PDF

May 16, 2026

Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps

Long-context inference in large language models is bottlenecked by the quadratic cost of full attention. Existing efficient alternatives often rely either on native sparse training or on heuristic token eviction, creating an undesirable trade-off among efficiency, training cost, and accuracy. In this work, we show that full-attention LLMs are already intrinsically sparse and can be transformed into highly sparse models with only minimal adaptation. Our approach is built on three observations: (1) only a small subset of attention heads truly requires full long-context processing; (2) long-range retrieval is governed primarily by a low-dimensional subspace, allowing relevant tokens to be retrieved efficiently with a 16-dimensional indexer; and (3) the useful token budget is strongly query-dependent, making dynamic top-

p

selection more suitable than fixed top-

k

sparsification. Based on these insights, we propose RTPurbo, which retains the full KV cache only for retrieval heads and introduces a lightweight token indexer for sparse attention. By exploiting the model's intrinsic sparsity, RTPurbo achieves sparsification with only a few hundred training steps. Experiments on long-context benchmarks and reasoning tasks show that RTPurbo preserves near-lossless accuracy while delivering substantial efficiency gains, including up to a 9.36

\times

prefill speedup at 1M context and about a 2.01

\times

decode speedup. These results suggest that strong sparse inference can be obtained from standard full-attention training without expensive native sparse pretraining.

Yanke Zhou, Yiduo Li, Hanlin Tang +6

Nanjing University · Alibaba Group

Self-Attention Long-Context Language Model Inference LLM Inference Acceleration Sparse Attention Multi-Head Attention

Ask PDF

Abstract

LLM Inference Acceleration

Explore similar work

Augmenting Attention with Exponentially Decaying Memory Improves Query-Aware KV Sparsity

RBS-Attention: Radius-Bounded Sparse Prefill for Long-Context Large Language Models

Full Attention Strikes Back: Transferring Full Attention into Sparse within Hundred Training Steps