# Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers for Long Contexts

arXiv: https://arxiv.org/abs/2506.05229 · HTML: https://mlsanity.com/paper/2506.05229/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 30 of its 35 references that are in the map.
- First posted on arXiv 5 Jun 2025.

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints](https://mlsanity.com/paper/2305.13245/) (2023)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration](https://mlsanity.com/paper/2306.00978/) (2023)
- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://mlsanity.com/paper/2405.04434/) (2024)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [Compressive Transformers for Long-Range Sequence Modelling](https://mlsanity.com/paper/1911.05507/) (2019)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [Retentive Network: A Successor to Transformer for Large Language Models](https://mlsanity.com/paper/2307.08621/) (2023)
- [Linear Transformers Are Secretly Fast Weight Programmers](https://mlsanity.com/paper/2102.11174/) (2021)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [Parallelizing Linear Transformers with the Delta Rule over Sequence Length](https://mlsanity.com/paper/2406.06484/) (2024)
- [Memory Networks](https://mlsanity.com/paper/1410.3916/) (2014)
- [DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models](https://mlsanity.com/paper/2309.14509/) (2023)
- [Repeat After Me: Transformers are Better than State Space Models at Copying](https://mlsanity.com/paper/2402.01032/) (2024)
- [BABILong: Testing the Limits of LLMs with Long Context Reasoning-in-a-Haystack](https://mlsanity.com/paper/2406.10149/) (2024)
- [Recurrent Memory Transformer](https://mlsanity.com/paper/2207.06881/) (2022)
- [The Illusion of State in State-Space Models](https://mlsanity.com/paper/2404.08819/) (2024)
- [Speculative Decoding: Exploiting Speculative Execution for Accelerating Seq2seq Generation](https://mlsanity.com/paper/2203.16487/) (2022)
- [What Formal Languages Can Transformers Express? A Survey](https://mlsanity.com/paper/2311.00208/) (2023)
- [Block-Recurrent Transformers](https://mlsanity.com/paper/2203.07852/) (2022)
- [Associative Recurrent Memory Transformer](https://mlsanity.com/paper/2407.04841/) (2024)

## Built on by (most cited first, 2 of 2)

- [Extending LLM Context via Associative Recurrent Memory](https://mlsanity.com/paper/2607.11614/) (2026)
- [BCMT: Blockwise Causal Memory Transformer](https://mlsanity.com/paper/2608.13578/) (2026)

## Abstract (arXiv)

Transformer models struggle with long-context inference due to their quadratic time and linear memory complexity. Recurrent Memory Transformers (RMTs) offer a solution by reducing the asymptotic cost to linear time and constant memory usage. However, their memory update mechanism leads to sequential execution, causing a performance bottleneck. We introduce Diagonal Batching, a scheduling scheme that unlocks parallelism across segments in RMTs while preserving exact recurrence. This approach eliminates the sequential constraint, enabling efficient GPU inference even for single long-context inputs without complex batching and pipelining techniques. Because the technique is purely a run-time computation reordering, existing RMT models adopt it with no retraining. Applied to a LLaMA-1B ARMT model, Diagonal Batching yields a 3.3x speedup over standard full-attention LLaMA-1B and a 1.8x speedup over the sequential RMT implementation on 131,072-token sequences. By removing sequential bottleneck, Diagonal Batching reduces inference cost and latency, thereby strengthening RMTs as a practical solution for real-world, long-context applications.