EP: CUDA-graph-safe HT internode dispatch+combine (worst-tokens mode) - #7
Closed
fergusfinn wants to merge 3 commits into
Closed
EP: CUDA-graph-safe HT internode dispatch+combine (worst-tokens mode)#7fergusfinn wants to merge 3 commits into
fergusfinn wants to merge 3 commits into
Conversation
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).
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
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.
fergusfinn
force-pushed
the
ht-cudagraph-worst-tokens
branch
from
June 10, 2026 14:06
d8aafc3 to
73a5d4c
Compare
Owner
Author
|
Moved to the org fork: superseded by doublewordai#3 (same branch, same description plus a note that it is CXI-agnostic). |
fergusfinn
added a commit
to doublewordai/vllm
that referenced
this pull request
Jun 10, 2026
With the UCCL ht-cudagraph-worst-tokens kernels (fergusfinn/uccl#7), DeepEP high-throughput dispatch+combine are stream-capturable: no host count sync, worst-case static recv shapes, per-expert counts in a device tensor. Behind VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH (default off): - allow FULL_DECODE_ONLY for deepep_high_throughput instead of forcing cudagraph_mode NONE - size num_worst_tokens per captured shape (padded step tokens x num_dispatchers) from the forward context; eager steps (prefill, mixed, uncaptured) keep the host-synced path with zero padding cost - build ExpertTokensMetadata directly from the device counts tensor (expert_num_tokens_cpu=None); fused_marlin_moe never reads counts - join dispatch AND combine back to the compute stream (async off) whenever worst-token mode is active: async comm-stream work in flight races subsequently replayed graphs on the UCCL ring buffers Validated 2-node EP8 DSv4-Flash (capture sizes 64..2048): conc 1024 ISL/OSL 1024/128 -> 1642 tok/s (+45% vs the thread-DBO baseline); conc 8192 1024/1024 -> 9243 tok/s (+13%), 8192/8192 requests, 0 failures. Known: shapes <64 tokens crash capture (suspected lazy NCCL/EPLB init mid-capture; floor the capture list via compilation config), and the blanket sync fence taxes prefill (event fence is follow-up).
fergusfinn
added a commit
to doublewordai/vllm
that referenced
this pull request
Jun 11, 2026
With the UCCL ht-cudagraph-worst-tokens kernels (fergusfinn/uccl#7), DeepEP high-throughput dispatch+combine are stream-capturable: no host count sync, worst-case static recv shapes, per-expert counts in a device tensor. Behind VLLM_DEEPEP_HT_WORST_TOKEN_DISPATCH (default off): - allow FULL_DECODE_ONLY for deepep_high_throughput instead of forcing cudagraph_mode NONE - size num_worst_tokens per captured shape (padded step tokens x num_dispatchers) from the forward context; eager steps (prefill, mixed, uncaptured) keep the host-synced path with zero padding cost - build ExpertTokensMetadata directly from the device counts tensor (expert_num_tokens_cpu=None); fused_marlin_moe never reads counts - join dispatch AND combine back to the compute stream (async off) whenever worst-token mode is active: async comm-stream work in flight races subsequently replayed graphs on the UCCL ring buffers Validated 2-node EP8 DSv4-Flash (capture sizes 64..2048): conc 1024 ISL/OSL 1024/128 -> 1642 tok/s (+45% vs the thread-DBO baseline); conc 8192 1024/1024 -> 9243 tok/s (+13%), 8192/8192 requests, 0 failures. Known: shapes <64 tokens crash capture (suspected lazy NCCL/EPLB init mid-capture; floor the capture list via compilation config), and the blanket sync fence taxes prefill (event fence is follow-up).
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.
Makes the high-throughput internode path fully CUDA-graph capturable and replay-safe when
num_worst_tokensis set, validated end-to-end on Isambard (GH200/CXI, 2 nodes x 4 ranks). All changes are inert whennum_worst_tokens == 0.Commits
notify_dispatchspin-waits for the host to reset host-mapped counters before writing sums; host code doesn't re-execute under replay, so the second replay deadlocked. The wait is now skipped in worst-tokens mode. Also fixes the combine NVL sender's last-slot bound, which fell back to the (padded) input row count and sent garbage tail tokens into flow-control timeouts — it now reads the real total from the handle'srecv_gbl_rank_prefix_sumon device.notify_dispatchoptionally writes per-expert recv counts to a deviceint32[num_local_experts]tensor, anddispatch()returns it in the counts slot under worst-tokens mode, so graph-mode consumers (masked grouped GEMM) never touch the host.ep/bench/test_ht_cudagraph.py— captures layout+dispatch+expert-standin+combine into one graph, 50 replays with fresh data and routing, validated against the normal host-synced path, an analytic oracle, and the device counts; plus eager/replay interleaving on one Buffer.Validation: all phases pass on 2x4 GH200 over Slingshot/CXI (jobs 5132248 and 5132662; eager 0.76 vs replay 0.72 ms/step at test size). Known cosmetic issue: the test prints
cudaErrorContextIsDestroyedat interpreter exit afterbuffer.destroy()— teardown ordering, after all checks pass.