自动调优 API
本文说明 Triton-distributed 的函数级和 context 自动调优接口，并给出最小 Ascend 使用示例。context 方式内部 Kernel 所使用的 triton.autotune 参数请参考 Triton-Ascend API 文档。
triton_dist.tune.autotune
接口定义
triton_dist.tune.autotune(config_space, key_fn, prune_fn=None)
装饰普通 Python Host 函数，从用户给定的配置空间中选择耗时最短的配置，并将配置字段作为关键字参数注入被装饰函数。
参数说明
config_space
类型：
list[dict]必选：是
候选配置列表。每个字典表示一组完整配置，其字段名必须能作为关键字参数传给被装饰函数。
config_space = [
{"cfg": config_0},
{"cfg": config_1},
]
配置空间不能为空。建议各候选使用一致的字段集合，并将稳定配置放在第一项，供 autotune=False 时使用。
key_fn
类型：
Callable必选：是
根据本次业务参数生成缓存 key。key 相同表示可以复用同一最佳配置，key 不同会独立查询或执行调优。
def key_fn(x, alpha, *args, **kwargs):
return (tuple(x.shape), str(x.dtype))
key 应覆盖可能改变最佳配置的 shape、stride、dtype、world size 或运行模式，并避免包含 Tensor 指针、rank 等不会改变最佳配置或不稳定的值。
prune_fn
类型：
Optional[Callable]必选：否
默认值：
None
在性能测试前过滤配置：
keep = prune_fn(config, *args, **kwargs)
返回
True：保留该配置；返回
False：裁掉该配置。
prune_fn 只负责过滤，不负责测量或选择最佳配置。未传入时，全部候选配置都参与调优。裁剪逻辑应轻量、无副作用，并保证所有 rank 结果一致且至少保留一个配置。
适合放入 prune_fn 的条件包括：配置超过 shape、padding 利用率过低、内存预算不足或缓冲区容量不满足等确定性约束。
返回值
返回可调用的 AutoTuner 对象。其最终返回值与原始 Host 函数一致。首次调优完成后，接口使用最佳配置再次调用原函数，并返回该次结果。
调用控制参数
result = tuned_func(
*args,
autotune=True,
autotune_verbose=False,
autotune_allow_arg_overwrite=False,
autotune_pg=None,
**kwargs,
)
autotune
类型：
bool默认值：
True
是否启用调优和缓存查询。设置为 False 时直接使用 config_space[0]，不执行 key_fn、prune_fn 和性能测试。
autotune_verbose
类型：
bool默认值：
False
是否在标准输出显示 INFO 级调优信息。发生实际调优时，详细日志同时写入磁盘缓存目录。
autotune_allow_arg_overwrite
类型：
bool默认值：
False
业务 kwargs 与配置字段重名时，是否允许业务参数覆盖配置。默认抛出 ValueError；设置为 True 时业务参数优先。
通常不建议开启，否则被覆盖字段无法按候选空间正常调优。
autotune_pg
类型：
torch.distributed.ProcessGroup | None默认值：
None
用于多 rank 同步的进程组。传入后，对各 rank 的配置耗时执行 MAX 归约。未传入时尝试使用 Triton-distributed 初始化的 world group；没有可用进程组时按单进程方式调优。
当前最终耗时归约使用 PyTorch 默认 WORLD group，因此 autotune_pg 应与 WORLD 包含相同 rank。
最小使用示例
以下示例省略 Kernel 实现，只展示 Host 接口的配置注入过程：
import torch
import triton
from triton_dist.tune import autotune
def configs():
return [
{"cfg": triton.Config({"BLOCK": block}, num_warps=warps)}
for block in (256, 512)
for warps in (4, 8)
]
def key_fn(x, alpha, *args, **kwargs):
return (tuple(x.shape), str(x.dtype))
def prune_fn(entry, x, alpha, *args, **kwargs):
block = entry["cfg"].all_kwargs()["BLOCK"]
return block <= x.numel()
@autotune(config_space=configs(), key_fn=key_fn, prune_fn=prune_fn)
def scale(x, alpha, cfg):
y = torch.empty_like(x)
meta = cfg.all_kwargs()
grid = (triton.cdiv(x.numel(), meta["BLOCK"]),)
scale_kernel[grid](x, y, alpha, x.numel(), **meta)
return y
y = scale(x, 2.5, autotune=True, autotune_verbose=True)
多 rank 调用只需增加进程组：
y = scale(
x,
2.5,
autotune=True,
autotune_pg=process_group,
)
调优、缓存与异常
每个保留配置执行 5 次预热和 10 次计时。Ascend 使用 torch.npu.Event 计算平均耗时。候选出现资源不足或运行异常时，其耗时记为无穷大，并继续测试其他配置。
缓存和日志位置：
~/.triton_dist/autotune/<function_name>/
磁盘缓存 key 包含被装饰函数源码 hash、硬件摘要和业务 key；缓存内容还记录软件依赖与各候选耗时。
支持的环境变量：
环境变量 |
默认值 |
说明 |
|---|---|---|
|
|
设置为 |
|
|
设置为 |
triton_dist.autotuner.contextual_autotune
接口定义
triton_dist.autotuner.contextual_autotune(
is_dist=False,
n_repeat=5,
n_warmup=3,
)
装饰一个包含 autotune Kernel 调用的 Python 函数，使内部 Kernel 在外层函数的真实执行上下文中逐配置运行和选优。
参数说明
is_dist
类型：
bool默认值：
False
是否启用分布式选优。设置为 True 时，将每个有效配置的耗时放到当前 NPU，并通过 PyTorch 默认 WORLD group 执行 all_reduce(MAX)。
使用前必须初始化 torch.distributed，且所有 WORLD rank 必须进入相同的 context 调优流程。
n_repeat
类型：
int默认值：
5
每个候选配置的计时次数。接口使用这些测量结果的平均值比较配置。
n_warmup
类型：
int默认值：
3
每个候选配置在正式计时前的预热次数。
返回值
返回可调用的 ContextualAutoTuner 对象。所有内部 Kernel 完成调优后，外层函数使用最佳配置再执行一次，并返回该次结果。
最小使用示例
import torch
import triton
import triton.language as tl
from triton_dist.autotuner import contextual_autotune
BLOCK = 512
@triton.autotune(
configs=[
triton.Config({"BLOCK": BLOCK}, num_warps=warps)
for warps in (4, 8)
],
key=["n"],
)
@triton.jit
def scale_kernel(x_ptr, y_ptr, alpha, n, BLOCK: tl.constexpr):
pid = tl.program_id(0)
offsets = pid * BLOCK + tl.arange(0, BLOCK)
mask = offsets < n
x = tl.load(x_ptr + offsets, mask=mask, other=0.0)
tl.store(y_ptr + offsets, alpha * x, mask=mask)
@contextual_autotune(is_dist=False, n_repeat=3, n_warmup=2)
def run_scale(x, alpha):
y = torch.empty_like(x)
grid = (triton.cdiv(x.numel(), BLOCK),)
scale_kernel[grid](x, y, alpha, x.numel())
return y
y = run_scale(x, 3.0)
分布式模式不接收单独的 ProcessGroup 参数，而是在已初始化的 WORLD group 上设置：
@contextual_autotune(is_dist=True, n_repeat=3, n_warmup=2)
def run_scale_dist(x, alpha):
return launch_scale_in_context(x, alpha)
执行行为
context 调优器会临时接管内部 Triton Autotuner 的运行过程：
执行外层函数并注册需要调优的内部 Kernel；
重复执行外层函数，使内部 Kernel 的候选配置依次完成预热和计时；
对有效配置选择平均耗时最短者；
is_dist=True时先对各 rank 耗时取最大值；将最佳配置写入内部 Kernel 的进程内 cache；
使用最佳配置执行并返回外层函数结果。
内部 Kernel 的候选裁剪由其自身的 autotune 配置负责，contextual_autotune 不提供单独的 prune_fn 参数。
日志、缓存与异常
日志位置：
./.autotune_logs/rank-<rank>.log
日志记录 Kernel 名称、key、配置编号、测量轮次、异常、平均耗时和最佳配置。
context 方式复用内部 Kernel 的进程内 autotune cache，不生成函数级 JSON 磁盘缓存。若全部候选均无效，则抛出 RuntimeError("cannot find valid config")。
同一时刻只允许一个活动的 ContextualAutoTuner，不支持嵌套或并发调用。
接口选择
需求 |
使用接口 |
|---|---|
调优完整 Host launcher，并需要自定义 key、裁剪和磁盘缓存 |
|
在外层通信或同步上下文中调优内部 Kernel |
|
完整流程见算子性能测试与调优：autotune 特性使用介绍，分布式实践见autotune 样例：优化方法实践。