Add HIP-graph capture coverage for the sharded-relay collectives by srinathb-meta · Pull Request #3850 · meta-pytorch/torchcomms

srinathb-meta · 2026-08-27T14:49:06Z

Summary:
The sharded-relay collectives had zero graph-capture test coverage, so the
four known graph hazards in the relay were all unverifiable:

G1 ScratchBufferCache allocates with hipMallocAsync/hipFreeAsync keyed on
(device, stream, key), which under capture becomes a graph allocation
node whose address is only valid while that graph executes.
G2 The one-shot IPC handshake advances a host-side epoch that reaches the
kernel as a by-value argument, so a replay would spin on a baked epoch
that every flag already satisfies.
G3 One-shot region creation does a synchronous hipMemset plus two
bootstrapAllGather calls, none of which are capturable.
G4 The reduce-scatter overlap side stream is forked from the caller stream,
whose captured form has no relationship to the uncaptured stream.

G2, G3 and G4 currently sit behind explicit capture bail-outs and G1 is
latent, so every case here passes as-is. That is the point: this is the
regression baseline the four fixes get measured against, and each fix brings
its own failing-before/passing-after case with it rather than landing one red
test now.

Covers all four collectives on a single group with A=4 (ranks 0-3 active,
4-7 helpers): capture, then replay with input that varies per replay, then
verify each replay. Two details are load-bearing and were both found the hard
way while getting this to fail for the right reasons:

Odd ranks sleep on the host before the launch. Without that skew a stale
baked epoch or stale staging is only a race, because the kernel still
pushes real data on every replay. With it the unskewed rank pushes, flags,
falls through the handshake and reduces data the skewed rank has not
written yet, so the mismatch is deterministic.
Cases warm up at a size large enough to fill the scratch cache for the
route the capture will take. Below the 1 MiB one-shot band an uncaptured
call goes one-shot while the captured call falls back to PureDirect, so
warming up at the size under test leaves that cache cold and the capture
allocates inside itself. Left unprimed the one-shot-band case hung for the
whole harness timeout, and it was hanging on G1 rather than on the one-shot
behaviour it is named for.

hipStreamSynchronize is replaced by a bounded poll: a graph-compat regression
shows up as a kernel that never retires, and a suite that hangs for the
harness timeout reports nothing useful and takes every later case down with
it.

Reviewed By: JinghanHuang

Differential Revision: D117629710

Summary: The A=4 sharded-relay all-to-all splits each off-diagonal segment into `[directA | relay | directB]` and routes the middle region two hops through one helper while both direct regions take one hop. `relay` was 128-element aligned but `directA` was taken as an exact `segmentCount / 3`, so the relay region started on an aligned offset only when the segment count happened to divide by three -- which a power-of-two segment never does, since 2^n is never divisible by 3. Forcing the route on at those sizes measures **0.75x** vs NCCL, i.e. the relay ran a third *slower* than the plain direct exchange it was supposed to beat. The `[63 MiB, 256 MiB)` routing window was masking this rather than reflecting a crossover. Inside it, 63 / 135 / 144 MB all divide into thirds cleanly and the relay measured 1.25-1.31x; at and above 256 MiB the route fell back to pure direct, where the helpers contribute nothing and the collective sat at 0.94-1.01x. Align `directA` down to the same chunk as the relay region. The two are equal by construction -- the schedule's two serialized phases each carry exactly one such chunk per link, which is what puts the optimum at a third -- so this is one shared count with `directB` absorbing the `/3` remainder and the alignment loss, and every region boundary is now 128-element aligned. With that fixed the relay wins at every size above the crossover, so: - the 256 MiB ceiling is removed, and - the lower bound drops from 63 MiB to 27 MiB fused / 9 MiB independent (re-measured; an independent call has the cross links to itself, so it crosses over earlier, and forcing the relay below these points regresses -- fused 13.5 MB goes 1.48x -> 1.21x). **Also evaluated and NOT included: striping each pair's relay region across all four helpers.** A source has only A-1 = 3 destinations, so pinning one helper per (source, dest) pair can reach at most 3 of its 4 cross links and the fourth idles, capping the schedule at 1.5x; striping balances all four for a 1.67x ceiling. Measured, striping needs 12 relay sends per phase instead of 3 and lands at 1.15-1.25x against the aligned single-helper schedule's 1.28-1.36x, so the op-count cost exceeds the better ceiling. Dropped. Perf -- single-group A=4 all-to-all (MI350X, 8 GPUs, bf16, best-of-100, speedup = 4-rank NCCL all-to-all with the other 4 ranks idle / relay), same-day before and after: ``` Msg Size before after 9 MB 0.99x 1.07x 13.5 MB 0.85x 1.12x 27 MB 1.03x 1.20x 31.5 MB 1.02x 1.21x 36 MB 1.13x 1.21x 63 MB 1.26x 1.28x 72 MB 1.25x 1.28x 135 MB 1.20x 1.20x 256 MB 1.06x 1.24x 512 MB 1.00x 1.33x 1 GB 1.01x 1.32x ``` Fused (2 groups) after, vs the checked-in `perf_data/bench_sharded_relay_fused_sweep_results.txt` baseline: ``` Msg Size before after 27 MB 1.06x 1.23x 36 MB 1.04x 1.24x 256 MB 1.01x 1.27x 512 MB 0.94x 1.24x 1 GB 0.99x 1.29x ``` No size regresses in either scenario. The A=2 routes, the pure-direct route and the other three collectives are untouched. Differential Revision: D117629683

…ex (3.2x -> 4.3x) Summary: The 2-active sharded-relay all-gather runs two serialized ncclGroups: group 1 scatters each active rank's chunks to the helpers, group 2 has the helpers forward them on. When nGroups == 1 the active ranks and the helpers are DISJOINT sets, so a cross link carries only the scatter in one direction and only the forward in the other -- which means that schedule leaves every cross link HALF-DUPLEX for the whole call. The forward direction is idle for all of group 1 and the scatter direction for all of group 2. Tile the relay into T tiles and issue tile t's forward in the SAME group as tile t+1's scatter. Both directions are then busy in every group. Per link and direction each of the T+1 groups carries one unit u against a count of ((H+1)*T + 1)*u, so the per-link cost is (T+1)/((H+1)*T + 1) of the count: count/4 at T = 1, which is exactly the existing two-group schedule, falling towards count/7 as T grows on an 8-GPU node. A 4x ceiling becomes 7x. count/7 is also the hard floor -- an active rank has to move its whole buffer out across its 7 links exactly once -- so there is nothing beyond this to find. Why this is gated on nGroups == 1: in a FUSED call every rank is active for one group and a helper for the others, so its scatter and its forward are egress on the SAME link direction. They add rather than overlap, the link is already full, and the extra boundaries are pure cost. That is the same reason the 16-stage pipeline removed in D115998361 lost, and fused numbers here are bit-identical because relayA2PipelineTiles() returns 1 and the original function runs. Details worth knowing: - Depth is chosen by minimizing (T + 1) * (perStageBytes + kRelayPipelineBoundaryBytes) over powers of two, i.e. the two competing terms stated directly rather than a size threshold per depth. With the boundary priced at 768 KiB (~25 us at the measured ~50 GB/s per link) this picks the measured-best depth at every size on the sweep. - Depth is capped at 4 for STABILITY, not throughput. Every group boundary is a cross-rank sync point, so a deeper pipeline gives co-resident independent jobs more room to skew against each other. At depth 8 the 4-job parallel sweep goes erratic and non-reproducible -- 0.74x-1.46x outliers scattered across the size axis, landing on different sizes each run -- while depth 4 is as stable as the unpipelined schedule. Depth 4 is also what the cost model picks at every measured single-group size, so the cap costs nothing there. - Helper staging drops from the whole chunk to two ping-pong units per active source (receive into k%2, forward (k-1)%2, disjoint by construction and separated by a group boundary). A forwarded tile is therefore still cache-resident when it is read back, instead of a full HBM round trip. - Every rank posts exactly one send and one recv per link per group, which is also the best case for p2p channel assignment. Perf -- single-group 2-active all-gather (MI350X, 8 GPUs, bf16, best-of-100, speedup = 2-rank NCCL all-gather with the other 6 ranks idle / relay), same-day before and after: ``` Msg Size before after relay ms 9 MB 2.04x 2.03x 0.111 -> 0.111 (depth 1, unchanged) 13.5 MB 2.29x 2.28x 0.140 -> 0.141 (depth 1, unchanged) 27 MB 2.67x 2.78x 0.223 -> 0.208 31.5 MB 2.76x 2.90x 0.251 -> 0.232 36 MB 2.79x 2.97x 0.277 -> 0.258 63 MB 3.02x 3.49x 0.436 -> 0.381 67.5 MB 3.05x 3.46x 0.463 -> 0.409 72 MB 3.05x 3.53x 0.487 -> 0.421 135 MB 3.08x 3.82x 0.872 -> 0.710 144 MB 3.06x 3.84x 0.927 -> 0.742 256 MB 3.15x 4.08x 1.596 -> 1.234 512 MB 3.21x 4.26x 3.118 -> 2.343 1 GB 3.20x 4.31x 6.181 -> 4.589 ``` Everything below 9 MB is on the pure-direct route and untouched. Fused (4 groups) is unchanged by construction, and re-measured to confirm: 2.67x / 3.02x / 3.11x / 3.22x at 27 MB / 63 MB / 135 MB / 1 GB. Parallel (4 co-resident jobs) is within its run-to-run noise, trending better at the large sizes, against a same-day depth-1 baseline: ``` Msg Size depth 1 after 27 MB 2.51x 2.27x 63 MB 2.81x 2.52x 135 MB 2.63x 2.63x 256 MB 2.61x 2.75x 512 MB 2.67x 3.14x 1 GB 2.70x 3.00x ``` That sweep swings ~13% run to run at a fixed depth (63 MB measured 0.467 ms and 0.530 ms on two depth-4 runs), so the mid-range deltas are not resolvable; what matters is that the erratic behaviour seen at depth 8 is gone. The 4-active all-gather, the pure-direct route and the other three collectives are untouched. Differential Revision: D117629684

