Ang artikulo ay naglalarawan ng isang bagong pag-aaral na may pamagat na CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs, na may pangunahing layunin na mapabuti ang efisensiya ng pag-train ng Transformer models, lalo na sa paglutas ng mga “memory-intensive” na operasyon na tila nagkakahiwalay ngunit kumakalat sa malaking oras.

May-akda ng artikulo, pinagkukunan: Machine Mind

Noong Mayo 22, inihahatid ni Tri Dao ang isang tweet ni Han Guo sa social media. Isinulat niya rin: “Pagkatapos ng ilang matematikal na pagpapakilala, natuklasan na ang lahat ng mga bagay sa Transformer ay isang serye ng GEMM + epilogue (multiplication ng matrix at epilogue). Sa pamamagitan ng ilang pinagbutihing primitibo, maaaring isulat ng LLM (at mga baguhan) ang mga kernel na parang liwanag para sa lahat ng Transformer operations!”

Si Tri Dao ay isa sa mga pangunahang may-akda ng FlashAttention series, at ang tweet na ito ay naglalayong patungkol sa isang papel na kanilang inilabas sa araw na iyon: CODA.

Pamagat ng papel: CODA: Pagsusulat muli ng mga Transformer Block bilang mga GEMM-Epilogue Program
Link ng papel: https://arxiv.org/abs/2605.19269
Code address: https://github.com/HanGuo97/coda-kernels

Ang pangalan na ito, binabasa bilang “finale,” at binibigkas bilang “CUDA.” Ang mga mananaliksik mula sa MIT, Princeton, Together AI, at Meta ay sinusubukang gamitin ang isang bagong programming abstraction upang sistematisahin ang paglutas ng mga maliit na pagkalkula na karaniwang pinagkakailangan ng oras sa pagtratrain ng Transformer.

Ang "tax sa pagpapahinga" sa pagtratrabaho ng malalaking modelo

Upang maunawaan kung ano ang problema na lutasin ng CODA, unang kailangan mong maunawaan kung saan nagkakaroon ng oras sa pag-train ng malalaking modelo.

Sa isang NVIDIA H100, ang pag-train ng isang 1B parameter model na LLaMA-3 style, ang karamihan ay mag-iisip na ang lahat ng oras ay ginugugol sa matrix multiplication at attention computation, dahil iyon ang “totoong computation.” Ang intuisyong ito ay karamihan ay tama: ang matrix multiplication (GEMM) at attention ay talagang nangunguna sa pangunahing computing power.

Ngunit kung bubuksan mo ang performance profiler at titingnan nang mabuti, makikita mo ang isang grupo ng mga "maliit na operator" na tahimik na nagpapalubos ng oras: normalization (RMSNorm), activation functions (SwiGLU, RoPE), residual addition, cross-layer reduction... Ang bawat isa ay may maliit na computational load, ngunit madalas na nagdadala ng malalaking intermediate tensors pabalik at paharap sa GPU memory.

Ito ang tinatawag na “memory bandwidth bottleneck”: parang isang chef na may kakayahang magluluto nang perpekto, ngunit kailangan niyang dalhin ang bawat sangkap mula sa malayong imbakan, at ibalik ito pagkatapos gamitin, kesa itago sa tabi ng kanyang mesa. Kahit gaano pa kalakas ang bilis ng kamay ng chef, ang oras na inaasam sa pagdadala ay tunay na pagkawala.

Mas masama pa, habang nagiging mas mabilis ang mga matrix computation dahil sa mababang precision formats tulad ng FP8 at FP4 ng NVIDIA, tumataas ang relatibong gastos ng mga “paghahatid” na ito: nagmabilis ang matrix multiplication, ngunit hindi nagkakasabay ang gastos sa pagdadala at pagdadala pabalik ng tensor.

May isang set ng datos na diretso sa papel: Habang tinuturuan ang isang modelo na may 1B na parameter sa H100 gamit ang TorchTitan, ang mga operasyon na hindi matrix multiplication ay nagsisilbing malaking bahagi ng kabuuang oras ng pagpapatakbo, at ang proporsyon na ito ay lalong magiging mas makabuluhan kasama ang pagpapakilala ng FP8 precision.

