Repository navigation
[Bugfix][MoE] Keep DeepEP high-throughput buffers alive across DBO streams - #60735
AndreasKaratzas merged 4 commits into
Conversation
…reams Under DBO, DeepEPHTPrepareAndFinalize runs DeepEP dispatch and combine on the shared ubatch comm stream with allocate_on_comm_stream=False. The dispatch outputs (recv x, x scales, topk ids and weights) are therefore allocated in the comm stream's caching-allocator pool but read by the experts on the compute stream, and combine reads the reduced expert output that was allocated on the compute stream. Nothing recorded these cross-stream uses, so a block could be freed and handed to the other ubatch's next DeepEP call (or to compute work) while kernels queued on the other stream still read it. With the DeepGEMM MoE this produced garbage expert ids and an illegal memory access in ep_scatter. Record the consuming stream on every tensor that crosses between the DBO compute and comm streams. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Anirudh Ravichandran <ani@primeintellect.ai>
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
|
@SageMoore could you take a look? DeepEP-HT outputs are allocated on the DBO comm stream but used on the compute stream without record_stream, so the allocator can reuse them too early. cc @LucasWilkinson |
|
This also fixes an intermittent ROCm failure. MI355 Distributed Features crashed in main nightly build 95634 on On MI355 (gfx950, DeepEP
|
|
/ci run |
|
✅ Triggered Buildkite CI #95673 for commit |
…lready orders On ROCm the comm stream is host-synced right after HT dispatch and combine, so the dispatch inputs and the combine input are no longer read when they can be freed. Only the dispatch outputs read on the compute stream still need recording. Each record_stream costs about 10 us on ROCm, so this removes about 4% TPOT with DBO + DeepEP HT on MI355. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Andreas Karatzas <akaratza@amd.com>
|
I pushed 6b323f5 on top, plus a merge of main. On ROCm, The skipped calls weren't free: each MI355, DeepSeek-V2-Lite-Chat, DP2 + EP,
6b323f5 also passed 10/10 under the compute-stream delay that faults main 10/10. One of the 9 main runs crashed with the same memory access fault without any delay. |
|
/ci run |
|
✅ Triggered Buildkite CI #95740 for commit |
…reams (vllm-project#60735) Signed-off-by: Anirudh Ravichandran <ani@primeintellect.ai> Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Overview
Under DBO,
DeepEPHTPrepareAndFinalizehas DeepEP allocate its dispatch outputs on the DBO comm stream, then reads them on the compute stream; the combine input crosses the other way. The caching allocator is never told about these cross-stream reads, so the other microbatch's next DeepEP call can reuse a block and overwrite it before it is read. With the DeepGEMM FP8 MoE this crashes the engine at startup with an illegal memory access inep_scatter. This PR callsrecord_streamon every tensor that crosses between the two streams.A GitHub search found no existing issue or PR for this. Related: #46990 added a ROCm-only host sync after HT dispatch/combine under DBO, possibly for the same hazard (see Details). #27666 and #38451 fixed other DBO + DeepEP HT problems.
Claims
VLLM_USE_DEEP_GEMM=1 VLLM_MOE_USE_DEEP_GEMM=1, the defaults). For GLM-5.3 FP8 on main, 2 of 2 runs crashed in the first microbatched step without the fix. On 0.29.0, so did all 7 runs that had neither the fix nor a serializing workaround (CUDA_LAUNCH_BLOCKING, no caching allocator, host sync). With the fix, full GSM8K scores 0.927 with CUTLASS linears and 0.912 with DeepGEMM linears (DBO off: 0.923).--enable-dbo, main runs Model Runner V1 unlessVLLM_USE_V2_MODEL_RUNNER=1is set. MRV2's opt-in DBO uses the same microbatch contexts and streams, so it gets the same fix. Only Model Runner V1 was run.record_streamcalls per MoE layer per microbatch, plus one CUDA event per recorded block when it is freed; the allocator reuses that block only after the event completes. The throughput cost was not measured.deepep_ll.pyis unchanged. The open [DBO] Add DBO support to the DeepEP V2 backend #52177 adds the same comm-stream switching to DeepEP v2, so it may need the same treatment.Validation
On main, without and with this PR. Setup: GLM-5.3 FP8 on 1 x 8 H200, DP4 x TP2, EP8,
--all2all-backend deepep_high_throughput --enable-dbo,cudagraph_mode NONE, fp8 KV, prefix caching, Model Runner V1. Env:VLLM_USE_DEEP_GEMM=1 VLLM_MOE_USE_DEEP_GEMM=1 VLLM_USE_DEEP_GEMM_E8M0=0 VLLM_DEEP_GEMM_WARMUP=skip; the logs showUsing DEEPGEMM Fp8 MoE backend. Build: main @ 4a705c7 + #60732 + #60734. Running GLM-5.x with DBO needs both of those; this PR is independent of them. Eval:lm_evalGSM8K 5-shot, 1,319 questions,num_concurrent=256, standard error ±0.007–0.008. Microbatched steps were counted with a diagnostic log that is not part of this PR.--linear-backend cutlass)DeepGEMM-linears DBO on and DBO off differ by 0.010, about one standard error of the difference. We did not run CUTLASS linears with DBO off, or DeepGEMM linears with DBO on without the fix, on main. On 0.29.0 the latter crashes the same way (see below).
The crash without the fix:
RuntimeError: DeepEP error: CPU recv timeoutinbuffer.dispatch. The other microbatch thread getsCUDA error: an illegal memory access was encounteredatcomm_stream.wait_event, and the run ends withtorch.cat(): expected a non-empty list of Tensors.CUDA Exception: Warp Out-of-range Addressin_fwd_kernel_ep_scatter_2<<<(8192,1,1),(256,1,1)>>>atdeep_gemm_utils.py:109, which istl.load(expert_map + expert_id)inapply_expert_map, inlined from:241(expert_id = tl.load(recv_topk + ...)). The faulting SASS loadsexpert_map + expert_id * 4after theexpert_id != -1check, so the expert id read from thetorch.whereover DeepEP's topk ids was out of range.Isolating the cause on vLLM 0.29.0. Same 1-node GLM-5.3 setup, with backports of #60732 and #60734, DeepGEMM MoE and CUTLASS linears unless noted. The record_stream rows use an env-gated version of this patch: same three call sites, using
torch.cuda.current_stream().ep_gatherint64 fix from fe284ea (#53914); DeepGEMM or CUTLASS linears; with a core dumpdeep_gemm_utils.py:110on 0.29)VLLM_DBO_COMM_SMS=0VLLM_RANDOMIZE_DP_DUMMY_INPUTS=1CUDA_LAUNCH_BLOCKING=1PYTORCH_NO_CUDA_MEMORY_CACHING=1Taken together, these runs narrow the cause:
record_streamalone each remove the crash.That points at allocator reuse of DeepEP's buffers across DBO streams, which only the DeepGEMM MoE turns into a crash.
Production shape (vLLM 0.29.0). Setup: 4 x 8 H200, one engine per GPU (TP1, DP32/EP32, external load balancing), GLM-5.3-BF16 quantized online with
fp8_per_block, DeepEP HT, eager, prefix caching, fp8 KV,--max-num-batched-tokens 4096,--gpu-memory-utilization 0.70. Each row is one run with the same backports as above. Prefill load: three waves of 1,536 unique prompts of about 9.7k tokens each,max_tokens=1, 6 in flight per rank. GSM8K (1,319 questions) was then run against rank 0.* Two earlier attempts with FlashInfer autotune on stalled during startup autotuning. Throughput is the third wave over all 32 GPUs. DeepGEMM kernels JIT-compile during the first waves (
VLLM_DEEP_GEMM_WARMUP=skip), and all three runs were still speeding up in wave 3, so these are not steady-state numbers. This table shows that DBO + DeepGEMM MoE works at production scale with the fix. It is not a performance claim for this PR. Without the fix this configuration does not start, so measuring the fix's own cost needs a MoE backend that runs either way, such as FLASHINFER_CUTLASS. We did not do that.CI coverage. This PR adds no test. The existing DBO test,
tests/v1/distributed/test_dbo.py(DeepSeek-V2-Lite-Chat, DP2), already runs the changed lines in itsdeepep_high_throughputcase, with a BF16 MoE (TritonExperts) rather than the DeepGEMM path. It runs in the 2 x H100 "Distributed Features" job (optional), the B200 "Distributed" job (also optional; xfailed on Blackwell for unstable DBO accuracy; we have not checked whether that is related) and the AMD mirrors. With this patch, on 2 x H200 it passes both cases (deepep_low_latency,deepep_high_throughput) in 9 minutes, with 200+ microbatched steps per worker in the HT case. To build a 2-GPU reproducer, we quantized that model's MoE online to block FP8 with--quantization online --quantization-config '{"moe": "fp8_per_block"}' --moe-backend deep_gemm. The logs confirm the DEEPGEMM MoE behindDeepEPHTPrepareAndFinalize. That variant passed without the fix every time, so it cannot catch a regression:--max-num-batched-tokens 16384 --max-num-seqs 256, 512 questionsIt would add DeepGEMM-MoE coverage of the patched path for 5–7 minutes on 2 GPUs. We can add it if reviewers want it, but it would not have caught this bug. The regression evidence is the manual GLM-5.3 runs above. Separately, quantizing V2-Lite's linears as well (
--quantization fp8_per_block) scored about 0 on GSM8K in this DBO setup, both with and without this PR. That is unrelated and was not investigated.Lint: pre-commit passes on the rebased commit: ruff check and format (0.14.0), typos (1.43.5), mypy (1.20.2), and the local SPDX, forbidden-import,
torch.cudaand boolean-context-manager hooks. Before the rebase, with torch types resolved, mypy reported 2 errors, which also appear on the unpatched file.Not run:
VLLM_USE_V2_MODEL_RUNNER=1). It uses the same code path.test_dbo.pywith DeepEP HT, so CI will exercise the change there.tests/kernels/moe/test_deepep*.pywith this patch. They do not enable DBO, so the change is a no-op in them.git patch-id), wheredeepep_ht.py,ubatching.py,gpu_ubatch_wrapper.pyandgpu/ubatch_utils.pyare unchanged.deep_gemm_utils.py,deep_gemm_moe.pyandmodular_kernel.pychanged in between ([Bugfix] Widen ep_gather output offset to int64 #58373, [Perf][MoE] Skip padded work in block-FP8 DeepGEMM experts #59128, [Perf][MoE] Eliminate staging from NCCL symmetric reduce-scatter #49194). On the rebased commit we re-ranpytest -v -s tests/v1/distributed/test_dbo.pyon 2 x H200: 2 passed (deepep_low_latency,deepep_high_throughput) in 9 minutes. The GLM-5.3 runs were not repeated on the newer base.Reproduction commands
Details
Root cause (main @ 240785b, before this PR):
UBatchWrappercreates one comm stream (gpu_ubatch_wrapper.py:76) and gives both microbatches that comm stream and the same compute stream (:288-291). MRV2'sUBatchRunnerdoes the same (gpu/ubatch_utils.py:327,458-462)._do_dispatchswitches to the comm stream (deepep_ht.py:126). It then callsget_dispatch_layoutanddispatchwithallocate_on_comm_stream=False(:139,168) and, under DBO,async_finish=False(:138,167). DeepEP therefore allocates recv x, scales, topk ids and topk weights on the current stream, which is the DBO comm stream's allocator pool. It runs its kernels on its own internal stream, which waits only onprevious_event. Withasync_finish=Falseit then makes the calling stream wait for that internal stream, and it records no other stream on its tensors: its HT dispatch and combine callrecord_streamonly whenasync_finish=True(csrc/deep_ep.cpp).dbo_switch_to_compute_sync()(:177) orders the compute stream after the dispatch, and_receiverreads the outputs there: first thetorch.whereover the topk ids (:225), then the experts. Once those kernels are enqueued the Python references are dropped; for example,expert_topk_idsis rebound by thetorch.where. The blocks go back to the comm pool, free for any later comm-stream work.previous_event, captured on the compute stream before it yielded ([Perf] Fix DBO overlap: capture DeepEP event before yield #38451). The DBO comm stream itself waits only on that microbatch's compute-done event (ubatching.py:86-87, recorded inyield_and_switch_from_compute_to_commbefore the yield). Both events were recorded before this microbatch enqueued its reads. So that call's outputs can land in the freed blocks and be written before the reads run. The DeepGEMM contiguous path then loadsexpert_map[expert_id]for whatever id it finds (deep_gemm_utils.py:107-110, called from_fwd_kernel_ep_scatter_2at:244-247), which is the faulting load in the core dump.TopKWeightAndReduceDelegate),_finalizereduces into a new compute-stream tensor (apply(output=None),:365), and combine reads it on the comm stream (:376).dbo_switch_to_compute()(:389) does not order the compute stream after the combine. The tensor is freed when_finalizereturns, so later compute-stream work, such as shared experts or the other microbatch, can reuse the block while combine is still reading it. On main, though, no experts class that can run behind DeepEP HT delegates: the delegating classes all use the batched activation format, DeepEP HT uses the standard one, and the standard CUTLASS FP8 and Marlin experts returnTopKWeightAndReduceNoOP. Every experts class that can run behind DeepEP HT reduces itself (DeepGEMM, Triton, FlashInfer CUTLASS, ...), so the combine input is the per-microbatch workspace (or the caller's output tensor when the shapes match), which is not freed there. The_finalizecall (1 of the 9) is therefore defensive on main, for a future delegating experts class. We did not observe this direction.dbo_switch_to_compute_sync()makes the compute stream wait on the comm stream. They are recorded anyway (4 of the 9 calls), so their safety does not depend on that ordering; this is also the version validated on 0.29.Fix:
_dbo_record_stream()callst.record_stream(current_stream())on each non-Nonetensor whendbo_enabled(). It is called in three places:_do_dispatch, after switching to comm, on the dispatch inputs (compute to comm);_receiver, before thetorch.where, on DeepEP's outputs (comm to compute);_finalize, after switching to comm, on the reduced expert output (compute to comm).When a recorded tensor is freed, the allocator records an event on the consuming stream and does not reuse the block until that event has completed. For the two compute-to-comm sites the actual reader is DeepEP's internal stream. Recording the DBO comm stream still covers it, because with
async_finish=Falsethe DBO comm stream waits for DeepEP's kernels before dispatch or combine returns, and the tensors are freed after that. DeepEP's handle and layout tensors andcombined_xare used only on the DBO comm stream and DeepEP's internal stream, so they are not recorded (but see "Not covered").CUDA graphs: DeepEP HT is never captured. vLLM forces
cudagraph_modeNONE for HT with DP > 1, and DBO needs DP > 1. PyTorch also supportsrecord_streamduring capture.Why FLASHINFER_CUTLASS does not crash, but may still race. The hazard is in the prepare/finalize, so it does not depend on the experts backend; what differs is what a corrupted id does. FlashInfer's CUTLASS MoE receives
ep_rank/ep_sizeand keeps only ids in this rank's range (expert >= start_expert && expert < end_expert,cutlass_fused_moe_kernels.cuh:223,449in FlashInfer 0.7.0.post1). A garbage id is therefore dropped instead of dereferenced, and garbage activations or weights would only produce wrong numbers. With FLASHINFER_CUTLASS, DBO matches DBO off on GSM8K: 0.922–0.927 on main in #60734, and 0.924 at EP32 above. So we have no evidence that the race fires there in practice. We also do not know why the window opens reliably for GLM-5.3 with DeepGEMM but not for DeepSeek-V2-Lite with DeepGEMM. #60734 notes an illegal memory access with DeepSeek-V3.2 and DBO that also uses FLASHINFER_CUTLASS and also disappears withCUDA_LAUNCH_BLOCKING=1. We have not tested whether this PR fixes it.Why ROCm already host-syncs. #46990 added
_sync_dbo_comm_if_needed(deepep_ht.py:61,71-76), a hostsynchronize()of the current stream after HT dispatch and combine under DBO, on ROCm only. Its comment cites reuse of "Buffer-owned communication workspace". Forcing the same sync on CUDA also removed the crash (0.29 table). We suspect ROCm hit the same allocator hazard, but we have no ROCm hardware to check, so this PR leaves the ROCm sync in place; on ROCm,record_streamonly adds bookkeeping.Alternatives considered
allocate_on_comm_stream=True. This moves the outputs to DeepEP's internal stream, but the compute stream still reads them, so on its own it does not order reuse. DeepEP also accepts it only together withprevious_eventandasync_finish=True, and vLLM disablesasync_finishunder DBO. Not tried.record_stream(this PR). It marks exactly the tensors that cross streams, adds no syncs or kernels, and keeps the fix inside the one prepare/finalize that switches streams under DBO.Not covered: blocks freed on the comm stream itself, for example
combined_xafteroutput.copy_in the finalize receiver, are reused by later comm-stream allocations. DeepEP's kernels for that later call wait onprevious_event, a compute-stream event, rather than on the comm stream, so in principle they could run before the earlier comm-stream copy finishes. We inferred this from DeepEP's source (its internal stream waits onprevious_eventwhen one is given, otherwise on the calling stream) and have not observed it.record_streamcannot express this ordering. Passingprevious_event=Noneunder DBO would make DeepEP wait on the comm stream instead. That is left for a follow-up because it changes the ordering #38451 tuned.Related:
BLOCK_SIZE_Kin the batched Triton FP8 MoE.The first three came out of the same GLM-5.3 bring-up and are independent of this PR. The
ep_gatherint64 fix in fe284ea (#53914) is already on main; the 0.29 no-fix runs crash with and without a backport of it.AI assistance: Claude Code was used to trace the root cause, write the patch, run the validations above, and draft this description. I reviewed every changed line, can explain and defend the change, and checked the test results above.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
🤖 Generated with Claude Code