[Rebase] Rebase to vllm releases/v0.22.0 (#3891)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Signed-off-by: ZhengWG <zwg0606@gmail.com>
Signed-off-by: david6666666 <530634352@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Signed-off-by: nainiu258 <cperfect02@163.com>
Signed-off-by: nainiu258 <101917677+nainiu258@users.noreply.github.com>
Signed-off-by: JasonJ2021 <jasonj@zju.edu.cn>
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: skf <54565339+skf-1999@users.noreply.github.com>
Co-authored-by: Zheng Wengang <zwg0606@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: nainiu258 <101917677+nainiu258@users.noreply.github.com>
Co-authored-by: Jason <72191212+JasonJ2021@users.noreply.github.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

.buildkite/test-merge.yml

@@ -107,7 +107,7 @@ steps:
     depends_on: upload-merge-pipeline
     steps:
       - label: "Entrypoint Test with H100"
-        timeout_in_minutes: 35
+        timeout_in_minutes: 60
         commands:
           - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
         agents:

.buildkite/test-ready.yml

@@ -183,7 +183,7 @@ steps:
     depends_on: upload-ready-pipeline
     commands:
       - |
-        timeout 15m bash -c "
+        timeout 20m bash -c "
                   pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model'
         "
     agents:

benchmarks/tts/model_configs.yaml

@@ -14,6 +14,7 @@ models:
     supported_tasks: [default_voice, voice_design]
     backend: openai-audio-speech
     endpoint: /v1/audio/speech
+    trust_remote_code: true
     task_extra_body:
       default_voice:
         voice: Vivian
@@ -27,6 +28,7 @@ models:
     supported_tasks: [voice_clone]
     backend: openai-audio-speech
     endpoint: /v1/audio/speech
+    trust_remote_code: true
     task_extra_body:
       voice_clone:
         task_type: Base
@@ -36,6 +38,7 @@ models:
     trust_remote_code: true
     backend: openai-audio-speech
     endpoint: /v1/audio/speech
+    trust_remote_code: true
     task_extra_body:
       voice_clone: {}
       default_voice:

docker/Dockerfile.ci

@@ -1,5 +1,8 @@
+# Target vllm version: v0.21.0 (releases/v0.21.0 branch).
+# The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
+# be updated on every rebase to match the new upstream vllm commit/tag.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.21.0
+ARG VLLM_BASE_TAG=v0.22.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docs/user_guide/quantization/gguf.md

@@ -74,7 +74,7 @@ Online:
 vllm serve Qwen/Qwen-Image \
   --omni \
   --port 8000 \
-  --quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
+  --diffusion-quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
 ```
 
 ## Parameters

recipes/README.md

@@ -47,6 +47,7 @@ recipes/
 | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
 | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
 | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
+|
 
 Within a single recipe file, include different hardware support sections such
 as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like

requirements/common.txt

@@ -1,4 +1,10 @@
 # Common dependencies for all platforms
+# transformers: keep aligned with upstream vLLM constraint.
+# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
+# 5.9.0 breaks models that define _keys_to_ignore_on_load_unexpected as a list
+# (e.g. OpenMOSS-Team/MOSS-TTS-Nano) because it changed
+# _adjust_missing_and_unexpected_keys from list+list to set|set.
+transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0, < 5.9.0
 av>=14.0.0
 omegaconf>=2.3.0
 diffusers==0.38.0

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -223,8 +223,12 @@ def _get_open_port() -> int:
         return s.getsockname()[1]
 
 
-def _wait_for_port(host: str, port: int, timeout: int = 1200) -> None:
-    """Block until the given host:port accepts connections or timeout expires."""
+def _wait_for_port(host: str, port: int, timeout: int = 1200, proc: subprocess.Popen | None = None) -> None:
+    """Block until the given host:port accepts connections or timeout expires.
+
+    If *proc* is provided, also monitors the process; raises RuntimeError
+    immediately if the server process exits before the port becomes available.
+    """
     start = time.time()
     while time.time() - start < timeout:
         try:
@@ -234,6 +238,10 @@ def _wait_for_port(host: str, port: int, timeout: int = 1200) -> None:
                     return
         except Exception:
             pass
+        if proc is not None:
+            ret = proc.poll()
+            if ret is not None:
+                raise RuntimeError(f"Server process exited with code {ret} before port {host}:{port} became ready")
         time.sleep(2)
     raise RuntimeError(f"Server did not start on {host}:{port} within {timeout}s")
 
@@ -331,7 +339,7 @@ class DiffusionServer:
             env=env,
             cwd=str(Path(__file__).parent.parent.parent.parent),
         )
-        _wait_for_port(self.host, self.port)
+        _wait_for_port(self.host, self.port, proc=self.proc)
         print(f"DiffusionServer ready on {self.host}:{self.port}")
 
     def __enter__(self):

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

@@ -5,7 +5,8 @@
             "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
             "use_omni": false,
             "extra_cli_args": [
-                "--no-enable-prefix-caching"
+                "--no-enable-prefix-caching",
+                "--skip-mm-profiling"
             ]
         },
         "benchmark_params": [

tests/diffusion/offloader/test_diffusion_cpu_offload.py

@@ -17,7 +17,7 @@ AUDIO_MODEL = {
 }
 
 IMAGE_VIDEO_MODELS = {
-    "riverclouds/qwen_image_random": {"cuda": 2500, "rocm": 2100},
+    "riverclouds/qwen_image_random": {"cuda": 2200, "rocm": 2100},
 }
 
 MODELS = {**AUDIO_MODEL, **IMAGE_VIDEO_MODELS}

tests/diffusion/test_diffusion_model_runner.py

@@ -301,8 +301,4 @@ def test_vllm_set_forward_context_implementation(monkeypatch):
                 "vllm_config",
             ),
         ),
-        ("ir_op_priority", None),
-        ("enable_torch_wrap", vllm_config.compilation_config.ir_enable_torch_wrap),
-        ("enable_torch_wrap_exit", vllm_config.compilation_config.ir_enable_torch_wrap),
-        ("ir_op_priority_exit", None),
     ], ERROR_MESSAGE

tests/diffusion/test_stage_diffusion_proc.py

@@ -85,7 +85,7 @@ class MockOmniRequestOutput:
 BASE_HEIGHT = 512
 BASE_WIDTH = 512
 BASE_INFER_STEPS = 10
-DELAY_BASE = 2
+DELAY_BASE = 0.01
 
 
 class MockDiffusionEngine:
@@ -93,7 +93,7 @@ class MockDiffusionEngine:
         def simulate_step_delay(height, width, num_inference_steps) -> float:
             return (height / BASE_HEIGHT) * (width / BASE_WIDTH) * (num_inference_steps / BASE_INFER_STEPS)
 
-        DELAY_BASE = 2
+        DELAY_BASE = 0.01
         delay_scale = simulate_step_delay(
             request.sampling_params.height, request.sampling_params.width, request.sampling_params.num_inference_steps
         )
@@ -139,7 +139,7 @@ async def test_proc_process_request_with_batching_async_output():
     assert len(results) == len(test_requests)
     base_time = DELAY_BASE
     time_gap_std = DELAY_BASE * 2 * 2 * 1  # height/width/steps infer time scale
-    eps = 0.5
+    eps = 0.1
     for i, (res, elapsed_time) in enumerate(results):
         assert res.request_id == test_requests[i]["request_id"]
         assert isinstance(res, MockOmniRequestOutput)

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -703,6 +703,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
         kv_connector_output=None,
         cudagraph_stats=None,
         req_id_to_index={"req-s1": 0},
+        routed_experts=None,
         routed_experts_dict=None,
     )
 
@@ -789,6 +790,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
         cudagraph_stats=None,
         req_id_to_index={"req-ar": 0},
         kv_extracted_req_ids=None,
+        routed_experts=None,
         routed_experts_dict=None,
     )
 

tests/distributed/omni_coordinator/test_omni_coordinator.py

@@ -68,13 +68,13 @@ def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates():
 
     sub_ctx = zmq.Context.instance()
     sub = sub_ctx.socket(zmq.SUB)
-    sub.connect(pub_addr)
+    sub.connect(coordinator.pub_zmq_addr)
     sub.setsockopt(zmq.SUBSCRIBE, b"")
 
     time.sleep(0.3)  # PUB/SUB slow-joiner
 
     client = OmniCoordClientForStage(
-        router_addr,
+        coordinator.router_zmq_addr,
         "tcp://stage:coalesce",
         "tcp://stage:coalesce-out",
         0,
@@ -131,7 +131,7 @@ def test_omni_coordinator_registration_broadcast():
 
     sub_ctx = zmq.Context.instance()
     sub = sub_ctx.socket(zmq.SUB)
-    sub.connect(pub_addr)
+    sub.connect(coordinator.pub_zmq_addr)
     sub.setsockopt(zmq.SUBSCRIBE, b"")
 
     # ZMQ PUB/SUB slow-joiner: allow SUB to connect before clients register.
@@ -139,9 +139,9 @@ def test_omni_coordinator_registration_broadcast():
 
     # Create 3 stage clients; each auto-registers on init.
     clients = [
-        OmniCoordClientForStage(router_addr, "tcp://stage:10001", "tcp://stage:10001-out", 0),
-        OmniCoordClientForStage(router_addr, "tcp://stage:10002", "tcp://stage:10002-out", 0),
-        OmniCoordClientForStage(router_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1),
+        OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:10001", "tcp://stage:10001-out", 0),
+        OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:10002", "tcp://stage:10002-out", 0),
+        OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:10003", "tcp://stage:10003-out", 1),
     ]
 
     msg = _wait_for_replica_list(sub, expected_count=3)
@@ -183,19 +183,19 @@ def test_omni_coordinator_heartbeat_timeout_handling():
 
     sub_ctx = zmq.Context.instance()
     sub = sub_ctx.socket(zmq.SUB)
-    sub.connect(pub_addr)
+    sub.connect(coordinator.pub_zmq_addr)
     sub.setsockopt(zmq.SUBSCRIBE, b"")
 
     time.sleep(0.3)
 
     # A and B: real clients that send heartbeats every 5s.
-    client_a = OmniCoordClientForStage(router_addr, "tcp://stage:a", "tcp://stage:a-out", 0)
-    client_b = OmniCoordClientForStage(router_addr, "tcp://stage:b", "tcp://stage:b-out", 0)
+    client_a = OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:a", "tcp://stage:a-out", 0)
+    client_b = OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:b", "tcp://stage:b-out", 0)
 
     # C: raw DEALER that sends only registration, no heartbeat.
     dealer_ctx = zmq.Context.instance()
     dealer_c = dealer_ctx.socket(zmq.DEALER)
-    dealer_c.connect(router_addr)
+    dealer_c.connect(coordinator.router_zmq_addr)
     reg_event = {
         "input_addr": "tcp://stage:c",
         "output_addr": "tcp://stage:c-out",
@@ -254,12 +254,12 @@ def test_omni_coordinator_replica_shutdown_handling():
 
     sub_ctx = zmq.Context.instance()
     sub = sub_ctx.socket(zmq.SUB)
-    sub.connect(pub_addr)
+    sub.connect(coordinator.pub_zmq_addr)
     sub.setsockopt(zmq.SUBSCRIBE, b"")
 
     time.sleep(0.3)  # PUB/SUB slow-joiner
 
-    client = OmniCoordClientForStage(router_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0)
+    client = OmniCoordClientForStage(coordinator.router_zmq_addr, "tcp://stage:shutdown", "tcp://stage:shutdown-out", 0)
 
     msg = _wait_for_replica_list(sub, expected_count=1)
     assert msg is not None

tests/e2e/online_serving/test_qwen_image_expansion.py

@@ -97,7 +97,7 @@ def _get_diffusion_feature_cases(model: str):
                     "--vae-patch-parallel-size",
                     "2",
                     "--vae-use-tiling",
-                    "--quantization-config",
+                    "--diffusion-quantization-config",
                     '{"method":"fp8"}',
                 ],
             ),

tests/entrypoints/test_stage_utils.py

@@ -4,7 +4,7 @@ import sys
 import pytest
 from pytest_mock import MockerFixture
 
-from vllm_omni.entrypoints.stage_utils import set_stage_devices
+from vllm_omni.entrypoints.stage_utils import _map_device_list, set_stage_devices
 
 
 def _make_dummy_torch(call_log):
@@ -130,3 +130,90 @@ def test_set_stage_devices_npu_platform(mocker: MockerFixture, monkeypatch: pyte
     set_stage_devices(stage_id=0, devices="0,1")
 
     assert os.environ["ASCEND_RT_VISIBLE_DEVICES"] == "4,5"
+
+
+# ---- _map_device_list unit tests ----
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_idempotency():
+    """Device IDs already in the visible set are returned as-is (idempotency)."""
+    result = _map_device_list(0, ["0", "1"], ["0", "1", "2", "3"])
+    assert result == ["0", "1"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_idempotency_single():
+    """Single device ID in visible set passes through."""
+    result = _map_device_list(1, ["1"], ["0", "1"])
+    assert result == ["1"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_index_mapping():
+    """Device IDs < num_visible not in set are treated as indices."""
+    result = _map_device_list(0, ["0", "1"], ["6", "7", "8", "9"])
+    assert result == ["6", "7"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_physical_fallback():
+    """Device IDs >= num_visible are treated as physical IDs and passed through."""
+    result = _map_device_list(1, ["2"], ["0", "1"])
+    assert result == ["2"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_physical_fallback_multi():
+    """Multiple device IDs >= num_visible all pass through as physical IDs."""
+    result = _map_device_list(1, ["2", "3"], ["0", "1"])
+    assert result == ["2", "3"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_physical_fallback_mixed():
+    """When some devices are >= num_visible but some < num_visible, index mapping applies for the subset."""
+    # "1" < num_visible(2) so index mapping applies: visible[1] = "5"
+    # "2" >= num_visible(2) so it's dropped via the partial mapping path
+    result = _map_device_list(1, ["1", "2"], ["0", "5"])
+    assert result == ["5"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_index_with_gaps():
+    """Index-based mapping works with non-contiguous visible device lists."""
+    result = _map_device_list(0, ["0", "1"], ["4", "5", "7"])
+    assert result == ["4", "5"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_index_mapping_no_idempotency():
+    """Device IDs < num_visible not in visible set map via index even when no ID matches literally."""
+    # "0","1" with visible ["2","3"]: idempotency check fails (none in set),
+    # but index mapping succeeds: visible[0]="2", visible[1]="3"
+    result = _map_device_list(0, ["0", "1"], ["2", "3"])
+    assert result == ["2", "3"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_partial_mapping():
+    """When only a subset can map, returns the mapped subset (no error)."""
+    result = _map_device_list(0, ["0", "1", "2"], ["5", "6"])
+    assert result == ["5", "6"]
+
+
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_raises_on_non_numeric():
+    """Non-numeric device IDs raise ValueError."""
+    with pytest.raises(ValueError, match="must be non-negative integers"):
+        _map_device_list(0, ["a"], ["0", "1"])

tests/helpers/env.py

@@ -160,9 +160,41 @@ def _print_gpu_processes() -> None:
     print("=" * 80)
 
 
+def _cleanup_stale_device_locks() -> None:
+    """Remove stale device-initialization lock files whose recorded PID is dead.
+
+    Lock files at ``/tmp/vllm_omni_device_*_init.lock`` may persist after a
+    crashed / killed test run and block subsequent orchestrator startups.
+    """
+    import glob as _glob
+
+    for lock_file in _glob.glob("/tmp/vllm_omni_device_*_init.lock"):
+        try:
+            with open(lock_file) as fh:
+                content = fh.read().strip()
+            if not content:
+                continue
+            pid = int(content)
+        except (OSError, ValueError):
+            continue
+
+        try:
+            os.kill(pid, 0)
+        except ProcessLookupError:
+            print(f"  Removing stale device lock {lock_file} (PID {pid} is dead)")
+            try:
+                os.unlink(lock_file)
+            except OSError:
+                pass
+        except PermissionError:
+            pass
+
+
 def run_pre_test_cleanup() -> None:
     print("Pre-test GPU status:")
 
+    _cleanup_stale_device_locks()
+
     num_gpus = current_omni_platform.device_count()
     if num_gpus > 0:
         try:

tests/helpers/mark.py

@@ -21,7 +21,7 @@ def cuda_marks(*, res: str, num_cards: int):
         return marks
     test_distributed = pytest.mark.distributed_cuda(num_cards=num_cards)
 
-    test_skipif = pytest.mark.skipif_cuda(
+    test_skipif = pytest.mark.skipif(
         not current_platform.is_cuda() or (current_platform.device_count() < num_cards),
         reason=f"Need at least {num_cards} CUDA GPUs to run the test.",
     )
@@ -52,7 +52,7 @@ def xpu_marks(*, res: str, num_cards: int):
         return marks
     test_distributed = pytest.mark.distributed_xpu(num_cards=num_cards)
 
-    test_skipif = pytest.mark.skipif_xpu(
+    test_skipif = pytest.mark.skipif(
         not current_platform.is_xpu() or (current_platform.device_count() < num_cards),
         reason=f"Need at least {num_cards} XPUs to run the test.",
     )

tests/helpers/runtime.py

@@ -305,40 +305,132 @@ class OmniServer:
             time.sleep(2)
         raise RuntimeError(f"Server failed to start within {max_wait} seconds")
 
+    @staticmethod
+    def _reap_zombie(proc: "psutil.Process") -> bool:
+        """Reap a zombie child process via ``os.waitpid``.
+
+        ``psutil.Process.wait()`` uses ``pidfd_open`` + ``poll()``, which
+        never fires for a zombie (the zombie is already dead and will not
+        change state).  Since the test process is the parent, we can reap
+        the zombie directly with ``os.waitpid(pid, os.WNOHANG)`` and
+        retrieve its exit code.
+
+        Returns True if the process was a zombie and was reaped.
+        """
+        if proc.status() != psutil.STATUS_ZOMBIE:
+            return False
+        try:
+            os.waitpid(proc.pid, os.WNOHANG)
+        except ChildProcessError:
+            pass
+        return True
+
+    @staticmethod
+    def _wait_or_reap(proc: "psutil.Process", timeout: int) -> None:
+        """Wait for *proc* to exit, handling zombie state transparently.
+
+        When the process is already a zombie (e.g. orchestrator thread
+        hung during shutdown and the main process exited), ``pidfd_open``
+        + ``poll()`` inside ``psutil`` will never see a state change.
+        Fall back to ``os.waitpid`` to reap the zombie.
+        """
+        if OmniServer._reap_zombie(proc):
+            return
+        try:
+            proc.wait(timeout=timeout)
+        except psutil.TimeoutExpired:
+            OmniServer._reap_zombie(proc)
+
     def _kill_process_tree(self, pid):
+        """Kill the process tree rooted at *pid*.
+
+        Terminate the parent **first** so the OmniServer can gracefully shut
+        down its stage-engine children through the orchestrator.  This avoids
+        the ``subprocess died unexpectedly`` ERROR that the APIServer monitor
+        thread logs when children are killed before the parent, which in turn
+        can cause CI watchdogs to false-trigger on the upstream ``Shutdown
+        initiated`` message.
+
+        When the parent does not exit within the grace period (e.g. CPU-
+        offloaded workers stuck in CUDA D-state), the method falls back to
+        killing children first so the parent can be reaped cleanly.
+        """
         try:
             parent = psutil.Process(pid)
             children = parent.children(recursive=True)
-            all_pids = [pid] + [child.pid for child in children]
-
-            for child in children:
-                try:
-                    child.terminate()
-                except psutil.NoSuchProcess:
-                    pass
-
-            _, still_alive = psutil.wait_procs(children, timeout=10)
-
-            for child in still_alive:
-                try:
-                    child.kill()
-                except psutil.NoSuchProcess:
-                    pass
 
+            # 1. Terminate the parent first — let it run its graceful
+            #    shutdown cascade (orchestrator → stage pools → engine cores).
             try:
                 parent.terminate()
-                parent.wait(timeout=10)
-            except (psutil.NoSuchProcess, psutil.TimeoutExpired):
+            except psutil.NoSuchProcess:
+                pass
+
+            # 2. Give the parent time to shut down its children cleanly.
+            parent_exited = False
+            try:
+                parent.wait(timeout=15)
+                parent_exited = True
+            except psutil.NoSuchProcess:
+                parent_exited = True
+            except psutil.TimeoutExpired:
+                parent_exited = OmniServer._reap_zombie(parent)
+
+            if not parent_exited:
+                # Parent is stuck — children (e.g. CPU-offloaded CFG workers)
+                # are likely in uninterruptible sleep.  Kill children first
+                # so the parent can be reaped without lingering as a zombie.
+                for child in children:
+                    try:
+                        child.kill()
+                    except psutil.NoSuchProcess:
+                        pass
+                psutil.wait_procs(children, timeout=5)
                 try:
                     parent.kill()
                 except psutil.NoSuchProcess:
                     pass
+                OmniServer._wait_or_reap(parent, timeout=5)
+            else:
+                # Parent exited cleanly — clean up any remaining children.
+                for child in children:
+                    try:
+                        if child.is_running():
+                            child.terminate()
+                    except psutil.NoSuchProcess:
+                        pass
 
+                gone, still_alive = psutil.wait_procs(children, timeout=10)
+
+                for child in still_alive:
+                    try:
+                        child.kill()
+                    except psutil.NoSuchProcess:
+                        pass
+
+                try:
+                    if parent.is_running() and not OmniServer._reap_zombie(parent):
+                        parent.kill()
+                        parent.wait(timeout=10)
+                except psutil.NoSuchProcess:
+                    pass
+
+            # 3. Final sweep — ``kill -9`` anything that escaped.
             time.sleep(1)
-            alive_processes = []
-            for check_pid in all_pids:
-                if psutil.pid_exists(check_pid):
-                    alive_processes.append(check_pid)
+            alive_processes: list[int] = []
+            for child in children:
+                try:
+                    if child.is_running():
+                        alive_processes.append(child.pid)
+                except psutil.NoSuchProcess:
+                    pass
+            # Only count the parent as alive if it is NOT a zombie
+            # (zombies are already dead — just waiting to be reaped).
+            try:
+                if parent.is_running() and parent.status() != psutil.STATUS_ZOMBIE:
+                    alive_processes.append(parent.pid)
+            except psutil.NoSuchProcess:
+                pass
 
             if alive_processes:
                 print(f"Warning: Processes still alive: {alive_processes}")

tests/test_arg_utils.py

@@ -164,7 +164,8 @@ def test_split_mixed_real_world():
     # orchestrator-only keys never reach engine
     assert "stage_init_timeout" not in engine
     assert "deploy_config" not in engine
-    assert "async_chunk" not in engine
+    # async_chunk is in SHARED_FIELDS so it flows to both
+    assert "async_chunk" in engine
 
 
 # ============================================================================
@@ -172,24 +173,47 @@ def test_split_mixed_real_world():
 # ============================================================================
 
 
-def test_user_typed_unclassified_warns(caplog):
+def test_user_typed_unclassified_warns(monkeypatch):
     """If the user types a flag we can't route, warn — don't silently drop."""
-    raw = {"bogus_flag": "value", "max_num_seqs": 64}
-    with caplog.at_level(logging.WARNING, logger="vllm_omni.engine.arg_utils"):
+    import io
+
+    buf = io.StringIO()
+    vllm_logger = logging.getLogger("vllm")
+    orig_handlers = list(vllm_logger.handlers)
+    vllm_logger.handlers.clear()
+    handler = logging.StreamHandler(buf)
+    handler.setLevel(logging.WARNING)
+    vllm_logger.addHandler(handler)
+    try:
+        raw = {"bogus_flag": "value", "max_num_seqs": 64}
         split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed={"bogus_flag"})
-    assert any("bogus_flag" in rec.message for rec in caplog.records), (
-        f"Expected warning mentioning 'bogus_flag', got: {[rec.message for rec in caplog.records]}"
-    )
+        assert "bogus_flag" in buf.getvalue(), f"Expected 'bogus_flag' in log output, got: {buf.getvalue()!r}"
+    finally:
+        vllm_logger.handlers.clear()
+        for h in orig_handlers:
+            vllm_logger.addHandler(h)
 
 
-def test_unclassified_without_user_typed_silent(caplog):
+def test_unclassified_without_user_typed_silent():
     """Without user_typed, unclassified keys drop silently (argparse defaults
     for server flags shouldn't spam logs on every launch)."""
-    raw = {"host": "0.0.0.0", "port": 8091, "max_num_seqs": 64}
-    with caplog.at_level(logging.WARNING, logger="vllm_omni.engine.arg_utils"):
+    import io
+
+    vllm_logger = logging.getLogger("vllm")
+    orig_handlers = list(vllm_logger.handlers)
+    vllm_logger.handlers.clear()
+    buf = io.StringIO()
+    handler = logging.StreamHandler(buf)
+    handler.setLevel(logging.WARNING)
+    vllm_logger.addHandler(handler)
+    try:
+        raw = {"host": "0.0.0.0", "port": 8091, "max_num_seqs": 64}
         split_kwargs(raw, engine_cls=_FakeEngineArgs, user_typed=None)
-    # No warnings because we don't know these were user-typed.
-    assert not any("host" in rec.message or "port" in rec.message for rec in caplog.records)
+        assert "host" not in buf.getvalue()
+    finally:
+        vllm_logger.handlers.clear()
+        for h in orig_handlers:
+            vllm_logger.addHandler(h)
 
 
 # ============================================================================
@@ -290,7 +314,7 @@ def test_internal_blacklist_keys_derived_from_orchestrator():
     # Spot-check expected entries
     assert "stage_init_timeout" in blacklist
     assert "deploy_config" in blacklist
-    assert "async_chunk" in blacklist
+    assert "async_chunk" not in blacklist  # now in SHARED_FIELDS
     # Shared fields must NOT appear — they flow to both orchestrator and engine
     assert "model" not in blacklist
     assert "log_stats" not in blacklist
@@ -339,18 +363,31 @@ def test_ambiguous_field_strict_raises():
         split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=True)
 
 
-def test_ambiguous_field_non_strict_routes_to_orchestrator(caplog):
+def test_ambiguous_field_non_strict_routes_to_orchestrator():
     """strict=False logs ERROR but routes the ambiguous field to orchestrator."""
 
     @dataclass
     class _AmbiguousEngine:
         deploy_config: str | None = None
 
-    with caplog.at_level(logging.ERROR, logger="vllm_omni.engine.arg_utils"):
+    import io
+
+    vllm_logger = logging.getLogger("vllm")
+    orig_handlers = list(vllm_logger.handlers)
+    vllm_logger.handlers.clear()
+    buf = io.StringIO()
+    handler = logging.StreamHandler(buf)
+    handler.setLevel(logging.ERROR)
+    vllm_logger.addHandler(handler)
+    try:
         orch, engine = split_kwargs({"deploy_config": "x"}, engine_cls=_AmbiguousEngine, strict=False)
-    assert orch.deploy_config == "x"
-    assert "deploy_config" not in engine
-    assert any("both OrchestratorArgs" in r.message for r in caplog.records)
+        assert orch.deploy_config == "x"
+        assert "deploy_config" not in engine
+        assert "both OrchestratorArgs" in buf.getvalue()
+    finally:
+        vllm_logger.handlers.clear()
+        for h in orig_handlers:
+            vllm_logger.addHandler(h)
 
 
 # Sentinel-default precedence invariants (#3035)
@@ -426,24 +463,20 @@ def test_help_text_preserves_default_after_nullify():
 
 _OMNIENGINEARGS_USER_INPUT_FIELDS = frozenset(
     {
-        "model_stage",
         "model_arch",
         "engine_output_type",
         "hf_config_name",
         "custom_process_next_stage_input_func",
         "subtalker_sampling_params",
-        "async_chunk",
         "omni_kv_config",
         "quantization_config",
         "worker_type",
         "task_type",
         "worker_cls",
-        "enable_sleep_mode",
         "omni_master_address",
         "omni_master_port",
         "stage_configs_path",
         "output_modalities",
-        "log_stats",
         "custom_pipeline_args",
     }
 )

tools/nightly/generate_nightly_perf_html.py

@@ -250,7 +250,7 @@ def _collect_diffusion_records(input_dir: str) -> list[dict[str, Any]]:
 
 
 def _sort_omni_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
-    by_date_desc = sorted(records, key=lambda r: (r.get("date") or ""), reverse=True)
+    by_date_desc = sorted(records, key=lambda r: r.get("date") or "", reverse=True)
     return sorted(
         by_date_desc,
         key=lambda r: (
@@ -264,8 +264,8 @@ def _sort_omni_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
 
 
 def _sort_diffusion_records(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
-    by_date_desc = sorted(records, key=lambda r: (r.get("date") or ""), reverse=True)
-    return sorted(by_date_desc, key=lambda r: (r.get("test_name") or ""))
+    by_date_desc = sorted(records, key=lambda r: r.get("date") or "", reverse=True)
+    return sorted(by_date_desc, key=lambda r: r.get("test_name") or "")
 
 
 def _ensure_parent_dir(path: str) -> None:

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -24,9 +24,11 @@ from vllm_omni.core.sched.omni_scheduling_coordinator import (
     OmniSchedulingCoordinator,
     uses_full_payload_input_coordinator,
 )
+from vllm_omni.core.sched.utils import omni_routed_experts_for_request
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
+from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.outputs import OmniConnectorOutput
 
@@ -321,7 +323,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     req = self.requests.get(req_id)
                     if req is not None and not req.is_finished():
                         outputs[req.client_index].append(
-                            EngineCoreOutput(
+                            OmniEngineCoreOutput(
                                 request_id=req_id,
                                 new_token_ids=[],
                                 kv_transfer_params={"kv_ready": True},
@@ -417,11 +419,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     stopped = True
 
             if stopped:
-                if (
-                    model_runner_output.routed_experts_dict is not None
-                    and req_id in model_runner_output.routed_experts_dict
-                ):
-                    routed_experts = model_runner_output.routed_experts_dict[req_id]
+                if model_runner_output.routed_experts is not None:
+                    routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
 
                 # Capture finish_reason BEFORE _handle_stopped_request, which may
                 # reset the status to WAITING for streaming requests that continue.
@@ -452,7 +451,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     stopped_preempted_reqs.add(request)
 
             # Extract sample logprobs if needed.
-            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
+            if request.sampling_params is not None and request.sampling_params.num_logprobs is not None and logprobs:
                 new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
 
             if num_nans_in_logits is not None and req_id in num_nans_in_logits:
@@ -463,7 +462,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
                 # Add EngineCoreOutput for this Request.
                 outputs[request.client_index].append(
-                    EngineCoreOutput(
+                    OmniEngineCoreOutput(
                         request_id=req_id,
                         new_token_ids=new_token_ids,
                         finish_reason=finish_reason,
@@ -501,7 +500,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
             for request in requests:
                 outputs[request.client_index].append(
-                    EngineCoreOutput(
+                    OmniEngineCoreOutput(
                         request_id=request.request_id,
                         new_token_ids=[],
                         finish_reason=request.get_finished_reason(),
@@ -633,6 +632,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         """
         req_id = session.request_id
         self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids)
+        session.discard_latest_async_tokens = False
         outstanding_async_tokens = getattr(session, "num_output_placeholders", 0)
         if outstanding_async_tokens > 0:
             # Async scheduling may already have sampled the previous

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -29,9 +29,11 @@ from vllm_omni.core.sched.omni_scheduling_coordinator import (
     uses_full_payload_input_coordinator,
 )
 from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
+from vllm_omni.core.sched.utils import omni_routed_experts_for_request
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
+from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.outputs import OmniConnectorOutput, OmniModelRunnerOutput
 
 logger = init_logger(__name__)
@@ -499,11 +501,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 stopped = True
 
             if stopped:
-                if (
-                    model_runner_output.routed_experts_dict is not None
-                    and req_id in model_runner_output.routed_experts_dict
-                ):
-                    routed_experts = model_runner_output.routed_experts_dict[req_id]
+                if model_runner_output.routed_experts is not None:
+                    routed_experts = omni_routed_experts_for_request(model_runner_output.routed_experts, request)
                 finish_reason = request.get_finished_reason()
                 finished = self._handle_stopped_request(request)
                 if not finished:
@@ -524,7 +523,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                     stopped_running_reqs.add(request)
 
             # Extract sample logprobs if needed.
-            if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
+            if request.sampling_params is not None and request.sampling_params.num_logprobs is not None and logprobs:
                 new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
 
             if new_token_ids and self.structured_output_manager.should_advance(request):
@@ -543,7 +542,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
             if new_token_ids or pooler_output is not None or kv_transfer_params or stopped:
                 outputs[request.client_index].append(
-                    EngineCoreOutput(
+                    OmniEngineCoreOutput(
                         request_id=req_id,
                         new_token_ids=new_token_ids,
                         finish_reason=finish_reason,
@@ -580,7 +579,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                         getattr(request, "external_req_id", None),
                     )
             outputs[request.client_index].append(
-                EngineCoreOutput(
+                OmniEngineCoreOutput(
                     request_id=request.request_id,
                     new_token_ids=[],
                     finish_reason=finish_reason,
@@ -607,7 +606,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
             for request in requests:
                 outputs[request.client_index].append(
-                    EngineCoreOutput(
+                    OmniEngineCoreOutput(
                         request_id=request.request_id,
                         new_token_ids=[],
                         finish_reason=request.get_finished_reason(),

vllm_omni/core/sched/utils.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared utilities for omni schedulers."""
+
+import numpy as np
+from vllm.v1.outputs import RoutedExpertsLists
+
+
+def omni_routed_experts_for_request(routed_experts: RoutedExpertsLists, request) -> np.ndarray | None:
+    """Extract per-request routed experts from RoutedExpertsLists using slot_mapping.
+
+    Matches upstream RoutedExpertsManager.get() pattern — filters routing_data
+    rows whose slot_mapping entries belong to this request's block_table.
+    """
+    if routed_experts is None:
+        return None
+    slots = getattr(request, "block_table", None)
+    if slots is None:
+        return None
+    slot_set = set(slots)
+    mask = np.isin(routed_experts.slot_mapping, list(slot_set))
+    data = routed_experts.routing_data[mask]
+    return data if data.size > 0 else None

vllm_omni/diffusion/diffusion_engine.py

@@ -867,7 +867,7 @@ class DiffusionEngine:
                 worker_thread.join(timeout=10)
             if worker_thread.is_alive():
                 logger.warning(
-                    "Worker thread did not terminate within 10s; scheduler and executor shutdown will be skipped."
+                    "Worker thread did not terminate within 10s; scheduler and executor shutdown will be deferred."
                 )
                 return
             else:

vllm_omni/diffusion/model_loader/hub_prefetch.py

@@ -64,6 +64,7 @@ import contextlib
 import errno
 import logging
 import os
+import time
 from collections.abc import Iterable, Iterator
 
 logger = logging.getLogger(__name__)
@@ -112,9 +113,43 @@ def _safe_repo_filename(model: str) -> str:
     return model.replace("/", "__").replace(os.sep, "__") + ".lock"
 
 
+def _dotfile_lock_acquire(lock_dir: str, model: str, timeout: float = 300.0, poll_interval: float = 0.5) -> bool:
+    """Acquire an exclusive lock via atomic directory creation.
+
+    ``os.makedirs(path, exist_ok=False)`` is atomic on POSIX filesystems
+    (including Lustre and NFS), making it a reliable fallback when
+    ``fcntl.flock`` is unsupported (``ENOLCK`` on FSx/Lustre mounts
+    without the ``flock`` mount option).
+
+    Returns ``True`` if the lock was acquired, ``False`` on timeout.
+    The caller must call ``os.rmdir(lock_path)`` to release.
+    """
+    lock_path = os.path.join(lock_dir, _safe_repo_filename(model) + ".dir")
+
+    deadline = time.monotonic() + timeout
+    while True:
+        try:
+            os.makedirs(lock_path, exist_ok=False)
+            logger.info("Acquired dotfile prefetch lock for %s at %s", model, lock_path)
+            return True
+        except FileExistsError:
+            if time.monotonic() >= deadline:
+                logger.warning(
+                    "Timed out waiting for dotfile prefetch lock %s after %.0fs; proceeding unlocked",
+                    lock_path,
+                    timeout,
+                )
+                return False
+            time.sleep(poll_interval)
+
+
 @contextlib.contextmanager
 def _repo_prefetch_lock(model: str) -> Iterator[None]:
-    """Hold an exclusive ``fcntl.flock`` keyed on the HF repo id.
+    """Hold an exclusive lock keyed on the HF repo id.
+
+    Tries ``fcntl.flock`` first; falls back to an atomic directory-creation
+    dotfile lock when flock is unsupported (e.g. FSx for Lustre without the
+    ``flock`` mount option, which returns ``ENOLCK``).
 
     Why we still need a node-wide lock on top of ``snapshot_download``'s
     own per-blob ``.lock`` files:
@@ -131,71 +166,93 @@ def _repo_prefetch_lock(model: str) -> Iterator[None]:
     * Multiple unrelated processes on the same node (concurrent CI jobs,
       multiple ``OmniServer`` instances, every ``DiffusionWorker``
       subprocess from the multiproc executor) can enter this code path
-      in parallel. ``flock`` on a per-repo file makes the whole
+      in parallel. An exclusive lock makes the whole
       ``snapshot_download(...)`` call a true node-wide critical
       section, so the *first* entrant fully populates the cache and
       every subsequent ``from_pretrained`` sees a warm, complete tree.
 
-    Best-effort: if we cannot create / acquire the lock (no ``fcntl``,
-    read-only FS, NFS without flock support) we log and run the download
-    anyway. Worst case behaviour reverts to "snapshot_download +
-    transformers v5 race" which is exactly the pre-fix state.
+    Best-effort: if we cannot acquire any lock (no ``fcntl``, read-only
+    FS, dotfile lock timeout) we log and run the download anyway. Worst
+    case behaviour reverts to "snapshot_download + transformers v5 race"
+    which is exactly the pre-fix state.
     """
+    lock_dir = None
+    lock_path = None
+    dotfile_held = None
+    fd = None
+    flock_held = False
+
+    # --- fcntl.flock path ---
     try:
         import fcntl  # type: ignore[import-not-found]
     except ImportError:  # pragma: no cover - non-POSIX (Windows)
-        logger.debug("fcntl unavailable on this platform; running prefetch unlocked")
-        yield
-        return
+        fcntl = None
 
-    try:
-        lock_dir = _node_lock_dir()
-    except OSError as exc:
-        logger.warning("Could not allocate lock dir for prefetch of %s (%s); skipping flock", model, exc)
-        yield
-        return
-
-    lock_path = os.path.join(lock_dir, _safe_repo_filename(model))
-
-    try:
-        fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644)
-    except OSError as exc:
-        logger.warning("Could not open prefetch lock %s (%s); proceeding without flock", lock_path, exc)
-        yield
-        return
-
-    locked = False
-    try:
+    if fcntl is not None:
         try:
-            fcntl.flock(fd, fcntl.LOCK_EX)
-            locked = True
-            logger.debug("Acquired prefetch flock for %s at %s", model, lock_path)
+            lock_dir = _node_lock_dir()
         except OSError as exc:
-            # FS without flock support (ENOLCK / EOPNOTSUPP / EINVAL).
-            # EACCES on some Lustre configs without ``flock`` mount opt.
-            if exc.errno in (errno.ENOLCK, errno.EOPNOTSUPP, errno.EACCES, errno.EINVAL):
-                logger.warning(
-                    "fcntl.flock not supported on '%s' (%s: %s); running prefetch unlocked",
-                    lock_path,
-                    type(exc).__name__,
-                    exc,
-                )
-            else:
-                raise
+            logger.warning("Could not allocate lock dir for prefetch of %s (%s); skipping flock", model, exc)
+            fcntl = None  # force dotfile fallback
+
+    if fcntl is not None and lock_dir is not None:
+        lock_path = os.path.join(lock_dir, _safe_repo_filename(model))
+        try:
+            fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644)
+        except OSError as exc:
+            logger.warning("Could not open prefetch lock %s (%s); falling back to dotfile lock", lock_path, exc)
+            fd = None
+
+        if fd is not None:
+            try:
+                fcntl.flock(fd, fcntl.LOCK_EX)
+                flock_held = True
+                logger.info("Acquired flock prefetch lock for %s at %s", model, lock_path)
+            except OSError as exc:
+                if exc.errno in (errno.ENOLCK, errno.EOPNOTSUPP, errno.EACCES, errno.EINVAL):
+                    logger.warning(
+                        "fcntl.flock not supported on '%s' (%s: %s); falling back to dotfile lock",
+                        lock_path,
+                        type(exc).__name__,
+                        exc,
+                    )
+                    os.close(fd)
+                    fd = None
+                else:
+                    raise
+
+    # --- dotfile fallback ---
+    if not flock_held:
+        if lock_dir is None:
+            try:
+                lock_dir = _node_lock_dir()
+            except OSError as exc:
+                logger.warning("Could not allocate lock dir (%s); running prefetch unlocked", exc)
+                yield
+                return
+
+        if _dotfile_lock_acquire(lock_dir, model):
+            dotfile_held = os.path.join(lock_dir, _safe_repo_filename(model) + ".dir")
+
+    try:
         yield
     finally:
-        if locked:
+        if flock_held and fd is not None:
             with contextlib.suppress(OSError):
                 fcntl.flock(fd, fcntl.LOCK_UN)
-        with contextlib.suppress(OSError):
-            os.close(fd)
+        if fd is not None:
+            with contextlib.suppress(OSError):
+                os.close(fd)
+        if dotfile_held is not None:
+            with contextlib.suppress(OSError):
+                os.rmdir(dotfile_held)
 
 
 def prefetch_subfolders(
     model: str,
     subfolders: Iterable[str],
     *,
-    local_files_only: bool = False,
+    local_files_only: bool | None = None,
     include_root_metadata: bool = True,
 ) -> None:
     """Materialise ``model``'s ``subfolders`` in the HF cache before loading.
@@ -206,17 +263,21 @@ def prefetch_subfolders(
         subfolders: Iterable of subfolder names (e.g. ``["text_encoder",
             "vae"]``) whose contents need to be fully present before any
             worker calls ``from_pretrained(subfolder=...)``.
-        local_files_only: When True, skip the prefetch entirely. The caller
-            has explicitly promised the cache is already populated (as happens
-            for local model checkouts), so hitting the network would defeat
-            the intent and may fail in air-gapped environments.
+        local_files_only: When ``True``, skip the prefetch entirely.
+            When ``None`` (default), auto-detect: skip if *model* is a
+            local directory, run otherwise.
         include_root_metadata: When True, also pull ``*.json`` at the repo
             root so ``model_index.json`` / ``config.json`` resolution during
             ``from_pretrained`` also hits a warm cache.
     """
+    if local_files_only is None:
+        local_files_only = os.path.isdir(model)
+
     if local_files_only or not model or os.path.isdir(model):
         return
 
+    logger.info("Prefetching %s subfolders: %s", model, list(subfolders))
+
     try:
         from huggingface_hub import snapshot_download
     except ImportError:  # pragma: no cover - huggingface_hub is a hard dep
@@ -251,6 +312,7 @@ def prefetch_subfolders(
                 repo_id=model,
                 allow_patterns=allow_patterns,
             )
+        logger.info("Prefetch complete for %s", model)
     except Exception as exc:
         # Best-effort: propagate only via logging. The subsequent
         # ``from_pretrained`` call will raise a clearer, call-site-specific
@@ -312,3 +374,31 @@ def _looks_like_auth_error(exc: BaseException) -> bool:
     # only load-bearing signal.
     msg = str(exc).lower()
     return "401 client error" in msg or "403 client error" in msg or "gatedrepo" in msg
+
+
+def retry_on_missing_shard(load_fn, *, max_retries: int = 3, base_delay: float = 5.0):
+    """Call *load_fn* with retry on the transformers v5 shard-resolution race.
+
+    When the prefetch lock cannot be acquired (e.g. flock unsupported on
+    the filesystem and dotfile lock times out), ``from_pretrained`` may
+    still hit the ``cached_files`` race. This wrapper retries with
+    exponential backoff when the OSError message matches the specific
+    "does not appear to have a file named" pattern.
+    """
+    for attempt in range(max_retries):
+        try:
+            return load_fn()
+        except OSError as exc:
+            if "does not appear to have a file" not in str(exc):
+                raise
+            if attempt == max_retries - 1:
+                raise
+            delay = base_delay * (attempt + 1)
+            logger.warning(
+                "from_pretrained failed with shard-resolution race (%s); retrying in %.1fs (attempt %d/%d)",
+                exc,
+                delay,
+                attempt + 1,
+                max_retries,
+            )
+            time.sleep(delay)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -26,7 +26,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
@@ -275,21 +275,22 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         self.device = get_local_device()
         model = od_config.model
         # Check if model is a local path
-        local_files_only = os.path.exists(model)
+        local_files_only = os.path.isdir(model)
 
         # See pipeline_qwen_image_edit_plus: guard against transformers v5
         # multi-worker race on partial subfolder shard sets (Buildkite #1043).
         prefetch_subfolders(
             model,
             ["scheduler", "text_encoder", "vae", "tokenizer"],
-            local_files_only=local_files_only,
         )
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
+        self.text_encoder = retry_on_missing_shard(
+            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model, subfolder="text_encoder", local_files_only=local_files_only
+            )
         )
         # Qwen2.5-VL ships a vision tower that text-to-image does not use.
         # Drop it while the model is still on CPU, before moving to GPU, so

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -27,7 +27,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
@@ -243,22 +243,23 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         model = od_config.model
 
         # Check if model is a local path
-        local_files_only = os.path.exists(model)
+        local_files_only = os.path.isdir(model)
 
         # See pipeline_qwen_image_edit_plus: guard against transformers v5
         # multi-worker race on partial subfolder shard sets (Buildkite #1043).
         prefetch_subfolders(
             model,
             ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
-            local_files_only=local_files_only,
         )
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
-        ).to(self.device)
+        self.text_encoder = retry_on_missing_shard(
+            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model, subfolder="text_encoder", local_files_only=local_files_only
+            ).to(self.device)
+        )
 
         self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
             self.device
@@ -266,8 +267,8 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
         self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.processor = Qwen2VLProcessor.from_pretrained(
-            model, subfolder="processor", local_files_only=local_files_only
+        self.processor = retry_on_missing_shard(
+            lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only)
         )
 
         self.stage = None

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -25,7 +25,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
 from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
@@ -208,7 +208,7 @@ class QwenImageEditPlusPipeline(
         model = od_config.model
 
         # Check if model is a local path
-        local_files_only = os.path.exists(model)
+        local_files_only = os.path.isdir(model)
 
         # Defend against a transformers v5 multi-worker race where a peer
         # worker's partially-written shard set makes our from_pretrained
@@ -219,15 +219,16 @@ class QwenImageEditPlusPipeline(
         prefetch_subfolders(
             model,
             ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
-            local_files_only=local_files_only,
         )
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
-        ).to(self.device)
+        self.text_encoder = retry_on_missing_shard(
+            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model, subfolder="text_encoder", local_files_only=local_files_only
+            ).to(self.device)
+        )
 
         self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
             self.device
@@ -236,8 +237,8 @@ class QwenImageEditPlusPipeline(
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
         self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.processor = Qwen2VLProcessor.from_pretrained(
-            model, subfolder="processor", local_files_only=local_files_only
+        self.processor = retry_on_missing_shard(
+            lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only)
         )
 
         self.stage = None

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -24,7 +24,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.autoencoder_kl_qwenimage import (
     AutoencoderKLQwenImage,
@@ -216,29 +216,30 @@ class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParalle
         self.device = get_local_device()
         model = od_config.model
         # Check if model is a local path
-        local_files_only = os.path.exists(model)
+        local_files_only = os.path.isdir(model)
 
         # See pipeline_qwen_image_edit_plus: guard against transformers v5
         # multi-worker race on partial subfolder shard sets (Buildkite #1043).
         prefetch_subfolders(
             model,
             ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
-            local_files_only=local_files_only,
         )
 
         # modules keep same as transformers & diffusers
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
-        ).to(self.device)
+        self.text_encoder = retry_on_missing_shard(
+            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
+                model, subfolder="text_encoder", local_files_only=local_files_only
+            ).to(self.device)
+        )
         self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
             self.device
         )
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.processor = Qwen2VLProcessor.from_pretrained(
-            model, subfolder="processor", local_files_only=local_files_only
+        self.processor = retry_on_missing_shard(
+            lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only)
         )
 
         # modules re-implemented for vLLM-Omni

vllm_omni/distributed/omni_coordinator/omni_coordinator.py

@@ -43,17 +43,20 @@ class OmniCoordinator:
             heartbeat_timeout: Seconds before a replica is considered
                 unhealthy if no heartbeat / update is received.
         """
-        self._router_zmq_addr = router_zmq_addr
-        self._pub_zmq_addr = pub_zmq_addr
         self._heartbeat_timeout = heartbeat_timeout
 
         # Dedicated ZMQ context for this coordinator instance.
         self._ctx = zmq.Context()
         self._router = self._ctx.socket(zmq.ROUTER)
-        self._router.bind(self._router_zmq_addr)
+        self._router.bind(router_zmq_addr)
+        # Recover the actual bound address (port=0 → real port).
+        self.router_zmq_addr = self._router.getsockopt_string(zmq.LAST_ENDPOINT)
+        self._pub_zmq_addr = pub_zmq_addr  # keep original for internal use
 
         self._pub = self._ctx.socket(zmq.PUB)
-        self._pub.bind(self._pub_zmq_addr)
+        self._pub.bind(pub_zmq_addr)
+        # Recover the actual bound address (port=0 → real port).
+        self.pub_zmq_addr = self._pub.getsockopt_string(zmq.LAST_ENDPOINT)
 
         self._replicas: dict[str, ReplicaInfo] = {}
         self._lock = threading.Lock()

vllm_omni/engine/arg_utils.py

@@ -144,7 +144,8 @@ class OmniEngineArgs(EngineArgs):
     force_cutlass_fp8: bool | None = None
     worker_type: str | None = None
     task_type: str | None = None
-    worker_cls: str = None
+    worker_cls: str = None  # type: ignore[assignment]  # Upstream default is "auto"; omni resolves
+    # in __post_init__ based on worker_type (ar/generation), so None is safe here.
     enable_sleep_mode: bool = False
     omni: bool = False
 
@@ -444,6 +445,48 @@ class OrchestratorArgs:
     # === Pre-built Objects ===
     parallel_config: Any = None
 
+    # === Diffusion model config ===
+    num_gpus: int | None = None
+    model_class_name: str | None = None
+    diffusion_load_format: str | None = None
+    diffusers_load_kwargs: str = "{}"
+    diffusers_call_kwargs: str = "{}"
+    ulysses_degree: int | None = None
+    ulysses_mode: str = "strict"
+    ring_degree: int | None = None
+    diffusion_quantization_config: str | None = None
+    use_hsdp: bool = False
+    hsdp_shard_size: int = -1
+    hsdp_replicate_size: int = 1
+    diffusion_attention_backend: str | None = None
+    diffusion_attention_config: str | None = None
+    cache_backend: str = "none"
+    cache_config: str | None = None
+    enable_cache_dit_summary: bool = False
+    step_execution: bool = False
+    vae_use_slicing: bool = False
+    vae_use_tiling: bool = False
+    enable_multithread_weight_load: bool = True
+    num_weight_load_threads: int = 4
+    enable_cpu_offload: bool = False
+    enable_layerwise_offload: bool = False
+    boundary_ratio: float | None = None
+    flow_shift: float | None = None
+    diffusion_kv_cache_dtype: str | None = None
+    diffusion_kv_cache_skip_steps: str | None = None
+    diffusion_kv_cache_skip_layers: str | None = None
+    cfg_parallel_size: int = 1
+    vae_patch_parallel_size: int = 1
+    default_sampling_params: str | None = None
+    max_generated_image_size: int | None = None
+    tts_max_instructions_length: int | None = None
+    enable_diffusion_pipeline_profiler: bool = False
+    enable_ar_profiler: bool = False
+    auxiliary_text_encoder: str | None = None
+    log_file: str | None = None
+    replica_id: int | None = None
+    omni_replica_address: str | None = None
+
     # === Multi-stage guards ===
     # --tokenizer is captured by the orchestrator and forwarded to stages
     # only when the stage does not define tokenizer/tokenizer_subdir itself.
@@ -459,6 +502,8 @@ SHARED_FIELDS: frozenset[str] = frozenset(
         "stage_id",  # orch: route (headless); engine: identity
         "log_stats",  # both want the flag
         "stage_configs_path",  # orch: load legacy YAML; engine: may reference for validation
+        "async_chunk",  # orch: read from CLI, redistribute; engine: per-stage flag
+        "tokenizer",  # orch: detect model type; engine: tokenization
     }
 )
 

vllm_omni/engine/async_omni_engine.py

@@ -2124,7 +2124,7 @@ class AsyncOmniEngine:
                             kwargs.get("diffusion_attention_config"),
                             attention_backend=kwargs.get("diffusion_attention_backend"),
                         )
-                quantization_config = kwargs.get("quantization_config")
+                quantization_config = kwargs.get("diffusion_quantization_config")
                 if quantization_config is not None:
                     if (
                         not hasattr(cfg.engine_args, "quantization_config")

vllm_omni/engine/output_processor.py

@@ -1,11 +1,7 @@
 from typing import Any
 
-import numpy as np
 import torch
 from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-    split_routed_experts,
-)
 from vllm.outputs import PoolingRequestOutput
 from vllm.sampling_params import RequestOutputKind
 from vllm.tokenizers import TokenizerLike
@@ -172,7 +168,6 @@ class OmniRequestState(RequestState):
         finish_reason: FinishReason | None,
         stop_reason: int | str | None,
         kv_transfer_params: dict[str, Any] | None = None,
-        routed_experts: np.ndarray | None = None,
     ) -> OmniRequestOutput | PoolingRequestOutput | None:
         """Create a request output from generation results.
 
@@ -199,7 +194,6 @@ class OmniRequestState(RequestState):
                 finish_reason,
                 stop_reason,
                 kv_transfer_params,
-                routed_experts,
             )
 
         finished = finish_reason is not None
@@ -234,16 +228,7 @@ class OmniRequestState(RequestState):
 
         external_req_id = self.external_req_id
 
-        # Split routing data into prompt and generation portions, matching
-        # upstream make_request_output behaviour.
-        prompt_routed_experts = None
-        gen_routed_experts = None
-        if routed_experts is not None:
-            prompt_len = len(self.prompt_token_ids) if self.prompt_token_ids else 0
-            num_gen = self.detokenizer.num_output_tokens() if self.detokenizer is not None else None
-            prompt_routed_experts, gen_routed_experts = split_routed_experts(routed_experts, prompt_len, num_gen)
-
-        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason, gen_routed_experts)
+        output = self._new_completion_output(new_token_ids, finish_reason, stop_reason)
 
         if self.parent_req is None:
             outputs = [output]
@@ -258,7 +243,6 @@ class OmniRequestState(RequestState):
             outputs,
             finished,
             kv_transfer_params,
-            prompt_routed_experts,
         )
 
     def _new_completion_output(
@@ -266,10 +250,11 @@ class OmniRequestState(RequestState):
         token_ids: list[int],
         finish_reason: FinishReason | None,
         stop_reason: int | str | None,
-        routed_experts: np.ndarray | None = None,
     ) -> Any:
         # Reuse base text/logprobs logic, then annotate with pooling_result.
-        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason, routed_experts)
+        # Note: upstream _new_completion_output no longer accepts routed_experts
+        # as a parameter; it reads from ``self.routed_experts_chunks`` internally.
+        base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
 
         # Inter-stage processors need the full cumulative token sequence.
         # In DELTA mode, base_output.token_ids only has the latest step's

vllm_omni/engine/stage_engine_core_proc.py

@@ -101,6 +101,11 @@ class StageEngineCoreProc(EngineCoreProc):
             set_death_signal(signal.SIGTERM)
             set_process_title(f"StageEngineCoreProc_{stage_label}_replica{omni_replica_id}_DP{dp_rank}")
             decorate_logs()
+            # Workaround for flashinfer/jit-cache version mismatch in CI.
+            # The parent process handles this gracefully via ring_globals.py,
+            # but the subprocess hits an unprotected import in TopKTopPSampler.
+            # Setting this env var allows the same graceful fallback to work.
+            os.environ.setdefault("FLASHINFER_DISABLE_VERSION_CHECK", "1")
             os.environ["VLLM_OMNI_REPLICA_ID"] = str(max(int(omni_replica_id), 0))
 
             engine_core = StageEngineCoreProc(

vllm_omni/engine/stage_init_utils.py

@@ -860,6 +860,45 @@ def build_stage0_input_processor(stage_vllm_config: Any) -> InputProcessor:
     return input_processor
 
 
+def _cleanup_stale_lock_if_dead(lock_file: str) -> bool:
+    """If *lock_file* exists and its recorded PID is dead, unlink the file.
+
+    Returns ``True`` if the stale lock was cleaned up (caller should retry),
+    ``False`` otherwise (lock holder appears alive, or file could not be read).
+    """
+    try:
+        with open(lock_file) as fh:
+            content = fh.read().strip()
+        if not content:
+            return False
+        pid = int(content)
+    except (OSError, ValueError):
+        return False
+
+    # Check whether the PID is still alive.
+    try:
+        os.kill(pid, 0)
+    except ProcessLookupError:
+        # PID does not exist — stale lock.
+        logger.info(
+            "Removing stale device lock %s (PID %s is dead)",
+            lock_file,
+            pid,
+        )
+        try:
+            os.unlink(lock_file)
+            return True
+        except OSError:
+            logger.debug("Failed to unlink stale lock %s", lock_file)
+            return False
+    except PermissionError:
+        # PID exists but we cannot signal it (different user) — treat as alive.
+        return False
+
+    # PID is alive — legitimate lock holder.
+    return False
+
+
 def acquire_device_locks(
     stage_id: int,
     engine_args_dict: dict[str, Any],
@@ -938,6 +977,7 @@ def acquire_device_locks(
         for device_id in devices_to_lock:
             lock_file = f"/tmp/vllm_omni_device_{device_id}_init.lock"
             lock_acquired = False
+            already_cleaned_stale = False  # only try stale cleanup once per device
 
             while not lock_acquired:
                 try:
@@ -952,6 +992,11 @@ def acquire_device_locks(
                         logger.debug("Acquired exclusive lock for device %s", device_id)
                     except BlockingIOError:
                         os.close(lock_fd)
+                        # Detect and clean stale locks from dead processes.
+                        if not already_cleaned_stale:
+                            already_cleaned_stale = True
+                            if _cleanup_stale_lock_if_dead(lock_file):
+                                continue  # retry flock immediately
                         if time.time() - wait_start > stage_init_timeout:
                             logger.warning(
                                 "Timeout waiting for device %s initialization lock, proceeding anyway",

vllm_omni/engine/stage_pool.py

@@ -22,6 +22,7 @@ from vllm_omni.engine.stage_client import (
     StagePoolDiffusionClient,
     StagePoolLLMClient,
 )
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
 from vllm_omni.metrics.stats import StageStats
 from vllm_omni.metrics.utils import count_tokens_from_outputs
@@ -533,6 +534,10 @@ class StagePool:
     ) -> int:
         """Submit a stage-entry request into this pool."""
         params = params_override if params_override is not None else req_state.sampling_params_list[self.stage_id]
+        # Convert plain vllm SamplingParams for single-stage diffusion models
+        # that receive sampling params from the user/caller directly.
+        if self.stage_type == "diffusion" and not isinstance(params, OmniDiffusionSamplingParams):
+            params = OmniDiffusionSamplingParams()
         submit_kwargs = dict(submit_kwargs or {})
         if self.stage_type == "diffusion":
             replica_id = await self._pick_or_select(
@@ -593,6 +598,8 @@ class StagePool:
     ) -> int:
         """Submit a streaming update to an already admitted request."""
         params = req_state.sampling_params_list[self.stage_id]
+        if self.stage_type == "diffusion" and not isinstance(params, OmniDiffusionSamplingParams):
+            params = OmniDiffusionSamplingParams()
         replica_id = self.get_bound_replica_id(request_id)
         if replica_id is None or self.clients[replica_id] is None:
             replica_id = await self._pick_or_select(request_id)

vllm_omni/entrypoints/cli/serve.py

@@ -436,7 +436,7 @@ class OmniServeCommand(CLISubcommand):
             "Equivalent to setting DiffusionParallelConfig.ring_degree.",
         )
         omni_config_group.add_argument(
-            "--quantization-config",
+            "--diffusion-quantization-config",
             type=json.loads,
             default=None,
             help=(

vllm_omni/entrypoints/omni_base.py

@@ -463,7 +463,7 @@ class OmniBase(PDDisaggregationMixin):
             peak_memory_mb=peak_memory_mb,
         )
 
-    def shutdown(self) -> None:
+    def shutdown(self, timeout: float | None = None) -> None:
         logger.info("[%s] Shutting down", self.__class__.__name__)
         self._shutdown_base()
 

vllm_omni/entrypoints/openai/api_server.py

@@ -55,13 +55,8 @@ from vllm.entrypoints.openai.engine.protocol import (
 from vllm.entrypoints.openai.models.protocol import BaseModelPath
 from vllm.entrypoints.openai.models.serving import OpenAIServingModels
 from vllm.entrypoints.openai.orca_metrics import metrics_header
-from vllm.entrypoints.openai.realtime.serving import OpenAIServingRealtime
 from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses
 from vllm.entrypoints.openai.server_utils import get_uvicorn_log_config
-from vllm.entrypoints.openai.speech_to_text.serving import (
-    OpenAIServingTranscription,
-    OpenAIServingTranslation,
-)
 from vllm.entrypoints.openai.utils import validate_json_request
 from vllm.entrypoints.pooling.classify.serving import ServingClassification
 from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding
@@ -74,6 +69,13 @@ from vllm.entrypoints.serve.disagg.serving import ServingTokens
 from vllm.entrypoints.serve.instrumentator.basic import base
 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
 from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
+from vllm.entrypoints.speech_to_text.realtime.serving import OpenAIServingRealtime
+from vllm.entrypoints.speech_to_text.transcription.serving import (
+    OpenAIServingTranscription,
+)
+from vllm.entrypoints.speech_to_text.translation.serving import (
+    OpenAIServingTranslation,
+)
 from vllm.entrypoints.utils import (
     create_error_response,
     load_aware_call,
@@ -355,7 +357,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
     warnings_module.filterwarnings("ignore", message=".*PydanticSerializationUnexpectedValue.*", category=UserWarning)
 
     # Add process-specific prefix to stdout and stderr.
-    decorate_logs("APIServer")
+    decorate_logs("APIServer", skip_if_decorated=True)
 
     listen_address, sock = setup_openai_server(args)
 
@@ -646,6 +648,19 @@ async def omni_init_app_state(
 
     state.vllm_config = vllm_config
 
+    # Propagate enable_in_reasoning to the API-server process. The engine core
+    # runs in a separate process, so the contextvar that backs
+    # `get_current_vllm_config_or_none()` is None on this stack. Tool parsers
+    # call `get_enable_structured_outputs_in_reasoning()` during request
+    # handling and need to see the real flag, otherwise they silently fall
+    # back to False and mismatch the engine-side bitmask gating.
+    if vllm_config is not None:
+        from vllm.tool_parsers.structural_tag_registry import (
+            set_enable_structured_outputs_in_reasoning,
+        )
+
+        set_enable_structured_outputs_in_reasoning(vllm_config.structured_outputs_config.enable_in_reasoning)
+
     # Get supported tasks
     supported_tasks: set[str] = {"generate"}
     if hasattr(engine_client, "get_supported_tasks"):
@@ -1531,6 +1546,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
                 prompt=request.prompt,
                 extra_body=extra_body,
                 request_id=f"img_gen-{random_uuid()}",
+                raw_request=raw_request,
             )
             if isinstance(generation_result, ErrorResponse):
                 return JSONResponse(
@@ -1937,6 +1953,7 @@ async def edit_images(
                 output_format=output_format,
                 output_compression=output_compression,
                 size=size_str,
+                raw_request=raw_request,
             )
             if stream and not isinstance(generation_result, ErrorResponse):
                 return StreamingResponse(

vllm_omni/entrypoints/openai/realtime_connection.py

@@ -9,8 +9,8 @@ from uuid import uuid4
 
 import numpy as np
 from vllm.entrypoints.openai.engine.protocol import UsageInfo
-from vllm.entrypoints.openai.realtime.connection import RealtimeConnection as VllmRealtimeConnection
-from vllm.entrypoints.openai.realtime.protocol import TranscriptionDelta, TranscriptionDone
+from vllm.entrypoints.speech_to_text.realtime.connection import RealtimeConnection as VllmRealtimeConnection
+from vllm.entrypoints.speech_to_text.realtime.protocol import TranscriptionDelta, TranscriptionDone
 from vllm.logger import init_logger
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni

vllm_omni/entrypoints/openai/serving_chat.py

@@ -2574,6 +2574,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         output_format: str = "png",
         output_compression: int = 100,
         size: str = "auto",
+        raw_request: Request | None = None,
     ) -> tuple[list[Image.Image], dict[str, Any], float, str | None] | ErrorResponse | AsyncIterator[str]:
         """Generate diffusion images and return raw images plus generation stats."""
         if request_id is None:
@@ -2635,6 +2636,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     output_format=output_format,
                     output_compression=output_compression,
                     size=size,
+                    raw_request=raw_request,
                 )
 
             result = None
@@ -2659,6 +2661,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         peak_memory_mb = result.peak_memory_mb
         cot_output = None
 
+        req_out = getattr(result, "request_output", None)
+        if req_out:
+            prompt_obj = getattr(req_out, "prompt", None)
+            if isinstance(prompt_obj, dict):
+                extra = prompt_obj.get("extra", {})
+                if isinstance(extra, dict):
+                    ar_text = extra.get("ar_generated_text")
+                    if isinstance(ar_text, str) and ar_text.strip():
+                        cot_output = ar_text
+
         req_out = getattr(result, "request_output", None)
         if req_out:
             prompt_obj = getattr(req_out, "prompt", None)
@@ -2679,6 +2691,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         output_format: str,
         output_compression: int,
         size: str,
+        raw_request: Request | None = None,
     ) -> AsyncIterator[str]:
         """Yield image edit SSE chunks from multi-stage diffusion outputs."""
         created = int(time.time())
@@ -2730,7 +2743,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             logger.error("EngineDeadError during streaming image edit: %s", exc)
             data = self.create_streaming_error_response(exc)
             yield f"data: {data}\n\n"
-            terminate_if_errored(engine=self.engine_client)
+            if raw_request is not None:
+                terminate_if_errored(
+                    server=raw_request.app.state.server,
+                    engine=self.engine_client,
+                )
+            else:
+                logger.warning(
+                    "[OmniOpenAIServingChat] Engine dead during streaming image edit, "
+                    "but cannot terminate server (no raw_request context)."
+                )
         except Exception as exc:
             logger.exception("Streaming image edit failed: %s", exc)
             chunk = ImageEditStreamError(

vllm_omni/entrypoints/stage_utils.py

@@ -114,6 +114,24 @@ def _map_device_list(stage_id: int, device_list: list[str], visible_device_list:
         return list(device_list)
 
     logical_ids = [int(device) for device in device_list]
+
+    # If ALL logical device IDs are >= num_visible, none can be 0-based indices
+    # into the visible list — they must be physical device IDs.  This happens
+    # during multi-replica stage initialization where split_devices_for_replicas
+    # assigns physical device IDs (e.g. "2") that are not part of the current
+    # process's CUDA_VISIBLE_DEVICES but are valid physical GPUs.  Pass them
+    # through directly so the replica subprocess inherits the correct device.
+    if logical_ids and all(d >= num_visible for d in logical_ids):
+        logger.warning(
+            "Stage %s has device IDs %s, none of which are < the visible device count %d "
+            "(visible=%s). Treating them as physical device IDs and passing through.",
+            stage_id,
+            device_list,
+            num_visible,
+            visible_device_list,
+        )
+        return list(device_list)
+
     mapped_devices = [visible_device_list[idx] for idx in logical_ids if idx < num_visible]
     mapping_pairs = [
         f"{logical_id}->{visible_device_list[logical_id]}" for logical_id in logical_ids if logical_id < num_visible

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import os
 from collections.abc import Iterable, Mapping, Sequence
-from dataclasses import replace
 from functools import partial
 from threading import Lock
 
@@ -365,6 +364,12 @@ class CosyVoice3Model(
         else:
             raise ValueError(f"Model stage not supported {self.model_stage}")
 
+    def get_language_model(self) -> "nn.Module":
+        """Return the language model for upstream MoE detection."""
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     def _create_llm_vllm_config(self, parent_config: VllmConfig) -> VllmConfig:
         """Create VllmConfig for the inner Qwen2 LLM.
 
@@ -542,17 +547,21 @@ class CosyVoice3Model(
         if self.model_stage != "cosyvoice3_talker":
             return None
 
-        sampler = getattr(self, "_talker_sampler", None)
-        if sampler is None:
-            sampler = Sampler()
-            self._talker_sampler = sampler
-
         if not self._cosyvoice3_ras_enabled(sampling_metadata):
+            sampler = getattr(self, "_talker_sampler", None)
+            if sampler is None:
+                sampler = Sampler()
+                self._talker_sampler = sampler
             return sampler(logits=logits, sampling_metadata=sampling_metadata)
 
         logits = logits.to(torch.float32)
-        sampling_for_processors = replace(sampling_metadata, no_penalties=True)
-        logits = sampler.apply_logits_processors(logits, sampling_for_processors, predict_bonus_token=False)
+        # Apply logits processors directly — RAS handles its own repetition
+        # logic.  We avoid instantiating Sampler() here because its import
+        # chain pulls in flashinfer / GPU deps that fail in CPU-only tests.
+        if sampling_metadata.allowed_token_ids_mask is not None:
+            logits.masked_fill_(sampling_metadata.allowed_token_ids_mask, float("-inf"))
+        for processor in sampling_metadata.logitsprocs.non_argmax_invariant:
+            logits = processor.apply(logits)
 
         sampling_cfg = dict(self.config.llm.get("sampling", {}))
         default_top_p = float(sampling_cfg.get("top_p", 0.8))

vllm_omni/model_executor/models/covo_audio/covo_audio.py

@@ -243,6 +243,12 @@ class CovoAudioForConditionalGeneration(
             else lambda *args, **kwargs: None
         )
 
+    def get_language_model(self) -> "nn.Module":
+        """Return the language model for upstream MoE detection."""
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     def fused_thinker_talker_preprocess(
         self,
         input_ids: torch.Tensor | None,

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py

@@ -143,6 +143,12 @@ class MingFlashOmniForConditionalGeneration(
                 f"For image generation use stage_type: diffusion with model_arch: MingImagePipeline."
             )
 
+    def get_language_model(self) -> "nn.Module":
+        """Return the language model for upstream MoE detection."""
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     def forward(
         self,
         input_ids: torch.Tensor,

vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py

@@ -400,6 +400,13 @@ class MossTTSNanoForGeneration(nn.Module):
         runtime_additional_information: list[dict[str, Any]] | None = None,
         **kwargs: Any,
     ) -> OmniOutput:
+        if self._lm is None or self._audio_tokenizer is None:
+            # load_weights() loads the real LM from HuggingFace.  When
+            # load_format=dummy the vllm DummyModelLoader skips calling
+            # model.load_weights(), so we do it lazily on first forward.
+            # Must happen before _make_dummy_hidden so self._device is set.
+            self.load_weights([])
+
         sr = getattr(self.config, "audio_tokenizer_sample_rate", 48000)
         sr_tensor = torch.tensor(sr, dtype=torch.int32)
         empty = torch.zeros((0,), dtype=torch.float32)
@@ -418,9 +425,6 @@ class MossTTSNanoForGeneration(nn.Module):
                 },
             )
 
-        if self._lm is None or self._audio_tokenizer is None:
-            raise RuntimeError("MOSS-TTS-Nano model not loaded.  Was load_weights() called?")
-
         outputs: list[torch.Tensor] = []
         srs: list[torch.Tensor] = []
         last_chunk_flags: list[bool] = []

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py

@@ -91,18 +91,17 @@ class Qwen2_5OmniForConditionalGeneration(
 
         elif self.model_stage == "talker":
             multimodal_config.skip_mm_profiling = True
-            # register the process function for the talker stage
             self.has_preprocess = True
             self.set_custom_preprocess(self.talker_preprocess)
             self.thinker = None
             # Initialize talker model wrapper (handles projection + LM)
-            self.talker = init_vllm_registered_model(
-                vllm_config=vllm_config,
-                prefix=maybe_prefix(prefix, "talker"),
-                hf_config=talker_config,
-                # Use registry architecture key
-                architectures=["Qwen2_5OmniTalkerModel"],
-            )
+            with self._mark_language_model(vllm_config=vllm_config):
+                self.talker = init_vllm_registered_model(
+                    vllm_config=vllm_config,
+                    prefix=maybe_prefix(prefix, "talker"),
+                    hf_config=talker_config,
+                    architectures=["Qwen2_5OmniTalkerModel"],
+                )
             self.model = self.talker
             self.token2wav = None
             # set suppress start id according to token2wav
@@ -130,12 +129,13 @@ class Qwen2_5OmniForConditionalGeneration(
             self.token2wav_config = getattr(config, "token2wav_config", None)
             self.token2wav = None
             if self.token2wav_config is not None:
-                self.token2wav = init_vllm_registered_model(
-                    vllm_config=vllm_config,
-                    prefix=maybe_prefix(prefix, "token2wav"),
-                    hf_config=self.token2wav_config,
-                    architectures=["Qwen2_5OmniToken2WavModel"],
-                )
+                with self._mark_language_model(vllm_config=vllm_config):
+                    self.token2wav = init_vllm_registered_model(
+                        vllm_config=vllm_config,
+                        prefix=maybe_prefix(prefix, "token2wav"),
+                        hf_config=self.token2wav_config,
+                        architectures=["Qwen2_5OmniToken2WavModel"],
+                    )
             # voice resources (loaded on demand)
             self._token2wav_conds: dict[str, torch.Tensor] = {}
             self._token2wav_ref_mels: dict[str, torch.Tensor] = {}
@@ -149,6 +149,18 @@ class Qwen2_5OmniForConditionalGeneration(
             (self.thinker.make_empty_intermediate_tensors) if self.model_stage == "thinker" else lambda: None
         )
 
+    def get_language_model(self) -> nn.Module:
+        """Return the language model for MoE detection in upstream load_model.
+
+        Upstream GPUModelRunner.load_model() calls get_language_model() to find
+        the MoE model inside VLM wrappers. For thinker (which IS MoE), delegate
+        to the child. For talker/code2wav (not MoE, not SupportsMultiModal),
+        return self.model directly to avoid NotImplementedError.
+        """
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     # -------------------- Device utilities --------------------
     @staticmethod
     def _module_device(module: nn.Module) -> torch.device:

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -278,6 +278,12 @@ class Qwen3OmniMoeForConditionalGeneration(
             return self.model.sampler
         return Sampler()
 
+    def get_language_model(self) -> torch.nn.Module:
+        """Delegate to the active stage's language model for upstream MoE resolution."""
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     def embed_input_ids(
         self,
         input_ids: torch.Tensor,

vllm_omni/model_executor/models/registry.py

@@ -2,6 +2,7 @@ from vllm.model_executor.models.registry import (
     _VLLM_MODELS,
     _LazyRegisteredModel,
     _ModelRegistry,
+    _resolve_module_name,
 )
 
 _OMNI_MODELS = {
@@ -261,7 +262,7 @@ OmniModelRegistry = _ModelRegistry(
     {
         **{
             model_arch: _LazyRegisteredModel(
-                module_name=f"vllm.model_executor.models.{mod_relname}",
+                module_name=_resolve_module_name(mod_relname),
                 class_name=cls_name,
             )
             for model_arch, (mod_relname, cls_name) in _VLLM_MODELS.items()

vllm_omni/model_executor/models/voxtral_tts/voxtral_tts.py

@@ -159,6 +159,13 @@ class VoxtralTTSForConditionalGeneration(
             raise ValueError("Invalid model stage")
 
     # -------------------- CUDA Graph for acoustic transformer --------------------
+
+    def get_language_model(self) -> "nn.Module":
+        """Return the language model for upstream MoE detection."""
+        if hasattr(self.model, "get_language_model"):
+            return self.model.get_language_model()
+        return self.model
+
     def _enable_acoustic_transformer_cudagraph(self):
         """Initialize and capture CUDA graphs for compute_mm_logits."""
         if self.model_stage != "audio_generation" or not hasattr(self, "_vllm_config"):

vllm_omni/patch.py

@@ -128,3 +128,25 @@ for module_name, module in list(sys.modules.items()):
         module.StreamingUpdate = OmniStreamingUpdate
     if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest:
         module.EngineCoreRequest = OmniEngineCoreRequest
+
+
+# Patch: add qwen3_omni_moe to vllm's chat template fallback registry.
+# Qwen/Qwen3-Omni-30B-A3B-Instruct stores its chat_template in a standalone
+# chat_template.json (not in tokenizer_config.json).  transformers < 5.9.0
+# does not load this file, so the tokenizer has no chat_template attribute.
+# vllm's resolve_chat_template falls back to MODEL_TYPE_TO_CHAT_TEMPLATE
+# which has "qwen" but not "qwen3_omni_moe".  Register the same fallback.
+def _patch_chat_template_registry():
+    try:
+        from vllm.transformers_utils.chat_templates.registry import (
+            _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK,
+            _get_qwen_chat_template_fallback,
+        )
+
+        if "qwen3_omni_moe" not in _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK:
+            _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = _get_qwen_chat_template_fallback
+    except ImportError:
+        pass
+
+
+_patch_chat_template_registry()

vllm_omni/platforms/cuda/platform.py

@@ -229,4 +229,4 @@ class CudaOmniPlatform(OmniPlatform, CudaPlatformBase):
         if envs.VLLM_USE_OINK_OPS:
             rms_norm = ["oink"] + default
 
-        return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)
+        return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm)

vllm_omni/platforms/interface.py

@@ -11,6 +11,7 @@ from vllm.config import CUDAGraphMode, VllmConfig
 from vllm.forward_context import BatchDescriptor
 from vllm.logger import init_logger
 from vllm.platforms import Platform
+from vllm.platforms.interface import PlatformEnum
 
 logger = init_logger(__name__)
 
@@ -239,6 +240,7 @@ class OmniPlatform(Platform):
 
 class UnspecifiedOmniPlatform(OmniPlatform):
     _omni_enum = OmniPlatformEnum.UNSPECIFIED
+    _enum = PlatformEnum.UNSPECIFIED
     device_type = ""
 
     @classmethod

vllm_omni/platforms/npu/worker/npu_ar_model_runner.py

@@ -16,11 +16,6 @@ from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_
 from vllm.distributed.parallel_state import get_pp_group, get_tp_group
 from vllm.forward_context import BatchDescriptor
 from vllm.logger import logger
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-    extract_routed_experts_for_current_batch,
-    get_global_experts_capturer,
-    issue_routing_d2h_copy,
-)
 from vllm.sequence import IntermediateTensors
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import (
@@ -257,7 +252,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
         intermediate_tensors: IntermediateTensors | None = None,
     ) -> OmniModelRunnerOutput | IntermediateTensors | None:
         if self.vllm_config.model_config.enable_return_routed_experts:
-            capturer = get_global_experts_capturer()
+            capturer = self.routed_experts_capturer
             if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
                 capturer.finalize_pending_copy()
         if self.ascend_config.profiling_chunk_config.enabled:
@@ -728,12 +723,7 @@ class NPUARModelRunner(OmniNPUModelRunner):
             deferred_state_corrections_fn()
 
         if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"):
-            issue_routing_d2h_copy(
-                input_batch_req_ids=self.input_batch.req_ids,
-                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
-                positions=self.positions,
-                positions_cpu=self._positions_cpu,
-            )
+            self._omni_routed_experts_d2h(scheduler_output)
 
         return None
 
@@ -911,17 +901,11 @@ class NPUARModelRunner(OmniNPUModelRunner):
             if self.speculative_config is not None:
                 self.finalize_kv_connector()
 
-        routed_experts_dict = None
+        routed_experts_lists = None
         if self.model_config.enable_return_routed_experts:
-            capturer = get_global_experts_capturer()
+            capturer = self.routed_experts_capturer
             if capturer is not None and hasattr(self.input_batch, "num_tokens_no_spec"):
-                routed_experts_dict = extract_routed_experts_for_current_batch(
-                    req_ids=req_ids_output_copy,
-                    requests=self.requests,
-                    req_id_to_index=self.input_batch.req_id_to_index,
-                    num_tokens_no_spec=self.input_batch.num_tokens_no_spec,
-                    max_model_len=self.max_model_len,
-                )
+                routed_experts_lists = self._omni_extract_routed_experts(scheduler_output)
 
         #  -------------------------------------- Omni-new -------------------------------------------------
         engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy)
@@ -1042,9 +1026,9 @@ class NPUARModelRunner(OmniNPUModelRunner):
             kv_connector_output=kv_connector_output,
             ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
             cudagraph_stats=cudagraph_stats,
-            routed_experts_dict=routed_experts_dict,
         )
         model_runner_output.kv_extracted_req_ids = kv_extracted_req_ids
+        model_runner_output.routed_experts = routed_experts_lists
         #  -------------------------------------- Omni-new -------------------------------------------------
 
         if self.ascend_config.profiling_chunk_config.enabled and hasattr(self, "_execution_start_time"):

vllm_omni/platforms/npu/worker/npu_generation_model_runner.py

@@ -15,11 +15,6 @@ from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
 from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.logger import logger
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-    extract_routed_experts_for_current_batch,
-    get_global_experts_capturer,
-    issue_routing_d2h_copy,
-)
 from vllm.sequence import IntermediateTensors
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
@@ -78,7 +73,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
         intermediate_tensors: IntermediateTensors | None = None,
     ) -> OmniModelRunnerOutput | IntermediateTensors | None:
         if self.vllm_config.model_config.enable_return_routed_experts:
-            capturer = get_global_experts_capturer()
+            capturer = self.routed_experts_capturer
             if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
                 capturer.finalize_pending_copy()
         if self.ascend_config.profiling_chunk_config.enabled:
@@ -387,12 +382,7 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
             deferred_state_corrections_fn()
 
         if self.vllm_config.model_config.enable_return_routed_experts and hasattr(self, "_positions_cpu"):
-            issue_routing_d2h_copy(
-                input_batch_req_ids=self.input_batch.req_ids,
-                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
-                positions=self.positions,
-                positions_cpu=self._positions_cpu,
-            )
+            self._omni_routed_experts_d2h(scheduler_output)
 
         return None
 
@@ -479,17 +469,11 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
         # [Omni] Copy req_id mappings to avoid async scheduling mutation.
         req_ids_output_copy = self.input_batch.req_ids.copy()
         req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
-        routed_experts_dict = None
+        routed_experts_lists = None
         if self.vllm_config.model_config.enable_return_routed_experts and hasattr(
             self.input_batch, "num_tokens_no_spec"
         ):
-            routed_experts_dict = extract_routed_experts_for_current_batch(
-                req_ids=req_ids_output_copy,
-                requests=self.requests,
-                req_id_to_index=self.input_batch.req_id_to_index,
-                num_tokens_no_spec=self.input_batch.num_tokens_no_spec,
-                max_model_len=self.max_model_len,
-            )
+            routed_experts_lists = self._omni_extract_routed_experts(scheduler_output)
         output = OmniModelRunnerOutput(
             req_ids=req_ids_output_copy,
             req_id_to_index=req_id_to_index_output_copy,
@@ -500,8 +484,8 @@ class NPUGenerationModelRunner(OmniNPUModelRunner):
             kv_connector_output=kv_connector_output,
             ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
             cudagraph_stats=cudagraph_stats,
-            routed_experts_dict=routed_experts_dict,
         )
+        output.routed_experts = routed_experts_lists
         #  -------------------------------------- Omni-new -------------------------------------------------
 
         if self.speculative_config is not None:

vllm_omni/platforms/rocm/platform.py

@@ -172,4 +172,4 @@ class RocmOmniPlatform(OmniPlatform, RocmPlatform):
         else:
             rms_norm = default
 
-        return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm)
+        return IrOpPriorityConfig.with_default(default, rms_norm=rms_norm, fused_add_rms_norm=rms_norm)

vllm_omni/quantization/factory.py

@@ -10,19 +10,73 @@ different QuantizationConfig subclass in the OMNI context (e.g. GGUF).
 from __future__ import annotations
 
 import inspect
+import sys
 from collections.abc import Callable, Mapping
+from types import ModuleType
 from typing import Any
 
 from vllm.logger import init_logger
-from vllm.model_executor.layers.quantization import (
+
+
+# ---------------------------------------------------------------------------
+# Stub the ``humming`` package so that vLLM's lazy import inside
+# ``get_quantization_config()`` (which unconditionally does
+# ``from .humming import HummingConfig``) does not crash when the real
+# ``humming`` wheel is not installed.  Only populate the bare-minimum
+# names that ``humming.py`` accesses at module level.
+# ---------------------------------------------------------------------------
+def _register_humming_stubs() -> None:
+    """Register stub ``humming`` sub-modules so that the optional
+    humming quantization backend can be imported without the real wheel."""
+    if "humming" in sys.modules:
+        return  # already present (real or stub)
+
+    # --- sub-modules ---
+    submodules: dict[str, tuple[str, ...]] = {
+        "humming": (),
+        "humming.config": ("GemmType",),
+        "humming.dtypes": ("DataType",),
+        "humming.layer": ("HummingLayerMeta", "HummingMethod"),
+        "humming.schema": (
+            "BaseInputSchema",
+            "BaseWeightSchema",
+            "HummingInputSchema",
+            "HummingWeightSchema",
+        ),
+        "humming.utils": (),
+        "humming.utils.weight": ("quantize_weight",),
+    }
+
+    registry: dict[str, ModuleType] = {}
+    for name, attrs in submodules.items():
+        mod = ModuleType(name)
+        for attr in attrs:
+            setattr(mod, attr, type(attr, (), {}))
+        registry[name] = mod
+
+    # wire parent references
+    registry["humming"].config = registry["humming.config"]
+    registry["humming"].dtypes = registry["humming.dtypes"]
+    registry["humming"].layer = registry["humming.layer"]
+    registry["humming"].schema = registry["humming.schema"]
+    registry["humming"].utils = registry["humming.utils"]
+    registry["humming.utils"].weight = registry["humming.utils.weight"]
+
+    for name, mod in registry.items():
+        sys.modules[name] = mod
+
+
+_register_humming_stubs()
+
+from vllm.model_executor.layers.quantization import (  # noqa: E402
     QUANTIZATION_METHODS,
     get_quantization_config,
 )
-from vllm.model_executor.layers.quantization.base_config import (
+from vllm.model_executor.layers.quantization.base_config import (  # noqa: E402
     QuantizationConfig,
 )
 
-from .component_config import ComponentQuantizationConfig
+from .component_config import ComponentQuantizationConfig  # noqa: E402
 
 logger = init_logger(__name__)
 

vllm_omni/worker/gpu_ar_model_runner.py

@@ -19,11 +19,6 @@ from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_
 from vllm.distributed.parallel_state import get_pp_group, get_tp_group
 from vllm.forward_context import set_forward_context
 from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-    extract_routed_experts_for_current_batch,
-    get_global_experts_capturer,
-    issue_routing_d2h_copy,
-)
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
 from vllm.v1.spec_decode.dflash import DFlashProposer
@@ -363,6 +358,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         if self.execute_model_state is not None:
             raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
 
+        if self.routed_experts_initialized:
+            self.routed_experts_capturer.clear_buffer()
+
         if not getattr(self, "_warmup_state_cleared", False):
             self._warmup_state_cleared = True
             if hasattr(self.model, "_clear_warmup_state"):
@@ -414,7 +412,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             )
 
         if self.routed_experts_initialized:
-            capturer = get_global_experts_capturer()
+            capturer = self.routed_experts_capturer
             if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
                 capturer.finalize_pending_copy()
 
@@ -679,7 +677,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 if not get_pp_group().is_last_rank:
                     # Return the intermediate tensors.
                     assert isinstance(hidden_states, IntermediateTensors)
-                    hidden_states.kv_connector_output = kv_connector_output
                     self.kv_connector_output = kv_connector_output
                     return hidden_states
 
@@ -692,7 +689,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                         kv_connector_output,
                     )
 
-                sample_hidden_states = hidden_states[logits_indices]
+                sample_hidden_states = hidden_states[logits_indices.to(hidden_states.device)]
                 # Try with sampling_metadata first; fall back to without for models that don't support it
                 try:
                     logits = self.model.compute_logits(
@@ -704,7 +701,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 # Rare case.
                 assert not self.is_pooling_model
 
-                sample_hidden_states = hidden_states[logits_indices]
+                sample_hidden_states = hidden_states[logits_indices.to(hidden_states.device)]
                 if not get_pp_group().is_last_rank:
                     all_gather_tensors = {
                         "residual": not is_residual_scattered_for_sp(self.vllm_config, num_tokens_padded)
@@ -753,12 +750,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             deferred_state_corrections_fn()
 
         if self.routed_experts_initialized and hasattr(self, "_positions_cpu"):
-            issue_routing_d2h_copy(
-                input_batch_req_ids=self.input_batch.req_ids,
-                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
-                positions=self.positions,
-                positions_cpu=self._positions_cpu,
-            )
+            self._omni_routed_experts_d2h(scheduler_output)
 
         return None
 
@@ -830,18 +822,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         if self.execute_model_state is None:
             kv_connector_output = self.kv_connector_output
             self.kv_connector_output = None
-            # Nothing to do (PP non-final rank case), output isn't used.
-            if not kv_connector_output:
-                return None  # type: ignore[return-value]
-
+            # receive sampled token ids from the last PP rank.
+            if self.use_async_scheduling and not get_pp_group().is_last_rank:
+                self._pp_receive_prev_sampled_token_ids_to_input_batch()
             # In case of PP with kv transfer, we need to pass through the
             # kv_connector_output
-            if kv_connector_output.is_empty():
-                return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
-
-            output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
-            output.kv_connector_output = kv_connector_output
-            return self.attach_omni_connector_output(output)
+            return self.attach_omni_connector_output(
+                OmniModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
+            )
 
         # Unpack ephemeral state.
         (
@@ -1108,15 +1096,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     self.accumulate_full_payload_output(rid, pooler_output[i], req_state)
 
         with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
-            routed_experts_dict = None
+            routed_experts_lists = None
             if self.routed_experts_initialized:
-                routed_experts_dict = extract_routed_experts_for_current_batch(
-                    req_ids=req_ids_output_copy,
-                    requests=self.requests,
-                    req_id_to_index=self.input_batch.req_id_to_index,
-                    num_tokens_no_spec=self.input_batch.num_tokens_no_spec,
-                    max_model_len=self.max_model_len,
-                )
+                routed_experts_lists = self._omni_extract_routed_experts(scheduler_output)
             output = OmniModelRunnerOutput(
                 req_ids=req_ids_output_copy,
                 req_id_to_index=req_id_to_index_output_copy,
@@ -1128,10 +1110,10 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
                 num_nans_in_logits=num_nans_in_logits,
                 cudagraph_stats=cudagraph_stats,
-                routed_experts_dict=routed_experts_dict,
             )
             output.kv_extracted_req_ids = kv_extracted_req_ids
             output.omni_connector_output = self.get_omni_connector_output()
+            output.routed_experts = routed_experts_lists
 
         if not self.use_async_scheduling:
             return output

vllm_omni/worker/gpu_generation_model_runner.py

@@ -8,7 +8,6 @@ from __future__ import annotations
 
 import gc
 import logging
-from copy import copy
 from dataclasses import replace
 
 import numpy as np
@@ -18,11 +17,6 @@ from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer
 from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.forward_context import set_forward_context
-from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
-    extract_routed_experts_for_current_batch,
-    get_global_experts_capturer,
-    issue_routing_d2h_copy,
-)
 from vllm.utils.math_utils import cdiv
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
@@ -112,6 +106,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         if self.execute_model_state is not None:
             raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
 
+        if self.routed_experts_initialized:
+            self.routed_experts_capturer.clear_buffer()
+
         if hasattr(self, "_omni_connector"):
             for request in getattr(scheduler_output, "pending_input_registrations", []):
                 self.register_chunk_recv(request)
@@ -123,7 +120,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                     self.flush_full_payload_outputs(flush_ids)
 
         if self.routed_experts_initialized:
-            capturer = get_global_experts_capturer()
+            capturer = self.routed_experts_capturer
             if capturer is not None and hasattr(capturer, "finalize_pending_copy"):
                 capturer.finalize_pending_copy()
 
@@ -369,12 +366,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             deferred_state_corrections_fn()
 
         if self.routed_experts_initialized and hasattr(self, "_positions_cpu"):
-            issue_routing_d2h_copy(
-                input_batch_req_ids=self.input_batch.req_ids,
-                num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
-                positions=self.positions,
-                positions_cpu=self._positions_cpu,
-            )
+            self._omni_routed_experts_d2h(scheduler_output)
 
         return None
 
@@ -390,18 +382,14 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         self.kv_connector_output = None
 
         if self.execute_model_state is None:
-            # Nothing to do (PP non-final rank case), output isn't used.
-            if not kv_connector_output:
-                return None  # type: ignore[return-value]
-
+            # receive sampled token ids from the last PP rank.
+            if self.use_async_scheduling and not get_pp_group().is_last_rank:
+                self._pp_receive_prev_sampled_token_ids_to_input_batch()
             # In case of PP with kv transfer, we need to pass through the
             # kv_connector_output
-            if kv_connector_output.is_empty():
-                return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
-
-            output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
-            output.kv_connector_output = kv_connector_output
-            return self.attach_omni_connector_output(output)
+            return self.attach_omni_connector_output(
+                OmniModelRunnerOutput.with_kv_conn_output_only(kv_connector_output)
+            )
 
         # Unpack ephemeral state.
         (
@@ -462,15 +450,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         # [Omni] Copy req_id mappings to avoid async scheduling mutation.
         req_ids_output_copy = self.input_batch.req_ids.copy()
         req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
-        routed_experts_dict = None
+        routed_experts_lists = None
         if self.routed_experts_initialized:
-            routed_experts_dict = extract_routed_experts_for_current_batch(
-                req_ids=req_ids_output_copy,
-                requests=self.requests,
-                req_id_to_index=self.input_batch.req_id_to_index,
-                num_tokens_no_spec=self.input_batch.num_tokens_no_spec,
-                max_model_len=self.max_model_len,
-            )
+            routed_experts_lists = self._omni_extract_routed_experts(scheduler_output)
         if self._should_accumulate_full_payload_output():
             for i, rid in enumerate(req_ids_output_copy):
                 req_state = self.requests.get(rid)
@@ -488,9 +470,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             num_nans_in_logits={},
             cudagraph_stats=cudagraph_stats,
             ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
-            routed_experts_dict=routed_experts_dict,
         )
         output.omni_connector_output = self.get_omni_connector_output()
+        output.routed_experts = routed_experts_lists
 
         if not self.use_async_scheduling:
             return output
@@ -544,9 +526,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         )
 
     @torch.inference_mode()
-    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
+    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> torch.Tensor:
         logger.warning("Dummy sampler run is not implemented for generation model")
-        return None
+        return torch.tensor([])
 
     @torch.inference_mode()
     def _dummy_run(

vllm_omni/worker/gpu_model_runner.py

@@ -36,6 +36,7 @@ from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
+    from vllm.v1.outputs import RoutedExpertsLists
 else:
     xgr = LazyLoader("xgr", globals(), "xgrammar")
     xgr_torch_compile = LazyLoader(
@@ -57,6 +58,52 @@ class OmniGPUModelRunner(GPUModelRunner):
         # when we initialize the metadata builders if enabled
         self.omni_prefix_cache = None
 
+    def _omni_routed_experts_d2h(self, scheduler_output) -> None:
+        """Issue routed-experts D2H copy matching upstream GPUModelRunner pattern.
+
+        Upstream does this inline in ``execute_model``:
+            buf = self.routed_experts_capturer.get_device_buffer()
+            total = scheduler_output.total_num_scheduled_tokens
+            self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True)
+            self.routed_experts_slot_mapping_cpu[:total].copy_(
+                self.routed_experts_slot_mapping_device[:total], non_blocking=True)
+        """
+        if not self.routed_experts_initialized:
+            return
+        buf = self.routed_experts_capturer.get_device_buffer()
+        total = scheduler_output.total_num_scheduled_tokens
+        self.routed_experts_cpu[:total].copy_(buf[:total], non_blocking=True)
+        if hasattr(self, "routed_experts_slot_mapping_device"):
+            self.routed_experts_slot_mapping_cpu[:total].copy_(
+                self.routed_experts_slot_mapping_device[:total],
+                non_blocking=True,
+            )
+
+    def _omni_extract_routed_experts(self, scheduler_output) -> "RoutedExpertsLists | None":
+        """Extract routed experts matching upstream GPUModelRunner pattern.
+
+        Upstream (sync path, sample_tokens):
+            total = scheduler_output.total_num_scheduled_tokens
+            output.routed_experts = RoutedExpertsLists(
+                routing_data=self.routed_experts_cpu[:total].numpy(),
+                slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(),
+            )
+
+        Returns RoutedExpertsLists (batch-level, with slot_mapping) so that
+        downstream schedulers can use slot_mapping to map back to requests.
+        """
+        from vllm.v1.outputs import RoutedExpertsLists
+
+        if not self.routed_experts_initialized:
+            return None
+        total = scheduler_output.total_num_scheduled_tokens
+        if total <= 0:
+            return None
+        return RoutedExpertsLists(
+            routing_data=self.routed_experts_cpu[:total].numpy(),
+            slot_mapping=self.routed_experts_slot_mapping_cpu[:total].numpy(),
+        )
+
     def initialize_metadata_builders(self, kv_cache_config, kernel_block_sizes):
         """Override to fix scheduler_metadata buffer size for FA3 + CUDA graph.
 
@@ -747,6 +794,15 @@ class OmniGPUModelRunner(GPUModelRunner):
             raise ValueError(f"Invalid hidden states type: {type(hidden_states)}")
         return text_hidden_states, multimodal_outputs
 
+    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # Models loaded with load_format=dummy (e.g. MossTTSNano) may
+        # produce CPU hidden_states while the upstream sampler warmup
+        # expects GPU tensors.  Skip the warmup in that case — a
+        # meaningful warmup requires real weights anyway.
+        if hidden_states.device != self.device:
+            return torch.tensor([])
+        return super()._dummy_sampler_run(hidden_states=hidden_states)
+
     @torch.inference_mode()
     def _dummy_run(
         self,
@@ -1100,7 +1156,7 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.eplb_step(is_dummy=True, is_profile=is_profile)
 
         logit_indices = np.cumsum(num_scheduled_tokens) - 1
-        logit_indices_device = torch.from_numpy(logit_indices).to(self.device, non_blocking=True)
+        logit_indices_device = torch.from_numpy(logit_indices).to(hidden_states.device, non_blocking=True)
         return hidden_states, hidden_states[logit_indices_device]
 
     # ------------------------------------------------------------------