Helion: Jagged Tensors & Autotuning/Config Management
Overview
Two high-value Helion areas that are easy to get wrong or miss entirely:
- Ragged/jagged tensors — iterates variable-length inner
dimensions with implicit masking, so you never hand-build masks.
- 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
Part 1 — Ragged / Jagged Tensors
Data layout
Jagged data is stored
prefix-packed: a flat
buffer holding all rows
concatenated, plus an
tensor of length
where row
is
x_data[x_offsets[i] : x_offsets[i+1]]
. Per-row length =
.
— the core API
is the jagged counterpart to
.
is an
N-D tensor of
per-lane end positions drawn from an enclosing tile context. It
lowers to a dense
loop but
masks out indices where
tile_k.index >= parent[lane]
— 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
Rules & gotchas for
- 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:
is valid, alone is not.
- Nest for multi-level ragged iteration (e.g. variable rows × variable
features). See .
- Use plain (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())
extra_mask=tile_k.index[None, :] < lengths[:, None]
Reference examples (in the Helion repo )
| File | Pattern |
|---|
| basic per-row reduction |
| nested (rows × features) |
| online/Welford reduction over ragged dim |
| per-row normalization |
| ragged + dense, manual style |
| , | jagged matmul / attention |
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"];
}
Tune once, save, reload
returns the best
.
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
methods:
/
,
,
(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")
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 var | Effect |
|---|
| cache root (default: torch cache dir ) |
| cache class: (default), (also keys on Helion/PyTorch/Triton versions), , |
| ignore cached config, re-tune, write result back |
| skip cache read AND write |
HELION_ASSERT_CACHE_HIT=1
| fail if no cached config (CI guard) |
| | | (default) — tuning depth |
HELION_AUTOTUNE_BUDGET_SECONDS
| wall-clock cap on tuning |
Deploy a handful of known configs ( + )
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.
controls
when it
re-selects (on top of shape specialization).
Seed tuning from previously cached/best configs
returns an
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): ...
AOT heuristics — zero-cost per-shape selection across GPUs (the headline feature)
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.
-
Decorate with
helion.experimental.aot_kernel
instead of
:
python
import helion.experimental
@helion.experimental.aot_kernel() # extras: batched=..., key=fn, collect_fn, measure_fn
def vector_add(x, y): ...
-
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
-
It emits
_helion_aot_<kernel>_<device>_sm<NN>.py
(e.g.
_helion_aot_vector_add_cuda_sm90.py
). Commit it.
-
Runtime lookup order:
→ file matching the current
compute capability →
older compatible capabilities (e.g. on
, tries
→
→
). 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
, commit the new
alongside the existing ones. See
for
worked examples (vector_add, softmax, layer_norm, rms_norm, cross_entropy, rope)
and
docs/deployment_autotuning.md
for the full workflow.
Share configs across machines (remote cache)
Implement
helion.autotuner.remote_cache.RemoteCacheBackend
(
/
, optional
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.
Control which dimensions specialize (affects cache entries)
| API | Where | Effect |
|---|
| (default) | decorator | specialize on exact shape/stride; best perf, one entry per shape |
| decorator | bucket dims into ; one kernel serves many sizes |
| inside kernel | that dim always a compile-time constant, every call |
torch._dynamo.mark_static(t, dims)
| before call | specialize dims on specific tensors only |
| decorator | custom extra grouping for re-selection |
With
, pin
if tensors routinely
exceed
elements to avoid an extra specialization boundary.
Advanced: manual routing / drop Helion at serving time
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
| Mistake | Fix |
|---|
| Hand-building masks for ragged loops | Use — masking is implicit |
| for a jagged child tile | Index with parent: |
| as the outermost loop | It needs an enclosing parent tile |
| Re-tuning on every deploy / new shape | AOT heuristics ( + ) |
| Assuming there's no autotune CLI | There is: python -m helion.experimental.aot_runner
|
| for varied shapes | One config fits all shapes; use / or AOT instead |
| Re-tuning per machine for the same GPU/shape | Remote cache backend warm-starts the team |
| Editing AOT files for a new GPU by hand | Re-run on that GPU; it emits the file |