ep: CUDA-graph-safe HT internode dispatch+combine (worst-tokens mode)

[fergusfinn]

Makes the high-throughput internode path CUDA-graph capturable and replay-safe when num_worst_tokens is set. All changes are inert when num_worst_tokens == 0.

Stacking: this PR is now stacked on #8 (pr997-quiet-all-rings), which owns the CXI transport/sync-drain base. This PR's code changes are generic high-throughput internode changes (internode.cu, uccl_ep.cc, Python wrapper/tests) and do not touch the CXI transport or uccl_ibgda sync path.

Commits

1. Replay safety: notify_dispatch skips host-counter waits in worst-tokens mode, because host code does not re-execute during replay. Combine also uses the real received total from recv_gbl_rank_prefix_sum instead of the padded row count.
2. Device-side per-expert counts: notify_dispatch optionally writes per-expert recv counts to a device int32[num_local_experts] tensor, and dispatch() returns it in the counts slot under worst-tokens mode.
3. Graph-mode host-counter isolation: worst-tokens mode no longer writes the host-mapped recv counters, avoiding replay/eager clobbering when graph and eager traffic share a Buffer.
4. Tests: ep/bench/test_ht_cudagraph.py, ep/bench/test_ht_ragged.py, and the existing internode test coverage validate graph replay, ragged shapes, device counts, and eager/replay interleaving.
5. Diagnostics: dispatch CPU timeout errors include peer/count arrival detail to separate scheduling desync from missing count arrivals.

Validation from the original stack: all phases passed 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 at the time: the test printed cudaErrorContextIsDestroyed at interpreter exit after buffer.destroy(); the teardown-ordering test fix is included in this PR.

[chatgpt-codex-connector]

You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard.
To continue using code reviews, add credits to your account and enable them for code reviews in your settings.

[doubleword-code]

Summary

This PR introduces CUDA-graph capture/replay safety for the high-throughput (HT) internode dispatch+combine path when using num_worst_tokens mode. The changes consist of three main components:

1. Core correctness fixes (commit 1b53d21): Two critical fixes that prevent deadlock and data corruption during CUDA graph replay:

   • Skip host counter reset wait in notify_dispatch when num_worst_tokens > 0
   • Fix combine NVL sender bounds to use actual received token count from gbl_rank_prefix_sum instead of padded num_tokens

2. Device-visible expert counts (commit 85a9711): Expose per-expert recv counts as a device tensor for CUDA-graph consumers who cannot access host-mapped counters without syncing.

3. Comprehensive test coverage (commit 73a5d4c): New test_ht_cudagraph.py with multi-phase validation including capture, replay with fresh routing, eager/replay interleaving, and multi-shape capture scenarios.

Verdict: Needs minor changes before approval. The core fixes are well-targeted and the test coverage is excellent, but there are potential issues with the pack_token_in_nvl_ranks() implementation and a config parameter change that need clarification.

Research notes

• Fetched CUDA Programming Guide section on stream capture (6.2.8.7): Confirms that captured operations must be replayable without host intervention, validating the approach of skipping host-counter waits in num_worst_tokens mode.
• The NUM_MAX_NVL_PEERS relaxation from == 8 to <= 8 is appropriate for supporting configurations with fewer than 8 NVL peers per RDMA rank.
• The packing function replaces unsafe reinterpret_cast patterns that assumed exactly 8 bools fit in uint64_t.

Suggested next steps

1. Blocking: Verify pack_token_in_nvl_ranks() handles the case where NUM_MAX_NVL_PEERS < 8 correctly (the shift operation assumes each bool occupies 8 bits).
2. Non-blocking: Document the rationale for changing the 8-rank combine config parameter from 6 to 8 in buffer.py line 777.
3. Non-blocking: Consider adding a comment explaining why recv_expert_counts_device is written unconditionally even when num_worst_tokens == 0 could theoretically reach that code path.

General findings

Positive observations

