PyTorch のプロファイリング（第 2 部）：nn.Linear から融合 MLP へ

[記事一覧に戻る](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::addmm（乗算と加算）演算の前に aten::t（トランスポーズ）演算があることに気づきます。これにより、nn.Linear は重みパラメータをトランスポーズしてから入力と乗算していることがすでに推測できます。これが aten::t 演算が表示される理由です。 注意すべき重要な点は、aten::t は実際にはデータをコピーしたり再編成したりするものではないということです。これは CPU 上で転置行列を表すためにテンソルのメタデータ（形状とストライド）を書き換えるだけであり、GPU でカーネルを起動することはありません。これを検証するには2つの方法があります。1つ目はトレースの GPU ラインを確認する方法、もう1つはプロファイラーテーブル内の aten::t 行を確認し、CUDA 上で要した時間をチェックする方法です。 ## なぜ個別の mul および add カーネルがないのか？ Figure 3: No aten::add in the profile of a linear layer 図3に示されるように、linear レイヤーのディスパッチチェーンには aten::add（バイアス加算）が存在しません。これは、バイアスの加算が「エピローグ」と呼ばれる仕組みを用いて行列乗算カーネルに *folded* （統合）されているためです。 **エピローグ**とは、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 が選択するカーネルです。この加算は独立したカーネルとして現れることはありません。なぜなら、それは**matmul カーネルの書き戻し（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 つのトレース（Eager モードとコンパイルモード）を注意深く読めば、Eager CPU のディスパッチチェーンにはコンパイル版よりも多くの要素が含まれていることに気づくでしょう。 図 4: aten::linear が aten::t (トランスポーズ) を経由して aten::addmm に至る Eager モードのディスパッチチェーン 図 5: トランスポーズを介さず、直接 aten::addmm が呼び出されるコンパイル版のディスパッチチェーン aten::linear 内部の Eager CPU のディスパッチチェーンは、まず aten::t に続き、次に aten::addmm です（図 4）。aten::t が実際に何を行うのかを理解するには、*ストライド (strides)* と *ビュー (views)* について少し立ち寄る必要があります。 テンソルは、メモリ上では数値の連続した 1 つのフラットな配列としてデータを格納しています。形状（shape）とストライドは、その配列の上に位置するメタデータであり、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 が行っていることです：新しいテンソルの割り当てやデータのコピーを行わず、書き換えられたストライドを持つ重みの *ビュー* を生成します。 図 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）。 Figure 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 ディスパッチが発生することを期待できます。また、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（バッチサイズ×シーケンス長 = 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 数、すなわち各々 2·8192·768·3072 ≈ 38.7 GFLOP を有していますが、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 です。 なぜこれが勝利となるのでしょうか？ eager モードでは、中間変数 h = gelu(g) は [8192, 3072] の bf16 テンソル（約 50 MB）であり、GeLU カーネルがこれを HBM に書き込み、mul カーネルが即座に読み戻します。融合により、このデータはレジスタ（チップ内部に存在し、HBM よりも近いメモリ）内に保持されます。Triton カーネルは g と u をそれぞれ 1 回だけ読み取り、gelu(g) * u を計算して結果を 1 回だけ書き込みます。これにより、中間データがグローバルメモリを経由する往復のラウンドトリップが完全に消滅します。 ## ハンドチューニングされたカーネルを使おう これまで私たちは PyTorch（eager モード）とコンパイラー（torch.compile）にカーネル選択を任せてきました。今度は、人間のエキスパートが手書きで調整したカーネルを組み込みます。ここでは「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つの具体的なことを指しており、どちらもトレース（trace）で確認できます。 Figure 13: コンパイル済み実行では、GEMM が実行される前に Dynamo、ガード、プロローグなどの事前演算のコストが発生する Figure 14: Liger カーネルには事前演算がない — 本来そこにあるべきボックスは空である - この融合は組み込まれています。LigerGEGLUMLP の forward は down_proj(LigerGELUMulFunction.apply(gate_proj(x), up_proj(x))) です。LigerGELUMulFunction は単一の Triton カーネル _geglu_tanh_forward_kernel を実行し、gelu(gate) * up を 1 パスで計算します。これは torch.compile から見たものと同じで、中間結果が HBM（High Bandwidth Memory）を経由して往復することはありません。図 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 をフルラウンドトリップする | 各 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 のプロファイリング（第 2 部）：nn.Linear から融合 MLP へ

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

調べる

選ぶ

サイト