ENID
Lensa ML

FlashAttention

Langkah 1 dari 6: Hierarki Memori

HBM besar dan lambat, SRAM kecil dan cepat

Data yang dibaca32 MB
Kapasitas (skala log), GPU ilustratif kelas A100SRAM, satu SM: 192 KiBSRAM, semua 108 SM: 21 MBHBM: 80.0 GB10 KB1.0 MB100 MB10.0 GBdata: 32 MBBandwidth (skala linear)HBM 2 TB/sSRAM sekitar 19 TB/s
Ukuran data
32 MB
Waktu baca dari HBM
15.8 µs
Pada kecepatan SRAM
tidak muat
Rasio bandwidth
9.5×

32 MB hanya muat di HBM. Setiap satu putaran penuh atas data ini menghabiskan 15,8 µs bandwidth, dan ukurannya 1,5× seluruh SRAM chip.

Geser Data yang dibaca dari 10 KB sampai 100 GB dan lihat di mana garis putus-putus memotong tiap tingkat memori. A100 mengerjakan sekitar 312 TFLOP/s matematika matriks fp16, tetapi membaca HBM pada sekitar 2 TB/s, jadi sebuah kernel butuh kira-kira 150 FLOP per byte agar unit aritmetikanya tetap sibuk. Softmax, masking, dan dropout hanya mengerjakan beberapa FLOP per elemen, sehingga dibatasi oleh memori. Sebagai gambaran, matriks skor satu head pada 4.096 token berukuran 33,6 MB dalam fp16, sudah lebih besar daripada seluruh SRAM di dalam chip.

Langkah 2 dari 6: Attention Standar

Matriks N×N bolak-balik ke HBM

Panjang sequence N4,096
Satu head, d = 128, fp16: tiga kernel, masing-masing lewat HBMHBMQ1.0 MBK1.0 MBS34 MBP34 MBV1.0 MBO1.0 MBS = QKᵀP = softmax(S)O = PVkernel berjalan di dalam chip, memakai SRAMLalu lintas HBM per head: 138 MBBaca/tulis S dan P 97%Q, K, V, O 3%
S per head (fp16)
34 MB
S, 32 head, satu layer
1.1 GB
Lalu lintas HBM per head
138 MB
Waktu per layer pada 2 TB/s
2.21 ms

Pada N = 4,096, S milik satu head berukuran 34 MB, lebih besar daripada seluruh SRAM di dalam chip, dan 32 head dalam satu layer butuh 1,1 GB. Memindahkannya mendominasi 2,21 ms yang dihabiskan layer untuk memori.

Geser Panjang sequence N dari 256 sampai 32.768. Implementasi standar menjalankan kernel terpisah, dan masing-masing menulis hasilnya ke HBM untuk dibaca kernel berikutnya. Per head, itu berarti sekitar 4N² + 4Nd elemen yang dipindahkan, sementara aritmetikanya sekitar 4N²d FLOP, yang bisa diselesaikan GPU dengan cepat. Masking dan dropout menambah putaran atas matriks N×N yang sama. Training juga menyimpan P untuk backward pass. Untuk model 7B pada 4.096 token, itu berarti 1,07 GB per layer di ke-32 head-nya.

Langkah 3 dari 6: Tiling

Blok Q, K, dan V yang muat di dalam chip

Ukuran blok64 baris
Matriks skor, 256 × 256 tokenTile K/V →Blok Q ↓abu-abu: dihitung, dipakai, dibuangSRAM satu SMBlok Q (64×128)16 KiBTile K (64×128)16 KiBTile V (64×128)16 KiBTile S (64×64)8.0 KiBBlok O (64×128)16 KiBterpakai 72 KiB dari 192 KiBtile fp16, d = 128; accumulator fp32 dan statistik tidak dihitung
Tile dalam S
4 × 4
Tile saat ini
belum ada
SRAM per langkah tile
72 KiB
S yang ditulis ke HBM
0 B, bukan 131 KB

Matriks skor 256×256 terbagi menjadi 16 tile 64×64. Tekan Putar atau Langkah untuk menjalankan loop: satu blok Q tetap di dalam chip sementara tile K/V mengalir melewatinya.

