FlashAttention

FlashAttention Menghitung Softmax Eksak per Tile Tanpa Menyimpan Matriks Skor Penuh

Scaled dot-product attention standar membentuk matriks skor dengan dua sumbu panjang yang sama-sama mengikuti posisi sequence. Untuk satu attention head, S = QK^T / sqrt(d) P = softmax(S) O = PV definisi matematisnya ringkas, tetapi implementasi GPU langsung dapat menulis intermediate besar S dan P ke high-bandwidth memory lalu membacanya kembali. FlashAttention mengubah aliran data tersebut. Query, key, dan value diproses dalam blok di memori on-chip, sementara state softmax per baris dibawa dari satu tile ke tile berikutnya agar seluruh tile skor tetap dapat digabungkan secara eksak.