2026年6月11日 09:00
PyTorchのMLPカーネル融合を解説
Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP
3行まとめ
- •Hugging FaceがPyTorch解説記事を公開
- •MLPをTritonカーネルに融合し高速化
- •torch.compileとLigerを比較検証
詳細
背景
Hugging Faceが、PyTorchのパフォーマンス最適化を解説する技術ブログ連載の第2回を公開した。今回はnn.Linear単体の動作分析から始め、3つの線形層とGeLU活性化関数で構成するMLP（多層パーセプトロン）全体の最適化までを、プロファイラのトレースを読みながら追跡する内容となっている。
内容
nn.Linearではバイアス加算が行列積（GEMM）のエピローグとしてcuBLASの単一カーネル内で処理される仕組みや、入力の転置状態に応じて異なるコンパイル済みカーネルが選択される挙動を解説している。torch.compileは単一の線形層では効果が小さい一方、MLP全体ではGeLUと乗算を1つのTritonカーネルに融合し、中間テンソルをHBMではなくレジスタに保持することで高速化する。計測では融合部分の実行時間はコンパイル版が89.4マイクロ秒、Hugging Face Hub上で提供される手書きのLigerカーネル版が92.8マイクロ秒で、後者は形状変化時の再コンパイルが不要という利点を持つ。
意義
記事は「先に挙動を推測してからトレースで確認する」習慣の重要性を強調しており、GPUカーネルレベルの動作理解を深めたい開発者向けの実践的な教材となっている。
なぜ重要か
カーネル融合はLLMの学習・推論コスト削減の中核技術であり、PyTorch最適化の実践手順を具体例で学べる解説。
元記事を読む — Hugging Face Blog