Dans cet article, nous proposons une nouvelle architecture basée sur un tableau systolique, le Flash Attention Systolic Array (FSA), pour une accélération efficace des modèles Transformer à l'aide de l'algorithme FlashAttention. Les accélérateurs existants basés sur un tableau systolique souffrent d'une faible utilisation en raison de la multiplication matricielle non consécutive et de l'opération softmax de FlashAttention. FSA résout ce problème en exécutant l'intégralité de l'algorithme FlashAttention au sein d'un seul tableau systolique grâce à un nouvel algorithme d'ordonnancement appelé SystolicAttention. SystolicAttention améliore l'utilisation du tableau grâce à un chevauchement élément par élément précis et garantit la stabilité numérique en conservant l'ordre d'origine des opérations en virgule flottante. Nous implémentons FSA comme un RTL synthétisable et le comparons à AWS NeuronCore-v2 et Google TPUv5e. Les résultats montrent qu'il atteint une utilisation d'attention FLOPs/s 1,77 fois supérieure (par rapport à NeuronCore-v2) et 4,83 fois supérieure (par rapport à TPUv5e) avec une surcharge d'environ 10 %.