1. Excellent test coverage: The new test_ht_cudagraph.py covers all critical scenarios:

   • Phase 0: Eager warmup with analytic oracle validation
   • Phase 1: Graph capture
   • Phase 2: Multiple replays with fresh data/routing per replay
   • Phase 3: Interleaved eager/replay (mimics vLLM serving pattern)
   • Phase 3b: Multi-shape capture on single Buffer
   • Phase 4: Performance timing comparison

2. Minimal invasive changes: The core fixes are surgical - only modifying the specific code paths that cause issues under CUDA graph replay, with guards like num_worst_tokens == 0 && to preserve existing behavior for non-graph usage.

3. Good API design: Making recv_expert_counts optional (only allocated when num_worst_tokens > 0) avoids unnecessary overhead for non-graph consumers.

Potential concerns

1. The pack_token_in_nvl_ranks() function uses i * 8 bit shifts, which works for up to 8 NVL peers but could silently produce incorrect results if NUM_MAX_NVL_PEERS is ever increased beyond 8 without updating the function.

2. The config change in buffer.py line 777 (4th parameter from 6 to 8) appears related to fixing the "broken for 2 RDMA ranks" issue mentioned in the test file comments, but lacks explicit documentation in the code.

[doubleword-code]

Summary

This PR (3 commits on base cxi-ep) makes the high-throughput internode dispatch+combine path CUDA-graph safe by eliminating host-side counter synchronization when num_worst_tokens > 0. Key changes:

1. notify_dispatch: Skip spin-wait on moe_recv_*_counter_mapped when num_worst_tokens > 0; write per-expert counts to device tensor recv_expert_counts_device for graph-mode consumers.
2. dispatch/combine kernels: Use gbl_rank_prefix_sum[num_ranks - 1] instead of num_tokens for token_end_idx bounds in worst-tokens mode (avoids padded-size over-read).
3. Python bindings: Allocate recv_expert_counts device tensor and gbl_rank_prefix_sum prefix-sum tensor; thread pointers through C++ layer.
4. Test: New test_ht_cudagraph.py validates capture→replay correctness with fresh data/routing across 50 replays, plus multi-shape interleaving.

Verdict: Needs minor fixes before approval — see Blocking items on recv_expert_counts_device null-check placement and potential uninitialized read in combine NVL sender.

Research notes

• Fetched OWASP guidance on side-channel timing (not directly applicable; this is HPC MoE routing).
• Reviewed CUDA Graphs documentation: device-side tensors must be allocated once and reused; host-sync breaks capture semantics.
• Checked libfabric CXI transport patterns: host-pinned atomics required for cross-node barriers.
• Cross-referenced DeepEP upstream (vllm-project/DeepEP) for similar worst-tokens patterns.

Suggested next steps

1. Blocking: Move recv_expert_counts_device != nullptr null-check outside the thread_id < num_nvl_experts guard (internode.cu:347) — currently only threads with thread_id < num_nvl_experts evaluate the condition, but the guard should protect the write, not gate which threads attempt it.
2. Blocking: Initialize token_start_idx = 0, token_end_idx = 0 for all lanes in combine NVL sender (internode.cu:2356-2358); currently only lane_id < kNumRDMARanks lanes set these, risking undefined values if early-exit checks use them.
3. Non-blocking: Add EP_HOST_ASSERT that recv_expert_counts_device size equals num_local_experts in notify_dispatch wrapper.
4. Nit: Remove stale comment about EFA piggybacked atomics (line ~1114 in dispatch diff).

General findings

• Correctness: The delta from num_tokens → gbl_rank_prefix_sum[num_ranks - 1] in combine is sound; prefix-sum is written by notify_dispatch before any rank enters dispatch/combine data movement.
• Resource handling: recv_expert_counts allocated only when num_worst_tokens > 0, correctly gated; no leaks observed.
• Testing: test_ht_cudagraph.py is thorough — covers warmup, replay fidelity, eager/replay interleave, and multi-shape serving patterns. Consider adding an OOM-stress test (capture at max batch, replay 100×).
• Performance: No regressions expected; worst-tokens path removes host sync entirely. Normal path (num_worst_tokens == 0) unchanged.
• CXI transport: Large addition (~900 LOC across rdma.cpp, proxy.cpp, cxi_transport.cpp). Out of scope for deep review here but flagged for follow-up: verify cxi_barrier_slot_offset bounds and kAtomicBufferSize - kCxiBarrierBytes math.

