[Feat] Support step-boundary abort in diffusion (#1769)

Signed-off-by: jader <yjader@foxmail.com>
Signed-off-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: jader <yjader@foxmail.com>

tests/diffusion/test_diffusion_scheduler.py

@@ -1,12 +1,15 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import queue
 import threading
+from types import SimpleNamespace
 from unittest.mock import Mock, patch
 
 import pytest
+import torch
 
-from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
 from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
@@ -14,6 +17,7 @@ from vllm_omni.diffusion.sched import (
     RequestScheduler,
     Scheduler,
     SchedulerInterface,
+    StepScheduler,
 )
 from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -29,9 +33,46 @@ def _make_request(req_id: str) -> OmniDiffusionRequest:
     )
 
 
-def _make_request_output(req_id: str, *, error: str | None = None) -> DiffusionOutput:
-    del req_id
-    return DiffusionOutput(output=None, error=error)
+def _make_request_output(req_id: str, *, error: str | None = None, finished: bool = True):
+    return SimpleNamespace(
+        req_id=req_id,
+        step_index=None,
+        finished=finished,
+        result=DiffusionOutput(output=None, error=error),
+    )
+
+
+def _make_step_output(
+    req_id: str,
+    step_index: int,
+    *,
+    finished: bool = False,
+    error: str | None = None,
+):
+    return SimpleNamespace(
+        req_id=req_id,
+        step_index=step_index,
+        finished=finished,
+        result=DiffusionOutput(output=None, error=error) if error is not None else None,
+    )
+
+
+def _make_step_request(
+    req_id: str,
+    *,
+    num_inference_steps: int = 4,
+    step_index: int | None = None,
+    sampling_params: OmniDiffusionSamplingParams | None = None,
+) -> OmniDiffusionRequest:
+    return OmniDiffusionRequest(
+        prompts=[f"prompt_{req_id}"],
+        sampling_params=sampling_params
+        or OmniDiffusionSamplingParams(
+            num_inference_steps=num_inference_steps,
+            step_index=step_index,
+        ),
+        request_ids=[req_id],
+    )
 
 
 def _new_ids(sched_output) -> list[str]:
@@ -43,7 +84,7 @@ def _cached_ids(sched_output) -> list[str]:
 
 
 class _StubScheduler(SchedulerInterface):
-    def __init__(self, request: OmniDiffusionRequest, output: DiffusionOutput) -> None:
+    def __init__(self, request: OmniDiffusionRequest, output) -> None:
         self._request = request
         self._output = output
         self.initialized_with = None
@@ -75,9 +116,10 @@ class _StubScheduler(SchedulerInterface):
             is_empty=False,
         )
 
-    def update_from_output(self, sched_output, output: DiffusionOutput) -> set[str]:
+    def update_from_output(self, sched_output, output) -> set[str]:
         del sched_output
         assert output is self._output
+        self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
         return {self._sched_req_id}
 
     def has_requests(self) -> bool:
@@ -185,9 +227,14 @@ class TestRequestScheduler:
         state_b = self.scheduler.get_request_state(req_id_b)
         assert state_b.status == DiffusionRequestStatus.FINISHED_ABORTED
 
+        first = self.scheduler.schedule()
+        assert first.finished_req_ids == {req_id_b}
         # A should still run normally.
-        output_a = self.scheduler.schedule()
-        assert _new_ids(output_a) == [req_id_a]
+        assert _new_ids(first) == [req_id_a]
+
+        # B is already marked finished aborted, scheduling again should not pull it.
+        second = self.scheduler.schedule()
+        assert second.finished_req_ids == set()
 
         # Abort running request.
         self.scheduler.finish_requests(req_id_a, DiffusionRequestStatus.FINISHED_ABORTED)
@@ -233,33 +280,33 @@ class TestDiffusionEngine:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine.scheduler = RequestScheduler()
         engine.scheduler.initialize(Mock())