…ex (3.2x -> 4.2x) Summary: Same change as the pipelined 2-active all-gather, applied to the 2-active all-to-all, whose relay has the identical shape: the exchange segment is scattered to the helpers in one ncclGroup and forwarded on in a second. With nGroups == 1 the active ranks and the helpers are DISJOINT sets, so a cross link carries only the scatter in one direction and only the forward in the other, and the two-group schedule leaves every cross link HALF-DUPLEX for the whole call. Tiling the relay and issuing tile t's forward in the same group as tile t+1's scatter fills both directions, taking the per-link cost from 2*count/(H+2) = count/4 towards count/7 -- the floor, since an active rank must move its exchange segment out across its 7 links exactly once. The depth selection, the boundary cost model, the depth-4 stability cap and the nGroups == 1 gate are all the shared relayA2PipelineTiles() introduced with the all-gather change; nothing new is added here. Fused calls keep the original schedule bit-for-bit because that helper returns 1 for nGroups > 1. The exchange segment splits the same way as the all-gather send buffer -- an offload region of H chunks of T tiles plus a direct region of T+1 chunks, the last absorbing the remainder -- and helper staging is again two ping-pong units per active source instead of the whole chunk, so a forwarded tile is read back out of cache. Perf -- single-group 2-active all-to-all (MI350X, 8 GPUs, bf16, best-of-100, speedup = 2-rank NCCL all-to-all with the other 6 ranks idle / relay), same-day before and after: ``` Msg Size before after relay ms 13.5 MB 0.92x 1.08x 0.184 -> 0.177 (depth 1; NCCL-side noise) 27 MB 2.08x 2.28x 0.151 -> 0.139 31.5 MB 2.06x 2.25x 0.167 -> 0.148 36 MB 2.43x 2.49x 0.167 -> 0.162 63 MB 2.68x 2.75x 0.248 -> 0.227 67.5 MB 2.56x 2.87x 0.267 -> 0.238 72 MB 2.63x 2.94x 0.276 -> 0.251 135 MB 2.93x 3.34x 0.459 -> 0.403 144 MB 2.96x 3.42x 0.486 -> 0.419 256 MB 3.05x 3.69x 0.822 -> 0.678 512 MB 3.11x 4.04x 1.591 -> 1.222 1 GB 3.16x 4.20x 3.119 -> 2.340 ``` Everything below 27 MB is on the pure-direct route for an independent call and is untouched. Fused (4 groups) re-measured to confirm no change: 2.35x / 2.67x / 2.82x / 3.18x at 27 MB / 63 MB / 135 MB / 1 GB, matching the checked-in `perf_data/bench_sharded_relay_fused_sweep_results.txt`. Parallel (4 co-resident jobs) improves as well, and stays smooth -- none of the depth-8 erraticism that set the depth cap: ``` Msg Size before after 27 MB 1.50x 2.04x 63 MB 1.91x 2.35x 135 MB 2.29x 2.46x 144 MB 2.31x 2.67x 256 MB 2.53x 2.64x 512 MB 2.53x 2.68x ``` (before from the checked-in parallel results file.) The 4-active all-to-all routes and the pure-direct route are untouched. Differential Revision: D117629685

…x (3.2x -> 4.4x) Summary: Third application of the duplex pipelining introduced for the 2-active all-gather, now on the 2-active allreduce. With nGroups == 1 the active ranks and the helpers are DISJOINT sets, so a cross link carries only the scatter in one direction and only the reduced forward in the other, and the two-group schedule leaves every cross link HALF-DUPLEX for the whole call. Tiling the relay and issuing tile t's forward in the same group as tile t+1's scatter fills both directions, taking the per-link cost from 2*count/(H+2) = count/4 towards count/7 -- the floor, since an active rank must move its buffer out across its 7 links exactly once. The one thing specific to allreduce is the reduce-at-helper. Both active ranks send the same logical tile index, so the helper's sum IS the final allreduced value and the return hop stays one chunk per active rank. In the pipelined schedule that reduce is issued between the group that receives tile k and the group that forwards it -- one launch per tile over u elements instead of one over the whole chunk. Because helper staging is two ping-pong units per active source rather than the whole chunk, each tile is summed and forwarded while still cache-resident instead of making a full HBM round trip. Depth selection, the boundary cost model, the depth-4 stability cap and the nGroups == 1 gate are the shared relayA2PipelineTiles() from the all-gather change. The small-message pure-direct route stays where it is, inside shardedRelayAllReduce2Active, and never pipelines: the dispatch only asks for a depth once selectAllReduceRoute() has chosen the relay. In-place and out-of-place both work. The relay region cannot alias dangerously even in place, because group k reads tile k out of sendBuff while writing the reduced tile k-1 into recvBuff, which are different offsets. Perf -- single-group 2-active allreduce (MI350X, 8 GPUs, bf16, best-of-100, speedup = 2-rank NCCL allreduce with the other 6 ranks idle / relay), same-day before and after: ``` Msg Size before after relay ms 9 MB 1.91x 1.94x 0.118 -> 0.116 (depth 1, unchanged) 13.5 MB 2.19x 2.21x 0.144 -> 0.143 (depth 1, unchanged) 27 MB 2.56x 2.78x 0.226 -> 0.210 31.5 MB 2.65x 2.90x 0.252 -> 0.230 36 MB 2.70x 2.92x 0.278 -> 0.256 63 MB 2.91x 3.30x 0.437 -> 0.386 67.5 MB 2.94x 3.34x 0.461 -> 0.406 72 MB 2.96x 3.41x 0.485 -> 0.424 135 MB 3.09x 3.84x 0.853 -> 0.688 144 MB 3.12x 3.92x 0.902 -> 0.718 256 MB 3.18x 4.14x 1.558 -> 1.195 512 MB 3.19x 4.30x 3.081 -> 2.287 1 GB 3.24x 4.39x 6.065 -> 4.470 ``` Everything below 6 MB is on the pure-direct route for an independent call and is untouched. Fused (4 groups) re-measured to confirm no change: 2.43x / 2.75x / 2.93x / 3.06x at 27 MB / 63 MB / 135 MB / 1 GB. Parallel (4 co-resident jobs) improves, against the checked-in baseline: ``` Msg Size before after 27 MB 1.48x 2.18x 31.5 MB 1.79x 2.25x 36 MB 1.61x 2.32x 63 MB 2.23x 2.31x 144 MB 2.30x 2.55x 256 MB 2.22x 2.71x 512 MB 2.65x 2.72x ``` The 4-active flat allreduce and the pure-direct route are untouched. Differential Revision: D117629686