Atur Ukuran blok, lalu tekan Putar atau Langkah untuk menjalankan loop. Paper-nya memilih tile K/V sebanyak ⌈M/4d⌉ baris dan blok Q sebanyak min(⌈M/4d⌉, d) baris, dengan M sebagai ukuran SRAM. Kernel asli menaruh tile K/V di loop luar. FlashAttention-2 menukar loop-nya, seperti yang ditampilkan di sini, sehingga setiap thread block memegang satu blok Q dan menulis output-nya satu kali. Matmul, mask, softmax, dropout, dan matmul kedua semuanya berjalan dalam satu kernel gabungan. Kernel sungguhan juga butuh accumulator fp32, jadi 64 atau 128 baris adalah ukuran yang umum.

Langkah 4 dari 6: Online Softmax

Maksimum dan jumlah berjalan, di-rescale per chunk

Skor x₅4.5
Ukuran chunk
1.2x13.0x20.4x32.1x44.5x51.8x63.6x70.9x8m = −∞, ℓ = 0Batang: e^(x − m) / ℓ sejauh ini. Garis luar: softmax penuh.0.020.120.010.050.540.040.220.01angka di bawah batang: softmax penuh atas seluruh 8 skor
Maksimum berjalan m
−∞
Jumlah berjalan ℓ
0
Faktor rescale
belum ada
Maks |online − penuh|
0/8 terbaca

Belum ada yang dibaca. Jalankan 4 chunk berisi 2 skor satu per satu, lalu lihat m dan ℓ berubah.

Geser Skor x₅, pilih Ukuran chunk, lalu jalankan chunk satu per satu. Pembaruannya adalah m_new = max(m, chunk max) dan ℓ_new = ℓ · e^(m − m_new) + Σ e^(x − m_new). Mengurangkan nilai maksimum menjaga eksponen tetap kecil: fp16 mengalami overflow di atas 65.504, yaitu e^11.1. Milakov dan Gimelshein (2018) menerbitkan normalizer satu putaran ini untuk softmax. FlashAttention menerapkan faktor rescale yang sama pada output berjalan, yaitu jumlah weight dikali V yang belum dinormalisasi, lalu membaginya dengan ℓ setelah tile K/V terakhir selesai.

Langkah 5 dari 6: Penghematan IO

Akses HBM lebih sedikit, memori linear

Panjang sequence N8,192
SRAM per SM (M)192 KiB
Dimensi head d
Memori tambahan per head (skala log)5122K8K32K128K1.0 KB1.0 MB1.0 GBpanjang sequence N (skala log)standar: N² skorFlashAttention: m dan ℓ, O(N)Akses HBMstandarFlash
Standar, Nd + N²
68.2M
Flash, N²d²/M
11.2M
Akses berkurang
6.1×
Memori tambahan per head
134 MB vs 66 KB

Dengan M = 98,304 elemen dan d = 128, rasio ilustratifnya sekitar M / d² = 6.1× lebih sedikit akses HBM, dengan tile K/V sekitar 192 baris. Memori tambahan turun dari 134 MB ke 66 KB per head pada N = 8,192.

Geser Panjang sequence N dan SRAM per SM, lalu ganti Dimensi head d. Teorema 2 dalam paper FlashAttention memberikan Θ(Nd + N²) akses HBM untuk attention standar dan Θ(N²d²/M) untuk FlashAttention. Keduanya tetap tumbuh seiring N², dan jumlah FLOP tidak turun. Backward pass bahkan menghitung ulang S dari Q dan K, yang menambah aritmetika, tetapi berjalan lebih cepat karena membaca lebih sedikit. Paper itu juga menunjukkan bahwa tidak ada algoritma attention eksak yang bisa lebih baik secara asimtotik pada suatu rentang ukuran SRAM.

Langkah 6 dari 6: Eksak, dan Ada di Mana-Mana

FlashAttention 1, 2, dan 3

Versi
Ukuran blok (cek ketepatan)4 key
202220232024FlashAttention-2, A100 (Ampere)Paralel atas panjang sequence dan juga head;lebih sedikit FLOP non-matmul, pembagian warp lebih baikDilaporkan: sekitar 2× FA1, 50-73% dari puncak FLOP/sOutput query 1, d = 8, Q, K, V acak (ilustratif)garis luar: attention standar; isian: tiled dengan online softmax
Tile K/V
4
Maks |tiled − standar|
2.2e-16
GPU sasaran
A100 (Ampere)

Dengan 4 tile berisi paling banyak 4 key, output tiled berbeda dari attention standar paling banyak 2,2e-16, murni pembulatan floating-point. FlashAttention-2 mengubah cara loop dijadwalkan di A100 (Ampere), bukan hasil ini.

