便利なソフトをAIが実際に動かしてわかりやすくレポート

🤖 AI自動レポート column

PyTorch 感覚で書いて GPU で爆速に——Helion が Triton カーネル開発を自動化する

タグ HelionGPUカーネルTritonPyTorchオートチューニング機械学習CUDAPythoncolumnコラムlinuxLinuxGitHubオープンソースpytorchhelion
📦 GitHubで見る

PyTorch 感覚で書いて GPU で爆速に——Helion が Triton カーネル開発を自動化する

ひとことでいうと

Helion は、GPU（画像処理や AI 計算に使う専用チップ）向けの計算プログラムを、普段の Python コードに近い書き方で作れるオープンソースのツールです。PyTorch（機械学習でよく使われる Python ライブラリ）の書き方をそのまま活かしながら、高速な GPU カーネル（GPU に直接命令する処理の最小単位）を自動でコンパイルしてくれます。面倒なインデックス計算や最適化の細かい設定を Helion が自動で担当するため、コードを書く量が大幅に減ります。GPU 向け機械学習の性能を追求したい開発者・研究者に向いており、BSD ライセンスで無料で利用できます。

こんな人におすすめ

1. GPU カーネルを手書きするのが大変な開発者

Triton（GPU カーネルを書くための低レベルな言語）を直接扱うには、ストライド計算やマスキングといった細かい処理を手動で実装する必要があります。Helion を使えば、行列の掛け算や softmax（機械学習でよく使う数値計算）のようなカーネルを短いコードで記述できます。さらに、何百通りもの実装パターンから最速の設定を自動で見つけてくれる機能も備わっています。

2. 機械学習システムの研究者

Helion はブロックサイズやループの順序、GPU 内部の並列化方式などを差分進化アルゴリズム（多数の候補を比較して最適解を見つける手法）で自動探索します。新しい GPU アーキテクチャに対してどのくらいの性能が出るかを評価する研究にも活用できます。単一の Helion カーネルから数百通りのコンフィグを体系的に試せるのが強みです。

3. 本番環境向けのカーネルチームのエンジニア

オートチューニング（最適設定を自動で見つける機能）を事前に済ませ、その結果をコードに固定することで、実際にサービスを動かすときのチューニングコストをゼロにできます。NVIDIA の Hopper・Blackwell や AMD の GPU にも対応しており、幅広いハードウェアで利用できます。

インストール・使い方

Step 1 — 動作環境を確認する

Helion を使うには、以下の環境が必要です。

  • OS: Linux（Windows・macOS は非対応）
  • Python: 3.10〜3.14
  • PyTorch: 2.9 以降
  • Triton: 3.5 以降
  • GPU: CUDA 対応の NVIDIA GPU（または AMD GPU）

まず PyTorch と Triton を公式の手順でインストールしておいてください。NVIDIA Blackwell GPU 向けに TileIR バックエンドを使う場合は、追加で Triton-to-tile-IR も必要です。必要な依存関係を先に整えておくことで、インストール後すぐに動作確認できます。

Step 2 — Helion をインストールする

ターミナル（文字でパソコンに命令を送る画面）を開き、以下のコマンドをコピー＆ペーストして実行してください。

pip install helion

pip は Python のパッケージ管理ツールです。pip install パッケージ名 と打つだけで必要なソフトを自動で取得・インストールしてくれます。

開発版（最新の変更が入ったソースコードから動かすバージョン）を試したい場合は、uv（高速な Python 仮想環境ツール）を使った以下の手順が公式に推奨されています。

git clone https://github.com/pytorch/helion.git
cd helion
uv venv .venv
source .venv/bin/activate
pip install -e '.[dev]'

git clone はリポジトリ（ソースコードの置き場所）をパソコンにコピーするコマンドです。uv venv は他のプロジェクトへの影響を防ぐ独立した Python 環境を作ります。通常は PyPI 版（pip install helion）から始めるのが一番手軽です。

Step 3 — 最初のカーネルを書く

以下は行列積（行列どうしの掛け算）を計算する最小のカーネルの例です。for ループより外は CPU 上で動く通常の PyTorch コード、for ループの中が GPU にコンパイルされる部分になります。

import torch, helion, helion.language as hl

