[PERF][distributed]: optimize deterministic ROCm collectives with HIP IPC

[maxiaosong1124]

Summary

Stacked on #356, this PR replaces the ROCm transport-only hot path with a single-node HIP IPC transport while preserving the same fixed balanced rank tree.

• Add direct HIP IPC AllReduce, AllGather, and ReduceScatter kernels for world sizes 2, 4, and 8, with the existing path retained for world size 1.
• Reduce only the destination shard for ReduceScatter, and use ReduceScatter + RCCL AllGather for large AllReduce payloads.
• Add reduce_scatter_many so FFN sequence-parallel backward can reduce the independent gate/up lanes under one ready/done generation without changing either lane's expression tree.
• Keep RCCL as the fallback for multi-node execution, builds without the IPC symbols, and measured message-size ranges where RCCL is faster.
• Preserve the existing CUDA implementation and ROCm Attention CP integration.

Design

Every rank owns one hipMalloc staging allocation. Handles are exchanged once and imported with HIP IPC. A system-scope release/acquire generation publishes the local input and waits for every peer; a second done generation prevents staging reuse while a peer can still read it. Close uses a third peer acknowledgement before releasing the allocation.

Reduction arithmetic remains:

((rank0 + rank1) + (rank2 + rank3)) + ...

FP16/BF16 use two-element vector instructions only when the output is aligned. Scalar tails and offset views use the scalar kernel. Empty AllGather inputs complete without launching a zero-sized grid.

The MI300X routing policy is:

• AllReduce <= 768 KiB: direct IPC fixed tree.
• AllReduce 768 KiB to 2.125 MiB: rank-major RCCL transport + local fixed tree.
• AllReduce >= 2.125 MiB: IPC ReduceScatter + RCCL shard AllGather.
• AllGather <= 256 KiB: direct IPC peer copy; larger messages use RCCL.
• ReduceScatter: direct IPC reduction of the local destination shard.

Determinism and correctness

• RCCL never performs strict floating-point reduction.
• FP32, FP16, and BF16 follow the same fixed rank tree as [FEAT][distributed]: add deterministic ROCm/RCCL transport collectives #356 and the CUDA deterministic collective.
• Repeated runs are bitwise identical on every rank for every benchmark row.
• Cross-TP tests cover TP=1/2/4/8, in-place and caller-provided outputs, FP16/BF16 offset outputs, empty AllGather, and two independent ReduceScatter lanes.

Validation

• ROCm extension build passed for gfx942 and gfx950.
• 41 passed in the CPU/reference, ROCm transport, and build-isolation suite.
• 3 passed in the 8-GPU MI300X AllReduce/AllGather/ReduceScatter cross-TP suite.
• Black, isort, flake8, pre-commit, and git diff --check passed.

Benchmark

Environment and method:

• AMD Instinct MI300X, 8 GPUs, ROCm/HIP 7.14.60850, BF16.
• 10 warmup iterations, 50 measured iterations, 5 samples.
• Each value is the median of the slowest-rank latency.
• Both [FEAT][distributed]: add deterministic ROCm/RCCL transport collectives #356 and this PR report all_rank_repeat_bitwise=true.

PYTHONPATH=. torchrun --standalone --nproc-per-node=8 \
  benchmarks/benchmark_rocm_collectives.py \
  --dtype bf16 \
  --operations all_reduce all_gather reduce_scatter \
  --size-bytes 4096 65536 1048576 16777216 \
  --warmup 10 \
  --iterations 50 \
  --samples 5

Lower latency is better. The two comparison columns express the latency change of this PR; “faster” means lower latency.

Operation	Size	Native RCCL	#356 deterministic	This PR	vs #356	vs native RCCL
AllReduce	4 KiB	28.31 us	71.63 us	21.43 us	70.08% faster / 3.34x	24.31% faster / 1.32x
AllGather	4 KiB	33.26 us	52.08 us	21.60 us	58.53% faster / 2.41x	35.06% faster / 1.54x
ReduceScatter	4 KiB	34.83 us	79.14 us	21.89 us	72.34% faster / 3.62x	37.16% faster / 1.59x
AllReduce	64 KiB	29.24 us	159.08 us	21.33 us	86.59% faster / 7.46x	27.07% faster / 1.37x
AllGather	64 KiB	40.62 us	53.59 us	27.91 us	47.91% faster / 1.92x	31.28% faster / 1.46x
ReduceScatter	64 KiB	36.09 us	180.11 us	21.95 us	87.81% faster / 8.21x	39.17% faster / 1.64x
AllReduce	1 MiB	58.60 us	79.77 us	91.04 us	14.13% slower / 0.88x	55.35% slower / 0.64x
AllGather	1 MiB	69.71 us	70.26 us	69.88 us	0.55% faster / 1.01x	0.23% slower / 1.00x
ReduceScatter	1 MiB	57.95 us	84.29 us	46.91 us	44.34% faster / 1.80x	19.04% faster / 1.24x
AllReduce	16 MiB	152.32 us	442.37 us	230.17 us	47.97% faster / 1.92x	51.11% slower / 0.66x
AllGather	16 MiB	400.58 us	400.75 us	400.97 us	0.05% slower / 1.00x	0.10% slower / 1.00x
ReduceScatter	16 MiB	83.29 us	408.87 us	148.35 us	63.72% faster / 2.76x	78.11% slower / 0.56x

Against #356, 10 of 12 rows improve, with 44%–88% lower latency on the improved reduction rows. Against native RCCL, every 4 KiB and 64 KiB row is 24%–39% faster, and 1 MiB ReduceScatter is 19% faster. Large strict deterministic reductions remain slower than native RCCL because native RCCL does not guarantee the fixed floating-point reduction order required by this implementation. The only material regression against #356 is the 1 MiB AllReduce crossover (+14.13% latency); the 1 MiB and 16 MiB AllGather differences are within noise.

Compute/communication fusion

This PR does not claim asynchronous overlap or compute/communication fusion. The API remains stream ordered and reports supports_async_overlap = False and supports_compute_communication_fusion = False.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 84e8d49f-4f48-456b-9ae2-d97663c723d6

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}