FlashAttention: GPU'yu Akıllı Kullanan Attention

Bir algoritma ders kitabını yenilemedi — bir donanım sezgisi yendi

2022 başında Stanford'da doktora öğrencisi Tri Dao bir tez yazdı. Konu: attention. Yaptığı şey teorik olarak yenilik değildi — attention'ın çıkışı bit bit aynıydı. Ama çalışma süresi 3–5 kat daha hızlı, bellek kullanımı 10 kat daha az.

Nasıl? Sinir ağına bakmamıştı. Donanıma bakmıştı.

GPU'nun gizli hiyerarşisi

Modern GPU'lar iki tip belleğe sahip:

• HBM (High Bandwidth Memory): Büyük (A100'de 40 GB), ama yavaş — bant genişliği 1.5 TB/s.
• SRAM: Çok küçük (her thread block için 100-200 KB), ama 100x hızlı — 19 TB/s.

CPU'ya benzer: RAM (büyük yavaş) vs. L1 önbellek (küçük hızlı).

Standart attention bu hiyerarşiyi görmez. Q, K, V matrislerini HBM'de tutar, attention matrisini de HBM'de inşa eder. $n$ token için $n \times n$ matris HBM'e yazılır ve okunur — toplam $O(n^2)$ HBM trafiği. GPU çoğu zamanı HBM'i beklemek için harcar.

Standart attention'ın matematiği

$$\text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d}}\right) V$$

Bu işi naif yapmak:

1. $S = QK^T$ hesapla, HBM'e yaz. ($n \times n$)
2. $S$ üzerinde softmax yap, HBM'e yaz.
3. $S \cdot V$ çarp, HBM'e yaz.

$n=8192$ için $S$ matrisi 256 MB; HBM bant genişliğinin sınırını zorlar.

FlashAttention'ın fikri

Dao'nun çözümü iki parçalı:

1. Tiling (Bloklara bölme)

$Q, K, V$ matrislerini küçük bloklara böl (mesela $128 \times 128$). Her bir blok çiftini SRAM'a yükle, ilgili attention parçasını SRAM içinde hesapla, sonucu HBM'e yaz.

Böylece büyük $n \times n$ attention matrisini hiç inşa etmiyorsun. Sadece sonuç $O(n \cdot d)$ HBM'e yazılıyor.

2. Online softmax

Asıl matematiksel zorluk: softmax tüm satıra ihtiyaç duyar — payda toplamı bilinmeli. Blok blok hesaplarken bunu nasıl yaparsın?

Online softmax algoritması: her blokta yerel maksimumu ve toplamı tut, yeni blok geldiğinde yerel değerleri ölçeklendirip birleştir. Numerik kararlılığı korur.

Pseudo-kod:

1. Her blok için max ve sum güncellemesi.
2. Önceki çıkışı yeni max'a göre yeniden ölçeklendir.
3. Yeni katkıyı ekle.

Bu trick olmasa tiling işe yaramazdı.

Performans

A100 GPU üzerinde 8K token için:

• Standart attention: ~2 saniye, 40 GB bellek
• FlashAttention: ~0.4 saniye, 4 GB bellek

A100 maksimum 60% kullanımdan 75%'e çıkar. HBM'e yazma 9 kat azalır. Çıktı bit bit aynıdır — sayısal olarak özdeş, yaklaşım değil.

Etki

Bu algoritma 2022 sonunda popüler olunca:

• PyTorch 2.0 içine entegre edildi (scaled_dot_product_attention).
• HuggingFace Transformers her modelde varsayılan oldu.
• GPT-4, Claude, Llama 2/3 dahil tüm büyük modellerin eğitiminde kullanıldı.

Tri Dao 2023'te FlashAttention-2'yi yayımladı: %2x daha hızlı (paralelliği iyileştirme). 2024'te FlashAttention-3 geldi: Hopper GPU (H100) için %1.5x ek hız.

"Akademik makale yerine kernel"

FlashAttention'ın anlamı sadece bir hızlanma değil. ML araştırmacılarına bir ders verdi: Donanım hiyerarşisini bilmek, model mimarisini bilmek kadar önemli.

Bu fark için yeni bir terim çıktı: IO-aware algoritma. Sadece FLOP'ları değil, bellek hareketini de hesaba kat.

Bu sezgiyi takiben pek çok yeni teknik geldi:

• PagedAttention (vLLM): KV-cache'i sayfa sayfa yönet.
• RingAttention: Dağıtık uzun bağlam.
• FlashDecoding: Üretim sırasında ek hız.

Klasik benzetme

FlashAttention'ı şöyle düşün: Bir bibliyografya yazıyorsun. Standart attention diyor ki "tüm kitapları tek bir masaya yığ, sonra tek tek bak." FlashAttention diyor ki "bir kasa kitap masaya çek, oradakileri işle, sonra kaldır, başka kasa çek." Aynı bibliyografya çıkar, ama masan asla taşmaz, sürekli kütüphane raflarına gitmezsin.

Tri Dao'nun ikinci darbesi

Tri Dao FlashAttention'dan sonra Mamba (2023) ile birlikte yine donanım dostu algoritma tasarladı — paralel scan. Bir araştırmacının iki yıl içinde modern AI'ın hem en hızlı transformer'ını hem en hızlı transformer-alternatifini yazması nadir.

Bağlam

Transformer için: [[transformer-attention-is-all-you-need]]. Donanım optimizasyonunun başka örneği: [[mamba-transformer-a-meydan-okuyan-uzun-bagimli-model]]. Modern GPU mimarisi için: [[gpu-paralel-hesaplama-mimarisi]].