-        engine.executor = Mock()
-        engine._rpc_lock = threading.Lock()
+        engine._rpc_lock = threading.RLock()
+        engine.abort_queue = queue.Queue()
 
         request = _make_request("engine")
-        expected = DiffusionOutput(output=None)
-        engine.executor.add_req.return_value = expected
+        runner_output = _make_request_output("engine")
+        engine.execute_fn = Mock(return_value=runner_output)
 
         output = engine.add_req_and_wait_for_response(request)
 
-        assert output is expected
-        engine.executor.add_req.assert_called_once_with(request)
+        assert output is runner_output.result
+        engine.execute_fn.assert_called_once()
 
     def test_supports_scheduler_interface_injection(self) -> None:
         request = _make_request("engine_iface")
-        expected = DiffusionOutput(output=None)
-        scheduler = _StubScheduler(request, expected)
+        runner_output = _make_request_output("engine_iface")
+        scheduler = _StubScheduler(request, runner_output)
 
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine.scheduler = scheduler
-        engine.executor = Mock()
-        engine.executor.add_req = Mock(return_value=expected)
-        engine._rpc_lock = threading.Lock()
+        engine._rpc_lock = threading.RLock()
+        engine.abort_queue = queue.Queue()
+        engine.execute_fn = Mock(return_value=runner_output)
 
         output = engine.add_req_and_wait_for_response(request)
 
-        assert output is expected
-        engine.executor.add_req.assert_called_once_with(request)
+        assert output is runner_output.result
+        engine.execute_fn.assert_called_once()
 
     def test_initializes_injected_scheduler(self) -> None:
         request = _make_request("init")
