FlashAttention-3: Attention Lebih Cepat dan Akurat dengan Asinkroni dan Presisi Rendah

• Bottleneck Attention pada Transformer selama ini membatasi kecepatan pelatihan dan inferensi LLM dengan konteks panjang, dan FlashAttention-3 adalah versi baru yang berupaya menguranginya lebih jauh dengan memanfaatkan fitur GPU Hopper
• Intinya adalah memanfaatkan asinkroni Tensor Core dan TMA untuk menumpangtindihkan komputasi dengan perpindahan data, serta menempatkan matmul dan softmax secara berselang-seling untuk mengurangi waktu GPU menganggur
• Di H100, FlashAttention-2 hanya memanfaatkan 35% dari FLOPS maksimum teoretis, tetapi FlashAttention-3 mencapai hingga 740 TFLOPS pada FP16, atau 75% dari batas teoretis H100
• Pada presisi rendah FP8, outlier pada activation LLM memperbesar error, sehingga incoherent processing berbasis Hadamard transform menurunkan error 2,6 kali dibandingkan attention FP8 acuan
• FlashAttention-3 telah dirilis di GitHub, dan semakin rendah biaya Attention, semakin besar peluang untuk memproses konteks yang lebih panjang serta meningkatkan efisiensi pelatihan dan inferensi LLM

---

Tujuan dan Performa FlashAttention-3

• Attention adalah layer inti Transformer, tetapi menjadi bottleneck utama pada model bahasa besar dan aplikasi berkonteks panjang
• FlashAttention dan FlashAttention-2 memelopori pendekatan untuk mempercepat Attention dengan mengurangi baca/tulis memori GPU, dan kini digunakan oleh sebagian besar library untuk mengakselerasi pelatihan dan inferensi Transformer
• Pendekatan ini berkontribusi pada peningkatan panjang konteks LLM dalam 2 tahun terakhir, dari 2–4K pada GPT-3 dan OPT menjadi 128K pada GPT-4, hingga 1M pada Llama 3
• FlashAttention-2 hanya mencapai pemanfaatan 35% dari FLOPS maksimum teoretis pada GPU H100, tetapi FlashAttention-3 meningkatkannya dengan memanfaatkan fitur baru GPU Hopper
• Performa FlashAttention-3 yang dipublikasikan adalah sebagai berikut
  • 1,5–2,0 kali lebih cepat daripada FlashAttention-2 pada FP16
  • Hingga 740 TFLOPS pada FP16
  • Pemanfaatan 75% dari FLOPS maksimum teoretis H100
  • Hampir 1,2 PFLOPS pada FP8
  • Error 2,6 kali lebih kecil dibandingkan attention FP8 acuan

Mengulas Kembali Cara Kerja FlashAttention

• FlashAttention mempercepat proses dengan menata ulang urutan komputasi Attention serta menggunakan tiling dan rekomputasi, sekaligus menurunkan penggunaan memori terhadap sequence length dari kuadratik menjadi linear
• Blok input dimuat dari HBM ke SRAM, Attention dijalankan pada blok tersebut, lalu output diperbarui ke HBM
• Karena matriks Attention perantara yang besar tidak ditulis ke HBM, baca/tulis memori berkurang, dan pada waktu eksekusi nyata dapat diperoleh peningkatan kecepatan 2–4 kali
• Dengan menggabungkan tiling dan softmax rescaling, output yang benar dapat diperoleh tanpa aproksimasi meskipun diproses per blok

Fitur GPU Hopper: WGMMA, TMA, FP8

• FlashAttention-2 dapat mencapai hingga 70% dari FLOPS maksimum teoretis pada GPU Ampere A100, tetapi belum memanfaatkan fitur baru GPU Hopper secara memadai
• FlashAttention-3 menggunakan tiga fitur Hopper
  • WGMMA: fitur warpgroup matrix multiply-accumulate yang memanfaatkan Tensor Core baru Hopper, dengan throughput lebih tinggi daripada mma.sync pada Ampere
  • TMA: unit hardware khusus yang mengakselerasi transfer data antara global memory dan shared memory, menangani perhitungan index dan out-of-bound predication sehingga mengurangi penggunaan register
  • FP8: dapat menggandakan throughput Tensor Core dibandingkan FP16, tetapi ada tradeoff dengan akurasi karena nilai floating point direpresentasikan dengan bit yang lebih sedikit