…duplex (2.9x -> 3.7x) Summary: Last of the four 2-active relays to get the duplex pipelining introduced for the all-gather. With nGroups == 1 the active ranks and the helpers are DISJOINT sets, so a cross link carries only the scatter in one direction and only the forward in the other, and the two-group schedule leaves every cross link HALF-DUPLEX for the whole call. Tiling the relay and issuing tile t's forward in the same group as tile t+1's scatter fills both directions, taking the per-link cost from 2*recvCount/(H+2) = recvCount/4 towards recvCount/7 -- the floor, since an active rank must move its foreign block out across its 7 links exactly once. Unlike allreduce, helpers stay pure PASSTHROUGH here and are not asked to reduce: slot 0 is a0's contribution to a1's output and slot 1 is a1's contribution to a0's output, which are different outputs with nothing in common to sum. The active rank reduces once at the end, and because foreignScratch still mirrors the output block layout that stays a single fused launch regardless of how many tiles the relay used. Depth selection, the boundary cost model, the depth-4 stability cap and the nGroups == 1 gate are the shared relayA2PipelineTiles() from the all-gather change. The small-message pure-direct route stays inside shardedRelayReduceScatter2Active and never pipelines. Helper staging drops from the whole chunk to two ping-pong units per active source. Reduce-scatter gains a little less than the other three (1.30x on relay latency at 1 GB versus 1.35x for all-gather) for two structural reasons worth recording: its relayed quantity is recvCount, i.e. half the message, so a given message lands at a shallower depth; and the closing full-buffer reduce is serial work the pipeline does not shrink. Perf -- single-group 2-active reduce-scatter (MI350X, 8 GPUs, bf16, best-of-100, speedup = 2-rank NCCL reduce-scatter with the other 6 ranks idle / relay), same-day before and after: ``` Msg Size before after relay ms 13.5 MB 1.06x 1.09x 0.183 -> 0.176 (depth 1, unchanged) 27 MB 2.38x 2.43x 0.149 -> 0.145 (depth 1, unchanged) 31.5 MB 2.45x 2.55x 0.166 -> 0.157 36 MB 2.59x 2.61x 0.178 -> 0.174 63 MB 2.86x 3.16x 0.266 -> 0.245 67.5 MB 2.90x 3.23x 0.285 -> 0.255 72 MB 2.91x 3.21x 0.297 -> 0.272 135 MB 3.21x 3.63x 0.497 -> 0.440 144 MB 3.20x 3.64x 0.524 -> 0.460 256 MB 3.00x 3.58x 0.915 -> 0.767 512 MB 2.95x 3.68x 1.777 -> 1.414 1 GB 2.88x 3.74x 3.499 -> 2.694 ``` Everything below 27 MB is on the pure-direct route for an independent call and is untouched. Fused (4 groups) re-measured to confirm no change: 2.40x / 2.77x / 3.15x / 2.91x at 27 MB / 63 MB / 135 MB / 1 GB. Parallel (4 co-resident jobs) is neutral to better, against the checked-in baseline: ``` Msg Size before after 27 MB 1.77x 2.00x 36 MB 1.86x 2.06x 63 MB 2.19x 2.46x 67.5 MB 2.36x 2.58x 135 MB 2.65x 2.58x 144 MB 2.58x 2.66x 256 MB 2.54x 2.50x 512 MB 2.45x 2.54x ``` A measurement note that cost me a false regression: the FIRST sweep after a rebuild reads 10-13% slow at 27-72 MB (cold scratch allocation and kernel warm-up), which showed up as reduce-scatter losing at 36 MB. Two warm repeats put it back at 0.173/0.173 ms against a 0.178 ms baseline. All numbers above are warm. The 4-active flat reduce-scatter and the pure-direct route are untouched. Differential Revision: D117629687

Summary: Extends the duplex pipelining from the 2-active relays to the 4-active all-to-all, and generalizes the depth selector so the remaining 4-active schedules can reuse it. The A=4 relay sends each (source, dest) pair's middle region two hops through the one helper the Latin permutation assigns it, in two serialized ncclGroups. With nGroups == 1 the active ranks and the helpers are disjoint, so those two hops sit on OPPOSITE directions of the same cross link and the two-group form leaves each of them half idle. Merging them by tiling makes every group carry one relay unit up and one down on each cross link, matched by one direct unit on the intra links, so with u = align(sc / (2*T + 1)) the cost is (T+1)*u: 2*sc/3 at T = 1 (identical to the current schedule) falling towards sc/2. The helper assignment is what makes the merged form clean. Each helper owns 3 tasks with three DISTINCT sources and three DISTINCT destinations, so per group every rank still posts exactly one send and one recv per link per direction -- 3 relay plus 3 direct each way for an active rank, 3 each way for a helper -- which is also the best case for p2p channel assignment. relayA2PipelineTiles() becomes relayPipelineTiles(), taking a RelayPipelineShape instead of hardcoding the 2-active geometry. A shape is just the two affine functions of depth that describe a schedule: how many units the count divides into, and how many units the busiest link direction ends up carrying. The 2-active shape is link {1,1} / total {H+1,1}; all-to-all at 4 active is link {1,1} / total {2,1}. Every shape reproduces the existing two-group schedule at depth 1, which is the check that it is written correctly. The four 2-active call sites move to relayShapeA2(numHelpers) and pick byte-identical depths -- verified numerically across the sweep and re-measured (2-active all-gather is unchanged to the millisecond at 4.589 ms / 1 GB). Perf -- single-group 4-active all-to-all (MI350X, 8 GPUs, bf16, best-of-100, speedup = 4-rank NCCL all-to-all with the other 4 ranks idle / relay), same-day before and after: ``` Msg Size before after relay ms 13.5 MB 1.12x 1.09x 0.114 -> 0.114 (depth 1, unchanged) 27 MB 1.20x 1.19x 0.167 -> 0.166 (depth 1, unchanged) 36 MB 1.21x 1.19x 0.201 -> 0.204 (depth 1, unchanged) 63 MB 1.28x 1.27x 0.304 -> 0.305 72 MB 1.28x 1.28x 0.338 -> 0.336 135 MB 1.20x 1.37x 0.576 -> 0.545 144 MB 1.20x 1.34x 0.611 -> 0.594 256 MB 1.24x 1.45x 1.045 -> 0.958 512 MB 1.33x 1.41x 2.031 -> 1.830 1 GB 1.32x 1.49x 4.012 -> 3.613 ``` The relayed quantity here is the segment, i.e. a quarter of the message, so the depth cost model only tiles from ~63 MB and only reaches depth 4 past ~256 MB; that is why everything below 135 MB is unchanged. At the large end the measured gain is 1.11x on relay latency against the 1.20x the link model allows, so ~92% of the available headroom. Parallel (4 co-resident jobs) also improves and stays smooth -- 1.23x / 1.13x / 1.33x / 1.45x at 63 MB / 135 MB / 256 MB / 1 GB, against a checked-in baseline that was ~1.0x at the large sizes. Fused (2 groups) is unchanged by construction and re-measured to confirm: 1.30x / 1.20x / 1.20x / 1.25x at 63 MB / 135 MB / 256 MB / 1 GB. Differential Revision: D117629688

meta-codesync · 2026-08-27T14:49:43Z

@srinathb-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D117629710.

Summary: Applies the duplex pipelining to the 4-active flat all-gather. With nGroups == 1 the active ranks and the helpers are disjoint, so the offload scatter (active -> helper) and the helper fan-out (helper -> active) sit on OPPOSITE directions of the same cross link and the two-group form leaves each of them half idle. Tiling the offload so tile t's fan-out shares a group with tile t+1's scatter fills both directions. This schedule is ASYMMETRIC, unlike the 2-active one: a helper fans each source's slice out to the A-1 = 3 other destinations, so a cross link carries one unit UP for every three DOWN. Merging is bounded by the down direction and the up direction stays underused, so the model only allows 1.12x here versus 1.75x at 2 active. With v = align(sc / (7*T + 1)) each group carries 3v on the busiest link and the direct region is split to match -- v for group 0, which has no fan-out to receive yet, and 3v for the rest -- giving (3*T + 1)*v: sc/2 at T = 1 (identical to the two-group schedule) falling towards 3*sc/7. That floor is the hard one, since an active rank must take in A-1 whole buffers across its 7 links. The measured gain is well ABOVE the link model: 1.22x at 1 GB and 1.28x at 512 MB against the model's 1.12x. The extra comes from helper staging, which drops from A*cs (512 MB at a 1 GB message) to two ping-pong tiles per source (~148 MB at depth 4) and so stops spilling out of cache on the fan-out re-reads. That also removes most of the large-size decay this path had: speedup used to fall 1.64x -> 1.55x -> 1.42x -> 1.36x from 135 MB to 1 GB and is now 1.77x -> 1.82x -> 1.81x -> 1.66x. Adds the all-gather shape to RelayPipelineShape; depth selection, the boundary cost model, the depth cap and the nGroups == 1 gate are the shared relayPipelineTiles(). Fused calls keep the original schedule bit-for-bit. Perf -- single-group 4-active all-gather (MI350X, 8 GPUs, bf16, best-of-100, speedup = 4-rank NCCL all-gather with the other 4 ranks idle / relay), same-day before and after: ``` Msg Size before after relay ms 9 MB 1.34x 1.35x 0.172 -> 0.171 (depth 1, unchanged) 13.5 MB 1.41x 1.45x 0.222 -> 0.223 (depth 1, unchanged) 27 MB 1.58x 1.57x 0.377 -> 0.378 31.5 MB 1.64x 1.61x 0.428 -> 0.423 36 MB 1.60x 1.62x 0.480 -> 0.477 63 MB 1.64x 1.70x 0.792 -> 0.771 67.5 MB 1.64x 1.72x 0.854 -> 0.817 72 MB 1.65x 1.72x 0.907 -> 0.868 135 MB 1.64x 1.77x 1.664 -> 1.542 144 MB 1.64x 1.77x 1.770 -> 1.641 256 MB 1.55x 1.82x 3.306 -> 2.808 512 MB 1.42x 1.81x 7.114 -> 5.563 1 GB 1.36x 1.65x 14.709 -> 12.105 ``` Fused (2 groups) is unchanged by construction and re-measured to confirm, matching the checked-in `perf_data/bench_sharded_relay_fused_sweep_results.txt` exactly: 1.28x / 1.31x / 1.32x / 1.31x at 63 MB / 135 MB / 256 MB / 1 GB. Parallel (2 co-resident jobs) improves too, comparing minima over 4 repetitions per depth (see the note below on why single runs are not usable here): ``` Msg Size depth 1 pipelined 135 MB 1.849 1.585 144 MB 1.992 1.667 256 MB 3.650 2.880 512 MB 7.131 6.505 1 GB 15.234 14.606 ``` **Note on parallel-sweep variance.** A single parallel run initially looked like this change destabilized the co-resident case (isolated 0.67x-0.80x points). It does not. That sweep has a run-scoped, bimodal ~20-30% spread on relay times which is present at depth 1 with EQUAL OR LARGER magnitude (max/min over 4 runs at 144 MB: 1.78 at depth 1 versus 1.24 at depth 4), is identical across co-resident jobs to within 3%, and does not appear in the NCCL baseline. It is not DVFS (fclk stays pinned at 1500 MHz, sclk moves ~6%) and is not the torch allocator (expandable_segments does not change it). Pinning NCCL_MIN/MAX_P2P_NCHANNELS makes it deterministic to +/-1% but costs 1.75x at 32 channels and 20x at 8, which localizes it to p2p channel resourcing inside RCCL rather than anything the relay schedule controls. The practical consequence is that parallel comparisons need repetitions and must compare minima; all parallel numbers here do. Differential Revision: D117629690

