Skip to content

cuda: a reduce-scatter in rank order over NCCL send and receive - #642

Merged
ashhart merged 1 commit into
ashhart:mainfrom
Bizuayeu:pr3/nccl-rank-scatter
Oct 16, 2026
Merged

ashhart merged 1 commit into
ashhart:mainfrom
Bizuayeu:pr3/nccl-rank-scatter

Conversation

@Bizuayeu

@Bizuayeu Bizuayeu commented Oct 10, 2026 •

Copy link
Copy Markdown
Contributor

Step 2 of #593, as claimed there: NCCL send/recv and an exact reduce-scatter in rank order. It doesn't depend on the weight format: it sums the fp32 partial rows a layer's attention and MLP (or experts) leave on every rank.

  • zig/src/cuda/nccl.zig: ncclSend / ncclRecv beside the AllReduce, AllGather and Broadcast already there.
  • zig/src/cuda/rank_scatter.zig: share cuts [rows][width] by whole rows; the first rows % world ranks take one more, as canon.partition cuts equal ranks. RankScatter.run puts the point-to-point sends and receives in one group: each peer gets its rows of this rank's partial, and this rank receives its rows from each peer. There's no transfer to itself, and a rank without rows only sends. rank_sum.cu then adds the terms rank 0 first, rounding each add. Every rank so holds the bits an all-gather followed by a rank-order sum gives, while receiving only its own rows. NCCL's ncclReduceScatter isn't used because it adds in its ring's order, not rank order. It's also unrelated to cluster/exchange.zig's reduceScatter, which cuts by columns and sums in a fixed binary tree.
  • tf-cuda-test rank-scatter <rank> <n> [id]: every rank's rows are checked against an all-gather plus a host rank-order sum. Rows are 1, 4, 64, 2,047 and 2,048 by 4,096, with three seeds each. The ranks join by NCCL_COMM_ID or by rank 0's unique id in a file. The id file is for hosts whose bootstrap interfaces share no subnet with rank 0's address, as on our three-host ring. The test also times both exchanges.

Receipts on GB10s (DGX Spark, sm_121, CUDA 13.3, NCCL 2.30.7), the ranks over direct RoCE links, each rank in a container with memlock unlimited:

  • zig build test and zig build test-cuda-host (300/300) pass. zig fmt --check is clean for the touched files. lean_check reports 542 problems, the same as main.

  • One GB10: tf-cuda-test nccl and rank-scatter 0 1 pass.

  • Every rank's rows are byte-equal to an all-gather plus a rank-order sum, for all five row counts and three seeds:

    ranks joined by 2,048 x 4,096 fp32: reduce-scatter all-gather alone
    2 NCCL_COMM_ID 1.01-1.05 ms 2.33-2.37 ms
    2 id file 1.00-1.05 ms 2.38-2.41 ms
    3 (683 / 683 / 682 and 683 / 682 / 682 rows) id file 1.70-1.81 ms 3.71-3.87 ms

    Times are the medians of 20 per rank. The reduce-scatter includes its sum; the all-gather column is the exchange alone, before any sum.

  • NCCL settings, from each rank's file: NCCL_NET=IB, NCCL_IB_GID_INDEX=3, NCCL_IB_ROCE_VERSION_NUM=2, NCCL_IB_ADDR_FAMILY=AF_INET, NCCL_SOCKET_FAMILY=AF_INET, NCCL_MIN_NCHANNELS=4 and NCCL_MAX_NCHANNELS=4. NCCL_SOCKET_IFNAME names the link's interface. NCCL_IB_HCA names the link's two rails at TP=2, and all four ports at TP=3, which also sets NCCL_IB_SUBNET_AWARE_ROUTING=1 and NCCL_NET_PLUGIN=none. At TP=2, NCCL reported 4 coll and 4 p2p channels (2 per peer) and GDR 0. The two other pairs' receipts in this thread measured 1.86-3.61 ms for the reduce-scatter. On one of those pairs, NCCL_MIN_NCHANNELS=4 NCCL_MAX_NCHANNELS=4 alone brought it to 0.99-1.06 ms, against 2.30-2.73 ms with the 64 channels NCCL picked by default (receipt); the rest of the list made no clear further difference there.

  • Against GLM-5.3-Flash and full GLM-5.3 on the native CUDA engine (tracker) #593's gates: the sums stay fp32 and are added in the order an all-gather plus a rank-order sum uses, so there's no new precision and no projection whose reduction order changes.

