PyTorch Helion、TPU バックエンドで異種ハードウェア向けカーネル開発を前進

導入

PyTorch エコシステムは、カスタム ML カーネル開発をより高レベルで、かつ複数ハードウェアに移植しやすい形へ進めようとしている。PyTorch Blog が紹介した Helion の TPU バックエンドは、その流れを象徴する取り組みだ。開発者は PyTorch に近い書き味の Helion DSL でカーネルを書き、コンパイラが TPU 向けの Pallas コードを生成する。

TPU は GPU と似た用途で使われる一方、プログラミングモデルは大きく異なる。GPU が多数のスレッドとハードウェア管理キャッシュに依存するのに対し、TPU では HBM とオンチップ VMEM の間のデータ移動を明示的に意識し、計算と転送を重ねる設計が重要になる。

主要ポイント

• Helion は性能移植性を狙う DSL：高レベルのカーネル記述から、GPU や TPU など各ハードウェアに応じた実装へ落とし込むことを目指している。
• TPU バックエンドは Pallas を生成：TPU に詳しくない PyTorch ユーザーでも、低レベルの Pallas を直接深く書かずに最適化カーネル開発へ入れる可能性がある。
• 最適化の中心はパイプライン化：Helion は HBM から VMEM へのロードと、MXU やベクトルユニットでの計算をできるだけ重ねる Pallas コード生成を行う。
• Flash Attention では内側ループが難所：Q タイルを処理する外側ループに加え、K/V タイルを走査する内側ループがあり、ここで emit_pipeline を使うか、VMEM に事前取得して展開するかなどを自動チューニングする。
• 性能結果も示された：記事によれば、対象の Flash Attention ワークロードで Helion 生成カーネルは TPU v7 上 838 TFLOPs、単一 tensor core の約 79% MFU を達成した。

意義と影響

この発表のポイントは、TPU を PyTorch の高性能カーネル開発における本格的なターゲットとして扱う姿勢にある。最新世代の TPU と高性能 GPU は、BF16 演算性能や HBM 帯域などの重要指標で比較可能になっている。しかし、開発体験には依然として差があり、TPU ではメモリ階層とソフトウェアパイプラインへの理解が性能を大きく左右する。

Helion はその複雑さをコンパイラとオートチューナー側へ寄せる。これにより、Pallas の専門家でなくても性能重視のカーネル開発に参加しやすくなり、さらに TPU と GPU の両方を対象に同じ高レベル記述を維持できる可能性が出てくる。

もちろん、記事の性能値は特定のワークロードと入力形状に基づくもので、すべての TPU カーネルで同じ効率が得られるわけではない。それでも、AI インフラの競争がチップ単体の性能だけでなく、その性能を開発者がどれだけ扱いやすく引き出せるかに移っていることを示す事例といえる。

出典：PyTorch Blog