…lined paths) Summary: The single-group relay pipeline was capped at depth 4 on the belief that depth 8 destabilized co-resident jobs. That was a measurement error. Re-measuring with repetitions shows depth 8 is better nearly everywhere and depth 16 has nothing left to give, so raise the cap to 8. **The measurement error.** The parallel (co-resident) sweep carries a run-scoped, bimodal ~20-30% spread on relay times. Comparing single runs against a single-run baseline attributed that spread to the pipeline. It is not the pipeline: - It is present at depth 1, the unpipelined schedule, with EQUAL OR LARGER magnitude. Max/min over 4 runs of the 4-active all-gather at 144 MB is 1.78 at depth 1, 1.19 at depth 2, 1.24 at depth 4. - All co-resident jobs agree to within 3% in every run, so it is not cross-job contention skew. Per-job output (BENCH_SWEEP_PER_JOB) is what showed this. - It hits the NCCL baseline too: the depth-1 control read 8.817 ms for NCCL at 256 MB against 5.162 and 5.112 ms in its other two runs. Relay times see it more often because the relay drives all 7 links over p2p while the NCCL ring baseline drives 3. - Each point is already a min over 100 iterations, so a slow point means a job was slow for every iteration -- run-scoped state, not jitter. A later finding, recorded here because it changes how much of this to believe: a CO-TENANT workload on the shared node reproduces these symptoms exactly. While another job had the GPUs, a 4-active single-group all-gather at 256 MB read 8.3-10.6 ms for the NCCL baseline against 5.08 ms on an idle node. So co-tenant interference is a confirmed contributor and the p2p-channel explanation below is suggestive rather than settled. The per-rep spread column added later in this stack is the way to tell them apart before trusting a number. What it is not: clocks (fclk stays pinned at 1500 MHz through a sweep, sclk moves ~6%, which would hit NCCL equally) or the host allocator (expandable_segments does not change it). Pinning NCCL_MIN/MAX_P2P_NCHANNELS removes it entirely -- +/-1% at 8 or 32 channels -- but costs 20x at 8 channels and 1.75x at 32, since the default resolves to 64 channels with 8 per peer and forcing it lower puts all 7 peers' traffic on the same channels. That places the variance in p2p channel resourcing inside RCCL, not in anything the relay schedule controls, and it is why the relay (7 links, p2p) sees it while the NCCL ring baseline (3 links, collective kernels) does not. The operative consequence, now recorded in the header: parallel comparisons need repetitions and must compare minima. **Why 8 and not more.** Depth 16 is a wash at best (2-active allreduce at 1 GB: 4.252 ms at depth 8, 4.253 ms at 16) and a loss where the per-stage op count is already high (4-active all-to-all at 512 MB: 1.836 ms -> 1.993 ms). The cost model does the rest: raising the cap changes the chosen depth ONLY at >= 256 MB for the 2-active shapes and >= 512 MB for the 4-active ones, which is exactly where depth 8 measures better. Everything at or below 144 MB keeps depth 4, where depth 8 measured worse (2-active all-gather at 135 MB: 0.708 ms at depth 4 versus 0.725 ms at 8). Perf -- single-group, relay ms, depth cap 4 -> 8 (MI350X, 8 GPUs, bf16, best-of-100): ``` 512 MB 1 GB 1 GB speedup all-gather A=4 5.583 -> 5.511 12.127 -> 10.929 1.65x -> 1.83x all-gather A=2 2.346 -> 2.249 4.584 -> 4.372 4.32x -> 4.52x allreduce A=2 2.315 -> 2.212 4.457 -> 4.250 4.40x -> 4.61x all-to-all A=2 1.226 -> 1.232 2.338 -> 2.249 4.20x -> 4.38x reduce-scat. A=2 1.420 -> 1.433 2.695 -> 2.595 3.74x -> 3.88x all-to-all A=4 1.835 -> 1.837 3.611 -> 3.505 1.49x -> 1.50x reduce-scat. A=4 2.365 -> 2.369 4.703 -> 4.709 1.18x -> 1.17x allreduce A=4 3.166 -> 3.159 6.301 -> 6.280 1.65x -> 1.66x ``` Parallel, comparing minima over 3 repetitions per depth (the only valid comparison there): gains up to 1.44x (2-active allreduce at 256 MB) and 1.12-1.19x on both allreduce variants; the largest loss is 4.3% (4-active all-gather at 512 MB), which is inside a sweep whose depth-4 spreads reach 2.38x. Fused is unaffected -- it always runs depth 1. Differential Revision: D117629691

Summary: Applies the duplex pipelining to the 4-active flat reduce-scatter, the mirror of the pipelined 4-active all-gather. With nGroups == 1 the active ranks and the helpers are disjoint, so the offload scatter (active -> helper) and the reduced return (helper -> active) sit on OPPOSITE directions of the same cross link, and the two-group form leaves each of them half idle. Tiling the offload so tile t's return shares a group with tile t+1's scatter fills both directions. Like all-gather at 4 active this shape is ASYMMETRIC, in the other direction: a source scatters the slice of ALL A-1 of its foreign blocks while the helper returns just one reduced slice per owner, so a cross link carries 3 units UP for every 1 DOWN and merging is bounded by the up side. It therefore reuses kRelayShapeA4Fanout unchanged. With w = align(recvCount / (7*T + 1)) the first T groups carry 3w on the busiest link and the last carries w, giving (3*T + 1)*w: recvCount/2 at T = 1 (identical to the two-group schedule) falling towards 3*recvCount/7. Two things are specific to reduce-scatter: - The helper's reduce is PER TILE, because it gates the return hop: after the group that receives tile k it sums that tile's A-1 contributions per owner, and the next group ships one reduced tile per owner. Its scratch is laid out buffer-major so a tile's A-1 inputs stay contiguous with stride w, which the fused multi-input reduce requires. - The OWNER's reduce stays a single fused pass over the whole block at the end. Reduce-scatter has no gather phase waiting on it, so unlike allreduce there is no need to retile dScratch -- it keeps the flat path's shard-major layout and the closing reduction is unchanged. Perf -- single-group 4-active reduce-scatter (MI350X, 8 GPUs, bf16, best-of-100, speedup = 4-rank NCCL reduce-scatter with the other 4 ranks idle / relay), same-day before and after: ``` Msg Size before after relay ms 27 MB 1.05x 1.03x 0.184 -> 0.188 (depth 1, unchanged) 31.5 MB 0.96x 1.00x 0.227 -> 0.217 (depth 1, unchanged) 36 MB 1.13x 1.01x 0.222 -> 0.246 (depth 1, unchanged) 63 MB 1.09x 1.16x 0.356 -> 0.345 67.5 MB 1.09x 1.16x 0.378 -> 0.367 72 MB 1.20x 1.14x 0.378 -> 0.393 135 MB 1.21x 1.25x 0.652 -> 0.640 144 MB 1.29x 1.24x 0.651 -> 0.694 256 MB 1.28x 1.30x 1.154 -> 1.148 512 MB 1.17x 1.28x 2.369 -> 2.203 1 GB 1.17x 1.30x 4.709 -> 4.249 ``` The relayed quantity is recvCount, a quarter of the message, so the depth model tiles later here than for all-gather; the win is concentrated at 512 MB (1.075x) and 1 GB (1.11x). The 36 MB / 72 MB / 144 MB cells move within this table's noise (they are depth-1 or depth-2 and the schedule there is unchanged or nearly so); repeated runs put 512 MB at 2.189-2.208 ms and 1 GB at 4.183-4.268 ms. One outlier worth recording: one full-range sweep out of ten read 14.353 ms at 512 MB. Six isolated repetitions of that single size read 2.193-2.218 ms and four further full-range sweeps were clean, so it did not reproduce, but it is larger than the ~25% environmental spread documented in sharded_relay_route.h and is noted here in case it resurfaces. Differential Revision: D117629692

