この記事では、CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programsという新しい研究を紹介しています。その中心的な目標は、Transformerモデルのトレーニング効率を最適化すること、特に、一見するとばらばらだが、累積すると時間のかかる「メモリ集中型」演算を解決することです。

5月22日、Tri DaoはソーシャルメディアでHan Guoのツイートをリツイートしました。彼はまた、「いくつかの数学的な書き換えの後、Transformerのすべての内容は一連のGEMM + epilogue（行列乗算と後処理）であることがわかりました。いくつかの最適化されたプリミティブを与えれば、LLM（および初心者）は、すべてのTransformer操作に対して光速カーネルを作成できます！」と書いています。Tri DaoはFlashAttentionシリーズの中核的な著者の一人であり、このツイートは彼らがその日に発表した論文CODAを指しています。論文タイトル：CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs。論文アドレス：https://arxiv.org/abs/2605.19269。コードアドレス：https://github.com/HanGuo97/coda-kernels。

この名前は、「終曲」のように聞こえ、「CUDA」のように発音します。MIT、プリンストン、Together AI、Metaの研究者たちは、新しいプログラミング抽象化を用いて、Transformerのトレーニングにおいてほとんど注目されず、時間を消費し続ける「ばらばらの計算」を体系的に消化しようとしています。

背景：大規模モデルのトレーニングにおける「怠惰税」。CODAがどのような問題を解決しようとしているのかを理解するには、まず大規模モデルのトレーニング時間がどこに使われているのかを理解する必要があります。NVIDIA H100上でLLaMA-3スタイルの10億パラメータモデルをトレーニングする場合、ほとんどの人は、時間が行列乗算と注意計算に使われていると直感的に考えるでしょう。結局のところ、それこそが「真の計算」だからです。この直感はほぼ正しいです。行列乗算（GEMM）と注意計算は確かに主要な計算能力を占めています。しかし、パフォーマンスアナライザを開いてよく見ると、正規化（RMSNorm）、活性化関数（SwiGLU、RoPE）、残差加算、層間削減など、一連の「小さな演算子」が静かに時間を消費していることに気づくでしょう。これらの演算は個々の計算量は大きくありませんが、大規模な中間テンソルを頻繁にメモリから出し入れします。

これが、いわゆる「メモリ帯域幅のボトルネック」です。それは、腕の良い料理人が、料理を作るたびに食材を遠くの倉庫から運び、使い終わったらまた送り返すようなものです。料理人の腕がどんなに速くても、運搬を待つ時間は無駄になります。さらに悪いことに、NVIDIAのFP8、FP4などの低精度フォーマットによって行列計算がますます高速化されるにつれて、これらの「運搬」操作の相対的なコストは上昇しています。行列乗算は高速化されましたが、テンソルを出し入れするコストは同じように短縮されていません。論文には、H100上でTorchTitanを使用して10億パラメータモデルをトレーニングする場合、行列乗算以外の操作がエンドツーエンドの実行時間の相当な部分を占めており、FP8精度の導入に伴い、この割合がさらに顕著になるという、非常にわかりやすいデータがあります。既存のプログラミングフレームワークは、これに対してほとんど何もできません。PyTorchは、Transformerの計算を一連の演算子シーケンスとして表現しており、演算子間には明確な境界があります。この境界は自動微分（autograd）には非常に適していますが、演算子間の融合最適化を妨げています。各演算子の境界は、多くの場合、不必要なメモリへの書き戻しです。

CODA：「後処理」に隠された宝物。CODAの出発点は、素朴な観察です。GPU上では、高性能な行列乗算（GEMM）カーネルは、構造的に2つの部分に分かれています。メインループ（mainloop）は、コアとなる行列のブロック乗算と加算計算を担当し、後処理（epilogue）は、結果をメモリに書き戻す前に、バイアス加算、型変換、簡単なスケーリングなどの後処理を担当します。後処理が存在する意義は、この時点で、行列乗算の出力がまだオンチップレジスタに「生きている」状態であり、グローバルメモリに書き込まれていないことです。これは短い黄金の窓です。この瞬間にさらに計算を行うことができれば、メモリへの書き込みと読み出しの往復を完全に省くことができます。