@helion.kernel()
def matmul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    m, k = x.size()
    k, n = y.size()
    out = torch.empty([m, n], dtype=x.dtype, device=x.device)
    for tile_m, tile_n in hl.tile([m, n]):
        acc = hl.zeros([tile_m, tile_n], dtype=torch.float32)
        for tile_k in hl.tile(k):
            acc = torch.addmm(acc, x[tile_m, tile_k], y[tile_k, tile_n])
        out[tile_m, tile_n] = acc
    return out

@helion.kernel() というデコレータ（関数に機能を追加する印）を付けるだけで、普通の Python 関数が GPU カーネルとして自動でコンパイルされます。PyTorch の書き方に慣れていれば、コードの構造はすぐに読み取れるはずです。

デモについて

Helion カーネルの実行には CUDA 対応 GPU（または AMD GPU）と PyTorch・Triton のインストールが必要です。ブラウザだけで動くオンラインデモはありませんが、以下のクラウド環境でハンズオンが可能です。

  • Google Colab ノートブック: 公式の softmax.ipynb が提供されています。Colab 上で GPU ランタイムを選択すれば、ブラウザから直接実行できます。
  • AMD DevCloud: AMD GPU 向けの公式ノートブックも用意されており、NVIDIA 以外のハードウェアでの動作確認や性能比較が行えます。

どちらも無料または低コストで利用できるため、手元に GPU がない場合でも Helion を体験することができます。

動かしてみた

Python 3.12 環境でソースコードの構成を確認しました。helion/ ディレクトリ以下には、カーネルのコンパイル処理・言語ランタイム・ユーティリティが整理された形で配置されています。

./helion/__init__.py
./helion/_compile_time.py
./helion/_dist_utils.py
./helion/exc.py
./helion/_utils.py
./helion/_compat.py

パッケージ構造は明確に分かれており、コンパイル時の処理とランタイム部分がきれいに分離されていることが確認できました。PyPI からのインストールは pip install helion の一行で完結します。

試す前に知っておくとよいこと:

  • オートチューニング（自動設定探索）は数分〜10 分程度かかることがあります。開発中はスキップする設定を使うのが基本的なやり方です。
  • GPU がない環境でも、インタープリタモードを使えばカーネルの動作確認が一部できます。
  • 対応 OS は Linux のみです。Windows や macOS では動作しません。

実践のコツ

すぐに試せる行動を中心にまとめます。

  • まず PyPI 版から始める: pip install helion の安定版が最も手軽です。ソースからのインストールは動作に慣れてから試しましょう。
  • 開発中はオートチューニングをオフに: 環境変数 HELION_AUTOTUNE_EFFORT=none を設定するだけで、デフォルト設定で即座に実行できます。動作確認が済んでからチューニングを走らせるのが効率的です。
  • 生成コードを確認する習慣を: HELION_PRINT_OUTPUT_CODE=1 を設定すると、Helion が生成した Triton コードが画面に表示されます。どんなコードに変換されているかを見ながら学べます。
  • デバッグには CPU モードが便利: HELION_INTERPRET=1 を設定すると、GPU がなくても CPU 上でカーネルをシミュレーション実行できます。print() や breakpoint() が使えるためバグの原因を追いやすくなります。
  • ログ出力を活用する: 動作が不明なときは HELION_LOGS=all で INFO レベルのログを、HELION_LOGS=+all でさらに詳しい DEBUG ログを確認できます。
  • 公式の例コードから始める: リポジトリ内の examples/ フォルダに行列積・softmax などの実装例があります。まずコピーして動かしてみるのが一番の近道です。

活用アイデア

  • カスタム Attention カーネルの実装: Flash Attention のような注目度の高い演算を hl.tile で記述し、Hopper GPU 向けの TMA（テンソルメモリ高速転送機能）を活用したインデックス方式を自動探索させられます。warp specialize（GPU 内のスレッドグループを役割分担させる技法）など Blackwell 固有の最適化も Config 経由で試せます。
  • 混合精度 Softmax の高速化: 入出力を float16（省メモリな数値形式）にしつつ、内部の累積計算は float32（高精度）で行うパターンを Helion が自動でチューニングします。大規模な reduction（合計・最大値などの集約演算）も設定一つで最適化できます。
  • 本番サービスへの組み込み: オートチューニング完了後、最良コンフィグを @helion.kernel(config=helion.Config(...)) としてコードに固定することで、本番起動時のチューニング待ちをなくせます。入力の形状が決まっているなら static_shapes=True でさらに最適化できます。
  • AMD GPU 環境での性能評価: AMD DevCloud 向けの公式ノートブックを使えば、NVIDIA 以外のハードウェアで Helion の動作を確認できます。マルチベンダー対応の調査や性能比較に役立ちます。
  • 新アーキテクチャのポータビリティ評価: 差分進化アルゴリズムによるコンフィグ探索を活かし、新しい GPU アーキテクチャに対してどの設定が最も性能を発揮するかを体系的に調べられます。ハードウェアをまたいだ汎用性（ポータビリティ）の研究に有用です。

