[Core] Add UniProc diffusion executor for single-GPU (#6308)

Signed-off-by: Rahul Steiger <rsteiger@nvidia.com>
Signed-off-by: Rahul Steiger <rsteiger@aws-cmh-slurm-1-vscode-04.cm.cluster>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Rahul Steiger <rsteiger@aws-cmh-slurm-1-vscode-04.cm.cluster>

docs/configuration/stage_configs.md

@@ -74,8 +74,10 @@ For fields whose deploy default is `null`, the deploy layer contributes no
 override. The effective value may still come from a platform section, an
 explicit CLI or stage override, or the downstream vLLM engine default.
 
-Note: for diffusion path, `distributed_executor_backend` currently defaults to
-`mp`, and `ray` / `external_launcher` are not fully supported yet.
+Note: for the diffusion path, an omitted `distributed_executor_backend` selects
+`uni` on a single GPU (in-process worker, no MessageQueue / `/dev/shm` output
+segments) and `mp` when `num_gpus > 1`. Set `mp` explicitly to keep a worker
+subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet.
 
 ### Stage fields
 

tests/diffusion/test_uniproc_executor.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Single-GPU (uniproc) diffusion executor: selection and RPC behaviour."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from vllm.v1.engine.exceptions import EngineDeadError
+
+from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
+from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeODConfig:
+    def __init__(self, num_gpus: int = 1, backend: str | None = None) -> None:
+        self.num_gpus = num_gpus
+        self.distributed_executor_backend = backend
+        self.worker_extension_cls = None
+        self.custom_pipeline_args = None
+
+
+def _od_config(num_gpus: int = 1, backend: str | None = None) -> _FakeODConfig:
+    return _FakeODConfig(num_gpus=num_gpus, backend=backend)
+
+
+def test_single_gpu_defaults_to_uniproc():
+    assert DiffusionExecutor.get_class(_od_config(num_gpus=1)) is UniProcDiffusionExecutor
+
+
+def test_explicit_mp_is_honored_on_single_gpu():
+    assert DiffusionExecutor.get_class(_od_config(num_gpus=1, backend="mp")) is MultiprocDiffusionExecutor
+
+
+def test_multi_gpu_defaults_to_multiproc():
+    assert DiffusionExecutor.get_class(_od_config(num_gpus=2)) is MultiprocDiffusionExecutor
+
+
+def test_explicit_uni_backend_is_honored():
+    assert DiffusionExecutor.get_class(_od_config(num_gpus=1, backend="uni")) is UniProcDiffusionExecutor
+
+
+@pytest.fixture
+def executor(monkeypatch: pytest.MonkeyPatch):
+    """A ``UniProcDiffusionExecutor`` with a mocked worker (no model load)."""
+    worker = MagicMock()
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls",
+        MagicMock(return_value="vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker"),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.executor.uniproc_executor.resolve_obj_by_qualname",
+        MagicMock(return_value=object),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase",
+        MagicMock(return_value=worker),
+    )
+    return UniProcDiffusionExecutor(_od_config(num_gpus=1)), worker
+
+
+def test_rejects_multi_gpu_config(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls",
+        MagicMock(return_value="unused"),
+    )
+    with pytest.raises(ValueError, match="single GPU only"):
+        UniProcDiffusionExecutor(_od_config(num_gpus=4))
+
+
+def test_shutdown_is_safe_on_partially_constructed_executor(monkeypatch: pytest.MonkeyPatch):
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.get_diffusion_worker_cls",
+        MagicMock(return_value="unused"),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.executor.uniproc_executor.resolve_obj_by_qualname",
+        MagicMock(return_value=object),
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase",
+        MagicMock(side_effect=RuntimeError("model load blew up")),
+    )
+    constructed = UniProcDiffusionExecutor.__new__(UniProcDiffusionExecutor)
+    with pytest.raises(RuntimeError, match="model load blew up"):
+        constructed.__init__(_od_config(num_gpus=1))
+
+    constructed.shutdown()
+
+
+def test_collective_rpc_calls_worker_directly(executor):
+    exec_, worker = executor
+    worker.execute_method.return_value = "result"
+
+    out = exec_.collective_rpc("some_method", args=(1, 2), kwargs={"k": "v"}, unique_reply_rank=0)
+
+    assert out == "result"
+    worker.execute_method.assert_called_once_with("some_method", 1, 2, k="v")
+
+
+def test_collective_rpc_returns_list_when_no_reply_rank(executor):
+    exec_, worker = executor
+    worker.execute_method.return_value = "result"
+
+    assert exec_.collective_rpc("some_method") == ["result"]
+
+
+def test_collective_rpc_propagates_worker_exceptions(executor):
+    exec_, worker = executor
+    worker.execute_method.side_effect = RuntimeError("boom")
+
+    with pytest.raises(RuntimeError, match="boom"):
+        exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+
+def test_recoverable_worker_failure_keeps_the_executor_alive(executor, monkeypatch):
+    exec_, worker = executor
+    worker.execute_method.side_effect = RuntimeError("CUDA out of memory")
+    monkeypatch.setattr(exec_, "_device_is_usable", lambda: True)
+    died = MagicMock()
+    exec_.register_failure_callback(died)
+
+    with pytest.raises(RuntimeError):
+        exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+    assert exec_.is_dead is False
+    died.assert_not_called()
+    exec_.check_health()
+
+
+def test_poisoned_cuda_context_marks_the_executor_dead(executor, monkeypatch):
+    exec_, worker = executor
+    worker.execute_method.side_effect = RuntimeError("an illegal memory access was encountered")
+    monkeypatch.setattr(exec_, "_device_is_usable", lambda: False)
+    died = MagicMock()
+    exec_.register_failure_callback(died)
+
+    with pytest.raises(RuntimeError):
+        exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+    assert exec_.is_dead is True
+    died.assert_called_once_with()
+    with pytest.raises(EngineDeadError):
+        exec_.check_health()
+
+
+def test_failure_is_latched_and_callbacks_fire_once(executor, monkeypatch):
+    exec_, worker = executor
+    worker.execute_method.side_effect = RuntimeError("boom")
+    monkeypatch.setattr(exec_, "_device_is_usable", lambda: False)
+    died = MagicMock()
+    exec_.register_failure_callback(died)
+
+    for _ in range(3):
+        with pytest.raises(RuntimeError):
+            exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+    died.assert_called_once_with()
+
+
+def test_device_probe_is_skipped_when_cuda_was_never_initialized(executor):
+    exec_, worker = executor
+    worker.execute_method.side_effect = RuntimeError("boom")
+
+    with pytest.raises(RuntimeError):
+        exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+    if not torch.cuda.is_initialized():
+        assert exec_.is_dead is False
+
+
+def test_device_probe_skipped_when_accelerator_unavailable(executor, monkeypatch):
+    exec_, worker = executor
+    worker.execute_method.side_effect = RuntimeError("boom")
+    monkeypatch.setattr(torch.accelerator, "is_available", lambda: False)
+    sync = MagicMock()
+    monkeypatch.setattr(torch.accelerator, "synchronize", sync)
+
+    with pytest.raises(RuntimeError, match="boom"):
+        exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+    sync.assert_not_called()
+    assert exec_.is_dead is False
+
+
+def test_device_probe_skipped_when_device_not_initialized(executor, monkeypatch):
+    exec_, worker = executor
+    worker.execute_method.side_effect = RuntimeError("boom")
+    monkeypatch.setattr(torch.accelerator, "is_available", lambda: True)
+    monkeypatch.setattr(torch.accelerator, "current_accelerator", lambda: SimpleNamespace(type="npu"))
+    monkeypatch.setattr(torch, "npu", SimpleNamespace(is_initialized=lambda: False), raising=False)
+    sync = MagicMock()
+    monkeypatch.setattr(torch.accelerator, "synchronize", sync)
+
+    with pytest.raises(RuntimeError, match="boom"):
+        exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+    sync.assert_not_called()
+    assert exec_.is_dead is False
+
+
+def test_device_probe_latches_sticky_fault_on_npu(executor, monkeypatch):
+    """NPU must not short-circuit the probe the way a CUDA-only guard did."""
+    exec_, worker = executor
+    worker.execute_method.side_effect = RuntimeError("boom")
+    monkeypatch.setattr(torch.accelerator, "is_available", lambda: True)
+    monkeypatch.setattr(torch.accelerator, "current_accelerator", lambda: SimpleNamespace(type="npu"))
+    monkeypatch.setattr(torch, "npu", SimpleNamespace(is_initialized=lambda: True), raising=False)
+    monkeypatch.setattr(
+        torch.accelerator,
+        "synchronize",
+        MagicMock(side_effect=RuntimeError("NPU context poisoned")),
+    )
+    died = MagicMock()
+    exec_.register_failure_callback(died)
+
+    with pytest.raises(RuntimeError, match="boom"):
+        exec_.collective_rpc("some_method", unique_reply_rank=0)
+
+    assert exec_.is_dead is True
+    died.assert_called_once_with()
+
+
+def test_check_health_ok_then_dead(executor):
+    exec_, _ = executor
+    exec_.check_health()
+
+    exec_._is_failed = True
+    with pytest.raises(EngineDeadError):
+        exec_.check_health()
+
+
+def test_shutdown_is_idempotent_and_closes_executor(executor, monkeypatch):
+    exec_, worker = executor
+    collect = MagicMock()
+    empty_cache = MagicMock()
+    monkeypatch.setattr("vllm_omni.diffusion.executor.uniproc_executor.gc.collect", collect)
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.is_available",
+        lambda: True,
+    )
+    monkeypatch.setattr(
+        "vllm_omni.diffusion.executor.uniproc_executor.current_omni_platform.empty_cache",
+        empty_cache,
+    )
+    exec_.register_failure_callback(MagicMock())
+
+    exec_.shutdown()
+    exec_.shutdown()
+
+    worker.shutdown.assert_called_once()
+    assert exec_.driver_worker is None
+    assert exec_._failure_callbacks == []
+    collect.assert_called_once_with()
+    empty_cache.assert_called_once_with()
+    with pytest.raises(RuntimeError, match="closed"):
+        exec_.collective_rpc("some_method", unique_reply_rank=0)

