[HIP] [Perf] Specialize partial tiles in two-stage custom all-reduce

[AndreasKaratzas]

Proposed modifications

1. Specialize the plain two-stage read kernel with a compile-time has_tail parameter.
2. Run the shared-memory loop only over complete tiles so every thread in a block executes the same number of barriers.
3. Reduce leftover packs directly with packed_reduce, preserving their original owner threads and output indices.
4. Preserve accumulation order and inter-device synchronization without adding a kernel launch.
5. Dispatch aligned messages to has_tail=false, removing tail predicates and restoring the generated aligned kernel from before [HIP] fix(custom_all_reduce): uniform trip count for the block barriers #5694.
6. Keep [HIP] fix(custom_all_reduce): uniform trip count for the block barriers #5694's write-mode and fused-kernel fixes intact, along with the existing naive handling of unequal rank partitions and TP6.
7. Add a Gloo-based regression with independent rank inputs, exact reference sums, and changing inputs in eager and graph execution.
8. Cover graph-owned intermediate registration, poisoned replay outputs, the reported shape and its neighbors, single-lane tails, and unequal partitions.
9. Include the flattened (1023,512), (1024,512), and (1025,512) cases in the existing default benchmark sweep, keeping the list sorted, deduplicated, and within its configured size cutoff.

Microbenchmarks and performance comparison

This follows merged #5694 and targets current main; it preserves that PR's correctness fixes while specializing the plain read path. The benchmark study used four gfx950 GPUs at TP4 and a common v0.1.22.post1 base (b4d9154d125e09efbe098d986e40fea3549c1244). The correct comparison baseline applied #5694's three barrier hunks (5e77166e0c627caa109046f6469de8220efa5ba6) to that base. These are the historical isolated-kernel measurements, not a newly timed end-to-end run of current main.

The study completed 44 variant runs: three alternating-order rounds for focused and broad sweeps, and five for eight additional BF16 shapes. Each case used five warmups and seven samples: 64 calls per eager sample, or 64 captured calls replayed ten times per graph sample. CUDA-event timing excludes input preparation and checking; eager results include Python launch/allocation gaps. Latencies are the median across rounds of each round's maximum per-rank sample median. Percentage changes are medians of paired within-round ratios and need not equal ratios of the displayed latencies.

For (1025, 512), lower latency is better:

Dtype	Mode	#5694 uniform loop (µs)	Tail specialization (µs)	Paired change
BF16	eager	16.461	16.330	-0.90%
BF16	graph	13.822	13.722	-0.72%
FP16	eager	16.440	16.254	-0.94%
FP16	graph	13.831	13.678	-0.95%
FP32	eager	26.883	26.753	-0.02%
FP32	graph	24.001	23.898	-0.48%

The full sweep includes BF16/FP16/FP32 at (1023,512), (1024,512), and (1025,512); 43 dtype/size configurations spanning 8 KiB–8 MiB and per-rank tails of 1/63/64/65/127 packs; and BF16 (M,2880) for M=120,121,127,157,160,161,256,512. Compared with the unsafe pre-#5694 release, the proposed specialization's largest observed case-level median increases were 0.56% graph and 1.92% eager; the initial whole-message naive fallback cost about 29% on the reported BF16 graph case. The unsafe release is a historical timing reference only.

The host was shared, the other four GPUs were busy, and three to five rounds do not establish a general sub-percent regression bound. Timed graph replays reuse inputs within a sample; the separate changing-input regression provides stronger race coverage. TP8/gfx942 configurations and end-to-end model throughput were not measured. This PR keeps the upstream write-mode dispatch, including its partial-tile handling.

Validation

The previously tested tail specialization passed 8,400 changing-input checks per rank at TP4 across BF16/FP16/FP32 and 600 per rank at TP2. The unmodified vLLM PR #58433 test file passed all four BF16/MXFP8 TP2/TP4 cases in 74.15 seconds. The unsafe release failed the reported shape at eager iteration 0, with 125 mismatches and a maximum absolute error of 22. The actual main-based branch and pristine main both clean-built for gfx950. All 48 gfx950 specialized read-kernel variants match the previously tested tail kernels byte-for-byte, including resource metadata; all 24 current-main read variants match the benchmarked #5694 baseline. All 24 PR write-kernel variants match pristine main. On the actual current-main PR branch, BF16/FP16/FP32 TP4 each passed six shapes with 100 eager and 100 graph iterations per shape. The unmodified vLLM integration suite passed again: 4 passed in 71.75 seconds, covering BF16/MXFP8 at TP2/TP4. Ruff, Black, Python syntax, and diff checks passed.

Follow-up validation repeated the unchanged vLLM test group in 50 fresh pytest processes: 200/200 cases passed, with zero failures, errors, skips, timeouts, or retries (68.2 minutes). Six runs emitted TCPStore shutdown warnings while all tests passed. The kernel and compiled collective library are unchanged by the benchmark-list follow-up.

