helion-jagged-and-autotuning

Helion：ジャグ配列テンソルと自動チューニング／設定管理

概要

間違えやすかったり、見落としやすかったりする、Helion の重要な領域が 2 つあります。

1. ラグド／ジャグ配列テンソル — hl.jagged_tile() は、暗黙的なマスキングを使って可変長の内側次元を反復処理するため、マスクを自分で組み立てる必要はありません。
2. 自動チューニングと設定管理 — 自動チューニングには時間がかかります。Helion には、チューニングを一度だけ行い、GPU アーキテクチャと入力形状をキーに結果を再利用する、ディスク上のキャッシュ → 保存済み設定 → AOT ヒューリスティックという階層型システムがあります。

最も見落とされやすい機能は、AOT ヒューリスティックです（helion.experimental.aot_kernel と python -m helion.experimental.aot_runner）。実行時に形状ごとの設定をゼロコストで選択でき、計算能力に応じたフォールバックも自動で行われます。「多数の GPU」「多数の形状」「デプロイ時に再チューニングしたくない」といった要望がある場合は、単なるキャッシュではなく AOT を使ってください。

---

パート 1 — ラグド／ジャグ配列テンソル

データレイアウト

ジャグ配列データはプレフィックスパック形式で格納されます。すべての行を連結したフラットな x_data バッファーと、長さが num_rows + 1 の x_offsets テンソルを使います。行 i は x_data[x_offsets[i] : x_offsets[i+1]] で表されます。行ごとの長さは ends - starts です。

hl.jagged_tile(parent) — 中核 API

hl.jagged_tile(parent) は、hl.tile() に対応するジャグ配列版です。parent は、外側のタイルコンテキストから取得する、レーンごとの終了位置を持つ N 次元テンソルです。これは密な hl.tile(parent.amax()) ループに変換されますが、tile_k.index >= parent[lane] となるインデックスは自動的にマスクされます。そのため、ラグドなループをそのまま記述できます。

import torch
import helion
import helion.language as hl


@helion.kernel()
def jagged_sum(x_data: torch.Tensor, x_offsets: torch.Tensor) -> torch.Tensor:
    b = x_offsets.size(0) - 1
    out = torch.zeros([b], dtype=x_data.dtype, device=x_data.device)

    for tile_b in hl.tile(b):
        starts = x_offsets[tile_b]
        ends = x_offsets[tile_b.index + 1]   # note: tile_b.index + 1 for the upper offset
        lengths = ends - starts

        acc = hl.zeros([tile_b], dtype=x_data.dtype)
        for tile_k in hl.jagged_tile(lengths):       # implicit masking, no manual mask
            idx = starts[:, None] + tile_k.index[None, :]
            acc = acc + x_data[idx].sum(dim=1)
        out[tile_b] = acc
    return out

jagged_tile のルールと注意点

• parent はランク 1 以上である必要があります（スカラーは不可）。すべての軸は、外側のタイルコンテキストから取得する必要があります。一般的なのは 1 次元の「行ごとの長さ」のケースです。
• カーネルの最外側のループにはできません。親タイルが必要です。
• ジャグ配列の子タイルは、親の軸を使ってインデックス指定する必要があります。x[tile_b, tile_k] は有効ですが、x[tile_k] だけでは無効です。
• 多段階のラグドな反復処理（たとえば、行数 × 特徴量数が可変の場合）では、jagged_tile をネストします。jagged_mean.py を参照してください。
• レーン間で内側の上限が一様な場合は、通常の hl.tile() を使います（ジャグ配列版ではありません）。
• 手動で同等の処理を書く場合は、for tile_k in hl.tile(lengths.amax()) を使い、hl.load に extra_mask=tile_k.index[None, :] < lengths[:, None] を渡します。

参考例（Helion リポジトリの examples/ 内）

ファイル	パターン
jagged_sum.py	行ごとの基本的なリダクション
jagged_mean.py	ネストした jagged_tile（行 × 特徴量）
jagged_softmax.py	ラグド次元に対するオンライン／Welford リダクション
jagged_layer_norm.py	行ごとの正規化
jagged_dense_add.py	ラグド配列と密な配列の加算。手動の extra_mask 形式
jagged_dense_bmm.py, jagged_hstu_attn.py	ジャグ配列の行列乗算／アテンション

---

パート 2 — 自動チューニングと設定管理

適切な仕組みを選ぶ

digraph config_mgmt {
    "Need a tuned config?" [shape=diamond];
    "One shape, one GPU?" [shape=diamond];
    "Fixed known shapes (a handful)?" [shape=diamond];
    "Many shapes / many GPUs / deploy without re-tuning?" [shape=diamond];
    "Share across machines/CI?" [shape=diamond];

    "Save config(); @kernel(config=...)" [shape=box];
    "@kernel(configs=[...]) + key=" [shape=box];
    "AOT heuristic (aot_kernel + aot_runner)" [shape=box];
    "Remote cache backend" [shape=box];
    "Default: on-disk autotune cache" [shape=box];

    "Need a tuned config?" -> "One shape, one GPU?";
    "One shape, one GPU?" -> "Save config(); @kernel(config=...)" [label="yes"];
    "One shape, one GPU?" -> "Fixed known shapes (a handful)?" [label="no"];
    "Fixed known shapes (a handful)?" -> "@kernel(configs=[...]) + key=" [label="yes"];
    "Fixed known shapes (a handful)?" -> "Many shapes / many GPUs / deploy without re-tuning?" [label="no"];
    "Many shapes / many GPUs / deploy without re-tuning?" -> "AOT heuristic (aot_kernel + aot_runner)" [label="yes"];
    "Many shapes / many GPUs / deploy without re-tuning?" -> "Share across machines/CI?" [label="no"];
    "Share across machines/CI?" -> "Remote cache backend" [label="yes"];
    "Share across machines/CI?" -> "Default: on-disk autotune cache" [label="no"];
}

一度チューニングして保存し、再読み込みする

Kernel.autotune(args) は、最適な helion.Config を返します。Config は JSON にシリアライズできます。

config = my_kernel.autotune(example_inputs)   # returns helion.Config
config.save("configs/my_kernel.json")          # atomic write

best = helion.Config.load("configs/my_kernel.json")
@helion.kernel(config=best)        # one config, applied to ALL shapes/dtypes/devices
def my_kernel(x, y): ...

その他の Config メソッド：.to_json()／.from_json(str)、.from_dict(d)、 .minimize(config_spec)（保存前にデフォルト値のキーを削除します）。

代表的な形状をいくつか使ってチューニングし、形状ごとに保存します。

for tag, args in datasets.items():
    my_kernel.autotune(args).save(f"configs/my_kernel_{tag}.json")

オンディスク自動チューニングキャッシュ（アーキテクチャ＋形状をキーとして自動管理）

デフォルトでは、初回の呼び出し時に自動チューニングが行われ、その結果はハードウェア、特殊化（形状／dtype／stride）、CUDA／ROCm ランタイム、バックエンド、カーネルソースのハッシュをキーとしてキャッシュされます。そのため、H100 と B200、または形状の異なる 2 つのケースでは別々のエントリーが作成され、以降の実行ではそれぞれに最適なものが自動選択されます。コードは不要です。

Env var	効果
HELION_CACHE_DIR	キャッシュのルート（デフォルト：torch のキャッシュディレクトリ /helion）
HELION_AUTOTUNE_CACHE	キャッシュクラス：LocalAutotuneCache（デフォルト）、StrictLocalAutotuneCache（Helion／PyTorch／Triton のバージョンもキーに含めます）、RemoteAutotuneCache、AOTAutotuneCache
HELION_FORCE_AUTOTUNE=1	キャッシュ済みの設定を無視して再チューニングし、結果を書き戻します
HELION_SKIP_CACHE=1	キャッシュの読み取りと書き込みをスキップします
HELION_ASSERT_CACHE_HIT=1	キャッシュ済みの設定がない場合に失敗します（CI ガード）
HELION_AUTOTUNE_EFFORT	none | quick | full（デフォルト）—チューニングの深さ
HELION_AUTOTUNE_BUDGET_SECONDS	チューニングに対する実時間の上限

既知の設定をいくつかデプロイする（configs= ＋ key=）

@helion.kernel(
    configs=[helion.Config.load("small.json"), helion.Config.load("large.json")],
    key=lambda x, y: helion.next_power_of_2(x.numel()),   # re-benchmark bucket
    static_shapes=False,
)
def my_kernel(x, y): ...

Helion は、各特殊化キーが初めて現れたときに、指定された設定の軽量ベンチマークを実行し、最速のものを選択します。key= は（形状の特殊化に加えて）どのタイミングで再選択するかを制御します。

キャッシュ済み／最良の設定からチューニングを開始する

helion.from_cache() は、最初から検索を始めるのではなく、過去の最良設定を起点とする autotuner_fn を返します。

@helion.kernel(autotuner_fn=helion.from_cache(max_configs=5))
def my_kernel(x, y): ...

AOT ヒューリスティック — GPU 間で形状ごとの選択をゼロコストで実行（目玉機能）