Where we use it: our GLM-5.3-Flash port reads nvidia/GLM-5.3-Flash-NVFP4 as stored (the NVFP4 family in #548). At TP=2, each prompt chunk's partial rows go through this exchange, and each rank finishes only its own rows; TP=3 is next.

Written with Weave (My Symbiotic AI)

@ashhart

ashhart commented Oct 11, 2026

Copy link
Copy Markdown
Owner

Thank you, @Bizuayeu. A reduce-scatter that adds in rank order, so every rank holds the bits an all-gather plus a rank-order sum gives, is exactly what GLM's TP path needs, and the receipts over RoCE are clear.

Our gate for it needs two linked GB10s, and both of ours were busy through tonight's 1.0.6 landing window, so it isn't in 1.0.6. We'll run tf-cuda-test rank-scatter across our two Sparks next and land it as soon as it holds there.

@Bizuayeu

Copy link
Copy Markdown
Contributor Author

Thank you, @ashhart. No rush; 1.0.6 coming first makes sense. #642 still merges cleanly on v1.0.6, so it should be ready to run as it is. If your two Sparks give anything other than byte-equal rows, or the run won't join, tell us and we'll rerun the same test on our pair with your settings.

@cwschroeder

Copy link
Copy Markdown
Contributor

Independent receipt for a9b86c3 on two ThinkStation PGX (GB10), sm_121, over a direct ConnectX-7 link. Zig 0.17.0, driver 580.178.04, CUDA 13.0.2 (nvcc 13.0.88), NCCL 2.28.9+cuda13.0. NCCL reports NET/IB RoCE, with GDR 0.

  • zig build test: 428/428 on each machine, including all 300 CUDA-host tests. Touched Zig files pass zig fmt --check; lean_check has the same 542 problems as the merge base.
  • One GB10: tf-cuda-test nccl and rank-scatter 0 1 pass.
  • Two ranks, joined by NCCL_COMM_ID and by id file: all five row counts x 4,096 and three seeds pass on both ranks, every owned slice byte-equal to an all-gather plus a host rank-order sum, including the uneven and zero-row cases. No failures, no skips, no code changes.
ranks joined by 2,048 x 4,096 fp32: reduce-scatter all-gather alone
2 NCCL_COMM_ID 2.125-2.129 ms 3.044-3.063 ms
2 id file 1.861-1.899 ms 3.037-3.121 ms

Ranges span the medians of 20 measurements per rank. Reduce-scatter includes its sum; all-gather measures the exchange alone. These runs used CUDA 13.0.2 and NCCL 2.28.9, with NCCL reporting GDR 0. No matched-configuration before/after comparison was run, so these timings establish neither a regression nor a performance bound. TP=3 was not tested.

@Bizuayeu

Copy link
Copy Markdown
Contributor Author

Thank you, @cwschroeder. A second pair on different hardware (PGX, CUDA 13.0.2, NCCL 2.28.9) giving byte-equal slices on every row count and seed, including the uneven and zero-row cases, is exactly the receipt this needed next to ours. For @ashhart: #642 still merges cleanly on v1.0.7.

@SvangenStudios

Copy link
Copy Markdown

Here's a receipt from a second GB10 pair, as asked in #593. The rank-scatter rows are byte-equal on both ranks in all four completed dual-rail runs.

Setup.

  • Hardware: two DGX Sparks (GB10, sm_121), DGX OS 7.6.0, kernel 7.0.0-1019-nvidia, driver 580.178.04 (CUDA 13.0).
  • Link: one QSFP cable directly between the outer ports. Its two PCIe functions are rocep1s0f1 (10.10.1.x) and roceP2p1s0f1 (10.10.2.x), RoCE v2, MTU 9000, 2X / 100 Gb/s each.
  • Code: this PR at a9b86c3, built with zig build -Dnvcc=… -Dsm=121 (Zig 0.17.0, CUDA 13.0.88). I copied the tf-cuda-test binary to the second Spark.
  • NCCL: 2.30.7+cuda13.3, the one in NGC PyTorch 26.07. Neither host has NCCL installed.

Host tests on this branch, each with a fresh --cache-dir: zig build test 37/37 steps and 428/428 tests, zig build test-cuda-host 22/22 steps and 300/300 tests, and zig fmt --check clean on rank_scatter.zig and nccl.zig.

One GB10, on each Spark: tf-cuda-test nccl PASS, and rank-scatter 0 1 PASS (reduce-scatter 0.28 ms, all-gather 0.27 ms).

Two ranks, joined by NCCL_COMM_ID=10.10.1.1:…, with NCCL_SOCKET_IFNAME=enp1s0f1np1 and NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1. NCCL reported NET/IB : Using [0]rocep1s0f1:1/RoCE [1]roceP2p1s0f1:1/RoCE, with GDR 0 (no GPUDirect RDMA).

  • Correctness: in all 4 completed runs, PASS rank-scatter rank r of 2: rows { 1, 4, 64, 2047, 2048 } x 4096 (seeds 1-3), every slice byte-equal… on both ranks.
  • Timing (medians of 20 per rank), 2,048 × 4,096 fp32:
Rails reduce-scatter all-gather alone
both (rocep1s0f1,roceP2p1s0f1), 4 runs 2.39–3.61 ms 1.94–2.30 ms
one (rocep1s0f1), 2 runs 6.70–6.92 ms 3.17–3.47 ms

Your table has 1.00–1.05 ms for the reduce-scatter and 2.33–2.41 ms for the all-gather at two ranks. Here the all-gather is similar or a little faster (1.94–2.30 ms), but the reduce-scatter takes 2.39–3.61 ms: longer than the all-gather on this pair, not about half. The reduce-scatter time includes its rank-order sum, and the all-gather column is the exchange alone. I haven't isolated the cause. If you set any NCCL variables on your pair, for example NCCL_NCHANNELS_PER_NET_PEER or NCCL_P2P_NET_CHUNKSIZE, I can run with the same ones and compare.

One setup note for anyone repeating this on the host: my first two-rank run, with both ranks as plain user processes, failed. Rank 0 stopped in NCCL's P2P preconnect with ibv_reg_mr_iova2 failed with error Cannot allocate memory, and rank 1 never reported a result. On that Spark, rank 0's login session had a locked-memory limit of 8 MiB, hard and soft (ulimit -l 8192). Running each rank in a container with --ulimit memlock=-1 --device /dev/infiniband --network host fixed it, and that's how the runs above were made.

@Bizuayeu

Copy link
Copy Markdown
Contributor Author

Thank you, @SvangenStudios. Byte-equal slices on a third setup (NCCL 2.30.7, two rails) is a useful receipt, and the memlock note will save the next person a failed first run.

Our pair did set NCCL variables, and the PR didn't say so. Each rank ran in a container with memlock unlimited, NCCL 2.30.7 (the same as yours), and these from its rank file:

NCCL_NET=IB
NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1   (rank 1: rocep1s0f0,roceP2p1s0f0)
NCCL_IB_GID_INDEX=3
NCCL_IB_ROCE_VERSION_NUM=2
NCCL_IB_ADDR_FAMILY=AF_INET
NCCL_SOCKET_IFNAME=<the link's interface>
NCCL_SOCKET_FAMILY=AF_INET
NCCL_MIN_NCHANNELS=4
NCCL_MAX_NCHANNELS=4

NCCL reported 4 coll channels, 4 p2p channels, 2 p2p channels per peer and GDR 0. We didn't set NCCL_NCHANNELS_PER_NET_PEER or NCCL_P2P_NET_CHUNKSIZE. The channel count is the difference I'd suspect first: the reduce-scatter is all point-to-point sends and receives, and 4 channels give it twice the per-peer connections a default two-rail setup may pick. If you can rerun with NCCL_MIN_NCHANNELS=4 NCCL_MAX_NCHANNELS=4, that would tell us whether it accounts for the gap. I'll add the settings to the PR description either way.

@SvangenStudios

Copy link
Copy Markdown

Thanks for the settings. Adding just NCCL_MIN_NCHANNELS=4 NCCL_MAX_NCHANNELS=4 brought the reduce-scatter to 0.99–1.06 ms on this pair, close to your 1.00–1.05 ms. Without those overrides, NCCL reported 64 channels and the reduce-scatter took 2.30–2.73 ms.

Setup. This is the same pair and the same commit (a9b86c3) as my earlier receipt. I rebuilt tf-cuda-test the same way and checked that the copy on the second Spark had the same sha256. Each rank ran in a container with memlock unlimited. Every run used NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1, NCCL_SOCKET_IFNAME, NCCL_COMM_ID and NCCL_DEBUG=INFO. There were three arms, run in the order B, N4, F in each of three rounds:

  • B: the baseline, without any of the settings below.
  • N4: adds NCCL_MIN_NCHANNELS=4 NCCL_MAX_NCHANNELS=4.
  • F: adds your full list (NCCL_NET=IB, NCCL_IB_GID_INDEX=3, NCCL_IB_ROCE_VERSION_NUM=2, NCCL_IB_ADDR_FAMILY=AF_INET, NCCL_SOCKET_FAMILY=AF_INET, plus the four-channel overrides).

Results. All 18 rank runs printed PASS rank-scatter … every slice byte-equal. Each run reported GDR 0, with both rails under NET/IB : Using. The times are medians of 20 measurements per rank, for 2,048 × 4,096 fp32, and each range spans three runs × two ranks. The reduce-scatter time includes its rank-order sum. The all-gather time is the exchange alone.

Arm NCCL reported reduce-scatter all-gather alone
B (baseline) 64 coll channels, 64 p2p channels, 2 p2p channels per peer 2.30–2.73 ms 2.10–2.22 ms
N4 4 coll channels, 4 p2p channels, 2 p2p channels per peer 0.99–1.06 ms 2.03–2.42 ms
F 4 coll channels, 4 p2p channels, 2 p2p channels per peer 0.99–1.02 ms 2.08–2.35 ms

Today's baseline was faster than yesterday's 2.39–3.61 ms, so the comparison above uses the baseline from these rounds.

The four-channel overrides alone brought the reduce-scatter to about 1 ms in all three rounds. N4 and F had overlapping ranges, and the full list gave no clear further improvement. With only three rounds in a fixed order, I can't conclude that the other settings have no effect.

The all-gather ranges overlap, with no consistent improvement across rounds and more variation in N4. I haven't profiled the parts of the reduce-scatter separately, or checked why the baseline reports 64 channels.

@Bizuayeu

Copy link
Copy Markdown
Contributor Author

Thanks, @SvangenStudios. That confirms the suspect the PR description named: with four channels your pair lands at our 1.00–1.05 ms, and the rest of our list made no clear difference on top of that. We haven't looked at why NCCL picks 64 channels on these pairs either.

I'll change the description's last sentence from "isn't confirmed yet" to point at your three-arm receipt. The byte equality doesn't depend on any of these settings, as your 18 runs show.

@cwschroeder

Copy link
Copy Markdown
Contributor

Thanks, @SvangenStudios and @Bizuayeu. The matched baseline and four-channel arms give us a concrete setting to compare on the PGX/GB10 pair used in my earlier receipt.

For that comparison I will keep binary, NCCL version, topology and call sizes fixed, check byte-equal slices on both ranks, and report reduce-scatter and all-gather separately. The HCA-startup retest requested in #636 is a separate check. Model throughput will need its own measurement.

@Bizuayeu

Copy link
Copy Markdown
Contributor Author

For whoever merges: #636 adds the same ncclSend and ncclRecv entries to zig/src/cuda/nccl.zig. If this lands first, #636 drops those two lines on its rebase; if #636 lands first, this one does. The rank-order reduce-scatter here is what #636's NCCL backend would need for an exact rank-order sum past two ranks (details in #636).

…_scatter.zig)

