feat(EPv2): [preview] internode dispatch/combine over CCO/GDA - #625
Merged
Merged
Conversation
…topk run_internode_test.sh drives the entry through torchrun, which puts the entry's own directory on sys.path and not the repo root. The v2 entry imports tests.python.ops.dispatch_combine_test_utils at module level for the shared test case, so that import could not resolve and the ci_cco.yml internode job failed with "No module named 'tests'" before reaching a kernel. Export PYTHONPATH from the script, which is what the ad-hoc two-node scripts had been doing by hand. The sys.path insert under the entry's __main__ guard cannot cover this: the import that fails runs at module import time, long before __main__. --topk replaces a hardcoded num_experts_per_token=8, so the internode bench can be swept at other topk values. Threaded through spawn alongside the harness's other options rather than read from argv in the worker, because the worker is a fresh interpreter. Co-authored-by: Cursor <cursoragent@cursor.com>
jhchouuu
force-pushed
the
dev/adpat_v1ops_to_cco_api
branch
from
September 1, 2026 03:33
3c605ae to
d8bf163
Compare
…rnel The cco internode kernel (ep_internode_kernel.hpp) is a separate copy of the shmem v1 kernel, so the #613 combine optimization landing in internode_v1.cpp via the main merge did not reach it. Mirror it here: - 16B load-first gather (CombineVecBytes=16 / CombineGather) with a per-launch CombineVecAligned() fallback to the 4B path, replacing WarpAccum<TokT,4> in CombineIntraNodeLLTyped and CombineInterNodeLLTyped. - Size the warp slice by the vector step instead of a fixed warpsPerToken=4: MultiWarpIter gets CombineVecStep granularity in the intra-LL path, and the inter-LL path sets warpsPerToken = hiddenDim / vecStep and snaps hiddenDimPerWarp to a whole number of steps, so every warp stays on the vector path instead of a scalar tail. Helpers are duplicated (not shared with internode_v1.cpp) to match the PR's deliberate two-copy split of the v1 kernel. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…1_ll The cco internode path takes its grid from dispatch_combine.py (the shmem resolve) via the launch redirect, which is not CU-aware: on MI308X (80 CU) it lands on 96/64/8 (MANUAL) or the InterNodeV1LL AUTO default 256/128/8 -- 1.2x and 3.2x the CU count, so the grid tail runs a second wave and small-token latency suffers. Add internode_tuning_configs.py: a per-device (PCI DID / CU count), per-shape, per-token table with an explicit rdma_block_num per phase, mirroring tuning_configs.py but looked up host-side with the live num_tokens (the redirect has no runtime schedule). lookup() clamps block_num <= CU count. The redirect's dispatch/combine wrapper pins the result when the caller left geometry unset; an explicit block_num still wins. MI308X EP16 h6144 (fp8 dispatch / bf16 combine), tuned 2-node on skyriver07+04, vs shmem-main (pr627) v1_ll, us disp/comb: tok4 shmem 47.4/58.4 -> cco 46.1/62.8 (disp -3%, comb structural gap) tok8 shmem 47.1/59.1 -> cco 47.3/61.6 tok16 shmem 49.3/62.5 -> cco 46.4/55.3 (disp -6%, comb -12%) tok32 shmem 56.1/78.9 -> cco 51.0/66.0 (disp -9%, comb -16%) dispatch beats shmem at every count; combine wins big at 16/32. The tok4/8 combine gap is the 4-kernel combine launch overhead, not geometry -- tracked separately. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…oupling The tok4/8 dispatch geometry (32/16/6, 64/16/4) was tuned for dispatch latency in isolation, but dispatch and combine are coupled: they share the CUDA-graph replay and the QPs, and a dispatch with rdma_block_num=16 leaves the combine that follows it ~18us slower at 4/8 tokens than rdma=32 does (combine ~77us vs ~56us, measured 2-node clean on skyriver07+04). The isolated dispatch win (~46 vs ~50us) is dwarfed by that combine loss. Hold small-token dispatch at 64/32/8. tok8 then measures d~50/c~56 (combine now beats shmem-main's ~59; dispatch ~ties), where 64/16/4 gave d~46/c~76. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Clean fp8-dispatch/bf16-combine sweep on skyriver07+04: tok8 combine at 32/21/6 measures Best 47 / Avg 58.7 vs 64/32/4's Best 53 / Avg 63.5 -- the earlier 64/32/4 was a contention-contaminated pick. Small blocks win at small token counts for combine, same as tok4. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The internode combine is 4 kernels and dispatch is 2, each launched as its own plan.launch across the ctypes ABI. At small token counts that per-pass host cost is a large fraction of the (short) GPU work. Add mori_jit_plan_launch_multi: one C ABI call takes N plan handles + a single shared arg buffer (all internode passes publish the same EpInterNodeCcoArgs schema and, in the redirect, get the same raw/devComm), loops vt->launch in C++, and does the N hipModuleLaunchKernel back to back. Python launch_multi() fills the arg struct once (reusing plans[0]'s cache) and hands the whole sequence over. The internode test redirect uses it on the fast path; MORI_INTERNODE_BATCH_LAUNCH (default on) falls back to per-pass launch for the A/B control, and the trace paths keep per-pass launch for fault attribution. Measured on the 2-node MI308X rig (fp8-dispatch/bf16-combine, hidden 6144, EP16): at small tokens where host overhead is exposed this shaves ~5-10us off both dispatch and combine (e.g. tok4 combine 73->64us); at tok32 the kernels are long enough to hide it and the two converge. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Three changes to the internode bench/tuning harness, all aimed at getting a stable, representative number out of a single 2-node run. Warmup/rounds defaults raised (MORI_EP_WARMUP 3->20, MORI_EP_ROUNDS 10->30, both still overridable). The CCO/GDA low-latency kernels reach steady state far more slowly than the shmem/IBGDA ones: at 10 measured rounds with 3 warmup their mean swings ~20% run to run and reads biased high, while shmem is already stable by round 3. 20 warmup + 30 timed rounds averages the intrinsic per-round fabric jitter down to a few percent so a single bench reproduces. Variance-aware tuning winner (MORI_EP_TUNING_TAIL, off by default). _build_phase_stats now also reports lat_std (round-to-round std of the per-round mean latency), and with the flag on the sweep picks, among configs within MORI_EP_TUNING_TAIL_TOL of the best mean, the one with the smallest lat_std -- the most reproducible of the near-best configs. Selecting on the whole distribution's std rather than a single worst sample is what makes this robust; the per-candidate summary now prints the std too. It is only meaningful with enough rounds (raise MORI_EP_ROUNDS) and, ideally, several tuning passes. Per-round barrier diagnostic (MORI_EP_PERROUND_SYNC, off by default). Syncs + barriers after each timed round to re-align ranks, isolating how much of the Best/Worst spread is cross-round drift vs the kernel; combine timing stays clean, dispatch does not (documented). Diagnostic only, not for real numbers. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…etry Full-scope sweep x3 + A/B validation on the same 2-node rig. The 4/8/32 rows held up -- the sweep's independent per-phase argmins do not reproduce once dispatch and combine run at different geometries (the dispatch<->combine coupling the module docstring describes), so e.g. tok8 stays on the current 32/21/6 combine, which A/B-beat the sweep's alternative. Only tok16 moved: one shared 80/rdma40/warp4 geometry for both phases beat the old 80/48/8 + 80/40/8 by ~4us total (disp 41.3 vs 43.6, comb 51.1 vs 53.1), reproducible across two A/B batches. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Add MORI_EP_DROP_ROUNDS to control how many leading timed rounds are discarded from the stats; default 1 keeps the historical behaviour (drop round 0 only) and is symmetric across backends so cco and shmem are measured identically. There is a real start transient worth knowing about: the dist.barrier() before the timed loop releases all ranks at once, so the first timed rounds' collectives fire simultaneously and hit peak fabric contention (thundering herd) before the rounds self-stagger. It is not a warmup/cache effect -- the 20 warmup rounds warm the kernels -- and both backends show it, but the CCO/GDA path is far more sensitive (first round ~1.6-2.5x steady, still elevated through round ~2) than shmem (recovered by round 1). Raising MORI_EP_DROP_ROUNDS (e.g. 3) excludes that ramp from the reported Best/Worst; it is left at 1 by default so both backends drop the same and the number matches the historical one. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Sync main + port #613/#627 combine opt to CCO v2, and internode perf/bench improvements Brings the CCO/GDA internode v1_ll path up to the shmem-main combine optimization and adds v2-specific launch, geometry, and bench improvements. * Merge origin/main and port the #613/#627 combine kernel optimization (kCombineStepsPerWarpSlice, WarpAccum, re-tuned MI308X schedule) into the CCO v2 internode kernel. * CU-bounded per-token launch geometry for the CCO internode path (block_num <= CU count, per-(world,hidden,topk) per-token schedule, with the dispatch<->combine coupling honored; tok16 re-tuned to a tighter shared geometry). * Batch the internode JIT launch through one ABI crossing (mori_jit_plan_launch_multi): the multi-pass combine / dispatch cross the ctypes ABI once instead of per-pass. Default on; MORI_INTERNODE_BATCH_LAUNCH=0 restores per-pass for A/B. * Reproducible internode bench defaults: higher warmup / measured rounds and a configurable leading-round drop (MORI_EP_DROP_ROUNDS), since the CCO/GDA LL kernels reach steady state slower than shmem. Plus variance-aware tuning winner selection (MORI_EP_TUNING_TAIL) and a per-round-barrier diagnostic (MORI_EP_PERROUND_SYNC).
The pre-commit black hook reformats two multi-target assignments in the v2 internode dispatch/combine test that exceed the 88-column default. Formatting only; no behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Formatting only, plus dropping a dead `da, ca` unpack that ruff flags as F841 (leftover from an earlier print format). No behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
jhchouuu
added a commit
that referenced
this pull request
Sep 8, 2026
…ransport's The missing arm existed after all. #625's entry is not a v2 op -- it keeps v1's op and redirects EpDispatchCombineOp._launch_multi to the v2 CCO plans -- so it is the same transport and effectively the same kernels as this branch, driven by v1's host code. That is the control that separates transport from our changes, and it runs on this rig. Six interleaved pairs at 4 tokens, max/mean: dispatch combine shmem / IBGDA 1.16 1.23 CCO/GDA, #625 base 1.27-1.31 1.46-1.56 CCO/GDA, this branch 2.31 1.75 COMBINE's tail is mostly the transport's (+27% CCO over shmem, +12% us on top). DISPATCH's tail is mostly OURS (+13% CCO, +76% us). The baseline never exceeded 1.44 in six runs with zero spiked rounds; this branch reaches 2.1-3.3 in five of six. This retracts the earlier "the tail belongs to the CCO path" -- that compared two transports and two kernels at once. This branch is still FASTER in the mean over the same pairs, dispatch -7% and combine -14%, so it trades mean for tail rather than simply regressing. Ruled out for the dispatch tail, all by direct check rather than argument: kernel logic (normalised diff over the 27 shared functions), arena region packing (per-region sizes match v1's, so the contended cache lines do), counter zeroing (present in both), and launch geometry (running ours at v1's 16/10/4 gives 1.67/2.31/1.21 against 1.22/1.22/2.00 -- no separation). Not resolved. The tail is bimodal in itself -- the same geometry gave 1.22 and 2.31 the same afternoon -- which is why these are six interleaved pairs and why short comparisons cannot separate mechanisms here. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
jhchouuu
added a commit
that referenced
this pull request
Sep 9, 2026
The 4-token internode bench reports a worst/mean of ~2.3 on some runs and ~1.2 on others with the same binary. The per-round series says why: a run lands in one of two regimes and holds it -- 37/45us or 50-70/72-79us dispatch/combine -- stepping between them over a single round 4-8 rounds into the timed loop, if it steps at all. The reported ratio is the mean blending two levels. Nothing here is a few bad rounds, which is the shape the previous notes assumed. Interleaved against #625 (v1 host, same v2 CCO kernels, same wire), both step on the same runs, so the trigger is environmental -- but ours degrades ~3x further (combine +60% vs +22%, dispatch +32% vs +2%). That sensitivity is ours; the trigger is not. Eliminated by direct measurement, all newly: the caching allocator (0 segments, 0 device allocs, 0 retries across the loop), launch-queue depth (a bare per-round drain does not remove the step), warm-up length (200 rounds steps where 20 does), RoCE congestion and loss (ECN marks, CNPs, out-of-sequence, retransmits, discards, out-of-buffer all +0), PFC pause (pri5 frames and transitions +0 on both nodes, slow runs included), and GPU clocks (sclk within 4%, fclk higher if anything on the slow runs). The per-pass split now prints on every rank -- the tail lives on whichever rank stalls, so a rank-0-only breakdown was reporting a quiet rank's median -- with the round series and the three worst rounds against the median. It puts the entire dispatch excess in dispatch_ll (+190us against a flat 7.0us copystaging) and the entire combine excess in combinesyncbarrier (+73us against a flat 29.4us combine_ll). On both legs it is the pass that waits, never the pass that works: combine's data movement has no tail of its own. Also adds, all behind MORI_EP_ROUND_SERIES: host time inside each call, host wall per round, the convert series and the loop's epoch bounds. These settle that the host is not the pacer -- 50-60us host hiccups leave the phase times untouched, and over 5000 rounds wall and the GPU sum agree to 0.2% -- and they close the per-round accounting so a hole can be located rather than inferred. --per-round-drain is added because --per-round-sync's gloo barrier costs ~1.3ms a round and injects more skew than it removes. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
jhchouuu
added a commit
that referenced
this pull request
Sep 9, 2026
BindCallingThreadToGpuNumaOnce() has existed since the shmem NUMA-affinity work,
and src/shmem/init.cpp:746 was its ONLY caller -- described there as "the single
bind site for the shmem/EP path". A job that used CCO instead of shmem, which is
every EP v2 job, ran completely unbound. ccoCommCreateImpl now calls it, at the
top, before the bootstrap and Context and any worker thread, since new threads
inherit the affinity. The caller has already hipSetDevice()'d -- the same
contract step 2 relies on when it caches comm->hipDev.
This is the cause of the bistable slow regime in docs/EP_INTERNODE_V2_TAIL.md.
Unbound, the scheduler can place two of the eight per-node ranks on the two SMT
siblings of one physical core (2 sockets x 96 cores, sibling of c is c+192), and
both then run at about half speed. Measured as an exact 2x on the host loop of
the affected ranks -- hdis 16 -> 30us, hcom 27 -> 50us -- and no change on the
others. That ~36us matches the measured 35-40us phase offset, and reaches the
peer through the 1:1 host-delay-to-peer-wait transfer already documented: the
late ranks stall their node's intra-node barrier, and the peer node then waits at
the cross-node rendezvous.
Eight runs each at 4 tokens, MORI_IGNORE_CPU_AFFINITY=1 as the control:
unbound bound
median total 88.0us 84.8us
worst total 126.0us 97.4us
runs with a doubled rank 1/4 0/8
It also improves the aligned case (median hwal 74.5 -> 72.8), which is the
cross-socket cost going away.
Consequence to act on: every EP v2 measurement against v1/shmem so far was made
across this difference, because the v1 harness calls
shmem_torch_process_group_init and was bound while the v2 path was not. The
"v2 degrades ~3x further than #625" result needs re-running before it means
anything.
Eight clean runs do not prove a zero residual rate -- the bind confines four
ranks to one socket's 192 CPUs, so a collision is unlikely, not impossible.
Internode correctness 30/30, 0 of 16 ranks disagreeing.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
jhchouuu
added a commit
that referenced
this pull request
Sep 9, 2026
The earlier "we degrade ~3x further than #625" result was measured with #625 bound (its harness calls shmem_torch_process_group_init, which was the only caller of BindCallingThreadToGpuNumaOnce) and this branch unbound. Withdrawn. Eight interleaved pairs after the CCO bind fix, ours matched to #625's environment (OMP_NUM_THREADS=4, the same 30-round loop): #625 median ours median paired median ours wins dispatch 39.52 37.85 -2.06us (-5.2%) 6/8 combine 53.41 47.05 -5.28us (-9.9%) 7/8 TOTAL 92.52 85.70 -7.21us (-7.8%) 7/8 spread #625 max/med 1.21 ours max/med 1.28 Faster on both legs, and the spread is comparable rather than 2x worse. Honest limits: #625's harness reports per-rank means only, so this compares means; two of its own eight runs are elevated, so it is not immune either. n=8. One of our eight was elevated (110.1) with NO host doubling and max per-rank hwal 86, so a smaller residual mechanism survives the bind -- it is not the 2x SMT one and is not yet characterised. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…646) feat(ep/v2): make the internode dispatch/combine a v2-native op The internode path reached its kernels through v1's op: v1 built the buffers and the arguments, and a _launch_multi redirect handed them to the v2 CCO kernels. That works, but it means the v2 op does not own its internode path -- the buffer layout, the geometry table and the launch sequence all live on the other side, and every v2 change has to be mirrored there. This makes the path v2-native and drops the coupling to zero: no v1 op, no v1 buffers, no redirect. WHAT REPLACES WHAT SymmArena One registered CCO window carved into 17 named, 256B-aligned regions (internode_regions.py), instead of a set of separately registered symmetric buffers. Region offsets are bound once on the plans; the launch path carries six arguments, not the whole schema. ccoDevComm Passed BY VALUE inside the argument struct. mori-shmem needed a device global filled by the host after every hipModuleLoad; this is the point of the CCO port and it is what lets a plan be launched without touching device memory first. Geometry Compiled per (block, rdma, warp) and resolved from a per-device table at launch (internode_tuning_configs.py). rdma_block_num splits the grid between the network blocks and the intra-node ones and the kernel branches on it, so it cannot be chosen per launch -- every bucket the table can name is built at construction, which is also why a geometry can never trigger a compile inside a timed region. LaunchGroup One ABI crossing per phase for the two- and four-kernel pass sequences, rather than one per kernel. The kernels themselves are shared with the v1 path and are not rewritten here. Two guards were added because CCO is not shmem: the local-node peer is skipped in the send loops (shmem dispatches per peer on transportTypes[pe] and uses P2P for a same-node peer; ccoGda has no such dispatch and RAIL leaves a local peer without a QP), and the one foldable device assert became a static_assert. Correctness is analytic rather than differential -- an identity expert, so combine[t] == U[t] * input[t] with U the number of distinct destination ranks -- which is the property that survives the kernel being rewritten underneath it. Green for bf16 and fp8->bf16 at 4/8/16/32 tokens, 0 of 16 ranks disagreeing. TWO FIXES THAT ARE NOT ABOUT EP cco_init.cpp ccoCommCreateImpl now calls BindCallingThreadToGpuNumaOnce(), and that helper now gives each rank its own PHYSICAL cores rather than the whole NUMA node. The helper already existed; src/shmem/init.cpp was its only caller, so every job that used CCO instead of shmem ran completely unbound. Unbound, two of the eight per-node ranks can land on the two SMT siblings of one core; measured with sched_getcpu, every colliding pair -- and only the colliding pair -- ran its host loop at exactly 2x, and each such run dragged the whole 16-rank collective from ~89us to 116-153us. Partitioning is by the kernel's own thread_siblings_list, since cpu 5 and cpu 197 are one core and a contiguous slice of the cpu list would separate them. Eight runs after: no collision, no doubled rank. This affects every CCO job, not just EP. plan_api.py The launch buffer persists between launches and the code already relied on that for bound ints; this extends it to the per-launch arguments, which every EP caller passes as .data_ptr() results and repeats every round. Not for blobs -- they cross by value and an unchanged address says nothing about the bytes. 14 _set_arg calls a round become 2, host wall 77 -> 73us, and in the aligned case that passes through to total time about 1:1. MEASUREMENT Against #625 (v1 host, same v2 kernels, same wire) with both sides NUMA-bound, eight interleaved pairs at 4 tokens: dispatch -5.2%, combine -9.9%, total -7.8% paired median, 7 of 8 pairs won, and the run-to-run spread comparable rather than worse. An earlier reading that this branch was ~3x more sensitive to the slow regime is withdrawn: it was measured with #625 bound and this branch not, because the v1 harness calls shmem_torch_process_group_init and was therefore bound all along. The internode tuning table was re-derived after the bind, since every A/B in it had raced a +-40us random term from CPU placement. Four stages -- sweep, keep only what reproduces across repeats, head-to-head against the shipped row, then a plain bench A/B -- and each of the last two rejected something the earlier ones had accepted. One row moved: 32-token combine 80/40/8 -> 64/48/6. docs/EP_INTERNODE_V2_TAIL.md is the writeup: what the "tail" actually is, the full list of what was eliminated and how, and the residual that survives (a ~10-20us offset in roughly 3 runs in 8, with no host doubling, mechanism not identified). TEST HARNESS tests/.../test_dispatch_combine_v2_internode.py gains --cmd test/bench/tuning and now reports what the v1 bench reports -- the per-round per-rank dump and the two performance tables, same formulas -- so the two can be read side by side. Defaults match v1 where they affect the comparison (scale_dim 32, experts_per_rank 256 // world_size, dist.barrier() at the measurement boundary); the alignment table in its docstring is machine-checked for the loop constants rather than asserted in prose.
…625 Review pass over the whole PR. Three broken CI steps, two real harness bugs, ~40 pieces of dead code, and a large amount of comment that recorded how the work was done rather than what the code does. Fixes - ci_cco.yml: the single-host internode step could never run -- it invoked `pytest <file>::test_internode_small_shapes`, and that file has no pytest functions at all; the v2 entry also refuses single-node outright ("world_size=8 with 8 GPUs per node is a single node"). Removed; the two-node job covers this path. - ci_cco.yml: the stress step passed `--cmd stress`, which the v2 entry's argparse rejects (test|bench|tuning). Now `--cmd bench --rounds 2000`, and run_internode_test.sh gained the `--rounds` pass-through it needs to forward it. - The non-greedy tuning summary read fixed_wins[0][1..2] where the tuple is (rank_key, metric, cand, cand_med, ...), so `--cmd tuning` crashed formatting its own result after doing all the work. - The correctness loop generated scales and then dispatched with None, which transports uninitialised staging whenever scale_dim > 0 -- the case _gen_round's own comment warns about, and the default is now 32. Dead code - EpInterNodeArgs 39 -> 37 fields. combineGridBarrier was only ever written (a store to a slot nothing reads); dispTokIdToSrcTokIdLocal was bound to 0 by the only producer, so both `== nullptr` guards were permanently true. - 14 profiler hooks that read args.profilerConfig, which EpInterNodeArgs does not have -- inert on the JIT path and would not compile if enabled. - The DEDUP=false half of DispatchInterNodeSend; the template has exactly one instantiation. Template parameter collapsed. - Dispatch launches reserved dynamic LDS via EpInterNodeDispatchSharedBytes, but no dispatch body declares any: the three dispatch kernels emit zero ds_ instructions. Now 0, helper removed. Combine keeps its reservation. - buildFlushDbrVal and flushAsyncImpl's qpn parameter, both orphaned by the snapshot-only rewrite; SymmMemObj::ccoWin/ccoWinOffset, written once and read nowhere; plan_api.launch_multi; EpInterNodeRegion::IsValid/Get; MaxXferBytesPerToken; the SendBufSlotOffset helpers; an empty namespace block; unused locals and an unused include. - Debug knobs with no consumer: MORI_EP_DEBUG_GEOM, MORI_INTERNODE_TRACE_ARGS, MORI_EP_SPLIT_PASSES, MORI_EP_HOST_PROFILE, MORI_EP_INJECT_HOST_US. - Harness flags for hypotheses that were tested and settled: --pre-barriers, --pre-sleep-ms, --barrier-kind, --realign-every, --per-round-drain, --per-round-sync. One unconditional gloo barrier remains, as in v1. Comments - docs/EP_INTERNODE_V2_TAIL.md 467 -> 78 lines. It was an investigation log: a section it labelled SUPERSEDED, a recommendation resting on a result the same file withdrew, a NEXT section proposing work already done, and recipes for six scripts that are not in the repo and one tool that was deleted before merge. Kept the durable part: the bimodal shape, the SMT-collision cause, the two env knobs, how to recognise it, and the A/B method. - flushAsyncImpl's 25-line comment -> 8. The hazard survives (postIdx is a reservation counter, ringDoorbellOrdered waits on equality, so advancing dbTouchIdx here strands an owner); the three-way post-mortem is in the PR. - internode_tuning_configs.py's 96-line header -> 30, keeping the coupling invariant and the three-guard protocol, dropping four dated re-tune stories. - Wrong comments corrected throughout: counts that did not match the structs, line references left over from the v1 file, WarpAccumLF named where the code calls WarpAccum, "eight passes" where a round runs six, an orphaned comment for a field that no longer exists, a dependency justified by an include the ported kernel does not have. Verified: library builds; all eight JIT kernels compile for gfx942; 8-GPU HIP op test passes; 35 pytest cases pass; ctest failure set identical before and after (10 pre-existing GDA failures on this host, environmental); black, clang-format and shell/YAML parse clean, with clang-format deviations reduced rather than added.
`Tensor.scatter_` is (dim, index, value); this passed (index, dim, value), so _rdma_algo_token_count raised TypeError as soon as it was reached. It is only reached on the non-LL path -- the LL branch returns before it -- which is why a 32-token bench (LL) was fine and every `--kernel-type v1` bench died in _report_tables after the kernels had already produced their numbers. That is the shape the CI bench and stress steps run. Also drops the clamp_(0, nodes - 1) that came with it: v1's compute_rdma_algo_token_count, which this reimplements, does not clamp, and its inputs carry no -1 entries. A clamp would silently fold an out-of-range expert onto node 0 and inflate the reported bandwidth rather than fail. Verified on 2 nodes x 8 MI308X (EP16), through tools/run_internode_test.sh exactly as ci_cco.yml drives it: `--cmd bench --kernel-type v1` at 128 and 4096 tokens now completes with rank0=0 rank1=0, where before it exited 1.
…explicit
The two internode families were already separate kernels -- separate JIT
modules, entry symbols and cache keys, four bodies rather than two branches --
but nothing said so at the API, and which one ran was decided by a token-count
rule with a private `_internode_force_ll` attribute as the only override. This
makes the choice a first-class config field and drops the private hook.
internode_kernel: "auto" | "v2" | "v2_ll" (default "auto")
internode_ll_max_tokens: int (default 512)
"v2" and "v2_ll" compile ONLY that family. "auto" compiles both, because only
"auto" chooses per launch; it picks v2_ll at or below internode_ll_max_tokens.
An explicit choice cannot fall back at runtime -- the other family was never
built -- which is the point: a benchmark number is comparable to another
harness's only if the kernel behind it is named.
The saving is 8 -> 6 kernels, not half: copystaging, combinesync,
combinesyncbarrier and combineall are shared by both families, so only the two
variant-specific passes drop out.
The crossover moves 2048 -> 512 and becomes configurable; it was a private
class attribute on the backend.
Renaming: the four device bodies lose their misleading V1 spelling
(EpDispatchInterNodeV1Kernel_body -> EpDispatchInterNodeV2_body, and the
LowLatency pair -> ...V2LL_body), as do the comments that called these kernels
"the v1 internode kernels". Comments that REFER to the v1 implementation for
comparison or provenance are left alone -- they are about the other tree, not
about this one. Entry tags (dispatch/dispatch_ll/...) are pass names, not
variant names, and do not change; the entry symbols are unchanged.
`--kernel-type` on the internode harness becomes auto|v2|v2_ll and now drives
the config instead of poking a private attribute; ci_cco.yml follows. The BENCH
line reports the family that RAN, with "(auto)" when it was not requested by
name, since under "auto" the request does not identify a kernel.
Verified on 2 nodes x 8 MI308X (EP16), ROCm 7.14:
- correctness, all three modes at 128 tokens: PASS
- auto flips at the crossover: 512 -> v2_ll, 1024 -> v2
- explicit overrides the shape: v2_ll at 1024, v2 at 128
- compile set read from a private JIT cache per run, not inferred:
auto -> 8 kernels (both dispatch/dispatch_ll and combine/combine_ll)
v2 -> 6 (no dispatch_ll, no combine_ll)
v2_ll -> 6 (no dispatch, no combine)
- full ci_cco.yml two-node matrix with --kernel-type v2: 4/4 green
- all 8 renamed kernels compile for gfx942; 35 pytest cases pass
A full adversarial review of the PR as it would land on main (28 files, +6480/-168) confirmed 35 findings. These are the five that are wrong rather than merely untidy. Four are pre-existing in #625; one I introduced. 1. DATA CORRUPTION, reproduced and fixed. DispatchInterNodeRecv indexed its 8-way sub-block from blockId, while its own combine twin uses bid. bid strides by rdmaBlockNum, so the two agree only when rdmaBlockNum % 8 == 0. Otherwise the eight blocks cooperating on a chunk collide on some sub-indices and never issue others: tokens silently dropped, others delivered twice. Latent only because every dispatch row in the shipped table happens to be a multiple of 8 -- the combine column already carries 21, which is not, and that side was already correct. A/B on 2x8 MI308X with MORI_EP_DISP_GEOM=64,21,8 at 4096 tokens (the loop must stride more than once, so 128 tokens does NOT reproduce it): pre-fix FAIL 3/3, 16 of 16 ranks disagree fixed PASS 3/3 2. SILENT HANG. rdma_block_num is defaulted (64) independently of block_num, and the only rdma < block clamp lived on the tuning-table hit path. The untuned shape and the MORI_EP_*_GEOM pin fed raw values through, so a caller pinning dispatch_block_num=32 got grid=32 with rdmaBlockNum=64: every block takes the RDMA leg, the intra-node half gets none, and the dispatch fan-in waits for 512 arrivals that only 256 warps can make. All 16 ranks spin forever with no host-side error. Clamped at a single choke point covering all three paths, plus a strict rdmaBlockNum < blockNum check in MakeEpInterNodeCfg so a bare C++ caller cannot build it either (== is equally broken: xgmiBlockNum == 0). 3. ACCEPTED CONFIG, CORRUPT DATA. #625 unlocked quant_type='fp8_direct_cast' on the internode path by making the scatter rule intranode-only. That path's fp8 combine staging does not produce correct tokens -- measured here, weights exact but ~51% of hidden elements outside a 30% tolerance, some exactly zero -- and no test anywhere runs it. Rejected explicitly at config time; the intranode path, where the feature is real and tested, is untouched. The half-built device path is left in place (if constexpr, so it costs nothing when the config cannot select it) rather than ripped out of the middle of the combine kernel. 4. reset() restored cross_device_flag to a literal 1, which is the intranode backend's epoch seed. The internode backend seeds it to 0, so reset() left that path disagreeing with its peers. Now keyed off is_internode. Pre-existing on main; only reachable on the internode path this PR adds. 5. MINE. Making an explicit --kernel-type compile only its own family meant CI, which names v2 everywhere, stopped compiling v2_ll at all -- previously both were built even though only one ran. ci_cco.yml now covers v2_ll as well. Also: run_internode_test.sh defaulted --kernel-type to "v1" and always forwarded it, so the v2 entry died in argparse whenever a caller omitted it. It is now forwarded only when given, and each entry keeps its own default and its own vocabulary. Verified on 2 nodes x 8 MI308X (EP16), ROCm 7.14: correctness for v2, v2_ll and auto at 128 tokens and for v2 at 4096; the ci_cco.yml matrix 5/5 green; the four host-side guards checked directly (internode quant rejected, intranode quant unchanged, a pinned block_num=32 now yielding (32,31,8) instead of (32,64,8), and reset() filling 0 vs 1 per path); 35 pytest cases; black, ruff, shell and YAML clean. The remaining 30 findings are not addressed here: resource leaks (the per-op DevCommHandle, arena on a construction exception), a raw indices pointer held across dispatch/combine without a reference, recv_scales() stride, and a set of coverage gaps. They want their own change rather than more weight on this one.
The previous commit rejected quant_type on the internode path because its fp8 combine staging returns wrong tokens. That left the device code unreachable, so remove it rather than leave a half-built feature that reads as a working one -- which is the same argument the rest of this cleanup runs on. I argued against this an hour ago on the grounds that if constexpr costs nothing and that cutting branches out of the middle of a combine kernel is risky. The first is true but irrelevant: the cost of unreachable code here is that it misleads the reader, not that it runs. The second turned out to be wrong once measured -- the code is regular, not tangled: five identical nine-line if-constexpr prologues that fall through to the unquantised call, one if/else in copystaging, and one self-contained EpCombineAllInternalFp8. 113 lines. Removed with it: the now-unused QuantType alias, the backend gate that still ADVERTISED fp8_direct_cast as internode-supported, and the blockwise arm of EpInterNodeCombineSharedBytes (blockwise never reached this path either). quantType stays in EpInterNodeKernelCfg. It is part of the request schema and of the rendered NTTP, and taking it out is a schema change for a field that costs nothing to carry; the comment now says it is always None here. NOT touched, and deliberately: the intranode fp8_direct_cast path, where the feature is real, tested and documented; and the fp8 *transport* dtypes, which are a different thing from quant_type entirely. Verified on 2 nodes x 8 MI308X (EP16): v2 at 128 and 4096, v2_ll at 128, and asymmetric fp8_e4m3_fnuz->bf16 dispatch at 128, all PASS -- that last one is the check that fp8 transport still works. Intranode 8-GPU test_op.py PASS, and an intranode fp8_direct_cast config still constructs and still selects scatter. All 8 JIT kernels compile for gfx942; 17 pytest; black/ruff/clang-format clean (clang-format complaints on the kernel drop 132 -> 123, all pre-existing).
Second pass over the 35 confirmed review findings, now that we own the PR. Everything here is pre-existing in #625 unless marked otherwise. Lifetime and safety (host) - The per-op DevCommHandle was never released: _close_backend closed only the JIT plans, so every internode op leaked a device communicator and its QPs for the Communicator's lifetime. Released, and the cached host_ptr dropped with it. - dispatch() kept the caller's indices tensor as a bare data_ptr and combine() dereferenced it on the next call. Nothing held a reference, so a caller who freed the tensor between the phases got a silently wrong combine. The tensor is now held for exactly as long as the pointer is live. - recv_scales() viewed out_scales at the 128 B-padded intranode stride, but the internode kernel writes it PACKED at scaleDim * scaleTypeSize. Reachable on the shipped test shape (56 x 4 = 224 B viewed as 256 B). The view now follows the stride the kernel writes; the intranode path keeps its padding. - Construction was not exception-safe: any failure after SymmArena() leaked the arena and the dev comm, and the tuning sweep swallows exactly those exceptions in a loop. It now unwinds what it took. - An internode dtype with no kernel entry (fp4) died with AttributeError AFTER allocating the arena; it is now rejected by the constructor's early gate. - _make_dev_comm did not validate world_size against the communicator. - combine() ignored want_weights: it folded whenever the preceding dispatch had carried weights, so combine(x, None) paid for a fold it was not asked for and dispatch(None) + combine(t) returned an unwritten buffer as the weights. The argument is now honoured, and the one combination the kernel cannot express is rejected instead of returning a wrong buffer. Bounds (device and arena) - combine_out_weights was sized by MaxNumTokensToRecv but EpCombineAll indexes it by the LOCAL token id; once max_total_recv_tokens clamps recv below m, both the kernel and the host view ran past the region. - The send buffer is addressed pe * m + slot, but slots are handed out in whole wavefronts, so a full last chunk reaches ceil(m/WAVE)*WAVE - 1, past m whenever m is not a multiple of WAVE. Rather than change a stride v1 also computes, the internode capacity is rounded up so the two agree by construction. (Inherited from v1, which has the same expression.) - Two unguarded args.weightsBuf dereferences in the dispatch staging copy and the intra-node dispatch, on a pointer the backend explicitly passes as null. - DispatchInterNodeLLRecv derived destPe from a raw expert id without the -1 sentinel its non-LL twin applies. -1 / numExpertPerRank truncates toward zero, so an unfilled expert slot read as PE 0 and could false-match the dedup ballot. - nodeRecvTokenNum is a uint64 region; one of its four readers viewed it as int32, so for node != 0 it read the high dword of a neighbouring slot. Fixed to match the other three. NOTE: I could not construct a failing case for this one -- an A/B that restored the int32 read still passed under every routing shape I tried. It is corrected on inspection, not on a reproduction. Shared layers - ccoGdaOptFlagsAggregateRequests has no caller, and after flushAsync became snapshot-only nothing would ring the doorbell for WQEs posted under it. Rejected where it enters, with what would have to be built to support it. - An unrecognised NIC provider string silently compiled ccoGda<MLX5> against a non-Mellanox QP; a typo now fails loudly where the value is read. - LaunchGroup fills only the lead plan's arg buffer, so bound defaults on non-lead plans were silently discarded. The group now validates that its plans agree, naming the plan and the field when they do not. Coverage - New --routing {uniform,skewed,local}. Routing was uniform with topk well above the node count, so essentially every token reached every node and no chunk was ever empty -- the whole empty-chunk half of the protocol was unreachable from the tests. skewed sends one token in eight remotely; local sends none. - ci_cco.yml runs test_internode_regions.py, which no workflow collected. Verified on 2 nodes x 8 MI308X (EP16), ROCm 7.14: correctness for v2, v2_ll and auto at 128 and for v2 at 4096; skewed and local routing for both families; the ci_cco.yml two-node matrix 5/5; all 8 JIT kernels compile for gfx942; 35 pytest; black, ruff and YAML clean.
Two hooks were failing on this PR's files. ruff (E731): five lambda assignments in the internode harness became defs. The tuning selector needed the most care -- `pick` was a conditional expression choosing between two lambdas, which does not survive a mechanical rewrite, so the branch is now on the enclosing if/elif/else and each arm defines its own. Verified both affected paths on hardware rather than by inspection: the bench reporting exercises r/bw/gm (tables render, two header rows) and --cmd tuning exercises geoms/pick (reaches TUNING RESULT). clang-format: 47 lines across four files, all pure reflow -- argument lists rewrapped, one macro joined onto a line, a static_assert message rewrapped. Confirmed non-semantic by reading every hunk that survived `git diff -w`. The other seven C++ files the hook processed were already clean, so there is no unrelated formatting churn in this commit. `pre-commit run --files <the PR's 28 files>` now passes every hook. Library builds, all 8 JIT kernels compile for gfx942, two-node bench and tuning green.
A review pass aimed only at documentation accuracy and module structure confirmed 27 findings, 9 of them major. Almost all were prose that described the code as it was before this PR. The worst three said the opposite of what the PR does: - MORI_JIT_V2_DESIGN.md section 9: "cross-node 没有：v2 只有 intranode 两个 body" -- while the PR's largest new file is a 1563-line internode kernel with eight registered Specs. A reader met the internode launch sequence in section 6 and was told a hundred lines later that it does not exist. - The package README was titled "intranode MoE dispatch / combine" and said "Intranode only (no GDA/RDMA)". - EP_INTERNODE_V2_TAIL.md claimed the affinity bind gives every rank disjoint physical cores unconditionally. It does not: the split needs more than one node-local GPU VISIBLE to the process (a launcher that slices HIP_VISIBLE_DEVICES per rank makes that one) and at least one physical core per GPU, and otherwise falls back to the whole-node bind -- logged at WARN, which the default ERROR level suppresses. That doc exists to stop readers believing a cross-regime A/B, so a false guarantee there causes exactly the error it warns against. It now states the precondition and how to check on a given run. The rest were counts and names that had drifted: "十个入口" for eleven, "八个 裸指针字段" for fourteen of thirty-seven, EpCfg 11 for 12, EpArgs 22/8 for 24/12, "seventeen regions" for sixteen-plus-scales, "the two Plan classes" for ten, a `HipBackend` class that does not exist, a tuning table scoped to v2_ll that both families read, "v2 is intranode-only" in the routing handle, and a module docstring for hip_backend.py describing only the intranode half of a file the PR nearly doubled. The internode files were added to the design doc's file/tuning/test tables and to the README layout table, and the new config fields are now documented where a user reads rather than only inline. Two API traps found by the same pass, fixed in code: - gpu_per_node was documented as what selects the internode path, but only the hip backend implements it and the DEFAULT backend is flydsl, which has no reference to it anywhere. A 2-node config would silently build intranode kernels and return wrong results. The flydsl backend now rejects it through the existing _unsupported gate, naming kernel_backend='hip'. - max_token_type_size is caller-settable and sizes every internode staging region, and is compiled in as the transport stride. A value below the widest transported element wrote past the region rather than failing. Now rejected, with both leg sizes in the message. Also corrected: the internode test's own module docstring showed a torchrun line (--nproc_per_node=8) that main() refuses, since --spawn defaults to 8. Verified: every replacement count and name read out of the code rather than recalled; 35 pytest; black, ruff and the full pre-commit suite clean over all 28 files.
Three changes from reading the config surface. 1. A pinned geometry field is now a manual override. It was not: on a shape the internode tuning table knows, _internode_geometry_buckets_raw returned the table's triples and dropped the config's six geometry fields entirely, so pinning dispatch_block_num on a tuned shape did nothing and said nothing. The override is per FIELD rather than all-or-nothing, so one knob can be fixed while the other five stay tuned. That needed the caller's original input: _resolve_geometry() fills every unpinned field with the untuned fallback, after which pinned and defaulted are indistinguishable. __post_init__ now records the pinned set before that runs. The rdma < block clamp still applies after the overlay, so a pin cannot reintroduce the all-ranks hang. Pinning block=48 on a row whose tuned rdma is 48 yields (48, 47, 8), not (48, 48, 8). MORI_EP_DISP_GEOM / MORI_EP_COMB_GEOM are unchanged: they override a whole leg and bypass the table. 2. num_qp_per_pe defaults to 2 rather than 1. Only the internode path reads it, and 1 starves that path by about 1.5x; the intranode path never reaches RDMA and ignores it. The old default was the safe value for a consumer that does not exist. It reaches the wire: _make_dev_comm sets reqs.gda_context_count = num_qp_per_pe, and ccoGda picks its QP as contextId % numQpPerPe, so a context count below the QP count silently collapses the stripes onto a subset. The internode harness now defaults --num-qp to 1, so the shipped default is what a plain run exercises and a deliberate 1 is still one flag away. 3. The field comments for all six geometry fields and for the two rdma ones said the opposite of what the code did. Corrected. No change to what is validated -- that surface was already two-layer and is listed in the PR discussion: __post_init__ rejects topology, shape and self-consistency errors (gpu_per_node divisibility, num_qp_per_pe, quant on internode, the two enum fields, both-or-neither dtypes, 16 B token alignment, max_token_type_size), and the backend's _unsupported gate rejects capability errors (dtype with no kernel, scatter, std_moe, quant, scale rows that are not whole dwords) before the arena or the device communicator is taken. Verified: the overlay checked field by field against the shipped MI308X table for no pin, a single pin, a two-field pin across phases, and an invalid pin that the clamp catches; 35 pytest; black, ruff and pre-commit clean over all 28 files. NOT pushed.
…ares against
The name did not carry the one fact a reader needs first: the field is only read
when internode_kernel == "auto". It is now internode_auto_ll_max_tokens, and the
harness flag is --auto-ll-max-tokens. The field is new in this PR and has never
shipped, so there is no compatibility path to keep.
The description was also wrong, and I wrote it. It said "the auto crossover, in
tokens per rank", which reads as a property of the configured shape and invites
confusion with max_num_inp_token_per_rank. What it is actually compared against
is THIS CALL's token count -- input.shape[0] in dispatch, cur_rank_num_token in
combine -- so one op alternates between the two families as the batch changes:
cfg = ...(max_num_inp_token_per_rank=4096, internode_auto_ll_max_tokens=512)
op.dispatch(x[128]) -> v2_ll
op.dispatch(x[4096]) -> v2
That is the whole reason "auto" is described elsewhere as the only mode that can
switch at runtime; the old wording contradicted that claim two lines later.
Renamed in the config, the backend, the tuning-table docstring, the package
README, the design doc and the harness. Verified: the per-call choice flips at
512/513, an explicit v2 or v2_ll ignores the threshold at every size, a negative
value still raises, --auto-ll-max-tokens is the flag argparse exposes, and no
occurrence of the old name survives. 17 pytest; black, ruff and pre-commit clean.
NOT pushed.
…e last commit A final pre-merge pass confirmed 11 findings. Four matter, and three of them are mine from the previous two commits. 1. THE FLYDSL GUARD WAS DEAD CODE. The previous commit added EpDispatchCombineOpFlyDSL._unsupported() to reject an internode config on the default backend, and claimed to have closed that trap. It did not: nothing calls the hook. The base's _gate() reads kernels.unsupported, and the flydsl KernelSet never passed one, so a 2-node config on the default backend was still accepted and still built intranode kernels. I verified that change by asserting the hook returned a string, not by constructing an op -- the probe I ran even printed "backend gate fires at op construction", which I had assumed rather than tested. It now gates BEFORE SymmArena, as the hip backend does and for the same reason: raising after the symmetric window is registered leaks it. 2. tuned() hijacked the internode tuning table. Its setdefault of the four intranode block/warp fields is indistinguishable, by __post_init__, from a caller pin -- and the per-field pin overlay added last commit then applied intranode geometry to every internode bucket (dispatch warp 16 where the table says 4 or 8). tuned() now returns early for an internode config: the intranode tables are keyed on one node's geometry and carry no rdma_block_num, so they had nothing to say about it even before the overlay. 3. reset() left the internode chunk-slot allocator advanced. Only EpCombineAll clears blockFlagCounter, so a reset taken between a dispatch and its combine -- which is what reset() is for -- left it non-zero, and the next dispatch addressed slots past its node's slice of the send buffer and signalled a chunk flag no receiver polls. reset() now zeroes it and interNodeBlocksBarrier. 4. The internode kernels are wave64 only, and nothing said so. The intra-warp prefix count is __popcll(mask << (warpSize - laneId)) on a uint64_t, which discards the lanes at or above laneId only when the shift width equals the container width. On wave32 the bits stay inside the 64-bit value and every set bit is counted: laneId 5 of a 7-lane mask yields 7 instead of 5, so the send slot is wrong rather than the code merely being slow. MakeEpInterNodeCfg now rejects waveSize != 64. Rewriting the idiom to be wave-agnostic is the better fix but needs a GPU run to land, and the machine is in use. Also: the README still documented num_qp_per_pe as defaulting to 1, with the rationale the previous commit reversed, and the harness flag as 2 when it is now 1. Both inverted, both corrected. Verified by construction rather than by inspection this time: _gate raises for a 2-node flydsl config; tuned() pins nothing on internode and its buckets are now identical to the plain constructor's, while intranode tuned() still tunes; the wave64 guard accepts gfx942 and rejects gfx1201 through the real spec entry point. All 8 JIT kernels compile for gfx942; 35 pytest; pre-commit clean. NOT pushed.
The series prints six invented tags -- disp, comb, conv, hdis, hcom, hwal -- and
defined none of them anywhere the reader of the output would look. I then used
hwal in comments and in a doc as though it were common knowledge.
Rank 0 now prints a legend ahead of the series, so the output is readable
without opening this file:
# series legend, one value per round, microseconds:
# disp dispatch kernel time (GPU events)
# comb combine kernel time (GPU events)
# conv the dtype cast between them, when the legs differ
# hdis host time inside the op.dispatch() call
# hcom host time inside the op.combine() call
# hwal host wall clock for the whole round
# A round whose hwal exceeds disp+conv+comb has a gap the other
# series do not account for. A rank whose hwal is ~2x its peers'
# is two ranks on the two SMT siblings of one physical core.
The tags stay short because sixteen ranks print these and the columns have to
line up. The two comments and the module docstring that used hwal bare now say
"host wall" and point at the legend.
Checked the rest of the PR for the same habit: these six are the only invented
abbreviations it prints; the table headers were already spelled out.
NOT pushed.
A readability pass over everything the PR adds. 109 renames and 28 reworded comments across nine files; no behaviour change, no logic edits. The harness took most of it, because it was the worst: single-letter locals in 1300 lines of measurement code, and a tuning function whose arithmetic ran on bt/ct_/bph/cph/bw/cw_/bm/cm/bmax/cmax/dmed. Those now say what they hold -- incumbent_metrics, candidate_phases, incumbent_worst, median_diff. The fixed_wins tuple is also unpacked by name rather than by position, which is where an index bug lived earlier in this PR. What was deliberately NOT renamed, and why: - Anything bound by name across a boundary: every field of EpInterNodeArgs, EpInterNodeCfg, EpInterNodeKernelCfg and EpInterNodeRequest (ctypes and the rendered JIT source read them as text), every arena region name string, and the two helpers the test calls directly (_internode_geom_for, _internode_use_ll). - The vocabulary of the codebase: pe, laneId, warpId, blockId, bid, destPe, topk, rdma, xgmi, ll, numa, qp, wqe, cq, lsa, gda, cco, fp8, bf16, dtype, cfg, args, idx. Expanding these would make this PR read differently from every file around it; each was checked against origin/main first. - Almost all of the device kernel. It is a port of src/ops/dispatch_combine/internode_v1.cpp and the two are read side by side whenever the port is checked, so a name that appears in both stays. Only five helpers this PR itself introduced were touched (EpInterNodeWin -> EpInterNodeWindow, EpInterNodeOff -> EpInterNodeOffset, and their `obj` parameter -> `region`); whitespace-insensitive, the non-comment device diff is 36 lines. - The six per-round series tags (disp, comb, conv, hdis, hcom, hwal). Sixteen ranks print them and the columns have to line up; the legend added in the previous commit is what makes them readable. Verified: ruff's F rules (undefined and unused names) pass over every changed file, which is what catches a half-finished rename; the eight JIT kernels compile for gfx942; the host library builds; 35 pytest; the argparse, the AST parse of the v1 harness, _median and _phase_stats all exercised directly; the fixed_wins pack and unpack re-checked field by field; pre-commit clean. NOT verified: the harness's GPU paths -- _bench, _tune and the correctness loop -- have not been executed since the rename, because another job holds both nodes. 1275 of the changed lines are in that file, so it wants a two-node run before this is trusted. NOT pushed.
The v2 entry had no stress case, so ci_cco.yml's soak step ran --cmd bench --rounds 2000 instead. That is not the same test. v1's stress cycles 128 pre-generated datasets in which EACH RANK draws a fresh token count in [1, max_tokens]; bench sends max_tokens every round from one buffer. The difference is the point of the step: - Ragged and unequal loads, which nothing in this PR tested. A review pass flagged exactly this: routing was uniform and every rank always sent max_tokens, so the empty-chunk half of the chunk protocol was unreachable from the tests. - Under internode_kernel="auto", the per-call switch between the v2 and v2_ll families, since the drawn counts fall on both sides of the crossover. Nothing else exercises that switch. - A deep queue: the soak drains every --stress-sync-interval rounds rather than every round, as v1 does, so a protocol bug is not hidden behind the synchronisation of a per-round check. No per-round verification, deliberately -- this looks for hangs, faults and drift over many rounds, and --cmd test is where correctness is checked. --stress-datasets and --stress-sync-interval both default to 128, matching v1, so the two soaks stay comparable. ci_cco.yml's step now runs --cmd stress with --kernel-type auto. NOT reproduced from v1: its CUDA-graph phase. Graph capture over the internode plan sequence is untested, and putting it in a soak would confuse a capture bug with a protocol one. Verified as far as is possible without hardware: the flag parses, the CI command line reaches torchrun through tools/run_internode_test.sh unchanged, ruff and black are clean, pre-commit passes over all 31 files. The soak itself has NOT been run -- another job holds both nodes.
Line-by-line against v1's stress_dispatch_combine, one call disagreed: v1 does
run_combine(op, combine_input, None, indices) -- no weights -- and the v2 soak
passed them. That is not cosmetic now that this PR made combine honour
want_weights: with weights the soak folds them every round and exercises a
different kernel path from the case it is supposed to mirror.
Everything else was already aligned: 128 datasets cycled, a drain every 128
rounds, a fresh per-rank token count in [1, max_tokens] per dataset, scales
carried through dispatch. The CUDA-graph phase is still deliberately not
reproduced.
Verified on 2 nodes x 8 MI308X after the change, one continuous run:
STRESS rounds=500000 datasets=128 tokens=[2,128] of 128 routing=uniform
kernel=auto families exercised: v2_ll
STRESS OK: 500000 rounds in 129.5s (0.26 ms/round, rank0=0 rank1=0)
No hang, no fault, no memory fault across 500k rounds on 16 GPUs.
One coverage note: at --max-tokens 128 every drawn count is below the auto
crossover, so all 500k rounds ran v2_ll and the per-call family switch was not
exercised. A 4096-token soak does reach both families (verified separately at
200 rounds: "families exercised: v2+v2_ll"); at ~5 ms/round a 500k-round run
there is about 40 minutes and has not been done.
…nsor The launch-argument cache skipped a write when the incoming int matched what the buffer last held, but the non-int branch wrote through it without clearing the entry. `weights_buf` and `scales_buf` take an int address, a Tensor, or None from the same caller, so the legal sequence `p, p, None, p` launched its fourth round against the null the third had stored. Reproduced on the published EpDispatchPlan ABI, no kernel launch: packed [4096, 4096, None, None] before, [4096, 4096, None, 4096] after. EP itself never hit this -- both of its optional buffers are always a Tensor or None, never an int, so they never enter the cache at all. Instrumenting the pop over a 30-round 16-rank two-node run counted zero fires, and an A/B of that run with and without this commit hangs and passes at the same rate. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Bare `wait` reaps every child and returns 0 regardless, so the two-node steps reported only node1's status: node2 could fail its correctness check, its bench, or its stress round and the job still went green. Half of EP16 runs over there. Reproduced with a stand-in remote leg exiting 17: bare `wait` ends the step at 0, `wait "$NODE2_PID"` at 17. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The v2 entry builds its 8 ranks per node by spawning, so --nproc-per-node 8 asks it for 64 and it refuses -- correctly, and it names the way out: pass --spawn 0 and let torchrun build them. That flag could not be given, because this wrapper rejects options it does not know, which left the new --nproc-per-node and the new entry unusable together. Verified with LOCAL_WORLD_SIZE=8 against the installed entry: the refusal still fires without --spawn and is cleared with --spawn 0, and the wrapper now forwards it. CI is unchanged -- it runs nproc=1 with the entry's own spawn. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
jhchouuu
changed the title
jhchouuu
changed the title
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
2 participants
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.
Summary
Ports the v1 internode dispatch/combine kernels onto the v2 CCO/GDA interface (JIT plans,
ep_internode_cfg/spec/kernel, plan wiring inep_plans.pyandplan_api.py), fixes a doorbell-ordering deadlock incco_scale_out.hppthat any GDA caller offlushAsynccan hit, and adds two-node CCO coverage toci_cco.yml.The internode path is selected by
gpu_per_node < world_sizeand is implemented only by thehipbackend; the FlyDSL backend rejects such a config rather than silently building intranode kernels.The CCO bug:
flushAsyncImplcorrupted the doorbell-ordering chainSymptom
The two-node internode bench/stress hung nondeterministically -- sometimes on round 1, sometimes after a couple hundred. All 16 GPUs sat at 100% utilisation with no forward progress, and nothing was compiling, so it was a genuine deadlock rather than a slow start.
Where it stopped
Attaching
rocgdbto a hung wave found exactly one wavefront parked inringDoorbellOrdered, waiting for its turn. Its registers heldmyPostIdx = 15541againstdbTouchIdx = 15543.That is the whole bug in two numbers.
ringDoorbellOrderedwaits on strict equality (dbTouchIdx == myPostIdx) against a counter that only ever increases, so once the counter is past the value being waited for, the condition can never hold again. The delta of 2 WQEs is exactly one put-with-signal.Root cause
flushAsyncImpltreatedwq->postIdxas "WQEs written" when it is really a reservation counter.putImpl/putValueImpl/signalImplall bumppostIdxbefore they write their WQEs and before they reachringDoorbellOrdered. A flush running concurrently with a posting warp therefore:postIdxcovering slots whose WQEs were not written yet, and rang the doorbell for them -- pointing the NIC at a stale SQ slot;dbTouchIdx = curPostIdx, stepping over the reservation of a warp that was still waiting for its turn. That warp is then never released, and the tokens behind its reservation are never sent;cq->needConsIdxfrom the flush, which races with a second flush on the same QP.Fix
flushAsyncImplbecomes a snapshot and nothing else: atomically loadpostIdx,__threadfence_system()to keep flush's release semantics, return it. Every posting path already rings its own doorbell, andringDoorbellOrderedadvancesdbTouchIdxin reservation order, so the queue drains on its own; draining to the snapshot isquietUntil's job.flushAsyncImplflushAsyncImplccoGdaOptFlagsAggregateRequestshad no caller and, after this change, nothing would ring for WQEs posted under it, so it is now rejected where it enters with a note on what re-enabling it would need.v2 and v2_ll are two kernels, and the choice is explicit
They always were separate JIT modules -- separate entry symbols, separate cache keys, four device bodies rather than two branches -- but nothing said so at the API, and which one ran was decided by a token-count rule with a private attribute as the only override.
Naming a family compiles only that family;
autocompiles both because onlyautochooses per launch. Read from a private JIT cache per run:auto-> 8 kernels,v2-> 6,v2_ll-> 6 (the four passes the families share are compiled once either way).internode_auto_ll_max_tokensis compared against this call's token count (input.shape[0]for dispatch,cur_rank_num_tokenfor combine), not against the configured capacity, so one op alternates between the families as the batch changes.Pinned geometry (
dispatch_block_num,combine_rdma_block_num, ...) now overrides the tuning table per field: a pinned field wins in every token bucket while the fields leftNonestay tuned.Bugs found and fixed while reviewing this PR
Several adversarial review passes over the full diff; these are the ones that were wrong rather than untidy.
DispatchInterNodeRecvindexed its 8-way sub-block fromblockIdwhere its own combine twin usesbid.bidstrides byrdmaBlockNum, so they agree only whenrdmaBlockNum % 8 == 0; otherwise tokens are dropped and others delivered twice. Latent because every dispatch row in the shipped table happens to be a multiple of 8 -- the combine column already carries 21. A/B withMORI_EP_DISP_GEOM=64,21,8at 4096 tokens: pre-fix FAIL 3/3 with 16 of 16 ranks disagreeing, fixed PASS 3/3.rdma_block_num >= block_numleaves the intra-node half with no blocks and makes the dispatch fan-in wait for arrivals that cannot happen. The clamp existed only on the tuning-table hit path; the untuned shape, the env pin and a caller-pinneddispatch_block_numall bypassed it. Clamped at one choke point, plus a strict check inMakeEpInterNodeCfg.quant_type='fp8_direct_cast'was unlocked on the internode path here, where its fp8 combine staging does not produce correct tokens (weights exact, ~51% of hidden elements outside a 30% tolerance). Rejected at config time and the now-unreachable device code removed (111 lines in the kernel). The intranode path, where the feature is real and tested, is untouched.__popcll(mask << (warpSize - laneId))on auint64_t, which is correct only when the shift width equals the container width.MakeEpInterNodeCfgnow rejectswaveSize != 64.DevCommHandlewas never released;dispatch()kept the caller'sindicestensor as a baredata_ptrthatcombine()dereferenced on the next call; construction leaked the arena on any exception afterSymmArena.combine_out_weightswas sized byMaxNumTokensToRecvbut indexed by the local token id;recv_scales()viewedout_scalesat the padded intranode stride while the internode kernel writes it packed; the send buffer is addressedpe * m + slotwhile slots are handed out in whole wavefronts._rdma_algo_token_counthadscatter_(index, dim, value), so every non-LL bench died in_report_tablesafter the kernels had produced their numbers; the non-greedy tuning summary read the wrong tuple fields; the correctness loop generated scales and then dispatchedNone.Performance
2 nodes x 8 MI355X (world=16), bf16, topk=6, hidden 7168, 2 QP/PE. Latency in µs, from the harness's
Averagerow.cco= ops-v2 JIT plans,shmem= v1 AOT kernels. Measured before the rename, so the kernel is what is now calledv2_ll.Dispatch
Combine
16-512 tokens is where cco wins by margins well clear of the noise; from 4096 up both saturate the network to within a percent; at 8192 cco is ~1% slower and that is not noise. The 1-token row is omitted -- its spread was 10-13%.
Small-token fp8 dispatch / bf16 combine on 2 x 8 MI308X, EP16, hidden 6144, topk 8, all four rows hitting the shipped tuning table:
An earlier revision of this section read the 512/1024 pair as a crossover and called the shipped default of 512 conservative. A later three-repetition measurement on MI308X at hidden 7168 / topk 6 withdraws that:
v2_llalso beatsv2by 4.7% at 4096 and 3.4% at 8192, so there is no single crossover on this shape and the threshold must not be raised on the strength of two points. The default stays at 512.CI
ci_cco.ymlgains a two-nodecco-internode-testjob modelled onci.yml'sinternode-test, driving correctness, bench and a stress soak. Correctness runs both kernel families; the soak runs--cmd stress --kernel-type auto, so it also covers the per-call family switch.gpu_per_nodeto equal CCO's LSA team size, so one host cannot emulate two nodes. The GPU-freetest_internode_regions.pyis collected instead.tools/run_internode_test.shgains--entry(defaulting to the existing shmem harness, soci.yml/nightly.ymlare unchanged),--topk,--roundsand--spawn.--kernel-typeis forwarded only when given, so each entry keeps its own default and its own vocabulary (v1|v1_ll|async_llagainstauto|v2|v2_ll).Test plan
Two nodes x 8 MI308X (world 16), ROCm 7.14, unless noted.
--cmd test, dispatch + combine against the analytic golden:v2,v2_llandautoat 128 tokens,v2at 4096, and asymmetric fp8 -> bf16--routing skewedand--routing local, which reach the empty-chunk path that uniform routing cannot--cmd stress, one continuous run of 500 000 rounds with per-rank token counts drawn in [1, max]: 129.5 s,rank0=0 rank1=0, no hang or fault--cmd benchand--cmd tuningon both familiesci_cco.ymltwo-node matrix reproduced throughtools/run_internode_test.shblockId/bidfix: FAIL 3/3 before, PASS 3/3 after, same geometry and token countauto8 kernels,v26,v2_ll6pre-commitclean over all 31 filesctestfailure set identical before and after (10 pre-existing GDA failures on this host, environmental)cco-internode-testjob green on the runner pair -- and green on work actually done, not on skipped steps: all three token cases report0 of 16 ranks disagreewith zeroFAILlines, and the soak printsSTRESS OK: 2000 rounds in 5.4sFollow-ups (deliberately not in this PR)
Tuning coverage.
internode_tuning_configs.pyships exactly one key,("mi308x", 16, 6144, 8)for fp8-dispatch/bf16-combine, and above 16 tokens thatkey is a single flat bucket spanning 32 through 8192. Every other shape -- including
hidden 7168 / topk 6, the shape this PR's own bf16 tables were measured at -- misses
the table, so
lookup()returnsNoneand the op keeps the untuned default ofblock 96 / rdma 64 / warp 8per phase.Two consequences worth a separate PR rather than more churn here:
block_num <= CU countlives insidelookup(), so it onlyapplies on a table hit. On an 80-CU MI308X the untuned default of 96 blocks
over-subscribes the CUs and runs the tail in a second wave, which is exactly what
the table's own header says a latency-bound small-token kernel cannot afford.
None, since the key space(device, world, hidden, topk)is combinatorial.Sweeping is cheap on this path -- geometry is deliberately not rendered into the
translation unit, so a 165-candidate sweep is one compile -- but accepting a
sweep winner is not. Dispatch and combine are coupled (same replay, same QPs, one
arena), so rows are best pairs rather than per-phase argmins, and two of the three
guards in the re-tuning protocol are manual. That work should not gate this port.
Remaining test-plan item. The unchecked 4096-token soak in the test plan above is where
autoactuallyalternates families; the 500 000-round soak that did run was at 128 tokens and so
exercised
v2_llthroughout.BNXT. On BNXT NICs the two-node steps this PR adds hit a pre-existing
quietUntil<BNXT>defect onmain(an unsignedfetch_maxlatching awrapped-below-zero
doneIdx,mainsince #449), which hangs roughly one run inthree. #653 fixes it; verified here 0 hangs in 12 against 2 in 4 unpatched.
It does not gate this PR -- the
cco-internode-testrunner is ionic, not BNXT -- butanyone running these kernels on BNXT wants #653.
Known limitations
quant_typeis rejected on the internode path (see above).recv_scales/recv_indicesare not compared against anything.BindCallingThreadToGpuNumaOnce. It is measured on the CCO path only.