CODA: Transformer Bloklarını GEMM-Epilogue Programları Olarak Yeniden Yazmak adlı yeni bir araştırma, Transformer modellerinin eğitim verimliliğini optimize etmeyi hedefliyor, özellikle görünüşte dağınık ancak biriktiğinde ciddi zaman kaybına neden olan "hafıza yoğun" işlemlerle ilgileniyor.

Yazan: Makine Kalbi

22 Mayıs'ta Tri Dao, Han Guo'nun bir tweet'ini sosyal medyada paylaştı. Ayrıca şunu yazdı: "Bazı matematiksel yeniden yazmalar sonucu, Transformer'ın tüm içeriğinin bir dizi GEMM + epilogue (matris çarpımı ve son ek) olduğu ortaya çıktı. Bazı optimize edilmiş primitifler verildiğinde, LLM'ler (ve yeni başlayanlar) tüm Transformer işlemlerini ışık hızında çekirdekler olarak yazabilir!"

Tri Dao, FlashAttention serisinin temel yazarlarından biridir ve bu tweet, o gün yayınladıkları CODA adlı makaleye işaret etmektedir.

Başlık: CODA: Transformer Bloklarını GEMM-Epilog Programları Olarak Yeniden Yazmak
Makale adresi: https://arxiv.org/abs/2605.19269
Kod adresi: https://github.com/HanGuo97/coda-kernels

Bu isim, “son beste” gibi okunur, “CUDA” gibi telaffuz edilir. MIT, Princeton, Together AI ve Meta’dan araştırmacılar, Transformer eğitimi sırasında nadiren dikkat çekse de sürekli zaman tüketen “dağınık hesaplamaları” sistematik bir şekilde ortadan kaldırmak için yeni bir programlama soyutlaması geliştirdi.

Büyük modeli eğitmek için "tembellik vergisi"

CODA'nın hangi sorunu çözmeye çalıştığını anlamak için büyük modellerin eğitim süresinin nerede harcandığını anlamak gerekir.

Bir NVIDIA H100 üzerinde bir LLaMA-3 tarzı 1 milyar parametreli model eğitmek, çoğu insanın doğrudan şu şekilde düşüneceğini düşünür: zamanın tamamı matris çarpımları ve dikkat hesaplamalarında harcanır, çünkü bunlar "gerçek hesaplama"dır. Bu sezgi büyük ölçüde doğrudur: matris çarpımları (GEMM) ve dikkat mekanizmaları gerçekten ana hesaplama yükünü oluşturur.

Ancak performans analizörünü dikkatlice incelediğinizde, normalleştirme (RMSNorm), aktivasyon fonksiyonları (SwiGLU, RoPE), kalan toplama, katmanlar arası toplama gibi bir dizi "küçük operatör"ün zamanı sessizce tükettiğini görebilirsiniz. Bunlar tek başına düşük hesaplama yüküne sahip olsa da, büyük ara tensörleri sık sık bellekten çıkarıp tekrar yükler.

Bu, "bellek bant genişliği darboğazı" olarak bilinir: Mükemmel bir şef olmakla birlikte, her yemeği hazırlamak için malzemeleri uzaktaki depodan getirmek ve kullandıktan sonra geri göndermek zorunda kalmak; bunun yerine malzemelerin hemen yanında bir masada hazır olması gerekir. Şefin el hızı ne kadar hızlı olursa olsun, taşıma süresi gerçek bir zaman kaybıdır.

Daha da kötüsü, NVIDIA'nın FP8, FP4 gibi düşük hassasiyetli formatlar sayesinde matris hesaplamaları giderek hızlanırken, bu "taşınma" işlemlerinin göreli maliyeti artmaktadır: Matris çarpımı hızlandı, ancak tensörlerin içeriye ve dışarıya taşınma maliyeti aynı oranda kısalmadı.