nccl.zig binds ncclSend and ncclRecv beside the collectives. rank_scatter.zig cuts a
[rows][width] fp32 partial by whole rows (the first rows % world ranks one more, as
canon.partition cuts equal ranks), sends each peer its rows and receives its own from each
peer in one group, then rank_sum.cu adds the terms rank 0 first, rounding each add. Every
rank so holds the bits an all-gather followed by a rank-order sum gives, with each rank
receiving only its own rows. NCCL's ncclReduceScatter is not used: it adds in its ring's
order, not rank order.

tf-cuda-test rank-scatter <rank> <n> [id] checks every rank's rows against an all-gather
plus a host rank-order sum (rows 1, 4, 64, 2047 and 2048 by 4096, three seeds), the ranks
joined by NCCL_COMM_ID or by rank 0's unique id in a file, and times both exchanges.
@Bizuayeu
Bizuayeu force-pushed the pr3/nccl-rank-scatter branch from a9b86c3 to 0a36982 Compare October 12, 2026 19:43
Bizuayeu added a commit to Bizuayeu/TensorFold that referenced this pull request Oct 13, 2026
… over hosts

The family ashhart#593 tracks, on the shared NVFP4 layer (`qlinear`, `experts.Grouped`) and hand-written `.cu` for the
rest: KDA and MLA/DSA attention, the hyper-connection glue, MoE routing, the MTP draft head and sampling. Rank 0
serves through `tensorfold-native serve` and the other ranks follow its steps; each rank holds its share of every
axis (uneven where a width doesn't divide), and a prompt chunk's rows are exchanged with a rank-order
reduce-scatter (ashhart#642). Drafts and concurrency go through the lanes core: `prefill_step` takes a prompt a chunk at a
time between the other streams' rounds, and `hidden_rows` lets several streams' verify windows share one forward.
The prompt cache keeps states on the family's 2,048-row chunk starts (`snap_grid`).

Checks in `tf-cuda-test`: `glm-ref` holds every projection whose reduction order is the kernels' own against float64
on synthetic inputs; the lanes checks (cover, concurrent, shared) run at TP over hosts and, with `glm-synth`'s
two-layer checkpoint, on one GPU.
@ashhart ashhart mentioned this pull request Oct 14, 2026
@ashhart
ashhart merged commit 1e3678c into ashhart:main Oct 16, 2026
2 checks passed
ashhart added a commit that referenced this pull request Oct 16, 2026
ashhart added a commit that referenced this pull request Oct 16, 2026
@ashhart

ashhart commented Oct 16, 2026

Copy link
Copy Markdown
Owner

Thank you, @Bizuayeu. A reduce-scatter that adds in rank order, so every rank holds the bits an all-gather and a rank-order sum would give while receiving only its own rows, is what an exact sum across tensor-parallel ranks needs. Thank you too, @cwschroeder and @SvangenStudios, for running it on your own pairs.

It ships in TensorFold 1.0.12 under your name (d309dad9c). Our 88943c573 moves the sum's launch into RankScatter.rankSum(), which run() calls as before, adds tf-cuda-test rank-sum (the sum kernel alone on one GPU for 2 to 5 ranks, against a host sum in rank order) and wraps the new lines. Your head is merged after them with no change to the tree (1e3678cda), so this pull request shows as merged.

On one GB10, tf-cuda-test nccl, rank-scatter 0 1 and rank-sum pass, and a build whose kernel adds its own term first fails rank-sum. Your share test passes on macOS and Linux.

Our own two-rank run is still owed, so your receipts and the other two pairs' carry the exchange itself. GLM-5.3-Flash at TP over hosts (#804) is its first caller.

Posted by Ash's AI Assistant Synthena

@Bizuayeu

Copy link
Copy Markdown
Contributor Author

@ashhart, one correction for the record: #804 doesn't call RankScatter. It left PR D at the 1.0.10 rebase, and GLM's split exchange gathers each rank's block with its own sends and adds them in rank order (exchangePieces in cuda_forward.zig), so rank_scatter.zig has no caller in main yet. The receipts in this PR (ours at TP=2 and TP=3, and the two other pairs') are its only runs on hardware. Its fitting caller is #636's NCCL backend, for a rank-order sum past two ranks.

@cwschroeder

Copy link
Copy Markdown
Contributor

Confirmed by a source check on main 46a08c18: RankScatter is exported by the CUDA root and exercised by the test runners, but no family calls it. GLM's split exchange uses its own grouped sends/receives.

Our two-GB10 receipt here covers the standalone rank-scatter test, not GLM end-to-end or #636's integrated collective. I will keep that scope explicit in our reports. Once #636 uses the rank-order reducer beyond two ranks, that caller will need its own integration gate.

jayleaton added a commit to jayleaton/TensorFold that referenced this pull request Oct 16, 2026
…-fabric

build.zig keeps main's kernel-contract and video steps beside the TP steps; nccl.zig keeps ashhart#642's send/recv with the TP abort and async-error symbols.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
jayleaton added a commit to jayleaton/TensorFold that referenced this pull request Oct 16, 2026
…t#861) into stack/13b-performance-device

The shared corrections reach the device head: sessions null-unlimited / zero-none budgets and the Mac NVMe-test guard
(ashhart#637), the TP graph-cache startup agreement (ashhart#636, graph_cache_env), main's ashhart#642 NCCL send/recv, and ashhart#861's lane core.
- sessions_gpu.zig: TF_DSV41_SESSION_RESIDENT / _RESIDENT_GIB default to null (unlimited). The old 0 default
  would now mean "no idle chats kept". The store's policy comes from one `policy` function, and its test fails with
  the old default.
- zig/build/cuda.zig: main's runner block with the DeepSeek-V4.1 image-aware runner, tool, server and sampler installs.
- zig/build/dsv41.zig: the host compile check's options carry slide_test_hooks.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
BlivionIaG pushed a commit to BlivionIaG/TensorFold that referenced this pull request Oct 20, 2026
BlivionIaG pushed a commit to BlivionIaG/TensorFold that referenced this pull request Oct 20, 2026
The contributor's commit landed above with the same code; this keeps the PR's head in the history with no change to the tree.
BlivionIaG pushed a commit to BlivionIaG/TensorFold that referenced this pull request Oct 20, 2026
BlivionIaG pushed a commit to BlivionIaG/TensorFold that referenced this pull request Oct 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants