Back to Details

helion-jagged-and-autotuning

Compare original and translation side by side

🇺🇸

Original

English
🇨🇳

Translation

Chinese

Helion: Jagged Tensors & Autotuning/Config Management

Helion：锯齿张量与自动调优/配置管理

Overview

概述

Two high-value Helion areas that are easy to get wrong or miss entirely:
  1. Ragged/jagged tensors — 
    hl.jagged_tile()
    iterates variable-length inner dimensions with implicit masking, so you never hand-build masks.
  2. Autotuning & config management — autotuning is slow; Helion has a layered system (on-disk cache → saved configs → AOT heuristics) for tuning once and reusing results keyed by GPU architecture and input shape.
The single most-missed feature: AOT heuristics (
helion.experimental.aot_kernel
  • python -m helion.experimental.aot_runner
    ) give zero-cost per-shape config selection at runtime, with automatic compute-capability fallback. If a request mentions "many GPUs," "many shapes," or "don't want to re-tune at deploy," reach for AOT, not just the cache.
Helion中两个极易出错或被完全忽略的高价值领域：
  1. 不规则/锯齿张量 —— 
    hl.jagged_tile()
    通过隐式掩码迭代可变长度的内部维度，你完全无需手动构建掩码。
  2. 自动调优与配置管理 —— 自动调优耗时较长；Helion提供了分层系统（磁盘缓存 → 已保存配置 → AOT启发式算法），实现一次调优，按GPU架构和输入形状为键复用结果。
最容易被忽略的核心特性：AOT启发式算法（
helion.experimental.aot_kernel
+ 
python -m helion.experimental.aot_runner
）可在运行时实现零成本的按形状配置选择，并支持自动计算能力降级。如果需求中提到“多GPU”、“多形状”或“不想在部署时重新调优”，请优先使用AOT，而不仅仅是缓存。

Part 1 — Ragged / Jagged Tensors

第一部分 — 不规则/锯齿张量

Data layout

数据布局

Jagged data is stored prefix-packed: a flat 
x_data
buffer holding all rows concatenated, plus an 
x_offsets
tensor of length 
num_rows + 1
where row 
i
is 
x_data[x_offsets[i] : x_offsets[i+1]]
. Per-row length = 
ends - starts
.
锯齿数据采用前缀打包存储：一个扁平的
x_data
缓冲区存储所有拼接后的行，外加一个长度为
num_rows + 1
的
x_offsets
张量，第
i
行对应
x_data[x_offsets[i] : x_offsets[i+1]]
。每行长度 = 
ends - starts
。

hl.jagged_tile(parent)
— the core API

hl.jagged_tile(parent)
—— 核心API

hl.jagged_tile(parent)
is the jagged counterpart to 
hl.tile()
. 
parent
is an N-D tensor of per-lane end positions drawn from an enclosing tile context. It lowers to a dense 
hl.tile(parent.amax())
loop but masks out indices where 
tile_k.index >= parent[lane]
automatically — you write the ragged loop directly.
python
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
hl.jagged_tile(parent)
是
hl.tile()
的锯齿版本。
parent
是一个N维张量，存储来自外层tile上下文的每个通道的结束位置。它最终会被编译为密集的
hl.tile(parent.amax())
循环，但会自动掩码掉
tile_k.index >= parent[lane]
的索引——你可以直接编写不规则循环。
python
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

Rules & gotchas for 
jagged_tile

