Skip to content

Fix DeepEP v2 ElasticBuffer construction inside CUDA graph capture - #39078

Open
whn09 wants to merge 1 commit into
sgl-project:mainfrom
whn09:fix/deepep-v2-prebuild-elastic-buffer
Open

whn09 wants to merge 1 commit into
sgl-project:mainfrom
whn09:fix/deepep-v2-prebuild-elastic-buffer

Conversation

@whn09

@whn09 whn09 commented Sep 11, 2026

Copy link
Copy Markdown
Contributor

Motivation

DeepEPv2Buffer.get_buffer() is reached from exactly one place — _DeepEPv2Impl.dispatch()
(token_dispatcher/deepep_v2.py:329, and :405 for combine). So the process-wide
ElasticBuffer is constructed on the first dispatch, and for a decode CUDA graph that
first dispatch happens inside torch.cuda.graph().

ElasticBuffer.__init__ reaches ncclDevCommCreate (DeepEP
csrc/kernels/backend/nccl.cu:188), which allocates, synchronises and calls into the driver
— none of it legal during stream capture. CUDA invalidates the capture and NCCL surfaces its
generic backend error:

Exception: Capture cuda graph failed: NCCL exception
(/opt/DeepEP/csrc/kernels/backend/nccl.cu:188): 3
  (GIN: DevComm setup failed on all available backends)

That message names GIN, so it invites a GIN-tuning response, and it is not a GIN problem:
the identical construction with the identical arguments succeeds when it happens one moment
earlier, outside the capture. I lost about a day to this before instrumenting it.

Observed rather than argued. A standalone probe builds the same buffer with the same
arguments, once outside a capture and once inside, on 8× H200 with TP=EP=8:

outside capture inside capture
--deepep-v2-mode direct, NCCL_GIN_TYPE=2 OK, 0.19 GiB cudaErrorStreamCaptureInvalidated
--deepep-v2-mode hybrid, NCCL_GIN_TYPE=5 OK, 0.19 GiB cudaErrorStreamCaptureInvalidated

So on this fabric it is neither the GIN type nor the mode — it is the ordering.

Scope, stated honestly, because it is narrower than "everyone hits this". Whether the
lazy build survives the capture depends on the GIN backend, and I have a counter-example of
my own: a single-node 8× B300 host on InfiniBand with NCCL_GIN_TYPE=3 and
--deepep-v2-mode direct captured decode graphs and served happily on an unpatched image
(3 launches, capture 61.59 s / 18.48 GB pool). Same sglang, same lazy build, no failure. It
fails on EFA hosts — both the type-2 CPU proxy and type-5 GDAKI, single- and two-node.

So the honest claim is not "this always breaks", it is: the current code depends on the GIN
backend tolerating an allocating, synchronising, collective constructor during stream capture,
which CUDA does not promise anyone.
Where it works today it works by luck, and the probe
above shows the graph is genuinely invalidated the moment you look. Every _get_buffer() call
site in the tree is inside dispatch()/combine(), so there is no configuration in which a
user can get the buffer built early themselves.

And the one place this can fire is the one place the workaround is unavailable. For
deepep_v2, prefill never captures at all — disable_tc_piecewise_cudagraph_if_incompatible
turns off tc_piecewise for any moe_a2a_backend != "none", and
disable_breakable_cudagraph_if_incompatible lists ("none", "deepep", "megamoe", "flashinfer") as validated, which deepep_v2 is not. So the only capture in a
deepep_v2 server is the decode graph, and "just pass --disable-cuda-graph" means running
a decode role without graphs. Nobody deploys that: it is ~3× end-to-end on this model and
3.27× on the B300/IB host above, and it is every token, not just TTFT. SGLang's own error
text says as much ("Not recommended. Huge performance loss").

So on EFA this is not "a bug with a workaround" — deepep_v2 decode is unusable, and where
the bug is currently latent the workaround is equally unacceptable.