CODAの中核的な洞察は、Transformerにおけるメモリ集中型の操作の多くが、代数的に再パラメータ化され、「後処理」のウィンドウに詰め込むことができるということです。これには少し数学的な技巧が必要です。最も一般的なGEMM-RMSNorm-GEMMパターンを例にとると、行列乗算の結果は、残差加算、RMS正規化を経て、別の行列乗算が行われます。従来の方法では、3つの独立した演算子が直列に実行され、中間結果が2回メモリに書き込まれます。CODAチームは、RMS正規化における行スケーリング因子rが、各行で共有されるスカラーであるため、後の行列乗算と交換法則を満たすことを発見しました。rの適用を「2番目のGEMMの前」から「2番目のGEMMの後処理」に遅らせることができます。遅延後、最初のGEMMの後処理は、局所的な「ブロック平均二乗根」（partial RMS）を計算するだけで済み、非常に軽量な補助削減カーネルによって結合され、完全なRMSNorm計算は消滅します。

同様の再パラメータ化は、SwiGLU、RoPE（回転位置エンコーディング）、交差エントロピー損失などの操作にも適用でき、逆伝播にも適用できます。論文には、前方後処理が「ブロックローカル」である限り、逆伝播は自動的に同じ構造を継承するという定理の証明があります。詳細については、元の論文を参照してください。

5種類の「積み木」と「レゴ言語」のセット。CODAは具体的な融合カーネルではなく、プログラミング抽象化のセットです。専門家によって最適化されたGEMMメインループを固定し、後処理の位置に、組み合わせ可能な5種類の基本プリミティブを公開します。要素ごとの変換（残差加算、活性化関数、RoPE）、ベクトルロードとストア（ブロードキャストRMSNorm重み）、行列ブロックロードとストア（逆伝播に使用する中間活性化の保存）、ブロック削減（局所平均二乗根、ブロックlog-sum-exp）、状態変換（オンライン正規化に必要なmaxおよびsum-exp統計）。これらの5種類の積み木を使用すると、標準的なTransformerの順伝播および逆伝播において、注意機構を除くほぼすべての操作をカバーできます。

さらに興味深いのは、この抽象化が「誰がコードを書くか」に対して寛容であることです。論文では、実験で2つの実装モードを評価しました。1つは人間のプログラマーが作成する方法、もう1つはClaude Codeを使用して生成する方法です。CODAのプリミティブの説明、いくつかの例、および実装ログが与えられ、AIがほとんどのカーネルコードを完成させ、人間が軽く監督します。どちらのモードでも、高いレベルのパフォーマンスが達成されました。Tri Daoはツイートで「LLMと初心者でも光速カーネルを作成できる」と述べていますが、これは論文の実験結果が現実のレベルで反映されたものです。

実験結果。CODAのベンチマークテストでは、cuBLASとtorch.compile、およびLLM用に最適化されたLiger KernelとFlashInferという、非常に厳しい相手が選択されています。論文では、各カーネルに対して2つの実装が評価されています。CODA (LLM)はClaude Codeによって生成され、研究者はプリミティブの説明、いくつかの例、および継続的に更新される実装テクニックのログを提供し、AIが主要なコードを完成させ、人間が軽く監督します。CODA (Human)は、人間のプログラマーが独立して作成し、同じ高レベルの再パラメータ化のアイデアを使用しますが、CODAプリミティブセット自体には依存しません。両方の結果は、cuBLAS + torch.compile、Liger Kernel、FlashInferなどの最適化ライブラリと比較されます。

単一演算子のレベルでは、GEMM-RMSNorm-GEMMという典型的なパターンを例にとると、CODAは対応する1B、7B、70Bの3つのモデルサイズの隠れ次元において、cuBLAS + PyTorchベースラインを上回っています。SwiGLU、RoPE、交差エントロピーなどの後処理の組み合わせも同様のパフォーマンスを示しています。LLMによって生成されたカーネルは、ほとんどのベンチマークで人間の手書きバージョンとほぼ同じであり、一部の構成ではわずかに上回っています。これは、GPUカーネルの最適化という、これまで非常に敷居の高い分野では、非常にまれな結論です。

