Diagonales Batching ermöglicht Parallelität in Rekurrenten Speicher-Transformern für lange Kontexte

Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers for Long Contexts

June 5, 2025

Autoren: Danil Sivtsov, Ivan Rodkin, Gleb Kuzmin, Yuri Kuratov, Ivan Oseledets

cs.AI

Zusammenfassung

Transformer-Modelle haben Schwierigkeiten mit Langzeitkontext-Inferenz aufgrund ihrer quadratischen Zeit- und linearen Speicherkomplexität. Recurrent Memory Transformers (RMTs) bieten eine Lösung, indem sie den asymptotischen Aufwand auf lineare Zeit und konstanten Speicherverbrauch reduzieren. Allerdings führt ihr Speicher-Aktualisierungsmechanismus zu einer sequenziellen Ausführung, was einen Leistungsengpass verursacht. Wir stellen Diagonal Batching vor, ein Scheduling-Schema, das Parallelität über Segmente in RMTs ermöglicht, während die exakte Rekurrenz erhalten bleibt. Dieser Ansatz beseitigt die sequenzielle Einschränkung und ermöglicht eine effiziente GPU-Inferenz selbst für einzelne Langzeitkontext-Eingaben ohne komplexe Batching- und Pipelining-Techniken. Da die Technik lediglich eine Neuordnung der Laufzeitberechnung ist, können bestehende RMT-Modelle sie ohne erneutes Training übernehmen. Angewendet auf ein LLaMA-1B ARMT-Modell, erzielt Diagonal Batching eine 3,3-fache Beschleunigung gegenüber dem standardmäßigen Full-Attention LLaMA-1B und eine 1,8-fache Beschleunigung gegenüber der sequenziellen RMT-Implementierung bei Sequenzen mit 131.072 Tokens. Durch die Beseitigung des sequenziellen Engpasses reduziert Diagonal Batching die Inferenzkosten und -latenz und stärkt damit RMTs als praktische Lösung für reale Anwendungen mit Langzeitkontext.

English

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.