title: Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers
    for Long Contexts
summary: 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.
paper_authors._id: 6842bc6855574a112d5733cc
paper_authors.name: Danil Sivtsov
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2025-06-07T05:48:36.994Z
paper_authors._id: 643e428edb5d62f97e5ba47b
paper_authors.avatarUrl: /avatars/00fddefc656ac114be87efcb88cbec2b.svg
paper_authors.fullname: Sivtsov Danil
paper_authors.type: user
paper_authors.user: svtdanny
paper_authors._id: 6842bc6855574a112d5733cd
paper_authors.name: Ivan Rodkin
paper_authors._id: 6842bc6855574a112d5733ce
paper_authors.name: Gleb Kuzmin
paper_authors._id: 6842bc6855574a112d5733cf
paper_authors.name: Yuri Kuratov
paper_authors.status: claimed_verified
paper_authors.statusLastChangedAt: 2025-06-07T05:48:33.509Z
paper_authors._id: 618b9540682ec1c38327e586
paper_authors.avatarUrl: https://cdn-avatars.huggingface.co/v1/production/uploads/618b9540682ec1c38327e586/v_ZBkfh8O9Zh6C2YQpuBX.jpeg
paper_authors.fullname: Yury Kuratov
paper_authors.type: user
paper_authors.user: yurakuratov
paper_authors._id: 6842bc6855574a112d5733d0
paper_authors.name: Ivan Oseledets
paper_title: Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers
    for Long Contexts
paper_summary: 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.
paper_ai_summary: Diagonal Batching enables parallel inference in Recurrent Memory Transformers, significantly improving speed and efficiency for long-context tasks.
paper_ai_keywords: 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