Sıfırdan yazılan CUDA attention kerneli, NVIDIA B200 üzerinde FlashAttention-4 performansının yüzde 94,4'üne ulaştı

FlashAttention-4 hızına doğru 14 adımlık bir tırmanış

Hacker News'in ana sayfasına ulaşan bir blog yazısı, NVIDIA'nın Blackwell B200 GPU'su için sıfırdan, az miktarda PTX içeren düz CUDA ile yoğun (dense) bir attention kernelinin nasıl inşa edileceğini adım adım anlatıyor. Yazarın belirttiğine göre bitmiş kernel, FA4 makalesinde kullanılan 4K, 8K ve 16K dizi uzunluğu (sequence length) şekillerinde FlashAttention-4 performansının yüzde 94,4'ünü sağlıyor. Belirtilen hedef rekor kırmak değil, öğretmek; son birkaç yüzdeyi sıkmak ise sonraki bir yazıya bırakılmış.

Materyal, her adımın tek bir optimizasyon eklediği ve yaklaşık 60 diyagramın açıklamayı taşıdığı 14 kernel sürümünden oluşan bir ilerleme halinde düzenlenmiş. Kapsam bilinçli olarak dar tutulmuş: head boyutu 128 olan, BF16 ile hesaplanan yoğun, nedensel olmayan (non-causal) attention. Bir capstone olarak son kernel, örnek videolar üretmek için bir video üretim modeline bağlanıyor.

Kernel işini nasıl bölümlendiriyor

Yazıya göre Q, K, V ve çıktı tensörünün tümü batch, head sayısı, dizi uzunluğu ve head boyutu şeklini paylaşıyor ve 128'e 128 parçalara (tile) kesiliyor. Her çıktı parçasına bir CTA atanıyor ve bu CTA'lar grid genelinde paralel çalışıyor. CTA içinde iş sıralı: kendi Q parçasını bir kez yüklüyor, sonra K/V parça çiftleri üzerinde döngüye giriyor. Her yineleme S = Q @ K^T hesaplıyor, satır maksimumlarını ve toplamlarını tutarken online-softmax güncellemesi uygulayarak P üretiyor ve P @ V'yi özel bir çıktı tamponunda biriktiriyor. Döngü bitince bu tampon normalize ediliyor ve global belleğe yazılıyor. Yazar yapıyı parçalı (tiled) bir matmule benzetiyor; ek zorluk ise softmax durumunun K/V döngüsü boyunca hayatta kalması gerektiği.

Darboğaz matmul değil, softmax

Yazının merkezindeki teknik nokta, attention'ın ortasında softmax bulunan iki matris çarpımına denk geldiği ve Blackwell'de verimi sınırlayanın softmax olduğu. Çok daha az FLOP gerçekleştirse de softmax, tensor core'lar yerine ALU ve MUFU birimlerinde çalışıyor. B200'de tensor core verimi kabaca ikiye katlanırken üstel birimleri büyük ölçüde aynı kaldı; FA4 makalesinin "asimetrik donanım ölçeklenmesi" adını verdiği bir uyumsuzluk bu. Yazarın parça boyutlarında softmax, matmuller kadar döngü tüketiyor. Bu nedenle ilerlemedeki optimizasyonlar iki yönden saldırıyor: softmax işini ucuzlatmak ve onu tensor core matematiğiyle örtüşştürmek.

CuTe yerine ham CUDA

FlashAttention-4'ün kendisi CuTe ile yazılmış, ancak yazar daha az soyutlama katmanı içerdiği için ham CUDA artı PTX'i takip etmeyi daha kolay buluyor. Sonuç, FA4'ün satır satır bir çevirisi değil bağımsız bir inşaat; FA4 ise ana referans ve optimizasyon fikirlerinin çoğunun kaynağı olarak gösteriliyor. Bir eşlik eden depo (repository), kodu komşu kernel sürümlerinin diff'lenebileceği şekilde düzenliyor ve her adımın neyi değiştirdiğini izole ediyor.

Giriş bilinçli olarak yumuşak tutulmuş: okuyucuların temel CUDA aşinalığı ile attention ve online softmax anlayışına sahip olması yeterli; tcgen05 tensor core kavramları gibi donanıma özgü mekanizmalar yalnızca gerektiğinde tanıtıldığı için önceki Blackwell bilgisi gerekmiyor. Yazar, yazıyı takip etmek için bir B200 gerektiğini de not ediyor; kendisi de sahip değil.

Neden önemli

Attention kerneleri eğitim ve çıkarımın (inference) kritik yolunda yer alıyor ve Blackwell ekonomilerini değiştiriyor: tensor core'lar üstel işlemleri yapan özel işlev birimlerinden çok daha hızlı ölçeklenirken, önceki GPU'lar için yazılmış kerneller verimi boşta bırakıyor. Yazara göre B200 attention'ına dair mevcut açıklamalar ya onu üreten ilerlemeyi belgelemeden bitmiş bir kerneli belgeliyor ya da yüksek seviyede kalıyor. FlashAttention-4'ün yüzde altısı içinde kalan, diff'lenebilir kodla birlikte adım adım bir anlatım, kernel mühendislerine Blackwell'in ödünleşimleri hakkında akıl yürütmek için çalışan bir şablon veriyor; zihinsel modeller de bu tek kernelin çok ötesine taşınıyor.