「実行時のチューニングなしで、多数の形状と多数の GPU に対応するカーネルをデプロイする」場合は、オフラインで代表的な形状に対してカーネルを実行し、それぞれをチューニングしたうえで、実行時にマイクロ秒単位で設定を選択する決定木ヒューリスティックを生成します。

1. helion.kernel の代わりに helion.experimental.aot_kernel でデコレートします。

   import helion.experimental
   
   @helion.experimental.aot_kernel()   # extras: batched=..., key=fn, collect_fn, measure_fn
   def vector_add(x, y): ...
   

2. ターゲット GPU 上で、形状スイープを実行するベンチマークを使って AOT ワークフローを実行します（低速です。異なる形状ごとにフル自動チューニングを行うため、各形状につき約 5～15 分かかります）。

   python -m helion.experimental.aot_runner -- python my_benchmark.py
   

3. _helion_aot_<kernel>_<device>_sm<NN>.py がカーネルのソースファイルと同じディレクトリに生成されます（例：_helion_aot_vector_add_cuda_sm90.py）。これをコミットします。

4. 実行時の検索順序：$HELION_HEURISTIC_DIR → 現在のコンピュートケーパビリティに一致するファイル → 互換性のある古いケーパビリティ（例：sm120 では sm120 → sm100 → sm90 の順に試します）。1 つのヒューリスティックで複数世代の GPU に対応できます。ファイルが見つからない場合は、デフォルト設定を使用し、警告を 1 回だけ表示します。

新しい GPU を追加するには、そのハードウェアで aot_runner を再実行し、新しい _helion_aot_*_sm<NN>.py を既存のファイルとともにコミットします。実行例については pretuned_kernels/（vector_add、softmax、layer_norm、rms_norm、cross_entropy、rope）を、ワークフロー全体については docs/deployment_autotuning.md を参照してください。

マシン間で設定を共有する（リモートキャッシュ）

helion.autotuner.remote_cache.RemoteCacheBackend（get／put、ウォームスタート用の list は任意）を実装し、次のようにします。

export HELION_REMOTE_CACHE_BACKEND=mypkg.cache.MyBackend   # required to load it
export HELION_AUTOTUNE_CACHE=RemoteAutotuneCache           # also write winners to remote

読み取り時のキャッシュフォールスルーと書き込み時のキャッシュスルーに対応します。ローカルディスクを信頼できる情報源として維持し、リモートの障害が発生しても適切に縮退して動作します。チーム／CI／マルチノード環境でのウォームスタートに使用します。

特殊化する次元を制御する（キャッシュエントリーに影響）

API	場所	効果
static_shapes=True（デフォルト）	デコレーター	正確な形状／stride で特殊化します。最高のパフォーマンスが得られますが、形状ごとに 1 つのエントリーが作成されます
static_shapes=False	デコレーター	次元を {0, 1, ≥2} にまとめます。1 つのカーネルでさまざまなサイズに対応します
hl.specialize(dim)	カーネル内	その次元を毎回、コンパイル時定数として扱います
torch._dynamo.mark_static(t, dims)	呼び出し前	特定のテンソルの次元のみを特殊化します
key=lambda ...	デコレーター	再選択時に使う追加のグループ化をカスタマイズします

static_shapes=False の場合、テンソルの要素数が頻繁に 2**31 を超えるなら、余分な特殊化の境界が生じないよう index_dtype=torch.int64 を指定してください。

上級者向け：手動ルーティング／サービング時に Helion を外す

bound = my_kernel.bind(example_inputs)            # BoundKernel, tied to input types
run_small = bound.compile_config(small_cfg)        # callable for a specific config
src = bound.to_triton_code(small_cfg)              # export raw Triton source

---

よくある間違い

間違い	修正方法
不規則なループ用のマスクを手作業で作成する	hl.jagged_tile(lengths) を使用します。マスキングは暗黙的に行われます
不規則な子タイルに x[tile_k] を使う	親タイルでインデックスを指定します：x[tile_b, tile_k]
hl.jagged_tile を最外側のループにする	囲む親タイルが必要です
デプロイや新しい形状のたびに再チューニングする	AOT ヒューリスティクス（aot_kernel + aot_runner）を使用します
自動チューニング用 CLI がないと思い込む	あります：python -m helion.experimental.aot_runner
形状が異なる場合に @kernel(config=...) を使う	1 つの設定ではすべての形状に対応できません。代わりに configs=/key= または AOT を使用します
同じ GPU／形状についてマシンごとに再チューニングする	リモートキャッシュバックエンドがチームのウォームスタートを実現します
新しい GPU 用に AOT ファイルを手作業で編集する	その GPU で aot_runner を再実行すると、sm<NN> ファイルが生成されます