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 14, 01:45 UTC
[Bug] Fix DBO IMA issue for DeepEPHT (#27666)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Oct 29, 2025 at 16:28 UTC
b5d90f740048d43376390a61ca5b77c287505d0e
2 files changed
+18
-3
vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
+9
-3
@@ -16,6 +16,7 @@ from vllm.utils.math_utils import round_up
16
from vllm.v1.worker.ubatching import (
17
dbo_current_ubatch_id,
18
dbo_enabled,
19
+ dbo_get_previous_event,
20
dbo_switch_to_comm,
21
dbo_switch_to_compute,
22
dbo_switch_to_compute_sync,
@@ -110,6 +111,10 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
111
# for the other ubatch before the dispatch kernel starts.
112
dbo_yield_and_switch_from_compute_to_comm()
113
114
+ # capture a DeepEP event and pass it as previous_event so
115
+ # DeepEP honors the dependency internally.
116
+ previous_event = dbo_get_previous_event(self.buffer.capture)
117
+
118
(
119
num_tokens_per_rank,
120
num_tokens_per_rdma_rank,
@@ -119,7 +124,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
124
) = self.buffer.get_dispatch_layout(
125
topk_idx=rank_topk_ids,
126
num_experts=num_experts,
122
- previous_event=None,
127
+ previous_event=previous_event,
128
async_finish=False,
129
allocate_on_comm_stream=False,
130
)
@@ -148,7 +153,7 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
153
# to this value.
154
expert_alignment=1,
155
config=self._get_dispatch_config(),
151
- previous_event=None,
156
+ previous_event=previous_event,
157
async_finish=self.async_prepare and not dbo_enabled(),
158
allocate_on_comm_stream=False,
159
)
@@ -339,13 +344,14 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
344
assert fused_expert_output.dtype == torch.bfloat16, (
345
f"Expected fused_expert_output bfloat16, got {fused_expert_output.dtype}"
346
)
347
+ previous_event = dbo_get_previous_event(self.buffer.capture)
348
combined_x, _, event = self.buffer.combine(
349
# HT combine only supports BF16
350
x=fused_expert_output,
351
handle=handle,
352
topk_weights=None,
353
config=self._get_combine_config(),
348
- previous_event=None,
354
+ previous_event=previous_event,
355
async_finish=do_async and not dbo_enabled(),
356
allocate_on_comm_stream=False,
357
)
vllm/v1/worker/ubatching.py
+9
@@ -185,6 +185,15 @@ def dbo_register_recv_hook(recv_hook):
185
next_ctx.recv_hook = recv_hook
186
187
188
+def dbo_get_previous_event(func, *args, **kwargs):
189
+ if len(_THREAD_ID_TO_CONTEXT) > 0:
190
+ ctx_idx = _THREAD_ID_TO_CONTEXT[threading.get_ident()]
191
+ ctx = _CURRENT_CONTEXTS[ctx_idx]
192
+ # execute callable on the ubatch compute stream to record/wait events there
193
+ with torch.cuda.stream(ctx.compute_stream):
194
+ return func(*args, **kwargs)
195
+
196
+
197
def make_ubatch_contexts(
198
num_micro_batches: int,
199
compute_stream: torch.cuda.Stream,