FlashAttention 4: Perhatian yang Lebih Cepat dan Efisien Memori untuk LLM
Pendahuluan
Perhatian produk titik-mengalikan berskala (Scaled‑dot‑product attention, SDPA) mendominasi waktu inferensi dan anggaran energi model bahasa besar (LLM): hampir semua eksekusi operator terjadi dalam satu primitif. Attention mengambil query (Q) dan mengalikannya dengan key (K). Kemudian hasilnya dinormalisasi melalui operasi softmax sebelum dikalikan dengan value (V) untuk menghasilkan keluaran. Ini dibatasi oleh memori: membaca Q/K/V berkali-kali dari memori berbandwidth tinggi, kemudian menulis ulang ubin-ubin sementara secara berulang.
Keluarga FlashAttention mengurangi hambatan ini dengan menyimpan lebih banyak data di chip untuk waktu yang lebih lama dengan memanfaatkan multiprosesor streaming GPU (SM). Pertama kali dibuka sebagai sumber terbuka pada tahun 2022, FlashAttention merancang ulang rutin perhatian untuk menghitung softmax secara langsung. Hanya beberapa baris dari Q × K yang disimpan di chip pada waktu tertentu.
Artikel ini berfungsi sebagai panduan mendalam yang ditujukan untuk praktisi mengenai FlashAttention 4 (FA4). Kami akan merangkum apa saja yang telah berubah sejak versi sebelumnya dan memberikan beberapa panduan mengenai adopsi/benchmarking. Tujuan kami adalah untuk membekali insinyur infrastruktur LLM, insinyur kernel, dan tim platform ML dengan informasi yang cukup agar mengetahui kapan – dan jika – FA4 sebaiknya digunakan dalam tumpukan mereka.
Poin Penting
- A4 dirancang sebagai kernel perhatian pertama Blackwell, ditujukan untuk penerapan pada GPU SM100. Ini menampilkan pipeline 5-tahap yang khusus untuk warp, memungkinkan tumpang tindih yang lebih besar dan penggunaan ulang di chip dibandingkan dengan generasi sebelumnya dari FlashAttention.
- Efisiensi ditingkatkan dalam softmax melalui dua teknik utama: menjalankan perangkat lunak exp2() pada inti CUDA itu sendiri untuk mengurangi kontensi pada SFU, dan penskalaan ulang adaptif secara online untuk menghindari perhitungan penskalaan ulang yang tidak perlu sambil menjaga kestabilan.
- Masih ada beberapa fitur yang hilang: FA4 menggunakan forward-first dan saat ini belum memiliki implementasi untuk backward pass, varlen, dan GQA/MQA, yang membatasi penggunaannya untuk pelatihan dan beberapa arsitektur model.
- Adopsi FA4 secara pragmatis: gunakan terlebih dahulu untuk inferensi Blackwell, pasang di belakang benchmarking + pemeriksaan kebenaran + cadangan, dan pertahankan FA3 (Hopper)/FA2 (Ampere/Ada) sebagai default stabil untuk kebutuhan produksi yang luas.
Mengapa kernel perhatian masih mendominasi biaya LLM
Meskipun GPU menyediakan FLOPS yang tinggi, model transformer seringkali dibatasi oleh bandwidth memori karena perhatian harus memuat Q, K, dan V dari memori global berkali-kali. Untuk menghitung output perhatian untuk setiap token query, masing-masing harus melakukan dot-product dengan hingga ribuan token key, menormalisasi dan menskalakan hasil ini, dan akhirnya mengambil nilai. Implementasi sederhana menghitung seluruh matriks Q×K dan menyimpannya di memori. Ini dengan cepat melebihi ukuran memori di dalam chip, sehingga kinerja terhambat oleh pembacaan dari luar chip.
FlashAttention mengatasi masalah ini dengan membagi matriks attention menjadi ubin-ubin kecil. Alih-alih mematerialisasi Q×K, ia mengalirkan ubin kecil dari Q dan K ke dalam memori bersama, menghitung softmax pada ubin tersebut, dan hanya menyimpan nilai maksimum dan jumlah yang berjalan untuk setiap baris query. Ini secara signifikan mengurangi lalu lintas memori, sehingga kernel menjadi terbatas oleh komputasi.
Evolusi FlashAttention Sekilas
Sangat berguna untuk memahami bagaimana FlashAttention telah berkembang dari waktu ke waktu. Setiap versi memanfaatkan fitur perangkat keras baru untuk meningkatkan kecepatan atau efisiensi. FA4 secara khusus dioptimalkan untuk GPU generasi Blackwell. Berikut adalah tabel setiap versi FlashAttention beserta peningkatan utama yang dibawanya, dan dampak kinerjanya:
| Versi (Tahun) | Arsitektur GPU Target | Inovasi Kunci | Sorotan Kinerja |
|---|---|---|---|
| FlashAttention (v1) – 2022 | Ampere (NVIDIA A100) dan sebelumnya | Perhatian tepat yang sadar IO (tiling + penyimpanan sementara di chip) untuk mengurangi lalu lintas memori di luar chip. | 2–4× lebih cepat dibandingkan perhatian PyTorch dasar; hingga 10–20× penghematan memori dibandingkan implementasi naif. |
| FlashAttention-2 – 2023 | Ampere (A100), Ada (RTX 30/40) | Paralelisme & pembagian pekerjaan yang lebih baik (okupansi lebih tinggi), penjadwalan warp yang ditingkatkan. Dukungan kepala kueri multi-/tergrup (MQA/GQA). | Sekitar ~2× lebih cepat daripada v1; mencapai ~50–73% dari FLOPs teoretis A100 (~225 TFLOPs/s). |
| FlashAttention-3 – 2024 | Hopper (H100/H800) | Memanfaatkan komputasi asinkron & perpindahan data (Tensor Memory Accelerator, spesialisasi warp). Pipeline GEMM/softmax saling terkait. Dukungan presisi rendah FP8 dengan perilaku numerik yang ditingkatkan. | 1,5–2× lebih cepat dibandingkan v2 di H100; ~740 TFLOPs/s FP16 (~75% pemanfaatan) dan hingga ~1,2 PFLOPs/s di FP8. |
| FlashAttention-4 – 2025 | Blackwell (misalnya, B200, SM 10.x) | Spesialisasi pipeline lanjutan untuk concurrency Blackwell. Aproksimasi perangkat lunak (exp) dan softmax online yang dioptimalkan. Arsitektur kernel disesuaikan untuk Blackwell (melalui CUDA/CUTLASS/DSL). | ~20–22% lebih cepat daripada perhatian cuDNN di Blackwell dalam tolok ukur; peningkatan kinerja yang dilaporkan. |
Apa yang Baru di FlashAttention 4
FlashAttention 4 melanjutkan paradigma komputasi berbasis ubin yang diperkenalkan oleh rilis FlashAttention sebelumnya: tensor Q/K/V diproses ubin demi ubin (blok demi blok) untuk memaksimalkan lokalitas data di dalam chip dan meminimalkan lalu lintas ke memori global. Setiap instance kernel (disebut array thread kooperatif, atau CTA) bekerja sama untuk menghitung satu atau lebih ubin output dengan menyalurkan ubin yang sesuai dari Q, K, dan V dari memori global, kemudian memproyeksikannya ke output perhatian akhir.
Bagian internal FA4 terdiri dari jadwal yang sangat ter-pipeline yang menumpuk beberapa tahap perhitungan attention untuk meningkatkan penyembunyian latensi dan memaksimalkan pemanfaatan sumber daya GPU. Berbeda dengan FlashAttention 3, yang pada dasarnya memiliki tumpang tindih load/compute 2 tahap, FA4 membagi perhitungan menjadi beberapa set tahap yang berjalan secara bersamaan, dengan bagian-bagian berbeda dari pipeline ditangani oleh kelompok warp yang terpisah. Setiap kelompok warp berspesialisasi pada satu tanggung jawab selama "siklus hidup" sebuah tile melalui fungsi attention.
Model spesialisasi warp ini menugaskan warp 32-thread ke tugas-tugas tertentu, seperti:
- Pergerakan data (loads): mengambil ubin query, key, dan value dari memori global ke penyimpanan di chip yang lebih cepat.
- Tahapan komputasi: Melakukan skor perhatian parsial menggunakan perangkat keras matriks kali-akumulasi dengan tensor core.
- Logika Softmax dan normalisasi: Operasi eksponensial + normalisasi untuk bobot softmax perhatian.
- Logika penskalaan ulang dan pengurangan: menerapkan penyesuaian numerik atau menggabungkan hasil parsial sebelum akumulasi.
- Epik/operasi penyimpanan: menulis kembali ubin output tensor yang telah selesai ke memori global.
Intuisi di balik model eksekusi dasar adalah untuk mempertahankan jalur asinkron untuk setiap kelompok warp, sehingga penjadwal warp GPU dapat dengan cepat beralih antara kelompok warp ketika operand sumber tersedia.
Perangkat Lunak Eksponensial melalui CUDA Cores
Salah satu trik baru yang sangat mencolok di FA4 adalah cara menghitung fungsi eksponensial yang digunakan dalam softmax. Biasanya, kernel memanfaatkan SFU khusus di GPU untuk melakukan operasi matematika mahal seperti exp(). Namun, karena jumlah SFU per SM relatif sedikit, mereka dapat dengan mudah menjadi hambatan dalam antrean.
FA4 menghindari keterbatasan ini dengan membiarkan inti CUDA normal melakukan eksponensial perkiraan dalam perangkat lunak. Secara spesifik, FA4 mengandung implementasi khusus dari aproksimasi polinomial exp2(x) dengan presisi perangkat keras. FA4 meniru perhitungan eksponensial menggunakan polinomial kubik, yang kemudian dapat dieksekusi secara paralel di banyak inti CUDA daripada terkonsentrasi pada beberapa SFU.
Intinya: tidak perlu lagi menunggu SFU untuk melakukan softmax. Dengan menangani exp() pada core tujuan umum, FA4 memungkinkan warp yang menghitung softmax untuk menjaga kecepatannya lebih baik dengan alur pemrosesan.
Penyesuaian Skala Softmax Online yang Lebih Cerdas (Baru di FA4)
Inovasi lain dalam FA4 terkait dengan bagaimana akumulasi softmax mempertahankan stabilitas numerik. Seperti semua versi FlashAttention yang melakukan softmax secara "online" – mengakumulasi nilai max() dan jumlah sebagian saat mereka melewati urutan (untuk mencegah overflow). Versi sebelumnya terus-menerus memperbarui faktor skala berjalan (biasanya logit maksimum sejauh ini). Mereka akan selalu melakukan reskala – kapan pun mereka melihat nilai maksimum baru, mereka akan melakukan reskala untuk menjaga stabilitas softmax. FA4 memperkenalkan kebijakan reskala adaptif. Ini hanya melakukan reskala jika maksimum baru yang Anda lihat secara signifikan lebih besar dari yang lama. Hal ini dilaporkan mengurangi frekuensi operasi reskala hingga 10 kali lipat. Mari kita pertimbangkan diagram berikut:
Diagram ini membandingkan dua versi penskalaan ulang yang digunakan dalam FlashAttention, memungkinkan perhitungan perhatian yang stabil dan efisien dalam penggunaan memori. Strategi asli menskalakan ulang output setiap blok jika muncul nilai maksimum, memastikan jumlah kecil tetap akurat dengan menormalkan kontribusi langkah demi langkah. Dengan penskalaan adaptif, kita melacak maksimum lokal dan global, melakukan penskalaan ulang hanya ketika terjadi lonjakan atau ukuran blok membesar. Ini mengurangi beban sambil tetap mencegah overflow dan memungkinkan penjumlahan nilai kecil yang stabil secara numerik.
Kompatibilitas dan status saat ini
Tabel ini memberikan gambaran tentang kompatibilitas yang diharapkan dari FlashAttention 4 (FA4) dan 'apa yang bekerja hari ini' di seluruh stack:
| Aspek | Status Saat Ini (Apa yang Berfungsi Hari Ini) | Catatan/Sumber |
|---|---|---|
| Target Perangkat Keras | Kernel FA4 dirancang untuk GPU Blackwell (SM10.x/B200) | FA4 dioptimalkan untuk GPU generasi Blackwell. FA3 dan versi sebelumnya mendukung arsitektur yang lebih lama; dukungan untuk Blackwell sedang ditambahkan. (Modal) |
| Dukungan Blackwell dalam Rilis Resmi | Sedang dikerjakan | Dukungan Blackwell belum digabungkan ke versi yang dirilis; pengguna sedang memantau dukungan untuk arsitektur sm_120. |
| Lolos Maju (FA4) | Fungsional/Tersedia | Kernel perhatian maju FA4 telah dikomit dan dapat dijalankan di Blackwell dengan build yang sesuai. |
| Lompatan Mundur (FA4) | Belum Lengkap/Terbatas | Proses backward untuk Blackwell kehilangan fitur penting seperti dukungan varlen & GQA. |
| Dukungan Urutan dengan Panjang Variabel | Tidak sepenuhnya didukung (untuk versi sebelumnya) | Dilaporkan hilang dalam implementasi mundur; sedang dalam pengembangan. |
| GQA/Perhatian Kuery Tergrup | Belum didukung (ke belakang) | Proses mundur tidak mendukung perhatian kueri-terkelompok dan varian serupa. |
| Integrasi Kerangka Kerja | Belum didukung | vLLM saat ini hanya mengenali FA2 atau FA3 dan menghasilkan error untuk FA4. (Forum vLLM) |
| Integrasi Kerangka Kerja (PyTorch SDPA) | Belum termasuk dalam rilis stabil | Backend attention produk dot-skala PyTorch (SDPA) belum mendukung FA4. |
| Persyaratan Toolkit/Driver CUDA | CUDA 12.8+ umumnya digunakan untuk build Blackwell | Build Blackwell/SM10 biasanya dikompilasi dengan toolkit CUDA terbaru (>=12.8), meskipun persyaratan spesifiknya berkembang seiring perkembangan. |
Forward pass FA4 sudah tersedia luas (dikontribusikan) di pohon sumber publik. Ada beberapa posting terbuka di proyek seperti repo FlashAttention milik Lab di mana pengguna meminta dukungan penuh untuk backward pass/varlen/group/GQA. Ini menunjukkan bahwa ini masih menjadi isu publik yang terbuka.
Dukungan urutan panjang-variabel (pengelompokan urutan dengan panjang berbeda secara efisien tanpa padding) dan GQA (Grouped-Query Attention) belum didukung dalam proses backward. GQA adalah trik perhatian di mana beberapa kepala perhatian menggunakan proyeksi key/value yang sama (kadang digunakan dalam LLM untuk mengurangi penggunaan memori). Versi sebelumnya dari FlashAttention memang mendukung konfigurasi gaya MQA/GQA. Namun, kernel forward utama FA4 kemungkinan akan mengasumsikan jumlah kepala yang sama digunakan untuk setiap query. Jadi kemungkinan besar ini belum mendukung grouped heads dengan benar. Model dengan grouped-query attention (seperti beberapa varian Llama2) tidak akan didukung oleh FA4 hingga fitur ini diimplementasikan.
Integrasi ekosistem: FA4 belum banyak diadopsi oleh banyak framework/lapisan tingkat tinggi karena masih baru. Salah satu contohnya adalah mesin penyajian vLLM (sangat dioptimalkan untuk inferensi LLM), yang saat ini belum mendukung FA4 – hanya FlashAttention 2 dan 3 yang didukung. Memilih FA4 di vLLM akan menghasilkan kesalahan “versi tidak didukung” karena pemilih versi mereka saat ini hanya memungkinkan Anda memilih antara versi 2 atau 3.
Cara Mengadopsi FlashAttention – Panduan Keputusan
Gunakan tabel ini untuk menentukan versi FlashAttention yang paling sesuai dengan kebutuhan Anda.
| Anda menggunakan Ampere/Ada (A100, RTX 30/40, dll.) | Gunakan FlashAttention-2 (FA2) | FA4 hanya untuk Blackwell (SM10.0) dan tidak akan berjalan/mengompilasi pada SM8.x. FA2 adalah jalur yang paling didukung pada Ampere/Ada dan banyak terintegrasi dalam tumpukan modern. |
|---|---|---|
| Anda berada di Hopper (H100/H200) | Gunakan FlashAttention-3 (FA3) | FA3 dirancang untuk Hopper dan biasanya memberikan peningkatan yang signifikan (sering disebut 1,5–2× dibandingkan FA2) dengan cakupan fitur yang matang (mundur, varlen, dukungan GQA/MQA dalam implementasi umum). FA4 tidak ditujukan untuk Hopper. |
| Anda tidak memiliki perangkat keras Blackwell (tidak ada B200/B100) | Jangan gunakan FA4 | Tidak ada manfaat untuk mencoba FA4: itu dirancang untuk SM10.0 dan tidak akan bisa dikompilasi/dijalankan dengan benar di SM8.x/9.x. Pilih versi FlashAttention terbaru yang didukung oleh GPU Anda sebagai gantinya (FA2 untuk Ampere/Ada, FA3 untuk Hopper). |
| Anda memiliki Blackwell (B200/B100) dan beban kerja Anda hanya untuk inferensi | Uji FA4 (maju) terlebih dahulu, dengan cadangan | FA4 menggunakan metode forward-first hari ini, jadi inferensi adalah yang paling sesuai. Harapkan manfaat terbesar pada urutan panjang dan pola perhatian standar. Tetap aktifkan fallback cuDNN/SDPA untuk menghindari kegagalan keras atau regresi jika suatu fitur hilang. |
| Anda memiliki Blackwell dan beban kerja Anda adalah pelatihan (membutuhkan peninjauan kembali) | Lebih suka jalur pelatihan matang sampai FA4 mundur ada | Pelatihan memerlukan proses mundur. Dengan FA4 yang hanya maju, proses mundur Anda kemungkinan besar akan kembali menggunakan kernel lain, yang bisa mendominasi waktu langkah. Pilihan: (a) latih di Hopper dengan FA3 jika tersedia, (b) terima proses mundur yang lebih lambat di Blackwell untuk saat ini, (c) tunggu dukungan proses mundur FA4. |
Model Anda membutuhkan batch variabel (batch tak rata melalui cu_seqlens) | Hindari FA4 kecuali Anda bisa menambahkan padding atau menerima fallback | Jika FA4 tidak mendukung varlen di stack Anda, Anda akan melihat kesalahan atau fallback diam-diam. Solusinya: tambahkan padding agar panjangnya seragam, kelompokkan berdasarkan panjang, atau nonaktifkan FA4 untuk kasus varlen. |
| Model Anda menggunakan GQA/MQA (kepala kueri yang dikelompokkan, varian Llama, dll.) | Validasi dengan hati-hati; siapkan diri untuk menonaktifkan FA4 | Jika GQA tidak didukung di FA4 untuk jalur Anda (terutama mundur, dan terkadang maju), Anda mungkin mengalami perilaku yang salah atau fallback. Gunakan backend yang sudah terbukti baik (FA3/cuDNN/SDPA) sampai GQA dikonfirmasi di lingkungan Anda. |
| Anda menginginkan FP4 atau fitur presisi sangat rendah | Jangan menganggap FA4 sudah mengaktifkan ini | Adopsi FA4 sebaiknya dimulai dengan BF16/FP16. Penggunaan FP4 tergantung pada dukungan kernel + framework dan bukan sesuatu yang dapat direncanakan untuk penerapan FA4 awal kecuali secara eksplisit didokumentasikan dalam toolchain Anda. |
| Anda membutuhkan stabilitas maksimum dan perilaku yang telah terbukti dalam pertempuran | Pilih opsi paling matang untuk GPU Anda (sering FA3 pada Hopper, FA2 pada Ampere, cadangan cuDNN/SDPA pada Blackwell) | FA4 baru dan terus berkembang; rilis awal mungkin memiliki fitur yang hilang dan bug pada kasus tertentu. Jika risiko produksi tidak dapat diterima, prioritaskan kernel yang sudah matang dan hanya perkenalkan FA4 di balik fitur flag dengan pemeriksaan ketelitian yang ketat. |
| Anda dapat bereksperimen dan mengoptimalkan secara agresif | Mengadopsi FA4 secara bertahap di Blackwell | Pendekatan terbaik: aktifkan FA4 untuk potongan sempit (inference, bucket urutan panjang, non-GQA, panjang tetap), ukur throughput/latensi, verifikasi numerik, kemudian luaskan cakupan seiring dukungan meningkat. |
Pertanyaan yang Sering Diajukan
Masalah apa yang coba dipecahkan oleh FlashAttention 4 (FA4)?
Ini mengurangi lalu lintas memori HBM yang berat pada perhatian dengan menghitung perhatian tile-per-tile di chip, menghindari pembentukan matriks QK⊤ penuh.
Apa perubahan arsitektur utama di FA4 dibanding FA3?
FA4 menggunakan pipeline khusus warp dengan kira-kira 5 tahap (load, compute, softmax/normalize, rescale/reduce, store) untuk meningkatkan overlap dan penggunaan ulang di chip pada Blackwell (SM10.x).
Mengapa FA4 menghitung eksponensial di inti CUDA alih-alih SFU?
Karena SFU langka dan mudah menjadi hambatan saat menghitung softmax, FA4 beralih untuk menghitung perkiraan perangkat lunak dari exp2() di inti CUDA sebagai gantinya.
Apa yang diubah oleh “adaptive online rescaling” pada softmax?
Alih-alih mengubah skala setiap kali muncul nilai maksimum baru, FA4 hanya mengubah skala ketika nilai maksimum tersebut meningkat secara signifikan. Ini mengurangi beban pengubahan skala sambil menjaga stabilitas numerik.
Kapan seorang insinyur sebaiknya mengadopsi FA4 dalam produksi saat ini?
Mulai dengan Blackwell hanya inferensi, panjang tetap, perhatian standar di balik fitur flag; pertahankan cuDNN/SDPA (dan FA3/FA2 di GPU lama) sebagai cadangan sampai dukungan backward/varlen/GQA matang.
Kesimpulan
FlashAttention 4 membangun pada tujuan FlashAttention untuk mengurangi permintaan bandwidth memori dari attention dengan melakukan lebih banyak pekerjaan Q/K/V di chip. Namun, berbeda dengan versi sebelumnya, FA4 melakukan ini dengan pipeline 5-tahap khusus warp untuk Blackwell dan dua perbaikan softmax utama (menggunakan eksponensial perangkat lunak pada inti CUDA untuk mengurangi kemacetan SFU; dan menggunakan penyesuaian ulang adaptif secara online untuk menghilangkan perhitungan stabilitas yang berlebihan). FA4 telah menunjukkan peningkatan throughput pass-maju yang cukup baik dibandingkan cuDNN pada panjang urutan yang panjang, tetapi tetap bersifat forward-first dan hanya untuk Blackwell. Selain itu, FA4 saat ini memiliki banyak kekurangan fitur (tidak ada backward, varlen, GQA, dll.), dan dukungan ekosistem belum sepenuhnya matang. Kesimpulan untuk saat ini adalah mempertimbangkan jalur FA4 untuk inferensi Blackwell sebagai kemungkinan jalur dengan potensi keuntungan tinggi yang membutuhkan adopsi melalui benchmarking yang ketat, keakuratan, dan cadangan. Pada saat yang sama, beban kerja pelatihan dan perangkat keras yang lebih luas sebaiknya tetap menggunakan jalur kode FA3/FA2 yang matang untuk produksi.
Referensi
- FlashAttention-3: Perhatian Cepat dan Akurat dengan Asinkron dan Presisi Rendah
- FlashAttention 4: Menembus Batas Petaflop pada Kernel Perhatian GPU
- Bagaimana cara menerapkan FA4 pada B200?
- Dukungan untuk proses balik FA4 pada Blackwell
Terima kasih telah belajar bersama Komunitas Linux-Console.net.