STEEL : l'attention fusionnée sensible à la parcimonie réduit la consommation énergétique par 9,17 sur le NPU XDNA d'AMD

Les chercheurs présentent STEEL, la première implémentation open-source de FlashAttention ciblant les moteurs de traitement neuronal (NPU) de type XDNA, afin de relever le défi du mappage des mécanismes d'attention sur les SoC de classe portable.

• STEEL introduit une formulation de flux de données pour l'attention de préremplissage afin d'exploiter le parallélisme spatial et la mémoire embarquée.
• Il utilise un placement de pipeline sensible à la parcimonie pour gérer le déséquilibre de charge provenant des masques causaux, réduisant ainsi la surcharge de synchronisation.
• Évalué sur le SoC AMD Ryzen AI 9 HX 370, STEEL réduit la consommation énergétique en moyenne par 9,17 par rapport aux références CPU et par 1,75 par rapport aux références GPU.
• Sur le matériel XDNA 1, STEEL obtient une réduction de latence de 9,6 fois par rapport à l'état de l'art précédent.
• Il offre un gain de vitesse moyen de 22,8 fois par rapport aux implémentations d'attention couche par couche sur XDNA 2.

Ce travail permet une inférence économe en énergie pour les agents basés sur des grands modèles de langage sur les appareils edge, atténuant les problèmes de fiabilité et de confidentialité associés au déport vers le cloud.