Ang mga umiiral na framework sa pagprograma ay malabo sa pagharap dito. Ang PyTorch ay nagpapahayag ng pagkalkula ng Transformer bilang isang serye ng mga operator, na may malinaw na mga hangganan sa pagitan nito. Ang mga hangganan na ito ay napakadali para sa automatic differentiation (autograd), ngunit ito ay tinatanggal ang optimization ng fusion sa pagitan ng mga operator: bawat hangganan ng operator, karaniwan ay isang hindi kinakailangang pagbabalik sa memorya.

CODA: Ang kayamanan ay nakatago sa "Wakas"

Ang layunin ng CODA ay isang simpleng obserbasyon.

Sa GPU, ang isang mataas na performance na matrix multiplication (GEMM) kernel ay nahahati sa dalawang bahagi: ang mainloop na responsable sa pangunahing matrix blocking multiply-add computation, at ang epilogue na responsable sa ilang pagsasapupot sa pagkakasulat ng resulta pabalik sa VRAM, tulad ng pagdaragdag ng bias, type conversion, at simpleng scaling.

Ang kahalagahan ng wakas ay ang output ng matrix multiplication ay nananatiling "nabubuhay" sa on-chip registers at hindi pa nabababa sa global video memory. Ito ay isang maikling golden window: kung makakagawa ka ng karagdagang computation sa panahong ito, maaari mong buong-buoing iwasan ang isang round-trip ng pag-write at pag-read sa video memory.

Ang pangunahing pag-unawa ng CODA ay: ang mga operasyon na naglalayong mag-imbak ng memorya sa Transformer, marami sa kanila ay maaaring algebraically na muling-parameterisahan at isama sa loob ng "huling" window para sa pagsasagawa.

Kailangan nito ng kaunting matematikal na kasanayan. Sa halimbawa ng pinakakaraniwang GEMM-RMSNorm-GEMM pattern: ang resulta ng isang matrix multiplication, na sinusunod ng residual addition, RMS normalization, at pagkatapos ay isa pang matrix multiplication. Sa tradisyonal na paraan, tatlo silang hiwalay na operator na ginagawa nang serial, at dalawang beses ang intermediate results ay isinusulat sa video memory.

Nakakita ang koponan ng CODA na ang pag-scalar na r sa RMS normalization, dahil ito ay isang scalar na ibinabahagi sa bawat row, ay sumusunod sa commutative property sa pagkakapareho ng matrix multiplication: maaaring i-delay ang pag-apply ng r mula sa “bago ang pangalawang GEMM” hanggang sa “wakas ng pangalawang GEMM”. Pagkatapos ng pag-delay, ang wakas ng unang GEMM ay kailangan lamang na kalkulahin ang lokal na “partial RMS”, na pinagsasama ng isang napakamaliit na auxiliary reduction kernel, at nawala ang buong RMSNorm computation.

Ang katulad na reparameterization ay maaaring mag-apply din sa mga operasyon tulad ng SwiGLU, RoPE (rotational positional encoding), at cross-entropy loss, at kahit sa backpropagation. Mayroon isang teorema sa papel na patotohanan: kung ang forward pass ay “block-local”, ang backpropagation ay awtomatikong mananatili sa parehong istruktura. Tingnan ang orihinal na papel para sa detalye.

Limang mga “block” at isang “Lego language”

Ang CODA ay hindi isang tiyak na fusion kernel, kundi isang set ng programming abstraction.

Ito ay nagpapahigpit sa naka-optimize na GEMM main loop na galing sa mga eksperto, at pagkatapos ay inilalabas ang limang uri ng komposableng pangunahing primitibo:

Element-wise transformation (residual addition, activation function, RoPE)
Vector loading and storage (broadcasting RMSNorm weights)
Matrix block loading and storage (saving intermediate activations for backpropagation)
Block reduction (local RMS, block log-sum-exp)
May estado ng pagbabago (mga istatistika ng max at sum-exp na kinakailangan para sa online normalization)

