FlashAttention: IO Farkındalığı ile Hızlı ve Bellek Açısından Verimli Tam Dikkat

Okuma rehberi

Problem ve motivasyon

Transformer modelleri, dizi uzunluğu $N$'ye bağlı olarak karesel zaman ve bellek karmaşıklığına sahiptir. Yaklaşık dikkat yöntemleri teorik FLOP sayılarını azaltsa da, modern GPU hesaplamaları aritmetik işlemlerden ziyade çoğunlukla bellek transferleri (IO) ile darboğaza girdiğinden, bu yöntemler genellikle gerçek zamanlı hızlanma elde edemez.

Standart dikkatte $\mathbf{S}=\mathbf{Q}{\mathbf{K}}^{\top }$, $\mathbf{P}=\mathrm{softmax}(\mathbf{S})$ ve $\mathbf{O}=\mathbf{PV}$ hesaplamaları; $N\times N$ boyutundaki ara matrisleri yavaş Yüksek Bant Genişlikli Belleğe (HBM) yazar ve bunları defalarca geri okur. Bu durum $\Theta (Nd+N^2)$ HBM erişimine ve $O(N^2)$ bellek ayak izine neden olarak pratik bağlam pencerelerini sınırlar.

Temel fikirler

• IO Farkındalığı: GPU bellek hiyerarşisini doğrudan hesaba katmak: Hızlı, küçük çip üstü SRAM (örneğin A100 üzerinde ~19 TB/s bant genişliği ile 192 KB/SM) karşısında büyük, yavaş HBM (1,5–2,0 TB/s bant genişliği ile 40–80 GB).
• Döşeme (Blok Softmax Hesaplaması - Tiling): $\mathbf{Q},\mathbf{K},\mathbf{V}$ girdilerini tamamen SRAM içine sığacak bloklara ayırmak. Çalışan softmax satır maksimumu ve normalleştirici istatistiklerini takip ederek, tam $N\times N$ matrisini HBM üzerinde somutlaştırmadan softmax indirgemesi bloklar arasında artımsal olarak hesaplanır.
• Geri Geçişte Yeniden Hesaplama: Geriye yayılım (backpropagation) için ara $N\times N$ dikkat matrisini saklamaktan kaçınmak. Bunun yerine yalnızca çıktı $\mathbf{O}$ ve softmax istatistikleri $(\ell ,m)$ saklanır; geri geçiş sırasında $\mathbf{S}$ ve $\mathbf{P}$ blok blok SRAM çipinde yeniden hesaplanır.
• Çekirdek Kaynaştırma (Kernel Fusion): Matris çarpımları, maskeleme, softmax ve dropout işlemlerini tek bir CUDA çekirdeğinde kaynaştırarak gereksiz HBM gidiş-dönüşlerini ortadan kaldırmak.

Yöntem (temel denklemlerle birlikte)

$\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{N\times d}$ verildiğinde, FlashAttention blok boyutlarını $B_c=\lceil M/(4d)\rceil$ ve $B_r=\min (\lceil M/(4d)\rceil ,d)$ olarak belirler; burada $M$, SRAM kapasitesidir.

Artımsal Softmax Takibi

$x^{(1)}$ ve $x^{(2)}$ alt vektörlerine ayrıştırılmış bir vektör için çalışan istatistikler şu şekilde güncellenir: $$m(x)=\max (m(x^{(1)}),m(x^{(2)})),$$ $$\ell (x)=e^{m(x^{(1)})-m(x)}\ell (x^{(1)})+e^{m(x^{(2)})-m(x)}\ell (x^{(2)}).$$

Anahtarların/değerlerin $j$. bloğu ve sorguların $i$. bloğu için:

1. SRAM üzerinde normalleştirilmemiş skorları hesapla: ${\mathbf{S}}_{ij}={\mathbf{Q}}_i{\mathbf{K}}_j^{\top }\in {\mathbb{R}}^{B_r\times B_c}$.
2. Blok yerel istatistiklerini hesapla: ${\tilde{m}}_{ij}=\mathrm{rowmax}({\mathbf{S}}_{ij})$, ${\tilde{\mathbf{P}}}_{ij}=\exp ({\mathbf{S}}_{ij}-{\tilde{m}}_{ij})$ ve ${\tilde{\ell }}_{ij}=\mathrm{rowsum}({\tilde{\mathbf{P}}}_{ij})$.
3. Çalışan küresel normalleştiricileri güncelle: $$m_i^{\mathrm{new}}=\max (m_i,{\tilde{m}}_{ij}),{\ell }_i^{\mathrm{new}}=e^{m_i-m_i^{\mathrm{new}}}{\ell }_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\ell }}_{ij}.$$
4. Çalışan çıktı bloğunu ${\mathbf{O}}_i\in {\mathbb{R}}^{B_r\times d}$ güncelle: $${\mathbf{O}}_i\leftarrow \mathrm{diag}({\ell }_i^{\mathrm{new}}{)}^{-1}\left(\mathrm{diag}({\ell }_i)e^{m_i-m_i^{\mathrm{new}}}{\mathbf{O}}_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\mathbf{P}}}_{ij}{\mathbf{V}}_j\right).$$

