[Feat] DiffusionEngine Support async batch infer (#2729)

Signed-off-by: Semmer <semmer@live.cn> Signed-off-by: jader <yjader@foxmail.com> Signed-off-by: asukaqaq-s <1311722138@qq.com> Co-authored-by: asukaqaq-s <1311722138@qq.com> Co-authored-by: jader <yjader@foxmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com> Co-authored-by: Gao Han <hgaoaf@connect.ust.hk>
2026-05-06 09:47:13 +12:00
1 parent 2b85474c23
commit e49fbd8a2d
@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import queue
+import threading
+import time
+from dataclasses import dataclass, field
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+
+@dataclass
+class DiffusionSchedulerOutput:
+    step_id: int
+    scheduled_new_reqs: list = field(default_factory=list)
+    scheduled_cached_reqs: Any = None
+    finished_req_ids: set = field(default_factory=set)
+    num_running_reqs: int = 0
+    num_waiting_reqs: int = 0
+
+    @property
+    def scheduled_req_ids(self):
+        ids = [req.request_id for req in self.scheduled_new_reqs]
+        if self.scheduled_cached_reqs:
+            ids.extend(self.scheduled_cached_reqs.sched_req_ids)
+        return ids
+
+    @property
+    def is_empty(self):
+        return len(self.scheduled_req_ids) == 0
+
+
+class MockScheduler:
+    def __init__(self):
+        self._waiting_queue = []
+        self._step_id = 0
+
+    def add_request(self, request):
+        self._waiting_queue.append(request)
+        return request.request_id
+
+    def has_requests(self):
+        return len(self._waiting_queue) > 0
+
+    def schedule(self) -> DiffusionSchedulerOutput:
+        if not self._waiting_queue:
+            return DiffusionSchedulerOutput(step_id=self._step_id)
+
+        batch = []
+        while self._waiting_queue:
+            req = self._waiting_queue.pop(0)
+            batch.append(SimpleNamespace(request_id=req.request_id))
+
+        output = DiffusionSchedulerOutput(step_id=self._step_id, scheduled_new_reqs=batch)
+        self._step_id += 1
+        return output
+
+    def update_from_output(self, sched_output, runner_output):
+        # assume all new req finished
+        return [req.request_id for req in sched_output.scheduled_new_reqs]
+
+
+@pytest.mark.asyncio
+async def test_async_add_req_and_wait_for_response():
+    from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+
+    engine = object.__new__(DiffusionEngine)
+    engine.scheduler = MockScheduler()
+    engine._out_queue = {}
+    engine.abort_queue: queue.Queue[str] = queue.Queue()
+    engine._rpc_lock = threading.Lock()
+
+    engine._finalize_finished_request = lambda rid, out, err: out.result
+
+    def mock_execute_batch(sched_output):
+        req_ids = sched_output.scheduled_req_ids
+
+        time.sleep(1)
+
+        class MockRunnerOutput:
+            def __init__(self, ids):
+                self.req_id = ids
+                self.step_index = 0
+                self.finished = True
+                self._results = {rid: SimpleNamespace(result_data=f"data_{rid}") for rid in ids}
+
+            def get_req_output(self, rid):
+                return SimpleNamespace(result=self._results[rid], step_index=0, finished=True)
+
+        return MockRunnerOutput(req_ids)
+
+    engine.execute_fn = mock_execute_batch
+
+    engine.stop_event = threading.Event()
+    engine.start_background_loop()
+
+    async def run_task(rid):
+        req = SimpleNamespace(request_id=rid)
+        start = time.time()
+        res = await engine.async_add_req_and_wait_for_response(req)
+        return rid, res, time.time() - start
+
+    task_ids = [f"req_{i}" for i in range(5)]
+    tasks = [run_task(rid) for rid in task_ids]
+    results = await asyncio.gather(*tasks)
+
+    engine.stop_event.set()
+    engine.worker_thread.join(timeout=5)
+    assert len(results) == 5
+    for rid, res, elapsed in results:
+        assert rid in res.result_data
+
+    eps = 0.5
+    latencies = [r[2] for r in results]
+    assert max(latencies) - min(latencies) < eps
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

+import asyncio
 import queue
 import threading
 from types import SimpleNamespace
@@ -20,6 +21,7 @@ from vllm_omni.diffusion.sched import (
    StepScheduler,
 )
 from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
+from vllm_omni.diffusion.worker.utils import RunnerOutput
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams

 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
@@ -34,7 +36,7 @@ def _make_request(req_id: str) -> OmniDiffusionRequest:


 def _make_request_output(req_id: str, *, error: str | None = None, finished: bool = True):
-    return SimpleNamespace(
+    return RunnerOutput(
        req_id=req_id,
        step_index=None,
        finished=finished,
@@ -49,7 +51,7 @@ def _make_step_output(
    finished: bool = False,
    error: str | None = None,
 ):
-    return SimpleNamespace(
+    return RunnerOutput(
        req_id=req_id,
        step_index=step_index,
        finished=finished,
@@ -91,6 +93,7 @@ class _StubScheduler(SchedulerInterface):
        self._sched_req_id = request.request_ids[0]
        self._state = None
        self._scheduled = False
+        self.max_num_running_reqs = 1

    def initialize(self, od_config) -> None:
        self.initialized_with = od_config
@@ -218,6 +221,46 @@ class TestRequestScheduler:
        assert third.num_running_reqs == 1
        assert third.num_waiting_reqs == 0

+    def test_batches_compatible_requests_up_to_max_num_seqs(self) -> None:
+        scheduler = RequestScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+        req_id_a = scheduler.add_request(_make_request("a"))
+        req_id_b = scheduler.add_request(_make_request("b"))
+
+        sched_output = scheduler.schedule()
+
+        assert _new_ids(sched_output) == [req_id_a, req_id_b]
+        assert sched_output.num_running_reqs == 2
+        assert sched_output.num_waiting_reqs == 0
+
+    def test_incompatible_waiting_head_blocks_later_compatible_request(self) -> None:
+        scheduler = RequestScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=3))
+
+        req_id_a = scheduler.add_request(_make_request("a"))
+        req_id_b = scheduler.add_request(
+            OmniDiffusionRequest(
+                prompts=["prompt_b"],
+                sampling_params=OmniDiffusionSamplingParams(width=768),
+                request_ids=["b"],
+            )
+        )
+        scheduler.add_request(_make_request("c"))
+
+        first = scheduler.schedule()
+
+        assert _new_ids(first) == [req_id_a]
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 2
+
+        scheduler.update_from_output(first, _make_request_output(req_id_a))
+        second = scheduler.schedule()
+
+        assert _new_ids(second) == [req_id_b]
+        assert second.num_running_reqs == 1
+        assert second.num_waiting_reqs == 1
+
    def test_abort_request_for_waiting_and_running(self) -> None:
        req_id_a = self.scheduler.add_request(_make_request("a"))
        req_id_b = self.scheduler.add_request(_make_request("b"))
@@ -348,18 +391,25 @@ class TestDiffusionEngine:
        assert req_id in finished
        assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED

-    def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None:
+    @pytest.mark.asyncio
+    async def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None:
        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine._loop_started = False
+        engine._init_lock = asyncio.Lock()
+        engine.main_loop = asyncio.get_running_loop()
+        engine.stop_event = threading.Event()
        engine.pre_process_func = None
-        engine.add_req_and_wait_for_response = mocker.Mock(
+        engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
            return_value=DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
        )

        with pytest.raises(DiffusionRequestAbortedError, match="Request req-abort aborted"):
-            engine.step(_make_request("req-abort"))
+            await engine.step(_make_request("req-abort"))

    def test_abort_queue_marks_request_finished_aborted(self) -> None:
        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine._rpc_lock = threading.RLock()
+        engine._cv = threading.Condition(engine._rpc_lock)
        engine.scheduler = RequestScheduler()
        engine.scheduler.initialize(SimpleNamespace())
        engine.abort_queue = queue.Queue()
@@ -519,7 +569,7 @@ class TestStepScheduler:
        sched_output = self.scheduler.schedule()
        finished = self.scheduler.update_from_output(
            sched_output,
-            SimpleNamespace(
+            RunnerOutput(
                req_id=req_id,
                step_index=None,
                finished=True,
@@ -576,6 +626,64 @@ class TestStepScheduler:
        assert finished == {req_id}
        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_ABORTED

+    def test_batches_compatible_step_requests(self) -> None:
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+        req_a = scheduler.add_request(_make_step_request("a"))
+        req_b = scheduler.add_request(_make_step_request("b"))
+
+        sched_output = scheduler.schedule()
+
+        assert _new_ids(sched_output) == [req_a, req_b]
+        assert sched_output.num_running_reqs == 2
+        assert sched_output.num_waiting_reqs == 0
+
+    def test_step_batch_allows_different_num_inference_steps(self) -> None:
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=2))
+
+        req_a = scheduler.add_request(_make_step_request("a", num_inference_steps=2))
+        req_b = scheduler.add_request(_make_step_request("b", num_inference_steps=4))
+
+        sched_output = scheduler.schedule()
+
+        assert _new_ids(sched_output) == [req_a, req_b]
+        assert sched_output.num_running_reqs == 2
+        assert sched_output.num_waiting_reqs == 0
+
+    def test_step_batch_rejects_different_sampling_key(self) -> None:
+        scheduler = StepScheduler()
+        scheduler.initialize(SimpleNamespace(max_num_seqs=3))
+
+        req_a = scheduler.add_request(_make_step_request("a"))
+        req_b = scheduler.add_request(
+            _make_step_request(
+                "b",
+                sampling_params=OmniDiffusionSamplingParams(
+                    height=768,
+                    num_inference_steps=4,
+                ),
+            )
+        )
+        scheduler.add_request(_make_step_request("c"))
+
+        sched_output = scheduler.schedule()
+
+        assert _new_ids(sched_output) == [req_a]
+        assert sched_output.num_running_reqs == 1
+        assert sched_output.num_waiting_reqs == 2
+
+        scheduler.update_from_output(
+            sched_output,
+            _make_step_output(req_a, step_index=4, finished=True),
+        )
+        second = scheduler.schedule()
+
+        assert _new_ids(second) == [req_b]
+        assert second.num_running_reqs == 1
+        assert second.num_waiting_reqs == 1
+
    def test_preempt_request_preserves_step_index(self) -> None:
        request = _make_step_request("preempt", num_inference_steps=3)
        req_id = self.scheduler.add_request(request)
@@ -77,12 +77,12 @@ class _StepPipeline:
        self.prepare_calls += 1
        state.timesteps = [torch.tensor(10), torch.tensor(5)]
        state.latents = torch.tensor([0.0])
+        state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
        return state

-    def denoise_step(self, state, **kwargs):
-        del state, kwargs
+    def denoise_step(self, input_batch, **kwargs):
        self.denoise_calls += 1
-        return torch.tensor([1.0])
+        return torch.full_like(input_batch.prompt_embeds, fill_value=0.5)

    def step_scheduler(self, state, noise_pred, **kwargs):
        del noise_pred, kwargs
@@ -145,6 +145,7 @@ class _DistributedStepPipeline(CFGParallelMixin):
        state.step_index = 0
        state.scheduler = self.scheduler
        state.do_true_cfg = self.mode == "cfg"
+        state.prompt_embeds = torch.tensor([[0.0, 0.0], [1.0, 1.0]])
        return state

    def denoise_step(self, state, **kwargs):
@@ -291,6 +292,7 @@ def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine:
    engine.scheduler = scheduler
    engine.execute_fn = execute_fn
    engine._rpc_lock = threading.RLock()
+    engine._cv = threading.Condition(engine._rpc_lock)
    engine.abort_queue = queue.Queue()
    return engine

@@ -327,10 +329,11 @@ def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master
            raise ValueError(f"Unsupported distributed test mode: {mode}")

        runner = _make_distributed_runner(mode, device)
-        output = DiffusionModelRunner.execute_stepwise(
+        result = DiffusionModelRunner.execute_stepwise(
            runner,
            _make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
        )
+        output = result.get_req_output("req-1")

        assert output.finished is True
        assert output.result is not None
@@ -354,14 +357,16 @@ class TestRunner:
        req = _make_step_request()
        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)

-        first = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
+        result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
+        first = result.get_req_output("req-1")
        assert first.req_id == "req-1"
        assert first.step_index == 1
        assert first.finished is False
        assert first.result is None
        assert "req-1" in runner.state_cache

-        second = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
+        result = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
+        second = result.get_req_output("req-1")
        assert second.req_id == "req-1"
        assert second.step_index == 2
        assert second.finished is True
@@ -393,8 +398,8 @@ class TestRunner:
            num_waiting_reqs=0,
        )

-        with pytest.raises(ValueError, match="batch_size=1"):
-            DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+        result = DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+        assert len(result) == 2

    def test_rejects_missing_cached_state(self):
        runner = _make_runner()
@@ -408,8 +413,8 @@ class TestRunner:
        req = _make_step_request()
        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)

-        output = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
-
+        result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
+        output = result.get_req_output("req-1")
        assert output.req_id == "req-1"
        assert output.step_index == 0
        assert output.finished is True
@@ -580,7 +585,7 @@ class TestEngine:
    )
    def test_step_engine_returns_error(self, execute_fn, expected_error, mocker: MockerFixture):
        scheduler = StepScheduler()
-        scheduler.initialize(mocker.Mock())
+        scheduler.initialize(SimpleNamespace())
        engine = _make_engine(scheduler, execute_fn=execute_fn)

        output = engine.add_req_and_wait_for_response(_make_engine_request("req-error", num_inference_steps=2))
@@ -590,7 +595,7 @@ class TestEngine:

    def test_step_execution_completes(self, mocker: MockerFixture):
        scheduler = StepScheduler()
-        scheduler.initialize(mocker.Mock())
+        scheduler.initialize(SimpleNamespace())
        engine = _make_engine(scheduler)
        request = _make_engine_request("req-step", num_inference_steps=2)

@@ -616,7 +621,7 @@ class TestEngine:

    def test_step_abort_stops_rescheduling_after_first_step(self, mocker: MockerFixture):
        scheduler = StepScheduler()
-        scheduler.initialize(mocker.Mock())
+        scheduler.initialize(SimpleNamespace())
        engine = _make_engine(scheduler)
        request = _make_engine_request("req-stop", num_inference_steps=4)

@@ -641,7 +646,7 @@ class TestEngine:

    def test_step_abort_after_reschedule_returns_aborted_output(self, mocker: MockerFixture):
        scheduler = StepScheduler()
-        scheduler.initialize(mocker.Mock())
+        scheduler.initialize(SimpleNamespace())
        engine = _make_engine(scheduler)
        request = _make_engine_request("req-mid", num_inference_steps=4)

@@ -668,7 +673,7 @@ class TestEngine:

    def test_finished_step_without_result_returns_error(self, mocker: MockerFixture):
        scheduler = StepScheduler()
-        scheduler.initialize(mocker.Mock())
+        scheduler.initialize(SimpleNamespace())
        engine = _make_engine(
            scheduler,
            execute_fn=lambda _: RunnerOutput(
@@ -19,6 +19,7 @@ from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.sched import RequestScheduler
 from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput

 pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]
@@ -34,7 +35,11 @@ def _tagged_output(tag: str) -> DiffusionOutput:

 def _mock_request(tag: str):
    """Return a lightweight request object identifiable by *tag*."""
-    return SimpleNamespace(request_ids=[tag])
+    return SimpleNamespace(
+        request_ids=[tag],
+        prompts=[f"prompt_{tag}"],
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
+    )


 def _make_executor(num_gpus: int = 1):
@@ -406,7 +411,7 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
        t.start()

        # Wait until request execution is truly inside the lock and blocking.
-        add_req_blocked.wait(5)
+        assert add_req_blocked.wait(5)

        # collective_rpc should time out at lock acquisition, not hang.
        with pytest.raises(TimeoutError):
@@ -2,8 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

 import asyncio
+import time
 from concurrent.futures import ThreadPoolExecutor
-from dataclasses import asdict
+from dataclasses import asdict, dataclass
 from types import SimpleNamespace

 import pytest
@@ -18,7 +19,7 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
    async def run_test():
        captured = {}

-        def step(request):
+        async def step(request):
            captured["request"] = request
            return [
                SimpleNamespace(
@@ -73,3 +74,75 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
        assert result.images == ["img-1", "img-2"]

    asyncio.run(run_test())
+
+
+@dataclass
+class MockOmniRequestOutput:
+    request_id: str = ""
+    status: str = "success"
+
+
+BASE_HEIGHT = 512
+BASE_WIDTH = 512
+BASE_INFER_STEPS = 10
+DELAY_BASE = 2
+
+
+class MockDiffusionEngine:
+    async def step(self, request):
+        def simulate_step_delay(height, width, num_inference_steps) -> float:
+            return (height / BASE_HEIGHT) * (width / BASE_WIDTH) * (num_inference_steps / BASE_INFER_STEPS)
+
+        DELAY_BASE = 2
+        delay_scale = simulate_step_delay(
+            request.sampling_params.height, request.sampling_params.width, request.sampling_params.num_inference_steps
+        )
+        delay = DELAY_BASE + delay_scale * DELAY_BASE
+        await asyncio.sleep(delay)
+        return [MockOmniRequestOutput(request_id=request.request_id)]
+
+
+@pytest.mark.asyncio
+async def test_proc_process_request_with_batching_async_output():
+    stage_proc = object.__new__(StageDiffusionProc)
+    stage_proc._engine = MockDiffusionEngine()
+
+    test_requests = [
+        {
+            "request_id": "req_1",
+            "prompt": "prompt1",
+            "params": {"height": BASE_HEIGHT * 2, "width": BASE_WIDTH * 2, "num_inference_steps": BASE_INFER_STEPS * 1},
+        },
+        {
+            "request_id": "req_2",
+            "prompt": "prompt2",
+            "params": {"height": BASE_HEIGHT * 2, "width": BASE_WIDTH * 2, "num_inference_steps": BASE_INFER_STEPS * 2},
+        },
+        {
+            "request_id": "req_3",
+            "prompt": "prompt3",
+            "params": {"height": BASE_HEIGHT * 2, "width": BASE_WIDTH * 2, "num_inference_steps": BASE_INFER_STEPS * 3},
+        },
+    ]
+
+    async def run_task(req_data):
+        start_time = time.time()
+        result = await stage_proc._process_request(
+            request_id=req_data["request_id"], prompt=req_data["prompt"], sampling_params_dict=req_data["params"]
+        )
+        end_time = time.time()
+        return result, end_time - start_time
+
+    coros = [run_task(req) for req in test_requests]
+    results = await asyncio.gather(*coros)
+
+    assert len(results) == len(test_requests)
+    base_time = DELAY_BASE
+    time_gap_std = DELAY_BASE * 2 * 2 * 1  # height/width/steps infer time scale
+    eps = 0.5
+    for i, (res, elapsed_time) in enumerate(results):
+        assert res.request_id == test_requests[i]["request_id"]
+        assert isinstance(res, MockOmniRequestOutput)
+        time_gap = elapsed_time - base_time
+        assert time_gap > time_gap_std - eps and time_gap < time_gap_std + eps
+        base_time = elapsed_time