Gamit ang limang uri ng bloke, ang halos lahat ng operasyon sa forward at backward pass ng isang standard na Transformer, maliban sa atensyon, ay maaaring sakop.

Mas interesante pa ang pagtanggap ng abstraksyon na ito sa "sino ang sumusulat ng code". Sa eksperimento, sinuri ng papel ang dalawang paraan ng implementasyon: isa ay isinulat ng mga programmer, at ang isa ay ginawa ng Claude Code—binigyan ng mga primitibo ng CODA, ilang halimbawa, at mga log ng implementasyon, ang AI ang nagawa ang karamihan sa core code, habang ang tao ay nagmungkahi ng kaunting pagmamasid.

Nakamit ng parehong mode ang mataas na antas ng performance. Sinabi ni Tri Dao sa tweet na “ang LLM at ang mga baguhan ay maaaring sumulat ng light-speed kernel,” na kung ano ang direktang pagpapakita ng mga resulta ng pananaliksik sa realidad.

Mga resulta ng eksperimento

Ang mga benchmark ng CODA ay nagpili ng mga kalaban na mas mahigpit: cuBLAS kasama ang torch.compile, at ang mga Liger Kernel at FlashInfer na espesyal na dinisenyo para sa LLM.

Sinusuri ng papel ang dalawang implementasyon para sa bawat kernel: ang CODA (LLM), na ginawa ng Claude Code, kung saan ang mga mananaliksik ang nagbigay ng mga paliwanag sa primitibo, ilang halimbawa, at isang patuloy na updated na log ng mga teknik sa implementasyon, habang ang AI ang nagawa ang pangunahing code at ang tao ay nagbigay ng mahinang pagmamasid; ang CODA (Human), na isinulat ng mga programmer na tao nang hiwalay, gumagamit ng parehong mataas na antas na reparameterization na ideya, ngunit hindi nakasalalay sa sariling CODA primitive set. Ang mga resulta ng dalawang grupo ay ihinahambing sa mga optimized library tulad ng cuBLAS + torch.compile, Liger Kernel, at FlashInfer.

Sa antas ng single operator, na may halimbawa ng karaniwang pattern na GEMM-RMSNorm-GEMM, nakamit ng CODA ang paglilipas sa cuBLAS + PyTorch baseline sa lahat ng tatlong laki ng modelo—1B, 7B, at 70B—sa kanilang mga hidden dimension. May katulad na pagganap din ang mga tail combination tulad ng SwiGLU, RoPE, at cross-entropy.

Ang mga kernel na ginawa ng LLM ay nagtatampok ng katulad na performans sa karamihan ng mga benchmark kumpara sa mga hand-written na bersyon, at sa ilang konfigurasyon ay kahit na umaabot sa mas mataas na antas. Isang napakakakaibang resulta ito sa larangan ng GPU kernel optimization, na karaniwang may mataas na hadlang.

Ang mga pagpapabilis sa backward pass ay lalong nakakatanda: ang backward kernel ng GEMM-Residual-PartialRMS-GEMM ay nakakapagpabilis ng 1.6 hanggang 1.8 beses kumpara sa baseline, habang ang SwiGLU backward ay may pagpapabilis na halos 1.4 hanggang 1.6 beses. Sa direksyong ito, ang pagkakaiba sa pagitan ng LLM at manual implementation ay maliit din. Hindi ito nakakagulat: ang backward pass ay natural na naglalaman ng mas maraming pag-access sa intermediate tensors, kaya mas malaki ang benepisyo ng tail fusion; at ang CODA’s primitive design ay sapat na malinaw upang payagan ang AI model na tamang i-combine ang mga ito.

Sa end-to-end benchmark ng buong Transformer layer, ang forward acceleration ng CODA ay umabot sa halos 5% hanggang 20% sa iba’t ibang sukat, at mas makabuluhan sa mas malalaking laki ng modelo (na tumutugma sa hidden dimension na 70B).

