autotune 样例：UDMA Reverse All2All 增量优化实践

文档定位

本文以 Ascend UDMA Reverse All2All 为例，介绍如何在不修改 Triton Kernel 主体的情况下，为已有 Host launcher 增量接入 triton_dist.tune.autotune。

仓库内代码位置：

原始算子：06-ascend-reverse-all2all-udma.py
autotune 样例：06-ascend-reverse-all2all-udma-autotune.py

本文仅说明从原始算子到 autotune 样例新增或调整的代码：配置空间、缓存 key、配置裁剪、Host 包装器、分布式选优以及最佳配置导出。Kernel 的 UDMA 通信逻辑保持不变。

修改概览

增量修改	作用
引入 `triton_dist.tune`	使用函数级 autotune 接口
增加 `TUNE_CONFIG_SPACE`	定义 Host launcher 的候选参数
增加 `_autotune_key`	按 shape 和 rank 数复用最佳配置
增加 `_prune`	在计时前过滤无效候选
增加 autotune Host 包装器	重复运行完整算子并注入候选参数
调用时传入 `autotune_pg`	按多 rank 最大耗时选择配置
读取 `best_configs`	使用并导出当前 shape 的最佳配置

1. 定义配置空间

在原始文件的常量区域增加以下配置：

import triton_dist.tune


TUNE_CONFIG_SPACE = [
    {
        "COMM_BLOCK_S": bs,  # 序列维通信分块。
        "COMM_BLOCK_D": bd,  # 特征维数据分块。
        "buffer_num": bn,    # 流水缓冲区数量。
    }
    for bs in [32, 64, 128, 256, 512]
    for bd in [64, 128]
    for bn in [2, 3, 4, 8]
]

# 搜索空间语义发生变化时更新版本，隔离旧缓存。
TUNE_SPACE_VERSION = "reverse-a2a-udma-autotune-v1"

该搜索空间共有 5 × 2 × 4 = 40 组配置。字典字段名与 Host launcher 的同名参数对应，调优器会在测试每个候选时自动注入这些值。

2. 增加缓存 key

def _autotune_key(A, C, peer_mem, signal_mem, rank, rank_size):
    return (
        TUNE_SPACE_VERSION,  # 配置空间版本。
        tuple(A.shape),      # 输入 shape。
        tuple(C.shape),      # 输出 shape。
        rank_size,           # 通信规模。
    )

key 相同表示可以复用同一最佳配置。样例不加入 rank 和 Tensor 地址，因为不同 rank 需要选择相同的全局配置，而内存地址也不应影响缓存命中。

样例固定使用 torch.bfloat16。如果算子扩展到多种 dtype 或运行模式，应将相应信息加入 key。

3. 增加配置裁剪

_prune 在实际计时前过滤确定无效的配置，减少首次搜索开销：

def _prune(config_record, A, *args, **kwargs):
    import math

    # A 后的位置参数依次为 C、peer_mem、signal_mem、rank、rank_size。
    rank_size = args[4]
    S = A.shape[0] // rank_size
    D = A.shape[2]
    bs = config_record["COMM_BLOCK_S"]
    bd = config_record["COMM_BLOCK_D"]
    bn = config_record["buffer_num"]
    num_blocks_s = math.ceil(S / bs)

    # 过滤缓冲区冗余、分块越界、利用率过低或数据块过大的配置。
    if bn > num_blocks_s:
        return False
    if bs == 128 and bd == 128:
        return True
    if bs > S or bd > D:
        return False
    if S / (num_blocks_s * bs) < 0.75:
        return False
    if D / (math.ceil(D / bd) * bd) < 0.75:
        return False
    if bs * bd * A.element_size() > 128 * 1024:
        return False
    return True

返回 True 表示保留配置并进入性能测试，返回 False 表示裁掉。_prune 只负责过滤，最佳配置仍由实际耗时决定。

4. 包装原有 Host launcher

原始 Kernel 不需要增加 triton.autotune。在 Kernel 外新增普通 Python Host 函数，并使用函数级装饰器：

@triton_dist.tune.autotune(
    config_space=TUNE_CONFIG_SPACE,
    key_fn=_autotune_key,
    prune_fn=_prune,
)
def _function_autotuned_hccl_reverse_a2a_udma(
    A,
    C,
    peer_mem,
    signal_mem,
    rank,
    rank_size,
    # 这三个参数由候选配置注入。
    buffer_num=2,
    COMM_BLOCK_S=128,
    COMM_BLOCK_D=128,
):
    S_total, H, D = A.shape
    S = S_total // rank_size
    vec_num = NPUUtils().get_aivector_core_num()

    # autotune 会重复调用本函数，每个候选开始前恢复相同状态。
    signal_mem.fill_(0)
    dist.barrier()

    kernel_hccl_reverse_a2a_pipelined[vec_num, 1, 1](
        A,
        C,
        peer_mem,
        signal_mem,
        rank,
        rank_size,
        buffer_num,
        S,
        H,
        D,
        A.stride(0),
        A.stride(1),
        A.stride(2),
        C.stride(0),
        C.stride(1),
        C.stride(2),
        COMM_BLOCK_S=COMM_BLOCK_S,
        COMM_BLOCK_D=COMM_BLOCK_D,
    )

