Operator Performance Testing and Tuning: autotune Feature Usage Guide

Document Positioning

This document introduces two auto-tuning methods provided by Triton-distributed:

• Function-level auto-tuning: triton_dist.tune.autotune;

• Context auto-tuning: triton_dist.autotuner.contextual_autotune.

The content is for those who already have callable Triton Kernels and want to further tune distributed operators, focusing on usage methods, tuning boundaries, distributed selection, and caching behavior on Ascend.

For Kernel-level triton.autotune used internally by the context method, refer to Triton-Ascend autotune Usage Guide, API Documentation, and Usage Examples.

Feature Overview

Tuning Method	Tuning Object	Config Source	Cache	Typical Scenario
triton_dist.tune.autotune	Regular Python Host launcher	config_space	In-process and disk cache	Tune complete function containing communication and one or more Kernels
contextual_autotune	Python function containing autotune Kernel	Internal Kernel configs	Internal Kernel in-process cache	Kernel must complete tuning in real outer context

Both methods are adapted for Ascend device timing and support taking maximum time across multiple ranks before selecting global configuration.

Quick Start

Function-Level Auto-Tuning

Below is a minimal example organized from repository Ascend scale test. Function-level decorator is added to Host launcher:

import torch
import triton
import triton.language as tl
from triton_dist.tune import autotune


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


def config_space():
    return [
        {"cfg": triton.Config({"BLOCK": block}, num_warps=warps)}
        for block in (256, 512, 1024, 2048)
        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.element_size() < 32 * 1024


@autotune(
    config_space=config_space(),
    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

When calling, no need to pass cfg, tuner will inject from config space:

x = torch.randn(1 << 18, device="npu", dtype=torch.float32)
y = scale(x, 2.5, autotune=True, autotune_verbose=True)
torch.testing.assert_close(y, 2.5 * x, atol=1e-5, rtol=1e-4)

On first encounter of a new key, framework filters and tests candidate configs; subsequent calls reuse best config. This example doesn’t put alpha in key because it doesn’t change input layout and tuning config selection.

Context Auto-Tuning

Context method decorates outer Python function, which internally calls Kernel with Kernel-level autotune:

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

x = torch.randn(1 << 18, device="npu", dtype=torch.float32)
y = run_scale(x, 3.0)
torch.testing.assert_close(y, 3.0 * x, atol=1e-5, rtol=1e-4)
assert scale_kernel.best_config is not None

Context tuner will drive internal Kernel to run per-config during complete run_scale call, and write best config to internal Kernel cache.

Working Principle

Function-Level Method

On first call with a new business key:

1. key_fn generates key from business parameters;

2. Query in-process cache and disk cache;

3. When cache misses, use prune_fn to filter candidate configs;

4. For each config, execute 5 warmup and 10 timing runs;

5. In multi-rank mode, perform MAX reduction on time from each rank;

6. Save best config, and execute launcher again with best config.

Function-level method measures complete device-side work triggered by decorated Host function. Config can affect both Kernel parameters and launch method or communication pipeline in launcher.

Context Method

Context tuner takes over internal Kernel’s autotune flow during outer function execution:

1. Execute outer function and discover internal Kernels needing tuning;

2. In same outer context, warmup and time each candidate config of internal Kernel;

3. After one internal Kernel completes selection, continue processing subsequent Kernels;

4. After all Kernels complete, execute outer function once more using their respective best configs.

n_warmup and n_repeat control warmup count and timing count for each candidate config respectively. When is_dist=True, takes maximum of candidate time from each rank for selection.

How to Choose

Use function-level method when:

• Want to directly tune regular Python launcher;

• Config affects communication, Kernel launch, or multiple execution steps simultaneously;

• Need to reuse disk cache across processes;

• Need to use key_fn and prune_fn for custom caching and pruning logic.

Use context method when:

• Function internally already has one or more autotune Kernels;

• Kernel cannot be independently tested outside outer communication, synchronization, or buffer state;

• Want to complete internal Kernel selection during real function call.

Ascend and Distributed Support

Function-level method uses torch.npu.Event for timing; context method creates Event, gets stream, and cleans benchmark cache through current Triton Ascend driver’s device interface.

Distributed selection uniformly adopts slowest rank time:

global_time(config) = max(time_rank_0, ..., time_rank_n)

Function-level method passes in HCCL ProcessGroup through autotune_pg:

result = scale(
    x,
    alpha,
    autotune=True,
    autotune_pg=process_group,
)

Context method enables distributed reduction through is_dist=True, and uses PyTorch default WORLD group. All participating ranks must execute same candidate order and call flow.

Caching and Logging

Tuning Method	Cache	Logging
Function-level	In-process cache and ~/.triton_dist/autotune/ disk cache	Located in same function directory as disk cache
Context	Internal Kernel’s in-process autotune cache	./.autotune_logs/rank-<rank>.log

Function-level disk cache records function source code, hardware info, business key, and software dependencies. Modifying function source or hardware generates different cache; when dependencies mismatch, defaults to warning, can set TRITON_DIST_AUTOTUNE_VERSION_CHECK=1 to force re-tune.

Set the following environment variable to ignore existing function-level results:

export TRITON_DIST_AUTOTUNE_ALWAYS_TUNE=1

Performance Testing Process

Recommend separating first tuning from steady-state performance testing:

1. Execute tuning with real shape, dtype, and rank count;

2. Use reference implementation to verify correctness of best config;

3. Retain disk cache or in-process cache;

4. Run again with same key, only measure steady-state time of best config;

5. For distributed operators, focus on both per-rank time and slowest rank time.

Usage Considerations

• Both methods repeatedly execute tuned function, input, output, and synchronization state must support repeated calls;

• autotune only compares time, doesn’t automatically verify correctness;

• Function-level prune_fn only filters configs, not responsible for performance comparison;

• Multi-rank key, config count, config order, and pruning results must be consistent;

• Context tuner only allows one active instance at a time, doesn’t support nesting or concurrent use;

• First call includes search and compilation overhead, cannot be directly used as steady-state performance result.

Summary

Function-level method is suitable for using Host launcher as complete tuning boundary; context method is suitable for tuning internal Kernel in real outer call. Both select based on actual timing results, and use slowest rank’s time as selection basis in distributed scenarios.

For complete interface, see Distributed General Host Interface: autotune, for distributed operator practice, see autotune Example: Optimization Method Practice.