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

## 文档定位

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

仓库内代码位置：

- 原始算子：[`06-ascend-reverse-all2all-udma.py`](https://gitcode.com/Ascend/Triton-distributed-ascend/blob/master/tutorials/ascend/04-ascend-reverse-all2all/04-ascend-reverse-all2all.py)
- autotune 样例：[`06-ascend-reverse-all2all-udma-autotune.py`](https://gitcode.com/Ascend/Triton-distributed-ascend/blob/master/tutorials/ascend/04-ascend-reverse-all2all/06-ascend-reverse-all2all-udma-autotune-bare.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. 定义配置空间

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

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

```python
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` 在实际计时前过滤确定无效的配置，减少首次搜索开销：

```python
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 函数，并使用函数级装饰器：

```python
@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. 按搜索空间准备资源

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

```python
# 最小 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 函数：

```python
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 的配置：

```python
# 稳态计时只使用最佳配置，不再遍历搜索空间。
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 环境配置后，从仓库根目录执行：

```bash
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` 可以忽略已有结果并重新调优。

## 执行流程小结

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

## 注意事项

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

接口参数见[分布式通用 Host 接口：autotune](../api/autotune_api.md)，整体流程见[算子性能测试与调优：autotune 特性使用介绍](../developer-guide/kernel-performance/operator_performance_testing_and_tuning_autotune_guide.md)。