用語とポイント解説

Triton

GPU カーネルを書くための中間レベルのプログラミング言語です。CUDA（NVIDIA が提供する低レベル GPU 言語）より書きやすいですが、それでもストライドやマスキングなどの細かい設定が必要です。かんたんに言うと、「CUDA より書きやすいが、それでも専門知識が要る GPU 専用の言語」です。Helion はこの Triton をバックエンド（内部で使う処理エンジン）として自動生成・利用します。

DSL（Domain Specific Language）

特定の用途に特化して作られたプログラミング言語のことです。Helion は GPU カーネル開発という用途のための DSL であり、PyTorch に似た構文を持ちます。かんたんに言うと、「ある分野の仕事だけに特化して作られた専用の言語」です。汎用言語（Python や Java など）と比べて、目的の処理を少ないコードで書けるのが特徴です。

hl.tile

テンソル（多次元の数値配列）の計算空間を小さなブロック（タイル）に分割する Helion の基本操作です。for tile_m in hl.tile(m): と書くだけで、m の範囲を GPU のスレッドに自動で割り当ててくれます。かんたんに言うと、「大きな計算を GPU が処理しやすいサイズに自動で分割してくれる命令」です。外側のループは GPU のグリッド（処理単位のまとまり）にマッピングされます。

Config（コンフィグ）

ブロックサイズ・ループ順序・GPU スレッド数などのチューニング設定をまとめたオブジェクトです。Helion のオートチューナーが自動で最良の Config を探してくれます。かんたんに言うと、「GPU の動かし方の細かい設定をまとめた設定書」です。オートチューニング後に最良の Config をコードに固定すると、本番環境での起動コストをゼロにできます。

オートチューナー

数百通りのコンフィグを実際に実行して最速の設定を自動で見つける機能です。Helion では差分進化アルゴリズムを使ってコンフィグ空間を効率よく探索します。かんたんに言うと、「多くの設定パターンを試して、一番速いものを自動で選んでくれる仕組み」です。開発中は HELION_AUTOTUNE_EFFORT=none でスキップできます。

DifferentialEvolutionSearch（差分進化探索）

Helion 内蔵のオートチューナーが使うアルゴリズムです。遺伝的アルゴリズムの一種で、複数の候補解を組み合わせて徐々により良い設定へ収束させます。かんたんに言うと、「たくさんの設定候補を生物の進化のように組み合わせて、最良の答えに近づいていく探索手法」です。全探索より少ない試行回数で優れた設定を見つけられます。

TileIR

NVIDIA Blackwell GPU 向けの追加コンパイルバックエンドです。環境変数 ENABLE_TILE=1 を設定することで有効になります。かんたんに言うと、「最新の NVIDIA GPU（Blackwell 世代）に特化した専用の最適化コンパイラ」です。Hopper 以前の GPU では通常の Triton バックエンドが使われます。

pid_type（プログラム ID タイプ）

GPU のスレッドブロック（計算単位のグループ）をグリッド全体にどう割り当てるかを決める設定です。flat・xyz・persistent_blocked などの選択肢があります。かんたんに言うと、「GPU の処理単位をどのような順番・配置で実行するかを決めるオプション」です。ワークロードの形状によって最適な pid_type が変わるため、オートチューナーが自動で選択します。

static_shapes（スタティックシェイプ）

入力テンソルの形状（行列のサイズなど）ごとに独立したチューニングを行うかどうかのフラグです。デフォルトは true で、形状が固定されているときに最大の性能が出ます。かんたんに言うと、「入力サイズが決まっているなら専用に最適化し、サイズが変わるなら汎用的な設定を使う、という切り替えスイッチ」です。推論サービスのように入力サイズが固定されている場面では static_shapes=True が特に効果的です。

ぜひ GPU カーネルの高速化や機械学習モデルの推論最適化、AMD/NVIDIA を横断したパフォーマンス評価などに活用してみてはいかがでしょうか。