算子性能测试与调优：autotune 特性使用介绍

文档定位

本文介绍 Triton-distributed 提供的两种自动调优方式：

• 函数级自动调优：triton_dist.tune.autotune；

• 上下文自动调优：triton_dist.autotuner.contextual_autotune。

内容面向已有可调用 Triton Kernel、进一步调优分布式算子，重点说明两种方式在 Ascend 上的使用方法、调优边界、分布式选优和缓存行为。

context 方式内部使用的 Kernel 级 triton.autotune 可参考 Triton-Ascend autotune 使用指南、API 文档 和 使用样例。

功能概述

调优方式	调优对象	配置来源	缓存	典型场景
triton_dist.tune.autotune	普通 Python Host launcher	config_space	进程内和磁盘缓存	调优包含通信、一个或多个 Kernel 的完整函数
contextual_autotune	内含 autotune Kernel 的 Python 函数	内部 Kernel 的配置	内部 Kernel 进程内缓存	Kernel 必须在真实外层上下文中完成调优

两种方式均已适配 Ascend 设备计时，并支持多 rank 耗时取最大值后选择全局配置。

快速上手

函数级自动调优

下面是基于仓库 Ascend scale 测试整理的最小示例。函数级装饰器添加在 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

调用时不需要传入 cfg，调优器会从配置空间中注入：

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)

首次遇到新的 key 时，框架过滤并测试候选配置；后续调用复用最佳配置。本例未将 alpha 放入 key，因为它不会改变输入布局和调优配置选择。

上下文自动调优

context 方式装饰外层 Python 函数，外层函数内部调用带有 Kernel 级 autotune 的 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

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 调优器会在完整的 run_scale 调用过程中驱动内部 Kernel 逐配置运行，并将最佳配置写入内部 Kernel cache。

工作原理

函数级方式

首次调用一个新业务 key 时：

1. key_fn 根据业务参数生成 key；

2. 查询进程内缓存和磁盘缓存；

3. 缓存未命中时，使用 prune_fn 过滤候选配置；

4. 对每个配置执行 5 次预热和 10 次计时；

5. 多 rank 模式下对各 rank 耗时执行 MAX 归约；

6. 保存最佳配置，并用最佳配置再次执行 launcher。

函数级方式测量的是被装饰 Host 函数触发的完整设备侧工作。配置既可以影响 Kernel 参数，也可以影响 launcher 中的发射方式或通信流水。

context 方式

context 调优器在外层函数执行期间接管内部 Kernel 的 autotune 流程：

1. 执行外层函数并发现需要调优的内部 Kernel；

2. 在相同外层上下文中，对内部 Kernel 的候选配置分别预热和计时；

3. 一个内部 Kernel 完成选优后继续处理后续 Kernel；

4. 所有 Kernel 完成后，再使用各自最佳配置执行一次外层函数。

n_warmup 和 n_repeat 分别控制每个候选配置的预热次数和计时次数。is_dist=True 时，对各 rank 的候选耗时取最大值后选优。

如何选择

使用函数级方式的情况：

• 希望直接调优普通 Python launcher；

• 配置同时影响通信、Kernel 发射或多个执行步骤；

• 需要跨进程复用磁盘缓存；

• 需要使用 key_fn 和 prune_fn 自定义缓存与裁剪逻辑。

使用 context 方式的情况：

• 函数内部已有一个或多个 autotune Kernel；

• Kernel 不能脱离外层通信、同步或缓冲区状态独立测试；

• 希望在真实函数调用过程中完成内部 Kernel 选优。

Ascend 与分布式支持

函数级方式使用 torch.npu.Event 计时；context 方式通过当前 Triton Ascend driver 的设备接口创建 Event、获取 stream 并清理 benchmark cache。

分布式选优均采用最慢 rank 耗时：

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

函数级方式通过 autotune_pg 传入 HCCL ProcessGroup：

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

context 方式通过 is_dist=True 启用分布式归约，并使用 PyTorch 默认 WORLD group。所有参与 rank 必须执行相同的候选顺序和调用流程。

缓存与日志

调优方式	缓存	日志
函数级	进程内缓存和~/.triton_dist/autotune/ 磁盘缓存	与磁盘缓存位于同一函数目录
context	内部 Kernel 的进程内 autotune 缓存	./.autotune_logs/rank-<rank>.log

函数级磁盘缓存会记录函数源码、硬件信息、业务 key 和软件依赖。修改函数源码或硬件后会生成不同缓存；依赖不一致时默认告警，可设置 TRITON_DIST_AUTOTUNE_VERSION_CHECK=1 强制重新调优。

设置以下环境变量可忽略已有函数级结果：

export TRITON_DIST_AUTOTUNE_ALWAYS_TUNE=1

性能测试流程

建议将首次调优与稳态性能测试分开：

1. 使用真实 shape、dtype 和 rank 数执行调优；

2. 使用参考实现校验最佳配置的正确性；

3. 保留磁盘缓存或进程内 cache；

4. 再次运行相同 key，只测量最佳配置的稳态耗时；

5. 分布式算子同时关注各 rank 耗时和最慢 rank 耗时。

使用注意事项

• 两种方式都会重复执行被调优函数，输入、输出和同步状态必须支持重复调用；

• autotune 只比较耗时，不自动校验正确性；

• 函数级 prune_fn 只负责过滤配置，不负责性能比较；

• 多 rank 的 key、配置数量、配置顺序和裁剪结果必须一致；

• context 调优器同一时刻只允许一个活动实例，不支持嵌套或并发使用；

• 首次调用包含搜索和编译开销，不能直接作为稳态性能结果。

小结

函数级方式适合以 Host launcher 为完整调优边界；context 方式适合在真实外层调用中调优内部 Kernel。两者都以实际计时结果选优，并在分布式场景使用最慢 rank 的耗时作为选择依据。

完整接口见分布式通用 Host 接口：autotune，分布式算子实践见autotune 样例：优化方法实践。