jagged_tile
的规则与常见陷阱

  • parent
    must be rank ≥ 1 (no scalars); every axis must come from an enclosing tile context. The 1-D "per-row length" case is the common one.
  • It cannot be the outermost loop of a kernel — it needs a parent tile.
  • A jagged child tile must be indexed with its parent axes: 
    x[tile_b, tile_k]
    is valid, 
    x[tile_k]
    alone is not.
  • Nest 
    jagged_tile
    for multi-level ragged iteration (e.g. variable rows × variable features). See 
    jagged_mean.py
    .
  • Use plain 
    hl.tile()
    (not jagged) when the inner bound is uniform across lanes.
  • The manual equivalent (when you need it): 
    for tile_k in hl.tile(lengths.amax())
    with 
    extra_mask=tile_k.index[None, :] < lengths[:, None]
    passed to 
    hl.load
    .
  • parent
    必须秩≥1（不能是标量）；每个轴都必须来自外层tile上下文。一维“逐行长度”是最常见的场景。
  • 它不能作为内核的最外层循环——需要一个父tile。
  • 子锯齿tile必须使用父轴进行索引：
    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]
    。

Reference examples (in the Helion repo 
examples/
)

参考示例（位于Helion仓库的
examples/
目录下）

FilePattern
jagged_sum.py
basic per-row reduction
jagged_mean.py
nested 
jagged_tile
(rows × features)
jagged_softmax.py
online/Welford reduction over ragged dim
jagged_layer_norm.py
per-row normalization
jagged_dense_add.py
ragged + dense, manual 
extra_mask
style
jagged_dense_bmm.py
, 
jagged_hstu_attn.py
jagged matmul / attention
文件模式
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
锯齿矩阵乘法 / 注意力机制

Part 2 — Autotuning & Config Management

第二部分 — 自动调优与配置管理

Pick the right mechanism

选择合适的机制

dot
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"];
}
dot
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"];
}

Tune once, save, reload

一次调优，保存与重载

Kernel.autotune(args)
returns the best 
helion.Config
. 
Config
is JSON-serializable.
python
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): ...
Other 
Config
methods: 
.to_json()
/ 
.from_json(str)
, 
.from_dict(d)
, 
.minimize(config_spec)
(drop default-valued keys before saving).
Tune across several representative shapes and save per-shape:
python
for tag, args in datasets.items():
    my_kernel.autotune(args).save(f"configs/my_kernel_{tag}.json")
Kernel.autotune(args)
返回最优的
helion.Config
。
Config
支持JSON序列化。
python
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)
（保存前剔除默认值的键）。
可针对多个代表性形状进行调优，并按形状保存：
python
for tag, args in datasets.items():
    my_kernel.autotune(args).save(f"configs/my_kernel_{tag}.json")

The on-disk autotune cache (automatic, keyed by arch + shape)

磁盘自动调优缓存（自动生效，按架构+形状为键）

By default the first call autotunes and caches the result keyed on hardware, specialization (shape/dtype/stride), CUDA/ROCm runtime, backend, and kernel source hash — so an H100 and a B200, or two different shapes, get separate entries and each is auto-selected on later runs. No code needed.
Env varEffect
HELION_CACHE_DIR
cache root (default: torch cache dir 
/helion
)
HELION_AUTOTUNE_CACHE
cache class: 
LocalAutotuneCache
(default), 
StrictLocalAutotuneCache
(also keys on Helion/PyTorch/Triton versions), 
RemoteAutotuneCache
, 
AOTAutotuneCache
HELION_FORCE_AUTOTUNE=1
ignore cached config, re-tune, write result back
HELION_SKIP_CACHE=1
skip cache read AND write
HELION_ASSERT_CACHE_HIT=1
fail if no cached config (CI guard)
HELION_AUTOTUNE_EFFORT
none
| 
quick
| 
full
(default) — tuning depth
HELION_AUTOTUNE_BUDGET_SECONDS
wall-clock cap on tuning
默认情况下，首次调用会执行自动调优，并以硬件、特化维度（形状/数据类型/步长）、CUDA/ROCm运行时、后端、内核源码哈希为键缓存结果——因此H100和B200，或是两种不同的形状，会有独立的缓存条目，后续运行时会自动选择对应配置。无需编写额外代码。
环境变量作用
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
调优的墙上时钟时间上限

