Artikel ini memperkenalkan penelitian baru berjudul CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs, yang bertujuan utama untuk mengoptimalkan efisiensi pelatihan model Transformer, khususnya dalam menangani operasi "berat memori" yang tampaknya terpisah tetapi secara kumulatif memakan waktu lama.

Penulis artikel, sumber: Machine Learning Indonesia

Pada 22 Mei, Tri Dao membagikan ulang sebuah tweet dari Han Guo. Ia juga menulis: "Setelah beberapa penulisan ulang matematis, ternyata seluruh aspek Transformer adalah serangkaian GEMM + epilogue (perkalian matriks ditambah epilog). Dengan beberapa primitif yang dioptimalkan, LLM (dan pemula) dapat menulis kernel berkecepatan cahaya untuk semua operasi Transformer!"

Tri Dao adalah salah satu penulis utama seri FlashAttention, dan tweet ini mengarah ke makalah yang mereka rilis pada hari itu: CODA.

• Judul makalah: CODA: Menulis Ulang Blok Transformer sebagai Program GEMM-Epilogue
• Alamat makalah: https://arxiv.org/abs/2605.19269
• Alamat kode: https://github.com/HanGuo97/coda-kernels

Nama ini, dibaca seperti "finale", dilafalkan seperti "CUDA". Para peneliti dari MIT, Princeton, Together AI, dan Meta mencoba menghilangkan secara sistematis perhitungan-perhitungan kecil yang jarang diperhatikan namun terus-menerus menghabiskan waktu dalam pelatihan Transformer, dengan menggunakan abstraksi pemrograman baru.

Biaya kemalasan dalam pelatihan model besar

Untuk memahami masalah yang diatasi oleh CODA, pertama-tama harus dipahami di mana waktu habis selama pelatihan model besar.

Melatih model berparameter 1B dengan gaya LLaMA-3 pada satu NVIDIA H100, kebanyakan orang akan secara intuitif menganggap: seluruh waktu dihabiskan untuk perkalian matriks dan perhitungan perhatian, karena itulah "komputasi sejati". Intuisi ini secara umum benar: perkalian matriks (GEMM) dan perhatian memang mendominasi daya komputasi utama.

Namun, jika Anda membuka profiler dan memeriksanya secara detail, Anda akan melihat sekelompok "operator kecil" yang secara diam-diam menghabiskan waktu: normalisasi (RMSNorm), fungsi aktivasi (SwiGLU, RoPE), penjumlahan residual, reduksi lintas lapisan... Mereka masing-masing memiliki beban komputasi kecil, tetapi sering kali memindahkan tensor menengah besar masuk dan keluar dari memori GPU.

Ini disebut sebagai "bottleneck bandwidth memori": seperti seorang koki hebat yang harus membawa bahan dari gudang jauh, lalu mengembalikannya setelah digunakan, alih-alih menempatkannya di meja di sampingnya. Seberapa cepat tangan koki itu bekerja, waktu menunggu pengangkutan tetap merupakan pemborosan nyata.

Yang lebih buruk lagi, semakin cepatnya perhitungan matriks akibat format presisi rendah seperti FP8 dan FP4 dari NVIDIA, biaya relatif operasi "pemindahan" ini justru meningkat: perkalian matriks menjadi lebih cepat, tetapi biaya memindahkan tensor masuk dan keluar tidak berkurang secara proporsional.

Sebuah set data dalam makalah ini sangat intuitif: saat melatih model dengan 1 miliar parameter di H100 menggunakan TorchTitan, operasi non-perkalian matriks menyumbang sebagian besar waktu eksekusi end-to-end, dan proporsi ini akan semakin menonjol dengan diperkenalkannya presisi FP8.

Framework pemrograman yang ada hampir tidak berdaya terhadap hal ini. PyTorch menyatakan perhitungan Transformer sebagai rangkaian urutan operator, dengan batas yang jelas di antara masing-masing operator. Batas-batas ini sangat ramah terhadap diferensiasi otomatis (autograd), tetapi justru menghalangi optimasi fusi lintas-operator: setiap batas operator sering kali merupakan penulisan ulang memori yang tidak perlu.

CODA: Kekayaan tersembunyi di dalam "Penutup"

Titik awal CODA adalah pengamatan sederhana.

Pada GPU, kernel perkalian matriks berperforma tinggi (GEMM) terbagi secara struktural menjadi dua bagian: mainloop yang bertanggung jawab atas perhitungan perkalian dan penjumlahan blok matriks inti, serta epilogue yang melakukan penanganan akhir sebelum hasil ditulis kembali ke memori video, seperti penambahan bias, konversi tipe, dan penskalaan sederhana.

Makna dari akhir ini adalah bahwa output dari perkalian matriks masih "hidup" di register on-chip dan belum ditulis ke memori global. Ini adalah jendela emas singkat: jika dapat melakukan lebih banyak perhitungan pada saat ini, maka dapat sepenuhnya menghindari satu siklus penulisan dan pembacaan kembali memori video.