Makaledeki bir veri grubu oldukça açık: H100 üzerinde TorchTitan ile 1 milyar parametreli bir modeli eğitirken, matris çarpma işlemleri dışındaki işlemler,端到端 çalışma süresinin önemli bir kısmını oluşturuyor ve FP8 hassasiyetinin eklenmesiyle bu oran daha da artıyor.

Mevcut programlama çerçeveleri bu konuda neredeyse etkisiz kalıyor. PyTorch, Transformer hesaplamalarını net sınırlara sahip bir operatör dizisi olarak ifade eder. Bu sınırlar otomatik türev (autograd) için çok uygun olmakla birlikte, operatörler arası entegrasyon optimizasyonunu tam olarak engeller: her operatör sınırı genellikle gereksiz bir bellek yazma işlemidir.

CODA: «Son» içinde hazineler saklı

CODA, basit bir gözlemle başlar.

Bir GPU üzerinde, yüksek performanslı bir matris çarpımı (GEMM) çekirdeği, temel matris blok çarpım-toplama hesaplamalarını yapan ana döngü (mainloop) ve sonuçlar belleğe yazılırken önyükleme, tip dönüşümü, basit ölçekleme gibi son işlemleri yapan bitiş (epilogue) olmak üzere iki parçaya ayrılmıştır.

Son aşamanın amacı, bu anda matris çarpımının çıktısının hala çip üzerindeki kaydedicilerde "yaşaması" ve henüz genel video belleğine yazılması değildir. Bu, geçici bir altın penceredir: Bu anda ekstra hesaplamalar yapılırsa, bir video belleğine yazma ve tekrar okuma döngüsü tamamen atlanabilir.

CODA'nın temel içgörüsü, Transformer'daki hafıza yoğun işlemlerin çoğu cebirsel olarak yeniden parametrize edilip bu "son bölüm" penceresine yerleştirilebilir.

Bunun biraz matematik becerisi gerektirdiğini unutmayın. En yaygın GEMM-RMSNorm-GEMM modelini örnek alalım: Bir matris çarpımının sonucu, kalan toplama ve RMS normalizasyonundan geçtikten sonra başka bir matris çarpımı yapılır. Geleneksel yöntemde, üç bağımsız operatör sıralı olarak çalışır ve ara sonuçlar iki kez GPU belleğine yazılır.

CODA ekibi, RMS normalizasyonundaki satır ölçekleme faktörü r'nin, her satır için paylaşılan bir skaler olduğu için, ardından gelen matris çarpımı ile değişme özelliğine sahip olduğunu fark etti: r'nin uygulanması, "ikinci GEMM'den önce" yerine "ikinci GEMM'in sonunda" ertelenebilir. Bu ertelemeyi yaptıktan sonra, ilk GEMM'in sonunda yalnızca yerel "parçalı RMS" (partial RMS) hesaplanır ve bu, çok hafif bir yardımcı toplama çekirdeği ile birleştirilir; tam RMSNorm hesaplaması ortadan kalkar.

Benzer bir yeniden parametrelendirme, SwiGLU, RoPE (döndürülmüş konum kodlaması), çapraz entropi kaybı gibi işlemler için de geçerlidir ve hatta geri yayılım için bile geçerlidir. Makalede bir teorem kanıtlanmıştır: ileri yayılımın son aşaması “blok bazlı yerel” ise, geri yayılım otomatik olarak aynı yapıyı miras alır. Ayrıntılar için orijinal makaleye başvurun.

Beş adet "blok" ve bir set "Lego dili"

CODA, belirli bir birleşik çekirdek değil, bir programlama soyutlamaları kümesidir.

Uzmanlar tarafından optimize edilen GEMM ana döngüsünü sabitler ve sona doğru beş tür birleştirilebilir temel primitif曝露 eder:

Öğe bazlı dönüşüm (artık toplama, aktivasyon fonksiyonu, RoPE)
Vektör yükleme ve depolama (RMSNorm ağırlıklarının yayınlanması)
Matris bloklarının yüklenmesi ve depolanması (geri yayılım için ara aktivasyonların saklanması)
Blok İndirgeme (Yerel Kök Ortalama Karesi, Blok Log-Toplam-Exp)
Durum değişikliği (çevrimiçi normalleştirme için gerekli maksimum ve sum-exp istatistikleri)