Ganti Versi dan geser Ukuran blok untuk mengulang cek ketepatan. Metode aproksimasi seperti attention sparse atau low-rank mengubah hasil demi menghemat komputasi, sedangkan setiap versi FlashAttention mereproduksi attention standar sampai batas pembulatan. Di PyTorch 2.0 dan yang lebih baru, torch.nn.functional.scaled_dot_product_attention memilih kernel FlashAttention ketika dtype-nya fp16 atau bf16 dan bentuk tensornya memungkinkan, dan beralih ke kernel lain jika tidak. Hugging Face Transformers menyediakannya lewat attn_implementation="flash_attention_2".

Apa itu FlashAttention?

FlashAttention adalah algoritma attention eksak yang menghitung softmax(QKᵀ/√d)V dalam tile kecil yang disimpan di memori GPU di dalam chip, sehingga matriks skor N×N penuh tidak pernah ditulis ke memori lambat di luar chip. Output-nya sama dengan attention standar. Keuntungannya datang dari byte yang dipindahkan lebih sedikit, bukan dari aritmetika yang lebih sedikit.

Attention adalah bagian transformer yang biayanya tumbuh seiring kuadrat panjang sequence, dan di GPU bentuk standarnya dibatasi oleh lalu lintas memori. FlashAttention mempercepat training GPT-2 sekitar 3× pada 1K token dan memangkas memori tambahan attention dari kuadratik menjadi linear terhadap panjang sequence. Kernel-nya kini dipakai di balik scaled_dot_product_attention milik PyTorch, Hugging Face Transformers, vLLM, dan sebagian besar stack training LLM.

Ide utama

  • Attention standar dibatasi oleh lalu lintas memori, karena ia menulis dan membaca matriks skor N×N lewat HBM beberapa kali.
  • FlashAttention memecah Q, K, dan V menjadi tile yang muat di SRAM dalam chip, lalu menggabungkan seluruh attention ke dalam satu kernel.
  • Online softmax menyimpan maksimum dan jumlah berjalan per baris, sehingga softmax bisa dihitung per tile dan di-rescale di akhir.
  • Akses HBM turun dari Θ(Nd + N²) ke Θ(N²d²/M), dan memori tambahan turun dari O(N²) ke O(N).
  • FlashAttention eksak: output-nya sama dengan attention standar, kecuali pembulatan floating-point.
  • FlashAttention-2 memperbaiki paralelisme dan pembagian kerja, dan FlashAttention-3 menyasar GPU Hopper dengan asinkroni dan FP8.

Pertanyaan umum

Apakah model harus di-train ulang untuk memakai FlashAttention?

Tidak. FlashAttention menghitung fungsi yang sama dengan attention standar, jadi weight yang ada tetap berfungsi tanpa perubahan. Mengganti kernel attention adalah pilihan saat runtime, sering kali cukup dengan satu flag. Output-nya bisa berbeda di bit terakhir karena urutan penjumlahannya berbeda.

Kenapa FlashAttention lebih cepat padahal FLOP-nya lebih banyak?

Di GPU modern, attention standar menghabiskan sebagian besar waktunya menunggu memori, bukan mengerjakan aritmetika. FlashAttention menghitung ulang matriks skor di backward pass, bukan menyimpannya. Ini menambah FLOP, tetapi menghapus jauh lebih banyak byte lalu lintas HBM. Karena memori yang menjadi hambatan, total waktunya turun.

Apakah FlashAttention membantu saat decoding LLM?

Prefill, saat seluruh prompt diproses sekaligus, paling diuntungkan karena membangun matriks skor besar. Saat decode, setiap langkah punya satu query per sequence, jadi matriks skornya hanya satu baris dan biaya utamanya adalah membaca KV cache. Varian seperti Flash-Decoding memecah pembacaan itu ke banyak thread block agar seluruh GPU terpakai.

GPU apa saja yang mendukung FlashAttention?

FlashAttention-2 berjalan di GPU NVIDIA Ampere, Ada, dan Hopper seperti A100, RTX 4090, dan H100, dan versi aslinya juga mencakup kartu Turing yang lebih lama. FlashAttention-3 ditulis untuk Hopper. GPU AMD punya port ROCm, dan PyTorch beralih ke kernel lain ketika tidak ada kernel FlashAttention yang berlaku.

Penjelasan terkait

Referensi

  1. Dao et al. (2022), "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"
  2. Dao (2023), "FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning"
  3. Shah et al. (2024), "FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision"
  4. Milakov and Gimelshein (2018), "Online normalizer calculation for softmax"

Terakhir diperbarui