AI 推論與編譯器
PyTorch HelionにTPUバックエンドが追加、同一の高水準カーネルコードでTritonとPallasに対応
PyTorchとGoogleはHelion向けのTPUバックエンドを構築し、PyTorchスタイルのカーネルコードをPallasへコンパイルできるようにした。これにより、TPU用の低水準カーネルを別途記述する負担が軽減される。チームはTPU v7上のFlash Attentionで最大838 TFLOPsを測定したが、機能はなお活発な開発段階にある。
PyTorchチームは7月23日、機械学習カーネルDSLのHelionが、Google PallasをバックエンドとしてTPUプログラムを生成できるようになったと発表した。Helionは従来、Pythonに組み込まれたPyTorchに近いテンソル処理およびタイル処理のコードを、主にTritonへコンパイルしていた。今回新たな経路が加わったことで、エンジニアは同一の高水準カーネルを維持しながら、ハードウェアに応じてGPUまたはTPUバックエンドを選択できるようになり、TPU向けに大量のPallasコードを直接書き直す必要がなくなると期待される。
この取り組みの要点は、単にコードジェネレーターを置き換えることではない。TPUは少数の大規模な作業単位と、明確に区分されたHBMおよびVMEMのメモリ階層を採用しており、性能はソフトウェアによる非同期データ転送のスケジューリングに左右される。これはGPUのSIMTモデルとは異なる。そのためHelionは、外側のタイル処理を`pallas_call`パイプラインへ変換し、内側では`emit_pipeline`と`unroll`を自動的に比較する。前者はKとVをブロック単位でHBMからVMEMへ転送するため、より長いシーケンスを処理できる。後者はデータをあらかじめVMEMに保持することで演算のバブルを解消できる一方、シーケンス長の増加に伴ってメモリ要件も大きくなる。
公式によるB=8、H=32、D=256のFlash Attentionテストでは、8Kシーケンスでunrollを使用した場合に892 TFLOPsを記録した。一方、32KシーケンスではVMEM不足のため、この戦略を採用できなかった。オートチューナーは形状に応じてパイプライン方式へ切り替えられる。別のTPU v7上の測定では838 TFLOPsに達し、単一tensor coreのMFUで約79％に相当した。より広範なカーネル群では、HelionがTorchTPU eagerに対して幾何平均で1.55倍、`torch.compile`に対して1.12倍高速だったとチームは報告している。
これにより、アクセラレーターをまたぐカーネルのポータビリティーは、実用上の要件に一歩近づいた。ただし、これらの数値はプロジェクトチームが選定したカーネル、形状、ハードウェアに基づくものであり、モデル全体の学習や推論性能へ直接一般化することはできない。今後は、TPUバックエンドの対応範囲、コンパイルとオートチューニングのコスト、さらに同一のHelionカーネルが異なる世代のTPUやGPU上でも、手書きカーネルに近い性能を安定して実現できるかが注目される。