Bu beş tür blok kullanılarak, dikkat mekanizmasının dışında bir standart Transformer'ın ileri ve geri yayılımındaki neredeyse tüm işlemler kapsanabilir.

Daha ilginç olanı, bu soyutun “kimin kod yazacağını” ne kadar kabul edici olduğudur. Makalede, iki uygulama modeli değerlendirildi: biri insan programcılar tarafından yazılan, diğeri ise Claude Code ile üretilen — CODA’nın ilkel açıklamaları, bazı örnekler ve uygulama günlükleri verilerek, AI’nın büyük ölçüde çekirdek kodunu oluşturması ve insanın hafif denetimi.

İki modun performansı da yüksek seviyede gerçekleşti. Tri Dao, tweet'te "LLM ve yeni başlayanlar ışın hızı çekirdeği yazabilir" dedi, bu da makaledeki deney sonuçlarının gerçek dünyadaki yansımasıdır.

Deneysel sonuçlar

CODA, benchmark seçiminde daha zorlu rakipleri kullanıyor: cuBLAS ve torch.compile, ayrıca LLM'ler için optimize edilmiş Liger Kernel ve FlashInfer.

Her bir çekirdek için iki uygulama değerlendirildi: CODA (LLM), Claude Code tarafından oluşturuldu, araştırmacılar temel açıklamalar, birkaç örnek ve sürekli güncellenen uygulama ipuçları günlüğü sağladı, AI ana kodu tamamladı, insanlar hafif bir denetim yaptı; CODA (Human), insan programcılar tarafından bağımsız olarak yazıldı, aynı yüksek seviyeli yeniden parametreleme fikrini kullandı, ancak CODA temel setine bağlı kalmadı. İki grup sonucu, cuBLAS + torch.compile, Liger Kernel ve FlashInfer gibi optimize edici kütüphanelerle karşılaştırıldı.

Tek bir operatör düzeyinde, GEMM-RMSNorm-GEMM tipik modeli örneği olarak alınarak, CODA, 1B, 7B ve 70B model boyutlarına karşılık gelen gizli boyutlarda cuBLAS + PyTorch temelini aştı. SwiGLU, RoPE, çapraz entropi gibi sonraki bileşenlerde de benzer performans gösterildi.

LLM tarafından oluşturulan çekirdekler, çoğu benchmark'ta insan tarafından yazılan sürümlerle eşit düzeyde performans gösteriyor ve bazı yapılandırmalarda hafifçe üstünlük sağlıyor. Bu, geleneksel olarak çok yüksek bir engel olan GPU çekirdek optimizasyonu alanında nadiren görülen bir sonuçtur.

Ters yayılım getirisi özellikle belirgindir: GEMM-Residual-PartialRMS-GEMM ters çekirdeği, temel duruma göre 1,6 ila 1,8 kat hızlanma sağlarken, SwiGLU tersi de yaklaşık 1,4 ila 1,6 katlık bir iyileşme göstermektedir. Bu yönde, LLM ile insan tarafından yapılan uygulamalar arasındaki fark da benzer şekilde küçüktür. Bu şaşırtıcı değildir: ters yayılım doğal olarak daha fazla ara tensör erişimi gerektirir, bu da son aşamada birleştirme getirisini artırır; CODA'nın primitif tasarımı yeterince net olduğundan, AI modelleri doğru bir şekilde birleştirme yapabilmektedir.

Tam Transformer katmanının端到端 benchmarkinde, CODA'nın ileri yönlü hızlandırması farklı boyutlarda %5 ile %20 arasındadır ve daha büyük model boyutlarında (70B boyutuna karşılık gelen gizli boyutlar) daha belirgin bir etki gösterir.

CODA'nın yeniden parametrelendirilmesi, RMSNorm ölçeklendirme faktörünün uygulanma zamanını ayarlar, ancak deneyler, bu sayının PyTorch referans uygulamasıyla karşılaştırıldığında benzer sayısal hatalara sahip olduğunu gösterir; bazı yapılandırmalarda hata daha da küçüktür — bunun nedeni GEMM ana döngüsünün kendisinin daha yüksek hassasiyete sahip toplayıcıları olmasıdır.

CODA'nın yapabileceği şeyler: Daha geniş bir bakış açısıya geçmeden önce CODA'nın yetenek sınırlarını netleştirin.

Kapsam: Dikkat ve kelime gömmeleri dışında, LLaMA mimarisine benzer standart Transformer'ların ileri ve geri yayılımında yer alan neredeyse tüm hesaplamalar, RMSNorm, kalan toplama, SwiGLU aktivasyonu, RoPE döner konum kodlaması, çapraz entropi kaybı ve yukarıdaki işlemlerin geri yayılım gradyanları dahil.
Hızlandırma etkisi: Karşılık gelen 1B ila 70B boyutunda gizli boyutlarda, tek işlemci düzeyinde cuBLAS + torch.compile temeline göre farklı oranlarda artış sağlanmıştır; geri yayılımda kazanım en belirgin olup (bazı çekirdeklerde %1,6'dan fazla); tam Transformer katmanında端到端 ileri yayılım hızlandırması yaklaşık %5 ila %20 arasındadır ve daha büyük model boyutlarında daha belirgin sonuçlar vermektedir.
CODA, CuTeDSL (NVIDIA CUTLASS'in Python DSL'i) üzerinde oluşturulmuştur ve hem insan programcılar hem de AI modelleri için yüksek performanslı çekirdek yazımını destekler.
Mevcut sınırlamalar: Şu anda yalnızca tek GPU senaryosu desteklenmektedir, dağıtık eğitim dahil değildir; yeniden parametrelendirme, standart Transformer mimarisine odaklanmaktadır, diğer mimariler için uygunluk doğrulanmamıştır.

Sonuç

CODA, izole bir çalışma değildir. Bu, GPU üzerinde gerçek optimizasyon alanlarının «ne hesaplanacağı» değil, «nasıl taşınacağı» olduğuna dair bir düşünce sınıfının somut bir uygulamasıdır.

FlashAttention, dikkat hesaplamalarını çip üzerindeki belleğe "taşıyor"; CODA, normalleştirme ve aktivasyon fonksiyonlarını da aynı şekilde "taşımaya" çalışıyor. Triton, özelleştirilmiş çekirdekler yazma engelini düşürüyor; ThunderKittens, TileLang vb. bu alanda farklı seviyelerde daha fazla keşif yapıyor. Bu çalışmalar, PyTorch operatör grafiğinin kolay ifade edilebilirliğini, el yazımı CUDA ile neredeyse eşdeğer yürütme verimliliğiyle birlikte, tek bir programlanabilir çerçevede gerçekten birleştirmeye doğru yöneliyor.

Tri Dao'nun tweet'inin son cümlesi tekrar düşünülmesi değerlidir: «LLM'ler ve yeni başlayanlar, tüm Transformer işlemlerine ışık hızında çekirdekler yazabilir.» Bunun ardında daha derin bir mantık yatıyor: Programlama soyutlaması yeterince iyi tasarlandığında, AI modelleri kendi eğitim altyapılarının optimizasyonuna dahil olabilir. Bu döngü, CODA'nın en ilgi çekici yönüdür.

Bu açıdan bakıldığında, «CODA» adı belki de başka bir anlam taşımaktadır. Klasik müzikte, Coda eserin sonunu kapatan bölümdür. Burada ise, GEMM çekirdeğinin «sonu»dur—ve bu sonu iyi yazmak, Transformer eğitimi sisteminin verimliliğindeki bir sonraki önemli adım olabilir.