• FlashAttention-3 memanfaatkan fitur Hopper menggunakan abstraksi dari NVIDIA CUTLASS
• Hanya dengan menulis ulang FlashAttention agar memakai fitur baru ini, performa forward pass FP16 meningkat dari sekitar 350 TFLOPS pada FlashAttention-2 menjadi sekitar 540–570 TFLOPS

Menumpangtindihkan GEMM dan softmax dengan Asinkroni

• Operasi utama Attention adalah GEMM antara Q-K dan P-V, serta softmax
• Pada akselerator modern, operasi non-matmul jauh lebih lambat daripada matmul, dan special function seperti exponential pada softmax diproses oleh unit terpisah dari floating point multiply-add atau matrix multiply-add
• H100 SXM5 menyediakan 989 TFLOPS untuk matrix multiply FP16, tetapi throughput special function hanya 3,9 TFLOPS, 256 kali lebih rendah
• Pada head dimension 128, meskipun FLOPS matmul 512 kali lebih banyak daripada exponential, exponential dapat menghabiskan 50% waktu dibandingkan matmul
• Pada FP8, FLOPS matmul menjadi dua kali lebih cepat tetapi kecepatan exponential tetap sama, sehingga menjalankan matmul dan softmax secara paralel menjadi lebih penting

• Pingpong scheduling antar-warpgroup

  • GPU warp scheduler secara otomatis melakukan sebagian overlap dengan menjalankan warp lain saat sebagian warp menunggu hasil GEMM
  • FlashAttention-3 menggunakan synchronization barrier untuk menumpangtindihkan GEMM dan softmax dari dua warpgroup secara manual dengan lebih baik
  • Warpgroup 1 terlebih dahulu menjalankan GEMM1 pada satu iteration dan GEMM0 pada iteration berikutnya
  • Setelah itu, saat warpgroup 2 menjalankan GEMM, warpgroup 1 memproses softmax
  • Pingpong schedule ini menyembunyikan softmax di balik waktu eksekusi GEMM dari warpgroup lain
  • Penjadwalan aktual tidak sepenuhnya rapi seperti diagram, tetapi pada FP16 attention forward pass dengan head dimension 128 dan sequence length 8K, pendekatan ini meningkatkan sekitar 570 TFLOPS menjadi 620 TFLOPS

• Overlap di dalam warpgroup

  • Bahkan di dalam satu warpgroup, sebagian softmax dapat dijalankan saat GEMM milik warpgroup tersebut sedang dieksekusi
  • Pipelining ini meningkatkan throughput FP16 attention forward dari sekitar 620 TFLOPS menjadi 640–660 TFLOPS
  • Sebagai gantinya, register pressure meningkat karena GEMM accumulator serta input/output softmax harus disimpan bersamaan
  • Secara keseluruhan, teknik ini memberikan tradeoff yang menguntungkan

Presisi Rendah FP8 dan Incoherent Processing

• Activation LLM dapat memiliki outlier dengan magnitude jauh lebih besar daripada feature lainnya
• Outlier membuat kuantisasi lebih sulit dan memperbesar quantization error
• FlashAttention-3 memanfaatkan incoherent processing yang digunakan dalam literatur kuantisasi seperti QuIP
• Dengan mengalikan query dan key dengan random orthogonal matrix, outlier disebarkan dan error kuantisasi berkurang
• Implementasinya menggunakan Hadamard transform dengan random sign
  • Jika head dimension disebut d, ini dapat dilakukan per attention head dalam waktu O(d log d), bukan O(d²)
  • Karena Hadamard transform bersifat memory-bandwidth bound, proses ini dapat digabungkan dengan operasi sebelumnya yang juga memory-bandwidth bound, seperti rotary embedding, sehingga dapat diproses tanpa biaya tambahan
• Dalam eksperimen yang membuat Q, K, V dari distribusi normal standar dan menyimulasikan outlier dengan memasukkan magnitude besar ke 0,1% entry, incoherent processing menurunkan error kuantisasi sebesar 2,6 kali

Benchmark dan Status Rilis

• FlashAttention-3 dibandingkan bukan hanya dengan FlashAttention-2, tetapi juga dengan implementasi Triton dan cuDNN yang sudah menggunakan fitur hardware baru GPU Hopper
• Pada FP16, FlashAttention-3 menunjukkan peningkatan kecepatan sekitar 1,6–1,8 kali dibandingkan FlashAttention-2
• Pada FP8, performanya hampir mencapai 1,2 PFLOPS
• Repositori GitHub FlashAttention-3 telah dibuka untuk publik
• Papernya juga dapat dilihat di repositori flash-attention yang sama

Optimasi Tersisa dan Integrasi Mendatang

