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.

`triton_dist.tune.autotune`

Interface Definition

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.

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.

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:

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

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 

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

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:

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:

~/.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

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 

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:

@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:

Executes the outer function and registers internal kernels that need tuning;
Repeatedly executes the outer function, allowing candidate configurations of internal kernels to complete warmup and timing sequentially;
Selects the one with the shortest average timing among valid configurations;
When is_dist=True, first takes the maximum value of timings from each rank;
Writes the best configuration to the internal kernel’s in-process cache;
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:

./.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; for distributed practices, see autotune Example: Optimization Method Practice.

Autotuning API

triton_dist.tune.autotune

Interface Definition

Parameter Description

config_space

key_fn

prune_fn

Return Value

Invocation Control Parameters

autotune

autotune_verbose

autotune_allow_arg_overwrite

autotune_pg

Minimal Usage Example

Tuning, Caching, and Exceptions

triton_dist.autotuner.contextual_autotune

Interface Definition

Parameter Description

is_dist

n_repeat

n_warmup

Return Value

Minimal Usage Example

Execution Behavior

Logs, Caching, and Exceptions

Interface Selection

`triton_dist.tune.autotune`

`config_space`

`key_fn`

`prune_fn`

`autotune`

`autotune_verbose`

`autotune_allow_arg_overwrite`

`autotune_pg`

Minimal Usage Example 

`triton_dist.autotuner.contextual_autotune`

`is_dist`

`n_repeat`

`n_warmup`

Minimal Usage Example 