Skip to content

ep(bench): unified in-process Python EP benchmark — review-comment follow-ups - #858

Open
Qinghua Zhou (seagater) wants to merge 23 commits into
feature/epfrom
qinghuazhou/ep_bench_update
Open

ep(bench): unified in-process Python EP benchmark — review-comment follow-ups#858
Qinghua Zhou (seagater) wants to merge 23 commits into
feature/epfrom
qinghuazhou/ep_bench_update

Conversation

@seagater

Copy link
Copy Markdown
Contributor

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)

  • Move the kernel-name parse into the backend files — each backend owns
    parse_kineto_kernels(...); the harness holds zero per-library kernel knowledge
    (shared sum_matching_kernel_us() in ep_bench_common).
  • remove sync here — dropped the redundant warmup stream.synchronize().
  • Move cudaGraph logic to run_ep_bench_python / unified timing — capture is
    owned by the harness (_capture_paired_graph); each setup_* returns a uniform
    {dispatch, combine, teardown, barrier, graph} dict and one shared loop times all
    four backends. Per-library specifics stay in the backend closures.
  • run multiple iterations inside the cuda graph — added --iters-per-graph N
    (sglang bench_moe_ep.py pattern): N dispatch→combine iterations per graph, reported
    per iteration.

Other fixes

  • DeepEP CUDA-graph gate keyed on transport (EP_DISABLE_GIN), not node count —
    verified graph capture at 1/2/4 nodes on GB200 NVL72.
  • Force the paired kineto pass whenever the harness captures a single graph
    (fixes mscclpp "captured 0" under --cuda-graph).
  • Removed the C++ LL bench, shell driver, and CUPTI timer.
  • Clarified EP_KINETO_SEPARATE=0 is the required (not "legacy") CUDA-graph mode.

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.

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

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

@seagater Qinghua Zhou (seagater) changed the title ep(bench): unified in-process Python EP benchmark — review-comment follow-upsQinghuazhou/ep bench update ep(bench): unified in-process Python EP benchmark — review-comment follow-ups Jul 30, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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) in ep_bench_common.
  • Refactors CUDA-graph timing to capture a single paired dispatch+combine graph and adds --iters-per-graph to 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.

Comment thread test/python/ep/run_ep_bench_python.py
Comment thread test/python/ep/run_ep_bench_python.py Outdated
Qinghua Zhou (seagater) and others added 2 commits July 30, 2026 10:40
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>
Comment thread test/python/ep/run_ep_bench_python.py Outdated
Comment thread test/python/ep/run_ep_bench_python.py
Comment thread test/python/ep/run_ep_bench_python.py Outdated
Comment thread test/python/ep/run_ep_bench_python.py
Comment thread test/python/ep/run_ep_bench_python.py Outdated
# 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):

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why need pre_replay here?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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‎.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why it require this? I think dispatch/combine is the sync point. No need additional sync before dispatch

Comment thread test/python/ep/run_ep_bench_python.py Outdated
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.
Comment thread test/python/ep/ep_bench_deepep.py Outdated
Comment on lines +113 to +118
# 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.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This sounds like a DeepEP bug, this is by design?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For non-cached dispatch, I think it will use allgather to get layout, it will break the cuda graph?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. does a blocking device→host readback of the exact received-token count, and
  2. 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,

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we really need this?

Comment thread test/python/ep/run_ep_bench_python.py Outdated
…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.
@Andyli1007

Andy li (Andyli1007) commented Aug 7, 2026 via email

Copy link
Copy Markdown

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants