Transformer のメモリ律速を解消──非 Attention 演算を GEMM 後処理として統合する CODA
Transformer ブロック内の正規化や活性化関数を GEMM のエピローグとして再定義し、GPU メモリへの冗長なアクセスを削減。(原題: CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs)
リリース: 2026-05-19 · 読了 5 分Transformer 学習において、正規化、活性化、残差接続などのメモリ律速(Memory-bound)な演算が、演算密度の高い GEMM 以外の実行時間の多くを占めている現状を指摘。
CODA は、これらの中間テンソルをグローバルメモリに書き出す前に、GEMM の出力タイルがオンチップ(SRAM/レジスタ)にある状態で処理を完結させる GPU カーネル抽象化である。
標準的な Transformer ブロックの順伝播および逆伝播における、Attention 以外のほぼすべての非線形演算をカバーする、構成可能なエピローグ・プリミティブ(縮約、変換、蓄積など)を提供。
人間および LLM が作成した CODA カーネルの両方で高いパフォーマンスを達成し、フレームワークの生産性とハードウェア効率の両立が可能であることを示した。
GPU の演算性能(FLOPS)向上に対しメモリ帯域の成長が追いつかない中、演算ごとにカーネルを分離する既存の設計はデータ移動のオーバーヘッドにより限界に達している。本手法を知らずに個別カーネルの最適化を続けても、システム全体の性能向上は頭打ちになる。
LLM によるカーネル生成(LLM-authored kernels)との親和性が高く、専門的な GPU プログラミングの知識がなくても、ハードウェアの限界に近い効率的な学習スタックを構築できる可能性がある。
LLM の学習・推論スタックを低レイヤーで最適化するエンジニアは、カスタム演算を単体カーネルとして実装するのではなく、GEMM のエピローグとして統合する設計パターンに移行すべき。CODA のような抽象化を用いることで、実装コストを抑えつつメモリ帯域のボトルネックを劇的に改善できる。
国内固有の追加文脈は限定的(汎用的に有用)。