| 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) |