与原始 launcher 相比，主要变化只有两点：调优参数改为可注入的函数参数；信号内存复位与 rank 同步移入包装器，保证每次候选测试的初始条件一致。

5. 按搜索空间准备资源

资源容量需要覆盖所有候选，而不是只满足默认配置：

# 最小 S 分块产生最多序列块，最大 buffer_num 需要最多缓冲区。
min_bs = min(c["COMM_BLOCK_S"] for c in TUNE_CONFIG_SPACE)
max_bn = max(c["buffer_num"] for c in TUNE_CONFIG_SPACE)
max_num_blocks_d = max(
    triton.cdiv(D, c["COMM_BLOCK_D"]) for c in TUNE_CONFIG_SPACE
)

max_num_blocks_s = triton.cdiv(S, min_bs)
signal_mem_size = _signal_mem_size(
    S,
    H,
    D,
    rank_size,
    max_bn,
    max_num_blocks_s,
    max_num_blocks_d,
)

这样可以避免较大候选因辅助内存不足而在调优过程中失败。

6. 发起分布式调优

在原有 benchmark 流程中，将直接调用 Kernel 的位置替换为 autotune Host 函数：

process_group = dist.new_group(ranks=list(range(rank_size)))

_function_autotuned_hccl_reverse_a2a_udma(
    A_local,
    C_local,
    peer_mem,
    signal_mem,
    rank,
    rank_size,
    autotune=True,              # 启用搜索或缓存查询。
    autotune_pg=process_group,  # 启用多 rank 同步和选优。
)

# 使用相同 key 读取本次选择的配置。
key = _autotune_key(
    A_local, C_local, peer_mem, signal_mem, rank, rank_size
)
best = _function_autotuned_hccl_reverse_a2a_udma.best_configs.get(key)

每组配置在各 rank 上以相同顺序执行。框架对耗时执行 all_reduce(MAX)，使用最慢 rank 的耗时作为该配置的全局结果，再选择全局耗时最小的配置。

当前最终耗时归约使用默认 WORLD group，因此样例创建的 process_group 与 WORLD 包含相同 rank。

7. 使用并导出最佳配置

调优结束后，样例把 best 传给稳态性能 launcher，并由 rank 0 保存每个 shape 的配置：

# 稳态计时只使用最佳配置，不再遍历搜索空间。
launch = _prepare_reverse_launch(
    A_local,
    C_local,
    peer_mem,
    signal_mem,
    rank,
    rank_size,
    best["buffer_num"],
    S,
    H,
    D,
    best["COMM_BLOCK_S"],
    best["COMM_BLOCK_D"],
)

if rank == 0:
    manifest.append({
        "S": S,
        "H": H,
        "D": D,
        "COMM_BLOCK_S": best["COMM_BLOCK_S"],
        "COMM_BLOCK_D": best["COMM_BLOCK_D"],
        "buffer_num": best["buffer_num"],
    })

autotune 只负责性能选优。样例仍使用 PyTorch HCCL 结果执行正确性校验，并将稳态性能写入 CSV。

8. 运行样例

完成 Ascend、CANN、HCCL 和 SHMEM 环境配置后，从仓库根目录执行：

source /usr/local/Ascend/ascend-toolkit/set_env.sh

export REV_PROFILE_WARMUP=5
export REV_PROFILE_ITERS=50
export REVERSE_A2A_MANIFEST=reverse_a2a_udma_autotune_manifest.json
export REVERSE_A2A_CSV=reverse_a2a_udma_autotune_perf.csv

torchrun --nproc-per-node=<rank_num> \
  06-ascend-reverse-all2all-udma-autotune.py

rank_num 应与参与运行的 NPU 数量一致。REV_PROFILE_WARMUP 和 REV_PROFILE_ITERS 控制最佳配置选出后的稳态性能测试，不改变函数级调优器内部的 5 次预热和 10 次计时。

首次遇到一个新 key 时会执行配置搜索；再次运行相同源码、硬件和 key 时，会从 ~/.triton_dist/autotune/ 查询磁盘缓存。设置 TRITON_DIST_AUTOTUNE_ALWAYS_TUNE=1 可以忽略已有结果并重新调优。

执行流程小结

根据当前 shape 生成 key 并查询缓存；
缓存未命中时，使用 _prune 过滤 40 组候选；
所有 rank 重复执行保留的候选，并按最慢 rank 耗时选优；
使用最佳配置执行稳态计时和正确性校验；
rank 0 输出 manifest 和 CSV 文件。

注意事项

被调优 Host 函数会重复执行，输入、通信缓冲区和同步状态必须可复位；
所有 rank 的 key、配置空间和裁剪结果必须一致；
搜索空间第一项应保证可执行，供 autotune=False 时直接使用；
首次调优耗时与最佳配置的稳态性能应分开统计；
修改配置语义时应更新 TUNE_SPACE_VERSION。

接口参数见分布式通用 Host 接口：autotune，整体流程见算子性能测试与调优：autotune 特性使用介绍。