Wawasan inti CODA adalah: banyak operasi yang memakan memori di Transformer sebenarnya dapat direparameterisasi secara aljabar dan dilakukan dalam jendela "penutup" ini.

Ini memerlukan sedikit keahlian matematika. Sebagai contoh pola GEMM-RMSNorm-GEMM yang paling umum: hasil perkalian matriks, kemudian penjumlahan residual, normalisasi RMS, dan dilanjutkan dengan perkalian matriks lainnya. Pendekatan tradisional menjalankan ketiga operator ini secara seri, dengan hasil sementara ditulis ke memori GPU dua kali.

Tim CODA menemukan bahwa faktor penskalaan baris r dalam normalisasi RMS, karena merupakan skalar yang dibagikan per baris, memenuhi sifat komutatif dengan perkalian matriks berikutnya: penerapan r dapat ditunda dari "sebelum GEMM kedua" menjadi "di akhir GEMM kedua". Setelah ditunda, di akhir GEMM pertama hanya perlu menghitung "akar rata-rata kuadrat sebagian" (partial RMS), yang digabungkan oleh kernel reduksi tambahan yang sangat ringan, sehingga perhitungan RMSNorm lengkap menjadi tidak diperlukan.

Reparameterisasi serupa juga berlaku untuk operasi seperti SwiGLU, RoPE (rotational positional encoding), cross-entropy loss, bahkan juga berlaku untuk backpropagation. Terdapat teorema dalam paper asli yang membuktikan: selama tahap akhir forward pass bersifat "block-local", backpropagation secara otomatis mewarisi struktur yang sama. Untuk detail lebih lanjut, silakan kunjungi paper aslinya.

Lima «blok» dan satu set «bahasa Lego»

CODA bukanlah sebuah kernel fusi spesifik, melainkan seperangkat abstraksi pemrograman.

Ia mengunci loop utama GEMM yang telah dioptimasi oleh ahli, lalu mengekspos lima jenis primitif dasar yang dapat dikombinasikan di posisi akhir:

• Transformasi elemen-per-elemen (penjumlahan residual, fungsi aktivasi, RoPE)
• Pemuatan dan penyimpanan vektor (broadcast RMSNorm weights)
• Pemuatan dan penyimpanan blok matriks (menyimpan aktivasi tengah untuk digunakan dalam propagasi balik)
• Block reduction (local RMS, block log-sum-exp)
• Ada perubahan status (statistik max dan sum-exp yang diperlukan untuk normalisasi online)

Dengan lima jenis blok ini, hampir seluruh operasi dalam forward dan backward propagation Transformer standar, selain perhatian, dapat dicakup.

Yang lebih menarik adalah tingkat toleransi sistem abstrak ini terhadap "siapa yang menulis kode". Dalam eksperimen, makalah mengevaluasi dua pola implementasi: satu oleh programmer manusia, dan satu lagi oleh Claude Code — dengan memberikan deskripsi primitif CODA, beberapa contoh, dan log implementasi, AI menyelesaikan sebagian besar kode inti, dengan pengawasan ringan oleh manusia.

Kinerja kedua mode mencapai tingkat yang tinggi. Tri Dao mengatakan dalam tweet, "LLM dan pemula dapat menulis kernel kecepatan cahaya," yang merupakan refleksi nyata dari hasil eksperimen dalam paper tersebut.

Hasil eksperimen

Benchmark untuk CODA memilih lawan yang cukup menantang: cuBLAS ditambah torch.compile, serta Liger Kernel dan FlashInfer yang dioptimalkan khusus untuk LLM.

Paper mengevaluasi dua implementasi untuk setiap kernel: CODA (LLM) yang dihasilkan oleh Claude Code, dengan peneliti menyediakan penjelasan primitif, beberapa contoh, dan log trik implementasi yang diperbarui secara terus-menerus; AI menghasilkan kode utama, sementara manusia memberikan pengawasan ringan; CODA (Human) ditulis secara independen oleh programmer manusia, menggunakan pendekatan reparameterisasi tingkat tinggi yang sama, tetapi tidak bergantung pada set primitif CODA itu sendiri. Hasil kedua kelompok dibandingkan dengan pustaka optimasi seperti cuBLAS + torch.compile, Liger Kernel, dan FlashInfer.

Pada tingkat operator tunggal, dengan contoh pola khas GEMM-RMSNorm-GEMM, CODA mencapai kinerja yang melampaui baseline cuBLAS + PyTorch pada dimensi tersembunyi untuk tiga skala model: 1B, 7B, dan 70B. Kelompok akhir seperti SwiGLU, RoPE, dan cross-entropy juga menunjukkan performa serupa.

