FlashAttention-2: Attention yang Lebih Cepat dengan Paralelisasi dan Pembagian Kerja yang Lebih Baik

• Muncul model bahasa dengan konteks yang lebih panjang seperti GPT-4(32k), MPT(65k), dan Calude(100k)
• Memperluas panjang konteks transformer itu sulit karena kebutuhan runtime & memori meningkat hingga pangkat empat
• FlashAttention yang dirilis tahun lalu digunakan secara luas karena mengurangi penggunaan memori dan meningkatkan kecepatan attention
• Saat dirilis, ia sudah 2-4x lebih cepat, tetapi masih ada ruang untuk perbaikan. Dibandingkan operasi perkalian matriks (GEMM) yang dioptimalkan, performanya masih belum secepat itu dan secara teoretis hanya mencapai 25-40% dari FLOPs/s maksimum (maksimum 124 TFLOPs/s pada GPU A100)
• FlashAttention-2 2x lebih cepat daripada versi sebelumnya dan memberikan performa hingga 230 TFLOP/s pada GPU A100
• Dalam pelatihan model bahasa bergaya GPT, performanya mencapai hingga 225 TFLOPS (72% utilisasi FLOP model)
• Algoritme disesuaikan untuk mengurangi non-matmul FLOPs
• Paralelisasi yang lebih baik, serta perubahan cara pembagian kerja di setiap thread block
• Jumlah Head Dimensions diperluas dari 128 menjadi 256