PyTorch の Fused MLP を活用した最適化手法（29 分読了）

[記事一覧に戻る](https://huggingface.co/blog) このシリーズの「PyTorch でのプロファイリング」[第 1 部](https://huggingface.co/blog/torch-profiler)では、`torch.add(torch.matmul(x, w), b)` を用いて PyTorch プロファイラーのトレースを読み取る方法を学びました。また、CPU ディスパッチチェーン、起動オーバーヘッド、オーバーヘッドバウンドと計算バウンドのレジームの違い、そして `torch.compile` の内部構造など、いくつかの他のトピックについても議論しました。 2 回目の反復（今回のブログ記事）では、一歩階段を登ります。手書きの matmul-add ペアを `nn.Linear` (bias=True) に置き換えます。これはあらゆる深層学習モデルが使用するビルディングブロックです。その後、活性化関数を間に挟みながらこれらを 3 つ積み重ね（例に特化した構成）、Multilayer Perceptron (MLP: 多層パーセプトロン) ブロックを形成します。 ** このブログ記事のスクリプトはここにあります：[02_linear.py](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/02_linear.py)、[03_simple_mlp.py](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_simple_mlp.py)、および [03_kernels_mlp.py](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_kernels_mlp.py)。前回の投稿と同様に、これらを別タブで開きながらコードを追いかけて読むと理解が深まります。スクリプトの実行には NVIDIA A100-SXM4-80GB GPU を使用しています。Hugging Face のインフラ上で GPU をセットアップし、[Dev Mode with Spaces](https://huggingface.co/docs/hub/spaces-dev-mode) を用いてスクリプトを試しに動かすのは非常に簡単です。また、[Hugging Face Jobs pipeline](https://huggingface.co/docs/huggingface_hub/en/guides/jobs) を利用してスクリプトを実行することも可能です。 始める前に、後で繰り返し参照することになる 2 つの概念を簡単に復習しておきましょう： - GPU カーネル（GPU kernel）とは、GPU の多数のスレッド上で並列実行されるプログラムのことです。 - CPU はこれらのカーネルのスケジューリングと起動を担当します。プロファイラトレースで目にする PyTorch のオーバーヘッドの大部分は、このスケジューリング作業に起因するものです。 ## matmul-add から Linear へ nn.Linear は、[Part 1](https://huggingface.co/blog/torch-profiler) で既にプロファイル済みの行列乗算と加算をラップしたモジュールです。唯一の違いは、重み（weight）とバイアス（bias）をパラメータとして保持し、PyTorch ユーザーが慣れ親しんだ forward メソッドを公開している点にあります。 # bias=True は、シリーズの第 1 部で確認した乗算と加算演算を真にエミュレートします linear_layer = nn.Linear(in_dim, out_dim, bias=True) y = linear_layer(x) ここで扱っている演算は以下のように記述できます: y = x @ w.T + b ここで、x は入力、w は重み、b はバイアスです。[02_linear.py](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/02_linear.py) を実行してプロファイルを確認してみましょう。 uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64 uvx trace-util traces -b traces [trace-util](https://x.com/ariG23498/status/2054811716727517374) は、トレースを [Hugging Face bucket](https://huggingface.co/storage) に同期し、ターミナル上で [Preffeto URLs](https://perfetto.dev/) を提供するユーティリティです。 Figure 1: Profiler trace of nn.Linear Figure 1 は、線形層のフォワード呼び出しのプロファイラトレースを示しています。この線形層のフォワード呼び出しは、以前のトレースと同様のスケジュール設定（wait=1, warmup=1, active=3）でトレースされています。そのため、CPU および GPU ラーンに 3 つの Profile Steps が表示されます。 ## トランスポーズは何をしているのか？ Figure 2: The transpose CPU row Figure 2 のようにプロファイラトレースを拡大すると、aten::t（トランスポーズ）演算子が aten::addmm（乗算と加算）演算子の前にあることに気づきます。これにより、nn.Linear は重みパラメータをトランスポーズしてから入力と乗算していることがすでに推測できます。これが aten::t 演算子が表示される理由です。 注意すべき重要な点は、aten::t は実際にはデータをコピーしたり再編成したりするものではないということです。これは CPU 上で転置行列を表すためにテンソルのメタデータ（形状とストライド）を書き換えるだけであり、GPU でカーネルを起動することはありません。これを検証するには2つの方法があります。1つ目はトレースの GPU ラインを確認する方法、2つ目はプロファイラーテーブル内の aten::t 行を確認し、CUDA 上で要した時間をチェックする方法です。 ## なぜ個別の mul および add カーネルがないのか？ Figure 3: リニア層のプロファイルに aten::add は存在しない 図3に示されるように、リニア層のディスパッチチェーンには aten::add（バイアス加算）は存在しません。これは、バイアスの加算がエピローグと呼ばれる手法を用いて行列乗算カーネルに*折り畳まれている*からです。 **エピローグ**とは、GEMM（GEneral Matrix Multiply：汎用行列乗算）カーネルが結果を HBM（High Bandwidth Memory：GPU のメインメモリ）へ書き戻す直前に行う小さな計算処理のことです。バイアスの加算、活性化関数の適用、定数によるスケーリングなどはすべて古典的なエピローグの例です。エピローグの目的は、メモリアクセスが演算コストを高めるため、HBM への読み込みまたは書き出しを2回行うことを回避することにあります。 nn.Linear は torch.nn.functional.linear を呼び出し、それがさらに aten::linear を呼び出します。aten::linear は入力を確認し、バイアスが渡されていることに気づくと、行列乗算と加算を別々に行うのではなく、代わりに aten::addmm(bias, x, weight) をディスパッチします。addmm は以下のように計算を行います。 out = x @ weight.T + bias GPU で実行される cuBLAS GEMM カーネルには、バイアス加算バリアントが組み込まれており、これが aten::addmm が選択するカーネルです。この加算は独立したカーネルとして現れることはありません。なぜなら、それは**行列乗算カーネルの書き戻し（writeback）の一部**であり、まさにエピローグ（epilogue）が果たす役割そのものだからです。 ここで注意すべき微妙な点があります。[Part 1 の --compile 下](https://huggingface.co/blog/torch-profiler#did-we-fuse-the-matmul-and-add-kernels-into-one) でご覧になったカーネル（addmm）は、すでに eager モードの nn.Linear が使用しているものです。torch.compile にはここで融合させるべきものが残っていません。これが次に検証する内容です。 ## --compile は単一の Linear に役立つか？ 順伝播呼び出しをコンパイルして、プロファイラトレースを確認してみましょう。（プロファイラトレースは[次のセクション](#where-did-the-transpose-go-kernel-layouts-and-pre-ops)で可視化されます） uv run 02_linear.py --batch 1024 --in_dim 32 --out_dim 64 --compile uvx trace-util traces -b traces 単一の nn.Linear の順伝播における eager モードとコンパイルされたモードのトレースを比較すると、以下のことがわかります。 - GPU 上では同じ cuBLAS GEMM カーネルが実行される。 - CPU 上では同じ aten::addmm オペレーション（op）が使用される。 - CPU ラーンには、compile に固有の追加行がいくつか存在する。 これは深く理解しておく価値があります。モデルが遅いと感じると、誰もが反射的に torch.compile を使おうとする傾向があります。しかし、バイアス付きの単一の GEMM 演算においては、compile の効果は非常に限定的です。これはバグではなく、compile が何らかの融合を行うためには複数の操作が必要であるというだけの話です。これを証明するために、[MLP（多層パーセプトロン）を見てみましょう](#stacking-two-linears-the-mlp)。 ## トランスポーズはどこへ行ったのか？カーネルレイアウトと事前演算 2 つのトレース（イーガーモードとコンパイルモード）を注意深く読めば、イーガー CPU ディスパッチチェーンにはコンパイル版よりも多くの要素が含まれていることに気づくでしょう。 図 4: aten::linear が aten::t（トランスポーズ）を経て aten::addmm を通るイーガーディスパッチチェーン 図 5: トランスポーズを介さず、直接 aten::addmm が呼び出されるコンパイルディスパッチチェーン aten::linear 内部のイーガー CPU ディスパッチチェーンは、まず aten::t に続き、次に aten::addmm です（図 4）。aten::t が実際に何を行うのかを理解するには、*ストライド*と*ビュー*について少し立ち寄る必要があります。 テンソルはメモリ上では、1 つの平坦で連続した数値の列としてデータを格納しています。形状（shape）とストライド（stride）は、その列の上に位置するメタデータであり、PyTorch がどのように走査するかを指示します：ストライドが (s0, s1) である場合、「行を 1 つ進むには s0 要素分ステップし、列を 1 つ進むには s1 要素分ステップする」という意味です。メタデータを変更すれば、コピーなしで同じ生データを異なる*ビュー*として取得できます。 >>> M = torch.tensor([[0, 1], ... [2, 3], ... [4, 5]]) >>> M.shape, M.stride() (torch.Size([3, 2]), (2, 1)) # 行あたり 2 ステップ、列あたり 1 ステップ >>> T = M.t() # トランスポーズ >>> T.shape, T.stride() (torch.Size([2, 3]), (1, 2)) # 形状とストライドが入れ替わり、データは不変 >>> T tensor([[0, 2, 4], [1, 3, 5]]) >>> T.flatten() # 強制的に実体化されるため、データ順序が入れ替わる tensor([0, 2, 4, 1, 3, 5]) M.t() は数値を一つも移動させませんでした。これはストライドを入れ替えた新しいビューを返すものであり、行ごとに読み込むと、元のバッファの 0, 1, 2, 3, 4, 5 が転置された順序で走査されることになります。基盤となるデータは同一であり、異なるのはメタデータだけです。 これはまさに linear レイヤー内部で aten::t が行うことです：新しいテンソルを割り当てたりデータをコピーしたりするのではなく、書き換えられたストライドを持つ重みの *ビュー* を生成します。 Figure 5 で確認できるように、compile は GPU カーネルを削除したわけではありません。そのビューをディスパッチするための *CPU オーバーヘッド* を削除しました。Inductor はコンパイル時にビューチェーンを追跡し、結果のストライドを一度計算して、それらのストライドがハードコードされた直接の aten::addmm 呼び出しを生成しました。GPU が同じ数学演算を行う間、数マイクロ秒分の CPU 作業が消滅します。 当然のことながら、入力データがコンパイラによって事前に計算されたストライドに違反すると、エラーが発生します。 両方のトレースにおける GPU ラーンを見ると、フォワードパスごとにカーネルはちょうど一つであり、その二回とも *同じ* カーネルです： cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8 もし転置カーネルが実行されなかったなら、GEMM が重み行列を転置された順序で読み込むように誰が教えたのでしょうか？その答えはカーネル名にあります。接尾辞を見てみましょう： cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8 ^^ この tn がレイアウト記述子です。cuBLAS と CUTLASS は、入力レイアウトの各組み合わせに対して *個別のカーネルバイナリ* を事前コンパイルしています。 n（非転置）と t（転置）は、カーネルが内部ループ内で入力をどのように走査するかを示します。ディスパッチャの役割は、入力ストライドを確認し、どの接尾辞の組み合わせが一致するかを判断して、適切な事前コンパイル済みカーネルを選択することです。 ** プロファイラトレースにおけるカーネル名は、カーネルのアイデンティティのハッシュダンプです。2 つの実行で同じカーネル名が表示される場合、GPU は同じ作業を行っています。名前が異なる場合（例：_tn_ と _nn_、bf16 と fp16、または s16816gemm と s161616gemm）、GPU は異なる作業を行っており、ディスパッチャが別の分岐を選択したことを意味します。この名前を読み解く方法を学ぶことは、トレースを比較する際に最も有用な習慣の一つです。 ## 3 つの Linear をスタック：MLP このセクションでは、Multilayer Perceptron（MLP）をプロファイルします。より興味深いものにするため、GeGLU アクティベーション変種を持つフィードフォワードネットワークをプロファイルします（これは実際には非常に頻繁に使用されています）。また、これは深層学習研究の歴史において書かれた最も優れた行の一つへのオマージュでもあります（図 6）。 図 6: [GLU Variants Improve Transformer](https://arxiv.org/abs/2002.05202) ペーパーの結論セクション。 class SimpleGeGLUMLP(nn.Module): def __init__(self, dim, hidden): super().__init__() self.gate_proj = nn.Linear(dim, hidden, bias=False) self.up_proj = nn.Linear(dim, hidden, bias=False) self.down_proj = nn.Linear(hidden, dim, bias=False) def forward(self, x): g = self.gate_proj(x) u = self.up_proj(x) h = F.gelu(g, approximate="tanh") m = h * u y = self.down_proj(m) return y 完全なスクリプトはここで見つけることができます：[03_simple_mlp.py](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_simple_mlp.py)。以下のように実行してください: uv run 03_simple_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072 uvx trace-util traces -b traces トレースを開く前に、一緒に何を期待すべきか考えてみましょう。forward 関数はかなりの計算量を行いますが、その大部分はすでに私たちに馴染みのあるものです。 各 nn.Linear レイヤーに対して 1 つずつ、合計 3 つの aten::linear ディスパッチ（dispatch）が発生することを期待できます。また、GeLU と乗算演算それぞれに対して 1 つずつ、計 2 つのポイントワイズカーネル起動も予想されます。トレースを見る前にこの期待を形成することは、プロファイリングの旅において最も有用な習慣の一つです：トレースを読むのは、ゼロから推測を作るためではなく、仮説を*確認または反証するため*に行うものです。 Figure 7: GeGLU MLP のプロファイラトレース Figure 8: 線形投影 CPU ラーンにおける「occupancy query」のハイライト Figure 7 から、私たちの直感が正しかったと自分自身を褒め称えることができます。1 つのフォワードパス（1 つの mlp_fwd）あたり、GPU は正確に 5 つのカーネルを実行します。Figure 8 では、線形投影レイヤーの CPU ラーンで確認できる「occupancy query」がハイライトされています。 Op CPU op GPU kernel launches gate_proj aten::linear ampere_bf16_s16816gemm_bf16_128x128_... occupancy query + cudaLaunchKernel up_proj aten::linear ampere_bf16_s16816gemm_bf16_128x128_... occupancy query + cudaLaunchKernel gelu aten::gelu vectorized_elementwise_kernel cudaLaunchKernel h * u aten::mul vectorized_elementwise_kernel cudaLaunchKernel down_proj aten::linear ampere_bf16_s16816gemm_bf16_128x256_... occupancy query + cudaLaunchKernel 3 つの GEMM（行列積演算）は、それぞれ起動前に `cudaOccupancyMaxActiveBlocksPerMultiprocessor` 呼び出しを余分に行っています。これについては Part 1 で別のセクションを設けており、[こちら](https://huggingface.co/blog/torch-profiler#why-does-matmul-have-an-extra-cuda-runtime-call)で確認できます。これは cuBLAS がグリッドサイズを決定しているためです。一方、ポイントワイズ演算（GeLU と mul）は、occupancy クエリを行わずに直接起動されます。つまり、「線形演算」は実際にはクエリと起動の両方を含みますが、「ポイントワイズ演算」は単なる起動のみです。 Figure 9: The table shows that some ops launch zero kernels `aten::t`, `aten::transpose`, `aten::reshape`, `aten::view`, `aten::as_strided`, および `aten::_unsafe_view` の各演算は、ゼロ個のカーネルを起動します。これらは表（Figure 9）で CUDA 時間が 0.000us と表示されていますが、これは CPU 上でテンソルのメタデータ（形状とストライド）を書き換えるだけだからです。表をスキャンする読者は、線形演算ごとに約 6 つの演算名を目にしますが、GPU に到達するのはそのうちの一つ（mm）だけです。 ## なぜ GEMM カーネルには 2 種類あるのか？ MLP は、行列積のために [batch, seq, dim] を [batch * seq, dim] にフラット化します。コマンドラインでの呼び出しではバッチサイズに 64、シーケンス長に 128 を使用したため、以下の 8192（batch * seq = 64 * 128）という値がここから来ています。 トレースより: Linear aten::mm input dims M·K·N cuBLAS kernel avg CUDA gate_proj [8192,768] x [768,3072] 8192·768·3072 …128x128…stages_32x5_tn 0.19ms up_proj [8192,768] x [768,3072] 8192·768·3072 …128x128…stages_32x5_tn 0.19ms down_proj [8192,3072] x [3072,768] 8192·3072·768 …128x256…stages_64x3_tn 0.17ms これら 3 つの GEMM（General Matrix Multiplication、一般行列積）はすべて同じ FLOP 数を持ち、それぞれ約 38.7 GFLOP（2·8192·768·3072）ですが、down_proj は約 10% 高速です。同じ計算量でありながら形状が異なる（N が 3072 ではなく 768）ため、cuBLAS はその形状に最適なタイル（128×256 で、より深い stages_64x3 パイプラインを持つもの）を選択し、リユース性を高めています。 タイル処理について深く学びたい場合は、[こちら](https://alvinwan.com/how-to-tile-matrix-multiplication/) が素晴らしい入門リソースです。 まさにこの理由から、表には 2 つの GEMM の行（Figure 9）が記載されています：128x128 の行は gate と up を示し、128x256 の行は down を示しています。 ## torch.compile は何をするのか？ フォワードメソッドをコンパイルしてトレースを可視化する前に、再び頭の中で「トレースで何を期待するか」を考える演習を行いましょう。これは楽しい実験であり、自分で何かをプロファイリングするたびに繰り返すべき重要な作業です。常に直感を頼りにし、何かが予想と一致しない瞬間には立ち止まってその理由を探る癖をつけましょう。 uv run 03_simple_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072 --compile uvx trace-util traces -b traces Figure 10: The profiler trace for the compiled GeGLU MLP Eager モードでは、各 nn.Linear が dispatcher ops の連鎖（aten::linear → aten::t → aten::transpose → aten::matmul → aten::reshape → aten::mm）に展開されます。これらは、ATen が実際の GEMM に到達する前に通る高レベルなラッパーです。torch.compile はこの連鎖を除去します。 コンパイルされたグラフが実行される頃には、linear も matmul も transpose も reshape も存在せず、それらのメタデータ演算は mm の呼び出し方に折り畳まれています。Figure 10 に示すように、3 つの裸の aten::mm 外部呼び出しが見られます。これが同じ GEMM であることの証拠は、カーネル名が eager モードとバイト単位で完全に一致していることです：ゲートおよびアップ用には ...128x128...stages_32x5_tn、ダウン用には ...128x256...stages_64x3_tn です。 ## The fused Triton kernel Figure 11: The fused Triton kernel これはコンパイルの教訓全体の要です。2 つの eager ポイントワイズカーネル（GeLU と mul）と reshape が、1 つのカーネル triton_poi_fused__unsafe_view_gelu_mul_0 に統合されました（Figure 11）。この名前を解読してみましょう： - triton: Inductor の Triton バックエンドによって生成されたもの（cuBLAS でも ATen でもありません）。 - poi: pointwise（Inductor はポイントワイズカーネルに poi、リダクションに red、永続的リダクションに per というタグを付けます）。 - fused__unsafe_view_gelu_mul: 統合された演算：_unsafe_view（reshape）、GeLU、および mul です。 - 0: グラフ内の一意な ID です。 なぜこれが勝利となるのか？イーグルモードでは、中間変数 h = gelu(g) は [8192, 3072] の bf16 テンソル（約 50 MB）であり、GeLU カーネルがこれを HBM に書き込み、mul カーネルが即座に読み戻します。融合により、このデータはレジスタ（チップ内部に存在し、HBM よりも近いメモリ）内に保持されます。Triton カーネルは g と u をそれぞれ 1 回だけ読み取り、gelu(g) * u を計算して結果を 1 回書き込みます。これにより、中間変数がグローバルメモリを経由する往復のラウンドトリップが完全に消滅します。 手動調整されたカーネルを使ってみましょう これまで PyTorch（イーグルモード）とコンパイラー（torch.compile）にカーネル選択を任せてきました。今度は人間のエキスパートが手書きで調整したカーネルを組み込みます。ここでは、kernels ライブラリを使用して [Hugging Face Hub](https://huggingface.co/kernels/kernels-community/liger-kernels) から簡単に取得できる LigerGEGLUMLP レイヤーを使用します。 from kernels import get_kernel kernels_layers = get_kernel("kernels-community/liger-kernels", version=1).layers kernels_geglu_mlp = kernels_layers.LigerGEGLUMLP(Config()).to(device, dtype=torch.bfloat16).eval() 完全なスクリプトはこちらです：[03_kernels_mlp.py](https://huggingface.co/datasets/ariG23498/profiling-pytorch/blob/main/03_kernels_mlp.py)。 uv run 03_kernels_mlp.py --batch 64 --seq 128 --dim 768 --hidden 3072 uvx trace-util traces -b traces Figure 12: The profiler trace for the LigerGEGLUMLP layer Figure 12 shows the profile for the LigerGEGLUMLP layer using the Liger kernels from the Hub. なぜカーネルライブラリを使用するのか Triton や CUDA でカーネルを書くことは一つの課題ですが、それらを配布することは別の課題です。カーネルは、GPU アーキテクチャ、CUDA バージョン、PyTorch バージョンの組み合わせが完全に一致するようにコンパイルされなければなりません。これは通常失敗するステップです（「自分のマシンでは動く」、nvcc の欠如、間違った Triton バージョンなど）。 [kernels](https://github.com/huggingface/kernels) ライブラリは、このビルドステップをあなたのマシンから外します。get_kernel("kernels-community/liger-kernels", version=1) は、Hugging Face Hub から事前ビルドされバージョンが固定されたカーネルパッケージをダウンロードし、ローカル（ここでは ~/.cache/...kernels-community--liger-kernels 以下）にキャッシュします。その利点は以下の通りです。 - カーネルは CI 内で一度だけコンパイルされ、多くのアーキテクチャとバージョンの組み合わせに対応しています。自分でコンパイルするのではなく、正しいバイナリをダウンロードできます。 - version=1 は正確なビルドを固定するため、スクリプトを実行する誰もが同じカーネルを使用します。「パッケージを更新したら遅くなった」という現象は発生しません。 - パッケージには .layers 属性が公開されており、ドロップインで使える nn.Modules（LigerGEGLUMLP など）が含まれています。モデルのモジュールをそれらに置き換えるだけで、モデルの他の部分は一切変更する必要がありません。 ## なぜチューニングされたカーネルが優れているのか 「チューニングされた」と言うとき、私たちは2つの具体的なことを指しており、どちらもトレースで確認できます。 Figure 13: コンパイル実行では、GEMM が実行される前に Dynamo、ガード、プロローグなどの事前演算のコストが発生します Figure 14: Liger カーネルには事前演算がありません — 通常そこにあるべきボックスは空です - この融合は組み込まれています。LigerGEGLUMLP の順伝播は、down_proj(LigerGELUMulFunction.apply(gate_proj(x), up_proj(x))) です。LigerGELUMulFunction は単一のTritonカーネル_geglu_tanh_forward_kernel を実行し、gelu(gate) * up を1回のパスで計算します。これは、中間データがHBM（High Bandwidth Memory）を経由して往復しないという、torch.compile で見たのと同じ結果です。図 13 と図 14 に示されているように、ここではコンパイラなしでこの恩恵を得られます（Dynamoガードなし、コンパイル遅延なし、再コンパイルリスクなし）。 - 起動パラメータはハードウェアに合わせて選択されています。カーネルがランダムにブロックサイズを推測することはありません。Liger の calculate_settings は、列数からこれらの値を選択します。 ここでトレードオフについて正直になる価値があります。なぜなら、生データだけを見ると誤解を招く可能性があるからです。Liger カーネルの実行時間は**92.8 µs**ですが、コンパイル実行におけるInductorの融合カーネルは**89.4 µs**です。一見すると手書きのカーネルがわずかに遅いように見えますが、この比較にはその価値を正当化するコストが見落とされています。 torch.compile は**静的形状**に対して特別化されます。Inductor の 89.4 µs のカーネルが高速なのは、まさに*この特定の*[8192, 3072]の問題のために生成されたからです。バッチサイズやシーケンス長、隠れ次元を変更すると、Dynamo は再トレースを行い、新しい特別化カーネルを得るために再びコンパイルコストを支払う必要があります。 つまり、真の選択は「遅い人手によるカーネル vs 高速なコンパイル済みカーネル」ではありません。それは**汎用的に高速なカーネル vs 特定の1つの入力形状に特化したカーネル**です。Liger カーネルは1セットのパラメータを起動パラメータとして受け取り、再コンパイルを行うことなく*あらゆる*形状に対して実行します。これは、形状ごとの特化がもたらす数マイクロ秒の速度向上を犠牲にする代わりに、変化する形状に対する堅牢性を獲得するものです。 ## 結論 以下の表は、各ステップで GPU に何の変更を加え、何が変更されなかったかをまとめています。 Setup | What changed | What stayed the same ---|---|--- Eager nn.Linear | ベースライン：バイアス加算はすでに GEMM エピローグ（addmm）に折り込まれており、*1 つの* cuBLAS カーネルであり、行列乗算と加算の組み合わせではない | — Compiled nn.Linear | いくつかの CPU ディスパッチ演算（aten::t ビューのブックキーピング）が消失 | 同じ単一の cuBLAS GEMM カーネルで、バイト単位でも同一。コンパイルでは融合は行われない Eager MLP | GPU カーネル5個：3 つの GEMM + GeLU + 乗算。[8192, 3072] の中間結果が HBM（High Bandwidth Memory）をフルラウンドトリップする | 各 GEMM は、独立した線形演算と同じバイアスなし cuBLAS カーネルである Compiled MLP | GeLU + 乗算 + リシェイプが**1 つの**融合された Triton カーネルに統合され、中間結果はレジスタ内に残る。コンパイル前処理（Dynamo、ガード）のコストを払う | 3 つの GEMM は変更されず、cuBLAS カーネル名も同一 Liger MLP | 同じ融合だが、ハードウェアチューニングされた起動パラメータを持つ手書きの Triton カーネルに組み込まれており、**Dynamo やガード、コンパイル遅延は不要** | 3 つの GEMM は依然として同じ cuBLAS カーネルである 一つだけ継続すべき習慣があるとすれば、それはトレースの前に行ってきたもの、すなわち**まず推測し、その後確認する**というものです。トレースに何が含まれると予想するかを述べ、それを開き、不一致が生じた場合は画面で最も興味深い出来事として扱ってください。 これは「PyTorch におけるプロファイリング」というシリーズの二番目のステップでした。次の投稿では、この MLP ブロックからアテンションブロックへと、そして最終的には完全なモデルへと、さらに階段を上って進んでいきます。 記事の初期草案に対するレビューを寄せてくださった [Noe Flandre](https://huggingface.co/NoeFlandre) 氏と [Pedro Gabriel Gengo Lourenço](https://huggingface.co/pedrogengo) 氏に感謝いたします！

PyTorch の Fused MLP を活用した最適化手法（29 分読了）

背景や根拠まで確認しますか？

調べる

選ぶ

サイト