perf(ep/v2): sync main + port #613 InterNodeV1LL combine opt to the cco kernel - #631
Merged
Merged
Conversation
- Opt-in via MORI_ENABLE_HOST_PROXY. Unset, every existing path is unchanged: no removed lines in shmem init, the native branch of symmetric memory registration is identical, and GpuStates only grows by appended fields. - GPU warps publish 128 B descriptors into a host-pinned ring; per-NIC CPU threads drain them into ibv_post_send, chaining by QP so a batch costs one post per QP. - Built on IBVerbsDeviceContext rather than the mlx5/bnxt/ionic DV providers -- proxy mode needs no vendor direct-verbs library, and inherits its QP setup, error handling and endpoint teardown. - One ring and one proxy thread per NIC; QPs are picked by the agreed-rail formula and symmetric memory carries per-NIC lkey/rkey. - Atomics: native IBV_WR_ATOMIC_FETCH_AND_ADD on mlx5; SEND_WITH_IMM emulation elsewhere, where the responder performs the add and replies for the fetch variants. Targets are validated against the heap range. - Signals paired with data use RDMA_WRITE on the same QP as the payload so PCIe ordering is preserved. - MORI_PROXY_ENABLED goes through _tunable_defines(), so it is part of the JIT cache key and reaches both the genco and bitcode compiles. - Validated on 2-node EP16 across Mellanox CX7 / MI300X, AMD Pensando AINIC (DigitalOcean and Crusoe) and Broadcom Thor2. Steady-state throughput matches main; proxy mode shows tighter worst-case latency. - Known limitation: the async_ll kernel type is not supported under the proxy.
* migrate mlx ci to 19-05 & 19-14 * fix mlx poll cq bug: clamp doneIdx against dbTouchIdx, not postIdx — the old window exceeded 65536, so a stale wqe_counter wrapped into a bogus advance and the recycle gate overwrote live SQ slots.
…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>
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
4 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.
Targets the #625 branch (
dev/adpat_v1ops_to_cco_api), notmain.What
Merge
maininto the CCO-internode branch. Brings in perf(ep): optimize InterNodeV1LL small-token latency on EP16 #613 (InterNodeV1LL small-token opt), feat(shmem): add a CPU host-proxy RDMA transport alongside IBGDA #558 (host-proxy RDMA transport), bugfix: mlx5 ci hung #622, etc. Clean merge, no conflicts.Port perf(ep): optimize InterNodeV1LL small-token latency on EP16 #613's combine-gather optimization to the cco kernel.
ep_internode_kernel.hppis a deliberate second copy of the v1 kernel, so perf(ep): optimize InterNodeV1LL small-token latency on EP16 #613 landing ininternode_v1.cppvia the merge does not reach it. Mirrored here in the two LL combine paths (CombineIntraNodeLLTyped,CombineInterNodeLLTyped):CombineGather/kCombineVecBytes=16) with a per-launchCombineVecAligned()fallback to the 4B path, replacingWarpAccum<TokT,4>.MultiWarpItergetsCombineVecStepgranularity in the intra-LL path; the inter-LL path setswarpsPerToken = hiddenDim / vecStep(was a fixed 4 → 6 for hidden 6144 bf16) and snapshiddenDimPerWarpto a whole number of steps, keeping every warp on the vector path.Helpers are duplicated rather than shared, matching the PR's two-copy split of the v1 kernel.
Validation (2-node MI308X + bnxt, ROCm 7.14, EP16)
v1_ll
test128-tok, 500 rounds → 1000/1000 Dispatch+Combine Pass, 0 fail. The 16B gather is correct (bf16 hidden 6144 is 16B-aligned; unaligned topk/dtype fall back to 4B).v1_ll
bench128-tok, clean idle before/after (2-rep median):Note on the perf delta
The port is correct and non-regressing but latency-neutral on this MI308X CCO path (combine ~-1.5%, within noise). #613's larger win on MI308X came substantially from its re-tuned InterNodeV1LL launch table (combine block/rdma/warp = 32/21/6), whereas the cco path uses the hardcoded
_KERNELSgeometry (256/128/8) — onlywarpsPerToken(4→6) tracks the tuned split. Small-token combine on MI308X is also XGMI-bandwidth-bound, so vectorizing the gather is neutral there. Realizing #613's latency win on the cco path needs a separate re-tune of the cco v1_ll launch geometry; this PR keeps the two kernel copies in sync and should help on the arch/geometry combos where the shmem path saw the gain.🤖 Generated with Claude Code