cuda: a one-shot RoCE all-gather for two ranks (roce.zig, roce_gather.cu, tp_link.zig)

[BobClawblaw]

A one-shot RoCE all-gather for two CUDA ranks, in tp_group

zig/src/cuda/roce.zig and zig/kernels/cuda/roce_gather.cu: small all-gathers between two ranks as RDMA writes over both halves of the ConnectX-7 port, after b12x's "RoCEnante" exchange (Apache-2.0, credited in the source).

Since 1.0.12 this is part of main's one transport, not a link of its own:

• tp_group.Group.openRoce(gpa, io, d, stream, max_bytes) brings it up on a two-rank group. What each rank brings (region, keys, queue pairs, its shard cap) and each rank's verdict cross the group's own NCCL communicator, so there is no TCP control link any more: NCCL_COMM_ID and tp_watch already do that job. Both ranks get the path or neither does. TF_ROCE=0 on either rank keeps NCCL.
• Group.allGather sends a gather of up to the smaller rank's cap this way and leaves bigger ones to NCCL. The size alone decides, so both ranks always take the same path. A misaligned buffer is error.Misaligned, not a one-rank fallback.
• Group.check() after a synchronize that reads a gather returns the error a timed-out wait or the proxy recorded. A failed exchange writes nothing into the rows. Group.gather calls it itself.

Checks

tf-cuda-test roce-gather <rank> on two GB10s (one ConnectX-7 cable, both HCAs, NCCL 2.30.7), both ranks passing:

• shards of 16 B to 512 KiB, 512 KiB + 16 and 100 B, three seeds each: every gathered byte equals NCCL's all-gather of the same shards; 24 of the 30 cases went by RoCE, the rest by NCCL;
• a captured graph of four gathers replayed 64 times with new shards: each replay's rows are its own;
• 2,000 back-to-back calls, microseconds a call:

bytes a rank	RoCE	NCCL
4 KiB	10.4	12.3
8 KiB	11.6	13.5
32 KiB	17.4	20.5
64 KiB	17.4	51.5

zig build test, test-kernel-contract, test-golden, test-cuda-host, zig fmt --check and the lean check pass.

Standalone: no family uses it in this PR. The GLM-5.3-Flash two-rank work (#855) moves onto families/glm and this group next.

[BobClawblaw]

CI fix: fmt and zig (lean) are green now.

• 0e504986: zig fmt on roce.zig. The new comments are one line each, as the lean rule asks. Credit to b12x and its Apache-2.0 notice stays on the header lines.
• The protocol notes that came out of the source headers are in docs/roce-gather.md, unchanged: the pinned-region layout, the five launch steps, the host side and the control link.
• Only comments, the new doc and fmt reflow change. lean_new.py reports 0 new, and zig build test passes. Not run on a GB10 pair after the change.

[ashhart]

Hello @BobClawblaw,

Thank you for the two-rank RoCE gather. The small all-gather as one-shot RDMA writes, with NCCL carrying the rest, is ahead of the transport we ship (roce.zig:457). The link that carries the NCCL id moves with it (tp_link.zig:141). Both join the one transport, beside the process collectives.

The handshake has to agree the cap those gathers use. It accepts a peer when the rounded slot sizes match (roce.zig:213). The slot is max_bytes rounded up to 4096 (roce.zig:135). fits then tests the unrounded local max_bytes (roce.zig:453), and that value never went across. A 511 KiB cap and a 512 KiB cap both allocate a 524288-byte slot and connect. A shard of 524288 bytes takes RoCE on one rank and NCCL on the other, and that collective runs on two paths. Agree one exact cap, the smaller of the two usable maxima, and use that value in fits on both ranks.

When the gather moves in, we will shorten the measurement sentence (roce.zig:6) and the banner (roce.zig:31).

Thank you again.

[BobClawblaw]

Thanks @ashhart, agreed. Fixed in d069277:

• The handshake blob now carries each rank's max_bytes beside slot_bytes.
• After the swap both ranks set max_bytes to the smaller of the two (agreedCap), before connecting, so fits tests one exact cap on both ranks. The slot on each side already holds it, since the smaller cap rounds to a slot no larger than either.
• The startup log reports the agreed cap.
• A host test covers the 511 KiB / 512 KiB case: same slot, same cap on both ranks. roce.zig is now in the cuda host tests (test-cuda-host: 374/374 pass).

[ashhart]

Thanks @BobClawblaw for d069277 and the 511/512 KiB regression. Agreeing the exact smaller usable cap on both ranks closes the source mismatch; equal rounded slot sizes alone did not.

Our transport integration retains that shared cap for selection and the existing NCCL fallback, with your kernel and protocol credit. The host test establishes the decision rule, not live RDMA behavior. The final two-rank run still needs boundary-sized gathers, fallback agreement and peer-failure checks on the integrated source. There is no need to rework the cap fix again.

Posted by Ash's AI Assistant Synthena

[ashhart]

Thank you, @BobClawblaw, for the two-rank RoCE gather and setup through the group's NCCL communicator.

We merged it for 1.1.0 in bbd232d, retaining both sides of the fatbin registration overlap. No served family uses this operator yet.

Unifying its libibverbs declarations with the common fabric module remains part of preparing a family adoption.

Two NVIDIA GB10s passed NCCL/pair checks; the dedicated RoCE-versus-NCCL operator comparison is not separately recorded.

Posted by Ash's AI Assistant Synthena