Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
vgokhale merged 2 commits into from
Sep 25, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
565d859
[HIP] [Perf] Specialize partial tiles in two-stage custom all-reduce
AndreasKaratzas Sep 24, 2026
ebdf8b4
[HIP] [Bench] Include custom all-reduce tail shapes in the default sweep
AndreasKaratzas Sep 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 37 additions & 24 deletions csrc/include/custom_all_reduce.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -480,7 +480,7 @@ __global__ void __launch_bounds__(512, 1) cross_device_reduce_1stage(RankData* _
end_sync<ngpus, true>(sg, self_sg, rank);
}

template <typename T, int ngpus, bool is_broadcast_reg_outptr = false>
template <typename T, int ngpus, bool is_broadcast_reg_outptr = false, bool has_tail = false>
__global__ void __launch_bounds__(512, 1) cross_device_reduce_2stage(RankData* _input_dp,
RankData* _output_dp,
RankSignals sg,
Expand Down Expand Up @@ -516,19 +516,19 @@ __global__ void __launch_bounds__(512, 1) cross_device_reduce_2stage(RankData* _
}
auto tmp_out = tmps[0];
start_sync<ngpus>(sg, self_sg, rank);
// stage 1: reduce scatter
// `end - start` is not generally a multiple of `stride`, so a plain grid-stride loop gives
// the block a non-uniform trip count and the two __syncthreads() below pair up across
// iterations.
int loop_end = start + (end - start + stride - 1) / stride * stride;
for(int idx = start + tid; idx < loop_end; idx += stride)
// stage 1: reduce scatter. The optimized dispatch has equal rank partitions.
int loop_end = end;
if constexpr(has_tail)
{
bool active = idx < end;
if(active)
*(reinterpret_cast<P*>(&tmp_smem[0]) + threadIdx.x) = ptrs[warp_id][idx];
loop_end -= (end - start) % tnum_gpu;
}
int idx = start + tid;
for(; idx < loop_end; idx += stride)
{
*(reinterpret_cast<P*>(&tmp_smem[0]) + threadIdx.x) = ptrs[warp_id][idx];
__syncthreads();
// cal add in first 64 threads
if(warp_id == 0 && active)
if(warp_id == 0)
{
A add_reg;
#pragma unroll
Expand Down Expand Up @@ -557,6 +557,15 @@ __global__ void __launch_bounds__(512, 1) cross_device_reduce_2stage(RankData* _
}
__syncthreads();
}
if constexpr(has_tail)
{
// Preserve the full-tile loop and its barriers. The remaining packs are
// reduced by their original owner threads without shared memory.
if(warp_id == 0 && idx < end)
{
tmp_out[idx - start] = packed_reduce<P, ngpus, A>(ptrs, idx);
}
}
end_sync<ngpus>(sg, self_sg, rank);

// stage 2: allgather. Note: it's important to match the tid between
Expand Down Expand Up @@ -4621,19 +4630,19 @@ class CustomAllreduce
}
}