Kernel yang dihasilkan oleh LLM sebanding dengan versi yang ditulis secara manual pada sebagian besar benchmark, dan dalam beberapa konfigurasi bahkan sedikit melampaui. Ini merupakan kesimpulan yang cukup langka di bidang optimasi kernel GPU yang secara historis memiliki ambang masuk sangat tinggi.

Manfaat dari backward propagation sangat menonjol: kernel backward untuk GEMM-Residual-PartialRMS-GEMM dapat mempercepat hingga 1,6 hingga 1,8 kali dibandingkan baseline, sementara backward SwiGLU juga meningkat sekitar 1,4 hingga 1,6 kali. Dalam arah ini, kesenjangan antara LLM dan implementasi manual juga sangat kecil. Hal ini tidak mengejutkan: backward propagation secara alami melibatkan lebih banyak akses ke tensor menengah, sehingga manfaat dari fusion di akhir lebih besar; dan desain primitif CODA cukup jelas, memungkinkan model AI untuk secara tepat melakukan kombinasi.

Dalam benchmark end-to-end lapisan Transformer lengkap, percepatan forward CODA berkisar sekitar 5% hingga 20% di berbagai skala, dengan efek yang lebih signifikan pada ukuran model yang lebih besar (sesuai dengan dimensi tersembunyi skala 70B).

Dalam hal presisi numerik, reparameterisasi CODA mengubah waktu penerapan faktor penskalaan RMSNorm, tetapi eksperimen menunjukkan bahwa kesalahan numeriknya sebanding dengan implementasi referensi PyTorch, dan dalam beberapa konfigurasi, kesalahannya bahkan lebih kecil—berkat akumulator dengan presisi lebih tinggi dalam loop utama GEMM.

Apa yang dapat dilakukan CODA: Daftar cepat untuk memperjelas batasan kemampuan CODA sebelum melihat perspektif yang lebih luas.

• Cakupan: Hampir seluruh perhitungan dalam forward dan backward propagation dari Transformer standar (seperti arsitektur LLaMA), selain perhatian dan embedding kata, termasuk RMSNorm, penjumlahan residual, aktivasi SwiGLU, encoding posisi rotasi RoPE, loss cross-entropy, serta perhitungan gradien backward dari operasi-operasi tersebut.
• Efek percepatan: Pada dimensi tersembunyi skala 1B hingga 70B, tingkat peningkatan pada level operator dibandingkan baseline cuBLAS + torch.compile bervariasi, dengan manfaat terbesar pada propagasi mundur (beberapa kernel mencapai lebih dari 1,6 kali); percepatan end-to-end forward pada lapisan Transformer lengkap sekitar 5% hingga 20%, dengan efek yang lebih menonjol pada model berukuran lebih besar.
• Siapa pun dapat menggunakan: CODA yang diimplementasikan berdasarkan CuTeDSL (DSL Python dari NVIDIA CUTLASS), mendukung dua cara penulisan kernel—oleh programmer manusia dan model AI—dengan kinerja tinggi pada kedua metode tersebut.
• Batasan saat ini: Saat ini hanya mendukung skenario satu GPU, tanpa melibatkan pelatihan terdistribusi; reparameterisasi terutama ditujukan untuk arsitektur Transformer standar, dan kelayakan untuk arsitektur lain masih perlu diverifikasi.

Penutup

CODA bukan pekerjaan yang terpisah. Ini adalah implementasi konkret dari sekelompok pemikiran: di GPU, ruang optimasi sejati sering kali bukan pada "apa yang dihitung", melainkan "bagaimana memindahkannya".

FlashAttention membuat perhitungan perhatian "berpindah" ke memori on-chip, sementara CODA berusaha membuat normalisasi dan fungsi aktivasi juga "berpindah" ke sana. Triton menurunkan hambatan untuk menulis kernel khusus, sementara ThunderKittens, TileLang, dan lainnya terus mengeksplorasi ruang ini pada berbagai tingkatan. Pekerjaan-pekerjaan ini secara bersama-sama menunjuk ke arah yang sama: menyatukan kemudahan ekspresi grafik operator PyTorch dengan efisiensi eksekusi yang mendekati CUDA tulisan tangan, dalam satu kerangka yang dapat diprogram.

Kalimat terakhir dari tweet Tri Dao patut direnungkan lagi: "LLM dan pemula dapat menulis kernel berkecepatan cahaya untuk semua operasi Transformer." Di balik ini ada logika yang lebih dalam: ketika abstraksi pemrograman dirancang dengan cukup baik, model AI sendiri dapat berpartisipasi dalam optimasi infrastruktur pelatihan mereka sendiri. Siklus inilah yang paling menarik dari CODA.

Dari sudut pandang ini, nama "CODA" mungkin memiliki makna tersembunyi. Dalam musik klasik, Coda adalah bagian penutup yang mengakhiri sebuah komposisi. Di sini, ia adalah "penutup" dari inti GEMM—dan menulis penutup ini dengan baik mungkin merupakan bab penting berikutnya dalam peningkatan efisiensi sistem pelatihan Transformer.