Le traitement des textes longs a longtemps constitué un goulot d’étranglement majeur pour les grands modèles linguistiques (LLM), en raison de la croissance quadratique du coût computationnel lié au mécanisme d’attention auto-attentionnelle. Plus la séquence s’allonge, plus les calculs deviennent prohibitifs, rendant l’extension des modèles difficile. En février de cette année, Moonshot AI avait proposé une architecture prometteuse, le Mixture of Block Attention (MoBA), inspirée des systèmes à experts mélangés (MoE). MoBA découpe le texte en blocs, puis utilise un « routeur » dynamique et économe pour ne calculer l’attention que sur les blocs les plus pertinents. Cette approche permet de réduire la complexité algorithmique de quadratique à quasi-linéaire, offrant ainsi une solution théoriquement évolutive pour les contextes longs. Cependant, malgré ses promesses, MoBA souffre de deux limitations majeures en pratique : d’une part, ses principes fondamentaux restent mal compris, et d’autre part, son implémentation actuelle n’est pas optimisée pour les architectures matérielles modernes, notamment les GPU. En particulier, lorsqu’on utilise de petits blocs — théoriquement plus efficaces — les surcoûts liés à la gestion de nombreux petits fragments (comme les accès mémoire non continus ou les opérations de tri) peuvent annuler tout gain computationnel. Ce problème a freiné la mise en œuvre industrielle de cette approche. Face à ce défi, une équipe conjointe du laboratoire H.A.N. du MIT, dirigé par le professeur Han Song, et d’NVIDIA a développé FlashMoBA, une solution révolutionnaire qui combine optimisation logicielle fine et conception matérielle adaptée. Cette recherche, publiée sur arXiv, dévoile pour la première fois les mécanismes sous-jacents au succès de MoBA, en les modélisant comme un problème de rapport signal-bruit (SNR). Elle montre que la précision du routeur dépend directement du rapport entre la dimension des têtes d’attention (d) et la taille des blocs (B). Ainsi, pour améliorer la performance, il suffit, dans un modèle de taille fixe, de réduire la taille des blocs. En outre, l’ajout d’un opérateur de convolution courte favorise la formation de groupes d’informations pertinentes au sein des blocs, amplifiant ainsi les signaux clés. Mais le vrai défi réside dans l’exécution sur GPU. Traiter des milliers de petits blocs revient à demander à un cargo de livrer des milliers de lettres individuelles : inefficace, coûteux, et lent. FlashMoBA surmonte cette limitation grâce à une architecture CUDA entièrement repensée, conçue pour minimiser les transferts de données vers la mémoire HBM et maximiser l’utilisation des unités de calcul. Deux innovations clés sont à l’origine de ses performances exceptionnelles : FlashTopK : une pipeline intégrée qui calcule simultanément les centres de blocs et sélectionne les meilleurs candidats (top-k) sans jamais générer de matrice de scores massive en mémoire. Cela évite à la fois les pics de mémoire et les goulets d’étranglement liés au tri. Gather-and-Densify : une stratégie en deux phases inspirée d’un bibliothécaire efficace. Elle rassemble d’abord les blocs pertinents dans le cache rapide du GPU (SRAM), puis les regroupe en opérations matricielles denses — parfaitement adaptées aux GPU. Cette approche élimine les accès mémoire aléatoires et permet d’exploiter pleinement la puissance de calcul. Les résultats sont impressionnants : sur des séquences de 64K tokens, FlashMoBA est 7,4 fois plus rapide que l’implémentation originale de MoBA, avec une réduction de 6,1 fois de la mémoire utilisée. Alors que MoBA échoue à 128K, FlashMoBA traite sans problème des séquences de 512K tokens, sans débordement mémoire. Dans les comparaisons directes avec FlashAttention-2, l’état de l’art actuel, FlashMoBA atteint une accélération de 14,7 fois sur les longues séquences. Et ce gain n’est pas au détriment de la qualité : des modèles entraînés de zéro montrent une amélioration significative de la précision, notamment en tâches de modélisation linguistique et de recherche dans des contextes longs. Cela s’explique par une réduction efficace du phénomène de « dilution de l’attention », où les modèles classiques perdent de vue les informations cruciales dans des textes très longs. En somme, FlashMoBA ne se contente pas d’optimiser une implémentation — il repense l’architecture de l’attention pour les longues séquences, en synergie parfaite entre théorie, algorithmique et matériel. Une avancée majeure pour le futur des grands modèles linguistiques.