Summary: Last of the eight relay variants to get the duplex pipelining, and the largest 4-active win. With nGroups == 1 the active ranks and the helpers are disjoint, so the offload's scatter and its reduced broadcast sit on OPPOSITE directions of the same cross link. What makes allreduce the best 4-active case is that its offload is SYMMETRIC -- the helper takes one chunk per active rank and returns one reduced chunk per active rank -- so the cross link carries the same in each direction and merging is not throttled by a heavy side, unlike all-gather and reduce-scatter where one direction carries 3 units for every 1. TWO dependencies are pipelined here, not one: offload scatter tile k in group k, helper sums the A chunks, reduced tile broadcast back in group k+1. direct reduce-scatter tile k in group k, owner folds its shard's tile k, and the all-gather of that tile goes out in group k+1. With y = align(count / (12*T)) each of the T+1 groups carries 2y on the busiest link, so the cost is 2*(T+1)*y: count/4 for the two-group schedule falling to 5*count/24 at depth 4 and 3*count/16 at depth 8. Two things this path needs that none of the others did: - **Its own depth selector.** Every other shape's depth-1 case reproduces the shipped two-group schedule, so it fits RelayPipelineShape. This one does not: at depth 1 the pipelined form moves count/3 per link against the two-group schedule's count/4, and at depth 2 it moves exactly count/4 with an extra boundary -- worse, then break-even. relayA4AllReducePipelineTiles() therefore offers only depths 4 and 8 and compares them against the two-group cost directly. - **A tile-major dScratch.** The owner has to fold one tile's A-1 peer contributions while the next group is already in flight, and the fused multi-input reduce requires those contiguous with stride tileSz. The flat path's shard-major layout cannot serve that, so the pipelined path relays it out; the flat path is untouched. A stage here is priced at TWICE kRelayPipelineBoundaryBytes, because this is the only schedule that issues two reduce kernels per stage (the owner's shard tile and the helper's offload tile) on top of the group boundary. That is not a fudge, it is what the measured depth curve requires -- it wants the two-group schedule at or below 72 MB, depth 4 from 135 MB to 256 MB, and depth 8 from 512 MB, which pins the per-stage price between 1.33 and 1.875 MiB: ``` relay ms two-group depth 4 depth 8 63 MB 0.441 0.483 0.559 67.5 MB 0.468 0.510 0.597 72 MB 0.488 0.532 0.599 135 MB 0.874 0.843 0.946 144 MB 0.927 0.886 0.985 256 MB 1.601 1.477 1.520 512 MB 3.159 2.825 2.766 1 GB 6.280 5.578 5.284 ``` At the shared 768 KiB the model crosses over near 54 MB instead and loses 7-8% at 63-72 MB; that regression is what this constant removes. Perf -- single-group 4-active allreduce (MI350X, 8 GPUs, bf16, best-of-100, speedup = 4-rank NCCL allreduce with the other 4 ranks idle / relay), same-day before and after: ``` Msg Size before after relay ms 27 MB 1.44x 1.41x 0.227 -> 0.228 (two-group, unchanged) 36 MB 1.46x 1.47x 0.283 -> 0.280 (two-group, unchanged) 63 MB 1.56x 1.56x 0.441 -> 0.440 (two-group, unchanged) 72 MB 1.58x 1.58x 0.488 -> 0.491 (two-group, unchanged) 135 MB 1.60x 1.66x 0.874 -> 0.842 144 MB 1.59x 1.66x 0.927 -> 0.892 256 MB 1.63x 1.76x 1.601 -> 1.479 512 MB 1.65x 1.89x 3.159 -> 2.767 1 GB 1.66x 1.98x 6.280 -> 5.273 ``` No size regresses. Fused (2 groups) is unchanged by construction -- the selector returns 1 for nGroups > 1, so a fused call runs the two-group schedule. Differential Revision: D117629693

Summary: The single-group pipelined reduce-scatter was the one relay variant far off the link roofline: 62% of the measured 51.6 GB/s per-link ceiling at 2 active ranks against 74-83% for the other seven. The cause was not the schedule, it was the owner's reduce, which ran as a single pass only after the last transfer had landed -- 0.608 ms of a 2.642 ms call at 1 GB. Two clean controls establish that. All-gather at the 256 MB label and reduce-scatter at 1 GB have identical busiest-link traffic (114.8 MB), the same shape and the same depth, yet 2.803 ms vs 4.193 ms. And ablating the reduce entirely recovers exactly the gap, while ablating the per-tile HELPER reduce recovers 0% at 2 active ranks and 2.4% at 4. Tiling that reduce and issuing it per stage on the caller's stream would achieve nothing -- stream order is total, so the T+1 smaller reduces serialize in the same places, for T extra launches. So each region's reduce is issued on a cached side stream as that region lands, gated by an event recorded at the group boundary, with one join before return. Only the join creates a side -> caller dependency, so group k+1 never waits on the reduce of stage k. That needs the arrivals of one group to be contiguous, so the shipped block is re-indexed from helper-major to STAGE-major: region 0 is direct chunk 0, and region k is relay stage k-1's H pieces followed by direct chunk k. The unit count is unchanged at (H+1)*T + 1, so the depth model is untouched and depth 1 still reproduces the two-group schedule. Falls back to today's single trailing pass when the message is too small to pay for the event traffic, when the pipeline is deeper than the event pool (a depth-T pipeline runs T+1 groups and each needs its own event, so this keeps a future raise of kRelayMaxPipelineTiles from indexing past the array rather than making it a segfault), when the caller's stream is being captured into a graph (RCCL rejects a group mixing captured and uncaptured streams, and the side stream is uncaptured), or if the stream or events cannot be created. Only the 2-active path changes. The 4-active flat path keeps its trailing reduce: its prize is 10.8% rather than 23%, it would need two further scratch re-layouts, and it has a separate ~0.86 ms structural residual that may reshape that schedule anyway. Differential Revision: D117629696

… -> 1.61x) Summary: The pipelined 4-active reduce-scatter was the one variant far off the link roofline: 56% of the measured 51.6 GB/s per-link ceiling against 73-83% for the other seven. The cause was not bytes and not arithmetic -- it was operation GRANULARITY. Both directions of every link carry the same 3w per group, but they were carrying it as different operation counts. The cross links to the helpers took three w-sized offload sends; the intra links to the other active ranks took a single 3w direct send. RCCL budgets channels per operation, so the single-op link received a third of the channels of its three-op peers and gated the whole group while the cross links finished early and idled. Issuing the direct region as w-sized pieces makes every operation in a group the same size, so all seven links are provisioned alike. This is also why the 2-active shapes never showed the problem: their direct chunk is already exactly u. Splitting is not good for its own sake -- uniformity is. Six pieces instead of three measured 4.304 ms against 3.539, worse than not splitting at all. The all-gather mirror has the same asymmetry but far less to gain, being already the closest variant to the roofline, and the extra operations do not amortize at moderate sizes. It is therefore size-gated at 256 MiB (kRelayUniformDirectOpMinBytes) where it measurably pays, while reduce-scatter takes it unconditionally because it wins everywhere its offload route is active. The other two 4-active shapes need no change: the all-to-all already ships u-sized direct chunks. Differential Revision: D117629697

Summary: Same operation-granularity fix as the previous commit, applied to the last shape that had the asymmetry. The pipelined 4-active allreduce moves oTile = 2*y per link per group on its offload while the direct exchange moves y, and RCCL budgets channels per operation -- so a cross link carrying its bytes as ONE operation got half the channels of an intra link carrying the same bytes as two. Splitting the offload into y-sized pieces makes every operation in a group the same size, so all seven links are provisioned alike. oTile is exactly 2*y, so the split is exact. This is the fourth and last variant to need it. The all-to-all already ships u-sized direct chunks, and both 2-active shapes ship a direct chunk of exactly u. Gated at kRelayUniformDirectOpMinBytes (256 MiB), the same threshold the all-gather mirror uses and for the same reason: the extra operations do not amortize at moderate sizes. Differential Revision: D117629698

…up to 1.89x where it was 1.08x) Summary: The seven single-group relay/direct crossovers were measured before the tiling+pipelining stack landed. Pipelining cut a relay's per-link bytes from count/4 toward count/7 and filled the cross links' idle direction, so every crossover should have moved down -- and none did. The symptom was visible in the checked-in sweep as a ~2x step across a SINGLE doubling at each threshold, which a real crossover cannot produce: you switch where the two curves MEET, so the speedup has to be continuous through it. Measured both curves over the same size axis (temporary route override, 1-27 MB, finer sizes, reps=10) and set each threshold to where they actually cross: allreduce A=2 6 MB -> 2 MB (relay 0.080 vs direct 0.090 at 2.25 MB) allreduce A=4 9 MB -> 1 MB (relay 0.072 vs direct 0.076 at 1.125 MB) reduce-scatter A=2 27 MB -> 3 MB (relay 0.074 vs direct 0.089 at 4.5 MB) all-to-all A=2 27 MB -> 3 MB (relay 0.077 vs direct 0.089 at 4.5 MB) all-gather A=2 9 MB -> 2 MB (relay 0.076 vs direct 0.090 at 2.25 MB) all-gather A=4 8 MB -> 3 MB (relay 0.120 vs direct 0.140 at 4.5 MB) all-to-all A=4 9 MB -> 10 MB (relay 0.096 vs direct 0.092 at 9 MB) The A=4 allreduce moves the furthest because its pure-direct alternative is the only 3-launch schedule in the set -- a direct reduce-scatter, then a reduce, then a direct all-gather -- and it measures 0.89-0.97x of NCCL at every size, so there is very little for the relay to beat. The A=4 all-to-all moves the other way: at 9 MB the relay was 4% behind, so that threshold goes UP. Only the independent (nGroups == 1) thresholds change. The fused thresholds are untouched: pipelining is nGroups == 1 gated, so it cannot have invalidated them. Resulting single-group speedups (relay ms before -> after, NCCL unchanged): allreduce A=2 4.5 MB 0.138 -> 0.092 0.98x -> 1.49x allreduce A=4 4.5 MB 0.110 -> 0.089 0.95x -> 1.17x reduce-scatter A=2 4.5 MB 0.086 -> 0.075 1.02x -> 1.26x 9 MB 0.131 -> 0.086 1.05x -> 1.64x 13.5 MB 0.177 -> 0.102 1.08x -> 1.89x all-to-all A=2 4.5 MB 0.089 -> 0.076 1.08x -> 1.28x 9 MB 0.134 -> 0.088 1.08x -> 1.74x 13.5 MB 0.175 -> 0.102 1.08x -> 1.87x all-gather A=2 4.5 MB 0.132 -> 0.088 1.04x -> 1.55x all-gather A=4 4.5 MB 0.140 -> 0.120 1.01x -> 1.16x all-to-all A=4 9 MB 0.095 -> 0.092 1.05x -> 1.11x Two of those were below 1x (allreduce A=2 and A=4 at 4.5 MB) and are now 1.49x and 1.17x. The reduce-scatter A=4 offload threshold is deliberately left at 48 MB: its relay curve is the one shape pipelining did not move, still losing at 27 MB (0.202 ms relay vs 0.187 ms direct). The eight tests that assert these boundaries all failed on this change, which is what they are for; each is retargeted to the new threshold so it keeps bracketing the real one. Differential Revision: D117629699