The three boundary sizes also passed 18/18 cases through the existing benchmark at TP4: BF16/FP16/FP32 in eager and graph modes. Profiler records confirmed the intended two-stage tail specialization for both neighbors and the aligned specialization for the middle size on all four ranks; 21 size-list/cutoff checks passed. These benchmark checks use the existing random-input tolerance. The default sweep now includes these shapes without -s, for example:

python op_tests/multigpu_tests/test_custom_allreduce.py -t 4 -d bf16 -m graph

Regression command (repeat for fp16 and fp32):

HIP_VISIBLE_DEVICES=0,1,2,3 python op_tests/multigpu_tests/test_custom_allreduce.py \
    --regression -t 4 -d bf16 --iterations 100

Standalone microbenchmark reproducer

Build/install each revision using the repository's usual source-build instructions, save the following as bench_allreduce_tail.py, and run it against each revision. The output records the loaded collective library's SHA256, per-rank timings, and correctness results. Alternate revision order between repeated runs.

HIP_VISIBLE_DEVICES=0,1,2,3 python bench_allreduce_tail.py \
    --output latency.json --case-set all --dtypes bf16 fp16 fp32 \
    --modes eager graph --calls 64 --replays 10 --samples 7 --warmup 5

"""Direct AITER TP4 latency/correctness study; host coordination uses Gloo."""

import argparse
import datetime
import hashlib
import json
import os
from pathlib import Path
import platform
import socket
import statistics
import time


