pytorch/helion

Helion – Pythonに埋め込まれたGPUカーネル記述用DSL

概要 – Helionを使うと、GPU上で動作する小さなプログラムであるカーネルコードを、通常のPyTorchのような構文でPython内で直接記述できます。内部では、コードはTritonカーネル（高速なGPUマシンコードにコンパイルされる低レベル言語）に変換されます。

重要性 – 効率的なGPUカーネルを書くのは困難です。タイルサイズ、グリッド次元、メモリアクセスパターンなど、多くの低レベルな調整項目を決定する必要があります。Helionはその作業のほとんどを自動化します。単一のHelion関数は1つのGPUカーネルにコンパイルされますが、最初の起動前にオートチューナーが実行され、数百のTritonバリアントを試して最速のものを選択します。その結果、作成が容易でありながら、手動で調整されたパフォーマンスに近いカーネルが得られます。

主な機能

• PythonファーストのAPI – @helion.kernel()でカーネルを定義し、内部で通常のPyTorchテンソルを使用します。外側のPythonコードはCPU上で実行され（割り当て、形状チェック）、内側のループはGPUコードになります。
• 自動タイル化とグリッド選択 – hl.tileが反復空間を分割し、Helionがタイル次元、ループ順序、グリッドサイズを決定します。
• 豊富なオートチューニング – 大規模な構成空間（ブロックサイズ、ループ順序、インデックス戦略、ワープ特殊化など）を検索し、最適な構成をキャッシュして後続の実行に利用できます。また、独自の構成を指定して高価な検索をスキップすることもできます。
• メモリアクセス戦略 – ポインタベース、ブロックポインタ、TensorDescriptor（TMA）インデックスを自動的に選択し、上書きも可能です。
• リダクション、マスキング、永続カーネルのサポート – softmax、sum、大規模リダクションなどの一般的なパターンを追加コードなしで処理します。
• バックエンド – 主にTritonですが、実験的なTileIR（NVIDIA Blackwell GPU用）とCuTeバックエンド（CUDA 13+）もサポートしています。
• デバッグおよび開発ツール – 生成されたTritonコードの表示、再現可能なスクリプトの出力、インタープリターモードでのカーネル実行によるステップバイステップのデバッグが可能です。

典型的なワークフロー

1. @helion.kernel()とhl.tileを使用してカーネル関数を記述します。
2. CUDAテンソルで関数を呼び出します。最初の呼び出しでオートチューニングがトリガーされます（行列乗算の例では約10分）。
3. 必要に応じて最適なhelion.Configを抽出し、ハードコードして将来のチューニングを回避します。
4. 保存した構成でカーネルをデプロイし、起動時間を予測可能にします。

対象ユーザー – カスタムGPUカーネルが必要だが、TritonやCUDAを手書きしたくない研究者やエンジニア。PyTorchコードベースに自然に統合でき、既存ライブラリに特化した実装がないワークロードを高速化できます。

インストール – pip install helion（Python 3.10〜3.14、PyTorch ≥ 2.9、Triton ≥ 3.5が必要）。開発環境ではuvツールとCuTeバックエンド用のオプションの追加パッケージを使用します。

コミュニティ – ディスカッションはGPU MODE Discord（#helion）で行われています。このプロジェクトはガバナンスモデルに従い、BSDライセンスで提供されています。