@@ -289,6 +336,59 @@ 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) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.pre_process_func = None
+        engine.add_req_and_wait_for_response = Mock(
+            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"))
+
+    def test_abort_queue_marks_request_finished_aborted(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = RequestScheduler()
+        engine.scheduler.initialize(Mock())
+        engine.abort_queue = queue.Queue()
+
+        req_id = engine.scheduler.add_request(_make_request("req-abort"))
+        engine.abort("req-abort")
+        engine._process_aborts_queue()
+
+        assert engine.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_ABORTED
+
+    def test_finalize_finished_request_returns_aborted_output(self) -> None:
+        engine = DiffusionEngine.__new__(DiffusionEngine)
+        engine.scheduler = RequestScheduler()
+        engine.scheduler.initialize(Mock())
+
+        req_id = engine.scheduler.add_request(_make_request("req-finalize"))
+        engine.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED)
+
+        output = engine._finalize_finished_request(req_id)
+
+        assert output.aborted is True
+        assert output.abort_message == "Request req-finalize aborted."
+
+    def test_initializes_step_scheduler_when_step_execution_enabled(self) -> None:
+        od_config = Mock(model_class_name="mock_model")
+        od_config.step_execution = True
+        fake_executor = Mock()
+        fake_executor_cls = Mock(return_value=fake_executor)
+
+        with (
+            patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", return_value=None),
+            patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", return_value=None),
+            patch("vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", return_value=fake_executor_cls),
+            patch.object(DiffusionEngine, "_dummy_run", return_value=None),
+        ):
+            engine = DiffusionEngine(od_config)
+
+        assert isinstance(engine.scheduler, StepScheduler)
+        assert engine.execute_fn is fake_executor.execute_step
+        fake_executor_cls.assert_called_once_with(od_config)
+
     def test_dummy_run_raises_on_output_error(self) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine.od_config = Mock(model_class_name="mock_model")
@@ -297,3 +397,240 @@ class TestDiffusionEngine:
 
         with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
             engine._dummy_run()
+
+
+class TestStepScheduler:
+    def setup_method(self) -> None:
+        self.scheduler: StepScheduler = StepScheduler()
+        self.scheduler.initialize(Mock())
+
+    def test_single_request_step_lifecycle(self) -> None:
+        request = _make_step_request("step", num_inference_steps=3)
+        req_id = self.scheduler.add_request(request)
+
+        first = self.scheduler.schedule()
+        assert _new_ids(first) == [req_id]
+        assert _cached_ids(first) == []
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 0
+
+        finished = self.scheduler.update_from_output(first, _make_step_output(req_id, step_index=1))
+        assert finished == set()
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.RUNNING
+        assert request.sampling_params.step_index == 1
+        assert self.scheduler.has_requests() is True
+
+        second = self.scheduler.schedule()
+        assert _new_ids(second) == []
+        assert _cached_ids(second) == [req_id]
+        assert second.num_running_reqs == 1
+        assert second.num_waiting_reqs == 0
+
+        finished = self.scheduler.update_from_output(second, _make_step_output(req_id, step_index=2))
+        assert finished == set()
+        assert request.sampling_params.step_index == 2
+
+        third = self.scheduler.schedule()
+        assert _new_ids(third) == []
+        assert _cached_ids(third) == [req_id]
+
+        finished = self.scheduler.update_from_output(
+            third,
+            _make_step_output(req_id, step_index=3, finished=True),
+        )
+        assert finished == {req_id}
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
+        assert request.sampling_params.step_index == 3
+        assert self.scheduler.has_requests() is False
+
+    def test_fifo_single_request_scheduling(self) -> None:
+        req_id_a = self.scheduler.add_request(_make_step_request("a", num_inference_steps=2))
+        req_id_b = self.scheduler.add_request(_make_step_request("b", num_inference_steps=2))
+
+        first = self.scheduler.schedule()
+        assert _new_ids(first) == [req_id_a]
+        assert _cached_ids(first) == []
+        assert first.num_running_reqs == 1
+        assert first.num_waiting_reqs == 1
+
+        finished = self.scheduler.update_from_output(first, _make_step_output(req_id_a, step_index=1))
+        assert finished == set()
+
+        second = self.scheduler.schedule()
+        assert _new_ids(second) == []
+        assert _cached_ids(second) == [req_id_a]
+        assert second.num_running_reqs == 1
+        assert second.num_waiting_reqs == 1
+
+        finished = self.scheduler.update_from_output(
+            second,
+            _make_step_output(req_id_a, step_index=2, finished=True),
+        )
+        assert finished == {req_id_a}
+
+        third = self.scheduler.schedule()
+        assert _new_ids(third) == [req_id_b]
+        assert _cached_ids(third) == []
+        assert third.num_running_reqs == 1
+        assert third.num_waiting_reqs == 0
+
+    def test_error_output_marks_finished_error(self) -> None:
+        req_id = self.scheduler.add_request(_make_step_request("err", num_inference_steps=3))
+
+        sched_output = self.scheduler.schedule()
+        assert _new_ids(sched_output) == [req_id]
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            _make_step_output(req_id, step_index=1, finished=True, error="worker failed"),
+        )
+
+        assert finished == {req_id}
+        state = self.scheduler.get_request_state(req_id)
+        assert state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert state.error == "worker failed"
+        assert self.scheduler.has_requests() is False
+
+    def test_missing_step_index_marks_finished_error(self) -> None:
+        req_id = self.scheduler.add_request(_make_step_request("missing", num_inference_steps=3))
+
+        sched_output = self.scheduler.schedule()
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            SimpleNamespace(
+                req_id=req_id,
+                step_index=None,
+                finished=True,
+                result=None,
+            ),
+        )
+
+        assert finished == {req_id}
+        state = self.scheduler.get_request_state(req_id)
+        assert state.status == DiffusionRequestStatus.FINISHED_ERROR
+        assert state.error == "Missing step_index in RunnerOutput"
+
+    def test_abort_request_for_waiting_and_running(self) -> None:
+        req_id_a = self.scheduler.add_request(_make_step_request("a", num_inference_steps=2))
+        req_id_b = self.scheduler.add_request(_make_step_request("b", num_inference_steps=2))
+
+        self.scheduler.finish_requests(req_id_b, DiffusionRequestStatus.FINISHED_ABORTED)
+        assert self.scheduler.get_request_state(req_id_b).status == DiffusionRequestStatus.FINISHED_ABORTED
+
+        running = self.scheduler.schedule()
+        assert _new_ids(running) == [req_id_a]
+
+        self.scheduler.finish_requests(req_id_a, DiffusionRequestStatus.FINISHED_ABORTED)
+        assert self.scheduler.get_request_state(req_id_a).status == DiffusionRequestStatus.FINISHED_ABORTED
+        assert self.scheduler.has_requests() is False
+
+    def test_has_requests_state_transition(self) -> None:
+        assert self.scheduler.has_requests() is False
+
+        req_id = self.scheduler.add_request(_make_step_request("has", num_inference_steps=2))
+        assert self.scheduler.has_requests() is True
+
+        sched_output = self.scheduler.schedule()
+        assert self.scheduler.has_requests() is True
+
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            _make_step_output(req_id, step_index=2, finished=True),
+        )
+        assert finished == {req_id}
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
+        assert self.scheduler.has_requests() is False
+
+    def test_scheduled_request_aborted_before_update_is_returned_finished(self) -> None:
+        req_id = self.scheduler.add_request(_make_step_request("abort-late", num_inference_steps=2))
+
+        sched_output = self.scheduler.schedule()
+        self.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED)
+
+        finished = self.scheduler.update_from_output(
+            sched_output,
+            _make_step_output(req_id, step_index=1),
+        )
+        assert finished == {req_id}
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_ABORTED
+
+    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)
+
+        first = self.scheduler.schedule()
+        assert self.scheduler.update_from_output(first, _make_step_output(req_id, step_index=1)) == set()
+        assert request.sampling_params.step_index == 1
+
+        second = self.scheduler.schedule()
+        assert _cached_ids(second) == [req_id]
+        assert self.scheduler.preempt_request(req_id) is True
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.PREEMPTED
+        assert request.sampling_params.step_index == 1
+
+        third = self.scheduler.schedule()
+        assert _cached_ids(third) == [req_id]
+        assert request.sampling_params.step_index == 1
+
+    @pytest.mark.parametrize(
+        ("sampling_params", "expected_steps"),
+        [
+            (
+                OmniDiffusionSamplingParams(
+                    timesteps=torch.tensor([1.0, 0.5, 0.0]),
+                    sigmas=[1.0, 0.5, 0.25, 0.0],
+                    num_inference_steps=5,
+                ),
+                3,
+            ),
+            (
+                OmniDiffusionSamplingParams(
+                    sigmas=[1.0, 0.5],
+                    num_inference_steps=5,
+                ),
+                2,
+            ),
+            (
+                OmniDiffusionSamplingParams(
+                    num_inference_steps=4,
+                ),
+                4,
+            ),
+        ],
+    )
+    def test_total_steps_priority(self, sampling_params: OmniDiffusionSamplingParams, expected_steps: int) -> None:
+        request = _make_step_request("priority", sampling_params=sampling_params)
+        req_id = self.scheduler.add_request(request)
+
+        for _ in range(expected_steps - 1):
+            sched_output = self.scheduler.schedule()
+            assert sched_output.scheduled_req_ids == [req_id]
+            next_step = request.sampling_params.step_index + 1
+            assert (
+                self.scheduler.update_from_output(
+                    sched_output,
+                    _make_step_output(req_id, step_index=next_step),
+                )
+                == set()
+            )
+
+        final_output = self.scheduler.schedule()
+        assert final_output.scheduled_req_ids == [req_id]
+        assert self.scheduler.update_from_output(
+            final_output,
+            _make_step_output(req_id, step_index=expected_steps, finished=True),
+        ) == {req_id}
+        assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
+
+    @pytest.mark.parametrize(
+        "sampling_params",
+        [
+            OmniDiffusionSamplingParams(num_inference_steps=0),
+            OmniDiffusionSamplingParams(num_inference_steps=3, step_index=3),
+            OmniDiffusionSamplingParams(num_inference_steps=3, step_index=-1),
+        ],
+    )
+    def test_rejects_invalid_initial_step_state(self, sampling_params: OmniDiffusionSamplingParams) -> None:
+        request = _make_step_request("invalid", sampling_params=sampling_params)
+
+        with pytest.raises(ValueError):
+            self.scheduler.add_request(request)

