RoCE transport above two ranks: roce.zig:148 refuses every world but 2, so four GB10s fall back to NCCL (rebased 2..16-rank branch, gate question) #1209

On main (c9fab624) the RoCE transport accepts only two ranks. At any other world, Roce.init refuses before it sets up the wire:

• zig/src/fabric/tp/roce.zig:148: if (b.world != 2) return error.Unsupported;

With the default TF_TP_ROCE_FALLBACK=nccl, every rank of a three- or four-rank session drops to NCCL (rank 0 logs tp: RoCE unavailable (Unsupported); every rank uses NCCL). With TF_TP_ROCE_FALLBACK=error, the session refuses to start.

Why it matters on hardware like ours

We run four GB10s on a switched ConnectX-7 fabric. At two ranks they use RoCE for the small decode exchanges. At four ranks they can't. #1206 (DeepSeek-V4.1 at world 4) says "RoCE keeps its two-rank limit, so four ranks use NCCL", #1202 calls the limit out of scope, and #981 asks what TP=3/4 would need. This issue is about that limit only.

A rebased branch

neko-legends/TensorFold:pr/tp-roce-multi-rank (f0dc5699, three commits on c9fab624, 22 files, +1158/−450; compare). There's no PR yet. We'd like your answer to the question at the end first.

1. fabric/tp: RoCE transport for 2 to 16 ranks

• roce_guard: Layout.of(world, slot) gives every rank the same region layout, with one flag and one receive slot per source rank. peerAt orders the posts starting at the next rank, so the senders don't all hit one receiver first. stripes() takes the source and destination rank.
• roce_verbs: one RC queue pair per HCA and peer. Record carries qpn[hca][peer], and connect() takes every rank's record and pairs HCAs by subnet at both ends. The in-flight byte cap is now checked before each peer's stripe, not once before the loop. Checked once, a world-4 post could have (world − 1) × the payload in flight: 3 MiB against the 2 MiB cap. Checked per peer, a post adds at most one payload past the cap. The record's magic is bumped, because its layout and the registered region size both changed.
• roce: init accepts 2..max_world (16), every rank agrees the setup, and the host gather takes a lens slice of world entries.
• mailbox.cu (tp_roce): waits for every peer's flag per HCA, copies each sender's shard to that rank's slot, and sums in rank order at any world. The device helpers move to roce_device.cuh so mailbox.cu stays under the file-size rule.
• roce_shm: one mapping for every rank's region, so the shared-memory wire runs any world on one host without RDMA. That's what the host tests use.
• exchange stays a two-rank route. Above two ranks it goes to NCCL (hybrid.zig).
• Tests: three- and four-rank worlds over the shm wire (2,000 gathers of every size, variable lengths with empty shards, one stopped rank timing out the other three, abort), the region layout and post order, four-rank HCA pairing by subnet and the refusal of non-one-to-one pairings (roce_verbs_test.zig), the in-process collectives at two, three and four ranks, and a world-4 post checks the in-flight cap before every peer's stripe, so 3 MiB is never posted at once (roce_post_test.zig). Without the per-peer check, that last test fails.

2. fabric/tp: the tf-tp-test harness takes any world: gpu_checks.Ctx carries the world. Above two ranks, NCCL's float sums are compared within rounding, because its ring order isn't rank order. The route's own sums stay bit-exact. gpu_bench uses nccl-tests' bus-bandwidth factors for any world (unchanged at two). roce-host and failfast transport run at the session's world.

3. One-line comments for tools/zig/lean_check.py.

Two ranks keep the same collective semantics. The exchange route and the bits of a two-rank sum don't change. The wire record and the registered region size do change, so both ranks need the same build.

The rebase, and the one conflict

The branch was first cut on d1018b31 (1.1.1). Commit 1 applies to c9fab624 without conflicts. None of the transport files changed on main in between. The only conflict was in zig/src/fabric/tp/gpu_main.zig. Since then, main has replaced the runner's args[1] string matching with gpu_cli.parse, which validates the command into a Command union before any resource opens (094a0906). It has also reworded the hca failure messages for the verbs-loader errors (9ca00ae9). Our side changed suite and failfast to pass cfg.world to checks.Ctx.init, and also added a bench subcommand. The resolution:

• main's gpu_cli parsing, usage line and switch (command) dispatch are kept unchanged, along with its hca messages.
• The world goes into Ctx.init in suite and failfast, and suite warns when a requested non-NCCL backend fell back to NCCL.
• Our bench subcommand is dropped. Adding it would have meant extending gpu_cli's Command, usage and invalid-argument table, and it isn't needed for this change. The two-host shell scripts are left as they are on main too.

Gates (x86_64 Linux host, Zig 0.17.0, no GPU)

The skip is the same case on both trees: the verbs test that needs TF_TP_ROCE_TEST_HCA. No gate fails on either tree. lean_check.py exits 1 on both because of the 438 problems already on main.

On four GB10s

We ran the same transport code at world 4 on four GB10s over the switched ConnectX-7 fabric. roce.zig, roce_guard.zig, roce_shm.zig, hybrid.zig, mailbox.cu and roce_device.cuh are byte-identical to this branch, and roce_verbs.zig differs only in comments. This was a DeepSeek-V4.1 world-4 boot on top of #1202's split. All four ranks came up over RoCE (tp: rank N of 4 up (roce, …) on every rank), with zero panics. A rank only logs that line after every rank passes the load-time probe in Session.openRoce. The probe checks a 4 KiB all-gather over RoCE against NCCL's, and a bf16 sum against the rank-order reference. If any rank fails it, every rank falls back to NCCL.

Not tested

• No injected failures on hardware at world 4: no failfast kill, fatal or transport run. Only the host tests' stopped-rank and abort cases cover that.
• Speed not compared against NCCL at world 4, so we make no claim that RoCE is faster than NCCL at four ranks.
• The boot didn't reach serving, for a separate reason: DeepSeek-V4.1 plans no prefill segment at world 4 (DeepSeek-V4.1 at four ranks: no prefill segment plans (BMQ 32 against 16 heads a rank; x3gm's block against the decode scratch) #1205). So no model tokens went through these collectives at world 4. The probe above and the host tests are the evidence.
• No world 3, and no world above 4, on hardware. The tp_roce kernel isn't compiled by the gates above (this host has no nvcc). It was built and run in the hardware boot above.

The gate question

Is a transport change of this size (2 to 16 ranks, one QP per HCA and peer, a new region layout, a two-rank wire-format change) acceptable as one PR? Or would you prefer something narrower? For example, world 4 only, or the layout and guard first with the verbs and kernel changes after. If you're fine with the shape, we'll open the PR against this issue. We'll include the four-rank tf-tp-test suite and the failure runs if you want them first.