PyTorch におけるプロファイリング（パート 2）：nn.Linear から融合 MLP へ

Hugging Face は、PyTorch におけるマルチレイヤーパセプトロン (MLP) の最適化プロセスを詳細に解説し、標準的な nn.Linear 層から融合カーネルへ移行して CPU のオーバーヘッドと GPU のメモリトラフィックを削減する方法を示しています。主なポイントは、torch.compile がポイントワイズ演算を単一の Triton カーネルに融合して高価な HBM の往復を回避できる一方で、手動でチューニングされたカーネルはコンパイル遅延や形状特化のオーバーヘッドなしに同等の性能向上を提供できる、ということです。

nn.Linear の仕組みとカーネルエピローグ

PyTorch では、nn.Linear は行列乗算と加算のラッパーです。bias=True の場合、演算は y = x @ w.T + b として実行されます。

エピローグによるバイアス折りたたみ

プロファイリングの結果、バイアス加算のための別個の aten::add カーネルは存在しないことが分かります。代わりに、バイアス加算は エピローグ を用いて行列乗算カーネルに折りたたまれます。エピローグとは、GEMM（General Matrix Multiply）カーネルが結果を High Bandwidth Memory（HBM）に書き戻す直前に実行する小さな計算です。バイアス加算を matmul カーネルの書き戻しに統合することで、PyTorch は HBM への二度目のロードと書き込みを回避し、メモリトラフィックを削減します。

CPU ディスパッチと転置ビュー

aten::t（転置）操作は、aten::addmm の前のイージャー CPU ディスパッチチェーンに現れます。しかし、これは GPU カーネルを起動しません。aten::t は単に CPU 上でテンソルのメタデータ（形状とストライド）を書き換え、同じ生データの新しいビューを作成するだけです。

torch.compile を単一の nn.Linear 層に適用しても、使用される GPU カーネルは変わりません—同じ cuBLAS GEMM カーネルが実行されます。その代わりに、コンパイル時に結果のストライドをハードコードすることで、転置ビューのディスパッチにかかる CPU オーバーヘッドを除去します。

GeGLU MLP のプロファイリング

GeGLU MLP は、3 つの線形射影（gate_proj、up_proj、down_proj）とポイントワイズ活性化シーケンス（GeLU の後に乗算）で構成されます。

イージャーモードの実行性能

イージャーモードでは、GeGLU MLP の順方向パスは 5 つの異なる GPU カーネルを起動します：3 つの GEMM と 2 つのポイントワイズカーネル（GeLU と乗算）。GeLU 演算で生成された中間テンソルは HBM に書き込まれ、続く乗算カーネルで再び読み込まれるため、メモリボトルネックが発生します。

GEMM カーネルのバリエーション

同じ FLOP 数であっても、すべての GEMM カーネルが同一というわけではありません。例えば、down_proj は gate_proj より速いことがあります。これは、cuBLAS が入力形状に基づいて異なるタイルサイズ（例：128x256 対 128x128）やパイプライン深度を選択し、データ再利用を最適化するためです。

torch.compile と融合カーネルによる最適化

torch.compile の影響

torch.compile は、GeLU、乗算、リシェイプ操作を単一の融合 Triton カーネル（例：triton_poi_fused__unsafe_view_gelu_mul_0）にまとめることで MLP を最適化します。この融合により、中間値をレジスタに保持し HBM への書き込みを回避することで大幅な性能向上が得られます。しかし、この最適化にはコストが伴います。コンパイラは TorchDynamo ガードやプロローグオーバーヘッドといった「前処理」を導入し、特定の入力形状に対してカーネルを特化させます。入力形状が変わると、モデルは再トレースと再コンパイルが必要になります。

kernels ライブラリによる手動チューニングカーネル

コンパイル遅延や形状特化を回避するために、Hugging Face の kernels ライブラリを通じて手動でチューニングされたカーネルを使用できます。例えば、LigerGEGLUMLP 層は、GeLU と乗算の融合が組み込まれた事前構築済みの Triton カーネルを使用します。

Compiled vs. Hand‑Tuned カーネルの比較

機能	torch.compile (Inductor)	手動チューニング (Liger)
融合	ポイントワイズ演算の動的融合	組み込み融合
CPU オーバーヘッド	Dynamo ガードとプロローグ	コンパイル前処理なし
形状処理	静的形状に特化（単一形状で最速）	汎用的な起動パラメータ（形状変化に強い）
デプロイ	ローカルコンパイル/Triton が必要	HF Hub 経由の事前ビルドバイナリ

最適化段階の概要

セットアップ	GPU の変更	CPU の変更
イージャー nn.Linear	バイアス加算が GEMM エピローグに折りたたまれる	ベースラインディスパッチ
コンパイル済み nn.Linear	変更なし（同じ cuBLAS カーネル）	aten::t ビューのブックキーピングを削除
イージャー MLP	5 つのカーネル；中間結果が HBM にヒット	ベースラインディスパッチ
コンパイル済み MLP	GeLU と乗算が 1 つの Triton カーネルに融合	Dynamo/ガードの前処理を追加
Liger MLP	コンパイルと同様の融合；チューニングされた起動パラメータ	コンパイル前処理なし；再コンパイルリスクなし