ep: graph-safe worst-token HT dispatch+combine (internode + intranode)

[fergusfinn]

Summary

Makes the DeepEP high-throughput dispatch/combine path CUDA-graph capturable
and replay-safe when num_worst_tokens is set, across both the internode
(RDMA, multi-node) and intranode (NVLink, single-node) kernels. The
num_worst_tokens == 0 path is byte-for-byte unchanged.

Single-node EP (where num_rdma_ranks == 1) routes through the intranode
path; multi-node EP routes through the internode path. Both need the same
graph-safety treatment, so they are fixed together here.

Internode (internode.cu)

1. Gate the host-counter handshake. notify_dispatch spin-waits on the host
   to reset moe_recv_*_counter_mapped to -1 then writes the sum — a per-call
   GPU↔Python handshake. Under graph replay the host never runs, so the 2nd
   replay deadlocks in the spin, and the mid-flight write clobbers a concurrent
   eager dispatch sharing the Buffer. Both are now gated on num_worst_tokens == 0.
2. Device-side per-expert recv counts. The notify kernel also writes
   per-expert counts to a device int32 tensor, returned in place of the host
   list when num_worst_tokens > 0, so graph-mode consumers never touch the host.
3. Real combine total via gbl_rank_prefix_sum. In worst-token mode
   num_tokens is the padded row count; combine now reads the true received total
   from the device prefix sum instead of shipping garbage tail tokens.

Intranode (intranode.cu)

The intranode notify_dispatch kernel wrote the host-pinned
moe_recv_counter/moe_recv_expert_counter unconditionally. The host reads
these only on the num_worst_tokens == 0 path, but the kernel still wrote them on
every call, including graph replays. The single host-pinned counter is shared by
every dispatch on a Buffer, so a graph-replayed (worst-token) decode notify's
counter write, landing mid-flight after a subsequent host-synced eager dispatch
reset the counter to -1 but before that eager dispatch's own kernel completes,
made the eager dispatch's CPU spin-wait return a stale num_recv_tokens.
Downstream this surfaced non-deterministically as a combine-receiver hang (waiting
on a tail for tokens never dispatched) or an illegal memory access once the
wrongly-sized receive buffer fed the MoE/attention path. num_worst_tokens is now
threaded into the intranode notify_dispatch and both counter writes are gated on
num_worst_tokens == 0, mirroring the internode fix.

Tests

• test_ht_cudagraph.py — capture layout→dispatch→expert→combine in one CUDA
  graph; N replays with fresh routing each replay, validated against an analytic
  oracle and the eager path; plus eager/replay interleave on one Buffer (the vLLM
  prefill-eager + decode-graphed condition).
• test_ht_ragged.py — serving-shaped ragged/pressure coverage.

Validation

DeepSeek-V4-Flash, deepep_high_throughput, full decode CUDA graphs, mixed
prefill+decode serving benchmark (random, ISL/OSL 1024/1024, --ignore-eos):

• Single-node 4×GH200 (EP4): without the patch the engine dies after ~250–670
  requests with an illegal memory access; with it, the concurrency sweep
  1024→2048→4096 completes 2048/4096/8192 requests at 100%, across two waves and
  two fresh boots, graphs confirmed active, gsm8k strict-match parity with eager
  (0.977 vs 0.973).
• 2-node EP8: 4096/4096 successful across two waves, graphs active, no faults.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6593bb8fa9

ℹ️ 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".

[chatgpt-codex-connector]

Use the fixed 8-rank combine config in ragged test

In the documented 2-node x 4-GPU run this test asserts num_ranks == 8, so this line selects Buffer.get_combine_config(8), which is Config(..., rdma_chunked_send_tokens=6, ...). For the 2-RDMA-rank internode combine path, ep/src/internode.cu:3114 asserts num_max_rdma_chunked_send_tokens >= num_warps_per_forwarder, and num_warps_per_forwarder is 8, so the first combine throws before the ragged/interleave coverage runs. Use the same Config(..., 8, ...) workaround already applied in test_ht_cudagraph.py.

Useful? React with 👍 / 👎.