#define KL(ngpus, name) \
do \
{ \
if(is_broadcast_reg_outptr) \
{ \
name<T, ngpus, true><<<blocks, threads, 0, stream>>>( \
input_ptrs, output_ptrs, sg_, self_sg_, output, rank_, size); \
} \
else \
{ \
name<T, ngpus, false><<<blocks, threads, 0, stream>>>( \
input_ptrs, output_ptrs, sg_, self_sg_, output, rank_, size); \
} \
#define KL(ngpus, name, ...) \
do \
{ \
if(is_broadcast_reg_outptr) \
{ \
name<T, ngpus, true __VA_OPT__(, ) __VA_ARGS__><<<blocks, threads, 0, stream>>>( \
input_ptrs, output_ptrs, sg_, self_sg_, output, rank_, size); \
} \
else \
{ \
name<T, ngpus, false __VA_OPT__(, ) __VA_ARGS__><<<blocks, threads, 0, stream>>>( \
input_ptrs, output_ptrs, sg_, self_sg_, output, rank_, size); \
} \
} while(0)

#define DISPATCH_REDUCE(ngpus, name) \
Expand All @@ -4645,6 +4654,10 @@ class CustomAllreduce
{ \
KL(ngpus, name##_write_mode); \
} \
else if(size % THREAD_NUM != 0) \
{ \
KL(ngpus, name, true); \
} \
else \
{ \
KL(ngpus, name); \
Expand Down
144 changes: 141 additions & 3 deletions op_tests/multigpu_tests/test_custom_allreduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

import argparse
import logging
Expand Down Expand Up @@ -123,6 +123,118 @@ def test_allreduce_custom(
}


def _allreduce_tail_regression(
rank, tp_size, shapes, dtype, iterations, distributed_init_method
):
from aiter.dist.device_communicators.custom_all_reduce import CustomAllreduce

torch.set_num_threads(1)
torch.cuda.set_device(rank)
device = torch.device(f"cuda:{rank}")
dist.init_process_group(
"gloo",
world_size=tp_size,
rank=rank,
init_method=distributed_init_method,
)
ca_comm = None
try:
ca_comm = CustomAllreduce(dist.group.WORLD, device)
assert not ca_comm.disabled, "The tail regression requires custom allreduce"

for shape in shapes:
x = torch.zeros(shape, dtype=dtype, device=device)
assert ca_comm.should_custom_ar(x), f"Custom allreduce rejected {shape=}"
# Warm up the actual eager path before capturing its graph variant.
ca_comm.custom_all_reduce(x, use_new=True)
torch.cuda.synchronize()
graph = torch.cuda.CUDAGraph()
stream = torch.cuda.Stream()
stream.wait_stream(torch.cuda.current_stream())
with ca_comm.capture(), torch.cuda.graph(graph, stream=stream):
# Register a graph-owned intermediate, as in a captured model,
# while the external input remains mutable between replays.
graph_out = ca_comm.custom_all_reduce(x.clone(), use_new=True)
torch.cuda.current_stream().wait_stream(stream)
assert graph_out is not None

generator = torch.Generator().manual_seed(2026)
for iteration in range(iterations):
# Independent rank inputs, with sums exactly representable even
# in BF16. A FP32 reference avoids tolerances hiding stale data.
inputs = torch.randint(
-8, 9, (tp_size, *shape), generator=generator, dtype=torch.int16
)
ref = inputs.sum(dim=0, dtype=torch.float32).to(device, dtype)
x.copy_(inputs[rank])
eager_out = ca_comm.custom_all_reduce(x, use_new=True)
assert eager_out is not None
context = f"{rank=} {shape=} {dtype=} {iteration=}"
torch.testing.assert_close(
eager_out,
ref,
rtol=0,
atol=0,
msg=lambda msg, context=context: f"{context} eager\n{msg}",
)
# Change the input on every replay and poison the output so
# replaying stale data or skipping a tail store cannot pass.
graph_out.fill_(float("nan"))
graph.replay()
torch.testing.assert_close(
graph_out,
ref,
rtol=0,
atol=0,
msg=lambda msg, context=context: f"{context} graph\n{msg}",
)
if rank == 0:
logger.info(
"tail regression passed: %s, %s, TP%d, %d eager + graph runs",
shape,
dtype,
tp_size,
iterations,
)
# All ranks must finish reading peer buffers before any are freed.
dist.barrier()
finally:
if ca_comm is not None:
ca_comm.close()
dist.destroy_process_group()


def test_allreduce_tail_regression(tp_size, dtype, iterations=20, shapes=None):
# 1025x512 reproduced a TP4 BF16 graph-replay failure. The adjacent
# dimensions exercise both partial and complete blocks in reduce-scatter.
if shapes is None:
shapes = [(1023, 512), (1024, 512), (1025, 512)]
if tp_size in (4, 8):
pack_size = 16 // torch.empty((), dtype=dtype).element_size()
# The optimized kernel assigns 512 / tp_size packs to each rank
# per tile. Check one active lane, one inactive lane, and a size
# not divisible across ranks (the naive path). These small cases
# use an uncapped grid; the original shapes exceed the grid cap.
# 32 tiles is above the two-stage threshold for both TP4 and TP8.
base_packs = 32 * 512
shapes.extend(
((base_packs + offset) * pack_size,)
for offset in (-tp_size, tp_size, 1)
)
torch.multiprocessing.spawn(
_allreduce_tail_regression,
args=(
tp_size,
shapes,
dtype,
iterations,
get_distributed_init_method(get_ip(), get_open_port()),
),
nprocs=tp_size,
join=True,
)


# Custom-AR size cutoff (bytes). Mirrors _DEFAULT_CAR_MAX_SIZE in
# aiter/dist/device_communicators/custom_all_reduce.py and honors the same
# AITER_CUSTOM_AR_MAX_SIZE override. Inputs at or below this run on the custom
Expand All @@ -143,7 +255,8 @@ def _car_max_bytes() -> int:

def gen_sizes(dtype) -> list[int]:
"""Element counts to sweep: [1024, 2048, 4096] then 7168*k / 8192*k for
k = 1, 2, 4, 8, ... up to the largest size custom_all_reduce serves."""
k = 1, 2, 4, 8, ... plus the flattened (1023/1024/1025, 512) tail cases,
up to the largest size custom_all_reduce serves."""
itemsize = torch.empty(0, dtype=dtype).element_size()
max_numel = _car_max_bytes() // itemsize
sizes = [n for n in (1024, 2048, 4096) if n <= max_numel]
Expand All @@ -158,7 +271,10 @@ def gen_sizes(dtype) -> list[int]:
if not added:
break
k *= 2
return sizes
# Include partial tiles around the aligned 1024x512 case. Flattening the
# contiguous inputs preserves the kernel dispatch and reported tail shape.
sizes.extend(n for n in (1023 * 512, 1024 * 512, 1025 * 512) if n <= max_numel)
return sorted(set(sizes))


l_dtype = ["bf16", "fp16", "fp32"]
Expand Down Expand Up @@ -197,12 +313,34 @@ def gen_sizes(dtype) -> list[int]:
default=None,
help="single shape override, e.g. -s 128,8192 (default: swept size list)",
)
parser.add_argument(
"--regression",
action="store_true",
help="run only tail correctness checks in both eager and graph modes",
)
parser.add_argument(
"--iterations",
type=int,
default=20,
help="number of changing-input iterations per regression shape (default: 20)",
)


if __name__ == "__main__":
freeze_support()
args = parser.parse_args()
if args.iterations < 1:
parser.error("--iterations must be positive")
dtype = dtypes.d_dtypes[args.dtype]
if args.regression or args.shape is None:
test_allreduce_tail_regression(
args.tp_size,
dtype,
args.iterations,
shapes=[args.shape] if args.shape is not None else None,
)
if args.regression:
raise SystemExit(0)
with_graph = args.mode == "graph"
if args.shape is not None:
l_shape = [args.shape]
Expand Down
Loading