[FEAT][distributed]: add deterministic ROCm/RCCL transport collectives

[Flink-ddd]

Summary

Add a standalone deterministic communication implementation for ROCm, separated from the ROCm GEMM work.

• Add RCCL transport-only AllGather, AllReduce, and ReduceScatter. Preserve deterministic floating-point order with a fixed balanced rank tree.
• Support world sizes 1, 2, 4, and 8. Integrate the ROCm collective into FFN TP/SP communication.
• Add an rccl_ag_rs backend for attention CP communication. Keep the existing CUDA IPC implementation unchanged. Exclude CUDA IPC sources and libcuda when building on ROCm.
• Add correctness tests, build-isolation tests, documentation, and a multi-GPU ROCm benchmark.

Design

RCCL is used only to transport rank-ordered input tensors through all_gather_into_tensor. Floating-point reduction is performed locally using:

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

This avoids depending on RCCL's implementation-defined floating-point reduction order and matches the rank tree used by the CUDA deterministic collectives.

For performance:

• AllGather writes directly into the caller-provided output tensor.
• AllReduce and ReduceScatter reuse a lazily grown byte workspace.
• Balanced-tree reduction is performed in-place inside that workspace.
• Python object signature validation runs only once per new (operation, shape, dtype) signature.

Compute/communication fusion

This PR does not claim compute/communication fusion or asynchronous overlap.

The strict path completes rank-ordered RCCL transport before evaluating the fixed reduction tree. Native RCCL reduction, arbitrary arrival-order merging, and ROCm fused collectives are intentionally not used because they do not provide the required fixed floating-point parenthesization.

A future HIP/XGMI implementation can replace the transport behind the same factory after preserving the deterministic contract.

Validation

• ROCm gfx942 extension build completed successfully.
• 64 targeted tests passed and 27 environment-dependent tests were skipped.
• Black, isort, flake8, and all pre-commit checks passed.
• git diff --check passed.

Multi-GPU MI300 performance numbers are not included because the development environment did not expose AMD GPUs. The PR includes benchmarks/benchmark_rocm_collectives.py for running the measurement.

Benchmark

HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
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 20 \
  --iterations 100 \
  --samples 10 \
  --output benchmarks/results/rocm_collectives_mi300x_tp8.json

The report contains:

• slowest-rank median/min/max latency;
• latency ratio against native RCCL;
• input bandwidth;
• retained workspace and peak allocated memory;
• repeated-run bitwise stability;
• maximum absolute difference from native RCCL.
  Native RCCL is a performance reference, not the strict arithmetic oracle.

[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: 62bc92d2-a60d-4a35-862e-2e5d26a53070

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.}