helion-jagged-and-autotuning

Helion：不规则张量与自动调优／配置管理

概述

Helion 中有两个很有价值、但容易出错或被完全忽略的领域：

1. 不规则／锯齿张量——hl.jagged_tile() 可迭代长度可变的内层维度，并采用隐式掩码，因此无需手动构建掩码。
2. 自动调优与配置管理——自动调优很慢；Helion 提供了一套分层系统（磁盘缓存 → 已保存的配置 → AOT 启发式），可调优一次并复用结果，结果以 GPU 架构和输入形状为键。

最容易被忽略的功能是：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 是一个来自外层 tile 上下文的 N 维每个 lane 的结束位置张量。它会降级为一个密集的 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（不能是标量）；每个轴都必须来自外层 tile 上下文。常见用法是一维“每行长度”。
• 它不能作为内核最外层循环——必须有父 tile。
• 锯齿子 tile 必须使用其父轴进行索引：x[tile_b, tile_k] 有效，但仅使用 x[tile_k] 无效。
• 对多级不规则迭代（例如可变行数 × 可变特征数），可嵌套使用 jagged_tile。参见 jagged_mean.py。
• 当所有 lane 的内层边界一致时，使用普通的 hl.tile()（而非锯齿版本）。
• 手动实现方式（有需要时）：使用 for tile_k in hl.tile(lengths.amax())，并将 extra_mask=tile_k.index[None, :] < lengths[:, None] 传给 hl.load。

参考示例（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")

磁盘上的自动调优缓存（自动按架构 + 形状作为键）

默认情况下，首次调用会自动调优，并根据硬件、特化信息（形状／数据类型／步幅）、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	调优的实际耗时上限

部署少量已知配置（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.experimental.aot_kernel 而不是 helion.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）。一个启发式文件可以适用于多代 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 中的完整工作流。

在多台机器之间共享配置（远程缓存）

实现 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（默认）	装饰器	根据精确形状／步幅进行特化；性能最佳，每种形状对应一个条目
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，以避免增加额外的特化边界。

高级：手动路由／在服务时移除 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）
以为没有自动调优命令行工具	有：python -m helion.experimental.aot_runner
对形状各异的场景使用 @kernel(config=...)	一个配置无法适用于所有形状；请使用 configs=/key= 或 AOT
为同一 GPU/形状在每台机器上重新调优	远程缓存后端可为团队预热调优结果
手动编辑 AOT 文件以适配新 GPU	在该 GPU 上重新运行 aot_runner；它会生成 sm<NN> 文件