[Kernel][MoE] DeepEP v2: async finalize to overlap shared experts with combine (#52781)
Signed-off-by: Mario <404mario@users.noreply.github.com> Co-authored-by: Mario <404mario@users.noreply.github.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
404mario committed
Sep 15, 2026 at 12:04 UTC
9d14fd3095eb169044e2da0a2a98dd1f2c77ce22
2 files changed
+150
-8
tests/kernels/moe/test_deepep_v2_async_finalize.py
new
+122
@@ -0,0 +1,122 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Contract tests for DeepEPV2PrepareAndFinalize.finalize_async.
4
+
5
+Hermetic: fake buffer/event, single process, no GPU collectives — only
6
+requires deep_ep to be importable.
7
+
8
+ T1 finalize_async -> combine is issued with async_with_compute_stream=True
9
+ and the join is deferred to the receiver, which waits
10
+ on the combine event (device-side) before the copy.
11
+ T2 DBO active -> combine falls back to async_with_compute_stream=False;
12
+ the receiver still produces the correct output.
13
+ T3 finalize (sync) -> combine is issued synchronously and the output is
14
+ copied immediately.
15
+"""
16
+
17
+import pytest
18
+import torch
19
+
20
+from vllm.utils.import_utils import has_deep_ep_v2
21
+
22
+requires_deep_ep_v2 = pytest.mark.skipif(
23
+ not has_deep_ep_v2(),
24
+ reason="Requires DeepEP v2 (ElasticBuffer)",
25
+)
26
+
27
+if has_deep_ep_v2():
28
+ import vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 as _dv2
29
+ from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import (
30
+ DeepEPV2PrepareAndFinalize,
31
+ )
32
+ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
33
+ TopKWeightAndReduceContiguous,
34
+ )
35
+
36
+
37
+class _FakeEvent:
38
+ def __init__(self, has_event: bool):
39
+ self.event = object() if has_event else None
40
+ self.waited = False
41
+
42
+ def current_stream_wait(self):
43
+ assert self.event is not None
44
+ self.waited = True
45
+
46
+
47
+class _FakeBuffer:
48
+ def __init__(self, out: torch.Tensor):
49
+ self.out = out
50
+ self.calls: list[dict] = []
51
+ self.last_event: _FakeEvent | None = None
52
+
53
+ def combine(
54
+ self,
55
+ x,
56
+ handle,
57
+ topk_weights,
58
+ async_with_compute_stream,
59
+ allocate_on_comm_stream,
60
+ ):
61
+ self.calls.append({"async_with_compute_stream": async_with_compute_stream})
62
+ # DeepEP returns a completion event only for an async combine.
63
+ self.last_event = _FakeEvent(has_event=async_with_compute_stream)
64
+ return self.out, None, self.last_event
65
+
66
+
67
+def _make_pf(out: torch.Tensor):
68
+ pf = DeepEPV2PrepareAndFinalize(
69
+ buffer=_FakeBuffer(out),
70
+ num_dispatchers=1,
71
+ dp_size=1,
72
+ rank_expert_offset=0,
73
+ num_experts=8,
74
+ num_topk=2,
75
+ )
76
+ pf.handles[0] = object()
77
+ return pf
78
+
79
+
80
+def _run(pf, output: torch.Tensor, do_async: bool):
81
+ empty = torch.empty(0, dtype=torch.bfloat16)
82
+ weights = torch.empty(0, 2)
83
+ ids = torch.empty(0, 2, dtype=torch.int64)
84
+ fn = pf.finalize_async if do_async else pf.finalize
85
+ return fn(output, empty, weights, ids, False, TopKWeightAndReduceContiguous())
86
+
87
+
88
+@requires_deep_ep_v2
89
+def test_finalize_async_defers_join_to_receiver():
90
+ out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16)
91
+ pf = _make_pf(out_src)
92
+ dst = torch.zeros_like(out_src)
93
+ recv = _run(pf, dst, do_async=True)
94
+ assert pf.buffer.calls[0]["async_with_compute_stream"] is True
95
+ assert not torch.equal(dst, out_src), "join must not happen before receiver"
96
+ recv()
97
+ assert pf.buffer.last_event.waited, "receiver must join via event wait"
98
+ assert torch.equal(dst, out_src)
99
+
100
+
101
+@requires_deep_ep_v2
102
+def test_finalize_async_under_dbo_uses_sync_combine(monkeypatch):
103
+ monkeypatch.setattr(_dv2, "dbo_enabled", lambda: True)
104
+ out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16)
105
+ pf = _make_pf(out_src)
106
+ dst = torch.zeros_like(out_src)
107
+ recv = _run(pf, dst, do_async=True)
108
+ assert pf.buffer.calls[0]["async_with_compute_stream"] is False
109
+ recv()
110
+ assert not pf.buffer.last_event.waited, "no event to wait on for sync combine"
111
+ assert torch.equal(dst, out_src)
112
+
113
+
114
+@requires_deep_ep_v2
115
+def test_finalize_sync():
116
+ out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16)
117
+ pf = _make_pf(out_src)
118
+ dst = torch.zeros_like(out_src)
119
+ ret = _run(pf, dst, do_async=False)
120
+ assert ret is None
121
+ assert pf.buffer.calls[0]["async_with_compute_stream"] is False
122
+ assert torch.equal(dst, out_src)
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py
+28
-8
@@ -21,6 +21,7 @@ from vllm.triton_utils import tl, triton
21
from vllm.utils.math_utils import round_up
22
from vllm.v1.worker.ubatching import (
23
dbo_current_ubatch_id,
24
+ dbo_enabled,
25
)
26
27
@@ -89,8 +90,10 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
90
use cudagraphs anyway
91
- Provides expert_tokens_meta for efficient batched expert kernels
92
92
- Both modes use async_with_compute_stream=False (synchronous from
93
- caller's perspective). The ElasticBuffer handles comm internally.
93
+ Dispatch always uses async_with_compute_stream=False. finalize_async
94
+ issues the combine with async_with_compute_stream=True (except under
95
+ DBO) so the modular kernel can overlap the shared-expert FFN with the
96
+ combine a2a; the returned receiver joins via a device-side event wait.
97
"""
98
99
@staticmethod
@@ -465,15 +468,31 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
468
f"got {fused_expert_output.dtype}"
469
)
470
471
+ # DBO drives its own hook/receiver schedule; keep the combine
472
+ # synchronous there (the receiver then only performs the copy).
473
+ combine_async = do_async and not dbo_enabled()
474
combined_x, _, event = self.buffer.combine(
475
x=fused_expert_output,
476
handle=handle,
477
topk_weights=None,
472
- async_with_compute_stream=False,
478
+ async_with_compute_stream=combine_async,
479
+ allocate_on_comm_stream=combine_async,
480
)
481
475
- output.copy_(combined_x, non_blocking=True)
476
- return None
482
+ if do_async:
483
+ # The combine ran on DeepEP's comm stream; the modular kernel
484
+ # issues the shared-expert FFN before calling the receiver, which
485
+ # joins via a device-side cudaStreamWaitEvent (no host sync, so
486
+ # this is safe inside a captured region).
487
+ def _receiver():
488
+ if event.event is not None:
489
+ event.current_stream_wait()
490
+ output.copy_(combined_x, non_blocking=True)
491
+
492
+ return _receiver
493
+ else:
494
+ output.copy_(combined_x, non_blocking=True)
495
+ return None
496
497
def finalize_async(
498
self,
@@ -484,16 +503,17 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular):
503
apply_router_weight_on_input: bool,
504
weight_and_reduce_impl: mk.TopKWeightAndReduce,
505
) -> Callable:
487
- self._finalize(
506
+ receiver = self._finalize(
507
output,
508
fused_expert_output,
509
topk_weights,
510
topk_ids,
511
apply_router_weight_on_input,
512
weight_and_reduce_impl,
494
- False,
513
+ True,
514
)
496
- return lambda: None
515
+ assert receiver is not None
516
+ return receiver
517
518
def finalize(
519
self,