Karmaşıklık Analizi

• IO Karmaşıklığı: Standart dikkat $\Theta (Nd+N^2)$ HBM erişimi gerektirir. FlashAttention ise $\Theta (N^2{d}^2{M}^{-1})$ HBM erişimi gerektirir. Yazarlar, hiçbir tam dikkat algoritmasının tüm $M\in [d,Nd]$ için asimptotik olarak $o(N^2{d}^2{M}^{-1})$ erişimden daha azını elde edemeyeceğini kanıtlamaktadır.
• FLOP ve Bellek: FlashAttention $O(N^{2}d)$ FLOP gerçekleştirir (yeniden hesaplama nedeniyle geri geçişte biraz daha fazla FLOP ile) ve girdiler ile çıktıların ötesinde yalnızca $O(N)$ yardımcı bellek gerektirir.

Genişletme: Blok-Seyrek (Block-Sparse) FlashAttention

Hesaplamayı önceden tanımlanmış bir seyreklik maskesine $\mathbf{M}\in \{0,1{\}}^{N/B_r\times N/B_c}$ göre yalnızca sıfır olmayan bloklarla sınırlandırmak, $\Theta (Nd+N^2{d}^2{M}^{-1}s)$ IO karmaşıklığı sağlar; burada $s$, sıfır olmayan blokların oranıdır.

Sonuçlar (en önemli sayılarla birlikte)

• Eğitim Hızlanması:
  • BERT-large (dizi uz. 512): MLPerf 1.1 eğitim hızı rekorunu %15 geride bırakır ($8\times$ A100 üzerinde 20,0 dakikaya karşılık 17,4 dakika).
  • GPT-2 (dizi uz. 1K): HuggingFace'e kıyasla 3,5 kata, Megatron-LM'e kıyasla 1,8 kata kadar hızlanma sağlar.
  • Long-Range Arena (LRA): Standart dikkate kıyasla 2,4 kat hızlanma sunar; Blok-seyrek FlashAttention 2,8 kat hızlanma elde eder.
• Daha Uzun Bağlam Yetenekleri:
  • 4K bağlamlı GPT-2: Megatron-LM'in 1K bağlamdaki hızından daha hızlı eğitilirken doğrulama perplexity değerini 0,7 puan iyileştirir (18,2'ye kıyasla 17,5).
  • Uzun Belge Sınıflandırma: Bağlamı 512'den 16K'ya çıkarmak, MIMIC-III üzerinde 4,3 mikro-$F_1$ kazancı (52,8'den 57,1'e) ve ECtHR üzerinde 8,5 puana kadar kazanç sağlar (8K'da 72,2'den 80,7'ye).
  • Path-X ve Path-256: FlashAttention, Path-X'i (dizi uz. 16K) %61,4 doğrulukla çözen ilk Transformer modelidir. Blok-seyrek FlashAttention, Path-256'da (dizi uz. 64K) %63,1 doğruluk elde eder.

Kısıtlamalar ve açık sorular

• CUDA Mühendisliği İş Yükü: IO farkındalıklı çekirdekleri uygulamak, standart PyTorch veya derleyici soyutlamaları yerine özel, alt seviye CUDA programlaması gerektirir; bu da farklı donanım mimarileri arasında taşınabilirliği zorlaştırır.
• Çoklu GPU Dikkati: Tek düğüm için optimal IO çizelgelemesi, ara bağlantılar (NVLink, InfiniBand) üzerinden dağıtık çoklu GPU dikkatine doğrudan genişletilemez.
• Tam Dikkat Alt Sınırı: Genel $M$ değerleri için optimal olmakla birlikte, $M$ cinsinden ayrıntılı parametrelendirilmiş karmaşıklık alt sınırlarını kanıtlamak henüz açık bir problemdir.