…om four to two (0.87x -> 1.07x) Summary: Below its crossover the A>2 allreduce ran a pure-direct reduce-scatter plus all-gather among the active ranks with the helpers idle. That is the most expensive shape in the relay set on the metric that actually decides a small-message call -- launches on the critical path: NCCL, any collective 1 all-gather / all-to-all pure-direct 1 (one group) reduce-scatter, A=2 allreduce pure-direct 2 (group + reduce) A>2 allreduce pure-direct 4 (sendbuff -> recvbuff staging copy, RS group, shard reduce, AG group) and it showed: A=4 allreduce measured 0.87-0.97x of NCCL at every size below the crossover, the worst small-message column of any variant. Replace it, for nGroups == 1, with a plain full exchange: ONE group in which each active rank ships its whole buffer to the other A-1, then ONE fused reduce over all A contributions. Two launches. It moves (A-1)*count per link instead of 2*(A-1)*count/A, but this path only ever runs below the crossover, where the bytes are not what is being paid for -- this is exactly the trade the A==2 path has always made, generalized to A > 2. No new size gate is needed: the route selector already bounds the pure-direct regime, so "offload disabled" is precisely "below the crossover" (1 MB for a single-group A=4 call after the preceding change). Restricted to nGroups == 1 because a fused call has every rank active in one group and a helper in the others, so its links carry several groups' traffic at once and the extra bytes are not free there; the fused route keeps the two-group schedule. In-place and out-of-place take different kernels: in-place seeds the reduce from the destination (multiReduce), out-of-place seeds it from sendbuff (seededMultiReduce), which also removes the staging copy outright rather than just moving it. DISPATCH_SEEDED_MULTI_REDUCE is copied into this TU following the existing per-TU macro convention; the kernel itself was already shared. Single-group A=4 allreduce, relay ms before -> after (NCCL flat): 4 KB 0.054 -> 0.045 (-17%) 0.87x -> 1.07x 9 KB 0.056 -> 0.044 (-21%) 0.93x -> 1.16x 18 KB 0.057 -> 0.044 (-23%) 0.89x -> 1.20x 36 KB 0.057 -> 0.047 (-18%) 0.94x -> 1.15x 72 KB 0.058 -> 0.047 (-19%) 0.92x -> 1.10x 144 KB 0.063 -> 0.050 (-21%) 0.97x -> 1.20x 288 KB 0.068 -> 0.052 (-24%) 0.92x -> 1.20x 576 KB 0.068 -> 0.057 (-16%) 0.97x -> 1.13x All eight of the A=4 allreduce's sub-1x cells are retired, and the ~11 us saved across two removed launches prices a launch at ~5.5 us -- which is also the size of the reduce-scatter deficit (2 launches against NCCL's 1), so the same accounting predicts where the remaining sub-1x cells come from. Differential Revision: D117629700

…e comm group (0.89x -> 1.08x) Summary: The pure-direct all-gather and all-to-all each moved their diagonal -- the piece a rank contributes to its own output -- with a standalone cudaMemcpyAsync issued before the exchange group. That made them two-launch schedules against NCCL's one, which is the whole story at these sizes: the relay time is flat from 4 KB to 576 KB, so it contains no bandwidth term at all and a launch is the only unit of cost. Move the diagonal into the exchange group as a P2P pair whose peer is the issuing rank. RCCL services that as a local copy inside the same kernel, so the diagonal costs no transfer and no second launch -- the loops simply stop skipping j == m. Measured saving is 4-8 us, slightly more than one launch, because it also removes the serialization between the copy and the group that followed it. Single-group, relay ms before -> after (NCCL flat to within 4%): all-to-all A=2 4 KB 0.044 -> 0.038 0.96x -> 1.10x 576 KB 0.051 -> 0.045 1.04x -> 1.20x all-to-all A=4 4 KB 0.044 -> 0.036 1.01x -> 1.25x 576 KB 0.049 -> 0.043 1.00x -> 1.23x all-gather A=2 4 KB 0.045 -> 0.038 0.89x -> 1.08x 576 KB 0.056 -> 0.050 1.01x -> 1.12x all-gather A=4 72 KB 0.051 -> 0.043 0.90x -> 1.05x 576 KB 0.057 -> 0.053 1.08x -> 1.19x Every size in between moves by the same 4-8 us. This retires all 18 remaining sub-1x cells in these two collectives, and not merely to parity: the whole <= 576 KB band is now 1.05-1.27x. Two deliberate restrictions: - nGroups == 1 only. Folding the diagonal in the FUSED case made ShardedRelayMultiGroupAllGatherTest.Correctness_4Groups_A2_FusedRoutingThresholds fail intermittently -- once in four runs, with mismatched output slots. In the fused case all eight ranks issue a self pair alongside a real one in the same group instead of just the active ranks. Whether that is an RCCL self-P2P ordering problem or something else was not chased: every cell this targets is single-group, the fused route has its own tuned reference, and shipping a one-in-four correctness flake to buy 6 us on a path that was already at 1.00x is not a trade worth making. The gate is commented so the next person does not quietly lift it. - All-gather additionally requires the offload to be disabled. With the offload on, the direct region is only d1 of the sc elements, so the diagonal does not correspond to any single operation in the group and must stay a separate copy. Differential Revision: D117629701

…s (1.5-1.7x kernel, crossover ~2 MB) Summary: Reduce-scatter is the only relay collective still below 1x at small sizes, and it cannot be fixed with ncclSend/ncclRecv. Measured: deleting its trailing reduce outright still leaves it at 0.90-1.04x, because one ncclGroup of P2P ops costs ~0.038 ms while NCCL's ENTIRE fused reduce_scatter kernel costs ~0.035 ms. The group itself is the floor, so the only way past it is a single kernel that moves the data AND reduces it. RCCL already ships exactly that kernel (ncclSymRun_ReduceScatter_LL) but it is unreachable here: comm->symmetricSupport requires ncclCuMemEnable(), and ncclIsCuMemSupported() returns 0 unconditionally on AMD (rocmwrap.cc:48), so the symmetric window registration that supplies peer pointers can never be enabled on this platform. That leaves the manual route, and this probe establishes whether it is viable before any production code is written. Four questions, four answers: 1. Cross-process IPC works. One-time setup for 8 ranks: hipIpcGetMemHandle 0.012 ms bootstrapAllGather(8x64B) 1.729 ms hipIpcOpenMemHandle x7 0.654 ms ~1.7 ms total, so a lazily-initialized, once-per-communicator registration amortizes to nothing. Note the memory must come from plain hipMalloc: the relay's ScratchBufferCache uses cudaMallocAsync, and mempool-backed allocations CANNOT be IPC-exported. 2. A kernel CAN store into a peer's IPC-mapped buffer and have the peer read it correctly. The probe verifies a full push/flag/spin/reduce handshake -- __threadfence_system() plus release/acquire atomics on peer memory -- against a position-dependent fill, so a block or offset permutation cannot pass. 3. ncclCommRegister is a NO-OP on this build: it returns success in 0.001 ms with a NULL handle. User-buffer registration is therefore unavailable, which settles the design -- we cannot read peer sendbuffs directly, so each rank must PUSH its contribution into peers' pre-registered staging. Caller buffers are never registered, so there is no per-call registration cost at all. 4. The one-shot kernel is intrinsically cheaper than NCCL's fused one, with a clear crossover (2-rank reduce-scatter, float, kernel time, median of 10): input oneshot nccl speedup 4 KB 0.0051 0.0077 1.50x 9 KB 0.0052 0.0086 1.64x 18 KB 0.0052 0.0089 1.71x 36 KB 0.0053 0.0081 1.52x 72 KB 0.0060 0.0087 1.46x 144 KB 0.0069 0.0105 1.52x 288 KB 0.0090 0.0137 1.52x 576 KB 0.0133 0.0198 1.49x 1152 KB 0.0219 0.0322 1.47x 2304 KB 0.0397 0.0354 0.89x 4608 KB 0.0758 0.0592 0.78x 9216 KB 0.1478 0.1075 0.73x 1.46-1.71x up to ~1 MB, turning over between 1 and 2.25 MB. It loses above that because the push adds two HBM trips (store into peer staging, then read it back to reduce) that NCCL's streaming kernel does not pay -- irrelevant when latency-bound, decisive when bandwidth-bound. That crossover is the shape the shipped size gate will take, though the gate itself must be tuned on the single-group sweep, not on this number. Deliberately measured in the STREAMED regime (20 back-to-back launches / 20), which isolates kernel cost and strips per-call launch overhead. That is the right question here -- is a one-shot kernel intrinsically cheaper, and where does it turn over -- and it is not the same question as what a caller pays. The per-call regime is what the checked-in single-group sweep already measures, so it is not duplicated here; expect the ratio to compress once the fixed per-call cost both paths pay is added back. No production code is touched: this adds one test target and nothing else, so the checked-in sweep snapshots are unaffected. Differential Revision: D117629702

…ge path (0.83x -> 1.18x) Summary: Reduce-scatter was the last relay collective below 1x, and it could not be fixed by trimming launches. Its pure-direct schedule was already minimal for ncclSend/ncclRecv -- one group plus one fused reduce -- and that is still one launch more than NCCL, which fuses transfer and reduction into a single kernel. Deleting the trailing reduce outright measured 0.90-1.04x: the ncclGroup ALONE costs ~0.038 ms while NCCL's entire fused reduce_scatter costs ~0.035 ms. The group is the floor, so the group had to go. Replace it, below 1 MiB of per-active-rank input, with a single kernel that both moves the data and reduces it: each active rank pushes its foreign block into the peer's pre-registered staging, fences, raises a per-block flag, spins on the peer's matching flag, then reduces its own contribution with what the peer staged. One launch, no ncclGroup, no separate reduce. RCCL ships this kernel already (ncclSymRun_ReduceScatter_LL) but it is unreachable on this platform: comm->symmetricSupport requires ncclCuMemEnable(), and ncclIsCuMemSupported() returns 0 unconditionally on AMD (rocmwrap.cc:48), so the symmetric window registration that supplies peer pointers can never be enabled. Hence the manual hipIpc* build-out, whose feasibility the preceding probe commit established. Registration scheme: ONE region per communicator, created lazily on the first eligible call and kept for the communicator's life -- [nRanks slots of 1 MiB][nRanks * 64 flags], 8 MiB per rank on an 8-GPU node. It must be plain hipMalloc; the relay's ScratchBufferCache uses cudaMallocAsync and mempool-backed memory cannot be IPC-exported. Callers' buffers are never registered, which is not just a simplification: ncclCommRegister is a NO-OP on this build (returns success with a null handle in 0.001 ms), so reading a peer's sendbuff directly is not available. Pushing into pre-registered staging sidesteps that entirely and makes the per-call registration cost exactly zero. One-time setup measured ~1.7 ms. Two safety properties that took the most care: - oneShotAcquire() is COLLECTIVE (it bootstrap-all-gathers the IPC handles), so every rank must call it on the same call, including ranks that are helpers here and will launch nothing. It is therefore invoked before any branch on myActiveGroup, and gated only on predicates every rank computes identically (sizes and nGroups). It also AGREES its own success across ranks by all-gathering an ok flag: a region only some ranks had would leave those ranks spinning for a peer that took the ncclSend path, which hangs rather than degrades. - The epoch advances once per host launch and flags are never cleared, so a stale flag always compares as not-yet-arrived. That also means a replayed CUDA graph would reuse an epoch and let the handshake pass before the data landed, so the path is disabled under graph capture. Blocks handshake pairwise -- block b waits only on the peer's block b -- so there is no global barrier and no co-residency requirement: every block writes and flags before it waits. Two tuning results, both measured on the single-group sweep: - Block count: strictly more is better, because the kernel is copy-throughput bound and the handshake is free. At 576 KB, 1 block is 0.267 ms against 0.047 at 32; 64 beats 32 (0.043 vs 0.047) and 128 is no better than 64. Settled at 64. - 16-byte vector access: the push was moving bf16 element-wise, 2 bytes per thread over XGMI. Moving the bulk as 16-byte units took 576 KB from 0.043 to 0.041 and 288 KB from 0.039 to 0.037. Alignment is checked rather than assumed (staging is always aligned, caller buffers are not), with a scalar path for the ragged tail and for misaligned callers. Single-group reduce-scatter A=2, relay ms before -> after: 4 KB 0.040 -> 0.034 0.92x -> 1.07x 9 KB 0.042 -> 0.034 0.86x -> 1.00x 18 KB 0.042 -> 0.035 0.87x -> 1.02x 36 KB 0.043 -> 0.035 0.85x -> 1.02x 72 KB 0.044 -> 0.035 0.83x -> 1.06x 144 KB 0.046 -> 0.036 0.85x -> 1.11x 288 KB 0.046 -> 0.037 0.94x -> 1.13x 576 KB 0.049 -> 0.041 0.93x -> 1.18x 6-10 us off every size, and all eight cells are now at or above 1x. The margin is thinnest at the small end because ~30 us of a 34 us call is per-call overhead outside the kernel that both paths pay; the win is bounded by the kernel's share, which is why 576 KB gains most. The gate is 1 MiB because that is the last power of two entirely inside the region where the one-shot kernel beats NCCL's fused one. The probe measured the kernel crossover between 1 MiB and 2.25 MiB (1.46-1.71x below it, 0.89x at 2.25 MiB, 0.73x at 9 MiB): above the crossover the push costs two HBM trips -- store into peer staging, then read it back to reduce -- that a streaming kernel does not pay. Scope: single-group, 2 active ranks only. A=4 reduce-scatter still has sub-1x cells and is the obvious next step, as is checking whether the one-shot beats the 1.05-1.27x the other three collectives now get from their group path. Differential Revision: D117629704

…83x -> 1.15x) Summary: The one-shot IPC path shipped for 2 active ranks and left A=4 reduce-scatter as the last variant with sub-1x cells. Its problem was identical -- below the offload crossover it runs a pure-direct exchange plus a reduce, two launches against NCCL's one -- so it wants the same fix rather than a different one. Generalize the kernel from 2 ranks to A: each rank pushes the block belonging to active index j into peer j's staging slot indexed by its OWN active index, raises a flag per peer, waits on A-1 flags, and reduces its own contribution together with all A-1 staged ones. A=2 becomes the A==2 case of that loop rather than a separate kernel, and the A==2 and A>2 dispatches now share one helper, since the schedule differs only in A. The staging region, the gate, and the handshake are unchanged. A=4 needs no more memory: slots are sized by the per-active-rank input, which the gate already bounds, and A=4 stages recvCount = input/4 per slot rather than input/2. Single-group reduce-scatter, relay ms before -> after: A=4 4 KB 0.043 -> 0.036 0.87x -> 1.00x 9 KB 0.044 -> 0.036 0.89x -> 1.15x 18 KB 0.044 -> 0.037 1.09x -> 1.14x 36 KB 0.046 -> 0.037 0.91x -> 1.11x 72 KB 0.046 -> 0.039 0.87x -> 1.02x 144 KB 0.046 -> 0.041 0.85x -> 1.04x 288 KB 0.048 -> 0.042 0.83x -> 0.98x 576 KB 0.048 -> 0.043 0.92x -> 1.01x 5-9 us off every size. A=2 pays 0-2 us for the generalization (0.033-0.042 against 0.033-0.041), which is at the noise floor and worth it: A=4 gains 5-9 us, and one kernel is better than two that can drift apart. 288 KB reads 0.98x here. Its relay time improved by 6 us; the residual is the NCCL baseline, which drifts run to run at these sizes -- the same cell measured 1.04x and 0.98x on consecutive runs with our time pinned at 0.042. That is why the tables above are read on the relay column. Differential Revision: D117629705

