Aventures dans l'implémentation de Flash Attention en CUDA C++

Cet article détaille le parcours de l'auteur dans l'implémentation et l'optimisation de Flash Attention en CUDA C++. En commençant par une implémentation de base, l'auteur affine progressivement le noyau en utilisant des techniques telles que l'échange de mémoire partagée, le pipeline à deux étages et l'utilisation plus efficace de ldmatrix. Grâce à un profilage itératif et à l'optimisation, l'implémentation finale atteint des performances proches de la limite théorique du matériel. L'article approfondit également les subtilités de l'implémentation du softmax en ligne et de la résolution des conflits de banques de mémoire partagée, fournissant des informations précieuses aux développeurs CUDA C++.