# Autotuning API

This document describes the function-level and context autotuning interfaces of Triton-distributed, and provides minimal Ascend usage examples. For the `triton.autotune` parameters used by internal kernels in the context mode, please refer to the [Triton-Ascend API documentation](https://triton-ascend.readthedocs.io/zh-cn/latest/python-api/generated/triton.autotune.html).

## `triton_dist.tune.autotune`

### Interface Definition

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

Decorates a regular Python Host function to select the configuration with the shortest execution time from a user-provided configuration space, and injects the configuration fields as keyword arguments into the decorated function.

### Parameter Description

#### `config_space`

- Type: `list[dict]`
- Required: Yes

List of candidate configurations. Each dictionary represents a complete configuration set, and its field names must be valid as keyword arguments for the decorated function.

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

The configuration space cannot be empty. It is recommended that all candidates use a consistent set of fields, and place a stable configuration as the first item for use when `autotune=False`.

#### `key_fn`

- Type: `Callable`
- Required: Yes

Generates a cache key based on the current business parameters. The same key indicates that the same best configuration can be reused; different keys will perform independent queries or execute tuning.

```python
def key_fn(x, alpha, *args, **kwargs):
    return (tuple(x.shape), str(x.dtype))
```

The key should cover shape, stride, dtype, world size, or execution mode that may change the best configuration, and avoid including tensor pointers, rank, or other values that do not change the best configuration or are unstable.

#### `prune_fn`

- Type: `Optional[Callable]`
- Required: No
- Default: `None`

Filters configurations before performance testing:

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

- Returns `True`: Keep this configuration;
- Returns `False`: Prune this configuration.

`prune_fn` is only responsible for filtering, not measuring or selecting the best configuration. When not provided, all candidate configurations participate in tuning. The pruning logic should be lightweight, side-effect-free, and ensure that all ranks produce consistent results and at least one configuration is retained.

Conditions suitable for `prune_fn` include: configuration exceeding shape, low padding utilization, insufficient memory budget, or buffer capacity not meeting deterministic constraints.

### Return Value

Returns a callable `AutoTuner` object. Its final return value is consistent with the original Host function. After the first tuning is completed, the interface calls the original function again using the best configuration and returns the result from that invocation.

### Invocation Control Parameters

```python
result = tuned_func(
    *args,
    autotune=True,
    autotune_verbose=False,
    autotune_allow_arg_overwrite=False,
    autotune_pg=None,
    **kwargs,
)
```

#### `autotune`

- Type: `bool`
- Default: `True`

Whether to enable tuning and cache queries. When set to `False`, directly uses `config_space[0]` without executing `key_fn`, `prune_fn`, and performance testing.

#### `autotune_verbose`

- Type: `bool`
- Default: `False`

Whether to display INFO-level tuning information on standard output. When actual tuning occurs, detailed logs are also written to the disk cache directory.

#### `autotune_allow_arg_overwrite`

- Type: `bool`
- Default: `False`

When business `kwargs` and configuration fields have name conflicts, whether to allow business parameters to overwrite configurations. By default, raises `ValueError`; when set to `True`, business parameters take precedence.

Generally not recommended to enable, as overwritten fields cannot be properly tuned according to the candidate space.

#### `autotune_pg`

- Type: `torch.distributed.ProcessGroup | None`
- Default: `None`

Process group for multi-rank synchronization. When provided, performs `MAX` reduction on the configuration timings from each rank. When not provided, attempts to use the world group initialized by Triton-distributed; if no available process group exists, tunes in single-process mode.

Currently, the final timing reduction uses PyTorch's default WORLD group, so `autotune_pg` should contain the same ranks as WORLD.

### [Minimal Usage Example](https://gitcode.com/Ascend/Triton-distributed-ascend/blob/master/python/triton_dist/test/ascend/test_autotune_ascend.py)

The following example omits the kernel implementation and only demonstrates the configuration injection process for the Host interface:

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

For multi-rank invocation, simply add a process group:

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

### Tuning, Caching, and Exceptions

Each retained configuration executes 5 warmup runs and 10 timing runs. Ascend uses `torch.npu.Event` to calculate average timing. When candidates encounter resource shortages or runtime exceptions, their timing is recorded as infinity, and testing continues with other configurations.

Cache and log location:

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

The disk cache key includes the hash of the decorated function's source code, hardware digest, and business key; the cache content also records software dependencies and the timing of each candidate.

Supported environment variables:

| Environment Variable | Default | Description |
| --- | --- | --- |
| `TRITON_DIST_AUTOTUNE_ALWAYS_TUNE` | `0` | When set to `1`, ignores existing results and retunes |
| `TRITON_DIST_AUTOTUNE_VERSION_CHECK` | `0` | When set to `1`, does not reuse old cache if dependencies are inconsistent |

## `triton_dist.autotuner.contextual_autotune`

### Interface Definition

```python
triton_dist.autotuner.contextual_autotune(
    is_dist=False,
    n_repeat=5,
    n_warmup=3,
)
```

Decorates a Python function containing autotune kernel invocations, allowing internal kernels to run and select configurations in the actual execution context of the outer function.

### Parameter Description

#### `is_dist`

- Type: `bool`
- Default: `False`

Whether to enable distributed selection. When set to `True`, places the timing of each valid configuration on the current NPU and performs `all_reduce(MAX)` through PyTorch's default WORLD group.

Must initialize `torch.distributed` before use, and all WORLD ranks must enter the same context tuning flow.

#### `n_repeat`

- Type: `int`
- Default: `5`

Number of timing runs for each candidate configuration. The interface uses the average of these measurements to compare configurations.

#### `n_warmup`

- Type: `int`
- Default: `3`

Number of warmup runs before formal timing for each candidate configuration.

### Return Value

Returns a callable `ContextualAutoTuner` object. After all internal kernels complete tuning, the outer function executes once more using the best configuration and returns the result from that execution.

### [Minimal Usage Example](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)
```

Distributed mode does not accept a separate ProcessGroup parameter, but is set on the already initialized 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)
```

### Execution Behavior

The context tuner temporarily takes over the execution process of internal Triton Autotuners:

1. Executes the outer function and registers internal kernels that need tuning;
2. Repeatedly executes the outer function, allowing candidate configurations of internal kernels to complete warmup and timing sequentially;
3. Selects the one with the shortest average timing among valid configurations;
4. When `is_dist=True`, first takes the maximum value of timings from each rank;
5. Writes the best configuration to the internal kernel's in-process cache;
6. Executes and returns the outer function result using the best configuration.

Candidate pruning for internal kernels is handled by their own autotune configuration; `contextual_autotune` does not provide a separate `prune_fn` parameter.

### Logs, Caching, and Exceptions

Log location:

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

Logs record kernel name, key, configuration number, measurement round, exceptions, average timing, and best configuration.

The context mode reuses the in-process autotune cache of internal kernels and does not generate function-level JSON disk cache. If all candidates are invalid, raises `RuntimeError("cannot find valid config")`.

Only one active `ContextualAutoTuner` is allowed at a time; nested or concurrent calls are not supported.

## Interface Selection

| Requirement | Interface to Use |
| --- | --- |
| Tune complete Host launcher and need custom key, pruning, and disk caching | `triton_dist.tune.autotune` |
| Tune internal kernels in outer communication or synchronization context | `contextual_autotune` |

For complete workflow, see [Operator Performance Testing and Tuning: autotune Feature Usage Guide](../developer-guide/kernel-performance/operator_performance_testing_and_tuning_autotune_guide.md); for distributed practices, see [autotune Example: Optimization Method Practice](../tutorial/autotune_optimization_example.md).