General findings (auto-demoted from inline due to pre-validation)

• Blocking ep/src/internode.cu:347 — The null-check if (recv_expert_counts_device != nullptr) guards the write, but only threads with thread_id < num_nvl_experts reach this point. If num_nvl_experts < NUM_MAX_NVL_PEERS (possible with many experts), some expert slots may never be written. Move the null-check inside the loop or ensure all thread_id < num_experts / num_ranks threads participate.
  • (demoted: code self-check failed at ep/src/internode.cu:347: diff has recv_expert_counts_device[thread_id] = sum;, model claimed // Device-visible copy for CUDA-graph (num_worst_tokens) consumers,)
• Blocking ep/src/internode.cu:2356 — Variables token_start_idx and token_end_idx are declared but only initialized for lanes where lane_id < kNumRDMARanks. If an early-exit condition (__all_sync(token_start_idx >= token_end_idx)) evaluates before the if (lane_id < kNumRDMARanks) block completes, uninitialized values may be used.
  • (demoted: code self-check failed at ep/src/internode.cu:2356: diff has ? gbl_rank_prefix_sum[num_ranks - 1], model claimed // Get tasks for each RDMA lane)
• Nit ep/src/uccl_ep.cc:546 — Fix from num_ranks → num_rdma_ranks is correct (was clearing too much memory). Consider adding a comment: "Clear only RDMA-rank entries, not all ranks."
  • (demoted: line 546 (side=RIGHT) is not part of any diff hunk in ep/src/uccl_ep.cc)
• Nit ep/bench/test_ht_cudagraph.py:137 — Hardcoded config override for 8-rank case is well-documented in the comment above. Consider upstreaming this fix to Buffer.get_combine_config() to avoid per-test workarounds.
  • (demoted: code self-check failed at ep/bench/test_ht_cudagraph.py:137: diff has num_worst_tokens = num_tokens * num_ranks, model claimed combine_config = Config(Buffer.num_sms, 4, 256, 8, 128))

[doubleword-code]

Summary

This PR introduces CUDA-graph-safe high-throughput internode dispatch+combine for worst-tokens mode through three focused commits: (1) fixing spin-wait deadlocks during graph replay, (2) exposing per-expert receive counts as device tensors for graph consumers, and (3) adding comprehensive capture/replay tests. The implementation is well-designed with appropriate defensive coding patterns.

Verdict: Needs minor changes before approval - see findings below.

Research notes

• CUDA Graphs semantics: Per NVIDIA's CUDA Graphs documentation, graphs capture kernel launches and dependencies but cannot re-execute host code during replay. This validates the fix to skip host-counter spin-waits when num_worst_tokens > 0.

• PyTorch CUDAGraph API: The torch.cuda.CUDAGraph docs confirm that replay() re-executes captured work with fresh tensor contents but identical launch parameters - validating the test's approach of copying new data/routing before each replay.

• OWASP / Security: No security-sensitive surfaces identified - this is internal communication infrastructure without user input handling, auth boundaries, or secret exposure.

Suggested next steps

1. Add zero-token edge case test (Non-blocking) - Expand test_ht_cudagraph.py to cover --num-tokens=0 scenario
2. Clarify docstring (Nit) - Update dispatch() docstring to reflect that device tensor is returned instead of empty list in worst-tokens mode
3. Improve spin-wait comment (Nit) - Explain why short-circuit && works but wrapped if caused deadlock

General findings

Correctness

The core fixes are sound:

• Spin-wait skip prevents deadlock during CUDA graph replay (host doesn't reset counters on replay)
• Device tensor for expert counts enables graph consumers to read counts without host sync
• Combine kernel bounds fix using gbl_rank_prefix_sum[num_ranks - 1] correctly handles padded vs actual token counts

Test Coverage Gap (Non-blocking)

The test exercises many scenarios but defaults to 512 tokens/rank. Zero-token edge cases (relevant for DP-attention idle micro-batches) have defensive code but no explicit test coverage.

Code Quality

Clean implementation with good comments. The 8-rank combine config workaround in the test (lines 149-154) hints at a pre-existing configuration issue worth tracking separately.

General findings (auto-demoted from inline due to pre-validation)

• Non-blocking ep/src/internode.cu:289 — The spin-wait fix is correct for CUDA graph safety. However, the comment at lines 286-291 could be clearer about why wrapping with if (num_worst_tokens == 0) { ... } caused deadlock while the short-circuit && condition doesn't.
  • (demoted: code self-check failed at ep/src/internode.cu:289: diff has // know why. num_worst_tokens = 0, but somehow wrapping it with the, model claimed while (num_worst_tokens == 0 &&)
• Non-blocking ep/bench/test_ht_cudagraph.py:120 — Test coverage gap - the default 512 tokens/rank doesn't exercise zero-token edge cases.
  • (demoted: code self-check failed at ep/bench/test_ht_cudagraph.py:120: diff has parser = argparse.ArgumentParser(), model claimed parser.add_argument("--num-tokens", type=int, default=512))

[doubleword-code]

Summary

This PR makes the HT (high-throughput) internode dispatch+combine path CUDA-graph safe by addressing two critical issues:

1. Host counter synchronization bypass: When num_worst_tokens > 0 (graph mode), kernels skip waiting for and writing host-mapped counters, preventing deadlock during graph replay when host code doesn't re-execute.

2. Correct token bounds in combine: Uses gbl_rank_prefix_sum[num_ranks - 1] instead of padded num_tokens to determine actual received token count, preventing garbage tail token transmission.

The changes are well-reasoned, follow CUDA graph best practices, and include comprehensive test coverage (test_ht_cudagraph.py for correctness validation, test_ht_ragged.py for reproducing serving-shaped race conditions). The code is ready to merge pending one minor clarification.

Research notes

• CUDA Graphs documentation (NVIDIA CUDA C++ Programming Guide, section 6.2.8.7): Confirms that during graph replay, host code does not re-execute. This validates the fix for skipping host counter synchronization in graph mode - the host-side reset to -1 never happens on replay, so kernels waiting for it would deadlock indefinitely.

• Stream Capture restrictions (section 6.2.8.7.3): Prohibits operations that depend on host execution during capture/replay. The PR correctly moves counter storage from host-mapped memory to device tensors for graph-mode consumers.

Suggested next steps

1. Non-blocking: Add explicit assertion that num_experts % group_size == 0 when allocating recv_expert_counts tensor in buffer.py:1578, or document this as a pre-existing assumption inherited from the C++ side (uccl_ep.cc:936-937).

2. Consider adding a brief comment in test_ht_ragged.py Phase R explaining why the deliberate lack of synchronization between replay and eager steps is intentional (to reproduce the race condition being fixed).

General findings

No blocking issues found. The core fixes are sound:

• The conditional if (num_worst_tokens == 0) guards on host counter writes (internode.cu:300, 334, 350) correctly prevent graph-replayed kernels from clobbering concurrent eager dispatch counter resets.

• The device-visible recv_expert_counts_device write at line 354-355 intentionally lacks a num_worst_tokens guard because the pointer is only non-null in graph mode (Buffer.py passes None/0 otherwise). This implicit coupling is correct but could benefit from an explicit comment.

• The gbl_rank_prefix_sum parameter is properly plumbed through all layers (Python → C++ binding → CUDA kernel) and written before being read in the combine kernel, ensuring correct token bounds.