tests/diffusion/test_diffusion_step_pipeline.py

@@ -1,10 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Tests for step-level diffusion runner and worker execution."""
+"""Tests for step-level diffusion execution across runner / worker / executor / engine."""
 
 import os
+import queue
+import threading
 from contextlib import contextmanager
 from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -12,6 +15,7 @@ import torch
 import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
 from tests.utils import hardware_test
 from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D
 from vllm_omni.diffusion.distributed.parallel_state import (
@@ -20,10 +24,13 @@ from vllm_omni.diffusion.distributed.parallel_state import (
     init_distributed_environment,
     initialize_model_parallel,
 )
+from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 from vllm_omni.diffusion.ipc import (
     pack_diffusion_output_shm,
     unpack_diffusion_output_shm,
 )
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.sched import StepScheduler
 from vllm_omni.diffusion.sched.interface import (
     CachedRequestData,
     DiffusionSchedulerOutput,
@@ -32,6 +39,8 @@ from vllm_omni.diffusion.sched.interface import (
 from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
 from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
 from vllm_omni.diffusion.worker.utils import RunnerOutput
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@@ -86,6 +95,23 @@ class _StepPipeline:
         return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
 
 
+class _InterruptingStepPipeline(_StepPipeline):
+    interrupt = True
+
+    def denoise_step(self, state, **kwargs):
+        del state, kwargs
+        self.denoise_calls += 1
+        return None
+
+    def step_scheduler(self, state, noise_pred, **kwargs):
+        del state, noise_pred, kwargs
+        raise AssertionError("step_scheduler should not run after interrupt")
+
+    def post_decode(self, state, **kwargs):
+        del state, kwargs
+        raise AssertionError("post_decode should not run after interrupt")
+
+
 class _IdentityNoiseTransformer(torch.nn.Module):
     def forward(self, x: torch.Tensor, **kwargs):
         del kwargs
@@ -188,6 +214,21 @@ def _make_step_request(num_inference_steps: int = 2):
     )
 
 
+def _assert_aborted_output(output: DiffusionOutput, request_id: str) -> None:
+    assert output.output is None
+    assert output.error is None
+    assert output.aborted is True
+    assert output.abort_message == f"Request {request_id} aborted."
+
+
+def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) -> OmniDiffusionRequest:
+    return OmniDiffusionRequest(
+        prompts=[f"prompt-{req_id}"],
+        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps),
+        request_ids=[req_id],
+    )
+
+
 def _make_runner():
     runner = object.__new__(DiffusionModelRunner)
     runner.vllm_config = object()
@@ -242,6 +283,18 @@ def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_
     )
 
 
+def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine:
+    engine = object.__new__(DiffusionEngine)
+    engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline")
+    engine.pre_process_func = None
+    engine.post_process_func = None
+    engine.scheduler = scheduler
+    engine.execute_fn = execute_fn
+    engine._rpc_lock = threading.RLock()
+    engine.abort_queue = queue.Queue()
+    return engine
+
+
 def _expected_output_for_mode(mode: str) -> torch.Tensor:
     if mode == "cfg":
         return torch.tensor([[3.0]])
@@ -322,6 +375,52 @@ class TestRunner:
         assert runner.pipeline.scheduler_calls == 2
         assert runner.pipeline.decode_calls == 1
 
+    def test_rejects_multi_request_step_batch(self):
+        runner = _make_runner()
+        req_1 = _make_step_request()
+        req_2 = _make_step_request()
+        req_2.request_ids = ["req-2"]
+
+        scheduler_output = DiffusionSchedulerOutput(
+            step_id=0,
+            scheduled_new_reqs=[
+                NewRequestData(sched_req_id="req-1", req=req_1),
+                NewRequestData(sched_req_id="req-2", req=req_2),
+            ],
+            scheduled_cached_reqs=CachedRequestData.make_empty(),
+            finished_req_ids=set(),
+            num_running_reqs=2,
+            num_waiting_reqs=0,
+        )
+
+        with pytest.raises(ValueError, match="batch_size=1"):
+            DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
+
+    def test_rejects_missing_cached_state(self):
+        runner = _make_runner()
+
+        with pytest.raises(ValueError, match="Missing cached state"):
+            DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(sched_req_id="req-missing"))
+
+    def test_interrupt_marks_request_finished_and_clears_state(self, monkeypatch):
+        runner = _make_runner()
+        runner.pipeline = _InterruptingStepPipeline()
+        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))
+
+        assert output.req_id == "req-1"
+        assert output.step_index == 0
+        assert output.finished is True
+        assert output.result is not None
+        assert output.result.error == "stepwise denoise interrupted"
+        assert "req-1" not in runner.state_cache
+        assert runner.pipeline.prepare_calls == 1
+        assert runner.pipeline.denoise_calls == 1
+        assert runner.pipeline.scheduler_calls == 0
+        assert runner.pipeline.decode_calls == 0
+
     def test_load_model_rejects_unsupported_step_execution(self, monkeypatch):
         class _RequestOnlyPipeline:
             pass
@@ -439,6 +538,153 @@ class TestWorker:
             DiffusionWorker.execute_stepwise(worker, scheduler_output)
 
 
+@pytest.mark.cpu
+class TestExecutor:
+    """MultiprocDiffusionExecutor.execute_step"""
+
+    def test_execute_step_passes_through_runner_output(self):
+        executor = object.__new__(MultiprocDiffusionExecutor)
+        executor._ensure_open = lambda: None
+        expected = RunnerOutput(req_id="req-step", step_index=1, finished=False, result=None)
+        executor.collective_rpc = Mock(return_value=expected)
+
+        request = _make_engine_request("req-step", num_inference_steps=2)
+        scheduler_output = _make_scheduler_output(request, sched_req_id="req-step")
+
+        output = MultiprocDiffusionExecutor.execute_step(executor, scheduler_output)
+
+        assert output is expected
+
+
+@pytest.mark.cpu
+class TestEngine:
+    """Step-execution paths in DiffusionEngine.add_req_and_wait_for_response"""
+
+    @pytest.mark.parametrize(
+        ("execute_fn", "expected_error"),
+        [
+            (
+                lambda _: RunnerOutput(
+                    req_id="req-error",
+                    step_index=1,
+                    finished=True,
+                    result=DiffusionOutput(error="boom"),
+                ),
+                "boom",
+            ),
+            (
+                lambda _: (_ for _ in ()).throw(RuntimeError("gpu on fire")),
+                "gpu on fire",
+            ),
+        ],
+    )
+    def test_step_engine_returns_error(self, execute_fn, expected_error):
+        scheduler = StepScheduler()
+        scheduler.initialize(Mock())
+        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))
+
+        assert output.output is None
+        assert expected_error in output.error
+
+    def test_step_execution_completes(self):
+        scheduler = StepScheduler()
+        scheduler.initialize(Mock())
+        engine = _make_engine(scheduler)
+        request = _make_engine_request("req-step", num_inference_steps=2)
+
+        call_count = {"n": 0}
+
+        def execute_fn(_):
+            call_count["n"] += 1
+            finished = call_count["n"] == 2
+            return RunnerOutput(
+                req_id="req-step",
+                step_index=call_count["n"],
+                finished=finished,
+                result=(DiffusionOutput(output=torch.tensor([2.0])) if finished else None),
+            )
+
+        engine.execute_fn = execute_fn
+
+        output = engine.add_req_and_wait_for_response(request)
+
+        assert call_count["n"] == 2
+        assert output.error is None
+        assert torch.equal(output.output, torch.tensor([2.0]))
+
+    def test_step_abort_stops_rescheduling_after_first_step(self):
+        scheduler = StepScheduler()
+        scheduler.initialize(Mock())
+        engine = _make_engine(scheduler)
+        request = _make_engine_request("req-stop", num_inference_steps=4)
+
+        step = {"n": 0}
+
+        def execute_fn(_):
+            step["n"] += 1
+            engine.abort("req-stop")
+            return RunnerOutput(
+                req_id="req-stop",
+                step_index=1,
+                finished=False,
+                result=None,
+            )
+
+        engine.execute_fn = execute_fn
+
+        output = engine.add_req_and_wait_for_response(request)
+
+        assert step["n"] == 1
+        _assert_aborted_output(output, "req-stop")
+
+    def test_step_abort_after_reschedule_returns_aborted_output(self):
+        scheduler = StepScheduler()
+        scheduler.initialize(Mock())
+        engine = _make_engine(scheduler)
+        request = _make_engine_request("req-mid", num_inference_steps=4)
+
+        step = {"n": 0}
+
+        def execute_fn(sched_output):
+            step["n"] += 1
+            if step["n"] == 2:
+                assert sched_output == _make_cached_scheduler_output("req-mid", step_id=1)
+                engine.abort("req-mid")
+            return RunnerOutput(
+                req_id="req-mid",
+                step_index=step["n"],
+                finished=False,
+                result=None,
+            )
+
+        engine.execute_fn = execute_fn
+
+        output = engine.add_req_and_wait_for_response(request)
+
+        assert step["n"] == 2
+        _assert_aborted_output(output, "req-mid")
+
+    def test_finished_step_without_result_returns_error(self):
+        scheduler = StepScheduler()
+        scheduler.initialize(Mock())
+        engine = _make_engine(
+            scheduler,
+            execute_fn=lambda _: RunnerOutput(
+                req_id="req-missing",
+                step_index=1,
+                finished=True,
+                result=None,
+            ),
+        )
+
+        output = engine.add_req_and_wait_for_response(_make_engine_request("req-missing", num_inference_steps=1))
+
+        assert output.output is None
+        assert output.error == "Diffusion execution finished without a final output."
+
+
 @pytest.mark.cpu
 class TestIPC:
     def test_pack_unpack_runner_output_shm(self):
@@ -458,6 +704,15 @@ class TestIPC:
 class TestSupportedPipelines:
     """Step-execution protocol checks for supported pipelines."""
 
+    def test_default_stage_config_includes_step_execution(self):
+        stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+            {
+                "step_execution": True,
+            }
+        )[0]
+
+        assert stage_cfg["engine_args"]["step_execution"] is True
+
     def test_qwen_image_supports_step_execution(self):
         from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
         from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -66,7 +66,9 @@ def _make_engine(num_gpus: int = 1):
     sched.initialize(Mock())
     engine.scheduler = sched
     engine.executor = executor
-    engine._rpc_lock = threading.Lock()
+    engine._rpc_lock = threading.RLock()
+    engine.abort_queue = queue.Queue()
+    engine.execute_fn = executor.execute_request
     return engine, executor, req_q, res_q
 
 
@@ -80,7 +82,7 @@ def _start_worker(req_q, res_q, count=2):
             req = req_q.get(timeout=10)
             method = req.get("method", "")
             args = req.get("args", ())
-            if method == "generate" and args and hasattr(args[0], "request_ids"):
+            if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_ids"):
                 tag = f"result_for_{args[0].request_ids[0]}"
             elif args:
                 tag = f"result_for_{args[0]}"
@@ -116,11 +118,11 @@ def _inject_interleave(executor):
     return a_enqueued, b_complete
 
 
-# ──────────────────── bug-reproduction: concurrent add_req ────────────────
+# ───────────────── concurrent request execution ─────────────────
 
 
-class TestConcurrentAddReqBug:
-    """Two concurrent ``add_req_and_wait_for_response()`` calls swap results."""
+class TestConcurrentRequestExecution:
+    """Concurrent request execution should not swap results."""
 
     def test_results_are_correctly_routed(self):
         engine, executor, req_q, res_q = _make_engine()
@@ -151,11 +153,11 @@ class TestConcurrentAddReqBug:
         assert results["B"].error == "result_for_B"
 
 
-# ──────────────── bug-reproduction: concurrent collective_rpc ─────────────
+# ───────────────── concurrent collective RPC ─────────────────
 
 
-class TestConcurrentCollectiveRpcBug:
-    """Two concurrent ``collective_rpc()`` calls swap results."""
+class TestConcurrentCollectiveRpc:
+    """Concurrent ``collective_rpc()`` calls should not swap results."""
 
     def test_results_are_correctly_routed(self):
         engine, executor, req_q, res_q = _make_engine()
@@ -192,11 +194,11 @@ class TestConcurrentCollectiveRpcBug:
         assert results["B"].error == "result_for_call_B"
 
 
-# ──────── bug-reproduction: add_req vs collective_rpc concurrently ────────
+# ──────────── concurrent request execution and collective RPC ────────────
 
 
-class TestConcurrentAddReqVsCollectiveRpcBug:
-    """``add_req`` and ``collective_rpc`` running concurrently swap results."""
+class TestConcurrentRequestExecutionAndCollectiveRpc:
+    """Request execution and ``collective_rpc()`` should not swap results."""
 
     def test_results_are_correctly_routed(self):
         engine, executor, req_q, res_q = _make_engine()
@@ -205,7 +207,7 @@ class TestConcurrentAddReqVsCollectiveRpcBug:
 
         results: dict[str, object] = {}
 
-        def _a():  # add_req path
+        def _a():  # request execution path
             results["A"] = engine.add_req_and_wait_for_response(_mock_request("A"))
 
         def _b():  # collective_rpc path
@@ -230,10 +232,10 @@ class TestConcurrentAddReqVsCollectiveRpcBug:
         assert results["B"].error == "result_for_call_B"
 
 
-# ─────────────── backward-compatibility (serial) tests ────────────────────
+# ─────────────────────── serial operation coverage ───────────────────────
 
 
-class TestSerialOperations:
+class TestSerialEngineOperations:
     """Verify correct behaviour for single-threaded (serial) usage.
 
     These tests must pass both **before** and **after** any concurrency fix
@@ -385,18 +387,18 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
 
         executor._result_mq.dequeue = _hanging_dequeue
 
-        # Thread running add_req — acquires the lock, enqueues, then
+        # Thread running request execution — acquires the lock, enqueues, then
         # blocks on dequeue forever (worker hang).
-        def _stalled_add_req():
+        def _stalled_request_execution():
             try:
                 engine.add_req_and_wait_for_response(_mock_request("stalled"))
             except Exception:
                 pass
 
-        t = threading.Thread(target=_stalled_add_req, daemon=True)
+        t = threading.Thread(target=_stalled_request_execution, daemon=True)
         t.start()
 
-        # Wait until add_req is truly inside the lock and blocking.
+        # Wait until request execution is truly inside the lock and blocking.
         add_req_blocked.wait(5)
 
         # collective_rpc should time out at lock acquisition, not hang.