[Draft] Introduce internode communication support via NCCL GIN - #67
Draft
Rachmanino wants to merge 30 commits into
Draft
[Draft] Introduce internode communication support via NCCL GIN#67Rachmanino wants to merge 30 commits into
Rachmanino wants to merge 30 commits into
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
👋 Hi! Thank you for contributing to the TileScale project. Please remember to run We appreciate you taking this step! Our team will review your contribution, and we look forward to your awesome work! 🚀 |
GIN (GPU-Initiated Networking) is the network mode of the NCCL Device API, first usable in NCCL 2.28.7. Most environments still ship 2.27.5, which has no `nccl_device/gin.h` and no `ncclDevCommCreate`, so the feature has to be compiled out rather than assumed present. The flag has to be defined twice by two independent mechanisms, because the host library and the JIT are separate compilations. `src/cuda/CMakeLists.txt` gates whether `codegen_cuda.cc` emits the `nccl_gin.h` include; `env.py` autodetects an include dir for the JIT, which `contrib/nvcc.py` and both `libgen.py` back ends turn into `-I` plus `-DTL_ENABLE_NCCL_GIN=1`. Keeping the two halves independent means a host build without GIN headers still works, and `TL_NCCL_PATH` can override the JIT half alone. `FindNCCLGin.cmake` insists on `nccl.h`, `nccl_device/gin.h` and the `ncclDevCommCreate` export together: a tree can advertise a new enough version while missing the device path, and only the export proves the library half is there. It also searches the versioned `libnccl.so.2` soname, since pip wheels ship no unversioned dev symlink.
The table is the contract between the host allocator, the device headers and the host-side TMA encoder in runtime.cc, and all three indexed it with raw integers. A one-sided edit does not fail: cached kernels keep reading the old slot while the allocator publishes the new one, which corrupts remote addressing and hangs with no error. `runtime.cc` was already wrong this way, reading `meta_data[2 + rank]` after the header grew past index 2. `meta_layout.h` now declares every offset once as `TL_META_*` and is free of CUDA constructs so plain host translation units can include it. It documents why the GIN context count is *not* in the table: the devcomm may grant fewer contexts than requested, so the count has to be read back on the device, and publishing the host's request would let a kernel index past the end. Multi-node makes the global/local distinction load-bearing, so the table now carries node_rank, num_nodes, local_rank and local_world_size, and the peer pointer array is indexed by *local* rank -- inter-node peers have no local virtual address and can never appear there. `get_remote_base_ptr` returns 0 for a non-local rank rather than reading a neighbour's slot, and `runtime.cc` reduces a global destination to its node-local rank before indexing. `get_rank`/`get_num_ranks` stay as aliases so single-node kernels are unaffected.
GIN one-sided ops are callable from inside a kernel, unlike host NCCL collectives, which is what makes them the right fit for TileScale's kernel-side data plane. This exposes `T.nccl_gin.put` / `put_signal` / `signal` / `wait_signal` / `flush`, lowering through `src/op/nccl_gin.cc` to `tl::gin::` helpers in the device header. Two properties of GIN signals drive the design and are easy to get wrong: Signal state is per context. A put issued on sender context `i` increments the receiver's signal through context `i`, so a CTA spread over `C` contexts sees only `1/C` of the rank's arrivals. `wait_signal` therefore divides the caller's grid-wide target by `context_span()` *on the device*, because only the device knows how many contexts the devcomm actually granted -- `ncclDevCommRequirements.ginContextCount` is documented as a hint, and asking for 8 while 4 are granted means indexing past the end of the context array, which hangs. Requesting the count on the host and trusting it produced a reading of 332 GB/s, above the PCIe ceiling for one GPU's egress, because the wait was 2x too weak and the kernel returned early. Signals are cumulative and a wait does not consume them, so the expected count must be a runtime argument rather than a compile-time constant. A constant target is satisfied instantly on every launch after the first: correct once, then a silent no-op that benchmarks nothing. The codegen change adds `tl::gin::` to the prefixes that mark a kernel as distributed, so a kernel using only GIN still gets the distributed includes, and emits the `nccl_gin.h` include only when the host build found GIN.
…IN devcomm GIN addresses memory as an `(ncclWindow_t, byte offset)` pair, not a raw pointer, because a remote rank's allocation has no local virtual address. The window handle is not per-peer: one local handle plus a peer index names any rank's bytes. So the whole allocator arena is registered once, collectively, at allocator init rather than per tensor -- `ncclCommWindowRegister` is a collective, and per-tensor registration would turn every `allocate_tensor` into a world-wide barrier plus a handle table. Because the arena is symmetric, `local_ptr - arena_base` is the offset valid on every rank, the same subtraction the intra-node peer-pointer path already performs. Window registration rejects `cudaMalloc` memory: measured on NCCL 2.28.9, a plain `cudaMalloc` pointer fails with "invalid argument" under both `NCCL_WIN_COLL_SYMMETRIC` and `flags=0`, while a VMM-mapped arena registers cleanly. The allocator therefore raises rather than silently degrading when a window is requested on the cudaMalloc backend, and `shared_memory.cc` exposes the VMM capability probe the allocator needs to check that up front. Ordering is the one non-obvious constraint. `ncclDevCommCreate` needs a communicator that still supports symmetric memory, and a torch `ProcessGroupNCCL` communicator loses that after its first collective -- the call then *segfaults* rather than returning an error, so it cannot be attempted and recovered from. Allocator construction runs several collectives before GIN setup, so reusing the caller's group crashes every time. `_init_arena_window` instead makes a private `dist.new_group(backend="nccl", device_id=...)`; `device_id` makes the communicator eager, so its pointer is valid immediately instead of being created by the very collective that would invalidate it. The devcomm is created on that group *before* window registration, which must then use the same comm, and teardown runs in reverse. `init_dist` now returns node topology alongside the group so the allocator can tell node-local peers from inter-node ones, and `get_allocator` takes it as `node_info`. Single-node callers are unaffected: `num_nodes == 1` takes the existing IPC/VMM path untouched and no window is registered unless `TILESCALE_USE_GIN` asks for one.
… GIN
Three collectives built on `T.nccl_gin`, sharing `internode_common.py` and
launched by `run_internode.sh`. All three verified on two physical nodes on
2026-08-03 against their torch references (`all_gather_into_tensor`,
`all_reduce`, `reduce_scatter_tensor`) with `LOCAL_WORLD_SIZE=1 NNODES=2`, so
every put crossed the RoCE fabric.
At 64 MB shards, bf16, on two idle nodes they all beat torch NCCL:
tilescale torch ratio
allgather 47.6 39.8 1.20x
allreduce 47.2 45.1 1.05x
reduce_scatter 46.2 42.8 1.08x
46-48 GB/s is 93-95% of one 400 Gbps NIC's 50 GB/s, i.e. the collectives are
at line rate, which is why little else moves them.
The shape that gets there, modelled on Triton-distributed's inter-node sender:
size the grid by peers and channels, not by payload, so one CTA issues one
large put. The first version launched a CTA per 8192-element block, turning an
8 MB shard into 512 separate 16 KB puts -- latency-bound, since RDMA is only
bandwidth-bound once messages are large. That plus a `T.Parallel` local copy
and a per-launch signal target is ~25x the first implementation.
`sweep_internode.sh` and the `--tune` mode exist because the interesting knob
is not stable across conditions. `--chunks` matters for the reduce variants,
whose CTAs also carry the reduction (allreduce 41.5 -> 47.2 GB/s from 4 to 64)
and is flat for allgather. `--gin-contexts` is insurance rather than a win: on
idle NICs one context already reaches line rate, but on a NIC shared with
another tenant one context collapsed to 23 GB/s while 2-4 held ~44. Tuning on
contended hardware therefore misattributes the win, so the tuner verifies every
config and re-times torch before and after each sweep.
`--algo oneshot` for allreduce is kept behind a flag: at W=2 it sends the same
bytes as two-shot in half the phases and still loses (30.3 vs 43.4 GB/s),
because it reduces over the full buffer instead of a shard. The balance shifts
with rank count and dtype.
`test_meta_layout_drift.py` parses the `_META_*` ints in `allocator.py` and the `TL_META_*` defines in `meta_layout.h` and compares them. The layout is declared twice by necessity, and a one-sided edit corrupts remote addressing silently instead of failing -- a `TL_META_PEER_BASE` experiment that moved the offset and moved it back cost hours, because cached kernels kept the old index. The test also asserts the offsets are a dense range, that the variable-length peer array stays at the tail, and that `runtime.cc` never indexes `meta_data` with a raw integer. `probe_gin_window.py` checks the host GIN path with no native build: it loads `nccl_window.py` by path, so it runs on a node that has a GIN-capable NCCL but no compiled tree. It covers window registration on a VMM arena, rejection of `cudaMalloc` memory, and the devcomm-before-registration ordering, which makes it the fastest way to tell whether a node can do GIN at all. `test_nccl_gin_put.py` is the minimal device-side put: exchange a buffer with a peer, wait on the signal, verify the bytes.
Rachmanino
force-pushed
the
wt/nccl-gin-internode
branch
from
August 3, 2026 15:13
5b7002d to
9a46006
Compare
Compose the GIN collectives with a TileLang GEMM for the two fused shapes TP inference needs. Both verified on two nodes at 8192x4096x4096 bf16 against a torch reference. AG-GEMM overlaps comm with compute by splitting the GEMM into launches per row range: the allgather runs on a side stream while the GEMM covers the rows this rank already owns, then a stream wait gates the peer-owned rows. That takes 1.362 -> 1.110 ms, from 0.80x to 0.98x of torch. Overlapping by waiting on the GIN signal inside the GEMM instead deadlocks: the 2048-CTA grid fills every SM with waiting CTAs and the side-stream comm kernel never gets scheduled to signal them. Separate streams do not imply co-residency. A persistent grid capped at 132 CTAs avoids the hang but is slower than serial, so the split-launch form is kept. GEMM-RS stays serial. Its dependency runs the other way -- the GEMM produces what is sent -- and the obvious reorder is wrong: reduce_scatter_kernel also folds in this rank's own block, so starting it early reduces stale data and mismatched 5.9M of 8.4M elements. Real overlap needs the collective split into a per-peer put plus a separate accumulate. The remaining gap to torch is the GEMM, not the network: our best tile config reaches 470 TFLOP/s against cuBLAS's 1775 on a B200, while our allgather is the faster half at 0.726 ms against torch's ~0.94.
…ernels Port examples/gemm_sm100/gemm_tcgen5mma_ws_persistent.py into a row-range kernel and use it for AG-GEMM and GEMM-ReduceScatter. Both now beat torch on two nodes at bf16: ag_gemm 0.906 ms 303 TFLOP/s vs torch 1.097 1.21x (was 0.80x) gemm_rs 0.546 ms 252 TFLOP/s vs torch 0.571 1.05x (was 0.70x) The naive T.copy/T.gemm loop it replaces reached 470 TFLOP/s where the warp-specialised form reaches 1334 on the same shape, so the fused kernels were compute-bound by roughly 3x rather than limited by the network. What the fast path relies on: tcgen05 MMA accumulating in tensor memory, double-buffered by wave parity; warp specialisation, with warp 0 issuing TMA loads, warp 1 the MMAs and warps 4-7 the epilogue; a persistent one-CTA-per-SM grid; and hand-computed mbarrier phase parity to overlap loads with MMAs. m_offset is a runtime argument so one compiled kernel covers both the whole output and a single rank's row block, which is what the split-launch overlap needs. With the GEMM this fast the fused kernels are comm-bound -- 0.725 ms allgather against 0.188 ms of GEMM -- so split-launch overlap no longer pays and serial is kept as the default.
Add a POSIX-FD export/import path for the allocator arena and use it whenever fabric handles are unavailable. Fabric handles need an IMEX channel. On a node without one (/dev/nvidia-caps-imex-channels absent, nvidia-imex inactive) the arena is still a VMM allocation, so NCCL windows register and GIN works, but it is created with CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR and cannot be exported as fabric: cuMemExportToShareableHandle fails with "invalid argument" and every rank dies during allocator construction. That made more than one rank per node impossible here, so 16 GPUs across 2 nodes could not start at all. A POSIX-FD handle is only meaningful inside the owning process, so it has to be duplicated into the importer. Instead of SCM_RIGHTS over a unix socket, which needs a rendezvous directory and a connect/accept dance between every pair of ranks, publish (pid, fd) through the process group that already exists and let peers duplicate the descriptor with pidfd_getfd. Two syscalls, no rendezvous. It needs ptrace-level access to the peer, which holds for ranks of one job. Verified with 8 ranks per node on two nodes: allocator construction completes and peer base pointers are populated, where before it raised at the export stage. This unblocks hierarchical collectives, whose intra-node half needs real peer pointers -- TILESCALE_SKIP_PEER_IPC=1 keeps a pure-GIN run alive but leaves those pointers zero, which is precisely what a 2D algorithm cannot use.
The flat allgather is only optimal with one GPU per node. At 16 GPUs on two nodes every rank pushes its shard to all 15 peers over GIN, including the 7 on the same machine, and bus bandwidth collapses to 2.7 GB/s against torch's 309. With one rank per node a flat algorithm *is* the hierarchical one, so the existing benchmark could not see the flaw. Decompose along the topology instead: a rail-aligned inter-node exchange with the same local index on other nodes, then an intra-node allgather over the existing peer-pointer path. Doing only that is not enough -- one inter kernel, a barrier, one intra kernel runs the NIC and NVLink strictly serially and lands at 260 GB/s, still below torch. The win comes from splitting phase 2 by data dependency: a sibling's own shard is already in the symmetric arena before the collective starts, so those slots are ordered against nothing and can be pulled on a side stream concurrently with the fabric transfer. Only the other nodes' slots need the barrier. That halves the NVLink bytes on the critical path and hides the rest. 16 GPUs, 2 nodes, 240 MB bf16: 349 GB/s against torch's 309 (1.13x) and triton-dist's 290 (1.20x). --no-overlap keeps the serial variant to show the difference. Note src_pe/dst_pe are global ranks: get_remote_base_ptr returns 0 for a peer it considers inter-node, so a local rank yields a null base and faults -- but only on nodes other than node 0, where the two numberings coincide.
The transpose of the 2D allgather, and the flat reduce-scatter collapses at 16 GPUs for the same reason: every rank pushes a slice to all 15 peers over GIN, including the 7 siblings on the same machine. Reduce within the node over NVLink first -- each rank pulls, from every sibling, the part of that sibling's input belonging to our rail index -- then exchange the per-node partials rail-aligned. No host barrier is needed here, unlike the allgather: the intra phase reads input buffers the collective never writes, and the rail phase sends only bytes this rank produced, with the GIN signal proving arrival. 16 GPUs, 2 nodes, 240 MB bf16: 2.5 -> 293.9 GB/s, against torch's 319 (0.92x). Two negative results are encoded in the flags rather than discarded. The intra kernel's grid is one CTA per (node slot, chunk), which at the rail kernel's --chunks 8 is 16 CTAs on 148 SMs and 0.31x of torch; it carries no GIN signal so --intra-chunks decouples it, and 1024 wins. Pipelining the two phases by chunk to overlap them is 4x *slower*, because each rail launch then moves a few MB from a couple of CTAs and RDMA needs large messages -- so the overlap cuts along the node axis instead, where the message stays whole. That is worth only 3.6%, since the put kernel returns once the RDMA is issued and the flight time was already absorbed by the wait; --no-overlap shows it. The remaining gap to torch is the scratch round-trip in the NVLink reduce, which needs a vectorised elementwise peer load to remove.
…ce-scatter Multicast looked unavailable on this cluster and the allocator's multimem path was dead code here, which matched the documented "VMM/multicast needs an IMEX channel" limitation. That was wrong. cuMulticastCreate does fail with CU_MEM_HANDLE_TYPE_FABRIC (NOT_PERMITTED, no IMEX), but succeeds with CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR, and export/import/addDevice all work. NCCL logging "NVLS multicast support is available on dev N" on the same node was the giveaway. So pick the handle type by capability instead of hardcoding fabric, and move the descriptor into peers with the same pidfd_getfd duplication already used for the VMM arena. supports_multicast_impl now requires supports_vmm_impl rather than supports_vmm_fabric_impl, which is what had been vetoing it. All 23 tests in test_multimem.py pass on 8 GPUs over the FD-shared multicast object. Note every multicast probe needs a live CUDA context, since it reaches cuCtxGetDevice; cuInit plus a primary-context retain is not enough. A probe run before the allocator exists reports a silent False on hardware that fully supports multicast, so Context.supports_multicast touches CUDA itself. Then use it for the intra-node half of the 2D reduce-scatter. The pull implementation structurally cannot reach torch: reducing on the consumer means moving (lws-1)/lws of the input over NVLink into a staging buffer and reading it back, ~210 MB of NVLink plus ~500 MB of avoidable HBM traffic per rank, which was exactly the 8% it was losing. multimem.ld_reduce reduces in the switch, so the rank reads only the bytes it keeps and no staging buffer exists. This is the same NVLS mechanism torch is using here, so it also makes the comparison fair. 16 GPUs, 2 nodes, 240 MB bf16: 293.9 -> 339.0 GB/s, against torch's 322 (1.05x). The tile width is not tunable -- bf16 multimem lowers to packed x2, so the fragment must be one contiguous pair per thread or layout inference rejects it. That pins each thread at one 4-byte load, so work per thread has to come from a tile loop instead: one tile per CTA gives 317 GB/s, eight give 339.
The 2D allgather finished inside the node by having every rank read its seven siblings' buffers: ~220 MB of NVLink reads per rank to place 15 shards. But a rank only ever *owns* two of those shards, and multimem.st against the multicast VA reaches all eight local ranks in one instruction, so publishing what we own costs ~31 MB of stores and the switch does the fan-out. 403.6 GB/s against torch's 309 (1.31x) and triton-dist's 290 (1.40x), up from 358.6 on the pull path. --intra pull keeps the portable version, which is still the only option without multicast. The rail phase now lands in a separate railbuf rather than straight into out: the output has to live in the multicast buffer, while a GIN put must target the registered arena window, and those are different allocations. --mc-tiles is what matters, for the same reason as in the reduce-scatter: the tile width is pinned at 2*threads by the packed-x2 fragment layout, so looping over tiles is the only way to give a thread more than 4 bytes of work. 8 tiles gives 355 GB/s, 32 gives 404, flat beyond that. Retuned reduce-scatter to 32 as well, which takes it from 339 to 397 GB/s (1.24x torch).
…GEMM-RS on them
The 2D allgather and reduce-scatter had grown to ~450 lines each with their kernels
inline, and the three remaining flat kernels needed the same machinery. Rather than
copy it twice more, move it to internode_2d: the rail-aligned fabric hop, the
multimem broadcast and reduce, their portable get_block fallbacks, and two host
drivers that own the buffers and the signal bookkeeping.
Every collective here is the same two halves recombined -- allgather is
rail-then-broadcast, reduce-scatter is reduce-then-rail, allreduce is the two
composed, and the fused GEMM kernels wrap one of them around the tcgen05 GEMM -- so
the five examples are now ~120 lines each against one shared module.
Two seams were worth care in the composition. Allgather2D takes an existing arena
tensor as its input, so allreduce hands the reduce-scatter's output straight over
with no full-shard copy. And the two halves must use *different* GIN signals: signal
state is cumulative and a wait does not consume it, so sharing one would let the
allgather's wait be satisfied by the reduce-scatter's arrivals and return without
waiting -- silently, and only under repetition.
GEMM-RS stays serial. The dependency runs the wrong way for the AG-GEMM trick, and
the flat version's attempt at it silently mismatched 5.9M of 8.4M elements.
Also adds run_2d_proxy.sh, which presents one node's 8 GPUs as 2 nodes of 4. Nothing
in these kernels inspects the physical topology, so this exercises the whole
structure -- rail GIN between the groups, per-group multicast, the barriers, the
signal bookkeeping -- and catches every correctness bug a real two-node run can. All
five collectives pass under it. It is not a performance proxy: the inter-node hop
loops back through the NIC and both groups share one NVSwitch.
Note --chunks defaults to 8 on the 2D path rather than the flat 64. At 64 the rail
kernel fails to lower ("Can't fetch the lanes of a scalable vector"), and gemm_rs_2d
returned a wrong answer on one rank instead of erroring. Not yet root-caused, so
large chunk counts are treated as unsupported here.
… compute Serial fusion only inherits the collective's advantage: at the default shape it is 0.156 ms of comm plus 0.183 of GEMM against unfused torch's 0.204 + 0.170, about 1.10x. The prize is overlap, where the floor becomes max(comm, gemm) -- and where cuBLAS being 8% faster than our GEMM in isolation stops mattering, because that GEMM is covering the network. The existing --mode overlap only hides this rank's own row block, 1/16 of the compute, so it is worth ~3%. Note instead that global rank is node*lws + local and the row block index *is* the global rank, so a whole node's rows are contiguous, and our own node's rows are complete after the intra-node broadcast alone -- they never touch the fabric. So --mode pipeline runs the fabric hop on its own stream, publishes and then GEMMs our node's rows while it is still in flight, and only then joins. One GEMM launch per node rather than per rank; half the compute overlapped at two nodes. Done with stream and event ordering over Allgather2D's steps, which are now exposed individually, plus one extra barrier. Deliberately not an in-kernel signal wait: that deadlocked before, because a 2048-CTA GEMM grid fills every SM with waiting CTAs and the comm kernel never gets scheduled, and capping the GEMM at 132 persistent CTAs to fix the hang was slower than serial. Doing that properly needs an SM partition and is not something to tune blind. serial stays the default because it is what has been verified on hardware. Every kernel the pipeline needs lowers at both 16-GPU and proxy shapes (checked for the tcgen05 row-count constraints and for the put-size lowering bug), but the cluster has had no two free nodes since this was written, so the mode itself is untested. Flip the default once run_2d_proxy.sh confirms it.
gemm_rs_2d mismatched on 11 of 16 ranks at two nodes while passing on the 8-GPU single-node proxy -- the signature of a skew-dependent race, and it was one. The GEMM writes the collective's input, and the intra-node reduce reads *every local rank's* copy of that input through the multicast VA. Stream order sequences our own GEMM before our own reduce and says nothing about a sibling's GEMM, so a rank could reduce whatever a slower sibling had written so far. On one node the ranks stay tightly enough synchronised to hide it; across two nodes they do not. ReduceScatter2D advertises "no barrier needed", which is true when the input is filled once before the loop as in the standalone example. A producer inside the loop breaks that precondition, so the promise now states it and the fused example fences. gemm_rs_2d is now correct on all 16 ranks: 0.308 ms / 111.6 TF against unfused torch's 0.333 / 103.3, so 1.08x. Also make --mode pipeline the AG-GEMM default, now that it has run on hardware: correct on all 16 ranks, 0.430 ms against serial's 0.481 and torch's 0.500 -- pipelining the fabric hop under our own node's GEMM is worth 12%, and the fused kernel is 1.16x torch.
Every 2D collective had the same shape: a fabric hop, then ~0.19 ms of multicast
broadcast or reduce that could not start until the whole hop had landed. The fabric
floor at this size is 0.33 ms per shard, so that tail was a third of the runtime.
Split the hop into --rail-groups groups, each on its own GIN signal, so group g's NVLink
work runs while group g+1 is still in flight. The critical detail is that the per-put
size is *unchanged*: splitting chunks=8 into 2 groups of 4 CTAs keeps every message at
2 MB and only halves per-group parallelism. An earlier attempt split the payload into
more, smaller messages instead and came out 4x slower, because RDMA is bandwidth-bound
only once messages are large.
allgather 403.6 -> 471.5 GB/s 1.33x -> 1.52x torch
reduce_scatter 397.9 -> 428.9 GB/s 1.24x -> 1.34x torch
allreduce 401.1 -> 452.2 GB/s 0.81x -> 0.91x torch
2 groups is the optimum. 4 needs --gin-contexts 2 to keep each group's grid a multiple
of the context count, and losing contexts costs more than the extra group gains (434.0);
raising --chunks to 16 to keep 4 groups at 4 contexts hits the put-size lowering bug.
Three things this needed:
- The publish offset must be compile-time, so a group gets its own compiled kernel. A
runtime offset defeats multimem's bounds prover: "multimem packed multicast region must
be provably in bounds".
- The wait must be its own kernel. Its grid has to match the sender's chunk count,
because a put from sender CTA b signals through context b % contexts -- a wider grid
parks CTAs on contexts nothing signals, a narrower one rounds the target down. But the
arithmetic wants a wide grid, so folding them together made the merged allreduce
*slower* than composing the halves.
- Signal *ranges* must be disjoint, not merely signal ids. Each half now occupies one
signal per group, so the allgather starts at SIGNAL_DATA + rs.signals_used; starting it
at SIGNAL_PHASE2 overlapped the reduce-scatter's range and its group-0 wait was
satisfied by the other half's arrivals, corrupting exactly the second half of the
output.
Also adds --algo merged for allreduce: one fabric hop carrying every node partial, so
each rank finishes all slots locally and there is no second hop. Same fabric bytes, one
serialisation and one barrier instead of three -- but it measures 1.296 ms against
composed's 1.044, because the NVLink publish cost is identical and with only `nodes`
slots the pipeline is too coarse. Kept behind the flag with the number, since it is the
right shape at more nodes.
The tails also stop copying their own partial into the inbox and read it as one more
term, saving a shard of HBM write plus read per group.
Allreduce was the one collective losing to torch. Composed as reduce-scatter then
allgather the two halves add up exactly -- 0.550 + 0.500 ms -- because the allgather
cannot start until the reduce-scatter has produced the shard it broadcasts. But that is
only true per group: the allgather's hop for group g needs nothing but the
reduce-scatter's sum for group g. So push each group across the second hop as soon as it
is summed, and publish our own copy of that slice while it is in flight.
allreduce 1.088 -> 0.944 ms, 401 -> 500 GB/s, 0.81x -> 1.01x torch at 240 MiB
Two size-dependent failures fell out of sweeping sizes, both of which had been hiding
behind a single benchmark point:
Pipeline depth has to be capped by slice size, not just by the divisibility rules. Each
group costs about six launches, and once a slice is small those dominate: 8 groups is
0.95 ms at a 15.7 MB shard and 0.61 ms at a 3.1 MB shard, where the fabric alone needs
0.13 -- half a millisecond of pure overhead and 0.50x torch. MIN_GROUP_BYTES backs the
depth off and logs it; 48 MiB allreduce went 0.612 -> 0.357 ms.
And the put-size lowering bug made ordinary buffer lengths simply unusable: 240 MB works
at --chunks 8 while 120 MB does not, because its 491520-element put lands in the bad set.
workable_chunks() halves the count until it lowers, which doubles the put size and moves
off the bad value -- larger puts being what RDMA wants anyway. Probed with plain
tilelang.compile, since ctx.compile is collective and a rank raising inside it would hang
rather than fail. 120 MiB now runs, at 1.43x.
Also fixes a signal-range collision the pipelining introduced: each half occupies one
signal per group, so the allgather must start at SIGNAL_DATA + rs.signals_used. At
SIGNAL_PHASE2 it overlapped the reduce-scatter's range and its group-0 wait was satisfied
by the other half's arrivals, corrupting exactly the second half of the output.
Scoreboard, 16 GPUs, bf16, torch drift shown in each run:
48 MiB 120 MiB 240 MiB
allgather 1.80x 1.43x 1.50x
reduce_scat 1.47x 1.44x 1.34x
allreduce 0.87x 1.06x 1.01x
ag_gemm 1.16x
gemm_rs 1.08x
Allreduce still loses at 48 MiB: it is the only collective with two fabric hops and so
pays roughly twice the launches, and at that size NCCL's single fused kernel wins on
overhead rather than bandwidth. --algo merged (one hop carrying every node partial) is
kept behind a flag with its measurement -- 1.296 ms, worse here, but the right shape at
more nodes.
Start-up on 16 ranks is ~19 s with a warm cache and nobody knew where it went.
TL_STAGE_TRACE now timestamps each Context stage, and _collective_stage reports work and
sync time separately so a slow rank is distinguishable from slow work.
5.7 s init_dist
12.0 s allocator: 6.4 devcomm, 1.8 GIN process group, 1.3 Device API check (sync),
1.5 first collective, 0.25 all eight multicast stages, 0.3 window + handles
1.5 s 22 kernels, all cache hits
So ncclDevCommCreate alone is a third of it. The obvious lever looked like the resource
request -- a GIN context is a QP per peer -- so TILESCALE_GIN_CONTEXTS/_SIGNALS/_COUNTERS
now override it. It does not help: 8, 4 and 1 contexts measure 6.53, 6.90 and 6.36 s.
Flat. The cost is DOCA GPUNetIO / IBGDA setup inside NCCL, so it is not tunable from here,
and that negative result is recorded so the lever is not tried again.
Worth knowing what this clears: ~14 of the 19 s is NCCL initialisation, and everything we
own -- multicast setup, window registration, peer handle exchange, cached compiles --
totals under 2 s. Multicast is 0.25 s across all eight of its collective stages, which is
reassuring given it is the newest and most collective-heavy part.
The timing costs nothing when TL_STAGE_TRACE is unset; perf_counter is not even called.
…g the rest --mc-tiles was tuned at 240 MiB and left fixed, which cost small buffers badly: at a 3 MB shard 32 tiles per CTA leaves the publish only 48 CTAs, and allreduce at 48 MiB measured 0.361 ms against a 0.13 ms fabric floor -- 0.87x torch. Scaling it to one tile per 480 KB of shard takes that to 0.282, and reduce_scatter at 48 MiB from 1.47x to ~2.0x. Allgather is the exception and now overrides it to 32: its intra-node half is two multicast publishes and no reduce, so it is switch-bound rather than occupancy-bound and prefers few fat CTAs at every size (48 MiB 0.164 ms at 32 against 0.233 at 4). Four of the five collectives now beat torch at 48, 120 and 240 MiB. Allreduce is at parity and I am not claiming better: it straddles 1.0x, being the only one with two fabric hops against the same floor. Two findings worth more than the numbers. Chasing a closed-form heuristic kept regressing something -- a 240 KB/tile divisor beat 480 KB at 48 MiB and lost at 120 (0.747 vs 0.521) -- because the optimum depends on the collective as well as the size, and --rail-groups is coupled to --gin-contexts by the grid-divisibility rule. So the constants in the tree are labelled as fits with their measurements, not as laws, and both knobs stay exposed. And the measurements themselves are noisy enough to swamp several of the decisions: with other tenants on the same NICs, allreduce at 240 MiB spans 0.944-1.049 ms at one config while torch moves 0.950-1.021 over the same runs. Timing torch immediately either side of every run is what keeps the ratio meaningful (0.96-1.01x across all six). Any difference under ~5% needs repetition. This is the case for a multi-process autotuner rather than more hand-fitting; design notes are being written up separately. Also removes the never-exercised G=1 crash in the fused allreduce path, and adds a single-launch mc_reduce (slots="all") for when there is no fabric transfer to hide behind.
pick_mc_tiles' result was written back into the shared argparse namespace. That is wrong for any caller that builds more than one collective -- allreduce builds both halves, and an autotuner would build one per candidate -- because the second inherits the first's rewrite and reports a config it did not run. Several knobs are rewritten from the requested value (pick_mc_tiles, workable_chunks, the MIN_GROUP_BYTES cap), so this would have made a tuning table lie. Resolved onto self.mc_tiles instead, and the multimem divisibility check now reads that rather than the raw argument, which in auto mode is 0 and would have divided by zero. Verified on the 8-GPU proxy: allgather, reduce_scatter and allreduce all correct. Also corrects an overstatement in workable_chunks. A collective compile *can* report a lowering failure -- _maybe_compile_once ships the root's traceback through its all_gather_object -- so the reason to probe with a plain tilelang.compile is not that ctx.compile would hang on a clean failure. It is that each node runs a different interpreter and NCCL, so "every rank rejects the same put size" is an assumption about two toolchains rather than a guarantee. Adds docs/distributed_autotune_design.md: why the knobs cannot be hand-fitted (the optimum depends on collective and size, and --rail-groups is coupled to --gin-contexts by the grid-divisibility rule), why TileLang's single-process AutoTuner does not transfer (its thread pool, per-candidate SIGALRM, skip-on-error path and rank-local winner are each a 16-rank deadlock or a wrong answer), and a lockstep in-process sweep protocol that treats measurement as noisy -- which it is: one allreduce config spans 0.944-1.049 ms while torch moves 0.950-1.021 over the same runs.
Two results from an idle-pair sweep: one genuine reversal, one retraction. The AG-GEMM mode ordering inverted. When the collective was slower, pipelining the fabric hop under our own node's GEMM won by 12% (0.430 ms against serial's 0.481). Rail-group pipelining then made the collective ~50% faster, leaving far less fabric to hide, and the pipeline's extra barrier and split GEMM launches no longer pay: serial 0.373 ms / 736.8 TF against pipeline 0.441 / 622.9, so 1.34x torch against 1.13x. serial is the default again, and the docstring says to re-measure whenever the collective changes -- the answer is a function of the comm/compute balance rather than a property of the kernel. The 32 MB deficit against triton-dist (158.2 vs 247.9 GB/s) was contention, not configuration. I had two plausible explanations -- a 2 MB shard tripping MIN_GROUP_BYTES down to one rail group, and mc_tiles=32 starving the publish at 32 CTAs -- and both were wrong. On an idle pair the default measures 0.154 ms / 203.7 GB/s / 1.66x torch, and every override is worse or equal: --rail-groups 2 gives 0.171, --mc-tiles 4 gives 0.167, both give 0.154. The 158.2 came from a session where torch simultaneously read 86.6 GB/s against 122 idle, so the comparison is unsettled rather than lost and a same-session re-run is queued. gemm_rs improved to 1.17x (0.283 ms / 121.4 TF) purely from the faster collective, with --block-k 32 and --gemm-block-n 128 both inside noise (0.279, 0.277), so its defaults stand. triton-dist's fused AG-GEMM still does not complete here, but the failure is now a specific nvshmem4py teardown error -- "Buffer ... freed implicitly" -- not the illegal memory access previously recorded.
…nd say why Triton-distributed's swizzle_tiled_m_with_padding renumbers GEMM tiles so a rank computes the block belonging to rank r+1 first and its own last, so the transfers that must happen start earliest. Adapted here to node-slot granularity -- compute the other node's rows, reduce and send them, compute our own node's rows while they fly -- because their per-rank rotation does not transfer: multimem.ld_reduce needs every local rank to have written a segment, so rotating per rank leaves each rank ready on a segment its siblings are not. It is correct on all 16 ranks and ~10% slower: at k-per-rank 2048, 0.365 ms / 376 TF against serial's 0.334 / 411; at 4096, 0.476 / 578 against 0.429 / 640. serial stays the default. That is despite the shape being right for overlap -- comm is ~0.237 ms and near-fixed since the reduce-scattered output is M*N whatever K is, while the GEMM grows from 0.031 ms at K/rank 512 to ~0.248 at 4096. The two dist.barrier calls (~30-50 us each) plus splitting one GEMM into two half-height launches (worse wave quantisation for the persistent tcgen05 grid) cost more than the fabric time hidden. This is the third overlap attempt to lose the same way, after the merged allreduce (1.296 vs 1.044) and the ag_gemm pipeline once the collective got faster (0.441 vs 0.373). The one overlap that did work -- rail-group pipelining, 403->471 GB/s -- adds no host barrier at all, enforcing order with a device-side GIN signal per group. So the rule is that overlap pays along axes the device can synchronise and loses along axes needing the host to referee, which promotes the deferred *_barrier_gpu work from tidy-up to prerequisite. All three attempts are kept behind flags so they can be re-measured in one line each when it lands. Also records the arithmetic-density curves: gemm_rs 128 TF at K/rank 512 is a comm-bound shape artefact, not a slow GEMM -- the same kernel reads 411 TF at 2048 and 640 at 4096, and ag_gemm saturates near 876 TF at M/rank 2048.
Four things, none of which change behaviour. Dropped rail_sum_kernel. The reduce-scatter tail was replaced by rs_sum_kernel when the per-group pipeline landed -- that version reads `partial` as a term instead of copying it into the inbox first, and carries no wait_signal so its grid is free of the sender's chunk count -- leaving rail_sum_kernel reachable only from the lowering-bug repro. debug_put_size_lowering.py now contrasts the put against rs_sum_kernel, which is live, so the contrast means something. Committed examples/distributed/nccl_gin_internode/, which had been sitting untracked even though testing/.../test_nccl_gin_put.py already points readers at it. A committed test referring to an uncommitted example is a trap for the next person. Moved tests/test_distributed_multinode.py to testing/python/distributed/. It is a real 372-line unit test for NodeTopology and BaseAllocator, mock-based and needing no GPU, and it was invisible to the suite in a top-level directory the repo does not otherwise use. Deleted five superseded scratch files: PHASE1_SUMMARY.md, PHASE2_STATUS.md and GIN_POSIX_FD_SUMMARY.md, whose content is in CLAUDE.md in corrected form, and check_gpu_availability.sh / check_all_nodes_gpu.sh, replaced by the gpu-window skill and its always-on listener.
It holds cluster hostnames, internal IPs and absolute paths from a private workspace. It was pulled in by a broad 'git add -A' and should never have been committed, let alone pushed to a public remote.
Rachmanino
force-pushed
the
wt/nccl-gin-internode
branch
from
August 5, 2026 07:20
6b523d2 to
69c55bd
Compare
…es worth keeping The flat kernels pushed every rank's shard to all world_size-1 peers over GIN, which is optimal with one GPU per node and collapses with eight: each rank puts 15 shards onto the NIC, seven of them for siblings on the same machine, and allgather measured ~2.7 GB/s against torch's ~309. The 2D versions reach 400-470. There is no configuration in which the flat ones are the right choice on this hardware, so keeping them as a "baseline" only invited someone to run them. Removed the five flat examples and sweep_internode.sh, which only knew how to drive them, plus tune_grid, report_tuning and per_launch_signals from internode_common -- the tuning helpers no 2D example used, since add_2d_args and bench_vs_torch replaced them. Kept the one thing only the flat AG-GEMM documented, promoted into internode_2d's trap list: never wait on a signal inside a large-grid kernel. The fused shape where some CTAs of a GEMM grid issue puts while the rest wait deadlocks once the grid stops being co-resident, because a waiting CTA holds an SM while the CTAs that would issue the puts are still queued behind it. Capping the grid at one CTA per SM fixes the hang and is slower than not overlapping. That is why the overlap here uses separate streams, or a device-side wait in a small-grid kernel like rail_wait_kernel. Note on method: the first attempt at removing those three helpers bounded each function by the next line starting with def/class/#, which silently swallowed the TL_DTYPES and TORCH_DTYPES dicts that followed one of them. py_compile did not catch it -- a missing imported name is a runtime failure. Redone using the AST's own line ranges, and there is now a checker that walks every ImportFrom against what each module actually defines.
… examples The reference opened by declaring multi-node execution out of scope, which stopped being true. It now covers the inter-node surface: - GIN as a requirements row and a kernel-side section: put, put_signal, signal, wait_signal, flush, with the three signal properties that decide how they are called -- signals are cumulative and unconsumed by a wait, signal state is per context so wait_signal divides by the device-side span, and the requested context count is only a hint. - init_dist's return_node_info form and the NodeTopology it yields, plus the multi-node launch variables and the warning that NCCL_IB_DISABLE defaults to 1 and must be cleared. - Why GIN needs a VMM allocator (window registration rejects cudaMalloc) and how the arena is shared without an IMEX channel, via a POSIX fd duplicated with pidfd_getfd. - Multicast no longer requires fabric handles, so the requirements table was wrong; plus the fd export/import pair, multicast_uses_fd, and that the multicast buffer is a separate allocation from the arena sized by a bump pointer with no free. - TILESCALE_USE_GIN / _NCCL_LIB / _GIN_CONTEXTS / _GIN_SIGNALS / _GIN_COUNTERS, noting the resource counts affect device memory but not start-up, since ncclDevCommCreate costs the same at 1, 4 and 8 contexts. - That both capability probes reach cuCtxGetDevice and so report False with no current context, which silently misreports a capable machine. - Two limitations callers meet in normal use: put_signal sizes that fail to lower, and barrier_blocks being single-node despite its docstring, because it is lowered with global ranks and takes get_remote_base_ptr of each participant. - The two multimem calling constraints: fp16/bf16 regions must be one contiguous pair per thread, and the multicast region must be provably in bounds at compile time. With those documented once, the examples no longer restate them. internode_2d's trap list points at the reference and keeps only the trap specific to it -- that splitting the fabric hop must preserve the per-put size -- and the kernel docstrings that retold API-level rules are cut to their contracts. The two fused examples' docstrings are compressed the same way, keeping every measurement and both negative results.
…bs do not compose Eleven configurations at 32 MB, the one size where triton-dist leads us, with torch steady at 121.9-123.1 GB/s throughout so they are comparable. --mc-tiles 16 alone takes 194.5 -> 216.9 GB/s, 1.78x torch, closing the gap to triton-dist's 234.3 from 0.87x to 0.93x. The tiles curve at this size is 183.7 / 210.4 / 216.9 / 194.5 for 2 / 8 / 16 / 32, so it peaks at 16 and the previous "allgather wants 32 at every size" was wrong -- that was inferred from 48 MB where only 4 and 32 had been compared. The default is now a threshold on shard size, fitted to three points and honest about being nothing finer. The more useful finding is negative: every combination of the individual winners is worse than the best single change. --chunks 4 gives 209.4 alone and 185.2 with --mc-tiles 8; adding --mc-threads 256 to a three-knob combination drops it to 183.6. That rules out coordinate descent for the planned autotuner, since each knob's optimum moves when another changes, so the search must sweep tuples -- materially more budget than the design notes assumed. Small-size settings also do not generalise upward: at 240 MB the combination gives 401.0 against the default's 460.3. docs/internode_optimization_log.md is new and carries this per kernel, with why ratios rather than milliseconds are reportable here, and the SM-specialisation idea from triton-dist that is the one untried mechanism for in-kernel overlap. Also fixes a false positive in the GPU-window guard, which matched any mention of an example filename and so blocked commits and code edits whose text contained one -- worse than a miss, since it blocked exactly the offline work the guard exists to encourage. It now requires an actual launcher or interpreter invocation.
Round 1 established that the knobs do not compose -- --chunks 4 measures 209.4 GB/s alone and 185.2 combined with --mc-tiles 8 -- so tuning has to sweep tuples, and doing that by hand one run at a time is slow and easy to get wrong. --sweep takes a cartesian spec: --sweep "mc_tiles=8,16,32;chunks=4,8;gin_contexts=1,2,4" over chunks, rail_groups, gin_contexts, mc_threads, mc_tiles, intra_chunks and threads, and prints a ranked table carrying each candidate's *effective* configuration after the internal rewrites (workable_chunks, the MIN_GROUP_BYTES cap, pick_mc_tiles) rather than what was asked for -- otherwise the table names a configuration that did not run. It runs in one process, which is the point: start-up is ~19 s per rank and 6.4 s of that is ncclDevCommCreate, which does not shrink, so a process per candidate spends nearly all its time initialising. Three constraints this had to respect rather than details: Buffers are allocated by the first candidate and reused. No buffer shape depends on a knob -- only grids and tile widths do -- and reuse is mandatory rather than merely faster, because the arena is a bump allocator with no free and the multicast buffer is sized exactly, so a second allocation exhausts it. GIN signal counters are now shared through a dict the driver owns. They live on the device and accumulate for the process lifetime, so a candidate restarting its count would have its wait satisfied by a previous candidate's arrivals and would silently measure nothing. This also removes the same latent hazard for two collectives in one process. torch is re-timed either side of every candidate and the drift printed, because its reading moves ~8% with other tenants -- larger than most of what is being compared, and why several earlier conclusions in the log were noise. Invalid tuples are skipped rather than fatal, since the validity rules are deterministic functions of the candidate and every rank therefore rejects the same ones, keeping the collective compiles in lockstep. Wired into allgather first. Not yet exercised on hardware: all four nodes are busy, and a single-node proxy cannot fully check a collective code path. Also fixes two false positives in the GPU-window guard, both of which blocked commits rather than launches: it matched any mention of an example filename, and then any ```bash fence sitting above a launcher name. Heredoc bodies are now stripped before matching, since a heredoc is data -- a doc, a patch, a plan file -- and not what the shell runs.
…s by position The sweep shipped in the previous commit ranked candidates by where they appeared in the list rather than by their configuration, and the first real run exposed it. Sweeping mc_tiles=8,16,32 declared 8 the winner at 222.0 GB/s. Reversing to 32,16,8 declared 32 the winner at 218.1. The control settles it: the *same* configuration measured three times degraded monotonically, 1.73x then 1.55x then 1.44x. That is clock droop under sustained load, roughly 20% across three candidates, which is far larger than anything being compared -- so a single-pass sweep confidently reports whichever candidate is listed first. It also means round 1's separate-process finding that mc_tiles=16 wins at 32 MB needs re-checking with this tool, since the in-process runs disagreed with it in both directions. Candidates are now measured round-robin over --sweep-passes passes (default 3), each keeping its best pass: round-robin spreads every candidate across the droop curve and the minimum takes its least-throttled sample. Correctness is still verified once per candidate. The control now reads 4.4% spread instead of 20%, and the report prints a "within the noise floor -- treat these as tied" note whenever the whole spread is under 5%, so a sweep that discriminates nothing says so rather than inventing a winner. Also fixes the multicast exhaustion the first run hit: the example builds an instance before sweeping in order to fill its inputs, so run_sweep now seeds from those buffers instead of letting candidate 1 allocate a second time against an exactly-sized multicast buffer.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.