… (fixes a bootstrap deadlock) Summary: The one-shot region cache was keyed on the ncclComm_t POINTER, which deadlocks. There is no relay-visible hook on communicator teardown, so a destroyed communicator leaves its cache entry behind. The allocator can then hand the same address to the NEXT communicator -- and because every rank is a separate process with its own heap, it does so on SOME ranks and not others. A rank that finds the stale entry sees tried/valid already set and skips creation; a rank that does not enters bootstrapAllGather. The participation mismatch hangs the bootstrap. Observed exactly that way: the allreduce suite sat 20 minutes (against a normal 5) with the stack in bootstrapAllGather rcclx::relay::oneShotAcquire tryOneShotAllReduce shardedRelayAllReduceFlat ...ShardedRelayAllReduceSingleGroupA4Test::runSingleGroupA4Case comm->commHash is identical across the ranks of one communicator and different for a different one, so it is consistent exactly where the pointer is not. Key on it instead, and evict any entry whose recorded comm pointer matches the current communicator under a different hash -- that is the recycled-address case, and its IPC mappings refer to memory that went away with the old comm. This is a latent bug in the two preceding one-shot commits, not a new one. It did not fail there because whether an address is recycled depends on allocation history: the reduce-scatter suite happened not to hit a divergent reuse, and the allreduce suite did, once its own small-message tests started acquiring regions. Nothing about the reduce-scatter path made it safe. Differential Revision: D117629706

