# 自动调优 API

本文说明 Triton-distributed 的函数级和 context 自动调优接口，并给出最小 Ascend 使用示例。context 方式内部 Kernel 所使用的 `triton.autotune` 参数请参考 [Triton-Ascend API 文档](https://triton-ascend.readthedocs.io/zh-cn/latest/python-api/generated/triton.autotune.html)。

## `triton_dist.tune.autotune`

### 接口定义

```python
triton_dist.tune.autotune(config_space, key_fn, prune_fn=None)
```

装饰普通 Python Host 函数，从用户给定的配置空间中选择耗时最短的配置，并将配置字段作为关键字参数注入被装饰函数。

### 参数说明

#### `config_space`

- 类型：`list[dict]`
- 必选：是

候选配置列表。每个字典表示一组完整配置，其字段名必须能作为关键字参数传给被装饰函数。

```python
config_space = [
    {"cfg": config_0},
    {"cfg": config_1},
]
```

配置空间不能为空。建议各候选使用一致的字段集合，并将稳定配置放在第一项，供 `autotune=False` 时使用。

#### `key_fn`

- 类型：`Callable`
- 必选：是

根据本次业务参数生成缓存 key。key 相同表示可以复用同一最佳配置，key 不同会独立查询或执行调优。

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

在性能测试前过滤配置：

```python
keep = prune_fn(config, *args, **kwargs)
```

- 返回 `True`：保留该配置；
- 返回 `False`：裁掉该配置。

`prune_fn` 只负责过滤，不负责测量或选择最佳配置。未传入时，全部候选配置都参与调优。裁剪逻辑应轻量、无副作用，并保证所有 rank 结果一致且至少保留一个配置。

适合放入 `prune_fn` 的条件包括：配置超过 shape、padding 利用率过低、内存预算不足或缓冲区容量不满足等确定性约束。

### 返回值

返回可调用的 `AutoTuner` 对象。其最终返回值与原始 Host 函数一致。首次调优完成后，接口使用最佳配置再次调用原函数，并返回该次结果。

### 调用控制参数

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

### [最小使用示例](https://gitcode.com/Ascend/Triton-distributed-ascend/blob/master/python/triton_dist/test/ascend/test_autotune_ascend.py)

以下示例省略 Kernel 实现，只展示 Host 接口的配置注入过程：

```python
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 调用只需增加进程组：

```python
y = scale(
    x,
    2.5,
    autotune=True,
    autotune_pg=process_group,
)
```

### 调优、缓存与异常

每个保留配置执行 5 次预热和 10 次计时。Ascend 使用 `torch.npu.Event` 计算平均耗时。候选出现资源不足或运行异常时，其耗时记为无穷大，并继续测试其他配置。

缓存和日志位置：

```text
~/.triton_dist/autotune/<function_name>/
```

磁盘缓存 key 包含被装饰函数源码 hash、硬件摘要和业务 key；缓存内容还记录软件依赖与各候选耗时。

支持的环境变量：

| 环境变量 | 默认值 | 说明 |
| --- | --- | --- |
| `TRITON_DIST_AUTOTUNE_ALWAYS_TUNE` | `0` | 设置为 `1` 时忽略已有结果并重新调优 |
| `TRITON_DIST_AUTOTUNE_VERSION_CHECK` | `0` | 设置为 `1` 时，依赖不一致则不复用旧缓存 |

## `triton_dist.autotuner.contextual_autotune`

### 接口定义

```python
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 完成调优后，外层函数使用最佳配置再执行一次，并返回该次结果。

### [最小使用示例](https://gitcode.com/Ascend/Triton-distributed-ascend/blob/master/python/triton_dist/test/ascend/test_contextual_autotune_ascend.py)

```python
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 上设置：

```python
@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 的运行过程：

1. 执行外层函数并注册需要调优的内部 Kernel；
2. 重复执行外层函数，使内部 Kernel 的候选配置依次完成预热和计时；
3. 对有效配置选择平均耗时最短者；
4. `is_dist=True` 时先对各 rank 耗时取最大值；
5. 将最佳配置写入内部 Kernel 的进程内 cache；
6. 使用最佳配置执行并返回外层函数结果。

内部 Kernel 的候选裁剪由其自身的 autotune 配置负责，`contextual_autotune` 不提供单独的 `prune_fn` 参数。

### 日志、缓存与异常

日志位置：

```text
./.autotune_logs/rank-<rank>.log
```

日志记录 Kernel 名称、key、配置编号、测量轮次、异常、平均耗时和最佳配置。

context 方式复用内部 Kernel 的进程内 autotune cache，不生成函数级 JSON 磁盘缓存。若全部候选均无效，则抛出 `RuntimeError("cannot find valid config")`。

同一时刻只允许一个活动的 `ContextualAutoTuner`，不支持嵌套或并发调用。

## 接口选择

| 需求 | 使用接口 |
| --- | --- |
| 调优完整 Host launcher，并需要自定义 key、裁剪和磁盘缓存 | `triton_dist.tune.autotune` |
| 在外层通信或同步上下文中调优内部 Kernel | `contextual_autotune` |

完整流程见[算子性能测试与调优：autotune 特性使用介绍](../developer-guide/kernel-performance/operator_performance_testing_and_tuning_autotune_guide.md)，分布式实践见[autotune 样例：优化方法实践](../tutorial/autotune_optimization_example.md)。