Öz-dikkat, CUDA çekirdekleri ve FlashAttention

Modern Transformer'lar, devasa token dizilerini azami donanım hızında işlemek için Öz-Dikkat matris matematiğine, Triton gibi yüksek seviyeli GPU derleyicilerine ve FlashAttention gibi belleğe duyarlı CUDA çekirdeklerine dayanır.
1. Bir Vision Transformer'da Öz-Dikkat Matematiksel Olarak Nasıl Çalışır
Bir Vision Transformer'da (ViT) bir görüntü (ör. 224×224 piksel) 16×16 yamalardan oluşan bir ızgaraya dilimlenir ve N = 196 düzleştirilmiş yama elde edilir. Her yama, D boyutlu bir vektöre (ör. 768) doğrusal olarak yansıtılır ve bir konum gömmesiyle birleştirilerek X ∈ R^(N × D) girdi matrisini oluşturur.
- Sorgu, Anahtar ve Değer Oluşturma. Model, yama matrisi X'i üç öğrenilmiş ağırlık matrisiyle (W_Q, W_K, W_V ∈ R^(D × d_k)) çarparak her görüntü yaması için üç farklı temsil oluşturur:
Q=XW_(Q), K=XW_(K), V=XW_(V)
- Sorgu (Q): Belirli bir görüntü yamasının neyi aradığını temsil eder (ör. "Ben bir kürk dokusuyum, yakınlarda gözler veya kulaklar arıyorum").
- Anahtar (K): Bir yamanın hangi bilgiyi içerdiğini duyurur (ör. "Keskin, kavisli bir kenar içeriyorum").
- Değer (V): Bir eşleşme olursa ileriye aktarılacak gerçek görsel özellik yükünü tutar.
- Dikkat Skorlarını Hesaplama. Her yamanın diğer her yamayla ne kadar güçlü ilişkili olduğunu ölçmek için model, tüm Sorgular ile Anahtarlar arasındaki iç çarpımı hesaplar:
Scores=(QK^(T))/(sqrt(d_(k)))
sqrt(d_(k)) ile bölmek, büyük iç çarpımların sonraki softmax fonksiyonunu çok küçük gradyanlı bölgelere itmesini önler. Bu, (i, j) girdisinin i yamasının j yamasına ne kadar dikkat etmesi gerektiğini skorladığı N × N boyutlu bir dikkat matrisi üretir.
- Softmax Ağırlıklandırma ve Toplama. Satır bazında Softmax uygulamak ham skorları toplamı 1 olan olasılıklara dönüştürür. Bu olasılıkları Değer matrisi V ile çarpmak her yama için harmanlanmış bir temsil üretir:
Attention(Q,K,V)=softmax((QK^(T))/(sqrt(d_(k))))V
Çok Başlı Öz-Dikkat'te bu süreç tam olarak aynı şekilde birden fazla başlık (ör. 12 başlık) üzerinde paralel çalışır; böylece farklı başlıklar aynı anda farklı uzamsal mesafelere, renklere veya nesne parçalarına odaklanabilir.
2. Uygulamalı Kod Karşılaştırması: Ham CUDA C++ ve OpenAI Triton
GPU programlamanın nasıl evrildiğini görmek için N eleman üzerinde basit bir eleman bazlı vektör toplamını (C = A + B) ele alalım.
Ham CUDA C++ (İş Parçacığı Düzeyinde Yönetim)
CUDA C++'ta programı tek bir skaler iş parçacığının bakış açısından yazarsınız. İş parçacığı kimliklerini elle hesaplamanız, dizi sınırlarını denetlemeniz ve ana makine tarafı başlatma yapılandırmalarının kalıp kodunu yazmanız gerekir.
#include <cuda_runtime.h>
// Device kernel: executed by thousands of individual threads
__global__ void vector_add_cuda(const float* A, const float* B, float* C, int N) {
// Manually compute global thread index from block and thread IDs
int idx = blockIdx.x * blockDim.x + threadIdx.x;
// Guard against out-of-bounds memory access
if (idx < N) {
C[idx] = A[idx] + B[idx];
}
}
// Host launch setup
void launch_cuda(const float* d_A, const float* d_B, float* d_C, int N) {
int threads_per_block = 256;
int blocks_per_grid = (N + threads_per_block - 1) / threads_per_block;
vector_add_cuda<<<blocks_per_grid, threads_per_block>>>(d_A, d_B, d_C, N);
}OpenAI Triton (Blok Düzeyinde Python Programlama)
Triton'da programı doğrudan Python'da bir veri bloğunun (döşemesinin) bakış açısından yazarsınız. Triton’un derleyicisi iş parçacığı senkronizasyonunu, bellek birleştirmeyi (coalescing) ve vektörleştirmeyi arka planda otomatik olarak halleder.
import triton
import triton.language as tl
@triton.jit
def vector_add_triton(A_ptr, B_ptr, C_ptr, N, BLOCK_SIZE: tl.constexpr):
# Identify which block of data this program instance handles
pid = tl.program_id(axis=0)
# Create a vector of offsets for the entire block at once
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
mask = offsets < N
# Load chunks of memory into fast SRAM, add them, and store back
a = tl.load(A_ptr + offsets, mask=mask)
b = tl.load(B_ptr + offsets, mask=mask)
tl.store(C_ptr + offsets, a + b, mask=mask)
# Launch directly from PyTorch tensors
def launch_triton(A, B, C):
N = A.numel()
grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
vector_add_triton[grid](A, B, C, N, BLOCK_SIZE=1024)3. FlashAttention: Transformer Bellek Darboğazını Çözmek
Standart Öz-Dikkat, yüksek çözünürlüklü görüntüleri veya uzun metin dizilerini işlerken ciddi bir duvara çarpar; çünkü S = QK^T dikkat matrisi, dizi uzunluğu N ile ikinci dereceden (N × N) büyür.
- Bellek Bant Genişliği Darboğazı (HBM ve SRAM). Modern NVIDIA GPU'larında iki ana bellek katmanı vardır:
- HBM (Yüksek Bant Genişlikli Bellek / VRAM): Büyüktür (ör. 80 GB) ancak okuma ve yazma açısından görece yavaştır.
- SRAM (Çip Üzerinde Paylaşılan Bellek): Küçüktür (tüm Streaming Multiprocessor'lar boyunca toplam ~20 MB) ancak HBM'den yaklaşık 10 ila 20 kat daha hızlıdır.
Standart dikkatte GPU, S = QK^T'yi hesaplar, devasa N × N matrisini yavaş HBM'ye yazar, Softmax uygulamak için geri okur, olasılıkları tekrar HBM'ye yazar ve V ile çarpmak için bir kez daha okur. Tensor Çekirdekleri zamanlarının çoğunu bellek aktarımlarını bekleyerek boşta geçirir.
- Döşeme (N × N Matrisini Hiç Gerçekleştirmeme). Tri Dao ve arkadaşları (2022) tarafından tanıtılan FlashAttention, tüm dikkat hesaplamasını hızlı çip üstü SRAM içinde tutmak için çekirdek birleştirme (kernel fusion) kullanır. Q, K ve V'yi SRAM'e sığan küçük bloklara dilimler, dikkati döşeme döşeme hesaplar ve HBM'ye yalnızca nihai çıktı matrisini yazar. Devasa N × N ara matris VRAM'e hiç kaydedilmez.
- Çevrimiçi Softmax Hilesi. Normalde Softmax, herhangi bir elemanı normalleştirmeden önce bir satırın tamamındaki maksimum değeri ve üstel toplamlarını bilmeyi gerektirir. FlashAttention, her yeni döşeme işlenirken kısmi çıktıları matematiksel olarak yeniden ölçeklendirerek blok blok çalışan maksimumları ve normalleştirme faktörlerini izleyen çevrimiçi bir Softmax algoritması kullanır.
- Geri Yayılım Sırasında Yeniden Hesaplama. FlashAttention ileri geçiş sırasında N × N dikkat matrisini kaydetmediğinden, eğitim sırasında akıllıca bir takas yapar: HBM'de yalnızca küçük normalleştirme istatistiklerini saklar ve geri yayılım sırasında dikkat döşemelerini hızlı SRAM içinde anında yeniden hesaplar. SRAM'de ekstra matematik yapmak, HBM'den büyük matrisleri okumaktan çok daha hızlı çıkar; bellek kullanımını ikinci dereceden O(N²)'den doğrusal O(N)'ye düşürürken 2× ila 4× duvar saati hızlanması sağlar.