[Refact] Refactor diffusion output streaming (#5166)

Signed-off-by: bjf-frz <frz123db@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

docs/configuration/stage_configs.md

@@ -377,7 +377,12 @@ Default: `"0"`
 
 #### `engine_args.max_num_seqs`
 
-The maximum number of sequences for concurrent processing in this stage. For LLM stages, this controls the vLLM scheduler's maximum concurrent sequences. For all stage types, this also controls how many tasks can be batched together in the task processing loop.
+The maximum number of sequences for concurrent processing in this stage. For
+LLM stages, this controls the vLLM scheduler's maximum concurrent sequences.
+For diffusion stages, this controls scheduler wave capacity in both request
+batch mode and step batch mode. `max_num_seqs: 1` is the serial/conservative
+path; values above `1` allow compatible requests to batch when the selected
+diffusion pipeline supports that batching mode.
 
 Default: `1`
 

docs/design/feature/diffusion_continuous_batching.md

@@ -139,7 +139,8 @@ compatibility gating and runner-side `StepInputBatch` packing.
 
 ## Current Limitations
 
-- Experimental feature; use `max_num_seqs=1` for the older conservative path.
+- Experimental feature; use `max_num_seqs=1` for the conservative single-request
+  step path.
 - Only native pipelines that already support `step_execution=True`.
 - Only homogeneous batches keyed by `StepBatchSamplingParamsKey` are supported.
 - `cache_backend`, KV transfer, and other request-mode extras are not wired

docs/design/feature/diffusion_step_execution.md

@@ -32,7 +32,7 @@ Current engine/runtime limitations:
 
 - Continuous batching with `max_num_seqs > 1` is experimental and documented in
   [Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
-  Keep `max_num_seqs=1` if you want the older conservative behavior.
+  Keep `max_num_seqs=1` for the conservative single-request step path.
 - `cache_backend` is not supported in step mode.
 - Request-mode extras such as KV transfer are not wired into step mode yet.
 - Unsupported pipelines now fail early during model loading instead of failing on the first request.

docs/design/module/dit_module.md

@@ -52,7 +52,7 @@ The diffusion module follows a **multi-process, distributed architecture** with
 
 ### Responsibilities
 
-The `DiffusionEngine` is the **orchestrator** of the diffusion inference system. It manages the lifecycle of worker processes and coordinates the execution flow.
+The `DiffusionEngine` is the **orchestrator** of the diffusion inference system. It owns request admission, scheduler/executor coordination, output-stream delivery, cancellation cleanup, and startup warmup.
 
 ### Key Components
 
@@ -60,68 +60,84 @@ The `DiffusionEngine` is the **orchestrator** of the diffusion inference system.
 
 ```python
 class DiffusionEngine:
-    def __init__(self, od_config: OmniDiffusionConfig):
+    def __init__(
+        self,
+        od_config: OmniDiffusionConfig,
+        scheduler: SchedulerInterface | None = None,
+    ):
         self.od_config = od_config
-        self.post_process_func = get_diffusion_post_process_func(od_config)
-        self.pre_process_func = get_diffusion_pre_process_func(od_config)
-        self._processes: list[mp.Process] = []
-        self._make_client()
+        self._init_process_hooks(od_config)
+        self.execution_mode = self._resolve_execution_mode(od_config)
+        self._init_executor(od_config)
+        self._init_scheduler(od_config, scheduler)
+        self._init_runtime_state()
+        self._init_execute_fn()
 ```
 
 **Key Features**:
 
 - **Pre/Post Processing**: Registers model-specific pre-processing and post-processing functions via registry pattern
 
-- **Worker Management**: Launches and manages multiple worker processes (one per GPU)
+- **Execution Mode Selection**: Selects one of two execution modes:
+  - `REQUEST_BATCH`: complete requests are scheduled through request-level execution. `max_num_seqs=1` is the serial request path; values above `1` allow compatible requests to use fused `pipeline.forward(batch)` when the pipeline supports it.
+  - `STEP_BATCH`: requests are advanced one denoising step at a time through `prepare_encode()`, `denoise_step()`, `step_scheduler()`, and `post_decode()`. `max_num_seqs` controls the maximum number of compatible active requests in one step wave.
 
-- **Process Isolation**: Uses multiprocessing for true parallelism
+- **Scheduler Selection**: Uses `StepScheduler` for step-batch execution and `RequestScheduler` for request-batch execution. Tests and custom integrations may inject a `SchedulerInterface` instance explicitly.
 
-#### 1.2 Worker Launch Process
+- **Executor Selection**: Builds the configured `DiffusionExecutor`, then routes scheduled work through `execute_batch` or `execute_step` according to the execution mode. `execute_batch` handles the serial `max_num_seqs=1` request path internally.
 
-The engine launches workers using a **spawn** method:
+- **Factory-Owned Warmup**: `DiffusionEngine.make_engine()` constructs the selected engine backend and then runs startup warmup through `run_startup_warmup()`.
+
+#### 1.2 Execution Modes
 
 ```python
-def _launch_workers(self, broadcast_handle):
-    # Creates one process per GPU
-    for i in range(num_gpus):
-        process = mp.Process(
-            target=worker_proc.worker_main,
-            args=(i, od_config, writer, broadcast_handle),
-            name=f"DiffusionWorker-{i}",
-        )
-        process.start()
+class DiffusionExecutionMode(str, Enum):
+    REQUEST_BATCH = "request_batch"
+    STEP_BATCH = "step_batch"
+```
+
+`streaming_output=True` requires step execution. If streaming is enabled without `step_execution=True`, the engine enables step execution so the output stream can receive intermediate chunks or the final step output.
+
+#### 1.3 Output Stream Lifecycle
+
+```python
+request_id = engine.add_request(request)
+async for output in engine.get_output_stream(request_id):
+    yield output
 ```
 
 **Design Decisions**:
 
-- **Spawn Method**: Ensures clean state for each worker (no shared memory issues)
+- **Unified Stream Semantics**: `step_streaming()` is the public async engine path. Streaming callers forward every yielded output. Non-streaming callers consume the same stream to completion and return only the final output.
 
-- **Pipe Communication**: Uses `mp.Pipe` for initialization handshake
+- **Queue Delivery**: Each admitted request gets an `asyncio.Queue[DiffusionOutput]`. The busy loop emits chunks/final outputs into that queue.
 
-- **Device Selection**: Each worker is assigned a specific GPU (`cuda:{rank}`)
+- **Cancellation Cleanup**: If a client disconnects or cancels stream consumption, `get_output_stream()` removes only the queue. Scheduler terminal state is still finalized when the scheduler reports the request as finished, so request IDs do not remain active indefinitely.
 
-#### 1.3 Request Processing Flow
+#### 1.4 Request Processing Flow
 
 ```python
-def step(self, requests: list[OmniDiffusionRequest]):
+async def step_streaming(self, request: OmniDiffusionRequest):
     # 1. Pre-process requests
-    requests = self.pre_process_func(requests)
+    request = self.pre_process_func(request)
 
-    # 2. Send to scheduler and wait for response
-    output = self.add_req_and_wait_for_response(requests)
-
-    # 3. Post-process results
-    result = self.post_process_func(output.output)
-    return result
+    # 2. Add request and consume the unified output stream
+    async for output in self.async_add_req_and_stream_response(request):
+        # 3. Post-process each delivered output
+        yield self.postprocess_output(request, output, ...)
 ```
 
 **Flow**:
 
 1. **Pre-processing**: Applies model-specific transformations
 
-2. **Scheduling**: Delegates to scheduler for distribution
+2. **Scheduling**: Delegates request state transitions to the selected scheduler
 
-3. **Post-processing**: Converts raw outputs to final format (e.g., PIL images)
+3. **Execution**: Calls the selected executor function for request, request-batch, or step execution
+
+4. **Output Delivery**: Emits intermediate chunks and final outputs through the request output stream
+
+5. **Post-processing**: Converts raw outputs to final format (e.g., PIL images)
 
 ---
 
@@ -915,36 +931,49 @@ def initialize_model_parallel(
 1. User Request
    └─> Omni.generate(prompt)
        └─> Prepare OmniDiffusionRequest
-           └─> DiffusionEngine.step(requests)
+           └─> DiffusionEngine.step_streaming(request)
 
 2. Pre-processing
-   └─> pre_process_func(requests)
+   └─> pre_process_func(request)
        └─> Model-specific transformations
 
 3. Scheduling
-   └─> scheduler.add_request(request)
-       └─> scheduler.schedule()
-           └─> DiffusionEngine submits scheduled request to executor.add_req(req)
+       └─> scheduler.add_request(request)
+           └─> scheduler.schedule()
+           └─> DiffusionEngine selects execute_batch / execute_step
 
 4. Worker Execution
    └─> WorkerProc._worker_busy_loop()
        └─> DiffusionWorker.execute_model(req)
-           └─> Pipeline.forward(req)
-               ├─> encode_prompt()
-               ├─> prepare_latents()
-               ├─> diffuse() [loop]
-               │   ├─> transformer.forward() [with cache backend hooks]
-               │   └─> scheduler.step()
-               └─> vae.decode()
+           └─> DiffusionModelRunner selects execution mode
+               ├─> REQUEST_BATCH
+               │   └─> Pipeline.forward(req_or_batch)
+               │       ├─> encode_prompt()
+               │       ├─> prepare_latents()
+               │       ├─> diffuse() [loop]
+               │       │   ├─> transformer.forward() [with cache backend hooks]
+               │       │   └─> scheduler.step()
+               │       └─> vae.decode()
+               └─> STEP
+                   ├─> pipeline.prepare_encode(state) [new requests]
+                   ├─> pipeline.denoise_step(input_batch)
+                   ├─> pipeline.step_scheduler(state, noise_pred)
+                   └─> pipeline.post_decode(state) [chunk boundary or final step]
 
 5. Result Collection
-   └─> Executor returns DiffusionOutput
+   └─> Executor returns RunnerOutput / BatchRunnerOutput
        └─> scheduler.update_from_output(...)
-           └─> DiffusionEngine pops finished request state
+           ├─> DiffusionEngine emits chunk/final output to request queue
+           └─> DiffusionEngine finalizes terminal scheduler state
 
 6. Post-processing
    └─> post_process_func(output)
        └─> Convert to PIL images / final format
+
+7. Stream Completion / Cancellation
+   ├─> Streaming callers receive every yielded output
+   ├─> Non-streaming callers drain the stream and return the final output
+   └─> Cancelled streams drop their output queue but still allow scheduler state cleanup
 ```
 
 ---

docs/user_guide/diffusion/request_batching.md

@@ -15,6 +15,14 @@ error, and abort state. The scheduler decides which requests can run together.
 
 ## Enablement
 
+Request-level execution is the default diffusion request mode when
+`step_execution=False`. `max_num_seqs` controls the maximum number of compatible
+requests in one scheduler wave:
+
+- `max_num_seqs=1` is the serial request path.
+- `max_num_seqs>1` allows fused request-level batching for pipelines that
+  declare request-batch support.
+
 Increase `max_num_seqs` above `1` to allow the request scheduler to keep more
 than one compatible request active:
 
@@ -51,11 +59,11 @@ stage_args:
 
 ## Compatibility
 
-Only pipelines that declare request-batch support use the fused request-batch
-path. The engine validates that the pipeline `forward()` uses the request-batch
-contract and returns `list[DiffusionOutput]`. Pipelines that do not support this
-contract do not use fused `pipeline.forward(batch)`; scheduled requests are
-executed through per-request worker calls.
+Only pipelines that declare request-batch support can use the fused
+multi-request `pipeline.forward(batch)` path. Pipelines that do not support this
+contract can still use request-level execution with `max_num_seqs=1`, which is
+the serial request path. Setting `max_num_seqs>1` for a pipeline without
+request-batch support fails early during engine initialization.
 
 The scheduler batches only compatible requests. Compatibility is based on
 shape-sensitive and guidance-sensitive sampling fields, including resolution,
@@ -69,7 +77,8 @@ step-wise runtime, see [Step Execution](step_execution.md).
 ## Tuning
 
 - `max_num_seqs` caps the number of active compatible requests in one scheduler
-  wave.
+  wave. `1` means serial request execution; values above `1` enable fused
+  request-level batching when the pipeline supports it.
 - `request_batch_max_wait_ms` is an upper bound on extra admission wait before a
   new wave starts. Keep it small for latency-sensitive serving; values such as
   `10` to `50` ms are a practical starting range for bursty HTTP ingress.

docs/user_guide/diffusion/step_execution.md

@@ -39,6 +39,11 @@ vllm serve Qwen/Qwen-Image --omni \
 For serving, `--step-execution` enables the step-wise runtime. Continuous
 batching only becomes relevant when `--max-num-seqs > 1`.
 
+`max_num_seqs` controls the step scheduler capacity:
+
+- `max_num_seqs=1` runs one request through the step-wise state machine.
+- `max_num_seqs>1` allows compatible active requests to share step waves.
+
 ## Supported Pipelines
 
 | Pipeline | Example models | Step execution |
@@ -51,7 +56,7 @@ batching only becomes relevant when `--max-num-seqs > 1`.
     the step-wise path can batch
     compatible requests together. This is experimental. Requests with
     incompatible sampling parameters are intentionally kept in separate batches,
-    and `max_num_seqs=1` remains the conservative default.
+    and `max_num_seqs=1` remains the conservative single-request step path.
 
 ## Current Limitations
 

tests/diffusion/test_diffusion_engine.py

@@ -15,7 +15,11 @@ from pytest_mock import MockerFixture
 
 import vllm_omni.diffusion.diffusion_engine as diffusion_engine_module
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, _move_tensor_tree_to_cpu
+from vllm_omni.diffusion.diffusion_engine import (
+    DiffusionEngine,
+    DiffusionExecutionMode,
+    _move_tensor_tree_to_cpu,
+)
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import (
     CachedRequestData,
@@ -221,7 +225,7 @@ class TestRequestBatchCapability:
             diffusion_engine_module.supports_request_batch(od_config)
         registry_load.assert_not_called()
 
-    def test_engine_disables_batch_dispatch_for_single_request_pipeline(
+    def test_engine_uses_request_batch_mode_for_single_request_pipeline(
         self,
         monkeypatch: pytest.MonkeyPatch,
         mocker: MockerFixture,
@@ -230,6 +234,7 @@ class TestRequestBatchCapability:
             model_class_name="SinglePipeline",
             custom_pipeline_args=None,
             streaming_output=False,
+            max_num_seqs=1,
         )
         fake_executor = SimpleNamespace(
             execute_request=mocker.Mock(return_value="per-request"),
@@ -255,15 +260,50 @@ class TestRequestBatchCapability:
             "_try_load_model_cls",
             lambda model_class_name: _SingleRequestPipeline,
         )
-        monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
 
         engine = DiffusionEngine(od_config)
-        output = engine.execute_fn(_make_request_mode_sched_output("req-a", "req-b"))
+        output = engine.execute_fn(_make_request_mode_sched_output("req-a"))
 
+        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
         assert engine.supports_request_batch is False
-        assert output == "per-request"
-        fake_executor.execute_request.assert_called_once()
-        fake_executor.execute_batch.assert_not_called()
+        assert output == "batch"
+        fake_executor.execute_batch.assert_called_once()
+        fake_executor.execute_request.assert_not_called()
+
+    def test_engine_rejects_multi_seq_request_batch_for_single_request_pipeline(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        mocker: MockerFixture,
+    ) -> None:
+        od_config = SimpleNamespace(
+            model_class_name="SinglePipeline",
+            custom_pipeline_args=None,
+            streaming_output=False,
+            max_num_seqs=2,
+        )
+        fake_executor_cls = mocker.Mock()
+
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+            lambda *args, **kwargs: fake_executor_cls,
+        )
+        monkeypatch.setattr(
+            diffusion_engine_module.DiffusionModelRegistry,
+            "_try_load_model_cls",
+            lambda model_class_name: _SingleRequestPipeline,
+        )
+
+        with pytest.raises(ValueError, match="max_num_seqs=1"):
+            DiffusionEngine(od_config)
+        fake_executor_cls.assert_not_called()
 
     @pytest.mark.parametrize("request_ids", [("req-a",), ("req-a", "req-b")])
     def test_engine_enables_batch_dispatch_for_request_batch_pipeline(
@@ -276,6 +316,7 @@ class TestRequestBatchCapability:
             model_class_name="BatchPipeline",
             custom_pipeline_args=None,
             streaming_output=False,
+            max_num_seqs=2,
         )
         fake_executor = SimpleNamespace(
             execute_request=mocker.Mock(return_value="per-request"),
@@ -301,16 +342,55 @@ class TestRequestBatchCapability:
             "_try_load_model_cls",
             lambda model_class_name: _BatchCapablePipeline,
         )
-        monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
 
         engine = DiffusionEngine(od_config)
         output = engine.execute_fn(_make_request_mode_sched_output(*request_ids))
 
+        assert engine.execution_mode == DiffusionExecutionMode.REQUEST_BATCH
         assert engine.supports_request_batch is True
         assert output == "batch"
         fake_executor.execute_batch.assert_called_once()
         fake_executor.execute_request.assert_not_called()
 
+    def test_make_engine_runs_startup_warmup(self, monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture) -> None:
+        od_config = SimpleNamespace(
+            model_class_name="SinglePipeline",
+            custom_pipeline_args=None,
+            engine_backend="default",
+            streaming_output=False,
+        )
+        fake_executor = SimpleNamespace(
+            execute_request=mocker.Mock(),
+            execute_batch=mocker.Mock(),
+            execute_step=mocker.Mock(),
+        )
+        fake_executor_cls = mocker.Mock(return_value=fake_executor)
+        warmup = mocker.Mock()
+
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+            lambda *args, **kwargs: None,
+        )
+        monkeypatch.setattr(
+            "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+            lambda *args, **kwargs: fake_executor_cls,
+        )
+        monkeypatch.setattr(
+            diffusion_engine_module.DiffusionModelRegistry,
+            "_try_load_model_cls",
+            lambda model_class_name: _SingleRequestPipeline,
+        )
+        monkeypatch.setattr(DiffusionEngine, "run_startup_warmup", warmup)
+
+        engine = DiffusionEngine.make_engine(od_config)
+
+        assert isinstance(engine, DiffusionEngine)
+        warmup.assert_called_once_with()
+
 
 class TestRequestBatchAdmission:
     pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@@ -470,11 +550,20 @@ def test_move_tensor_tree_moves_nested_cuda_tensors_to_cpu() -> None:
     assert moved["items"][1][0] == "keep"
 
 
+async def _consume_final_output(generator):
+    final_output = None
+    async for output in generator:
+        final_output = output
+    if final_output is None:
+        raise RuntimeError("Diffusion execution finished without output.")
+    return final_output
+
+
 @pytest.mark.asyncio
-async def test_async_add_req_and_wait_for_response():
+async def test_async_add_req_and_stream_response():
     engine = object.__new__(DiffusionEngine)
     engine.scheduler = MockScheduler()
-    engine._out_queue = {}
+    engine._out_streams = {}
     engine.abort_queue: queue.Queue[str] = queue.Queue()
     engine._rpc_queue = queue.Queue()
     engine._rpc_lock = threading.RLock()
@@ -485,6 +574,7 @@ async def test_async_add_req_and_wait_for_response():
     engine._loop_started = False
     engine.main_loop = None
     engine.supports_request_batch = False
+    engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
 
     engine._finalize_finished_request = lambda rid, out, err: out.result
 
@@ -512,7 +602,7 @@ async def test_async_add_req_and_wait_for_response():
     async def run_task(rid):
         req = SimpleNamespace(request_id=rid)
         start = time.time()
-        res = await engine.async_add_req_and_wait_for_response(req)
+        res = await _consume_final_output(engine.async_add_req_and_stream_response(req))
         return rid, res, time.time() - start
 
     task_ids = [f"req_{i}" for i in range(5)]

tests/diffusion/test_diffusion_engine_cleanup.py

@@ -5,15 +5,17 @@ import asyncio
 import queue
 import threading
 from types import SimpleNamespace
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 
 import pytest
 
+from vllm_omni.diffusion import diffusion_engine as diffusion_engine_module
 from vllm_omni.diffusion.data import DiffusionOutput
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import DiffusionRequestStatus, RequestScheduler
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.outputs import OmniRequestOutput
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
@@ -33,8 +35,7 @@ def _make_engine() -> DiffusionEngine:
     engine.executor = SimpleNamespace(shutdown=Mock())
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
-    engine._out_queue = {}
-    engine._out_queue_streaming = {}
+    engine._out_streams = {}
     engine._closed = False
     engine._shutdown_complete = False
     engine.abort_queue = queue.Queue()
@@ -44,28 +45,13 @@ def _make_engine() -> DiffusionEngine:
     return engine
 
 
-def test_close_completes_pending_async_waiters() -> None:
-    engine = _make_engine()
-    event_loop = asyncio.new_event_loop()
-    try:
-        future = event_loop.create_future()
-        engine._out_queue["pending-req"] = future
-
-        engine.close()
-
-        assert future.done()
-        assert future.result().error == "DiffusionEngine is closed."
-    finally:
-        event_loop.close()
-
-
-def test_close_completes_pending_streaming_waiters() -> None:
+def test_close_completes_pending_output_streams() -> None:
     engine = _make_engine()
     event_loop = asyncio.new_event_loop()
     try:
         engine.main_loop = event_loop
         queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
-        engine._out_queue_streaming["pending-stream"] = queue
+        engine._out_streams["pending-stream"] = queue
 
         engine.close()
 
@@ -76,21 +62,101 @@ def test_close_completes_pending_streaming_waiters() -> None:
         event_loop.close()
 
 
-def test_handle_finished_requests_ignores_already_drained_waiter() -> None:
+def test_emit_finished_outputs_finalizes_already_drained_waiter() -> None:
     class RacingOutQueue(dict):
-        def __contains__(self, key):
-            return True
-
-        def pop(self, key, default=None):
+        def get(self, key, default=None):
             return default
 
     engine = _make_engine()
-    engine._out_queue = RacingOutQueue()
-    engine._finalize_finished_request = Mock(side_effect=AssertionError("should not finalize drained waiters"))
+    request_id = engine.scheduler.add_request(_make_request("pending-req"))
+    engine.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
+    engine._out_streams = RacingOutQueue()
 
-    engine._handle_finished_requests({"pending-req"})
+    engine._emit_finished_outputs({request_id})
 
-    engine._finalize_finished_request.assert_not_called()
+    assert engine.scheduler.get_request_state(request_id) is None
+
+
+def test_emit_step_outputs_finalizes_finished_request_without_stream() -> None:
+    engine = _make_engine()
+    engine.execution_mode = DiffusionExecutionMode.STEP_BATCH
+    request_id = engine.scheduler.add_request(_make_request("step-drained"))
+    engine.scheduler.finish_requests(request_id, DiffusionRequestStatus.FINISHED_ABORTED)
+
+    engine._emit_outputs({request_id}, [request_id], SimpleNamespace(get_request_output=lambda _request_id: None))
+
+    assert engine.scheduler.get_request_state(request_id) is None
+
+
+def test_init_accepts_custom_scheduler(monkeypatch: pytest.MonkeyPatch) -> None:
+    od_config = SimpleNamespace(
+        custom_pipeline_args=None,
+        model_class_name="CustomSchedulerPipeline",
+        streaming_output=False,
+    )
+    custom_scheduler = RequestScheduler()
+    fake_executor = SimpleNamespace(
+        execute_request=Mock(),
+        execute_batch=Mock(),
+        execute_step=Mock(),
+    )
+
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
+        lambda *args, **kwargs: None,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
+        lambda *args, **kwargs: None,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
+        lambda *args, **kwargs: Mock(return_value=fake_executor),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.diffusion_engine.supports_request_batch",
+        lambda *args, **kwargs: False,
+    )
+
+    engine = DiffusionEngine(od_config, scheduler=custom_scheduler)
+
+    assert engine.scheduler is custom_scheduler
+
+
+@pytest.mark.asyncio
+async def test_step_compatibility_wrapper_returns_final_batch() -> None:
+    engine = _make_engine()
+    first = [OmniRequestOutput.from_diffusion(request_id="req", images=[], finished=False)]
+    final = [OmniRequestOutput.from_diffusion(request_id="req", images=[], finished=True)]
+
+    async def _step_streaming(_request):
+        yield first
+        yield final
+
+    engine.step_streaming = _step_streaming  # type: ignore[method-assign]
+    with patch.object(diffusion_engine_module.logger, "warning_once") as warning_once:
+        output = await engine.step(_make_request("req"))
+
+    assert output is final
+    warning_once.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_async_wait_compatibility_wrapper_returns_final_output() -> None:
+    engine = _make_engine()
+    first = DiffusionOutput(output="chunk", finished=False)
+    final = DiffusionOutput(output="final", finished=True)
+
+    async def _stream_response(_request):
+        yield first
+        yield final
+
+    engine.async_add_req_and_stream_response = _stream_response  # type: ignore[method-assign]
+    with patch.object(diffusion_engine_module.logger, "warning_once") as warning_once:
+        output = await engine.async_add_req_and_wait_for_response(_make_request("req"))
+
+    assert output is final
+    warning_once.assert_called_once()
 
 
 def test_abort_request_id_aborts_scheduler_request() -> None:

tests/diffusion/test_diffusion_engine_metrics.py

@@ -80,27 +80,23 @@ class TestMetricKeys:
             "format_diffusion_outputs() — should only use 'preprocess_time_ms'"
         )
 
-    def test_metric_key_naming_consistency(self) -> None:
-        """Metric keys should map to the explicit timing fields."""
+    def test_timing_metric_key_naming_consistency(self) -> None:
+        """Timing metrics should be attached by the engine orchestration path."""
         source = _read_source(_FORMATTER_PATH)
         formatter_source = _get_function_source(source, None, "format_diffusion_outputs")
+        engine_source = _read_source(_ENGINE_PATH)
+        step_streaming_source = _get_function_source(engine_source, "DiffusionEngine", "step_streaming")
         lines = formatter_source.split("\n")
 
-        found_exec = False
-        found_total = False
         for line in lines:
             if '"diffusion_engine_exec_time_ms"' in line:
-                found_exec = True
-                assert "timings.exec_time_s" in line, (
-                    "diffusion_engine_exec_time_ms should measure executor time only (timings.exec_time_s)"
-                )
+                raise AssertionError("diffusion_engine_exec_time_ms should be attached in step_streaming()")
             if '"diffusion_engine_total_time_ms"' in line:
-                found_total = True
-                assert "timings.total_time_ms" in line, (
-                    "diffusion_engine_total_time_ms should measure full step time (timings.total_time_ms)"
-                )
-        assert found_exec, "diffusion_engine_exec_time_ms key not found in format_diffusion_outputs()"
-        assert found_total, "diffusion_engine_total_time_ms key not found in format_diffusion_outputs()"
+                raise AssertionError("diffusion_engine_total_time_ms should be attached in step_streaming()")
+
+        assert '"diffusion_engine_exec_time_ms": exec_total_time * 1000' in step_streaming_source
+        assert '"diffusion_engine_total_time_ms": step_total_ms' in step_streaming_source
+        assert '"postprocess_time_ms": postprocess_time * 1000' in step_streaming_source
 
 
 class TestDummyRunAllocation:

tests/diffusion/test_diffusion_engine_rpc_routing.py

@@ -30,7 +30,7 @@ from typing import Any
 import pytest
 
 from vllm_omni.diffusion.data import DiffusionOutput
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, _RpcTask
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode, _RpcTask
 from vllm_omni.diffusion.sched import RequestScheduler
 from vllm_omni.diffusion.sched.interface import RequestBatchSamplingParamsKey
 from vllm_omni.diffusion.worker.utils import RunnerOutput
@@ -91,7 +91,7 @@ class _ConcurrencyTrackingExecutor:
                 # args[0] is the request-like object for execute_model.
                 req = args[0] if args else None
                 tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
-                return DiffusionOutput(error=f"result_for_{tag}")
+                return DiffusionOutput(error=f"result_for_{tag}", finished=True)
             tag = args[0] if args else method
             return DiffusionOutput(error=f"rpc_result_for_{tag}")
         finally:
@@ -149,12 +149,12 @@ def _make_engine_with_loop(
     engine.scheduler = sched
     engine.step_execution = False
     engine.supports_request_batch = False
+    engine.execution_mode = DiffusionExecutionMode.REQUEST_BATCH
     engine.execute_fn = engine.executor.execute_request
 
     engine._rpc_lock = threading.RLock()
     engine._cv = threading.Condition(engine._rpc_lock)
-    engine._out_queue = {}
-    engine._out_queue_streaming = {}
+    engine._out_streams = {}
     engine._closed = False
     engine.abort_queue = queue.Queue()
     engine._rpc_queue = queue.Queue()
@@ -187,6 +187,15 @@ def _stop_engine(engine: DiffusionEngine) -> None:
     assert not engine.worker_thread.is_alive(), "Busy loop thread did not stop"
 
 
+async def _consume_final_output(generator):
+    final_output = None
+    async for output in generator:
+        final_output = output
+    if final_output is None:
+        raise RuntimeError("Diffusion execution finished without output.")
+    return final_output
+
+
 # ─────────────────────── single-thread invariant ───────────────────────────
 
 
@@ -199,7 +208,7 @@ async def test_executor_only_called_from_busy_loop_thread():
     busy_tid = engine.worker_thread.ident
     try:
         # Per-request path
-        await engine.async_add_req_and_wait_for_response(_make_request("req1"))
+        await _consume_final_output(engine.async_add_req_and_stream_response(_make_request("req1")))
         # Raw RPC path (sync from a worker thread)
         result = await asyncio.to_thread(engine.collective_rpc, "ping", args=("a",), unique_reply_rank=0)
         assert result.error == "rpc_result_for_a"
@@ -228,7 +237,7 @@ async def test_executor_calls_never_overlap_under_load():
             return await engine.async_collective_rpc("ping", args=(f"x{i}",), unique_reply_rank=0)
 
         async def _request(i: int):
-            return await engine.async_add_req_and_wait_for_response(_make_request(f"r{i}"))
+            return await _consume_final_output(engine.async_add_req_and_stream_response(_make_request(f"r{i}")))
 
         tasks = [_rpc(i) for i in range(20)] + [_request(i) for i in range(5)]
         results = await asyncio.gather(*tasks)
@@ -510,7 +519,9 @@ async def test_rpc_and_request_results_do_not_swap():
     engine = _make_engine_with_loop(loop, rpc_delay=0.02)
     try:
         rpc_task = asyncio.create_task(engine.async_collective_rpc("ping", args=("rpc1",), unique_reply_rank=0))
-        req_task = asyncio.create_task(engine.async_add_req_and_wait_for_response(_make_request("req1")))
+        req_task = asyncio.create_task(
+            _consume_final_output(engine.async_add_req_and_stream_response(_make_request("req1")))
+        )
         rpc_res, req_res = await asyncio.gather(rpc_task, req_task)
     finally:
         _stop_engine(engine)

tests/diffusion/test_diffusion_model_runner.py

@@ -98,6 +98,29 @@ class _ChunkStepPipeline:
         return output
 
 
+class _FinalOnlyStepPipeline:
+    device = torch.device("cpu")
+    supports_step_execution = True
+
+    def prepare_encode(self, state):
+        state.prompt_embeds = torch.zeros(1, 1, 1)
+        state.latents = torch.zeros(1, 1)
+        state.timesteps = torch.tensor([1.0])
+        state.step_index = 0
+        return state
+
+    def denoise_step(self, input_batch, states):
+        del states
+        return torch.ones_like(input_batch.latents)
+
+    def step_scheduler(self, state, noise_pred):
+        state.latents = noise_pred
+        state.step_index += 1
+
+    def post_decode(self, state):
+        return DiffusionOutput(output=state.latents.clone())
+
+
 def _make_request():
     sampling_params = SimpleNamespace(
         generator=None,
@@ -228,6 +251,35 @@ def test_execute_stepwise_streaming_returns_chunks_at_boundaries(monkeypatch):
     assert fourth.get_request_output("req").finished is True
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_execute_stepwise_streaming_decodes_final_only_pipeline(monkeypatch):
+    """Step streaming still returns a final output when no chunk boundaries exist."""
+    runner = _make_runner(cache_backend=None, cache_backend_name=None)
+    runner.pipeline = _FinalOnlyStepPipeline()
+    runner.od_config.streaming_output = True
+    runner.od_config.step_execution = True
+    req = _make_request()
+    req.request_id = "req"
+    req.sampling_params.num_inference_steps = 1
+
+    monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "reset_peak_memory_stats", lambda: None)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_reserved", lambda: 0)
+    monkeypatch.setattr(model_runner_module.current_omni_platform, "max_memory_allocated", lambda: 0)
+    scheduler_output = SimpleNamespace(
+        finished_req_ids=set(),
+        scheduled_new_reqs=[SimpleNamespace(request_id="req", req=req)],
+        scheduled_cached_reqs=SimpleNamespace(request_ids=[]),
+    )
+
+    output = DiffusionModelRunner.execute_stepwise(runner, scheduler_output).get_request_output("req")
+
+    assert output.finished is True
+    assert output.result is not None
+    assert torch.equal(output.result.output, torch.ones(1, 1))
+
+
 @pytest.mark.core_model
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 def test_execute_model_skips_cache_summary_without_active_cache_backend(monkeypatch):

tests/diffusion/test_diffusion_output_formatter.py

@@ -11,7 +11,6 @@ import torch
 from vllm_omni.diffusion import output_formatter
 from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.output_formatter import (
-    DiffusionStepTimings,
     format_diffusion_outputs,
     format_empty_diffusion_outputs,
     normalize_diffusion_postprocess_output,
@@ -42,15 +41,6 @@ def _config(model_class_name: str = "mock_model") -> SimpleNamespace:
     return SimpleNamespace(model_class_name=model_class_name)
 
 
-def _timings() -> DiffusionStepTimings:
-    return DiffusionStepTimings(
-        preprocess_time_s=0.01,
-        exec_time_s=0.02,
-        postprocess_time_s=0.03,
-        total_time_ms=60.0,
-    )
-
-
 def test_formatter_preserves_single_video_audio_actions_and_metadata(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
@@ -80,7 +70,6 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
         ),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert len(results) == 1
@@ -102,14 +91,10 @@ def test_formatter_preserves_single_video_audio_actions_and_metadata(
     assert result.stage_durations == {"execute": 1.25}
     assert result.peak_memory_mb == 321.0
     assert result.metrics == {
-        "preprocess_time_ms": 10.0,
-        "diffusion_engine_exec_time_ms": 20.0,
-        "diffusion_engine_total_time_ms": 60.0,
         "image_num": 1,
         "resolution": 512,
         "width": None,
         "height": None,
-        "postprocess_time_ms": 30.0,
     }
 
 
@@ -162,7 +147,6 @@ def test_formatter_normalizes_payload_metadata_envelope(
         diffusion_output=DiffusionOutput(output=None),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.multimodal_output == {
@@ -198,7 +182,6 @@ def test_formatter_maps_trajectory_payload_to_request_output(
         diffusion_output=DiffusionOutput(output=None),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == ["image-0"]
@@ -231,7 +214,6 @@ def test_formatter_preserves_text_envelope_metadata(monkeypatch: pytest.MonkeyPa
         diffusion_output=DiffusionOutput(output=None),
         output_data={"payload": {"text": "caption"}},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == []
@@ -269,7 +251,6 @@ def test_formatter_preserves_image_with_reasoning_metadata(
         diffusion_output=DiffusionOutput(output=None),
         output_data=None,
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert postprocess_output.primary_key == "image"
@@ -298,7 +279,6 @@ def test_formatter_preserves_audio_output_with_model_sample_rate_fallback(
         diffusion_output=DiffusionOutput(output=["waveform"]),
         output_data=["waveform"],
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == []
@@ -330,7 +310,6 @@ def test_formatter_preserves_audio_model_video_audio_and_actions(
         diffusion_output=DiffusionOutput(output=None),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == ["frame-0"]
@@ -365,7 +344,6 @@ def test_formatter_preserves_audio_only_postprocess_dict(
         diffusion_output=DiffusionOutput(output=None),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == []
@@ -397,7 +375,6 @@ def test_formatter_preserves_single_prompt_multiple_audio_outputs(
         diffusion_output=DiffusionOutput(output=["waveform-0", "waveform-1"]),
         output_data=["waveform-0", "waveform-1"],
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert result.images == []
@@ -425,7 +402,6 @@ def test_formatter_preserves_single_prompt_audio_and_action_payloads(
         diffusion_output=DiffusionOutput(output=None),
         output_data={"raw": "output"},
         postprocess_output=postprocess_output,
-        timings=_timings(),
     )
 
     assert len(results) == 1

tests/diffusion/test_diffusion_scheduler.py

@@ -11,7 +11,7 @@ import torch
 from pytest_mock import MockerFixture
 
 from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
-from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine, DiffusionExecutionMode
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import (
     BaseScheduler,
@@ -532,13 +532,11 @@ class TestDiffusionEngine:
         assert output is runner_output.result
         engine.execute_fn.assert_called_once()
 
-    def test_initializes_injected_scheduler(
+    def test_initializes_default_request_scheduler(
         self,
         monkeypatch: pytest.MonkeyPatch,
         mocker: MockerFixture,
     ) -> None:
-        request = _make_request("init")
-        scheduler = _StubScheduler(request, DiffusionOutput(output=None))
         od_config = SimpleNamespace(model_class_name="mock_model", streaming_output=False)
         fake_executor_cls = mocker.Mock(return_value=mocker.Mock())
 
@@ -554,11 +552,11 @@ class TestDiffusionEngine:
             "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
             lambda *args, **kwargs: fake_executor_cls,
         )
-        monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
 
-        DiffusionEngine(od_config, scheduler=scheduler)
+        engine = DiffusionEngine(od_config)
 
-        assert scheduler.initialized_with is od_config
+        assert isinstance(engine.scheduler, RequestScheduler)
+        assert engine.scheduler.max_num_running_reqs == 1
         fake_executor_cls.assert_called_once_with(od_config)
 
     def test_scheduler_alias_keeps_default_request_scheduler(self) -> None:
@@ -573,16 +571,19 @@ class TestDiffusionEngine:
         assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
 
     @pytest.mark.asyncio
-    async def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None:
+    async def test_step_streaming_raises_aborted_error(self, mocker: MockerFixture) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
         engine._check_and_start_background_loop = mocker.AsyncMock()
         engine.pre_process_func = None
-        engine.async_add_req_and_wait_for_response = mocker.AsyncMock(
-            return_value=DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
-        )
+
+        async def _stream(_request):
+            yield DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
+
+        engine.async_add_req_and_stream_response = mocker.Mock(return_value=_stream(None))
 
         with pytest.raises(DiffusionRequestAbortedError, match="Request req-abort aborted"):
-            await engine.step(_make_request("req-abort"))
+            async for _ in engine.step_streaming(_make_request("req-abort")):
+                pass
 
     def test_abort_queue_marks_request_finished_aborted(self) -> None:
         engine = DiffusionEngine.__new__(DiffusionEngine)
@@ -619,12 +620,13 @@ class TestDiffusionEngine:
         engine.scheduler.initialize(SimpleNamespace())
         engine._rpc_lock = threading.RLock()
         engine._cv = threading.Condition(engine._rpc_lock)
-        engine._out_queue_streaming = {}
+        engine._out_streams = {}
+        engine.execution_mode = DiffusionExecutionMode.STEP_BATCH
         engine.main_loop = asyncio.get_running_loop()
 
         req_id = engine.scheduler.add_request(_make_request("stream-engine"))
         queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
-        engine._out_queue_streaming[req_id] = queue
+        engine._out_streams[req_id] = queue
         sched_output = engine.scheduler.schedule()
 
         chunk = RunnerOutput(
@@ -634,7 +636,7 @@ class TestDiffusionEngine:
             result=DiffusionOutput(output="chunk-0", finished=False, chunk_index=0, total_chunks=2),
         )
         finished_req_ids = engine.scheduler.update_from_output(sched_output, chunk)
-        engine._handle_step_streaming_runner_output(finished_req_ids, sched_output.scheduled_request_ids, chunk)
+        engine._emit_outputs(finished_req_ids, sched_output.scheduled_request_ids, chunk)
 
         notified_chunk = await asyncio.wait_for(queue.get(), timeout=1)
         assert notified_chunk.output == "chunk-0"
@@ -648,7 +650,7 @@ class TestDiffusionEngine:
             result=DiffusionOutput(output="chunk-1", finished=True, chunk_index=1, total_chunks=2),
         )
         finished_req_ids = engine.scheduler.update_from_output(sched_output, final_chunk)
-        engine._handle_step_streaming_runner_output(finished_req_ids, sched_output.scheduled_request_ids, final_chunk)
+        engine._emit_outputs(finished_req_ids, sched_output.scheduled_request_ids, final_chunk)
 
         notified_final = await asyncio.wait_for(queue.get(), timeout=1)
         assert notified_final.output == "chunk-1"
@@ -662,15 +664,15 @@ class TestDiffusionEngine:
         engine.scheduler.initialize(SimpleNamespace())
         engine._rpc_lock = threading.RLock()
         engine._cv = threading.Condition(engine._rpc_lock)
-        engine._out_queue_streaming = {}
+        engine._out_streams = {}
         engine.main_loop = asyncio.get_running_loop()
 
         req_id = engine.scheduler.add_request(_make_request("stream-abort"))
         queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
-        engine._out_queue_streaming[req_id] = queue
+        engine._out_streams[req_id] = queue
         engine.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED)
 
-        engine._handle_empty_streaming_requests({req_id})
+        engine._emit_finished_outputs({req_id})
 
         output = await asyncio.wait_for(queue.get(), timeout=1)
         assert output.aborted is True
@@ -699,9 +701,9 @@ class TestDiffusionEngine:
             "vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
             lambda *args, **kwargs: fake_executor_cls,
         )
-        monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
         engine = DiffusionEngine(od_config)
 
+        assert engine.execution_mode == DiffusionExecutionMode.STEP_BATCH
         assert isinstance(engine.scheduler, StepScheduler)
         assert engine.execute_fn is fake_executor.execute_step
         fake_executor_cls.assert_called_once_with(od_config)

tests/diffusion/test_inline_stage_diffusion_client.py

@@ -54,9 +54,12 @@ def client(mock_engine):
 
 @pytest.mark.asyncio
 async def test_inline_dispatch_request_success(client, mock_engine):
-    # Setup mock engine step to return a successful result
     mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()])
-    mock_engine.step.return_value = [mock_result]
+
+    async def _step_streaming(_request):
+        yield [mock_result]
+
+    mock_engine.step_streaming = _step_streaming
 
     sampling_params = OmniDiffusionSamplingParams()
     await client.add_request_async("req-1", "A test prompt", sampling_params)
@@ -70,7 +73,6 @@ async def test_inline_dispatch_request_success(client, mock_engine):
 
     assert output is not None
     assert output.request_id == "req-1"
-    mock_engine.step.assert_called_once()
 
 
 @pytest.mark.asyncio
@@ -104,8 +106,11 @@ async def test_inline_dispatch_request_streaming_success(client, mock_engine):
 
 @pytest.mark.asyncio
 async def test_inline_dispatch_request_error(client, mock_engine):
-    # Setup mock engine step to raise an exception
-    mock_engine.step.side_effect = RuntimeError("Engine failure")
+    async def _step_streaming(_request):
+        raise RuntimeError("Engine failure")
+        yield  # pragma: no cover
+
+    mock_engine.step_streaming = _step_streaming
 
     sampling_params = OmniDiffusionSamplingParams()
     await client.add_request_async("req-err", "A test prompt", sampling_params)

tests/diffusion/test_multiproc_engine_concurrency.py

@@ -91,7 +91,7 @@ def _make_engine(num_gpus: int = 1):
     engine._loop_started = False
     engine._rpc_queue = queue.Queue()
     engine.abort_queue = queue.Queue()
-    engine.execute_fn = executor.execute_request
+    engine.execute_fn = executor.execute_batch
     return engine, executor, req_q, res_q
 
 

tests/diffusion/test_stage_diffusion_proc.py

@@ -27,7 +27,7 @@ DELAY_BASE = 0.01
 
 
 class MockDiffusionEngine:
-    async def step(self, request):
+    async def step_streaming(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)
 
@@ -37,7 +37,7 @@ class MockDiffusionEngine:
         )
         delay = DELAY_BASE + delay_scale * DELAY_BASE
         await asyncio.sleep(delay)
-        return [MockOmniRequestOutput(request_id=request.request_id)]
+        yield [MockOmniRequestOutput(request_id=request.request_id)]
 
 
 @pytest.mark.asyncio

tests/e2e/offline_inference/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -158,6 +158,7 @@ async def test_async_omni_generate():
             custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
             worker_extension_cls=WORKER_EXTENSION_CLASS,
             enforce_eager=True,
+            max_num_seqs=1,
         )
         after.callback(engine.shutdown)
 
@@ -183,6 +184,7 @@ async def test_async_omni_generate_with_logprobs():
             custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
             worker_extension_cls=WORKER_EXTENSION_CLASS,
             enforce_eager=True,
+            max_num_seqs=1,
         )
         after.callback(engine.shutdown)
 
@@ -220,6 +222,7 @@ async def test_async_omni_generate_concurrent():
             custom_pipeline_args={"pipeline_class": CUSTOM_PIPELINE_CLASS},
             worker_extension_cls=WORKER_EXTENSION_CLASS,
             enforce_eager=True,
+            max_num_seqs=1,
         )
         after.callback(engine.shutdown)
 

tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py

@@ -283,7 +283,7 @@ class vLLMOmniHttpServerLocal:
             "worker_extension_cls": WORKER_EXTENSION_CLASS,
             "trust_remote_code": _cfg_get(self.model_config, "trust_remote_code", True),
             "max_model_len": _cfg_get(self.config, "max_model_len", 1058),
-            "max_num_seqs": _cfg_get(self.config, "max_num_seqs", 256),
+            "max_num_seqs": _cfg_get(self.config, "max_num_seqs", 1),
             "enable_chunked_prefill": _cfg_get(self.config, "enable_chunked_prefill", False),
             "max_num_batched_tokens": _cfg_get(self.config, "max_num_batched_tokens", 8192),
             "enable_prefix_caching": _cfg_get(self.config, "enable_prefix_caching", False),
@@ -476,7 +476,8 @@ def init_server():
             "pipeline_model_parallel_size": 1,
             "gpu_memory_utilization": 0.8,
             "max_num_batched_tokens": 8192,
-            "max_num_seqs": 256,
+            # The custom log-prob pipeline only supports serial requests.
+            "max_num_seqs": 1,
             "max_model_len": 1058,
             "dtype": "bfloat16",
             "load_format": "auto",

vllm_omni/diffusion/diffusion_engine.py

@@ -11,6 +11,7 @@ import threading
 import time
 from collections.abc import AsyncGenerator, Iterable
 from dataclasses import dataclass, field
+from enum import Enum
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
@@ -33,7 +34,6 @@ from vllm_omni.diffusion.io_support import (
     supports_multimodal_input,
 )
 from vllm_omni.diffusion.output_formatter import (
-    DiffusionStepTimings,
     format_diffusion_outputs,
     format_empty_diffusion_outputs,
     normalize_diffusion_postprocess_output,
@@ -57,6 +57,7 @@ logger = init_logger(__name__)
 
 __all__ = [
     "DiffusionEngine",
+    "DiffusionExecutionMode",
     "_RpcTask",
     "_move_tensor_tree_to_cpu",
     "get_dummy_run_num_frames",
@@ -106,6 +107,13 @@ def supports_request_batch(od_config: OmniDiffusionConfig) -> bool:
     return bool(getattr(model_cls, "supports_request_batch", False))
 
 
+def _max_num_seqs(od_config: OmniDiffusionConfig) -> int:
+    try:
+        return max(1, int(getattr(od_config, "max_num_seqs", 1)))
+    except (TypeError, ValueError):
+        return 1
+
+
 def _move_tensor_tree_to_cpu(value: object) -> object:
     if isinstance(value, torch.Tensor):
         return value.cpu() if value.device.type != "cpu" else value
@@ -130,6 +138,11 @@ class _RpcTask:
     future: concurrent.futures.Future = field(default_factory=concurrent.futures.Future)
 
 
+class DiffusionExecutionMode(str, Enum):
+    REQUEST_BATCH = "request_batch"
+    STEP_BATCH = "step_batch"
+
+
 class DiffusionEngine:
     """The diffusion engine for vLLM-Omni diffusion models."""
 
@@ -149,27 +162,66 @@ class DiffusionEngine:
         """Initialize the diffusion engine.
 
         Args:
-            config: The configuration for the diffusion engine.
+            od_config: The configuration for the diffusion engine.
+            scheduler: Optional scheduler override for tests or custom engine
+                integrations. When omitted, the engine selects a scheduler
+                from the resolved execution mode.
         """
         self.od_config = od_config
 
+        self._init_process_hooks(od_config)
+        self.execution_mode = self._resolve_execution_mode(od_config)
+        self._init_executor(od_config)
+        self._init_scheduler(od_config, scheduler)
+        self._init_runtime_state()
+        self._init_execute_fn()
+        self._log_execution_mode(od_config)
+
+    def _init_process_hooks(self, od_config: OmniDiffusionConfig) -> None:
         self.post_process_func = get_diffusion_post_process_func(od_config)
         self.pre_process_func = get_diffusion_pre_process_func(od_config)
         # Cache whether the model-specific postprocess accepts request-level
         # sampling params so step() can support both legacy and extended hooks.
         self._post_process_accepts_sampling_params = _func_accepts_parameter(self.post_process_func, "sampling_params")
 
+    def _resolve_execution_mode(self, od_config: OmniDiffusionConfig) -> DiffusionExecutionMode:
         self.step_execution = bool(getattr(od_config, "step_execution", False))
-        if self.od_config.streaming_output and not self.step_execution:
+        if od_config.streaming_output and not self.step_execution:
             logger.warning("streaming_output=True requires step_execution=True; enabling step execution.")
-            self.od_config.step_execution = True
+            od_config.step_execution = True
             self.step_execution = True
 
+        if self.step_execution:
+            self.supports_request_batch = False
+            return DiffusionExecutionMode.STEP_BATCH
+
+        self.supports_request_batch = supports_request_batch(od_config)
+        if not self.supports_request_batch and _max_num_seqs(od_config) > 1:
+            raise ValueError(
+                f"{getattr(od_config, 'model_class_name', None)!r} does not support request-level batching. "
+                "Use max_num_seqs=1 for serial request execution, or choose a pipeline with "
+                "supports_request_batch=True."
+            )
+        return DiffusionExecutionMode.REQUEST_BATCH
+
+    def _init_executor(self, od_config: OmniDiffusionConfig) -> None:
         executor_class = DiffusionExecutor.get_class(od_config)
         self.executor = executor_class(od_config)
-        self.scheduler: BaseScheduler = scheduler or (StepScheduler() if self.step_execution else RequestScheduler())
+
+    def _init_scheduler(
+        self,
+        od_config: OmniDiffusionConfig,
+        scheduler: BaseScheduler | None = None,
+    ) -> None:
+        if scheduler is not None:
+            self.scheduler = scheduler
+        elif self.execution_mode == DiffusionExecutionMode.STEP_BATCH:
+            self.scheduler = StepScheduler()
+        else:
+            self.scheduler = RequestScheduler()
         self.scheduler.initialize(od_config)
-        self.supports_request_batch = False if self.step_execution else supports_request_batch(od_config)
+
+    def _init_runtime_state(self) -> None:
         self.main_loop: asyncio.AbstractEventLoop | None = None
         self.stop_event: threading.Event | None = None
         self.worker_thread: threading.Thread | None = None
@@ -181,33 +233,26 @@ class DiffusionEngine:
         # busy loop rather than via mutual exclusion.
         self._rpc_lock = threading.RLock()
         self._cv = threading.Condition(self._rpc_lock)
-        self._out_queue: dict[str, asyncio.Future] = {}
-        self._out_queue_streaming: dict[str, asyncio.Queue[DiffusionOutput]] = {}
+        self._out_streams: dict[str, asyncio.Queue[DiffusionOutput]] = {}
         self._closed = False
         self._shutdown_complete = False
         self.abort_queue: queue.Queue[str] = queue.Queue()
         self._rpc_queue: queue.Queue[_RpcTask] = queue.Queue()
-        if self.step_execution:
-            self.execute_fn = self.executor.execute_step
-        elif self.supports_request_batch:
-            self.execute_fn = self.executor.execute_batch
-        else:
-            self.execute_fn = self.executor.execute_request
 
-        if self.supports_request_batch:
+    def _init_execute_fn(self) -> None:
+        if self.execution_mode == DiffusionExecutionMode.STEP_BATCH:
+            self.execute_fn = self.executor.execute_step
+        else:
+            self.execute_fn = self.executor.execute_batch
+
+    def _log_execution_mode(self, od_config: OmniDiffusionConfig) -> None:
+        if self.execution_mode == DiffusionExecutionMode.REQUEST_BATCH:
             logger.info(
                 "[RequestBatch] engine init max_num_seqs=%s max_wait_ms=%s",
                 getattr(od_config, "max_num_seqs", None),
                 getattr(od_config, "request_batch_max_wait_ms", None),
             )
 
-        try:
-            self._dummy_run()
-        except Exception as e:
-            logger.error(f"Dummy run failed: {e}")
-            self.close()
-            raise e
-
     async def _check_and_start_background_loop(self):
         if self._closed:
             raise RuntimeError("DiffusionEngine is closed.")
@@ -227,24 +272,6 @@ class DiffusionEngine:
             self.worker_thread.start()
             self._loop_started = True
 
-    async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
-        await self._check_and_start_background_loop()
-
-        diffusion_engine_start_time = time.perf_counter()
-
-        # Apply pre-processing if available
-        preprocess_time = 0.0
-        if self.pre_process_func is not None:
-            preprocess_start_time = time.perf_counter()
-            request = self.pre_process_func(request)
-            preprocess_time = time.perf_counter() - preprocess_start_time
-            logger.debug("Pre-processing completed in %.4f seconds", preprocess_time)
-
-        exec_start_time = time.perf_counter()
-        output = await self.async_add_req_and_wait_for_response(request)
-        exec_total_time = time.perf_counter() - exec_start_time
-        return self.postprocess_output(request, output, diffusion_engine_start_time, preprocess_time, exec_total_time)
-
     async def step_streaming(self, request: OmniDiffusionRequest) -> AsyncGenerator[list[OmniRequestOutput], None]:
         await self._check_and_start_background_loop()
 
@@ -261,19 +288,50 @@ class DiffusionEngine:
         generator = self.async_add_req_and_stream_response(request)
         async for output in generator:
             exec_total_time = time.perf_counter() - exec_start_time
-            yield self.postprocess_output(
-                request, output, diffusion_engine_start_time, preprocess_time, exec_total_time
+            postprocess_start_time = time.perf_counter()
+            formatted_outputs = self.postprocess_output(request, output)
+            postprocess_time = time.perf_counter() - postprocess_start_time
+            step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000
+            logger.debug(
+                "DiffusionEngine.step_streaming breakdown: preprocess=%.2f ms, "
+                "add_req_and_wait=%.2f ms, postprocess=%.2f ms, total=%.2f ms",
+                preprocess_time * 1000,
+                exec_total_time * 1000,
+                postprocess_time * 1000,
+                step_total_ms,
             )
+            for request_output in formatted_outputs:
+                request_output.metrics.update(
+                    {
+                        "preprocess_time_ms": preprocess_time * 1000,
+                        "diffusion_engine_exec_time_ms": exec_total_time * 1000,
+                        "diffusion_engine_total_time_ms": step_total_ms,
+                        "postprocess_time_ms": postprocess_time * 1000,
+                    }
+                )
+            yield formatted_outputs
+
+    async def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
+        """Deprecated compatibility wrapper over ``step_streaming()``.
+
+        Use ``step_streaming()`` for new callers. This method drains the
+        unified output stream and returns only the final output batch, matching
+        the historical non-streaming ``step()`` behavior.
+        """
+        logger.warning_once(
+            "DiffusionEngine.step() is deprecated; use step_streaming() and consume the final output batch instead."
+        )
+        final_output: list[OmniRequestOutput] | None = None
+        async for output in self.step_streaming(request):
+            final_output = output
+        return final_output or []
 
     def postprocess_output(
         self,
         request: OmniDiffusionRequest,
         output: DiffusionOutput,
-        diffusion_engine_start_time: float,
-        preprocess_time: float,
-        exec_total_time: float,
     ) -> list[OmniRequestOutput]:
-        """Convert a DiffusionOutput to a list of OmniRequestOutput, attaching profiling metrics."""
+        """Convert a DiffusionOutput to a list of OmniRequestOutput."""
         if output.aborted:
             raise DiffusionRequestAbortedError(output.abort_message or "Diffusion request aborted.")
         if output.error:
@@ -297,7 +355,6 @@ class DiffusionEngine:
         if self.od_config.enable_cpu_offload:
             output_data = _move_tensor_tree_to_cpu(output_data)
 
-        postprocess_start_time = time.perf_counter()
         if self.post_process_func is not None:
             # Some video pipelines need request-level controls during
             # postprocess (for example worker-side frame interpolation).
@@ -309,18 +366,6 @@ class DiffusionEngine:
             outputs = output_data
 
         postprocess_output = normalize_diffusion_postprocess_output(outputs)
-        postprocess_time = time.perf_counter() - postprocess_start_time
-        logger.debug("Post-processing completed in %.4f seconds", postprocess_time)
-
-        step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000
-        logger.debug(
-            "DiffusionEngine.step breakdown: preprocess=%.2f ms, "
-            "add_req_and_wait=%.2f ms, postprocess=%.2f ms, total=%.2f ms",
-            preprocess_time * 1000,
-            exec_total_time * 1000,
-            postprocess_time * 1000,
-            step_total_ms,
-        )
 
         return format_diffusion_outputs(
             request=request,
@@ -328,12 +373,6 @@ class DiffusionEngine:
             diffusion_output=output,
             output_data=output_data,
             postprocess_output=postprocess_output,
-            timings=DiffusionStepTimings(
-                preprocess_time_s=preprocess_time,
-                exec_time_s=exec_total_time,
-                postprocess_time_s=postprocess_time,
-                total_time_ms=step_total_ms,
-            ),
         )
 
     def _busy_loop(self):
@@ -357,16 +396,13 @@ class DiffusionEngine:
                     # Only RPC / abort work pending; loop back to drain it.
                     continue
 
-                if self.supports_request_batch:
+                if self.execution_mode == DiffusionExecutionMode.REQUEST_BATCH:
                     self._wait_for_request_batch_admission_locked()
 
                 sched_output = self.scheduler.schedule()
 
             if sched_output.is_empty:
-                if self.od_config.streaming_output:
-                    self._handle_empty_streaming_requests(sched_output.finished_req_ids)
-                else:
-                    self._handle_finished_requests(sched_output.finished_req_ids, None)
+                self._emit_finished_outputs(sched_output.finished_req_ids, None)
                 continue
 
             try:
@@ -390,14 +426,7 @@ class DiffusionEngine:
             self._process_aborts_queue()
             self._process_rpc_queue()
             finished_req_ids = self.scheduler.update_from_output(sched_output, runner_output)
-            if self.od_config.streaming_output:
-                self._handle_step_streaming_runner_output(
-                    finished_req_ids,
-                    sched_output.scheduled_request_ids,
-                    runner_output,
-                )
-            else:
-                self._handle_finished_requests(finished_req_ids, runner_output)
+            self._emit_outputs(finished_req_ids, sched_output.scheduled_request_ids, runner_output)
 
         # Engine is stopping: fail any RPCs still queued so callers don't hang.
         self._fail_pending_rpcs(RuntimeError("DiffusionEngine is shutting down."))
@@ -509,74 +538,62 @@ class DiffusionEngine:
             if not task.future.done():
                 task.future.set_exception(exc)
 
-    def _handle_finished_requests(
+    def _emit_finished_outputs(
         self,
         finished_ids: set[str],
         runner_output: BaseRunnerOutput | None = None,
         missing_result_error: str = "Diffusion execution finished without a final output",
-    ):
+    ) -> None:
         for rid in finished_ids:
-            with self._cv:
-                fut = self._out_queue.pop(rid, None)
-            if fut is None:
-                continue
             if runner_output is not None:
                 _output = runner_output.get_request_output(rid)
             else:
                 _output = None
             out = self._finalize_finished_request(rid, _output, missing_result_error)
-            self._complete_future(fut, out)
+            self._put_output(rid, out)
 
-    def _handle_empty_streaming_requests(
+    def _emit_outputs(
         self,
         finished_ids: set[str],
-        missing_result_error: str = "Diffusion streaming request finished without execution output.",
-    ) -> None:
-        """Mirrors `_handle_finished_requests()` in non-streaming mode when used for empty scheduler output."""
-        for rid in finished_ids:
-            out = self._finalize_finished_request(rid, None, missing_result_error=missing_result_error)
-            self._put_streaming_output_with_cv(rid, out)
-
-    def _handle_step_streaming_runner_output(
-        self,
-        finished_req_ids: set[str],
         scheduled_request_ids: list[str],
         runner_output: BaseRunnerOutput,
     ) -> None:
-        """
-        Deliver partial step-execution outputs in streaming mode.
+        """Emit output chunks for every request through the unified output stream."""
+        if self.execution_mode != DiffusionExecutionMode.STEP_BATCH:
+            self._emit_finished_outputs(finished_ids, runner_output)
+            return
 
-        Step execution returns one ``RunnerOutput`` per scheduled request per
-        engine tick. Most denoise steps have ``result=None``; chunk boundaries
-        return a ``DiffusionOutput`` that must be delivered even before the
-        scheduler marks the request finished.
-        """
         delivered_finished_req_ids: set[str] = set()
 
         # finished_ids may have some requests that are not scheduler in this round.
         # First handle this-round requests.
         for request_id in scheduled_request_ids:
             req_output = runner_output.get_request_output(request_id)
-            if request_id in finished_req_ids:
+            if request_id in finished_ids:
                 # This entire request is finished (this is the last chunk)
                 out = self._finalize_finished_request(
                     request_id,
                     req_output,
-                    missing_result_error="Diffusion streaming execution finished without a final output.",
+                    missing_result_error="Diffusion step execution finished without a final output.",
                 )
-                self._put_streaming_output_with_cv(request_id, out)
+                self._put_output(request_id, out)
                 delivered_finished_req_ids.add(request_id)
             elif req_output is not None and req_output.result is not None:
-                # This is a non-terminal chunk. So it is not in scheduler's finished_req_ids, but still need delivering.
-                self._put_streaming_output_with_cv(request_id, req_output.result)
+                # This is a non-terminal chunk. It is not in scheduler's
+                # finished_ids, but it still belongs on the request stream.
+                self._put_output(request_id, req_output.result)
 
         # Then handle other requests that are finished in this round.
-        for request_id in finished_req_ids - delivered_finished_req_ids:
+        for request_id in finished_ids - delivered_finished_req_ids:
             out = self._finalize_finished_request(
                 request_id,
-                missing_result_error="Diffusion streaming request finished without execution output.",
+                missing_result_error="Diffusion step request finished without execution output.",
             )
-            self._put_streaming_output_with_cv(request_id, out)
+            self._put_output(request_id, out)
+
+    def _has_output_stream(self, request_id: str) -> bool:
+        with self._cv:
+            return request_id in self._out_streams
 
     @staticmethod
     def resolve_engine_class(config: OmniDiffusionConfig) -> type[DiffusionEngine]:
@@ -622,44 +639,31 @@ class DiffusionEngine:
 
         Args:
             config: The configuration for the diffusion engine.
+            scheduler: Optional scheduler override. When omitted, the selected
+                engine chooses the scheduler from its execution mode.
 
         Returns:
             An instance of the resolved ``DiffusionEngine`` (sub)class.
         """
         engine_class = DiffusionEngine.resolve_engine_class(config)
-        return engine_class(config, scheduler=scheduler)
+        engine = engine_class(config, scheduler=scheduler)
+        engine.run_startup_warmup()
+        return engine
 
     def add_request(self, request: OmniDiffusionRequest) -> str:
         with self._cv:
             if self._closed:
                 raise RuntimeError("DiffusionEngine is closed.")
-            if not self.od_config.streaming_output:
-                fut = self.main_loop.create_future()
-                request_id = self.scheduler.add_request(request)
-                self._out_queue[request_id] = fut
-            else:
-                queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
-                request_id = self.scheduler.add_request(request)
-                self._out_queue_streaming[request_id] = queue
+            queue: asyncio.Queue[DiffusionOutput] = asyncio.Queue()
+            request_id = self.scheduler.add_request(request)
+            self._out_streams[request_id] = queue
             self._cv.notify_all()
 
         return request_id
 
-    async def get_result(self, request_id: str) -> DiffusionOutput:
-        fut = self._out_queue.get(request_id)
-
-        if fut is None:
-            raise RuntimeError(f"Request {request_id} not found in output queue.")
-        try:
-            return await fut
-        except Exception as e:
-            logger.error(f"Wait for response failed: {e}")
-            raise
-
-    async def get_streaming_result(self, request_id: str) -> AsyncGenerator[DiffusionOutput, None]:
-        """Mirrors `get_result()` in non-streaming mode."""
+    async def get_output_stream(self, request_id: str) -> AsyncGenerator[DiffusionOutput, None]:
         with self._cv:
-            queue = self._out_queue_streaming.get(request_id)
+            queue = self._out_streams.get(request_id)
         if queue is None:
             raise RuntimeError(f"Request {request_id} not found in output queue.")
         try:
@@ -672,22 +676,31 @@ class DiffusionEngine:
             logger.error(f"Wait for response failed: {e}")
             raise
         finally:
-            # In streaming mode, an output queue is maintained until the terminal chunk is met.
-            # So unlike the non-streaming mode where output Future is popped in `_handle_finished_requests` (immediately
-            # after the request is returned), the streaming mode needs to pop the output queue here (one layer above).
             with self._cv:
-                if self._out_queue_streaming.get(request_id) is queue:
-                    self._out_queue_streaming.pop(request_id, None)
-
-    async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
-        # No lock needed: add_request is already protected by self._cv, and
-        # all executor calls are serialized inside the busy loop.
-        request_id = self.add_request(request)
-        return await self.get_result(request_id)
+                if self._out_streams.get(request_id) is queue:
+                    self._out_streams.pop(request_id, None)
 
     def async_add_req_and_stream_response(self, request: OmniDiffusionRequest) -> AsyncGenerator[DiffusionOutput, None]:
         request_id = self.add_request(request)
-        return self.get_streaming_result(request_id)
+        return self.get_output_stream(request_id)
+
+    async def async_add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
+        """Deprecated compatibility wrapper over ``async_add_req_and_stream_response()``.
+
+        Use ``async_add_req_and_stream_response()`` for new callers. This
+        method drains the unified output stream and returns only the final
+        ``DiffusionOutput``, matching the historical non-streaming behavior.
+        """
+        logger.warning_once(
+            "DiffusionEngine.async_add_req_and_wait_for_response() is deprecated; "
+            "use async_add_req_and_stream_response() and consume the final output instead."
+        )
+        final_output: DiffusionOutput | None = None
+        async for output in self.async_add_req_and_stream_response(request):
+            final_output = output
+        if final_output is None:
+            raise RuntimeError("Diffusion execution completed without an output.")
+        return final_output
 
     def add_req_and_wait_for_response(self, request: OmniDiffusionRequest) -> DiffusionOutput:
         with self._rpc_lock:
@@ -759,6 +772,14 @@ class DiffusionEngine:
             if is_start:
                 raise RuntimeError(f"Could not {action} profiler: {e}") from e
 
+    def run_startup_warmup(self) -> None:
+        try:
+            self._dummy_run()
+        except Exception as e:
+            logger.error(f"Dummy run failed: {e}")
+            self.close()
+            raise e
+
     def _dummy_run(self):
         """A dummy run to warm up the model."""
         num_inference_steps = 1
@@ -891,8 +912,8 @@ class DiffusionEngine:
     ) -> Any:
         """Async variant of :meth:`collective_rpc` for event-loop callers.
 
-        Mirrors :meth:`async_add_req_and_wait_for_response`: enqueue a task
-        keyed by a future and ``await`` the result without blocking the loop.
+        Enqueue a task keyed by a future and ``await`` the result without
+        blocking the loop.
         """
         await self._check_and_start_background_loop()
         task = self._submit_rpc(method, timeout, args, kwargs, unique_reply_rank)
@@ -905,51 +926,26 @@ class DiffusionEngine:
             task.future.cancel()
             raise TimeoutError(f"RPC call to {method} timed out.") from exc
 
-    def _complete_future(self, fut: asyncio.Future, output: DiffusionOutput) -> None:
-        if fut.done():
-            return
-
-        def _set_result() -> None:
-            if not fut.done():
-                fut.set_result(output)
-
-        try:
-            loop = fut.get_loop()
-        except AttributeError:
-            loop = self.main_loop
-
-        try:
-            running_loop = asyncio.get_running_loop()
-        except RuntimeError:
-            running_loop = None
-
-        if loop is not None and loop.is_running() and loop is not running_loop:
-            loop.call_soon_threadsafe(_set_result)
-        else:
-            _set_result()
-
-    def _put_streaming_queue_output(
+    def _put_queue_output(
         self,
         queue: asyncio.Queue[DiffusionOutput],
         output: DiffusionOutput,
     ) -> None:
-        """Append to streaming output queue in a safe event loop. Mirrors `_complete_future()` in non-streaming mode."""
         loop = self.main_loop
         if loop is not None and loop.is_running():
             loop.call_soon_threadsafe(queue.put_nowait, output)
         else:
             queue.put_nowait(output)
 
-    def _put_streaming_output_with_cv(self, request_id: str, output: DiffusionOutput) -> None:
+    def _put_output(self, request_id: str, output: DiffusionOutput) -> None:
         with self._cv:
-            queue = self._out_queue_streaming.get(request_id)
+            queue = self._out_streams.get(request_id)
         if queue is None:
             return
-        self._put_streaming_queue_output(queue, output)
+        self._put_queue_output(queue, output)
 
     def close(self) -> None:
-        pending_futures: list[asyncio.Future] = []
-        pending_streaming_queues: list[asyncio.Queue[DiffusionOutput]] = []
+        pending_streams: list[asyncio.Queue[DiffusionOutput]] = []
         with self._cv:
             if self._closed and self._shutdown_complete:
                 return
@@ -957,17 +953,13 @@ class DiffusionEngine:
                 self._closed = True
                 if self.stop_event is not None:
                     self.stop_event.set()
-                pending_futures = list(self._out_queue.values())
-                pending_streaming_queues = list(self._out_queue_streaming.values())
-                self._out_queue.clear()
-                self._out_queue_streaming.clear()
+                pending_streams = list(self._out_streams.values())
+                self._out_streams.clear()
                 self._cv.notify_all()
 
         closed_output = DiffusionOutput(error="DiffusionEngine is closed.")
-        for fut in pending_futures:
-            self._complete_future(fut, closed_output)
-        for streaming_queue in pending_streaming_queues:
-            self._put_streaming_queue_output(streaming_queue, closed_output)
+        for stream in pending_streams:
+            self._put_queue_output(stream, closed_output)
 
         worker_thread = self.worker_thread
         if worker_thread is not None:

vllm_omni/diffusion/executor/abstract.py

@@ -81,7 +81,7 @@ class DiffusionExecutor(ABC):
 
     @abstractmethod
     def execute_batch(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
-        """Execute request-mode work as a single batched RPC."""
+        """Execute request-mode work through the request-batch path."""
         pass
 
     @abstractmethod

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -356,14 +356,18 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         return BatchRunnerOutput.from_list(runner_outputs)
 
     def execute_batch(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
-        """Execute request-mode work through a single batched worker RPC.
+        """Execute request-mode work through the unified request-batch path.
 
-        The worker builds DiffusionRequestBatch from scheduler output and returns
-        BatchRunnerOutput with one RunnerOutput per scheduled request.
+        A scheduler wave with one request is the conservative serial case and
+        uses the single-request worker RPC. Waves with multiple requests use the
+        fused request-batch RPC and require pipeline request-batch support.
         """
         from vllm_omni.diffusion.worker.utils import BatchRunnerOutput
 
         self._ensure_open()
+        if len(scheduler_output.scheduled_new_reqs) <= 1:
+            return self.execute_request(scheduler_output)
+
         result = self.collective_rpc(
             "execute_model_batch",
             args=(scheduler_output, self.od_config),

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -139,8 +139,13 @@ class InlineStageDiffusionClient(StageClientBase):
                         result.request_id = request_id
                     self._output_queue.put_nowait(result)
             else:
-                results = await self._engine.step(request)
-                result = results[0]
+                # Non-streaming callers share the streaming engine path but
+                # only publish the final output.
+                result = None
+                async for results in self._engine.step_streaming(request):
+                    result = results[0]
+                if result is None:
+                    raise RuntimeError("Diffusion execution finished without output.")
                 if not result.request_id:
                     result.request_id = request_id
                 self._output_queue.put_nowait(result)

vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

@@ -3,7 +3,7 @@
 
 """DreamZero pipeline for vllm-omni.
 
-Entry point for DiffusionEngine.step() -> pipeline.forward(req)
+Entry point for DiffusionEngine.step_streaming() -> pipeline.forward(req)
 """
 
 from __future__ import annotations
@@ -1163,7 +1163,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin):
 
     @torch.no_grad()
     def forward(self, req: DiffusionRequestBatch, **kwargs) -> DiffusionOutput:
-        """Full inference step. Called by DiffusionEngine.step()."""
+        """Full inference step. Called by DiffusionEngine.step_streaming()."""
         extra_args = req.sampling_params.extra_args or {}
         robot_obs = extra_args.get("robot_obs")
         if robot_obs is None:

vllm_omni/diffusion/output_formatter.py

@@ -27,14 +27,6 @@ from vllm_omni.outputs.output_metadata import (
 )
 
 
-@dataclass(frozen=True)
-class DiffusionStepTimings:
-    preprocess_time_s: float
-    exec_time_s: float
-    postprocess_time_s: float
-    total_time_ms: float
-
-
 @dataclass(frozen=True)
 class DiffusionPostprocessOutput:
     outputs: DiffusionPayload
@@ -134,23 +126,18 @@ def format_diffusion_outputs(
     diffusion_output: DiffusionOutput,
     output_data: DiffusionPostprocessRawOutput,
     postprocess_output: DiffusionPostprocessOutput,
-    timings: DiffusionStepTimings,
 ) -> list[OmniRequestOutput]:
     """Convert a finished diffusion model output into API-facing outputs."""
 
     primary_payload = _primary_payload(postprocess_output)
     outputs = _ensure_list(primary_payload)
     metrics = {
-        "preprocess_time_ms": timings.preprocess_time_s * 1000,
-        "diffusion_engine_exec_time_ms": timings.exec_time_s * 1000,
-        "diffusion_engine_total_time_ms": timings.total_time_ms,
         "image_num": int(request.sampling_params.num_outputs_per_prompt),
         "resolution": (
             int(request.sampling_params.resolution) if request.sampling_params.resolution is not None else None
         ),
         "width": request.sampling_params.width,
         "height": request.sampling_params.height,
-        "postprocess_time_ms": timings.postprocess_time_s * 1000,
     }
 
     # Only the primary payload determines the response type. Some image models

vllm_omni/diffusion/output_metadata.py

@@ -1,10 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Compatibility imports for diffusion output metadata.
-
-The shared output payload and metadata schema lives in
-``vllm_omni.outputs.output_metadata``.  Keep this module as a stable import
-path for existing diffusion callers.
-"""
-
-from vllm_omni.outputs.output_metadata import *  # noqa: F403

vllm_omni/diffusion/stage_diffusion_proc.py

@@ -155,7 +155,7 @@ class StageDiffusionProc:
         sampling_params_dict: dict,
         kv_sender_info: dict[str, Any] | None = None,
     ) -> OmniRequestOutput:
-        """Build a diffusion request and run DiffusionEngine.step()."""
+        """Build a diffusion request and consume DiffusionEngine.step_streaming() to completion."""
         sampling_params = self._reconstruct_sampling_params(sampling_params_dict)
 
         request = OmniDiffusionRequest(
@@ -165,8 +165,13 @@ class StageDiffusionProc:
             kv_sender_info=kv_sender_info,
         )
 
-        results = await self._engine.step(request)
-        result = results[0]
+        # Non-streaming callers share the streaming engine path but only
+        # return the final output.
+        result = None
+        async for results in self._engine.step_streaming(request):
+            result = results[0]
+        if result is None:
+            raise RuntimeError("Diffusion execution finished without output.")
         if not result.request_id:
             result.request_id = request_id
         return result

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -714,7 +714,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                                 req, noise_pred[offset : offset + row_num] if noise_pred is not None else None
                             )
                             if self.od_config.streaming_output:
-                                should_decode = req.chunk_denoise_completed
+                                should_decode = req.chunk_denoise_completed or req.request_denoise_completed
                             else:
                                 should_decode = req.denoise_completed