Flash Attention

Flash Attention (FlashAttention), Transformer modellerinde tam dikkat işlemini, dikkat matrisinin tamamını GPU'nun ana belleğine hiç yazmadan hesaplayan bir algoritma ve GPU çekirdekleri ailesidir. N belirteç üzerinde standart öz dikkat, N × N boyutunda bir skor matrisi oluşturur, matrisi yüksek bant genişlikli belleğe (HBM) yazar ve softmax ile değerlerle çarpma işlemleri için matrisi yeniden okur; bu nedenle dikkatin çalışma süresinin çoğu verileri taşımaya harcanır. Stanford'dan Tri Dao ve çalışma arkadaşları bu yöntemi 2022 tarihli FlashAttention: G/Ç Farkındalığıyla Hızlı ve Bellek Açısından Verimli Tam Dikkat makalesinde tanıttı. Çalışmayı hızlı yonga üstü SRAM'e sığan parçalara böler; böylece bellek kullanımı dizi uzunluğuyla doğrusal olarak artar ve çıktı, kayan nokta yuvarlaması dışında standart dikkatle eşleşir. Çekirdekleri PyTorch içinde ve vLLM gibi büyük dil modeli (LLM) çıkarım motorlarında çalışır.

Flash Attention Nasıl Çalışır?#

Bir GPU, işlem birimlerinin yanında az miktarda çok hızlı belleğe, yonga dışında ise büyük ve daha yavaş bir belleğe sahiptir. FlashAttention makalesinde, 1,5-2,0 TB/sn hızında 40-80 GB HBM'e ve 108 akış çoklu işlemcisinin her birinde tahmini 19 TB/sn hızında 192 KB SRAM'e sahip bir NVIDIA A100 tanımlanır. Dikkat, okuduğu bayt başına az işlem gerçekleştirdiğinden hızı, verilerin iki bellek düzeyi arasında ne sıklıkta taşındığına bağlıdır; yöntem bu aktarımları sayar ve "G/Ç farkındalığı" da bunu ifade eder.

FlashAttention, sorgu parçacıklarını yükler ve ardından anahtar ile değer parçacıklarını SRAM üzerinden akıtır. Her parçacık için skorları hesaplar, her sorgu satırı için geçerli maksimumu ve toplamı günceller, maksimum büyüdüğünde kısmi çıktıyı yeniden ölçeklendirir ve yeni parçacığın katkısını ekler. Bu çevrimiçi softmax, tüm satırı aynı anda görmeden tam softmax sonucunu verir; yalnızca son çıktı ve satır başına iki sayı HBM'e geri yazılır. Geri yayılım geçişinde N × N matrisi depolamak yerine Q, K ve V'den skor parçacıkları yeniden hesaplanır; bu daha fazla işlem gerektirir, ancak çok daha az bellek okuduğu için yine de daha hızlı çalışır. Tüm sorguları her seferinde tek bir anahtar-değer parçacığına karşı işleyen aynı döngünün PyTorch sürümü:

import torch

torch.manual_seed(0)
n, d, block = 4096, 64, 256
q, k, v = torch.randn(3, n, d).unbind(0)

# Standard attention builds the full n x n score matrix in memory
reference = torch.softmax(q @ k.T / d**0.5, dim=-1) @ v

# FlashAttention-style: stream key/value tiles, keeping a running max and sum per query (online softmax)
row_max, row_sum, out = torch.full((n, 1), -torch.inf), torch.zeros(n, 1), torch.zeros(n, d)
for start in range(0, n, block):
    scores = q @ k[start : start + block].T / d**0.5  # only an n x block tile exists at once
    new_max = torch.maximum(row_max, scores.max(dim=-1, keepdim=True).values)
    rescale, probs = torch.exp(row_max - new_max), torch.exp(scores - new_max)
    row_sum = row_sum * rescale + probs.sum(dim=-1, keepdim=True)
    out = out * rescale + probs @ v[start : start + block]
    row_max = new_max
out = out / row_sum

print(f"Largest score tile: {tuple(scores.shape)} vs full matrix: ({n}, {n})")
print(f"Max difference from standard attention: {(out - reference).abs().max().item():.1e}")