… -> 1.57x) Summary: After the one-shot path took reduce-scatter above 1x, allreduce was the only collective left still paying two launches at small sizes. Both its pure-direct schedules -- the A==2 full exchange and the A>2 one -- are a group followed by a reduce, which is one more launch than NCCL and exactly the shape reduce-scatter had before the one-shot replaced it. Allreduce turns out to be the srcStride == 0 case of the same kernel: every peer is owed the WHOLE buffer rather than a per-peer block, and every rank keeps the whole result rather than a shard. So generalize the push-reduce kernel with a per-peer source stride and an owner offset reduce-scatter srcStride = rc, ownOffset = mySlot*rc allreduce srcStride = 0, ownOffset = 0 and both collectives share one kernel. Bytes moved are unchanged -- (A-1)*count pushed per rank, exactly what the group form already sent -- so this trades a launch for nothing. The staging region and the gate are unchanged too; allreduce stages a whole contribution per slot instead of a shard, which is what the gate already bounds. Single-group allreduce, relay ms before -> after (NCCL baseline flat to 2 us): A=2 4 KB 0.042 -> 0.033 1.03x -> 1.34x 9 KB 0.042 -> 0.034 1.08x -> 1.38x 18 KB 0.043 -> 0.034 1.10x -> 1.36x 36 KB 0.043 -> 0.036 1.23x -> 1.52x 72 KB 0.045 -> 0.035 1.22x -> 1.57x 144 KB 0.047 -> 0.036 1.18x -> 1.50x 288 KB 0.049 -> 0.040 1.12x -> 1.36x 576 KB 0.054 -> 0.046 1.14x -> 1.27x A=4 4 KB 0.043 -> 0.038 1.07x -> 1.24x 9 KB 0.046 -> 0.039 1.09x -> 1.30x 18 KB 0.044 -> 0.039 1.14x -> 1.33x 36 KB 0.047 -> 0.041 1.19x -> 1.33x 72 KB 0.049 -> 0.040 1.09x -> 1.31x 144 KB 0.049 -> 0.042 1.22x -> 1.44x 288 KB 0.052 -> 0.046 1.21x -> 1.37x 576 KB 0.056 -> 0.056 1.16x -> 1.19x 5-11 us off nearly every size. This is the largest one-shot gain of the three, because allreduce had the most to give up: two launches AND the highest baseline cost of the four collectives. Deliberately NOT extended to all-gather or all-to-all. After folding their diagonal into the comm group they are already ONE launch, so a one-shot there saves nothing and would add a staging round trip -- push into the peer's staging, then read it back into recvbuff -- where today the peer's data lands in recvbuff directly. Writing straight into a peer's recvbuff is not available, because ncclCommRegister is a no-op on this build, so the extra trip is unavoidable and the trade is negative. They are already 1.07-1.27x. Differential Revision: D117629707

…destroy Summary: The one-shot region's lifetime was never tied to anything, and both ways of managing it without a teardown hook fail: - Never freeing leaks 8 MiB plus seven IPC mappings per communicator. A suite that builds and destroys a comm per case reached ~140 mappings and comm setup/teardown degraded past the 900s re_timeout -- the allreduce suite went from 333s to over 900s while the tests themselves accounted for only 64s of it. Attributed by ablation: one-shot disabled, nothing else changed, back to 333s. - Bounding with an LRU is worse. Evicting a LIVE comm's region makes the next call recreate it, resetting the epoch to zero; a peer still waiting on the old epoch waits forever. That showed up as an intermittent hang, and the open/close thrash put the cost straight back. Tie the region to the communicator instead, which is what it was always scoped to: one region per comm, created lazily on that comm's first eligible call, and freed from RCCL's commFree(). commFree is the right hook -- it is the per-comm resource teardown reached from every destroy path, and it documents that it must not sync among ranks, which suits a release that only closes this rank's peer mappings and frees its own allocation. It sits beside the existing per-comm cleanups there (tempBuff, algoFactory, the low-precision buffer pool). Regions are INDEPENDENT: each owns its staging, its flags and its epoch. Two properties follow, and both were previously broken: - The create decision is derived from THIS comm's state alone, which starts uniformly absent on every rank of it, so all of a comm's ranks always agree on whether to enter createRegion's collective bootstrap. A process-global region decided this from whichever comm arrived first, which let some ranks decline while others entered the bootstrap and waited for them -- a hang, reachable with two overlapping communicators. - Flags are per region, so concurrent collectives on DIFFERENT communicators can no longer alias each other's handshake. Lifetime now matches the comm exactly, which removes both earlier failure modes by construction: nothing accumulates (release frees it) and nothing is evicted while live (so an epoch never rewinds under a peer still waiting on it). The map is keyed on the comm pointer, which the allocator recycles, so each region also records comm->commHash. A pointer whose hash no longer matches is a region whose release was missed; it is torn down and rebuilt rather than handed back as mappings into memory that went away with the old comm. Every rank of the new comm sees the same mismatch, because commHash is agreed across it. Each region also carries an identity tag: a sentinel word derived from the agreed commHash and the owner's rank, stamped at the end of the region and verified after opening each peer's handle. It answers the one question an IPC mapping cannot answer on its own -- whether the memory behind a handle is the peer's CURRENT region or a previous incarnation at an address the allocator recycled. Any rank can compute any other rank's expected value, so the check costs no extra communication, and a mismatch is voted on collectively so every rank declines the path together rather than some ranks reading stale mappings. Release is this rank's own teardown only: an hipIpcCloseMemHandle per peer plus one hipFree. No cross-rank handshake is needed because NCCL requires comms to be used collectively -- every rank has finished its collectives on a comm before any rank destroys it, so no peer can still be reading the region. The sentinel is the backstop if a recycled address ever turns up anyway. Still assumed, and now the only remaining precondition: relay collectives on a GIVEN communicator are serialized. Two in flight on one comm share that region's flag array, and the later one's epoch would satisfy the earlier one's wait. Adding epoch-indexed flag banks would close it if that ever becomes reachable. Differential Revision: D117629708

Summary: The sharded-relay collectives had zero graph-capture test coverage, so the four known graph hazards in the relay were all unverifiable: G1 ScratchBufferCache allocates with hipMallocAsync/hipFreeAsync keyed on (device, stream, key), which under capture becomes a graph allocation node whose address is only valid while that graph executes. G2 The one-shot IPC handshake advances a host-side epoch that reaches the kernel as a by-value argument, so a replay would spin on a baked epoch that every flag already satisfies. G3 One-shot region creation does a synchronous hipMemset plus two bootstrapAllGather calls, none of which are capturable. G4 The reduce-scatter overlap side stream is forked from the caller stream, whose captured form has no relationship to the uncaptured stream. G2, G3 and G4 currently sit behind explicit capture bail-outs and G1 is latent, so every case here passes as-is. That is the point: this is the regression baseline the four fixes get measured against, and each fix brings its own failing-before/passing-after case with it rather than landing one red test now. Covers all four collectives on a single group with A=4 (ranks 0-3 active, 4-7 helpers): capture, then replay with input that varies per replay, then verify each replay. Two details are load-bearing and were both found the hard way while getting this to fail for the right reasons: - Odd ranks sleep on the host before the launch. Without that skew a stale baked epoch or stale staging is only a race, because the kernel still pushes real data on every replay. With it the unskewed rank pushes, flags, falls through the handshake and reduces data the skewed rank has not written yet, so the mismatch is deterministic. - Cases warm up at a size large enough to fill the scratch cache for the route the capture will take. Below the 1 MiB one-shot band an uncaptured call goes one-shot while the captured call falls back to PureDirect, so warming up at the size under test leaves that cache cold and the capture allocates inside itself. Left unprimed the one-shot-band case hung for the whole harness timeout, and it was hanging on G1 rather than on the one-shot behaviour it is named for. hipStreamSynchronize is replaced by a bounded poll: a graph-compat regression shows up as a kernel that never retires, and a suite that hangs for the harness timeout reports nothing useful and takes every later case down with it. Reviewed By: JinghanHuang Differential Revision: D117629710

srinathb added 6 commits August 26, 2026 20:50

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 27, 2026

srinathb-meta force-pushed the export-D117629710 branch from 3919a45 to d6cb03b Compare August 27, 2026 15:07

srinathb and others added 17 commits August 27, 2026 20:59

meta-codesync Bot added the meta-exported label Aug 28, 2026

srinathb-meta force-pushed the export-D117629710 branch from d6cb03b to ad8044c Compare August 28, 2026 08:59

Navigation Menu

Uh oh!

Add HIP-graph capture coverage for the sharded-relay collectives - #3850

Add HIP-graph capture coverage for the sharded-relay collectives#3850
srinathb-meta wants to merge 23 commits into
meta-pytorch:mainfrom
srinathb-meta:export-D117629710

srinathb-meta commented Aug 27, 2026 •

edited by meta-codesync Bot

Loading

Uh oh!

meta-codesync Bot commented Aug 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

srinathb-meta commented Aug 27, 2026 • edited by meta-codesync Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

meta-codesync Bot commented Aug 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

srinathb-meta commented Aug 27, 2026 •

edited by meta-codesync Bot

Loading