• Paper ini juga mencakup optimasi selain yang dibahas di blog, seperti variable length sequence, persistent kernel, dan in-kernel transpose untuk FP8
• Merancang algoritme sesuai hardware eksekusi dapat menghasilkan peningkatan efisiensi besar dan membuka kemampuan model baru seperti konteks panjang
• Pekerjaan mendatang mencakup optimasi inference LLM dan generalisasi teknik ini ke arsitektur hardware lain
• FlashAttention-3 diharapkan akan diintegrasikan ke rilis PyTorch mendatang

[GN⁺]

Opini Hacker News

• Jika melihat komentar kode, tampaknya Tri Dao sudah mengerjakan FA3 sejak April 2022, tepat setelah pengumuman Hopper/H100
  Agak menarik bahwa butuh lebih dari 2 tahun sampai kodenya dirilis hari ini; mungkin karena solusi yang lebih baik sedang disiapkan
  Riwayat makalah Tri belakangan ini cenderung mengarah ke struktur keluarga SSM dan Mamba. FlashAttention memiliki kompleksitas waktu kuadratik terhadap panjang sekuens, tetapi algoritma terbaru bersifat subkuadratik, jadi bukan sekadar melakukan komputasi yang sama dengan lebih efisien, melainkan mengurangi jumlah komputasinya sendiri secara jauh lebih besar
  Dalam sebuah makalah panjang tahun ini, Dao dan Gu menunjukkan bahwa Mamba/SSM juga dapat diformulasikan agar mudah diakselerasi dengan operasi primitif hardware yang sama seperti yang menguntungkan Transformer

  • Sampai Strong Exponential Time Hypothesis (SETH) terbukti atau terbantahkan, biaya kuadratik akan diperlukan atau kita harus mengorbankan sesuatu. Pada akhirnya ini adalah biaya pencarian menyeluruh
    Jika SETH dibuktikan atau dibantah, masalah P versus NP juga akan terselesaikan, jadi sulit berharap itu terjadi dalam waktu dekat
    Intinya adalah apakah kasus penggunaan tertentu sanggup menanggung biaya tersebut

• Saya penasaran sejauh mana algoritma FlashAttention terikat pada hardware
  Misalnya, dalam pengumuman kali ini dikatakan memanfaatkan fitur asinkron GPU H100, jadi tampaknya kartu yang bukan seri H tidak akan mendapatkan peningkatan kecepatan itu
  Selain itu, library FlashAttention yang sebenarnya membutuhkan CUDA, tetapi algoritmanya tampaknya sudah di-porting ke Metal[^0]. Jika algoritmanya mendekati fungsi murni, bukankah seharusnya bisa diimplementasikan di GPU/framework machine learning apa pun?
  [0]: https://github.com/philipturner/metal-flash-attention

  • Ada banyak jawaban bagus, tetapi singkatnya, “dalam praktiknya, cukup banyak” terikat pada hardware. Contoh di bawah ini cukup bagus

    https://github.com/karpathy/nanoGPT/blob/master/model.py#L45
    nanoGPT milik Karpathy memeriksa keberadaan torch.nn.functional.scaled_dot_product_attention lalu memanggil FlashAttention
    https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
    Jika melihat dokumentasinya, pada praktiknya orang akan lebih sering ingin memanggil FA2, dan FA2 mengoptimalkan kernel perangkat untuk membagi operasi Softmax pada matriks segitiga, serta mengurangi pengiriman bolak-balik batch floating-point yang tidak perlu antara GPU dan CPU
    https://arxiv.org/pdf/2307.08691
    Makalah FA2 hampir sepenuhnya dijelaskan dari sudut pandang hardware tempat ia berjalan

  • Peningkatan algoritmik FlashAttention terutama berupa membagi dan menggabungkan bagian Softmax dari attention, dan itu sendiri bukanlah ide yang sepenuhnya baru. Kontribusi yang sangat besar ada pada implementasi metode dan detail-detail tersebut secara efisien di hardware Nvidia
  • FlashAttention pada awalnya hampir tidak memiliki ketergantungan pada hardware
    Versi terbaru bergantung pada tingkat abstraksinya. ThunderKittens[0] memberikan peningkatan kecepatan sekitar 1,3–2 kali dibanding FA2, mirip seperti yang disebutkan dalam tulisan, tetapi tetap relatif umum diterapkan di berbagai GPU
    Setiap hardware baru mungkin memiliki fitur khusus hardware yang dapat mengeluarkan performa tambahan. Biasanya vendor mengadopsi fitur yang membuat mereka unggul, tetapi seperti yang sudah terjadi pada CUDA, API dan library menjadi terfragmentasi
    [0]: https://hazyresearch.stanford.edu/blog/2024-05-12-tk
  • Secara konseptual agak terikat; dari sudut pandang implementasi praktik, sangat terikat. Implementasi Python standar pun secara internal mengompilasi kernel yang disesuaikan untuk hardware tertentu
  • Dari sudut pandang praktik, sebagai tambahan, hardware AMD masih kekurangan implementasi flash-attention-2 yang benar-benar baik. ROCm perlahan mulai layak dipakai, tetapi masih belum sampai pada level yang bisa dibandingkan dengan CUDA

