ep(bench): unified in-process Python EP benchmark — review-comment follow-ups - #858
ep(bench): unified in-process Python EP benchmark — review-comment follow-ups#858Qinghua Zhou (seagater) wants to merge 23 commits into
Conversation
Make the mscclpp cuda-graph path capture dispatch+combine into a SINGLE combined CUDA graph (dispatch_fn replays the pair, combine_fn is a no-op), matching the NCCL-EP / DeepEP single-graph path and how a real serving stack replays a fused MoE step. Measured on GB200 (rank-major, 1/2/4 nodes): per-kernel kineto times are unchanged vs the previous two-graph capture (dispatch ~16/23/31 us, combine ~15/21/26 us). Only the host per-phase split changes -- the host combine timer folds into dispatch (~2.4 us combine) -- and end-to-end Total(D+C) host latency is marginally lower (one graph launch instead of two).
Drop test/python/ep/run_ep_bench.py (the old shell-out driver), test/python/ep/mscclpp_ep_bench.cu (the pure-C++ LL benchmark), and the in-process CUPTI kernel timer (cupti_kernel_timer.cpp + the standalone CMakeLists.txt that built them). The unified in-process Python driver (run_ep_bench_python.py) is the single entry point, and its kernel-only timing uses the torch-kineto path (EP_KERNEL_TIMER=kineto) exclusively; the --kernel-timing flag and all CUPTI plumbing are removed.
The EP benchmark runs on a single NVL72 domain (same-rack MNNVL, NVLink transport, EP_DISABLE_GIN=1, NCCL_IB_DISABLE=1). DeepEP hybrid mode adds a hierarchical RDMA + NVLink tier for multi-rail networks that this workload never exercises; enabling it only reserves extra buffer capacity and QPs (65/129 vs 17). Set allow_hybrid_mode=0 so the DeepEP config matches the transport being benchmarked. Verified bit-for-bit-equivalent kernel timing vs hybrid mode on a 2-node t=4096 h8704 run (kernel total 1349 vs 1344 us).
The previous gate restricted DeepEP CUDA-graph capture to a single node (num_ranks <= local_world), assuming its symmetric-memory kernels always crash under graph capture internode (CUDA 719). That crash is specific to the RDMA/IB scale-out (GIN) path, NOT the node count: on an all-NVLink / MNNVL fabric (EP_DISABLE_GIN=1, one NVL72 domain) capture works at any scale. Verified DeepEP CUDA-graph at 1/2/4 nodes (4/8/16 GPUs) on GB200 NVL72 -- all exit 0, no CUDA 719. Gate now keys off EP_DISABLE_GIN so graph capture is enabled on the NVLink path and only disabled for the RDMA/GIN scale-out.
Reword the misleading "legacy paired loop" comments. EP_KINETO_SEPARATE=0 is not legacy: it is REQUIRED whenever a backend captures dispatch+combine in one CUDA graph (single replay runs both phases, so the skew-free separate pass cannot isolate combine). SEPARATE=1 remains the correct mode for eager runs (it collapses combine recv-spin skew). Also correct the main-loop force comment: DeepEP manages EP_KINETO_SEPARATE itself inside setup_deepep (single-graphs on the NVLink/MNNVL path at any node count, keeps the separate pass on the RDMA/GIN eager fallback), so it is not forced in the loop; only nccl/flashinfer are forced under --cuda-graph. Comment-only, no behavior change.
Address review comment (move the kernel name parse for specific libraries to the DeepEP/mscclpp files). Each backend module now owns a KINETO_KERNEL_MATCH dict declaring the dispatch/combine kernel-name substrings for that library (mscclpp single dispatch/combine kernels, NCCL-EP single kernels, DeepEP *_impl + epilogue, FlashInfer moeA2ADispatchKernel/moeA2ACombineKernel), with the library-specific naming documented next to the backend it describes. run_ep_bench_python threads the matcher for the active backend through run_backend into _kineto_kernel_us; the shared _parse is now generic (takes a substr tuple) and no longer hard-codes per-library kernel knowledge. Behavior-preserving: verified non-zero kernel-only dispatch/combine for all four backends on a 1-node rank-major run.
Address review comment (remove sync here). The torch.cuda.synchronize() after priming a dispatch for the combine kineto pass is redundant: _run_pass already warms the op and synchronizes before it starts profiling, so the intermediate sync did nothing. Verified kernel-only dispatch/combine still correct for all four backends on a 1-node rank-major run.
Address review comment (remove sync here): the stream.synchronize() between dispatch_fn() and combine_fn() in the warmup paired loop is unnecessary -- the paired dispatch->combine already runs in order on the same stream, and the loop still syncs + barriers after combine each iteration. Also restore the kineto combine-pass prime sync that a prior commit removed (that one is kept). Verified kernel-only dispatch/combine correct for all four backends on a 1-node rank-major run.
Follow-up to the review comment (move the kernel name parse for specific libraries to the backend files). Instead of a per-backend name-substring constant (which read as boilerplate since every library happens to embed the phase word in its kernel names), each backend module now exposes a parse_kineto_kernels(key_averages) -> (dispatch_us, combine_us) that owns its librarys kernel-name knowledge and delegates the summation to a shared sum_matching_kernel_us() in ep_bench_common. run_ep_bench_python threads the active backends parser through run_backend into _kineto_kernel_us; the shared _parse helper is gone and the harness holds zero per-library kernel knowledge. Behavior-preserving: verified non-zero kernel-only dispatch/combine for all four backends on a 1-node rank-major run.
mscclpp now captures dispatch+combine into a SINGLE combined CUDA graph (combine_fn is a no-op), like nccl/flashinfer/deepep. But it was missing from the force-EP_KINETO_SEPARATE=0 tuple, so under --cuda-graph it kept the default separate two-pass whose combine pass runs the no-op and records nothing -> kineto captured 0 LL kernels. Add mscclpp to the tuple so the paired single-pass attributes per-phase kernel time by kernel name. Verified: mscclpp cuda-graph now reports Dispatch/Combine (was captured 0); all four backends report valid kernel-only times in both eager and cuda-graph.
…-only The separate two-pass timing method was labeled as if it were DeepEP-specific code. Reword to make explicit it is a generic method (adopted from DeepEP bench_kineto) that applies to every backend via the backend-supplied dispatch_fn/combine_fn closures; the loop has no per-library logic. Comment-only.
…helper Address review comment (move the cuda-graph logic out of the per-backend files). All four backends now capture dispatch+combine as ONE combined graph (one replay runs both phases, combine_fn a no-op), so the prime/sync/capture boilerplate was duplicated four times. Extract it into capture_dispatch_combine_graph() and have each backend pass its own op closures. The library-specific bits stay local in those closures: DeepEP cached do_cpu_sync=False dispatch args, NCCL-EP capture-stream refetch, mscclpp dispatch-output-to-combine handoff (shared via a holder), and FlashInfer best-effort capture plus external MPI barrier. The helper lives in ep_bench_common (not run_ep_bench_python) because the backend modules import from common; a helper they call cannot live in the importer without a circular import. Behavior-preserving: verified eager and cuda-graph kernel-only dispatch/combine for all four backends on a 1-node rank-major run.
…hon) Address review comment (move the cuda-graph logic to run_ep_bench_python; unified timing where dispatch_fn/combine_fn come from the lib and the harness does capture -> iterate -> results). The backends no longer build CUDA graphs themselves; each setup_* returns a dict with eager dispatch/combine ops plus an optional "graph" spec (capture-safe dispatch/combine, an optional pre-replay barrier, and an on-capture-failure reset). run_ep_bench_python owns a single _capture_paired_graph() that captures dispatch+combine as ONE graph for any backend, wraps replay, forces the paired kineto pass when it captures, and drops the graph before teardown. Per-library specifics stay in the backend closures: DeepEP cached do_cpu_sync=False dispatch (+ GIN transport gate), NCCL-EP capture-stream refetch, mscclpp dispatch-output-to-combine handoff, FlashInfer barrier-outside-graph (pre_replay) + best-effort rebuild (on_fail). This also generalizes the earlier mscclpp fix: the harness forces EP_KINETO_SEPARATE=0 whenever it actually captures a single graph. Verified: all four backends capture (no eager fallback) and report correct kernel-only dispatch/combine in both eager and cuda-graph on a 1-node rank-major run.
…UDA graph Address review comment (run multiple iterations inside the cuda graph: with torch.cuda.graph(g_all): for i in range(100): dispatch(); combine()). Follows the sglang bench_moe_ep.py pattern: --iters-per-graph N captures N dispatch->combine iterations INSIDE the single graph so one replay runs them all, then divides the measured host time back to per-iteration. This amortizes per-replay launch overhead and keeps the spin-waiting dispatch/ combine kernels from being inflated by per-replay launch skew. Kernel-only kineto is already per-iteration (its per-launch average divides by the kernel count, which scales with N). The harness _capture_paired_graph loops the ops inside the capture; run_backend divides host times by iters_per_graph. Default 1 preserves current behavior. Verified on a 1-node rank-major run: all four backends capture at N=1 and N=10, report consistent per-iteration numbers, and N=10 drops/tightens the host-observed times (mscclpp host Total 53.3->43.7us) as expected.
…azhou/ep_bench_update # Conflicts: # test/python/ep/ep_bench_flashinfer.py # test/python/ep/ep_bench_mscclpp.py # test/python/ep/mscclpp_ep_bench.cu # test/python/ep/run_ep_bench.py # test/python/ep/run_ep_bench_python.py
Capture 10 dispatch->combine iterations per graph by default (was 1) so the CUDA-graph path amortizes launch overhead and avoids per-replay launch skew out of the box; reported times remain per-iteration. Grouping only applies under --cuda-graph, so the non-1 default is auto-clamped to 1 for eager runs instead of erroring. Verified: default --cuda-graph groups by 10 and eager runs cleanly, both reporting consistent per-iteration mscclpp numbers.
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
There was a problem hiding this comment.
Pull request overview
This PR updates the unified in-process Python EP benchmark harness to support four backends (mscclpp EP, NCCL-EP, DeepEP V2, FlashInfer) under a shared dispatch→combine timing methodology, while revising CUDA-graph capture and kernel-time attribution responsibilities.
Changes:
- Moves kineto kernel-name parsing into each backend module (
parse_kineto_kernels) and centralizes the shared summation helper (sum_matching_kernel_us) inep_bench_common. - Refactors CUDA-graph timing to capture a single paired dispatch+combine graph and adds
--iters-per-graphto replay multiple iterations per graph (reporting per-iteration times). - Removes the older C++ LL bench + in-process CUPTI components and related build/driver files.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| test/python/ep/run_ep_bench.py | Deleted legacy external-process driver for C++/NCCL-EP benches (superseded by unified Python harness). |
| test/python/ep/run_ep_bench_python.py | Main harness updates: unified backend contract, CUDA-graph capture ownership, iters-per-graph, backend kernel parsing hooks. |
| test/python/ep/mscclpp_ep_bench.cu | Deleted pure-C++ LL benchmark implementation. |
| test/python/ep/ep_bench_nccl.py | Adopts new backend contract and exposes NCCL-EP-specific kineto parsing + graph capture spec. |
| test/python/ep/ep_bench_mscclpp.py | Adopts new backend contract; provides mscclpp-specific kineto parsing + graph capture spec. |
| test/python/ep/ep_bench_flashinfer.py | Adopts new backend contract; provides FlashInfer-specific kineto parsing + graph capture spec with pre-replay barrier. |
| test/python/ep/ep_bench_deepep.py | Adopts new backend contract; updates DeepEP graph gating (transport-based) and provides kineto parsing + graph capture spec. |
| test/python/ep/ep_bench_common.py | Adds shared sum_matching_kernel_us() helper used by backend-specific parsers. |
| test/python/ep/cupti_kernel_timer.cpp | Deleted in-process CUPTI timer implementation. |
| test/python/ep/CMakeLists.txt | Deleted standalone build for the removed C++ bench + CUPTI timer. |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
| # replays dispatch+combine as ONE combined graph, so this single helper captures | ||
| # the paired op from any backend; the backend only supplies its capture-safe ops. | ||
| # ============================================================================ | ||
| def _capture_paired_graph(dispatch_op, combine_op, prime=True, pre_replay=None, on_fail=None, iters_per_graph=1): |
There was a problem hiding this comment.
Why need pre_replay here?
There was a problem hiding this comment.
This is a special handle for the FlashInfer. "FlashInfer's dispatch/combine each need an MPI host barrier to align ranks, but an MPI barrier cannot live inside a CUDA graph -- so the captured ops are barrier-free and the harness runs the barrier BEFORE each replay (pre_replay)." in ep_bench_flashinfer.py.
There was a problem hiding this comment.
Why it require this? I think dispatch/combine is the sync point. No need additional sync before dispatch
Address review comment: _kineto_kernel_us reads as a confusing name. Rename to torch_profiler_kernel_us, which describes what it does (times the dispatch/combine kernels with torch.profiler). Pure rename, no behavior change.
…p8 note) - Default --graph-group-size 10 -> 50 (reviewer: "Maybe increase to 50 by default?"): capture 50 dispatch->combine iterations per graph by default, further amortizing launch overhead / launch skew. - Unify the internal name to graph_group_size everywhere (reviewer: "Different with iteration_per_group?" and "Why hard code to 1 here?"): the harness used iters_per_graph internally while the CLI arg is --graph-group-size, which read as two different concepts. Rename _capture_paired_graph / run_backend params and the effective-value local to graph_group_size; the "=1" default now reads as "no grouping" (1 iteration captured), which is why it is 1 when a backend is not graph-captured. - Correct the FP8 wording (reviewer: "check if nccl support fp8 right?"): NCCL-EP DOES support FP8 (nccl_ep device code has token_data_type 0=FP8/uint8, calculate_fp8_scales, use_fp8). The previous "NCCL-EP path is bf16 only" text implied the library cannot; reword the --dispatch-dtype help and the guard to say the NCCL-EP path IN THIS BENCHMARK is BF16-only (the harness does not plumb NCCL-EP dispatch scales yet), not the library. Verified: default --cuda-graph captures with graph_group_size=50, per-iteration numbers unchanged; the fp8 guard prints the corrected message.
| # DeepEP CUDA-graph capture is limited by TRANSPORT, not node count. On the | ||
| # all-NVLink / MNNVL path (EP_DISABLE_GIN=1, a single NVL72 domain) the | ||
| # symmetric-memory kernels ARE graph-capturable at any node count -- verified | ||
| # capturing at 1/2/4 nodes on a GB200 NVL72. Only the RDMA/IB scale-out path | ||
| # (GIN enabled, multi-rack) crashes under graph capture (CUDA 719 in | ||
| # symmetric.hpp), so disable capture only when GIN is active. |
There was a problem hiding this comment.
This sounds like a DeepEP bug, this is by design?
There was a problem hiding this comment.
Update the description. Seems the limitation is related to the do_cpu_sync=True.
| # illegal inside a CUDA graph. Replay the CACHED dispatch instead: pass the | ||
| # primed handle (topk_idx reused from it), which forces do_cpu_sync=False and | ||
| # skips the host-side count read, leaving a pure on-stream kernel launch. | ||
| # illegal inside a CUDA graph. Use the CACHED dispatch instead: pass the |
There was a problem hiding this comment.
For non-cached dispatch, I think it will use allgather to get layout, it will break the cuda graph?
There was a problem hiding this comment.
Non-cached dispatch keeps [do_cpu_sync=True] and that is what breaks CUDA-graph capture — not the layout exchange (there's no allgather in the dispatch path; the inter-rank exchange is on-stream inside the GIN kernel). With [do_cpu_sync=True] the runtime:
- does a blocking device→host readback of the exact received-token count, and
- allocates the dispatch outputs to that runtime-determined (dynamic) size (recv_src_metadata.shape[0] in [deepseek-ai\DeepEP\tree\main\deep_ep\buffers\elastic.py]
Both are illegal under graph capture: capture records only async on-stream work (a blocking D2H sync aborts it), and a graph requires fixed, pre-known output shapes/addresses (data-dependent sizing isn't allowed).
| graph_spec = { | ||
| "dispatch": _dispatch, | ||
| "combine": _combine, | ||
| "pre_replay": comm.Barrier, |
There was a problem hiding this comment.
Do we really need this?
…t a bug Reword the DeepEP CUDA-graph gate comment: the RDMA/IB scale-out (GIN/IBGDA) path is not graph-capturable because DeepEP internode transport drives NVSHMEM/IBGDA put-signal operations that are illegal inside a CUDA graph (CUDA 719 in symmetric.hpp). This is a documented DeepEP internode limitation, not a harness bug; we disable capture when GIN is active and run that path eagerly.
…omment DeepEP V2 (ElasticBuffer) scale-out uses NCCL GIN (GPU-Initiated Networking, backed by GDAKI/DOCA GPUNetIO on this stack), not the legacy NVSHMEM/IBGDA Buffer path. Fix the earlier comment that misattributed the graph-capture gate to NVSHMEM/IBGDA put-signal ops and symmetric.hpp. The real on-stream blocker is the dispatch CPU sync for exact recv-token counts (do_cpu_sync); cached dispatch forces it False, which is why the NVLink/MNNVL path is capture-safe.
Per review, keep a single CLI flag for the number of dispatch->combine iterations captured inside one CUDA graph. Rename the arg dest to iters_per_graph, drop the --graph-group-size alias, and update the validation message, help text, comments, and the captured-graph log line accordingly. The internal _capture_paired_graph/run_backend graph_group_size parameter (which receives the value) is unchanged.
749d40d to
41ec252
Compare
|
Unscribe
________________________________
From: Qinghua Zhou ***@***.***>
Sent: Thursday, August 6, 2026 9:20 PM
To: microsoft/mscclpp ***@***.***>
Cc: Subscribed ***@***.***>
Subject: Re: [microsoft/mscclpp] ep(bench): unified in-process Python EP benchmark — review-comment follow-ups (PR #858)
@seagater commented on this pull request.
________________________________
In test/python/ep/ep_bench_deepep.py<#858 (comment)>:
+ # DeepEP CUDA-graph capture is limited by TRANSPORT, not node count. On the
+ # all-NVLink / MNNVL path (EP_DISABLE_GIN=1, a single NVL72 domain) the
+ # symmetric-memory kernels ARE graph-capturable at any node count -- verified
+ # capturing at 1/2/4 nodes on a GB200 NVL72. Only the RDMA/IB scale-out path
+ # (GIN enabled, multi-rack) crashes under graph capture (CUDA 719 in
+ # symmetric.hpp), so disable capture only when GIN is active.
Update the description. Seems the limitation is related to the do_cpu_sync=True.
—
Reply to this email directly, view it on GitHub<#858?email_source=notifications&email_token=AC7VMRKPCZJA7JSW6FP2VD35IVKIRA5CNFSNUABKM5UWIORPF5TWS5BNNB2WEL2QOVWGYUTFOF2WK43UKJSXM2LFO4XTIOBXHE4TCNZVHEYKM4TFMFZW63VKON2WE43DOJUWEZLEUVSXMZLOOSWGM33PORSXEX3DNRUWG2Y#discussion_r3733401969>, or unsubscribe<https://github.com/notifications/unsubscribe-auth/AC7VMRPIRJUFF4NGGSM5EAT5IVKIRAVCNFSNUABFKJSXA33TNF2G64TZHM2TSNRTGMZDONRXHNEXG43VMU5TKMBSGIZDKOBVHE4KC5QC>.
Triage notifications, keep track of coding agent tasks and review pull requests on the go with GitHub Mobile for iOS<https://github.com/notifications/mobile/ios/AC7VMRKF7KSVD64EH23XMIT5IVKIRA5CNFSNUABKM5UWIORPF5TWS5BNNB2WEL2QOVWGYUTFOF2WK43UKJSXM2LFO4XTIOBXHE4TCNZVHEYKM4TFMFZW63VKON2WE43DOJUWEZLEUVSXMZLOOSVGM33PORSXEX3JN5ZQ> and Android<https://github.com/notifications/mobile/android/AC7VMRLGQYLLV5HE6RZSPXL5IVKIRA5CNFSNUABKM5UWIORPF5TWS5BNNB2WEL2QOVWGYUTFOF2WK43UKJSXM2LFO4XTIOBXHE4TCNZVHEYKM4TFMFZW63VKON2WE43DOJUWEZLEUVSXMZLOOSXGM33PORSXEX3BNZSHE33JMQ>. Download it today!
You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>
|
Summary
Follow-up changes on top of the merged unified Python EP benchmark (#836),
addressing the outstanding review comments and refreshing the CUDA-graph
methodology. Benchmarks all four EP backends — mscclpp EP, NVIDIA NCCL-EP,
DeepEP V2, FlashInfer — through one shared in-process harness with identical
dispatch→combine timing.
Changes (by review comment)
parse_kineto_kernels(...); the harness holds zero per-library kernel knowledge(shared
sum_matching_kernel_us()in ep_bench_common).stream.synchronize().owned by the harness (
_capture_paired_graph); eachsetup_*returns a uniform{dispatch, combine, teardown, barrier, graph}dict and one shared loop times allfour backends. Per-library specifics stay in the backend closures.
--iters-per-graph N(sglang bench_moe_ep.py pattern): N dispatch→combine iterations per graph, reported
per iteration.
Other fixes
verified graph capture at 1/2/4 nodes on GB200 NVL72.
(fixes mscclpp "captured 0" under --cuda-graph).
Validation
All four backends: correct kernel-only dispatch/combine in eager + CUDA-graph,
genuine capture (no eager fallback), consistent per-iteration numbers at
--iters-per-graph 1 and 10, on GB200 NVL72.