逆伝播の利点は特に顕著です。GEMM-Residual-PartialRMS-GEMMの逆カーネルは、ベースラインと比較して1.6〜1.8倍の高速化を達成し、SwiGLUの逆カーネルも約1.4〜1.6倍の高速化を達成しています。この方向では、LLMと人間の実装の差もごくわずかです。これは驚くことではありません。逆伝播は本質的により多くの中間テンソルのアクセスを伴うため、後処理の融合による利点が大きくなります。また、CODAのプリミティブ設計は十分に明確であるため、AIモデルは組み合わせを正しく完了できます。

完全なTransformerレイヤーのエンドツーエンドのベンチマークでは、CODAの順伝播の高速化は、さまざまな規模で約5%〜20%であり、より大きなモデルサイズ（70B規模の隠れ次元に対応）では効果がより顕著です。数値精度に関しては、CODAの再パラメータ化により、RMSNormスケーリング因子の適用タイミングが調整されましたが、実験では、その数値誤差はPyTorchの参照実装と同程度であり、一部の構成では誤差がさらに小さくなっています。これは、GEMMメインループ自体がより高精度のアキュムレータを備えているためです。

CODAで何ができるか：速見表。より大きな視点に入る前に、CODAの能力の境界を明確にしておきましょう。カバー範囲：標準的なTransformer（LLaMAアーキテクチャなど）の順伝播および逆伝播において、注意機構と単語埋め込みを除くほぼすべての計算。これには、RMSNorm、残差加算、SwiGLU活性化、RoPE回転位置エンコーディング、交差エントロピー損失、および上記の操作の逆勾配計算が含まれます。高速化効果：1B〜70B規模の隠れ次元に対応して、単一演算子のレベルでは、cuBLAS + torch.compileベースラインと比較して、さまざまな程度の改善が見られます。その中で、逆伝播の利点が最も顕著です（一部のカーネルでは1.6倍以上）。完全なTransformerレイヤーのエンドツーエンドの順伝播の高速化は約5%〜20%であり、より大きなモデルサイズでは効果がより顕著です。誰が使用できるか：CODAはCuTeDSL（NVIDIA CUTLASSのPython DSL）に基づいて実装されており、人間のプログラマーとAIモデルの両方のカーネル作成方法をサポートしており、どちらの方法でも高性能を達成できます。現在の制限：現在、単一GPUのシナリオのみをサポートしており、分散トレーニングは含まれていません。再パラメータ化は主に標準的なTransformerアーキテクチャを対象としており、他のアーキテクチャへの適用可能性は検証されていません。

結論。CODAは孤立した作業ではありません。これは、ある種のアイデアの具体的な実装です。GPU上では、真の最適化スペースは多くの場合、「何を計算するか」ではなく、「どのように運ぶか」にあります。FlashAttentionは注意計算をオンチップメモリに「住まわせ」、CODAは正規化と活性化関数も「住まわせよう」としています。Tritonはカスタムカーネルの作成の敷居を下げ、ThunderKittens、TileLangなどは、さまざまなレベルでこのスペースをさらに探求しています。これらの作業はすべて、同じ方向を指しています。PyTorch演算子グラフの表現の利便性と、手書きのCUDAに近い実行効率を、プログラミング可能なフレームワークに真に統合することです。

Tri Daoのツイートの最後の言葉は、もう一度味わう価値があります。「LLMと初心者でも、すべてのTransformer操作に対して光速カーネルを作成できる。」この背景には、より深いロジックがあります。プログラミング抽象化が十分に設計されていれば、AIモデル自体が自身のトレーニングインフラストラクチャの最適化に参加できます。このサイクルこそが、CODAが最も興味深い点です。この観点から見ると、「CODA」という名前には別の意味があるかもしれません。クラシック音楽では、Codaは楽曲の最後に全体をまとめるセクションです。ここでは、GEMMカーネルの「後処理」です。そして、この後処理をうまく書くことこそが、Transformerトレーニングシステムの効率を向上させるための次の重要な章となるでしょう。

[机器之心]