• Saya ingin bertanya kepada orang-orang di bidang compiler. Apakah ada kemungkinan compiler bisa menemukan sendiri optimisasi seperti FlashAttention? TVM dan tinygrad tampaknya mengarah ke sana, tetapi sulit dipercaya bahwa itu mungkin

  • Secara teori mungkin. Berkat sifat aljabar dalam matematika, penataan ulang dalam skala besar bisa dilakukan, lalu tinggal ditambah polyhedral loop tiling yang relatif terstruktur
    Namun biayanya besar, jadi hasil pencarian itu harus di-cache
    Optimisasi e-graph tampaknya cocok untuk area ini. Namun dibutuhkan perubahan paradigma besar dalam cara menangani pass optimisasi, sehingga hampir tidak pernah di-deploy selain di beberapa alat niche. Misalnya, ini tidak cocok dengan call graph tradisional, sehingga untuk men-deploy e-graph yang melintasi ke luar/di antara basic block dan loop for, control flow harus banyak diubah, dan break serta return juga tidak didukung
  • Ini tampak sebagai masalah yang sangat sulit, tetapi bukan mustahil
    Saya tidak begitu tahu sejauh mana state-of-the-art optimisasi compiler saat ini dalam hal penempatan data dan memaksimalkan utilisasi prosesor
    Dulu saya pernah melihat video tentang optimisasi yang mengatakan bahwa optimisasi kecil memang meningkatkan kecepatan, tetapi dampaknya kecil dibanding variasi kecepatan yang berasal dari perbedaan layout memori akibat optimisasi itu, atau bahkan akibat perubahan acak
    Presentasi itu lebih berfokus pada membedakan sinyal dari noise, tetapi noise itu sendiri adalah tanda bahwa compiler bahkan belum mampu menangani bentuk yang jauh lebih sederhana daripada masalah yang dibahas di sini
    Struktur CPU dan memori saja sudah kompleks, dengan cache dan pola akses yang memengaruhi kecepatan; jika ditambah struktur GPU, rasanya ini masih menjadi wilayah yang cukup belum banyak dijelajahi
    Suatu saat mungkin saja bisa. Karena ini bidang AI, muncul juga pertanyaan apakah AI yang cukup pintar bisa melakukan ini, tetapi itu tergantung pada standar “cukup”
    Sebagai tes tingkat sangat tinggi untuk model AI, bisa dibayangkan memberi sesuatu seperti micrograd dan memintanya membuat sesuatu yang lebih cepat daripada torch sambil mempertahankan antarmuka yang sama. Kita belum mendekati itu, tetapi kalau suatu saat bisa, itu akan menarik
  • Saya rasa tidak. Ini harus dipikirkan seperti algoritma yang berbeda. Alih-alih hanya mempertimbangkan matematika, algoritmanya dirancang dengan mempertimbangkan bentuk hardware
    TVM masuk akal. Secara ketat, ia melakukan hal yang berbeda, tetapi areanya cukup dekat
    Namun saya tidak tahu mengapa tinygrad terasa seperti itu
  • https://github.com/uwplse/tensat
  • Memanggil operator tingkat tinggi dari bahasa pembungkus seperti Python cukup merepotkan

• Jika ada yang ingin mem-porting ini ke ROCm / AMD MI300x, silakan hubungi [e-mail]. Saya sama sekali tidak akan mengirim spam
  Saya bisa menyumbangkan waktu komputasi untuk pekerjaan ini

  • Ternyata ini perusahaan server akselerator AMD! Keren, semoga ada yang mengambilnya :)
  • Bukan bermaksud tidak sopan, tetapi saya penasaran dengan maksud tawaran ini. Apakah ada orang yang akan melakukan porting ini gratis hanya dengan mendapat akses hardware? Apa keuntungan bagi orang itu?