Modifications

  1. BaseRunner.warmup() gets a third pre-initialize slot,
    _pre_initialize_deepep_v2_elastic_buffer(), immediately after the two that already exist
    for precisely this reason — _pre_initialize_flashinfer_allreduce_workspace() ("must run
    before CG capture to keep broadcasts/barriers outside the capture context") and
    _pre_initialize_fi_a2a_workspace() ("must run before CG capture (it syncs the stream +
    barriers cross-rank, uncapturable)"). Same class of problem, same slot, so no new concept
    is introduced. warmup() is called from both decode_cuda_graph_runner.capture() and
    prefill_cuda_graph_runner.capture(), so one insertion covers both.

    It returns immediately unless get_moe_a2a_backend().is_deepep_v2(), and walks
    model.modules() for the first DeepEPv2Dispatcher. One call is both necessary and
    sufficient: the buffer is process-wide and its construction is collective, so every rank
    must build exactly once, in lockstep.

  2. New DeepEPv2Dispatcher.prepare_buffer(), so the runner is not reaching through two
    private attributes (dispatcher._impl._get_buffer()).

    It delegates to _impl._get_buffer() rather than re-deriving the constructor arguments,
    and that is the load-bearing detail: get_buffer() caches on an exact 7-element key
    (group, hidden_size, router_topk, num_max_dispatch_tokens_per_rank, use_fp8_dispatch, allow_hybrid_mode, world_size). A pre-build assembled from recomputed arguments that
    differ in any one element is a cache miss, so the dispatcher would go on to build a
    second buffer inside the capture — reintroducing this exact bug while appearing to have
    fixed it.

I also considered the existing model.prepare_before_cuda_graph_capture hook (used by
qwen3_5.py). It did not seem right here: that hook is for model-owned resources and would
need re-implementing per architecture, whereas the ElasticBuffer is owned by the dispatcher
and is process-wide. Happy to move it if you prefer.

Accuracy Tests

Before/after on 8× H200, --moe-a2a-backend deepep_v2 --deepep-v2-mode hybrid, TP=EP=8,
decode CUDA graphs enabled:

result
unpatched Capture cuda graph failed: NCCL exception (nccl.cu:188): 3 at startup
this PR Initialized DeepEP v2 ElasticBuffer: ... num_bytes=... logged before Capture cuda graph begin, capture completes, server serves

Numerics after the fix were checked separately against the same model on
--moe-a2a-backend deepep (v1) by comparing per-token logprobs on a fixed 12-prompt set:
10 of 12 rows bit-identical, the 2 differing rows being the only two longer than the v2
arm's prefill chunk. Not part of this PR's claim — the change adds no math — but it is why I
am confident the pre-built buffer is the same buffer the dispatch path would have built.

I have not added a unit test. Reproducing it needs a real multi-rank ElasticBuffer, a live
NCCL DevComm and an active capture, which is well outside what CI exercises today (nothing in
CI exercises deepep_v2 at all, which is also how #37211 stayed broken). If you would like
the standalone probe contributed as a manual/registered test I am glad to send it — it builds
the buffer twice and asserts the in-capture attempt raises, and it runs in ~30 s.

Speed Tests and Profiling

Not applicable in the usual sense: unpatched, the configuration cannot start. Against
--disable-cuda-graph, which is the only other way to run it, keeping decode graphs is worth
roughly 3× end-to-end.

The pre-build itself is one collective allocation that used to happen anyway, moved earlier,
so steady-state performance is unchanged.

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

cc @MengYu10151 (a3ae667, #35634 — the DeepEPv2 ElasticBuffer backend). I cannot add the
run-ci label or trigger CI myself, so this needs someone with write access to start it.

Related: #37211 fixes the other thing that stops deepep_v2 from completing a forward pass
(ep_scatter_from_psum missing two kernel args) and is still open. The two are independent —
this one is capture-time, that one is forward-time — but a deepep_v2 server needs both.

DeepEPv2Buffer.get_buffer() is only ever called from _DeepEPv2Impl.dispatch(), so
the process-wide ElasticBuffer is constructed on the first dispatch -- which for a
decode graph is inside torch.cuda.graph(). ElasticBuffer.__init__ reaches
ncclDevCommCreate, which allocates, synchronises and calls into the driver, all
illegal during stream capture, so CUDA invalidates the capture and NCCL reports its
generic backend failure.

Add a third _pre_initialize_* slot in BaseRunner.warmup(), alongside the flashinfer
allreduce and fi_a2a workspaces that exist for exactly this reason, and a public
DeepEPv2Dispatcher.prepare_buffer() for it to call.
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.

1 participant