fix(ep): support 12-rank HT combine - #10
Closed
fergusfinn wants to merge 9 commits into
Closed
Conversation
Post and wait QUIET across every D2H ring for the CXI path so sibling-ring work cannot remain queued behind a completed quiet. Cache the CXI transport selection when each proxy is constructed, and raise the 8-rank combine RDMA send chunk to satisfy the combine forwarder invariant when NUM_MAX_NVL_PEERS=4.
Two fixes, both inert when num_worst_tokens == 0: 1. notify_dispatch spin-waits for the host to reset the host-mapped recv counters to -1 before writing its sums. Host code does not re-execute under CUDA graph replay, so the second replay deadlocks on the stale counter. Skip the wait when num_worst_tokens > 0; the writes stay unconditional. 2. The combine NVL sender bounds its last (rank, channel) slot by num_tokens, which under num_worst_tokens is the padded input row count, so it sends garbage tail tokens and hits flow-control timeouts. Plumb the dispatch handle's recv_gbl_rank_prefix_sum (device tensor) into the kernel and use its last element as the real total.
…mode CUDA-graph consumers of num_worst_tokens dispatch cannot read the host-mapped per-expert counters without a sync, and the host list is deliberately skipped in that mode. Have notify_dispatch also write the counts to an optional device int32[num_local_experts] buffer, and have dispatch() return that tensor in the num_recv_tokens_per_expert_list slot when num_worst_tokens > 0 (downstream: masked grouped GEMM reads it directly inside the graph).
2-node test: captures layout+dispatch+expert-standin+combine into one torch.cuda.CUDAGraph with num_worst_tokens, replays 50x with fresh data and routing per replay, validates against the normal host-synced path, an analytic oracle, and the device expert-counts tensor; also checks eager/replay interleaving on one Buffer. Validated on Isambard (GH200/CXI, 2x4 ranks): all phases pass.
The host-synced protocol (host writes -1, kernel waits for -1 then writes its sum, host polls) is serialized by the host polls themselves. A replayed worst-tokens graph contains many notify kernels writing sums with no polls pacing them, and callers may issue the next host-synced eager dispatch while the graph is still executing: the eager call's -1 reset lands mid-flight, a graph sum clobbers it, and the eager kernel's wait-for--1 spins forever (observed as 'DeepEP error: timeout (dispatch CPU)' on all ranks under mixed replay/eager serving load). Worst-tokens mode has no use for the host counters - consumers read the device-side counts - so gate the stores, not just the waits, on num_worst_tokens==0. Adds ep/bench/test_ht_ragged.py: extreme-skew eager dispatch shapes (including the observed serving wedge shape and zero-token ranks) and a no-sync replay/eager alternation phase that reproduces the clobber window organically.
Both HT cudagraph tests aborted (SIGABRT, exit 134) at interpreter exit after all checks passed: buffer.destroy() tears down the CUDA context while captured graphs and their static tensors are still alive, so the caching allocator's deferred frees hit a destroyed context inside TensorImpl destructors, which throw from a noexcept path. Free every CUDA object first (drop refs / move phase locals into a nested frame), gc + empty_cache + synchronize, then destroy the buffer and the process group. Validated on 2x4 GH200: both tests now exit 0.
…ispatch test The graph-mode contract returns an int32 device tensor of per-local-expert recv counts where the legacy API returned an empty host list; validate it against the reference counts instead of asserting emptiness.
Report host-mapped recv counters and per-peer count-arrival state when dispatch times out. Sentinel-fill the count-exchange recv region before notify_dispatch so timeout readback can distinguish missing peer slots from arrived counts.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2a05fc0c89
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| 2: Config(Buffer.num_sms, 24, 256, 6, 128), | ||
| 4: Config(Buffer.num_sms, 6, 256, 6, 128), | ||
| 8: Config(Buffer.num_sms, 6, 256, 6, 128), | ||
| 12: Config(Buffer.num_sms, 36, 288, 20, 128), |
There was a problem hiding this comment.
With the default EP build, NUM_MAX_NVL_PEERS is 8, so enabling a default 12-rank config makes 12-rank HT runs get past the Python Unsupported number of EP ranks check only to fail in the C++ path: Config::get_nvl_buffer_size_hint/get_rdma_buffer_size_hint assert that ranks above the NVL peer count are divisible by NUM_MAX_NVL_PEERS (ep/include/ep_config.hpp:65-66 and :103), and the internode kernels also divide ranks by that compile-time value. If this entry is intended only for 2x6 or 3x4 builds, it needs to be gated on the compiled/local NVL peer count (or the C++ support updated); otherwise stock 12-rank launches now advertise a default config that aborts at runtime.
Useful? React with 👍 / 👎.
fergusfinn
changed the title
fergusfinn
force-pushed
the
codex/ep12-ht-config
branch
from
June 18, 2026 22:25
2a05fc0 to
ba3efa2
Compare
fergusfinn
force-pushed
the
codex/ep12-ht-config
branch
from
June 18, 2026 22:34
ba3efa2 to
2df5384
Compare
fergusfinn
changed the base branch from
ht-cudagraph-worst-tokens
to
upstream-base
June 19, 2026 15:34
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.
Summary
Validation