• FlashAttention-3 is optimized for Hopper GPUs (e.g. H100).
  Seberapa bagus performa FA3 di GPU konsumen seperti 3090 dan 4090?

  • Ini khusus Hopper. Peningkatannya sangat terkait dengan fitur Hopper seperti warp group dan TMA
    Di 4090, Anda mungkin bisa mendapat peningkatan kecepatan dengan memakai implementasi Triton untuk attention FP8: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html

• TMA (Tensor Memory Accelerator). This is a special hardware unit that accelerates the transfer of data between global memory and shared memory, taking care of all index calculation and out-of-bound predication. This frees up registers, which is a valuable resource to increase tile size and efficiency.
  Sepemahaman saya, memang benar TMA mengurangi penggunaan register, tetapi yang lebih penting adalah ia membuat pembuatan alamat ditangani oleh hardware. Semakin cepat operasi di sekitarnya, address generation bisa menjadi bottleneck

• Ini adalah salah satu peningkatan paling penting di seluruh AI. Dengan hardware yang sama, kita bisa memakai lebih banyak hal dengan lebih cepat, dan bagi sebagian besar pengguna AI ini memberi manfaat hampir tanpa trade-off

  • Ya, untuk pengguna yang punya H100

• Saya penasaran mengapa FlashAttention menjadi sekitar 5x lebih lambat saat memakai masking variabel dibanding saat tidak memakainya. Jika dukungan masking yang baik kurang, efek optimisasinya hampir hilang

  • Anda melihat benchmark itu di mana?

• Akan bagus jika ada pakar yang menjawab beberapa pertanyaan ini :)
  Apakah FlashAttention adalah pengganti drop-in untuk operasi attention di LLM? Apakah bisa dipakai di mana pun operasi “attention” digunakan, atau LLM harus dilatih secara terpisah agar memakai FA?
  Apa hubungan FA dengan strategi seperti GQA dan sliding window attention? Apakah konsepnya saling ortogonal, atau perlu implementasi FA tersendiri untuk tiap strategi?
  Baru-baru ini llama.cpp menambahkan dukungan FlashAttention; apakah ini berarti mereka mulai memakai sesuatu seperti kernel CUDA yang disediakan FlashAttention?
  Terakhir, artikel ini membandingkan FlashAttention dan Triton. Bukankah Triton semacam lapisan abstraksi? Tidak bisakah FA diimplementasikan dengan Triton? Saya kurang paham ungkapan “FlashAttention vs Triton”.

  • 1. Hampir benar. Secara matematis ekuivalen. Masalah di sisi software hanya hal-hal seperti pengelolaan versi dependensi atau format data di memori, dan FlashAttention 2 sudah masuk ke HuggingFace serta beberapa library populer. FlashAttention 3 juga kemungkinan besar segera masuk, tetapi untuk menjalankannya diperlukan GPU H100.
    2. FlashAttention 2 menambahkan dukungan GQA pada pembaruan versi sebelumnya:
       https://github.com/Dao-AILab/flash-attention
    3. Di sini yang dibandingkan adalah implementasi FlashAttention yang ditulis dengan CUDA C++ murni ini dengan implementasi Triton untuk algoritma serupa yang ditulis dengan Triton: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
  • FlashAttention dapat menjadi pengganti drop-in untuk operasi attention di LLM.
    FlashAttention adalah metode untuk menghitung bagian Softmax(QK^T)V dari attention, sedangkan GQA adalah metode untuk menghitung matriks Q, K, V. Untuk sliding window attention, saya tidak sepenuhnya yakin, tetapi itu adalah cara mengubah attention mask yang mengontrol query mana bisa melihat key mana.
    Saya belum pernah memakai llama.cpp, tetapi penjelasan bahwa mereka mulai menggunakan kernel CUDA sepertinya secara umum benar.
    Pertanyaan terakhir merujuk pada implementasi FlashAttention yang sebelumnya ditulis dengan Triton.

• Saya penasaran karena artikel ini mengatakan operasi seperti sigmoid sangat lambat.
  LLM modern banyak memakai fungsi aktivasi yang melibatkan sigmoid atau Softmax, seperti SiLU, Swish, dan SOLU.
  Apakah ReLU memiliki penalti performa yang lebih kecil? Jika begitu, mungkinkah lebih baik kembali ke ReLU lama yang bagus?

  • ReLU secara harfiah adalah fungsi linear yang dipotong menjadi 0 pada suatu titik, jadi jumlah komputasinya jauh lebih sedikit daripada yang melibatkan fungsi eksponensial. Namun, tampaknya sulit mendapatkan hasil yang kompetitif dengan fungsi aktivasi sesederhana itu.