Sa aspeto ng numerical precision, ang reparameterization ng CODA ay binago ang panahon ng pag-apply ng scaling factor ng RMSNorm, ngunit ang mga eksperimento ay nagpakita na ang numerical error nito ay katumbas ng PyTorch reference implementation, at sa ilang konfigurasyon, mas maliit pa — dahil sa mas mataas na precision ng accumulator sa pangunahing loop ng GEMM.

Ano ang kayang gawin ng CODA: Isang mabilis na talaan upang malinaw ang mga kakayahan ng CODA bago magsimula sa mas malawak na pananaw.

Sakop: Halos lahat ng pagkalkula sa forward at backward pass ng standard Transformer (tulad ng arkitekturang LLaMA), maliban sa atensyon at word embedding, kabilang ang RMSNorm, residual addition, SwiGLU activation, RoPE rotational position encoding, cross-entropy loss, at ang backward gradient computation ng mga operasyong ito.
Accelerated performance: Sa mga lihim na dimensyon na 1B hanggang 70B, may pagpapabuti sa iba’t ibang antas sa antas ng isang operator kumpara sa cuBLAS + torch.compile baseline, kung saan ang pinakamalaking benepisyo ay nasa backward propagation (mga ilang kernel ay maaaring umabot sa higit sa 1.6 beses); ang end-to-end forward acceleration sa buong Transformer layer ay humigit-kumulang 5% hanggang 20%, na mas nakikita sa mas malalaking sukat ng modelo.
Sino ang makakagawa ng: CODA, na isinagawa batay sa CuTeDSL (Python DSL ng NVIDIA CUTLASS), na sumusuporta sa dalawang paraan ng paggawa ng kernel—para sa mga programmer at para sa AI models—at parehong paraan ay nakakamit ng mataas na performance.
Kasalukuyang limitasyon: Kasalukuyang suportado lamang ang single GPU scenario, walang pagkakaayos ng distributed training; ang reparameterization ay pangunahing nakatuon sa standard Transformer architecture, at ang applicability sa iba pang arkitektura ay nasa pagpapatunay pa.

Wakas

Hindi nag-iisa ang CODA. Ito ay isang konkretong pagpapatupad ng isang uri ng ideya: sa GPU, ang tunay na puwang para sa pagpapabuti ay hindi nasa "ano ang kalkulahin", kundi sa "paano ilipat".

Ang FlashAttention ay isinaksak ang pagkalkula ng atensyon sa on-chip memory, habang sinusubukan ng CODA na isaksak din ang normalization at activation functions. Ang Triton ay bumaba sa hadlang sa pagsulat ng custom kernels, habang ang ThunderKittens, TileLang, at iba pa ay patuloy na nag-aaral ng espasyong ito sa iba’t ibang antas. Ang mga gawaing ito ay nagtutulung-tulungan upang patunayan ang isang direksyon: ang pagkakaisa ng kahusayan sa pagpapahayag ng PyTorch operator graph at ang pagkakaroon ng efficiency na malapit sa hand-written CUDA sa isang isang programa na framework.

Ang huling pangungusap ni Tri Dao ay值得再回味: "Ang LLM at ang mga baguhan ay maaaring sumulat ng mga kernel na may bilis ng liwanag para sa lahat ng Transformer operations." May mas malalim na lohika dito: kapag ang mga abstrakson sa pag-program ay sapat na mabuti, ang mga modelo ng AI mismo ay maaaring makilahok sa pag-optimize ng kanilang sariling imprastruktura sa pag-train. Ang siklo na ito, ang pinakamalalim na aspeto ng CODA.

Sa pananaw na ito, ang pangalan na «CODA» ay maaaring may mas malalim na kahulugan. Sa klasikong musika, ang Coda ay ang seksyon na nagtatapos ng buong komposisyon. Dito, ito ay ang «tapos» ng GEMM kernel—at ang pagpapalawig nang maayos sa tapos na ito ay maaaring ang susunod na mahalagang kabanata sa pagpapabuti ng efisensiya ng Transformer training system.