研究チームは、Transformer 学習システムにおけるボトルネックの解消を目指す新しい GPU カーネル抽象化「CODA」を開発しました。従来の Transformer 訓練は密度の高い線形代数を中心に構築されていますが、エンドツーエンドの処理時間の非自明な部分を占めるのは、正規化、活性化関数、残差更新、リダクションなどの周辺メモリ拘束演算です。これらの演算は、大規模な中間テンソルを主に演算量の少ないままグローバルメモリ間を移動させる必要があるため、最適化された学習スタックにおいてデータ移動が最大のボトルネックとなっています。 CODA は、これらの演算を GEMM（一般化行列乗算）加上後処理プログラムとして表現します。このアプローチの核心は、多くの Transformer 演算がフレームワーク固有のカーネルとして露出している際、GEMM の出力タイルをメモリに書き出す前にチップ上で処理できるという代数的再パラメータ化の観点に基づいています。GEMM のメインループを固定し、スケーリング、リダクション、ペアワイズ変換、累積などの合成可能な後処理プリミティブを最小限に制限するインターフェースを提供します。これにより、人間が手作業で書いた GEMM のような高性能な構造を維持しつつ、標準的な Transformer ブロックの順伝播および逆伝播における注意機構以外のほぼ全計算をカバーできる十分な表現力を備えています。 代表的な Transformer ワークロードにおける実証では、人間が開発した CODA カーネルに加え、LLM（大規模言語モデル）が生成したカーネルも高いパフォーマンスを発揮しました。これは、GEMM 加上後処理プログラミングが、フレームワークレベルの生産性とハードウェアレベルの効率を両立する実用的な道筋であることを示唆しています。本研究成果は、機械学習（cs.LG）分野の arXiv に論文として発表され、Transformer 学習の効率化における新たなパラダイムを提供するものと期待されます。データ移動を最小化し、計算資源をより有効活用するこの技術は、大規模モデルのトレーニング速度向上とコスト削減に貢献する可能性があります。