Implementasi Flash Attention Cepat di CUDA C++

Penulis mengimplementasikan algoritma Flash Attention untuk GPU NVIDIA 5090 menggunakan CUDA C++.

Flash Attention ditulis mulai dari versi dasar hingga lima versi optimalisasi dengan teknik berbeda.

Versi 1 (dasar) mencapai 68% dari Speed-of-Light (SOL), versi 2 (swizzling memori bersama) naik menjadi 86%.

Versi 3 (pipeling dua tahap) mencapai 90% SOL, versi 4 (ldmatrix.x4 untuk K dan V) mencapai 93%, dan versi 5 (pipeling lebih baik) mencapai 94%.

Implementasi memanfaatkan instruksi cp.async untuk memindahkan data global→shared, ldmatrix untuk shared→register, dan mma untuk komputasi.

Online softmax diimplementasikan secara in-place dengan pembaruan keadaan atensi yang asosiatif.

Profiling dengan Nsight Compute mengidentifikasi konflik bank memori bersama yang diatasi dengan swizzling alamat.

Pipelining dua tahap digunakan untuk mengoverlap memori dan komputasi menggunakan cp.async.commit_group dan cp.async.wait_group.

Teknik ldmatrix.x4 mengurangi jumlah instruksi ldmatrix dan meningkatkan performa.

Implementasi akhir mendekati batas teoretis tetapi masih sedikit di bawah performa cuDNN.

Usulan pengembangan selanjutnya termasuk backward pass, kuantisasi, penggunaan TMA, dan PagedAttention.