tests/e2e/features/custom_pipeline/test_async_omni_collective_rpc.py

@@ -166,8 +166,8 @@ async def test_sleep_memory_reclaimed_custom_pipeline():
 
     The fix moves custom_pipeline init outside the CUDA context so all weights
     go through the caching allocator and are therefore fully reclaimed by
-    sleep(level=1).  A non-zero ``CuMemAllocator.get_current_usage()`` after
-    sleep is the direct signal that the bypass is still occurring.
+    sleep(level=1). The tracked allocations must transition to the allocator's
+    asleep state and report physical reclamation.
     """
     with ExitStack() as after:
         engine = AsyncOmni(
@@ -187,62 +187,72 @@ async def test_sleep_memory_reclaimed_custom_pipeline():
         free_before, total = torch.cuda.mem_get_info()
         used_before_gib = (total - free_before) / 1024**3
 
-        # Measure CuMemAllocator-tracked usage before sleep.  In inline mode
-        # the worker runs in a thread pool inside this process, so the allocator
-        # singleton is shared and can be read directly.
+        # Measure CuMemAllocator-tracked usage before sleep.  In inline / uni
+        # mode the worker runs in this process, so the allocator singleton is
+        # shared and can be read directly.
         allocator = None
         tracked_before = 0
+        tracked_ptrs: set[int] = set()
         try:
             from vllm.device_allocator.cumem import CuMemAllocator
 
             allocator = CuMemAllocator.get_instance()
-            tracked_before = allocator.get_current_usage()
+            tracked_ptrs = {ptr for ptr, data in allocator.pointer_to_data.items() if not data.is_asleep}
+            tracked_before = sum(allocator.pointer_to_data[ptr].handle[1] for ptr in tracked_ptrs)
         except Exception:
             pass
 
-        # Put the engine to sleep; all weights should be offloaded via the pool.
-        acks = await engine.sleep(level=1)
-        await asyncio.sleep(0.5)  # allow the CUDA driver to settle
-        torch.accelerator.synchronize()
+        slept = False
+        try:
+            # Put the engine to sleep; this engine's weights should be
+            # offloaded via the pool.
+            acks = await engine.sleep(level=1)
+            slept = True
+            await asyncio.sleep(0.5)  # allow the CUDA driver to settle
+            torch.accelerator.synchronize()
 
-        # Measure after sleep.
-        free_after, _ = torch.cuda.mem_get_info()
-        used_after_gib = (total - free_after) / 1024**3
-        drop_gib = used_before_gib - used_after_gib
+            # Measure after sleep.
+            free_after, _ = torch.cuda.mem_get_info()
+            used_after_gib = (total - free_after) / 1024**3
+            drop_gib = used_before_gib - used_after_gib
 
-        # --- Primary assertion: allocator reports zero tracked memory. ---
-        # If this fails it means weights were allocated outside the CuMem pool
-        # (safetensors direct-to-GPU bypass) — the exact regression this test
-        # is designed to catch.
-        if allocator is not None:
-            tracked_after = allocator.get_current_usage()
-            assert tracked_after == 0, (
-                f"CuMemAllocator still tracks {tracked_after / 1024**3:.3f} GiB "
-                f"after sleep(level=1) on custom_pipeline path "
-                f"(was {tracked_before / 1024**3:.3f} GiB before sleep). "
-                "Weights were allocated outside the CuMem pool via the "
-                "safetensors direct-to-GPU fast path — loader-context fix "
-                "may not be applied."
+            # get_current_usage() includes unmapped sleeping handles, so it
+            # does not decrease on sleep. Verify that allocations mapped
+            # before the call instead transitioned to the asleep state.
+            if allocator is not None:
+                assert tracked_before > 0, (
+                    "Expected custom_pipeline weights to use the CuMem pool, "
+                    "but no mapped allocations were tracked before sleep."
+                )
+                still_awake = [
+                    ptr
+                    for ptr in tracked_ptrs
+                    if ptr in allocator.pointer_to_data and not allocator.pointer_to_data[ptr].is_asleep
+                ]
+                assert not still_awake, (
+                    f"{len(still_awake)} CuMem allocation(s) remained mapped "
+                    "after sleep(level=1) on the custom_pipeline path."
+                )
+
+            # Physical VRAM or ACK freed_bytes confirms reclamation at the
+            # driver level.
+            total_freed_bytes = sum(
+                (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0))
+                for ack in acks
+                if ack is not None
+            )
+            freed_gib = total_freed_bytes / 1024**3
+            assert freed_gib > 0 or drop_gib > 0, (
+                f"Expected GPU memory to be reclaimed after sleep(level=1) on "
+                f"custom_pipeline + enable_sleep_mode=True. "
+                f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, "
+                f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB."
             )
 
-        # --- Secondary assertion: physical VRAM or ACK freed_bytes confirms
-        # reclamation at the driver level.
-        total_freed_bytes = sum(
-            (ack.freed_bytes if hasattr(ack, "freed_bytes") else ack.get("freed_bytes", 0))
-            for ack in acks
-            if ack is not None
-        )
-        freed_gib = total_freed_bytes / 1024**3
-        assert freed_gib > 0 or drop_gib > 0, (
-            f"Expected GPU memory to be reclaimed after sleep(level=1) on "
-            f"custom_pipeline + enable_sleep_mode=True. "
-            f"CuMemAllocator tracked before={tracked_before / 1024**3:.3f} GiB, "
-            f"ACK freed={freed_gib:.3f} GiB, global VRAM drop={drop_gib:.3f} GiB."
-        )
-
-        # Engine must report it is sleeping.
-        assert await engine.is_sleeping()
-
-        # Wake up and confirm the engine is functional again.
-        await engine.wake_up()
-        assert not await engine.is_sleeping()
+            assert await engine.is_sleeping()
+            await engine.wake_up()
+            slept = False
+            assert not await engine.is_sleeping()
+        finally:
+            if slept:
+                await engine.wake_up()

tests/e2e/features/custom_pipeline/test_async_omni_qwen_image_generate.py

@@ -17,9 +17,12 @@ from functools import lru_cache
 
 import numpy as np
 import pytest
+import torch.distributed as dist
 from transformers import AutoTokenizer
 
 from tests.helpers.mark import hardware_test
+from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor
+from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -131,6 +134,19 @@ async def _generate_once(
     return last_output
 
 
+def _assert_live_uniproc_worker(engine: AsyncOmni) -> None:
+    """Default single-GPU path must use the in-process executor, not mp IPC."""
+    clients = engine.engine.stage_clients
+    assert len(clients) == 1
+    client = clients[0]
+    assert isinstance(client, InlineStageDiffusionClient)
+    executor = client._engine.executor
+    assert isinstance(executor, UniProcDiffusionExecutor)
+    assert executor.driver_worker is not None
+    assert dist.is_initialized()
+    assert dist.get_world_size() == 1
+
+
 def _assert_valid_image_output(output: OmniRequestOutput) -> None:
     assert output.final_output_type == "image"
     assert output.images, "Expected at least one generated image"
@@ -158,6 +174,7 @@ async def test_async_omni_generate():
         )
         after.callback(engine.shutdown)
 
+        _assert_live_uniproc_worker(engine)
         output = await _generate_once(
             engine,
             "a beautiful sunset over the ocean with vibrant orange and purple clouds "

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -225,7 +225,7 @@ def test_default_stage_config_whitelist_none_fallback():
     engine_args = stage_cfg["engine_args"]
 
     assert engine_args["trust_remote_code"] is False
-    assert engine_args["distributed_executor_backend"] == "mp"
+    assert engine_args["distributed_executor_backend"] is None
     assert engine_args["dtype"] == "auto"
     assert engine_args["enforce_eager"] is False
 

vllm_omni/config/omni_config.py

@@ -663,7 +663,7 @@ class _DiffusionConfigProjection:
     dtype: Any = "auto"
     trust_remote_code: bool = False
     revision: str | None = None
-    distributed_executor_backend: str = "mp"
+    distributed_executor_backend: str | None = None
     dist_timeout: int | None = None
     nccl_port: int | None = None
     master_port: int | None = None

vllm_omni/diffusion/data.py

@@ -700,8 +700,10 @@ class OmniDiffusionConfig:
     # in both directions. See docs/features/session_state_manager.md.
     enable_session_state_manager: bool = False
 
-    # Distributed executor backend
-    distributed_executor_backend: str = "mp"
+    # Distributed executor backend. ``None`` means auto: ``"uni"`` at
+    # ``num_gpus == 1``, ``"mp"`` otherwise. Explicit ``"mp"`` keeps a
+    # worker subprocess (process isolation and RPC timeouts).
+    distributed_executor_backend: str | None = None
     nccl_port: int | None = None
 
     # Engine backend selection, resolved by ``DiffusionEngine.resolve_engine_class``

vllm_omni/diffusion/executor/abstract.py

@@ -24,10 +24,15 @@ class DiffusionExecutor(ABC):
     def get_class(od_config: OmniDiffusionConfig) -> type[DiffusionExecutor]:
         executor_class: type[DiffusionExecutor]
         distributed_executor_backend = od_config.distributed_executor_backend
-        # Keep backward-compatible behavior for callers/configs that omit this
-        # field and rely on the historical diffusion default backend.
+        # Mirror vLLM's `world_size == 1 -> "uni"` default
+        # (vllm/config/parallel.py). A single-GPU diffusion deployment has
+        # nothing to distribute: spawning a worker only adds MessageQueues,
+        # /dev/shm output segments, and a second model load. Explicit "mp"
+        # keeps process isolation and RPC timeouts. Multi-GPU still defaults
+        # to "mp".
         if distributed_executor_backend is None:
-            distributed_executor_backend = "mp"
+            num_gpus = od_config.num_gpus or 1
+            distributed_executor_backend = "uni" if num_gpus == 1 else "mp"
 
         if isinstance(distributed_executor_backend, type):
             if not issubclass(distributed_executor_backend, DiffusionExecutor):
@@ -42,6 +47,10 @@ class DiffusionExecutor(ABC):
             from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
 
             executor_class = MultiprocDiffusionExecutor
+        elif distributed_executor_backend == "uni":
+            from vllm_omni.diffusion.executor.uniproc_executor import UniProcDiffusionExecutor
+
+            executor_class = UniProcDiffusionExecutor
         elif distributed_executor_backend == "external_launcher":
             raise NotImplementedError("external_launcher backend is not yet supported.")
         elif isinstance(distributed_executor_backend, str):

vllm_omni/diffusion/executor/uniproc_executor.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""In-process diffusion executor for single-GPU deployments.
+
+Counterpart to :class:`MultiprocDiffusionExecutor` for ``num_gpus == 1``.
+The worker is constructed and driven in the engine process, so none of the
+multiproc machinery is created:
+
+* no spawned worker subprocess (and no second model load / CUDA context)
+* no ``MessageQueue`` shared-memory ring buffers or zmq ``ipc://`` sockets
+* no POSIX ``/dev/shm`` segments for output tensors (``ipc.py`` pack/unpack)
+
+Mirrors vLLM's ``UniProcExecutor`` (``vllm/v1/executor/uniproc_executor.py``).
+
+``timeout`` on :meth:`collective_rpc` is accepted for interface parity but is
+not enforced: the call runs on the calling thread, so a hung worker blocks it.
+The multiproc executor enforces deadlines via its result-queue dequeue.
+"""
+
+from __future__ import annotations
+
+import gc
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+import torch
+from vllm.logger import init_logger
+from vllm.utils.import_utils import resolve_obj_by_qualname
+from vllm.v1.engine.exceptions import EngineDeadError
+
+from vllm_omni.diffusion.data import DiffusionOutput
+from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+from vllm_omni.platforms import current_omni_platform
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
+    from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
+
+logger = init_logger(__name__)
+
+
+class UniProcDiffusionExecutor(DiffusionExecutor):
+    """Runs a single diffusion worker inline, with no IPC of any kind."""
+
+    def _init_executor(self) -> None:
+        from vllm_omni.diffusion.worker.diffusion_worker import WorkerWrapperBase
+
+        self._closed = False
+        self._is_failed = False
+        self.driver_worker = None
+        self._failure_callbacks: list[Callable[[], None]] = []
+        self._warned_about_timeout = False
+
+        num_gpus = self.od_config.num_gpus or 1
+        if num_gpus != 1:
+            raise ValueError(
+                f"UniProcDiffusionExecutor supports a single GPU only, got num_gpus={num_gpus}. "
+                'Use distributed_executor_backend="mp" for multi-GPU deployments.'
+            )
+
+        # The worker still initializes a one-rank distributed environment so
+        # that group lookups (`get_dp_group()` and friends) resolve. At world
+        # size 1, vLLM's GroupCoordinator skips the device communicator and
+        # the shared-memory MessageQueue, and every collective short-circuits.
+        worker_cls_path = current_omni_platform.get_diffusion_worker_cls()
+        self.driver_worker = WorkerWrapperBase(
+            gpu_id=0,
+            od_config=self.od_config,
+            worker_extension_cls=self.od_config.worker_extension_cls,
+            custom_pipeline_args=getattr(self.od_config, "custom_pipeline_args", None),
+            base_worker_class=resolve_obj_by_qualname(worker_cls_path),
+        )
+        logger.info("Diffusion worker initialized in-process (uniproc executor)")
+
+    @property
+    def is_dead(self) -> bool:
+        return self._closed or self._is_failed
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise RuntimeError("DiffusionExecutor is closed.")
+        if self.driver_worker is None:
+            raise RuntimeError("Diffusion worker is not initialized.")
+
+    def collective_rpc(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple = (),
+        kwargs: dict | None = None,
+        unique_reply_rank: int | None = None,
+        exec_all_ranks: bool = False,
+    ) -> Any:
+        """Invoke ``method`` on the single in-process worker.
+
+        ``unique_reply_rank`` / ``exec_all_ranks`` only select which ranks
+        participate, which is degenerate at world size 1. The return shape
+        still matches :class:`MultiprocDiffusionExecutor` (a bare result when a
+        reply rank is named, a one-element list otherwise).
+        """
+        self._ensure_open()
+        if timeout is not None and not self._warned_about_timeout:
+            self._warned_about_timeout = True
+            logger.info(
+                "Diffusion RPC deadlines are not enforced by the in-process executor: "
+                "the worker runs on the calling thread, so a hung call blocks it."
+            )
+        try:
+            result = self.driver_worker.execute_method(method, *(args or ()), **(kwargs or {}))
+        except Exception:
+            if not self._device_is_usable():
+                self._mark_failed()
+            raise
+        return result if unique_reply_rank is not None else [result]
+
+    def execute_request(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
+        from vllm_omni.diffusion.sched.interface import validate_new_request_data_identity
+        from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
+
+        self._ensure_open()
+        runner_outputs: list[RunnerOutput] = []
+        for new_req in scheduler_output.scheduled_new_reqs:
+            validate_new_request_data_identity(new_req)
+            try:
+                args: tuple = (new_req.req, self.od_config, scheduler_output.kv_prefetch_job)
+                if new_req.diffusion_kv_metadata is not None:
+                    args += (new_req.diffusion_kv_metadata,)
+                result = self.collective_rpc(
+                    "execute_model",
+                    args=args,
+                    unique_reply_rank=0,
+                    exec_all_ranks=True,
+                )
+                if not isinstance(result, DiffusionOutput):
+                    raise RuntimeError(f"Unexpected response type: {type(result)!r}")
+                runner_outputs.append(
+                    RunnerOutput(
+                        request_id=new_req.request_id,
+                        step_index=None,
+                        finished=True,
+                        result=result,
+                    )
+                )
+            except Exception as exc:
+                runner_outputs.append(
+                    RunnerOutput(
+                        request_id=new_req.request_id,
+                        step_index=None,
+                        finished=True,
+                        result=DiffusionOutput(error=str(exc)),
+                    )
+                )
+        return BatchRunnerOutput.from_list(runner_outputs)
+
+    def execute_batch(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
+        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),
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        if not isinstance(result, BatchRunnerOutput):
+            raise RuntimeError(f"Unexpected response type for execute_batch: {type(result)!r}")
+        return result
+
+    def execute_step(self, scheduler_output: DiffusionSchedulerOutput) -> BaseRunnerOutput:
+        from vllm_omni.diffusion.worker.utils import BaseRunnerOutput
+
+        self._ensure_open()
+        result = self.collective_rpc(
+            "execute_stepwise",
+            args=(scheduler_output,),
+            unique_reply_rank=0,
+            exec_all_ranks=True,
+        )
+        if isinstance(result, BaseRunnerOutput):
+            return result
+        raise RuntimeError(f"Unexpected response type for execute_step: {type(result)!r}")
+
+    def _device_is_usable(self) -> bool:
+        """Whether the accelerator context survived the failure we just caught.
+
+        Running the worker inline means a worker fault no longer kills a
+        process, so nothing else notices it. Recoverable faults (OOM, an
+        invalid request) leave the context intact and must not take the
+        deployment down. A sticky fault (illegal memory access, ECC error)
+        poisons the context: every later launch fails, so the engine needs
+        to be restarted. ``synchronize()`` re-raises a pending sticky error
+        on the registered accelerator (CUDA, ROCm, NPU).
+        """
+        if not torch.accelerator.is_available():
+            return True
+        accelerator = torch.accelerator.current_accelerator()
+        device_mod = getattr(torch, accelerator.type, None) if accelerator is not None else None
+        is_initialized = getattr(device_mod, "is_initialized", None)
+        if callable(is_initialized) and not is_initialized():
+            return True
+        try:
+            torch.accelerator.synchronize()
+        except Exception as exc:
+            logger.error("Device context is unusable after a worker failure: %s", exc)
+            return False
+        return True
+
+    def _mark_failed(self) -> None:
+        if self._is_failed:
+            return
+        self._is_failed = True
+        logger.error("Diffusion worker failed unrecoverably in-process; marking the executor dead.")
+        for callback in self._failure_callbacks:
+            try:
+                callback()
+            except Exception:
+                logger.exception("Diffusion executor failure callback raised")
+
+    def register_failure_callback(self, callback: Callable[[], None]) -> None:
+        """Register a callback invoked when the inline worker fails fatally.
+
+        The multiproc executor fires these from its process monitor. There is
+        no process to monitor here, so they are fired from ``collective_rpc``
+        instead — without them a poisoned device context would leave
+        ``check_health`` reporting healthy while every request fails.
+        """
+        self._failure_callbacks.append(callback)
+
+    def check_health(self) -> None:
+        if self._is_failed:
+            raise EngineDeadError()
+        self._ensure_open()
+
+    def shutdown(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        worker = self.driver_worker
+        self.driver_worker = None
+        self._failure_callbacks.clear()
+        if worker is not None:
+            try:
+                worker.shutdown()
+            except Exception as exc:
+                logger.warning("Diffusion worker shutdown encountered an error: %s", exc)
+        # Unlike the multiprocess executor, tearing down an inline worker does
+        # not exit its process. Drop the final model reference and return cached
+        # allocations to the driver so a subsequent engine can use the device.
+        del worker
+        gc.collect()
+        try:
+            if current_omni_platform.is_available():
+                current_omni_platform.empty_cache()
+        except Exception as exc:
+            logger.warning("Failed to release accelerator cache during diffusion worker shutdown: %s", exc)

vllm_omni/engine/async_omni_engine.py

@@ -1065,9 +1065,7 @@ class AsyncOmniEngine:
             "custom_pipeline_args": kwargs.get("custom_pipeline_args", None),
             "worker_extension_cls": kwargs.get("worker_extension_cls", None),
             "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]),
-            "distributed_executor_backend": (
-                "mp" if kwargs.get("distributed_executor_backend") is None else kwargs["distributed_executor_backend"]
-            ),
+            "distributed_executor_backend": kwargs.get("distributed_executor_backend"),
             "enable_sleep_mode": kwargs.get("enable_sleep_mode", False),
             "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False),
             "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32),