カーネル融合

PyTorch のようなライブラリで書いたコードは、1行ずつ読めます。この2つの配列を足す、それを掛ける、指数を取る、という具合です。書かれたとおりに実行すると、それぞれの行が別々のGPUカーネルになります。各カーネルは GPU のメモリから入力を読み、ほんの少しの計算をして、結果を書き戻します。そして起動のたびに、数マイクロ秒の固定の立ち上げコストがかかります。

配列が小さいと、この立ち上げコストが時間の大半を占めます。次のような1行は

$r = e^{-40\,\overline{(q - q_{\text{ref}})^2}}$

4つのカーネル（引き算、二乗、平均、指数）になり、それぞれ一瞬で終わっては、次のカーネルが始まるのを待ちます。

カーネル融合は、このような連鎖を1つのカーネルにまとめます。$q$ と $q_{\text{ref}}$ を一度だけ読み、4つの演算をすべて GPU の高速なレジスタ上で行い、$r$ を一度だけ書き出します。起動の回数も、メモリとの往復も減ります。

torch.compile はこれを自動で行います。Python のコードをたどって演算のグラフにし、どれをまとめられるかを判断し、GPU 用の新しい融合カーネルを（Triton という言語で）書き出します。初回は数分かかるので、学習ループのように何百万回も実行されるコードで元が取れます。たどれないコード、たとえば別の GPU ライブラリの呼び出しはそのまま残され、コンパイルされた部分がその周りで動きます。

カーネル融合とCUDAグラフは、同じ無駄に両側から取り組みます。グラフは CPU が各カーネルを送るコストをなくし、融合はカーネルの大半をなくします。

このプロジェクトでは

2人のファイターと視覚 では、世界のコピー 2,048 個に対して、対戦相手環境の1ステップで約13,700個のカーネルが動いていました。MuJoCo Warp の物理計算が約1,800個、PyTorch で書かれた報酬、観測、リセットのための小さなカーネルが11,700個です。その PyTorch のコードを（物理計算の呼び出しは除いて）コンパイルすると、RTX 5090 での1ステップが 79.5 ms から 60.2 ms に、コピー 1,024 個では 58.7 ms から 41.7 ms に短くなりました。代わりに、各ランの開始時に約4分のコンパイルがかかります。