Le Batching Diagonal Débloque le Parallélisme dans les Transformers à Mémoire Récurrente pour les Contexte Longs

Diagonal Batching Unlocks Parallelism in Recurrent Memory Transformers for Long Contexts

June 5, 2025

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

cs.AI

Résumé

Les modèles Transformer rencontrent des difficultés avec l'inférence en contexte long en raison de leur complexité temporelle quadratique et de leur utilisation de mémoire linéaire. Les Transformers à Mémoire Récurrente (RMT) offrent une solution en réduisant le coût asymptotique à un temps linéaire et une utilisation de mémoire constante. Cependant, leur mécanisme de mise à jour de la mémoire entraîne une exécution séquentielle, créant un goulot d'étranglement de performance. Nous introduisons le Diagonal Batching, un schéma d'ordonnancement qui débloque le parallélisme entre les segments dans les RMT tout en préservant la récurrence exacte. Cette approche élimine la contrainte séquentielle, permettant une inférence GPU efficace même pour des entrées uniques en contexte long, sans recourir à des techniques complexes de batching et de pipelining. Comme cette technique consiste uniquement en un réordonnancement des calculs à l'exécution, les modèles RMT existants peuvent l'adopter sans nécessiter de réentraînement. Appliqué à un modèle LLaMA-1B ARMT, le Diagonal Batching permet une accélération de 3,3x par rapport au LLaMA-1B standard avec attention complète, et de 1,8x par rapport à l'implémentation séquentielle des RMT sur des séquences de 131 072 tokens. En supprimant le goulot d'étranglement séquentiel, le Diagonal Batching réduit le coût et la latence de l'inférence, renforçant ainsi les RMT comme une solution pratique pour les applications réelles en contexte long.

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.