def cases_for(args):
    cases = []
    for dtype in args.dtypes:
        element_size = 4 if dtype == "fp32" else 2
        candidates = []
        if args.case_set in ("all", "regression"):
            candidates += [(f"shape_{m}x512", (m, 512)) for m in (1023, 1024, 1025)]
        if args.case_set in ("all", "representative"):
            candidates += [(f"bytes_{n}", (n // element_size,))
                           for n in (8192, 65536, 262144, 524288, 1048576, 4194304, 8388608)]
        if args.case_set in ("all", "boundary"):
            candidates += [(f"tile_boundary_{offset:+}", ((1048576 + offset) // element_size,))
                           for offset in (-64, 0, 64, 63 * 64, 64 * 64, 65 * 64)]
        if args.case_set == "upstream":
            candidates += [(f"upstream_{m}x2880", (m, 2880))
                           for m in (120, 121, 127, 157, 160, 161, 256, 512)]
        seen = set()
        for name, shape in candidates:
            numel = 1
            for n in shape:
                numel *= n
            if numel in seen:
                continue
            seen.add(numel)
            cases.append({"name": name, "dtype": dtype, "shape": shape,
                          "bytes": numel * element_size, "numel": numel,
                          "full_packed_tile": numel * element_size % (512 * 16) == 0})
    return cases


def worker(rank, port, args):
    import torch
    import torch.distributed as dist
    torch.set_num_threads(1)
    torch.cuda.set_device(rank)
    device = torch.device("cuda", rank)
    dist.init_process_group("gloo", init_method=f"tcp://127.0.0.1:{port}",
                            rank=rank, world_size=4,
                            timeout=datetime.timedelta(seconds=60))
    from aiter.dist.device_communicators.custom_all_reduce import CustomAllreduce
    from aiter.jit import core
    ar = None
    report = {"config": vars(args), "results": []}
    try:
        ar = CustomAllreduce(dist.group.WORLD, device, max_size=16 * 1024 * 1024)
        assert not ar.disabled
        collective = Path(core.get_module("module_custom_all_reduce").__file__)
        report["metadata"] = {
            "python": platform.python_version(), "torch": torch.__version__,
            "device": torch.cuda.get_device_name(rank),
            "arch": torch.cuda.get_device_properties(rank).gcnArchName,
            "visible_devices": os.environ.get("HIP_VISIBLE_DEVICES"),
            "collective": str(collective),
            "collective_sha256": hashlib.sha256(collective.read_bytes()).hexdigest(),
        }
        dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
        cases = cases_for(args)
        # Alternate case order as well as variant order across paired rounds.
        if args.reverse_cases:
            cases.reverse()
        for case in cases:
            x = torch.empty(case["shape"], dtype=dtype_map[case["dtype"]], device=device)
            indices = torch.arange(case["numel"], dtype=torch.int32, device=device).reshape(case["shape"])

            def set_input(seed):
                local = ((indices * (rank + 1) + seed * (rank + 3)) % 31) - 15
                x.copy_(local)
                ref = sum((((indices * (r + 1) + seed * (r + 3)) % 31) - 15)
                          for r in range(4))
                return ref.to(x.dtype)

            def check(outputs, expected):
                # Check every captured call: racing intermediate outputs also count.
                mismatch = torch.zeros((), dtype=torch.int64, device=device)
                max_error = torch.zeros((), dtype=torch.float32, device=device)
                finite = torch.ones((), dtype=torch.bool, device=device)
                for out in outputs:
                    mismatch += torch.count_nonzero(out != expected)
                    finite &= torch.isfinite(out).all()
                    max_error = torch.maximum(max_error, (out.float() - expected.float()).abs().max())
                return {"mismatches": mismatch.item(), "finite": finite.item(),
                        "max_abs_error": max_error.item()}

            for mode in args.modes:
                expected = set_input(1)
                for _ in range(args.warmup):
                    warm = ar.custom_all_reduce(x)
                    assert warm is not None
                torch.cuda.synchronize()
                dist.barrier()
                graph = None
                outputs = []
                if mode == "graph":
                    stream = torch.cuda.Stream()
                    stream.wait_stream(torch.cuda.current_stream())
                    graph = torch.cuda.CUDAGraph()
                    with ar.capture():
                        with torch.cuda.graph(graph, stream=stream):
                            outputs = [ar.custom_all_reduce(x) for _ in range(args.calls)]
                    torch.cuda.current_stream().wait_stream(stream)
                    for _ in range(args.warmup):
                        graph.replay()
                    torch.cuda.synchronize()
                    dist.barrier()

                timings = []
                wall_timings = []
                correctness = []
                start = torch.cuda.Event(enable_timing=True)
                end = torch.cuda.Event(enable_timing=True)
                for sample in range(args.samples):
                    expected = set_input(100 + sample)
                    if mode == "graph":
                        for out in outputs:
                            out.fill_(float("nan"))
                    torch.cuda.synchronize()
                    dist.barrier()
                    before = time.perf_counter()
                    start.record()
                    if mode == "graph":
                        for _ in range(args.replays):
                            graph.replay()
                        n_ops = args.calls * args.replays
                    else:
                        outputs = [ar.custom_all_reduce(x) for _ in range(args.calls)]
                        n_ops = args.calls
                    end.record()
                    end.synchronize()
                    wall_timings.append((time.perf_counter() - before) * 1e6 / n_ops)
                    timings.append(start.elapsed_time(end) * 1000 / n_ops)
                    correctness.append(check(outputs, expected))
                    dist.barrier()

                local = {"rank": rank, "samples_us": timings,
                         "median_us": statistics.median(timings),
                         "wall_samples_us": wall_timings,
                         "correctness": correctness}
                gathered = [None] * 4
                dist.all_gather_object(gathered, local)
                if rank == 0:
                    result = {**case, "mode": mode, "ranks": gathered,
                              "max_rank_median_us": max(r["median_us"] for r in gathered),
                              "exact": all(c["mismatches"] == 0 for r in gathered for c in r["correctness"]),
                              "finite": all(c["finite"] for r in gathered for c in r["correctness"])}
                    report["results"].append(result)
                    destination = Path(args.output)
                    destination.parent.mkdir(parents=True, exist_ok=True)
                    temporary = destination.with_suffix(".tmp")
                    temporary.write_text(json.dumps(report, indent=2) + "\n")
                    temporary.replace(destination)
                    print(f"RESULT {case['dtype']} {case['name']} {mode} "
                          f"us={result['max_rank_median_us']:.4f} "
                          f"exact={result['exact']} finite={result['finite']}", flush=True)
                del outputs, graph
                torch.cuda.synchronize()
                dist.barrier()
            del x, indices, expected, warm
            torch.cuda.empty_cache()
        dist.barrier()
    finally:
        if ar is not None:
            ar.close()
        dist.destroy_process_group()


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--output", required=True)
    parser.add_argument("--dtypes", nargs="+", choices=("bf16", "fp16", "fp32"),
                        default=["bf16", "fp16", "fp32"])
    parser.add_argument("--modes", nargs="+", choices=("eager", "graph"), default=["eager", "graph"])
    parser.add_argument("--case-set", choices=("all", "regression", "representative", "boundary", "upstream"), default="all")
    parser.add_argument("--calls", type=int, default=64)
    parser.add_argument("--replays", type=int, default=10)
    parser.add_argument("--samples", type=int, default=7)
    parser.add_argument("--warmup", type=int, default=5)
    parser.add_argument("--reverse-cases", action="store_true")
    parser.add_argument("--list-cases", action="store_true")
    args = parser.parse_args()
    assert min(args.calls, args.replays, args.samples, args.warmup) > 0
    return args


if __name__ == "__main__":
    args = parse_args()
    if args.list_cases:
        print(json.dumps(cases_for(args), indent=2))
    else:
        import torch.multiprocessing as mp
        with socket.socket() as sock:
            sock.bind(("127.0.0.1", 0))
            port = sock.getsockname()[1]
        mp.spawn(worker, args=(port, args), nprocs=4, join=True)

History and motivation

This work began when vLLM PR #58433 intermittently failed TP4 tests in Buildkite build #90750 despite passing on retry. The (1025, 512) case exposed block barriers inside a loop whose iteration count differed between threads at a partial tile. A local regression reproduced corruption in the original release. Routing partial messages to the naive kernel restored correctness but increased BF16 graph latency by about 29%. Merged AITER PR #5694 fixes the barrier defect with uniform loop bounds. This follow-up retains its write-mode and fused fixes while specializing complete tiles and leftover packs in the plain read kernel. The motivation is to preserve the correctness fix with minimal performance cost and keep the original test shapes and assertions.

AI assistance was used for implementation, investigation, and validation.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5807 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.