A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 17:15 UTC
[Perf] Fix DBO overlap: capture DeepEP event before yield (#38451)
Signed-off-by: root <conway.zhu@cohere.com>
czhu-cohere committed
Mar 31, 2026 at 13:38 UTC
517b769b5858a8d8d233d277f54461acfc9def63
1 file changed
+7
-5
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py
+7
-5
@@ -107,15 +107,17 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
107
) -> Callable:
108
has_scales = token_scales is not None
109
110
+ # Capture a DeepEP event on the compute stream before yielding.
111
+ # This must happen before the yield so the event only covers this
112
+ # ubatch's compute work. If captured after, the compute stream tail
113
+ # may include the other ubatch's work, preventing overlap.
114
+ previous_event = dbo_get_previous_event(self.buffer.capture)
115
+
116
# We yield before launching the dispatch kernel since the dispatch
117
# kernel will block the CPU so we want to queue up all the compute
118
# for the other ubatch before the dispatch kernel starts.
119
dbo_yield_and_switch_from_compute_to_comm()
120
115
- # capture a DeepEP event and pass it as previous_event so
116
- # DeepEP honors the dependency internally.
117
- previous_event = dbo_get_previous_event(self.buffer.capture)
118
-
121
(
122
num_tokens_per_rank,
123
num_tokens_per_rdma_rank,
@@ -357,11 +359,11 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
359
topk_ids=topk_ids,
360
apply_router_weight_on_input=apply_router_weight_on_input,
361
)
362
+ previous_event = dbo_get_previous_event(self.buffer.capture)
363
dbo_yield_and_switch_from_compute_to_comm()
364
assert fused_expert_output.dtype == torch.bfloat16, (
365
f"Expected fused_expert_output bfloat16, got {fused_expert_output.dtype}"
366
)
364
- previous_event = dbo_get_previous_event(self.buffer.capture)
367
combined_x, _, event = self.buffer.combine(
368
# HT combine only supports BF16
369
x=fused_expert_output,