En büyük skor parçacığı 4.096 × 256 boyutundadır; bu, 4.096 × 4.096 matrisinin on altıda biridir ve sonuç, standart dikkatten en fazla 4.3e-07, yani FP32 yuvarlama hatası kadar farklıdır. Gerçek bir çekirdek ayrıca sorguları, her GPU iş parçacığı bloğuna bir tane düşecek şekilde parçacıklara böler ve parçacıkların SRAM'den hiç çıkmaması için döngünün tamamını tek bir çekirdekte birleştirir.

FlashAttention-2, -3 ve -4#

Her sürüm daha yeni bir GPU neslini hedefler. FlashAttention-2 (2023), işin iş parçacıkları arasında bölünme biçimini yeniden düzenledi ve ilk sürümden yaklaşık iki kat hızlı çalışarak A100'ün tepe veriminin %50-73'üne ulaştı. FlashAttention-3 (2024), H100'ün eşzamansız Tensor Core birimlerini kullanır ve FP8 desteği sunarak FP16'da 740 TFLOPs/sn, FP8'de ise yaklaşık 1,2 PFLOPs/sn hızına ulaşır. FlashAttention-4 (2026), Blackwell GPU'ları hedefler ve B200'de BF16 için 1.613 TFLOPs/sn'ye, yani tepe verimin %71'ine kadar ulaştığını bildirir. Çekirdekler, FP16 ve BF16 girdilerle çalışan açık kaynaklı flash-attention kütüphanesinde sunulur; PyTorch'un scaled_dot_product_attention işlevi, GPU ve girdiler desteklediğinde FlashAttention-2 çekirdeğini otomatik olarak seçer, aksi durumda bellek açısından verimli veya standart uygulamalara geçer. Ultralytics paketinde çalışan Meta'nın SAM 2 modeli, Hiera görsel kodlayıcısında ve bellek dikkatinde bu işlev üzerinden dikkat hesaplar; böylece PyTorch, donanım ve girdi hassasiyeti izin verdiğinde bu katmanları FlashAttention'a yönlendirir.

Flash Attention ve İlgili Terimler#

• Standart dikkat: Aynı matematiksel işlemdir; skor ve olasılık matrislerinin tamamı HBM'e yazılarak hesaplanır.
• Seyrek dikkat: Yalnızca seçili sorgu-anahtar çiftlerinin skorlarını hesaplar. FlashAttention makalesi ayrıca boş parçacıkları atlayan blok-seyrek bir varyant da tanıttı.
• Doğrusal dikkat: Maliyetin dizi uzunluğuyla doğrusal artması için softmax yerine bir özellik haritası kullanır; bu da modelin çıktılarını değiştirir.
• PagedAttention: Sunum sırasında KV önbelleğini sabit boyutlu sayfalarda depolayan bir bellek yöneticisidir. FlashAttention dâhil dikkat çekirdekleri, anahtarları ve değerleri bu sayfalardan okur.
• Halka dikkati: Çok uzun bir diziyi cihazlara böler; her cihaz kendi parçası üzerinde bloklar hâlinde dikkat hesaplarken anahtar-değer blokları bir halka üzerinde dolaşır.

Eğitime ve Uzun Bağlamlara Etkisi#

Bellek kullanımı dizi uzunluğuyla doğrusal arttığından FlashAttention, uzun bağlam pencerelerinin eğitimini uygulanabilir kıldı. İlk makalede, MLPerf 1.1 eğitim rekoruna kıyasla 512 dizi uzunluğunda BERT-large için uçtan uca %15 hızlanma, 1.024 dizi uzunluğunda GPT-2 için 3 kat hızlanma ve Long Range Arena kıyaslamasında 2,4 kat hızlanma bildirildi. Aynı tasarruflar, bir Görü Transformer'ının (ViT) veya difüzyon Transformer'ının yüksek çözünürlüklü görselleri ve videoları uzun belirteç dizilerine dönüştürdüğü görü alanında da geçerlidir. Daha uzun diziler doğruluğu da artırdı: makalede, 128 × 128 boyutundaki bir görselin 16.384 piksellik bir diziye dönüştürülmesiyle elde edilen iki noktanın bir yolla birbirine bağlı olup olmadığını soran Path-X görevinde, şanstan daha iyi sonuç veren ilk Transformer'lar tanıtıldı.