Derin öğrenme alanındaki en etkili optimizasyonlardan biri olan FlashAttention, 2022'de yayınlanmasından bu yana donanım gücü arttıkça performansı daha da artıran dört ana sürümle karşımıza çıktı. Ancak makaleleri okumakla, bu optimizasyonların neden yapıldığını anlamak arasında büyük bir fark var. Bu makale, FlashAttention'ın temel prensiplerinden başlayarak, v1 sürümünü kağıtta anlatıldığı gibi yeniden uygulamayı, profillemeyi ve darboğazları tespit etmeyi amaçlıyor. Yazar, algoritma üzerinde iterasyonlar yaparak ve profilleme yoluyla v2, v3 ve v4 gibi sonraki sürümlerin neden gerekli olduğunu keşfederek bir "performans arkeolojisi" yapıyor.
Geleneksel dikkat mekanizmalarında, Q, K ve V tensörleri kullanılarak yapılan hesaplamalar sırasında, özellikle scores = torch.matmul(q, k.transpose(3, 2)) ve output = torch.matmul(scores, v) adımlarında ciddi bir bellek darboğazı oluşur. Bu, tam dikkat matrisini depolamak için O(S²) bellek karmaşıklığına sahip (B, N_h, S, S) boyutunda bir tensörün ortaya çıkmasına neden olur. Özellikle S (dizi uzunluğu) D_h'den (baş boyutu) çok daha büyük olduğunda (örneğin S=8192, D_h=64 veya 128), bu durum GPU belleğinde gigabaytlarca yer kaplayabilir. FlashAttention algoritması, bu devasa asimetriyi hedefleyerek bellek kullanımını optimize etmeyi ve performansı artırmayı mümkün kılar. Makale, bu bellek sorununu ve FlashAttention'ın bunu nasıl çözdüğünü adım adım inceliyor.
FlashAttention, derin öğrenme modellerindeki dikkat mekanizmasının bellek ve performans darboğazlarını aşarak modern yapay zeka uygulamalarının ölçeklenebilirliğini ve verimliliğini önemli ölçüde artırmıştır.