diff --git a/tests/kernels/moe/test_deepep_v2_async_finalize.py b/tests/kernels/moe/test_deepep_v2_async_finalize.py new file mode 100644 index 0000000000..193d8906f3 --- /dev/null +++ b/tests/kernels/moe/test_deepep_v2_async_finalize.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Contract tests for DeepEPV2PrepareAndFinalize.finalize_async. + +Hermetic: fake buffer/event, single process, no GPU collectives — only +requires deep_ep to be importable. + + T1 finalize_async -> combine is issued with async_with_compute_stream=True + and the join is deferred to the receiver, which waits + on the combine event (device-side) before the copy. + T2 DBO active -> combine falls back to async_with_compute_stream=False; + the receiver still produces the correct output. + T3 finalize (sync) -> combine is issued synchronously and the output is + copied immediately. +""" + +import pytest +import torch + +from vllm.utils.import_utils import has_deep_ep_v2 + +requires_deep_ep_v2 = pytest.mark.skipif( + not has_deep_ep_v2(), + reason="Requires DeepEP v2 (ElasticBuffer)", +) + +if has_deep_ep_v2(): + import vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 as _dv2 + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_v2 import ( + DeepEPV2PrepareAndFinalize, + ) + from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( + TopKWeightAndReduceContiguous, + ) + + +class _FakeEvent: + def __init__(self, has_event: bool): + self.event = object() if has_event else None + self.waited = False + + def current_stream_wait(self): + assert self.event is not None + self.waited = True + + +class _FakeBuffer: + def __init__(self, out: torch.Tensor): + self.out = out + self.calls: list[dict] = [] + self.last_event: _FakeEvent | None = None + + def combine( + self, + x, + handle, + topk_weights, + async_with_compute_stream, + allocate_on_comm_stream, + ): + self.calls.append({"async_with_compute_stream": async_with_compute_stream}) + # DeepEP returns a completion event only for an async combine. + self.last_event = _FakeEvent(has_event=async_with_compute_stream) + return self.out, None, self.last_event + + +def _make_pf(out: torch.Tensor): + pf = DeepEPV2PrepareAndFinalize( + buffer=_FakeBuffer(out), + num_dispatchers=1, + dp_size=1, + rank_expert_offset=0, + num_experts=8, + num_topk=2, + ) + pf.handles[0] = object() + return pf + + +def _run(pf, output: torch.Tensor, do_async: bool): + empty = torch.empty(0, dtype=torch.bfloat16) + weights = torch.empty(0, 2) + ids = torch.empty(0, 2, dtype=torch.int64) + fn = pf.finalize_async if do_async else pf.finalize + return fn(output, empty, weights, ids, False, TopKWeightAndReduceContiguous()) + + +@requires_deep_ep_v2 +def test_finalize_async_defers_join_to_receiver(): + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + recv = _run(pf, dst, do_async=True) + assert pf.buffer.calls[0]["async_with_compute_stream"] is True + assert not torch.equal(dst, out_src), "join must not happen before receiver" + recv() + assert pf.buffer.last_event.waited, "receiver must join via event wait" + assert torch.equal(dst, out_src) + + +@requires_deep_ep_v2 +def test_finalize_async_under_dbo_uses_sync_combine(monkeypatch): + monkeypatch.setattr(_dv2, "dbo_enabled", lambda: True) + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + recv = _run(pf, dst, do_async=True) + assert pf.buffer.calls[0]["async_with_compute_stream"] is False + recv() + assert not pf.buffer.last_event.waited, "no event to wait on for sync combine" + assert torch.equal(dst, out_src) + + +@requires_deep_ep_v2 +def test_finalize_sync(): + out_src = torch.full((4, 8), 7.0, dtype=torch.bfloat16) + pf = _make_pf(out_src) + dst = torch.zeros_like(out_src) + ret = _run(pf, dst, do_async=False) + assert ret is None + assert pf.buffer.calls[0]["async_with_compute_stream"] is False + assert torch.equal(dst, out_src) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py index d834c6eff7..b3a146274d 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_v2.py @@ -21,6 +21,7 @@ from vllm.triton_utils import tl, triton from vllm.utils.math_utils import round_up from vllm.v1.worker.ubatching import ( dbo_current_ubatch_id, + dbo_enabled, ) @@ -89,8 +90,10 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): use cudagraphs anyway - Provides expert_tokens_meta for efficient batched expert kernels - Both modes use async_with_compute_stream=False (synchronous from - caller's perspective). The ElasticBuffer handles comm internally. + Dispatch always uses async_with_compute_stream=False. finalize_async + issues the combine with async_with_compute_stream=True (except under + DBO) so the modular kernel can overlap the shared-expert FFN with the + combine a2a; the returned receiver joins via a device-side event wait. """ @staticmethod @@ -465,15 +468,31 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): f"got {fused_expert_output.dtype}" ) + # DBO drives its own hook/receiver schedule; keep the combine + # synchronous there (the receiver then only performs the copy). + combine_async = do_async and not dbo_enabled() combined_x, _, event = self.buffer.combine( x=fused_expert_output, handle=handle, topk_weights=None, - async_with_compute_stream=False, + async_with_compute_stream=combine_async, + allocate_on_comm_stream=combine_async, ) - output.copy_(combined_x, non_blocking=True) - return None + if do_async: + # The combine ran on DeepEP's comm stream; the modular kernel + # issues the shared-expert FFN before calling the receiver, which + # joins via a device-side cudaStreamWaitEvent (no host sync, so + # this is safe inside a captured region). + def _receiver(): + if event.event is not None: + event.current_stream_wait() + output.copy_(combined_x, non_blocking=True) + + return _receiver + else: + output.copy_(combined_x, non_blocking=True) + return None def finalize_async( self, @@ -484,16 +503,17 @@ class DeepEPV2PrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): apply_router_weight_on_input: bool, weight_and_reduce_impl: mk.TopKWeightAndReduce, ) -> Callable: - self._finalize( + receiver = self._finalize( output, fused_expert_output, topk_weights, topk_ids, apply_router_weight_on_input, weight_and_reduce_impl, - False, + True, ) - return lambda: None + assert receiver is not None + return receiver def finalize( self,