TPU 上の Helion：異種混在ハードウェアカーネル作成に向けて

Helionは、PyTorchの高性能でポータブルなMLカーネルを記述するための高レベルDSLです。Googleと協力してTPUバックエンドを構築し、HelionカーネルをPallasにコンパイルすることで、PyTorchユーザーがTPU上で効率的にMLワークロードを実行できる新たな道を開きました。

• GoogleのTPU v7（Ironwood）は、NVIDIA B200に匹敵する性能と潜在的なTCO削減を両立するMLコンピューティングプラットフォームです。
• HelionはPyTorchのハイレベルDSLであり、PallasにコンパイルされるTPUバックエンドが開発され、TPUカーネル開発を容易にします。
• Helionで生成されたフラッシュアテンションカーネルは、TPU v7上で約79%のMFU（最大浮動小数点演算性能）を達成しました。
• Helionのオートチューナーは、異なる入力形状に対して最適なパイプラインスキーマを選択し、TPUのVMEMと計算能力を最大限に活用します。