Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers for Long Contexts

TL;DR

This paper introduces Diagonal Batching, a run-time scheduling scheme that enables parallel inference in Recurrent Memory Transformers, significantly improving speed and efficiency for long-context processing without retraining.

Contribution

Diagonal Batching is a novel scheduling method that unlocks parallelism in RMTs, eliminating sequential constraints and enabling efficient GPU inference for long sequences.

Findings

3.3x speedup over standard full-attention LLaMA-1B

1.8x speedup over sequential RMT on 131,072 tokens

Reduces inference cost and latency for long-context models

Abstract

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.