Deploy a handful of known configs (
configs=
+ 
key=
)

部署少量已知配置（
configs=
+ 
key=
）

python
@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 runs a lightweight benchmark of the listed configs the first time each specialization key is seen and picks the fastest. 
key=
controls when it re-selects (on top of shape specialization).
python
@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=
参数控制何时重新选择配置（在形状特化的基础上）。

Seed tuning from previously cached/best configs

从先前缓存/最优配置中种子化调优

helion.from_cache()
returns an 
autotuner_fn
that seeds the search from prior best configs instead of starting cold:
python
@helion.kernel(autotuner_fn=helion.from_cache(max_configs=5))
def my_kernel(x, y): ...
helion.from_cache()
返回一个
autotuner_fn
，它会从先前的最优配置中初始化搜索，而非从零开始：
python
@helion.kernel(autotuner_fn=helion.from_cache(max_configs=5))
def my_kernel(x, y): ...

AOT heuristics — zero-cost per-shape selection across GPUs (the headline feature)

AOT启发式算法 —— 跨GPU零成本按形状选择配置（核心特性）

For "ship a kernel that serves many shapes on many GPUs without runtime tuning": offline, sweep the kernel over representative shapes, tune each, and distill a decision-tree heuristic that picks a config in microseconds at runtime.
  1. Decorate with 
    helion.experimental.aot_kernel
    instead of 
    helion.kernel
    :
    python
    import helion.experimental

@helion.experimental.aot_kernel()   # extras: batched=..., key=fn, collect_fn, measure_fn
def vector_add(x, y): ...
  2. On the target GPU, run the AOT workflow against a benchmark that exercises the shape sweep (this is slow — full autotune per distinct shape, ~5–15 min each):
    bash
    python -m helion.experimental.aot_runner -- python my_benchmark.py
  3. It emits 
    _helion_aot_<kernel>_<device>_sm<NN>.py
    next to the kernel source (e.g. 
    _helion_aot_vector_add_cuda_sm90.py
    ). Commit it.
  4. Runtime lookup order: 
    $HELION_HEURISTIC_DIR
    → file matching the current compute capability → older compatible capabilities (e.g. on 
    sm120
    , tries 
    sm120
    → 
    sm100
    → 
    sm90
    ). One heuristic can serve multiple GPU generations. No file found → default config + one-time warning.
To add a new GPU: get on that hardware, re-run 
aot_runner
, commit the new 
_helion_aot_*_sm<NN>.py
alongside the existing ones. See 
pretuned_kernels/
for worked examples (vector_add, softmax, layer_norm, rms_norm, cross_entropy, rope) and 
docs/deployment_autotuning.md
for the full workflow.
适用于“发布一个内核，在多GPU上支持多种形状，无需运行时调优”的场景：离线阶段，针对代表性形状范围遍历内核，逐个调优，然后提炼出一个决策树启发式算法，可在运行时以微秒级速度选择配置。
  1. 使用
    helion.experimental.aot_kernel
    替代
    helion.kernel
    装饰器：
    python
    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分钟）：
    bash
    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
    ）。一个启发式算法可支持多代GPU。若未找到文件 → 使用默认配置并给出一次性警告。
要添加新GPU支持：在对应硬件上重新运行
aot_runner
，将新生成的
_helion_aot_*_sm<NN>.py
文件与现有文件一起提交即可。参见
pretuned_kernels/
目录下的完整示例（vector_add、softmax、layer_norm、rms_norm、cross_entropy、rope），以及
docs/deployment_autotuning.md
中的完整工作流说明。

Share configs across machines (remote cache)

跨机器共享配置（远程缓存）

