Helion、コンパイル時ヒューリスティクスでGPUカーネルを高速化

導入

GPUカーネルの性能は、単純な演算量だけでなく、ブロックサイズ、warps数、パイプライン段数、メモリ配置など多数の設定に左右される。候補をコンパイルして実測する自動チューニングは有効だが、カーネルごと、形状ごとに探索を繰り返すため、開発やデプロイ前の待ち時間が大きくなりやすい。PyTorchの機械学習向けDSLであるHelionは、これまで性能最適化を自動チューニングに大きく依存していた。今回の変更では、コンパイル時のプログラム解析から設定を推定し、探索を省略しても実用的な初期値を得られるようにした。

仕組みの要点

• 演算の種類で方針を分ける。 デバイスIRを走査し、行列積を含むカーネルはmatmul向け、行列積を含まずリダクションを含むものはreduction向け、それ以外はpointwise向けのヒューリスティクスに送る。matmulは標準的なGEMMだけでなく、hl.dotやtorch.baddbmmのような演算も対象にする。
• IRと生存期間を解析する。 どの軸がタイル化され、どのブロックサイズが変更可能か、各タイルが処理のどの段階まで生きているかを記録する。値が最後に使われる位置を調べ、レジスタや共有メモリなど、どこに置かれるかを近似する。
• ハードウェア資源を見積もる。 候補となるブロックサイズごとに、1つのCTAが必要とするレジスタ、共有メモリ、Blackwellではテンソルメモリの量を推定する。そこからSMあたりのCTA数、レジスタスピルの可能性、設定が起動可能かどうかを判断する。
• 設定同士のトレードオフを扱う。 パイプライン段数を増やすとロード遅延を隠しやすくなる一方、共有メモリを消費する。warpsを増やせばスレッドごとのレジスタ圧力を下げられるが、計算量が小さい場合は小さな行列積を多くのCTAで並列実行した方がよいこともある。

評価と意味

報告によると、ヒューリスティクスを有効にした場合のコンパイル時間との差は数十ミリ秒程度で、通常はTritonのコンパイル時間に埋もれる。H100上で9種類のlinear attention、96を超えるカーネルと形状の組み合わせを評価したところ、Flash Linear Attentionの手書きTriton実装に対して、forwardは1.18倍、backwardを含む場合は1.12倍だった。従来のHelionの自動チューニングなしの既定値に対しては2.13倍となった。

vLLMの6つの実運用向けリダクションカーネルを対象にした153を超える組み合わせの評価では、生成された設定は手書きのCUDA/C++オペレーターに対して1.27倍、従来の既定値に対して2.2倍だった。これは、コンパイル時の規則が常に完全な探索を上回るという意味ではない。重要なのは、性能と待ち時間の間に新しい選択肢を与える点にある。

自動チューニングを無効にすれば主ヒューリスティクスが既定値を生成し、有効にすれば補助ヒューリスティクスが約5〜20個の候補を探索の種として提供する。複数の行列積があるカーネルでは、ループで持ち越される累積値、動的な計算量、出力領域などから優先順位を付け、共有設定の衝突を解決する。Helionの設計は、性能最適化を全面的な試行錯誤から、プログラム構造とハードウェア資源モデルに基づくコンパイル時の判断へ移す試みといえる。

出典：PyTorch Blog