Autotuning Demo: UDMA Reverse All2All Optimization

Document Purpose

This document uses Ascend UDMA Reverse All2All as an example to introduce how to incrementally integrate triton_dist.tune.autotune into an existing Host launcher without modifying the Triton Kernel body.

Code locations in the repository:

Original operator: 06-ascend-reverse-all2all-udma.py
autotune example: 06-ascend-reverse-all2all-udma-autotune.py

This document only explains the code added or adjusted from the original operator to the autotune example: configuration space, cache key, configuration pruning, Host wrapper, distributed selection, and best configuration export. The Kernel’s UDMA communication logic remains unchanged.

Modification Overview

Incremental Modification	Purpose
Introduce `triton_dist.tune`	Use function-level autotune interface
Add `TUNE_CONFIG_SPACE`	Define candidate parameters for Host launcher
Add `_autotune_key`	Reuse best configuration by shape and rank count
Add `_prune`	Filter invalid candidates before timing
Add autotune Host wrapper	Repeatedly run complete operator and inject candidate parameters
Pass `autotune_pg` when calling	Select configuration based on multi-rank maximum time
Read `best_configs`	Use and export best configuration for current shape

1. Define Configuration Space

Add the following configuration in the constants section of the original file:

import triton_dist.tune


TUNE_CONFIG_SPACE = [
    {
        "COMM_BLOCK_S": bs,  # Sequence dimension communication block.
        "COMM_BLOCK_D": bd,  # Feature dimension data block.
        "buffer_num": bn,    # Number of pipeline buffers.
    }
    for bs in [32, 64, 128, 256, 512]
    for bd in [64, 128]
    for bn in [2, 3, 4, 8]
]

# Update version when search space semantics change to isolate old cache.
TUNE_SPACE_VERSION = "reverse-a2a-udma-autotune-v1"

This search space has 5 × 2 × 4 = 40 configurations. Dictionary field names correspond to same-named parameters of the Host launcher, and the tuner will automatically inject these values when testing each candidate.

2. Add Cache Key

def _autotune_key(A, C, peer_mem, signal_mem, rank, rank_size):
    return (
        TUNE_SPACE_VERSION,  # Configuration space version.
        tuple(A.shape),      # Input shape.
        tuple(C.shape),      # Output shape.
        rank_size,           # Communication scale.
    )

Same key means the same best configuration can be reused. The example doesn’t include rank and Tensor addresses, because different ranks need to select the same global configuration, and memory addresses shouldn’t affect cache hits.

The example uses fixed torch.bfloat16. If the operator extends to multiple dtypes or runtime modes, corresponding information should be added to the key.

3. Add Configuration Pruning

_prune filters definitely invalid configurations before actual timing, reducing first search overhead:

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

    # Position arguments after A are C, peer_mem, signal_mem, rank, rank_size in sequence.
    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)

    # Filter configurations with redundant buffers, block overflow, too low utilization, or oversized data blocks.
    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

Return True means keep the configuration and enter performance testing, return False means prune it. _prune only handles filtering; best configuration is still determined by actual timing.

4. Wrap Existing Host Launcher

The original Kernel doesn’t need to add triton.autotune. Add a new regular Python Host function outside the Kernel and use function-level decorator:

@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,
    # These three parameters are injected by candidate configurations.
    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 will repeatedly call this function, restoring same state before each candidate.
    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,
    )

Compared to the original launcher, there are mainly two changes: tuning parameters become injectable function parameters; signal memory reset and rank synchronization are moved into the wrapper to ensure consistent initial conditions for each candidate test.

5. Prepare Resources According to Search Space

Resource capacity needs to cover all candidates, not just satisfy the default configuration:

# Minimum S block produces most sequence blocks, maximum buffer_num needs most buffers.
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,
)

This avoids larger candidates failing during tuning due to insufficient auxiliary memory.

6. Initiate Distributed Tuning

In the existing benchmark flow, replace the location directly calling the Kernel with the autotune Host function:

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,              # Enable search or cache query.
    autotune_pg=process_group,  # Enable multi-rank synchronization and selection.
)

# Read configuration selected this time using same 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)

Each configuration executes in the same order on all ranks. The framework performs all_reduce(MAX) on timing, using the slowest rank’s timing as the global result for that configuration, then selects the configuration with minimum global timing.

Currently, final timing reduction uses the default WORLD group, so the process_group created in the example contains the same ranks as WORLD.

7. Use and Export Best Configuration

After tuning completes, the example passes best to the steady-state performance launcher, and rank 0 saves each shape’s configuration:

# Steady-state timing only uses best configuration, no longer traverses search space.
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 only handles performance selection. The example still uses PyTorch HCCL results to perform correctness validation and writes steady-state performance to CSV.

8. Run Example

After completing Ascend, CANN, HCCL, and SHMEM environment configuration, execute from repository root directory:

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 should match the number of participating NPUs. REV_PROFILE_WARMUP and REV_PROFILE_ITERS control steady-state performance testing after best configuration is selected, and don’t change the function-level tuner’s internal 5 warmups and 10 timing runs.

When encountering a new key for the first time, configuration search will execute; when running the same source code, hardware, and key again, disk cache will be queried from ~/.triton_dist/autotune/. Setting TRITON_DIST_AUTOTUNE_ALWAYS_TUNE=1 can ignore existing results and re-tune.

Execution Flow Summary

Generate key based on current shape and query cache;
When cache misses, use _prune to filter 40 candidates;
All ranks repeatedly execute retained candidates, selecting based on slowest rank timing;
Execute steady-state timing and correctness validation using best configuration;
Rank 0 outputs manifest and CSV files.

Notes

The tuned Host function will execute repeatedly; inputs, communication buffers, and synchronization state must be resettable;
All ranks’ keys, configuration spaces, and pruning results must be consistent;
First item in search space should guarantee executability for direct use when autotune=False;
First tuning time and best configuration steady-state performance should be counted separately;
Update TUNE_SPACE_VERSION when modifying configuration semantics.

See Distributed General Host Interface: autotune for interface parameters, and Operator Performance Testing and Tuning: autotune Feature Usage Guide for overall flow.