Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers
  for Long Contexts
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.
6842bc6855574a112d5733cc
Danil Sivtsov
claimed_verified
2025-06-07T05:48:36.994Z
643e428edb5d62f97e5ba47b
/avatars/00fddefc656ac114be87efcb88cbec2b.svg
Sivtsov Danil
user
svtdanny
6842bc6855574a112d5733cd
Ivan Rodkin
claimed_verified
2025-08-26T09:51:11.640Z
63c1ac8cc58fcfeac186bda2
https://cdn-avatars.huggingface.co/v1/production/uploads/63c1ac8cc58fcfeac186bda2/MajMhLgsvCBN81tmO9PZc.jpeg
Ivan Rodkin
user
irodkin
6842bc6855574a112d5733ce
Gleb Kuzmin
6842bc6855574a112d5733cf
Yuri Kuratov
claimed_verified
2025-06-07T05:48:33.509Z
618b9540682ec1c38327e586
https://cdn-avatars.huggingface.co/v1/production/uploads/618b9540682ec1c38327e586/v_ZBkfh8O9Zh6C2YQpuBX.jpeg
Yury Kuratov
user
yurakuratov
6842bc6855574a112d5733d0
Ivan Oseledets
claimed_verified
2025-06-12T12:43:09.299Z
6169a581d05945bfd8718dfa
/avatars/1892ab06a7ddb557232777de3cbec470.svg
Ivan Oseledets
user
oseledets
Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers
  for Long Contexts
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.
Diagonal Batching enables parallel inference in Recurrent Memory Transformers, significantly improving speed and efficiency for long-context tasks.
Transformer models long-context inference quadratic time complexity linear memory complexity Recurrent Memory Transformers RMTs memory update mechanism sequential execution Diagonal Batching run-time computation reordering parallelism GPU inference LLaMA-1B ARMT model full-attention LLaMA-1B inference cost latency