SuperAPI AI 工具集

helion-jagged-and-autotuning

智能体与模型

查看项目

更新于 10-16 05:05

技能介绍

用途:面向使用 Helion(import helion; import helion.language as hl)编写内核的开发者,解决两类易错问题——嵌套/不规则(ragged/jagged)张量迭代,以及自动调优与配置管理。

主要能力:一、锯齿张量。数据以前缀打包存储:扁平 x_data 加长度 num_rows+1 的 x_offsets,第 i 行为 x_data[x_offsets[i]:x_offsets[i+1]]。核心 API 为 hl.jagged_tile(parent),parent 是来自外层 tile 上下文的每 lane 结束位置张量,它降为稠密 hl.tile(parent.amax()) 循环并自动掩码 tile_k.index >= parent[lane],无需手写掩码。限制:parent 秩须≥1,各轴须来自外层 tile;不能作为内核最外层循环;锯齿子 tile 须用父轴索引(x[tile_b, tile_k] 合法,x[tile_k] 不合法);内层边界统一时用普通 hl.tile()。

二、自动调优与配置。Kernel.autotune(args) 返回 helion.Config,可 save/load、to_json/from_json/from_dict、minimize。默认首次调用自动调优并按硬件、特化(形状/类型/步长)、运行时、后端与源码哈希缓存到磁盘。环境变量包括 HELION_CACHE_DIR、HELION_AUTOTUNE_CACHE、HELION_FORCE_AUTOTUNE、HELION_SKIP_CACHE、HELION_ASSERT_CACHE_HIT、HELION_AUTOTUNE_EFFORT、HELION_AUTOTUNE_BUDGET_SECONDS。少量已知形状可用 configs=[...] 加 key=;可用 helion.from_cache() 预热搜索。AOT 启发式用 helion.experimental.aot_kernel 装饰,并以 python -m helion.experimental.aot_runner -- python my_benchmark.py 离线生成 helion_aot<kernel>_<device>_sm<NN>.py,运行时按计算能力回退查找,实现零成本按形状选配置。跨机共享可实现 RemoteCacheBackend 并设置 HELION_REMOTE_CACHE_BACKEND 与 HELION_AUTOTUNE_CACHE=RemoteAutotuneCache。static_shapes、hl.specialize、torch._dynamo.mark_static、key= 控制特化粒度。

适用场景:多形状、多 GPU 部署且不希望运行时重新调优;团队/CI 共享配置;锯齿数据上的归约、softmax、归一化、矩阵乘与注意力等。

限制:AOT 离线扫描较慢(每个不同形状约 5–15 分钟);未找到启发式文件时回退默认配置并一次性告警;新增 GPU 需在该硬件上重跑 aot_runner。具体收费与平台支持暂未提供。