Implement 
helion.autotuner.remote_cache.RemoteCacheBackend
(
get
/
put
, optional 
list
for warm-start), then:
bash
export HELION_REMOTE_CACHE_BACKEND=mypkg.cache.MyBackend   # required to load it
export HELION_AUTOTUNE_CACHE=RemoteAutotuneCache           # also write winners to remote
Read-through/write-through; local disk stays source of truth, remote outage degrades gracefully. Use for team/CI/multi-node warm-starts.
实现
helion.autotuner.remote_cache.RemoteCacheBackend
（
get
/
put
方法，可选
list
用于热启动），然后：
bash
export HELION_REMOTE_CACHE_BACKEND=mypkg.cache.MyBackend   # required to load it
export HELION_AUTOTUNE_CACHE=RemoteAutotuneCache           # also write winners to remote
采用读穿透/写穿透模式；本地磁盘仍是权威数据源，远程服务中断时会优雅降级。适用于团队/CI/多节点热启动场景。

Control which dimensions specialize (affects cache entries)

控制哪些维度参与特化（影响缓存条目）

APIWhereEffect
static_shapes=True
(default)		decoratorspecialize on exact shape/stride; best perf, one entry per shape
static_shapes=False
decoratorbucket dims into 
{0, 1, ≥2}
; one kernel serves many sizes
hl.specialize(dim)
inside kernelthat dim always a compile-time constant, every call
torch._dynamo.mark_static(t, dims)
before callspecialize dims on specific tensors only
key=lambda ...
decoratorcustom extra grouping for re-selection
With 
static_shapes=False
, pin 
index_dtype=torch.int64
if tensors routinely exceed 
2**31
elements to avoid an extra specialization boundary.
API位置作用
static_shapes=True
（默认）		装饰器基于精确形状/步长特化；性能最佳，每个形状对应一个条目
static_shapes=False
装饰器将维度分桶为
{0, 1, ≥2}
；一个内核支持多种尺寸
hl.specialize(dim)
内核内部该维度在每次调用时始终为编译时常量
torch._dynamo.mark_static(t, dims)
调用前仅对特定张量的维度进行特化
key=lambda ...
装饰器自定义额外分组，用于重新选择配置
当
static_shapes=False
时，如果张量元素数通常超过
2**31
，请固定
index_dtype=torch.int64
，以避免额外的特化边界。

Advanced: manual routing / drop Helion at serving time

高级用法：手动路由 / 在服务时移除Helion依赖

python
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
python
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

Common Mistakes

常见错误

MistakeFix
Hand-building masks for ragged loopsUse 
hl.jagged_tile(lengths)
— masking is implicit
x[tile_k]
for a jagged child tile		Index with parent: 
x[tile_b, tile_k]
hl.jagged_tile
as the outermost loop		It needs an enclosing parent tile
Re-tuning on every deploy / new shapeAOT heuristics (
aot_kernel
+ 
aot_runner
)
Assuming there's no autotune CLIThere is: 
python -m helion.experimental.aot_runner
@kernel(config=...)
for varied shapes		One config fits all shapes; use 
configs=
/
key=
or AOT instead
Re-tuning per machine for the same GPU/shapeRemote cache backend warm-starts the team
Editing AOT files for a new GPU by handRe-run 
aot_runner
on that GPU; it emits the 
sm<NN>
file
错误修复方法
为不规则循环手动构建掩码使用
hl.jagged_tile(lengths)
——掩码是隐式的
对子锯齿tile使用
x[tile_k]
索引		加上父tile索引：
x[tile_b, tile_k]
将
hl.jagged_tile
作为最外层循环		它需要一个外层的父tile
每次部署/遇到新形状都重新调优使用AOT启发式算法（
aot_kernel
+ 
aot_runner
）
认为没有自动调优CLI有：
python -m helion.experimental.aot_runner
对多变的形状使用
@kernel(config=...)
单个配置适用于所有形状；请改用
configs=
/
key=
或AOT
相同GPU/形状下每台机器都重新调优远程缓存后端为团队提供热启动
手动编辑新GPU的AOT文件在该GPU上重新运行
aot_runner
；它会生成
sm<NN>
文件