Merge remote-tracking branch 'origin/main' into dev/vllm-align

benchmarks/diffusion/backends.py

@@ -0,0 +1,198 @@
+import base64
+import mimetypes
+import os
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any
+
+import aiohttp
+from tqdm import tqdm
+
+
+@dataclass
+class RequestFuncInput:
+    prompt: str
+    api_url: str
+    model: str
+    width: int | None = None
+    height: int | None = None
+    num_frames: int | None = None
+    num_inference_steps: int | None = None
+    seed: int | None = None
+    fps: int | None = None
+    timestamp: float | None = None
+    slo_ms: float | None = None
+    extra_body: dict[str, Any] = field(default_factory=dict)
+    image_paths: list[str] | None = None
+    request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+
+
+@dataclass
+class RequestFuncOutput:
+    success: bool = False
+    latency: float = 0.0
+    error: str = ""
+    start_time: float = 0.0
+    response_body: dict[str, Any] = field(default_factory=dict)
+    peak_memory_mb: float = 0.0
+    slo_achieved: bool | None = None
+
+
+def _guess_mime_type(path: str) -> str:
+    mime, _ = mimetypes.guess_type(path)
+    return mime or "application/octet-stream"
+
+
+def _encode_image_as_data_url(path: str) -> str:
+    with open(path, "rb") as f:
+        encoded = base64.b64encode(f.read()).decode("utf-8")
+    mime = _guess_mime_type(path)
+    return f"data:{mime};base64,{encoded}"
+
+
+async def async_request_chat_completions(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput:
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    extra_body = dict(input.extra_body)
+    if input.width and input.height:
+        extra_body.setdefault("height", input.height)
+        extra_body.setdefault("width", input.width)
+    if input.num_frames:
+        extra_body.setdefault("num_frames", input.num_frames)
+    if input.num_inference_steps:
+        extra_body.setdefault("num_inference_steps", input.num_inference_steps)
+    if input.seed is not None:
+        extra_body.setdefault("seed", input.seed)
+    if input.fps:
+        extra_body.setdefault("fps", input.fps)
+
+    if input.image_paths and len(input.image_paths) > 0:
+        content = []
+        if input.prompt:
+            content.append({"type": "text", "text": input.prompt})
+        for img_path in input.image_paths:
+            if not os.path.exists(img_path):
+                output.error = f"Image file not found: {img_path}"
+                output.success = False
+                if pbar:
+                    pbar.update(1)
+                return output
+            content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": _encode_image_as_data_url(img_path)},
+                }
+            )
+        messages = [{"role": "user", "content": content}]
+    else:
+        messages = [{"role": "user", "content": input.prompt}]
+
+    payload = {
+        "model": input.model,
+        "messages": messages,
+    }
+    if extra_body:
+        payload["extra_body"] = extra_body
+
+    try:
+        async with session.post(input.api_url, json=payload) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+                if "peak_memory_mb" in resp_json:
+                    output.peak_memory_mb = resp_json["peak_memory_mb"]
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+async def async_request_openai_images(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput:
+    """
+    Send request to OpenAI's /v1/images/generations endpoint.
+    """
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    # Build size string from width/height
+    width = input.width or 1024
+    height = input.height or 1024
+    size = f"{width}x{height}"
+
+    payload: dict[str, Any] = {
+        "model": input.model,
+        "prompt": input.prompt,
+        "n": 1,
+        "size": size,
+        "response_format": "b64_json",
+    }
+
+    # Add optional parameters
+    if input.seed is not None:
+        payload["seed"] = input.seed
+    if input.num_inference_steps is not None:
+        payload["num_inference_steps"] = input.num_inference_steps
+
+    # Add any extra body parameters
+    if input.extra_body:
+        for key, value in input.extra_body.items():
+            if key not in payload:
+                payload[key] = value
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": "Bearer EMPTY",
+    }
+
+    try:
+        async with session.post(input.api_url, json=payload, headers=headers) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+                # Check for usage/memory info if available
+                if "usage" in resp_json and "peak_memory_mb" in resp_json.get("usage", {}):
+                    output.peak_memory_mb = resp_json["usage"]["peak_memory_mb"]
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+backends_function_mapping = {
+    "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
+    "openai": (async_request_openai_images, "/v1/images/generations"),
+}

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -6,80 +6,60 @@
 Benchmark online serving for diffusion models (Image/Video Generation).
 If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
 
+Supports multiple backends:
+    - vllm-omni: Uses /v1/chat/completions endpoint (default)
+    - openai: Uses /v1/images/generations endpoint
+
 Usage:
-    # Video
+    # Video (vllm-omni backend)
     t2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --dataset vbench --task t2v --num-prompts 10 \
+        --backend vllm-omni --dataset vbench --task t2v --num-prompts 10 \
         --height 480 --width 640 --fps 16 --num-frames 80
 
     i2v:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --dataset vbench --task i2v --num-prompts 10
+        --backend vllm-omni --dataset vbench --task i2v --num-prompts 10
 
 
-    # Image
+    # Image (vllm-omni backend)
     t2i:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --dataset vbench --task t2i --num-prompts 10 \
+        --backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024
 
     i2i:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --dataset vbench --task i2i --num-prompts 10
+        --backend vllm-omni --dataset vbench --task i2i --num-prompts 10
+
+    # Image (openai backend)
+    t2i:
+    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+        --backend openai --dataset vbench --task t2i --num-prompts 10 \
+        --height 1024 --width 1024 --port 3000
 
 """
 
 import argparse
 import ast
 import asyncio
-import base64
 import glob
 import json
-import mimetypes
 import os
 import time
 import uuid
 from abc import ABC, abstractmethod
 from collections.abc import AsyncGenerator
-from dataclasses import dataclass, field, replace
+from dataclasses import replace
 from typing import Any
 
 import aiohttp
 import numpy as np
 import requests
+from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
 from tqdm.asyncio import tqdm
 
 
-@dataclass
-class RequestFuncInput:
-    prompt: str
-    api_url: str
-    model: str
-    width: int | None = None
-    height: int | None = None
-    num_frames: int | None = None
-    num_inference_steps: int | None = None
-    seed: int | None = None
-    fps: int | None = None
-    timestamp: float | None = None
-    slo_ms: float | None = None
-    extra_body: dict[str, Any] = field(default_factory=dict)
-    image_paths: list[str] | None = None
-    request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
-
-
-@dataclass
-class RequestFuncOutput:
-    success: bool = False
-    latency: float = 0.0
-    error: str = ""
-    start_time: float = 0.0
-    response_body: dict[str, Any] = field(default_factory=dict)
-    peak_memory_mb: float = 0.0
-    slo_achieved: bool | None = None
-
-
 class BaseDataset(ABC):
     def __init__(self, args, api_url: str, model: str):
         self.args = args
@@ -684,92 +664,6 @@ async def iter_requests(
         yield req
 
 
-def _guess_mime_type(path: str) -> str:
-    mime, _ = mimetypes.guess_type(path)
-    return mime or "application/octet-stream"
-
-
-def _encode_image_as_data_url(path: str) -> str:
-    with open(path, "rb") as f:
-        encoded = base64.b64encode(f.read()).decode("utf-8")
-    mime = _guess_mime_type(path)
-    return f"data:{mime};base64,{encoded}"
-
-
-async def async_request_chat_completions(
-    input: RequestFuncInput,
-    session: aiohttp.ClientSession,
-    pbar: tqdm | None = None,
-) -> RequestFuncOutput:
-    output = RequestFuncOutput()
-    output.start_time = time.perf_counter()
-
-    extra_body = dict(input.extra_body)
-    if input.width and input.height:
-        extra_body.setdefault("height", input.height)
-        extra_body.setdefault("width", input.width)
-    if input.num_frames:
-        extra_body.setdefault("num_frames", input.num_frames)
-    if input.num_inference_steps:
-        extra_body.setdefault("num_inference_steps", input.num_inference_steps)
-    if input.seed is not None:
-        extra_body.setdefault("seed", input.seed)
-    if input.fps:
-        extra_body.setdefault("fps", input.fps)
-
-    if input.image_paths and len(input.image_paths) > 0:
-        content = []
-        if input.prompt:
-            content.append({"type": "text", "text": input.prompt})
-        for img_path in input.image_paths:
-            if not os.path.exists(img_path):
-                output.error = f"Image file not found: {img_path}"
-                output.success = False
-                if pbar:
-                    pbar.update(1)
-                return output
-            content.append(
-                {
-                    "type": "image_url",
-                    "image_url": {"url": _encode_image_as_data_url(img_path)},
-                }
-            )
-        messages = [{"role": "user", "content": content}]
-    else:
-        messages = [{"role": "user", "content": input.prompt}]
-
-    payload = {
-        "model": input.model,
-        "messages": messages,
-    }
-    if extra_body:
-        payload["extra_body"] = extra_body
-
-    try:
-        async with session.post(input.api_url, json=payload) as response:
-            if response.status == 200:
-                resp_json = await response.json()
-                output.response_body = resp_json
-                output.success = True
-                if "peak_memory_mb" in resp_json:
-                    output.peak_memory_mb = resp_json["peak_memory_mb"]
-            else:
-                output.error = f"HTTP {response.status}: {await response.text()}"
-                output.success = False
-    except Exception as e:
-        output.error = str(e)
-        output.success = False
-
-    output.latency = time.perf_counter() - output.start_time
-
-    if output.success and input.slo_ms is not None:
-        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
-
-    if pbar:
-        pbar.update(1)
-    return output
-
-
 def calculate_metrics(
     outputs: list[RequestFuncOutput],
     total_duration: float,
@@ -848,9 +742,9 @@ async def benchmark(args):
     if args.base_url is None:
         args.base_url = f"http://{args.host}:{args.port}"
 
-    # Setup dataset (vLLM-Omni supports diffusion via /v1/chat/completions)
-    api_url = f"{args.base_url}/v1/chat/completions"
-    request_func = async_request_chat_completions
+    # Setup API URL and request function based on backend
+    request_func, api_url = backends_function_mapping[args.backend]
+    api_url = f"{args.base_url}{api_url}"
 
     if args.dataset == "vbench":
         dataset = VBenchDataset(args, api_url, args.model)
@@ -920,9 +814,16 @@ async def benchmark(args):
     # Calculate metrics
     metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)
 
+    # Add configuration info to metrics for JSON output
+    metrics["backend"] = args.backend
+    metrics["model"] = args.model
+    metrics["dataset"] = args.dataset
+    metrics["task"] = args.task
+
     print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))
 
     # Section 1: Configuration
+    print("{:<40} {:<15}".format("Backend:", args.backend))
     print("{:<40} {:<15}".format("Model:", args.model))
     print("{:<40} {:<15}".format("Dataset:", args.dataset))
     print("{:<40} {:<15}".format("Task:", args.task))
@@ -978,6 +879,13 @@ if __name__ == "__main__":
     parser.add_argument("--host", type=str, default="localhost", help="Server host.")
     parser.add_argument("--port", type=int, default=8091, help="Server port.")
     parser.add_argument("--model", type=str, default="default", help="Model name.")
+    parser.add_argument(
+        "--backend",
+        type=str,
+        default="vllm-omni",
+        choices=["vllm-omni", "openai"],
+        help="Backend to target the benchmark to.",
+    )
     parser.add_argument(
         "--dataset",
         type=str,

docs/.nav.yml

@@ -58,12 +58,6 @@ nav:
     - contributing/model/README.md
     - contributing/model/adding_omni_model.md
     - contributing/model/adding_diffusion_model.md
-    - Advanced Features:
-      - contributing/features/cfg_parallel.md
-      - contributing/features/sequence_parallel.md
-      - contributing/features/tensor_parallel.md
-      - contributing/features/cache_dit.md
-      - contributing/features/teacache.md
   - CI: contributing/ci
   - Design Documents:
     - design/index.md
@@ -72,6 +66,11 @@ nav:
       - design/feature/disaggregated_inference.md
       - design/feature/ray_based_execution.md
       - design/feature/omni_connectors/
+      - design/feature/cfg_parallel.md
+      - design/feature/sequence_parallel.md
+      - design/feature/tensor_parallel.md
+      - design/feature/cache_dit.md
+      - design/feature/teacache.md
     - Module Design:
       - design/module/ar_module.md
       - design/module/dit_module.md

docs/cli/serve.md

@@ -0,0 +1,9 @@
+# vllm-omni serve
+
+## JSON CLI Arguments
+
+--8<-- "docs/cli/json_tip.inc.md"
+
+## Arguments
+
+--8<-- "docs/generated/argparse_omni/omni_serve.inc.md"

docs/contributing/metrics.md

@@ -1,25 +1,27 @@
 
-# Metrics vLLM-Omni:
+# Metrics
 
 You can use these metrics in production to monitor the health and performance of the vLLM-omni system. Typical scenarios include:
+
 - **Performance Monitoring**: Track throughput (e.g., `e2e_avg_tokens_per_s`), latency (e.g., `e2e_total_ms`), and resource utilization to verify that the system meets expected standards.
+
 - **Debugging and Troubleshooting**: Use detailed per-request metrics to diagnose issues, such as high transfer times or unexpected token counts.
 
 ## How to Enable and View Metrics
 
-### 1. Start the Service with Metrics Logging
+### Start the Service with Metrics Logging
 
 ```bash
 vllm serve /workspace/models/Qwen3-Omni-30B-A3B-Instruct --omni --port 8014 --log-stats
 ```
 
-### 2. Send a Request
+### Send a Request
 
 ```bash
 python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
 ```
 
-### 3. What You Will See
+### What You Will See
 
 With `--log-stats` enabled, the server will output detailed metrics logs after each request. Example output:
 
@@ -69,9 +71,13 @@ With `--log-stats` enabled, the server will output detailed metrics logs after e
 
 
 These logs include:
+
 - **Overall summary**: total requests, wall time, average tokens/sec, etc.
+
 - **E2E table**: per-request latency and token counts.
+
 - **Stage table**: per-stage batch and timing details.
+
 - **Transfer table**: data transfer and timing for each edge.
 
 You can use these logs to monitor system health, debug performance, and analyze request-level metrics as described above.
@@ -87,6 +93,8 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
 
 ## Parameter Details
 
+### Summary Metrics
+
 | Field                     | Meaning                                                                                       |
 |---------------------------|----------------------------------------------------------------------------------------------|
 | `e2e_requests`            | Number of completed requests.                                                                |
@@ -98,7 +106,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
 
 ---
 
-## E2E Table (per request)
+### E2E Table (per request)
 
 | Field                     | Meaning                                                               |
 |---------------------------|-----------------------------------------------------------------------|
@@ -110,7 +118,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
 
 ---
 
-## Stage Table (per stage event / request)
+### Stage Table (per stage event / request)
 
 | Field                     | Meaning                                                                                         |
 |---------------------------|-------------------------------------------------------------------------------------------------|
@@ -125,7 +133,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
 
 ---
 
-## Transfer Table (per edge / request)
+### Transfer Table (per edge / request)
 
 | Field                | Meaning                                                                   |
 |----------------------|---------------------------------------------------------------------------|
@@ -135,31 +143,31 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
 | `in_flight_time_ms`  | In-flight time in ms.                                                     |
 
 
-## Expectation of the Numbers (Verification)
+### Expectation of the Numbers (Verification)
 
 **Formulas:**
+
 - `e2e_total_tokens = Stage0's num_tokens_in + sum(all stages' num_tokens_out)`
+
 - `transfers_total_time_ms = sum(tx_time_ms + rx_decode_time_ms + in_flight_time_ms)` for every edge
 
 **Using the example above:**
 
-### e2e_total_tokens
+**e2e_total_tokens**
+
 - Stage0's `num_tokens_in`: **4,860**
 - Stage0's `num_tokens_out`: **67**
 - Stage1's `num_tokens_out`: **275**
 - Stage2's `num_tokens_out`: **0**
 
-So,
-```
-e2e_total_tokens = 4,860 + 67 + 275 + 0 = 5,202
-```
-This matches the table value: `e2e_total_tokens = 5,202`.
+so `e2e_total_tokens = 4,860 + 67 + 275 + 0 = 5,202`, which matches the table value `e2e_total_tokens`.
+
+**transfers_total_time_ms**
 
-### transfers_total_time_ms
 For each edge:
+
 - 0->1: tx_time_ms (**78.701**) + rx_decode_time_ms (**111.865**) + in_flight_time_ms (**2.015**) = **192.581**
+
 - 1->2: tx_time_ms (**18.790**) + rx_decode_time_ms (**31.706**) + in_flight_time_ms (**2.819**) = **53.315**
 
-Sum: 192.581 + 53.315 = **245.896**
-
-The table shows `transfers_total_time_ms = 245.895`, which matches the calculation (difference is due to rounding).
+192.581 + 53.315 = **245.896** = transfers_total_time_ms, which matches the calculation (difference is due to rounding)

docs/contributing/model/adding_diffusion_model.md

@@ -1,4 +1,4 @@
-# Adding a Diffusion Model to vLLM-Omni
+# Adding a Diffusion Model
 
 This guide walks you through adding a new diffusion model to vLLM-Omni. We use **Qwen-Image** as the primary example, with references to other models (LongCat, Flux, Wan2.2) to illustrate different patterns.
 
@@ -680,7 +680,7 @@ vLLM-Omni automatically compiles blocks in `_repeated_blocks` when `torch.compil
 
 ### Tensor Parallelism
 
-See detailed guide: [How to add Tensor Parallel support](../features/tensor_parallel.md)
+See detailed guide: [How to add Tensor Parallel support](../../design/feature/tensor_parallel.md)
 
 **Quick setup:**
 
@@ -694,7 +694,7 @@ omni = Omni(model="your-model", tensor_parallel_size=2)
 
 ### CFG Parallelism
 
-See detailed guide: [How to add CFG-Parallel support](../features/cfg_parallel.md)
+See detailed guide: [How to add CFG-Parallel support](../../design/feature/cfg_parallel.md)
 
 **Quick setup:**
 
@@ -708,7 +708,7 @@ omni = Omni(model="your-model", cfg_parallel_size=2)
 
 ### Sequence Parallelism
 
-See detailed guide: [How to add Sequence Parallel support](../features/sequence_parallel.md)
+See detailed guide: [How to add Sequence Parallel support](../../design/feature/sequence_parallel.md)
 
 **Quick setup:**
 
@@ -724,7 +724,7 @@ omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
 
 #### TeaCache
 
-See detailed guide: [How to add TeaCache support](../features/teacache.md)
+See detailed guide: [How to add TeaCache support](../../design/feature/teacache.md)
 
 **Quick setup:**
 
@@ -744,7 +744,7 @@ omni = Omni(model="your-model",
 
 #### Cache-DiT
 
-See detailed guide: [How to add Cache-DiT support](../features/cache_dit.md)
+See detailed guide: [How to add Cache-DiT support](../../design/feature/cache_dit.md)
 
 **Quick setup:**
 

docs/design/feature/cache_dit.md

@@ -1,4 +1,4 @@
-# Support Cache-DiT
+# Cache-DiT
 
 This section describes how to add cache-dit acceleration to a new diffusion pipeline. We use the Qwen-Image pipeline and LongCat-Image pipeline as reference implementations.
 

docs/design/feature/cfg_parallel.md

@@ -1,4 +1,4 @@
-# Support CFG-Parallel
+# CFG-Parallel
 
 This section describes how to add CFG-Parallel (Classifier-Free Guidance Parallel) to a diffusion pipeline. We use the Qwen-Image pipeline as the reference implementation.
 

docs/design/feature/sequence_parallel.md

@@ -1,4 +1,4 @@
-# Support Sequence Parallel
+# Sequence Parallel
 
 This section describes how to add Sequence Parallel (SP) to a diffusion transformer model. We use the Qwen-Image transformer and Wan2.2 transformer as reference implementations.
 

docs/design/feature/teacache.md

@@ -1,4 +1,4 @@
-# Support TeaCache
+# TeaCache
 
 This section describes how to add TeaCache to a diffusion transformer model. We use the Qwen-Image transformer as the reference implementation.
 

docs/design/feature/tensor_parallel.md

@@ -1,4 +1,4 @@
-# Support Tensor Parallel
+# Tensor Parallel
 
 This section describes how to add Tensor Parallel (TP) to a diffusion transformer model. We use the Z-Image transformer as the reference implementation.
 

docs/mkdocs/hooks/generate_argparse.py

@@ -0,0 +1,257 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import ast
+import logging
+import re
+import sys
+from argparse import SUPPRESS, Action, ArgumentParser, HelpFormatter, _ArgumentGroup
+from collections.abc import Iterable
+from importlib.machinery import ModuleSpec
+from pathlib import Path
+from typing import Literal
+from unittest.mock import MagicMock
+
+from pydantic_core import core_schema
+
+
+class _FlexibleArgumentParser(ArgumentParser):
+    """Fallback parser for docs when vllm is unavailable.
+
+    Accepts the 'deprecated' kwarg used by vllm CLI and emits warnings
+    if a deprecated argument is actually provided.
+    """
+
+    _deprecated: set[Action] = set()
+    _deprecated_warned: set[str] = set()
+
+    if sys.version_info < (3, 13):
+
+        def parse_known_args(self, args=None, namespace=None):
+            namespace, args = super().parse_known_args(args, namespace)
+            for action in _FlexibleArgumentParser._deprecated:
+                if (
+                    hasattr(namespace, dest := action.dest)
+                    and getattr(namespace, dest) != action.default
+                    and dest not in _FlexibleArgumentParser._deprecated_warned
+                ):
+                    _FlexibleArgumentParser._deprecated_warned.add(dest)
+                    logger.warning("argument '%s' is deprecated", dest)
+            return namespace, args
+
+    def add_argument(self, *args, **kwargs):
+        deprecated = kwargs.pop("deprecated", False)
+        action = super().add_argument(*args, **kwargs)
+        if deprecated:
+            _FlexibleArgumentParser._deprecated.add(action)
+        return action
+
+    class _FlexibleArgumentGroup(_ArgumentGroup):
+        def add_argument(self, *args, **kwargs):
+            deprecated = kwargs.pop("deprecated", False)
+            action = super().add_argument(*args, **kwargs)
+            if deprecated:
+                _FlexibleArgumentParser._deprecated.add(action)
+            return action
+
+    def add_argument_group(self, *args, **kwargs):
+        group = self._FlexibleArgumentGroup(self, *args, **kwargs)
+        self._action_groups.append(group)
+        return group
+
+
+logger = logging.getLogger("mkdocs")
+
+# Define root and doc output directories
+ROOT_DIR = Path(__file__).parent.parent.parent.parent
+ARGPARSE_DOC_DIR = ROOT_DIR / "docs/generated/argparse_omni"
+
+# Ensure the repo root is in sys.path for dynamic imports
+sys.path.insert(0, str(ROOT_DIR))
+
+
+class PydanticMagicMock(MagicMock):
+    """`MagicMock` that's able to generate pydantic-core schemas and avoids _mock_name=None errors."""
+
+    def __init__(self, *args, **kwargs):
+        name = kwargs.pop("name", None)
+        super().__init__(*args, **kwargs)
+        self.__spec__ = ModuleSpec(name, None)
+        self._mock_name = name or "mock"
+
+    def __get_pydantic_core_schema__(self, source_type, handler):
+        return core_schema.any_schema()
+
+
+# --- Static extraction for CLI argument docs ---
+
+
+def extract_omni_serve_subparser_init():
+    """
+    Statically parse vllm_omni/entrypoints/cli/serve.py to extract the subparser_init method
+    and return a callable that adds arguments to a parser. This avoids import and mock issues.
+    """
+    serve_path = ROOT_DIR / "vllm_omni" / "entrypoints" / "cli" / "serve.py"
+    with open(serve_path, encoding="utf-8") as f:
+        source = f.read()
+    tree = ast.parse(source, filename=str(serve_path))
+    # Find class OmniServeCommand
+    for node in tree.body:
+        if isinstance(node, ast.ClassDef) and node.name == "OmniServeCommand":
+            for item in node.body:
+                if isinstance(item, ast.FunctionDef) and item.name == "subparser_init":
+                    # We'll exec this function body in a dummy context
+                    func_code = ast.Module(body=[item], type_ignores=[])
+                    code = compile(func_code, filename=str(serve_path), mode="exec")
+                    # Prepare dummy context
+                    local_vars = {}
+
+                    # Provide a dummy subparsers with add_parser returning our parser
+                    class DummySubparsers:
+                        def add_parser(self, name, **kwargs):
+                            return _FlexibleArgumentParser(prog=name)
+
+                    dummy_subparsers = DummySubparsers()
+                    # Provide globals for exec
+                    # Extract DESCRIPTION from the source file (assume it's a triple-quoted string at the top)
+                    m = re.search(r'DESCRIPTION\s*=\s*([ru]?""".*?""")', source, re.DOTALL)
+                    DESCRIPTION = m.group(1)[3:-3] if m else ""
+                    exec_globals = {
+                        "_FlexibleArgumentParser": _FlexibleArgumentParser,
+                        "FlexibleArgumentParser": _FlexibleArgumentParser,
+                        "make_arg_parser": lambda parser: parser,  # no-op for doc
+                        "VLLM_SUBCMD_PARSER_EPILOG": "",
+                        "logger": logger,
+                        "DummySubparsers": DummySubparsers,
+                        "argparse": __import__("argparse"),
+                        "DESCRIPTION": DESCRIPTION,
+                    }
+                    exec(code, exec_globals, local_vars)
+                    # Get the function
+                    subparser_init = local_vars["subparser_init"]
+
+                    # Return a callable that mimics OmniServeCommand().subparser_init
+                    def parser_factory():
+                        class DummySelf:
+                            name = "serve"
+
+                        return subparser_init(DummySelf(), dummy_subparsers)
+
+                    return parser_factory
+    raise RuntimeError("Could not statically extract OmniServeCommand.subparser_init")
+
+
+OmniServeCommand = type("OmniServeCommand", (), {"subparser_init": staticmethod(extract_omni_serve_subparser_init())})
+
+
+class MarkdownFormatter(HelpFormatter):
+    """Custom formatter that generates markdown for argument groups."""
+
+    def __init__(self, prog: str, starting_heading_level: int = 3):
+        super().__init__(prog, max_help_position=sys.maxsize, width=sys.maxsize)
+
+        self._section_heading_prefix = "#" * starting_heading_level
+        self._argument_heading_prefix = "#" * (starting_heading_level + 1)
+        self._markdown_output = []
+
+    def start_section(self, heading: str):
+        if heading not in {"positional arguments", "options"}:
+            heading_md = f"\n{self._section_heading_prefix} {heading}\n\n"
+            self._markdown_output.append(heading_md)
+
+    def end_section(self):
+        pass
+
+    def add_text(self, text: str):
+        if text:
+            self._markdown_output.append(f"{text.strip()}\n\n")
+
+    def add_usage(self, usage, actions, groups, prefix=None):
+        pass
+
+    def add_arguments(self, actions: Iterable[Action]):
+        for action in actions:
+            if len(action.option_strings) == 0 or "--help" in action.option_strings:
+                continue
+
+            option_strings = f"`{'`, `'.join(action.option_strings)}`"
+            heading_md = f"{self._argument_heading_prefix} {option_strings}\n\n"
+            self._markdown_output.append(heading_md)
+
+            if choices := action.choices:
+                choices = f"`{'`, `'.join(str(c) for c in choices)}`"
+                self._markdown_output.append(f"Possible choices: {choices}\n\n")
+            elif (metavar := action.metavar) and isinstance(metavar, (list, tuple)):
+                metavar = f"`{'`, `'.join(str(m) for m in metavar)}`"
+                self._markdown_output.append(f"Possible choices: {metavar}\n\n")
+
+            if action.help:
+                self._markdown_output.append(f"{action.help}\n\n")
+
+            if (default := action.default) != SUPPRESS:
+                # Make empty string defaults visible
+                if default == "":
+                    default = '""'
+                self._markdown_output.append(f"Default: `{default}`\n\n")
+
+    def format_help(self):
+        """Return the formatted help as markdown."""
+        return "".join(self._markdown_output)
+
+
+# Function to create parser using subparser_init style CLI class
+
+
+def create_parser_subparser_init(subcmd_class):
+    """
+    Create an argparse parser using subparser_init style CLI class, with MarkdownFormatter.
+    Compatible with both AST-extracted static parser_factory and real class instance method.
+    """
+
+    class DummySubparsers:
+        def add_parser(self, name, **kwargs):
+            return _FlexibleArgumentParser(prog=name)
+
+    dummy_subparsers = DummySubparsers()
+    # If subparser_init is a zero-arg static function (AST-extracted), call it directly
+    subparser_init = getattr(subcmd_class, "subparser_init", None)
+    if subparser_init is not None:
+        import inspect
+
+        sig = inspect.signature(subparser_init)
+        if len(sig.parameters) == 0:
+            parser = subparser_init()
+        else:
+            parser = subcmd_class().subparser_init(dummy_subparsers)
+    else:
+        parser = subcmd_class().subparser_init(dummy_subparsers)
+    parser.formatter_class = MarkdownFormatter
+    return parser
+
+
+def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
+    """
+    Entry point for doc generation. Builds doc directory and outputs markdown for each CLI command.
+    """
+    logger.info("Generating vllm-omni argparse documentation")
+    logger.debug("Root directory: %s", ROOT_DIR.resolve())
+    logger.debug("Output directory: %s", ARGPARSE_DOC_DIR.resolve())
+    if not ARGPARSE_DOC_DIR.exists():
+        ARGPARSE_DOC_DIR.mkdir(parents=True)
+
+    # Register all CLI parsers; you can easily add more commands here
+    parsers = {
+        "omni_serve": create_parser_subparser_init(OmniServeCommand),
+        # "another_cmd": create_parser_subparser_init(AnotherCommandClass),
+    }
+
+    for stem, parser in parsers.items():
+        doc_path = ARGPARSE_DOC_DIR / f"{stem}.inc.md"
+        with open(doc_path, "w", encoding="utf-8") as f:
+            f.write(super(type(parser), parser).format_help())
+        logger.info("Argparse generated: %s", doc_path.relative_to(ROOT_DIR))
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    on_startup("build", False)

examples/offline_inference/bagel/README.md

@@ -177,6 +177,74 @@ Example configuration for TP=2 on GPUs 0 and 1:
 | `max_inflight`        | `1`     | Maximum inflight requests        |
 | `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB)   |
 
+## Using Mooncake Connector
+
+[Mooncake](https://github.com/kvcache-ai/Mooncake) is a high-performance distributed KV cache transfer engine that enables efficient cross-node data movement via TCP or RDMA, making it ideal for multi-node disaggregated inference.
+
+By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can switch to the Mooncake connector for better performance on multi-GPU setups and to enable multi-node deployment.
+
+### Prerequisites
+
+Install the Mooncake transfer engine:
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+### Step 1: Start the Mooncake Master
+
+On the **primary node**, start the Mooncake master service (run in a separate terminal or background with `&`):
+
+```bash
+# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
+# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+### Step 2: Run Offline Inference with Mooncake
+
+Use the provided Mooncake stage config [`bagel_multiconnector.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml). Before launching, update the `metadata_server` and `master` addresses in the YAML to match your Mooncake master node's IP (use `127.0.0.1` for single-node testing).
+
+```bash
+cd examples/offline_inference/bagel
+
+# Text to Image with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --prompts "A cute cat" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+
+# Image to Text with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2text \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Describe this image" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+
+# Text to Text with Mooncake
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --prompts "What is the capital of France?" \
+                  --stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+```
+
+For more details on the Mooncake connector and multi-node setup, see the [Mooncake Store Connector documentation](../../../docs/design/feature/omni_connectors/mooncake_store_connector.md).
+
+------
+
 ## FAQ
 
 - If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.

examples/offline_inference/qwen3_tts/README.md

@@ -87,7 +87,20 @@ Examples:
 python end2end.py --query-type Base --mode-tag icl
 ```
 
+## Batched Decoding
+
+The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
+
+```
+python end2end.py --query-type CustomVoice \
+    --txt-prompts benchmark_prompts.txt \
+    --batch-size 4 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
+```
+
+**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
+
 ## Notes
 
 - The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs.
-- Use `--output-dir` (preferred) or `--output-wav` to change the output folder.
+- Use `--output-dir` to change the output folder.

examples/offline_inference/qwen3_tts/benchmark_prompts.txt

@@ -0,0 +1,12 @@
+Hello, welcome to the voice synthesis benchmark test.
+She said she would be here by noon, but nobody showed up.
+The quick brown fox jumps over the lazy dog near the riverbank.
+I can't believe how beautiful the sunset looks from up here on the mountain.
+Please remember to bring your identification documents to the appointment tomorrow morning.
+Have you ever wondered what it would be like to travel through time and visit ancient civilizations?
+The restaurant on the corner serves the best pasta I have ever tasted in my entire life.
+After the meeting, we should discuss the quarterly results and plan for the next phase.
+Learning a new language takes patience, practice, and a genuine curiosity about other cultures.
+The train leaves at half past seven, so we need to arrive at the station before then.
+Could you please turn down the music a little bit, I'm trying to concentrate on my work.
+It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.

examples/offline_inference/qwen3_tts/end2end.py

@@ -248,6 +248,13 @@ def main(args):
     Args:
         args: Parsed CLI args from parse_args().
     """
+    if args.batch_size < 1 or (args.batch_size & (args.batch_size - 1)) != 0:
+        raise ValueError(
+            f"--batch-size must be a power of two (got {args.batch_size}); "
+            "non-power-of-two values do not align with CUDA graph capture sizes "
+            "of Code2Wav."
+        )
+
     query_func = query_map[args.query_type]
     if args.query_type in {"CustomVoice", "VoiceDesign"}:
         query_result = query_func(use_batch_sample=args.use_batch_sample)
@@ -260,6 +267,33 @@ def main(args):
         query_result = query_func()
 
     model_name = query_result.model_name
+
+    # Load prompts from text file if provided.
+    # Use the default query as a template so task-specific fields
+    # (e.g. ref_audio for Base) are preserved; only override text.
+    if args.txt_prompts:
+        with open(args.txt_prompts) as f:
+            lines = [line.strip() for line in f if line.strip()]
+        if not lines:
+            raise ValueError(f"No valid prompts found in {args.txt_prompts}")
+        template = query_result.inputs
+        if isinstance(template, list):
+            template = template[0]
+        template_info = template["additional_information"]
+        inputs = []
+        for text in lines:
+            additional_information = {**template_info, "text": [text]}
+            inputs.append(
+                {
+                    "prompt_token_ids": [0] * _estimate_prompt_len(additional_information, model_name),
+                    "additional_information": additional_information,
+                }
+            )
+    else:
+        inputs = query_result.inputs
+        if not isinstance(inputs, list):
+            inputs = [inputs]
+
     omni = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
@@ -267,32 +301,35 @@ def main(args):
         stage_init_timeout=args.stage_init_timeout,
     )
 
-    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    output_dir = args.output_dir
     os.makedirs(output_dir, exist_ok=True)
 
-    omni_generator = omni.generate(query_result.inputs, sampling_params_list=None)
-    for stage_outputs in omni_generator:
-        for output in stage_outputs.request_output:
-            request_id = output.request_id
-            audio_data = output.outputs[0].multimodal_output["audio"]
-            # async_chunk mode returns a list of chunks; concatenate them.
-            if isinstance(audio_data, list):
-                audio_tensor = torch.cat(audio_data, dim=-1)
-            else:
-                audio_tensor = audio_data
-            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-            sr_val = output.outputs[0].multimodal_output["sr"]
-            audio_samplerate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val[-1])
-            # Convert to numpy array and ensure correct format
-            audio_numpy = audio_tensor.float().detach().cpu().numpy()
+    batch_size = args.batch_size
+    for batch_start in range(0, len(inputs), batch_size):
+        batch = inputs[batch_start : batch_start + batch_size]
+        omni_generator = omni.generate(batch, sampling_params_list=None)
+        for stage_outputs in omni_generator:
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+                audio_data = output.outputs[0].multimodal_output["audio"]
+                # async_chunk mode returns a list of chunks; concatenate them.
+                if isinstance(audio_data, list):
+                    audio_tensor = torch.cat(audio_data, dim=-1)
+                else:
+                    audio_tensor = audio_data
+                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+                sr_val = output.outputs[0].multimodal_output["sr"]
+                audio_samplerate = sr_val.item() if hasattr(sr_val, "item") else int(sr_val[-1])
+                # Convert to numpy array and ensure correct format
+                audio_numpy = audio_tensor.float().detach().cpu().numpy()
 
-            # Ensure audio is 1D (flatten if needed)
-            if audio_numpy.ndim > 1:
-                audio_numpy = audio_numpy.flatten()
+                # Ensure audio is 1D (flatten if needed)
+                if audio_numpy.ndim > 1:
+                    audio_numpy = audio_numpy.flatten()
 
-            # Save audio file with explicit WAV format
-            sf.write(output_wav, audio_numpy, samplerate=audio_samplerate, format="WAV")
-            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+                # Save audio file with explicit WAV format
+                sf.write(output_wav, audio_numpy, samplerate=audio_samplerate, format="WAV")
+                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
 
 def parse_args():
@@ -341,9 +378,9 @@ def parse_args():
         help="Threshold for using shared memory in bytes (default: 65536)",
     )
     parser.add_argument(
-        "--output-wav",
+        "--output-dir",
         default="output_audio",
-        help="[Deprecated] Output wav directory (use --output-dir).",
+        help="Output directory for generated wav files (default: output_audio).",
     )
     parser.add_argument(
         "--num-prompts",
@@ -401,6 +438,12 @@ def parse_args():
         choices=["icl", "xvec_only"],
         help="Mode tag for Base query x_vector_only_mode (default: icl).",
     )
+    parser.add_argument(
+        "--batch-size",
+        type=int,
+        default=1,
+        help="Number of prompts per batch (default: 1, sequential).",
+    )
 
     return parser.parse_args()
 

examples/online_serving/bagel/README.md

@@ -45,6 +45,129 @@ For larger models or multi-GPU environments, you can enable Tensor Parallelism (
 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/your/custom_bagel.yaml
 ```
 
+#### Using Mooncake Connector
+
+By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can use the [Mooncake](https://github.com/kvcache-ai/Mooncake) connector to transfer KV cache between stages, which also enables multi-node deployment.
+
+**1. Install Mooncake**
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+**2. Start Mooncake Master** on the primary node:
+
+```bash
+# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
+# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+**3. Launch the server** with the Mooncake stage config:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
+```
+
+> **Note**: Before launching, edit [`bagel_multiconnector.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. For single-node testing, `127.0.0.1` works.
+
+The client-side usage is identical to the default setup -- the Mooncake connector is transparent to the API. See the requests section below.
+
+For more details on the Mooncake connector configuration, see the [Mooncake Store Connector documentation](../../../docs/design/feature/omni_connectors/mooncake_store_connector.md).
+
+#### Multi-Node Deployment
+
+You can deploy each stage on a **separate node** for better resource utilization. In this example, the orchestrator (Stage 0 / Thinker) and Stage 1 (DiT) run on different machines, connected via Mooncake.
+
+Replace `<ORCHESTRATOR_IP>` below with the actual IP address of your orchestrator node (e.g., `10.244.227.244`).
+
+> [!WARNING]
+> **Before launching**, edit [`bagel_multiconnector.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. Mismatched addresses will cause silent connection failures.
+
+**1. Start Mooncake Master** (on the orchestrator node):
+
+```bash
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=<ORCHESTRATOR_IP> \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003
+```
+
+**2. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
+    --port 8000 \ # API server port for client requests
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
+    --stage-id 0 \
+    -oma <ORCHESTRATOR_IP> \
+    -omp 8091
+```
+
+**3. Launch Stage 1 (DiT)** on the remote node in headless mode:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
+    --stage-id 1 \
+    --headless \
+    -oma <ORCHESTRATOR_IP> \
+    -omp 8091
+```
+
+**Mooncake Master arguments:**
+
+| Argument | Description |
+| :------- | :---------- |
+| `--rpc_port` | Mooncake RPC port for control-plane coordination between stages |
+| `--enable_http_metadata_server` | Enable the HTTP metadata server for service discovery |
+| `--http_metadata_server_host` | IP address to bind the metadata server (use the orchestrator node's IP) |
+| `--http_metadata_server_port` | Port for the HTTP metadata server |
+| `--metrics_port` | Port for Prometheus-compatible metrics endpoint |
+
+**vllm serve arguments:**
+
+| Argument | Description |
+| :------- | :---------- |
+| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) |
+| `--headless` | Run without the API server (worker-only mode) |
+| `-oma` | Orchestrator master address |
+| `-omp` | Orchestrator master port for Stage 1 to connect to Stage 0 for task coordination |
+
+> [!IMPORTANT]
+> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless).
+> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior.
+
+**Network Requirements**
+
+All nodes must have network connectivity to each other. Ensure the following ports are open **between all participating nodes**:
+
+| Port | Protocol | Service | Direction |
+| :--- | :------- | :------ | :-------- |
+| 50051 | TCP | Mooncake Master RPC | Worker → Orchestrator |
+| 8080 | TCP | Mooncake HTTP Metadata Server | Worker → Orchestrator |
+| 8091 | TCP | Orchestrator Master (`-omp`) | Worker → Orchestrator |
+| 8000 | TCP | API Server (`--port`) | Client → Orchestrator |
+| 9003 | TCP | Metrics (optional) | Monitoring → Orchestrator |
+
+> **Tip**: If nodes are behind a firewall or in different VPCs/security groups, make sure the above ports are allowed in ingress/egress rules. All nodes should be reachable via their IP addresses (no NAT). Using nodes on the same subnet or VPC is recommended to minimize latency for Mooncake KV cache transfers.
+
 ### Send Multi-modal Request
 
 Get into the bagel folder:

examples/online_serving/bagel/run_server_stage_cli.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+# Bagel multi-stage online serving startup script
+# Starts stage 0 as master with API server, and stage 1 in headless mode
+
+MODEL="${MODEL:-ByteDance-Seed/BAGEL-7B-MoT}"
+PORT="${PORT:-8091}"
+MASTER_ADDRESS="${MASTER_ADDRESS:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-8092}"
+STAGE_CONFIGS_PATH="$(dirname "$0")/../../../vllm_omni/model_executor/stage_configs/bagel.yaml"
+
+echo "Starting Bagel multi-stage server..."
+echo "Model: $MODEL"
+echo "API Port: $PORT"
+echo "Master Address: $MASTER_ADDRESS"
+echo "Master Port: $MASTER_PORT"
+echo "Stage Configs: $STAGE_CONFIGS_PATH"
+
+# Start stage 1 (DiT) in headless mode first
+echo "Starting Stage 1 (DiT) in headless mode..."
+vllm serve "$MODEL" --omni \
+    --stage-configs-path "$STAGE_CONFIGS_PATH" \
+    --stage-id 1 \
+    --headless \
+    -oma "$MASTER_ADDRESS" \
+    -omp "$MASTER_PORT" &
+
+# Start stage 0 (Thinker) as master with API server
+echo "Starting Stage 0 (Thinker) as master..."
+vllm serve "$MODEL" --omni \
+    --port "$PORT" \
+    --stage-configs-path "$STAGE_CONFIGS_PATH" \
+    --stage-id 0 \
+    -oma "$MASTER_ADDRESS" \
+    -omp "$MASTER_PORT"

mkdocs.yml

@@ -59,6 +59,7 @@ hooks:
   - docs/mkdocs/hooks/generate_api_readme.py
   - docs/mkdocs/hooks/url_schemes.py
   - docs/mkdocs/hooks/generate_examples.py
+  - docs/mkdocs/hooks/generate_argparse.py
 
 # Exclude include files from navigation warnings
 exclude_docs: |

requirements/common.txt

@@ -14,3 +14,4 @@ imageio[ffmpeg]>=2.37.2
 sox>=1.5.0
 prettytable>=3.8.0
 aenum==3.1.16
+pyzmq>=25.0.0

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -3,21 +3,23 @@
 
 import pytest
 
-from vllm_omni.entrypoints import omni as omni_module
+from vllm_omni.entrypoints import utils as utils_module
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+MODEL = "riverclouds/qwen_image_random"
+
 
 def test_default_stage_config_includes_cache_backend(monkeypatch):
     """Ensure cache_backend/cache_config are preserved in default diffusion stage."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
+    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
+    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
 
     omni = AsyncOmni(
-        model="dummy-model",
+        model=MODEL,
         cache_backend="cache_dit",
         cache_config='{"Fn_compute_blocks": 2}',
         vae_use_slicing=True,
@@ -41,13 +43,13 @@ def test_default_stage_config_includes_cache_backend(monkeypatch):
 
 def test_default_cache_config_used_when_missing(monkeypatch):
     """Ensure default cache_config is applied when cache_backend is set."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
+    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
+    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
 
     omni = AsyncOmni(
-        model="dummy-model",
+        model=MODEL,
         cache_backend="cache_dit",
     )
 
@@ -59,13 +61,13 @@ def test_default_cache_config_used_when_missing(monkeypatch):
 
 def test_default_stage_devices_from_sequence_parallel(monkeypatch):
     """Ensure devices list reflects sequence parallel size when no parallel_config is provided."""
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
-    monkeypatch.setattr(omni_module, "resolve_model_config_path", lambda model: None)
+    monkeypatch.setattr(utils_module, "load_stage_configs_from_model", lambda model, base_engine_args=None: [])
+    monkeypatch.setattr(utils_module, "resolve_model_config_path", lambda model: None)
     monkeypatch.setattr(AsyncOmni, "_start_stages", lambda self, model: None)
     monkeypatch.setattr(AsyncOmni, "_wait_for_stages_ready", lambda self, timeout=0: None)
 
     omni = AsyncOmni(
-        model="dummy-model",
+        model=MODEL,
         ulysses_degree=2,
         ring_degree=2,
     )

tests/entrypoints/test_omni_diffusion.py

@@ -18,6 +18,37 @@ warnings.filterwarnings(
     category=DeprecationWarning,
 )
 
+MODEL = "riverclouds/qwen_image_random"
+
+
+class _FakeStageRequestStats:
+    """Fake StageRequestStats object with necessary attributes aligned with real StageRequestStats."""
+
+    def __init__(self, **kwargs):
+        # Required fields (with defaults for testing)
+        self.batch_id = kwargs.get("batch_id", 0)
+        self.batch_size = kwargs.get("batch_size", 1)
+        self.num_tokens_in = kwargs.get("num_tokens_in", 0)
+        self.num_tokens_out = kwargs.get("num_tokens_out", 1)
+        self.stage_gen_time_ms = kwargs.get("stage_gen_time_ms", 10.0)
+        self.rx_transfer_bytes = kwargs.get("rx_transfer_bytes", 0)
+        self.rx_decode_time_ms = kwargs.get("rx_decode_time_ms", 0.0)
+        self.rx_in_flight_time_ms = kwargs.get("rx_in_flight_time_ms", 0.0)
+        self.stage_stats = kwargs.get("stage_stats", None)
+
+        # Optional fields
+        self.stage_id = kwargs.get("stage_id", None)
+        self.final_output_type = kwargs.get("final_output_type", None)
+        self.request_id = kwargs.get("request_id", None)
+        self.postprocess_time_ms = kwargs.get("postprocess_time_ms", 0.0)
+        self.diffusion_metrics = kwargs.get("diffusion_metrics", None)
+        self.audio_generated_frames = kwargs.get("audio_generated_frames", 0)
+
+        # Allow additional attributes for flexibility
+        for key, value in kwargs.items():
+            if not hasattr(self, key):
+                setattr(self, key, value)
+
 
 class _FakeEngineArgs(dict):
     """Fake engine args that can be used both as object attributes and as **kwargs."""
@@ -60,8 +91,10 @@ class _FakeQueue:
     def put_nowait(self, item):
         self._queue.put_nowait(item)
 
-    def get(self):
-        return self._queue.get()
+    def get(self, timeout=None):
+        if timeout is None:
+            return self._queue.get()
+        return self._queue.get(timeout=timeout)
 
     def get_nowait(self):
         return self._queue.get_nowait()
@@ -70,6 +103,52 @@ class _FakeQueue:
         return self._queue.empty()
 
 
+class _FakeZmqQueue:
+    """Fake ZmqQueue that wraps _FakeQueue and matches ZmqQueue interface."""
+
+    def __init__(
+        self,
+        ctx=None,
+        socket_type=None,
+        *,
+        bind: str | None = None,
+        connect: str | None = None,
+        recv_timeout_ms: int | None = None,
+        send_timeout_ms: int | None = None,
+    ):
+        """Initialize fake ZMQ queue with same signature as real ZmqQueue."""
+        self._queue = _FakeQueue(maxsize=0)
+        # Determine endpoint from bind or connect
+        path = bind if bind is not None else connect
+        self.endpoint = path or f"fake://zmq-endpoint-{id(self)}"
+        self._recv_timeout_ms = recv_timeout_ms
+        self._send_timeout_ms = send_timeout_ms
+
+    def put(self, obj: Any) -> None:
+        """Send an object to the queue."""
+        self._queue.put(obj)
+
+    def put_nowait(self, obj: Any) -> None:
+        """Send an object to the queue without blocking."""
+        self._queue.put_nowait(obj)
+
+    def get(self, timeout: float | None = None) -> Any:
+        """Receive an object from the queue with optional timeout in seconds."""
+        return self._queue.get(timeout=timeout)
+
+    def get_nowait(self) -> Any:
+        """Receive an object from the queue without blocking."""
+        return self._queue.get_nowait()
+
+    def empty(self) -> bool:
+        """Check if the queue is empty without blocking."""
+        return self._queue.empty()
+
+    def close(self) -> None:
+        """Close the queue."""
+        pass
+
+
 class _FakeStage:
     """Lightweight Stage stub for multi-process pipeline version with queue support."""
 
@@ -272,6 +351,19 @@ def _setup_multiprocessing_mocks(monkeypatch):
     monkeypatch.setattr(mp, "get_context", _mock_get_context, raising=False)
     monkeypatch.setattr(mp, "Process", fake_process_class, raising=False)
 
+    # Mock ZmqQueue to use _FakeZmqQueue
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.zmq_utils.ZmqQueue",
+        _FakeZmqQueue,
+        raising=False,
+    )
+    # Also mock where ZmqQueue is imported/used
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.ZmqQueue",
+        _FakeZmqQueue,
+        raising=False,
+    )
+
 
 def _setup_ipc_mocks(monkeypatch):
     """Helper function to set up IPC function mocks."""
@@ -324,6 +416,65 @@ def _setup_log_mocks(monkeypatch):
     )
 
 
+def _setup_connector_mocks(monkeypatch, omni_module=None):
+    """Helper function to set up connector mocks for stage-to-stage forwarding.
+
+    If omni_module is provided, mocks directly on the module. Otherwise, uses string path.
+    """
+
+    # Mock initialize_orchestrator_connectors to return fake connectors
+    def _fake_initialize_orchestrator_connectors(config_path, worker_backend=None, shm_threshold_bytes=None):
+        # Create fake connectors for all stage-to-stage edges
+        # Each connector is just a mock object that will be passed to try_send_via_connector
+        fake_connectors = {}
+        # Add connectors for common edges (0->1, 1->2, etc.)
+        for i in range(10):  # Support up to 10 stages
+            fake_connectors[(str(i), str(i + 1))] = MagicMock()
+        return None, fake_connectors
+
+    if omni_module is not None:
+        # Mock directly on the omni module where it's used (after import)
+        monkeypatch.setattr(omni_module, "initialize_orchestrator_connectors", _fake_initialize_orchestrator_connectors)
+    else:
+        # Mock via string path (before import)
+        monkeypatch.setattr(
+            "vllm_omni.entrypoints.omni.initialize_orchestrator_connectors",
+            _fake_initialize_orchestrator_connectors,
+            raising=False,
+        )
+
+
+def _setup_connector_adapter_mock(monkeypatch, omni_module):
+    """Helper function to mock try_send_via_connector on the omni module.
+
+    This must be called AFTER importing omni module, to mock the function where it's actually used.
+    """
+
+    # Mock try_send_via_connector to always succeed
+    def _fake_try_send_via_connector(
+        connector,
+        stage_id,
+        next_stage_id,
+        req_id,
+        next_inputs,
+        sampling_params,
+        original_prompt,
+        next_stage_queue_submit_fn,
+        metrics,
+    ):
+        # Simulate successful send by calling the submit function
+        task = {
+            "request_id": req_id,
+            "engine_inputs": next_inputs,
+            "sampling_params": sampling_params,
+        }
+        next_stage_queue_submit_fn(task)
+        return True
+
+    # Mock directly on the omni module where it's used
+    monkeypatch.setattr(omni_module, "try_send_via_connector", _fake_try_send_via_connector)
+
+
 @pytest.fixture(autouse=True)
 def mock_get_config(monkeypatch):
     """Auto-mock get_config and related model loading functions to avoid model path validation."""
@@ -449,10 +600,19 @@ def mock_get_config(monkeypatch):
 def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config):
     """Test that stage configs are auto-loaded when stage_configs_path is None."""
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [
-            _FakeStageConfig(fake_stage_config),
-            _FakeStageConfig(fake_stage_config),
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        cfg0 = dict(fake_stage_config)
+        cfg0["stage_id"] = 0
+        cfg1 = dict(fake_stage_config)
+        cfg1["stage_id"] = 1
+        return None, [
+            _FakeStageConfig(cfg0),
+            _FakeStageConfig(cfg1),
         ]
 
     # Remove modules from cache BEFORE setting mocks
@@ -471,10 +631,11 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
-    # Mock load_stage_configs_from_model
+    # Mock load_and_resolve_stage_configs
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -489,13 +650,13 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
     # Import the module after mocks are set
     import vllm_omni.entrypoints.omni as omni_module
 
-    # Patch the imported function and class in the module
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    # Patch the imported class in the module
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
     # Verify: auto-loaded stage_configs and stage_list have consistent count
     assert isinstance(omni.stage_configs, list)
     assert len(omni.stage_configs) == 2
@@ -514,8 +675,13 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
 def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
     """Test that generate raises ValueError when sampling_params_list length doesn't match."""
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(fake_stage_config)]
 
     import sys
 
@@ -531,14 +697,10 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-        _fake_loader,
-        raising=False,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -550,12 +712,12 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
     with pytest.raises(ValueError):
         omni.generate(prompts=["hi"], sampling_params_list=[])
 
@@ -563,11 +725,18 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
 def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
     """Test multi-stage generation pipeline with queue polling."""
     stage_cfg0 = dict(fake_stage_config)
+    stage_cfg0["stage_id"] = 0
     stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["stage_id"] = 1
     stage_cfg1["processed_input"] = ["processed-for-stage-1"]
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
 
@@ -575,6 +744,8 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
         "vllm_omni.entrypoints.omni_stage",
+        "vllm_omni.distributed.omni_connectors.adapter",
+        "vllm_omni.distributed.omni_connectors",
     ]:
         if module_name in sys.modules:
             del sys.modules[module_name]
@@ -585,12 +756,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
     _setup_log_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-        _fake_loader,
-        raising=False,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -602,8 +768,11 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, omni_module)
+    _setup_connector_adapter_mock(monkeypatch, omni_module)
 
     # Mock uuid.uuid4() to return a predictable value for request ID generation
     test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -612,7 +781,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
 
     # Generate the expected request ID format: "0_<uuid>"
     expected_request_id = f"0_{test_uuid}"
@@ -625,7 +794,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
         {
             "request_id": expected_request_id,
             "engine_outputs": [{"stage": 0, "text": "s0"}],
-            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+            "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0),
         }
     )
     # Stage 1 output (will be collected after stage 0 forwards to it)
@@ -637,7 +806,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
         {
             "request_id": expected_request_id,
             "engine_outputs": [{"stage": 1, "text": "s1"}],
-            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+            "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0),
         }
     )
 
@@ -662,11 +831,18 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
 def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
     """Test single-stage generation pipeline with multiple inputs in one batch."""
     stage_cfg0 = dict(fake_stage_config)
-    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg0["stage_id"] = 0
     stage_cfg0["final_output"] = False
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["stage_id"] = 1
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
 
@@ -674,6 +850,8 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
         "vllm_omni.entrypoints.omni_stage",
+        "vllm_omni.distributed.omni_connectors.adapter",
+        "vllm_omni.distributed.omni_connectors",
     ]:
         if module_name in sys.modules:
             del sys.modules[module_name]
@@ -684,7 +862,7 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
     _setup_log_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -696,8 +874,11 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, omni_module)
+    _setup_connector_adapter_mock(monkeypatch, omni_module)
 
     # Mock uuid.uuid4() to return a predictable value for request ID generation
     test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -706,7 +887,7 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
 
     # Generate the expected request ID format: "0_<uuid>"
     expected_request_id = f"0_{test_uuid}"
@@ -718,28 +899,28 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
         {
             "request_id": expected_request_id,
             "engine_outputs": [{"stage": 0, "text": "s0"}],
-            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+            "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0),
         }
     )
     omni.stage_list[0]._out_q.put_nowait(
         {
             "request_id": expected_request_id,
             "engine_outputs": [{"stage": 0, "text": "s0"}],
-            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+            "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0),
         }
     )
     omni.stage_list[1]._out_q.put_nowait(
         {
             "request_id": expected_request_id,
             "engine_outputs": [{"stage": 1}],
-            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+            "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0),
         }
     )
     omni.stage_list[1]._out_q.put_nowait(
         {
             "request_id": expected_request_id,
             "engine_outputs": [{"stage": 1}],
-            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+            "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0),
         }
     )
 
@@ -768,12 +949,19 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
 def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
     """Test that generate returns empty list when all stages have final_output=False."""
     stage_cfg0 = dict(fake_stage_config)
-    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg0["stage_id"] = 0
     stage_cfg0["final_output"] = False
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["stage_id"] = 1
     stage_cfg1["final_output"] = False
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
 
@@ -781,6 +969,8 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
         "vllm_omni.entrypoints.omni_stage",
+        "vllm_omni.distributed.omni_connectors.adapter",
+        "vllm_omni.distributed.omni_connectors",
     ]:
         if module_name in sys.modules:
             del sys.modules[module_name]
@@ -791,7 +981,7 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
     _setup_log_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -803,8 +993,11 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, omni_module)
+    _setup_connector_adapter_mock(monkeypatch, omni_module)
 
     # Mock uuid.uuid4() to return a predictable value for request ID generation
     test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -813,7 +1006,7 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
 
     # Generate the expected request ID format: "0_<uuid>"
     expected_request_id = f"0_{test_uuid}"
@@ -847,12 +1040,19 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
 def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_config):
     """Test that generate uses default sampling params when sampling_params_list is None."""
     stage_cfg0 = dict(fake_stage_config)
-    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg0["stage_id"] = 0
     stage_cfg0["final_output"] = False
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["stage_id"] = 1
     stage_cfg1["final_output"] = False
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
 
@@ -860,6 +1060,8 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
         "vllm_omni.entrypoints.utils",
         "vllm_omni.entrypoints.omni",
         "vllm_omni.entrypoints.omni_stage",
+        "vllm_omni.distributed.omni_connectors.adapter",
+        "vllm_omni.distributed.omni_connectors",
     ]:
         if module_name in sys.modules:
             del sys.modules[module_name]
@@ -870,7 +1072,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
     _setup_log_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -882,8 +1084,11 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, omni_module)
+    _setup_connector_adapter_mock(monkeypatch, omni_module)
 
     # Mock uuid.uuid4() to return a predictable value for request ID generation
     test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -892,7 +1097,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
 
     # Generate the expected request ID format: "0_<uuid>"
     expected_request_id = f"0_{test_uuid}"
@@ -902,14 +1107,14 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
         {
             "request_id": expected_request_id,
             "engine_outputs": [{"stage": 0}],
-            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+            "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0),
         }
     )
     omni.stage_list[1]._out_q.put_nowait(
         {
             "request_id": expected_request_id,
             "engine_outputs": [{"stage": 1}],
-            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+            "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0),
         }
     )
     # Use the default sampling params
@@ -919,8 +1124,13 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
 def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
     """Test that _wait_for_stages_ready handles timeout correctly."""
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(fake_stage_config)]
 
     import sys
 
@@ -936,9 +1146,10 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -961,13 +1172,13 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
 
     from vllm_omni.entrypoints.omni import Omni
 
     # Use very short timeout
-    omni = Omni(model="any", init_timeout=0.01)
+    omni = Omni(model=MODEL, init_timeout=0.01)
     # Verify that no stages are ready
     assert len(omni._stages_ready) == 0
 
@@ -975,8 +1186,13 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
 def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     """Test that generate handles error messages from stages correctly."""
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(fake_stage_config)]
 
     import sys
 
@@ -992,9 +1208,10 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -1006,8 +1223,8 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
 
     # Mock uuid.uuid4() to return a predictable value for request ID generation
     test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -1016,7 +1233,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
 
     # Generate the expected request ID format: "0_<uuid>"
     expected_request_id = f"0_{test_uuid}"
@@ -1034,7 +1251,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
         {
             "request_id": expected_request_id,
             "engine_outputs": [{"stage": 0, "text": "result"}],
-            "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0},
+            "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0),
         }
     )
 
@@ -1050,8 +1267,13 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
 def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
     """Test that close() sends shutdown signal to all input queues."""
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(fake_stage_config)]
 
     import sys
 
@@ -1067,9 +1289,10 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -1081,12 +1304,12 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
 
     # Call close
     omni.close()

tests/entrypoints/test_omni_llm.py

@@ -19,6 +19,8 @@ warnings.filterwarnings(
     category=DeprecationWarning,
 )
 
+MODEL = "riverclouds/qwen_image_random"
+
 
 class _FakeEngineArgs(dict):
     """Fake engine args that can be used both as object attributes and as **kwargs."""
@@ -61,8 +63,10 @@ class _FakeQueue:
     def put_nowait(self, item):
         self._queue.put_nowait(item)
 
-    def get(self):
-        return self._queue.get()
+    def get(self, timeout=None):
+        if timeout is None:
+            return self._queue.get()
+        return self._queue.get(timeout=timeout)
 
     def get_nowait(self):
         return self._queue.get_nowait()
@@ -71,6 +75,52 @@ class _FakeQueue:
         return self._queue.empty()
 
 
+class _FakeZmqQueue:
+    """Fake ZmqQueue that wraps _FakeQueue and matches ZmqQueue interface."""
+
+    def __init__(
+        self,
+        ctx=None,
+        socket_type=None,
+        *,
+        bind: str | None = None,
+        connect: str | None = None,
+        recv_timeout_ms: int | None = None,
+        send_timeout_ms: int | None = None,
+    ):
+        """Initialize fake ZMQ queue with same signature as real ZmqQueue."""
+        self._queue = _FakeQueue(maxsize=0)
+        # Determine endpoint from bind or connect
+        path = bind if bind is not None else connect
+        self.endpoint = path or f"fake://zmq-endpoint-{id(self)}"
+        self._recv_timeout_ms = recv_timeout_ms
+        self._send_timeout_ms = send_timeout_ms
+
+    def put(self, obj: Any) -> None:
+        """Send an object to the queue."""
+        self._queue.put(obj)
+
+    def put_nowait(self, obj: Any) -> None:
+        """Send an object to the queue without blocking."""
+        self._queue.put_nowait(obj)
+
+    def get(self, timeout: float | None = None) -> Any:
+        """Receive an object from the queue with optional timeout in seconds."""
+        return self._queue.get(timeout=timeout)
+
+    def get_nowait(self) -> Any:
+        """Receive an object from the queue without blocking."""
+        return self._queue.get_nowait()
+
+    def empty(self) -> bool:
+        """Check if the queue is empty without blocking."""
+        return self._queue.empty()
+
+    def close(self) -> None:
+        """Close the queue."""
+        pass
+
+
 class _FakeStage:
     """Lightweight Stage stub for multi-process pipeline version with queue support."""
 
@@ -273,6 +323,19 @@ def _setup_multiprocessing_mocks(monkeypatch):
     monkeypatch.setattr(mp, "get_context", _mock_get_context, raising=False)
     monkeypatch.setattr(mp, "Process", fake_process_class, raising=False)
 
+    # Mock ZmqQueue to use _FakeZmqQueue
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.zmq_utils.ZmqQueue",
+        _FakeZmqQueue,
+        raising=False,
+    )
+    # Also mock where ZmqQueue is imported/used
+    monkeypatch.setattr(
+        "vllm_omni.entrypoints.omni_stage.ZmqQueue",
+        _FakeZmqQueue,
+        raising=False,
+    )
+
 
 def _setup_ipc_mocks(monkeypatch):
     """Helper function to set up IPC function mocks."""
@@ -325,6 +388,65 @@ def _setup_log_mocks(monkeypatch):
     )
 
 
+def _setup_connector_mocks(monkeypatch, omni_module=None):
+    """Helper function to set up connector mocks for stage-to-stage forwarding.
+
+    If omni_module is provided, mocks directly on the module. Otherwise, uses string path.
+    """
+
+    # Mock initialize_orchestrator_connectors to return fake connectors
+    def _fake_initialize_orchestrator_connectors(config_path, worker_backend=None, shm_threshold_bytes=None):
+        # Create fake connectors for all stage-to-stage edges
+        # Each connector is just a mock object that will be passed to try_send_via_connector
+        fake_connectors = {}
+        # Add connectors for common edges (0->1, 1->2, etc.)
+        for i in range(10):  # Support up to 10 stages
+            fake_connectors[(str(i), str(i + 1))] = MagicMock()
+        return None, fake_connectors
+
+    if omni_module is not None:
+        # Mock directly on the omni module where it's used (after import)
+        monkeypatch.setattr(omni_module, "initialize_orchestrator_connectors", _fake_initialize_orchestrator_connectors)
+    else:
+        # Mock via string path (before import)
+        monkeypatch.setattr(
+            "vllm_omni.entrypoints.omni.initialize_orchestrator_connectors",
+            _fake_initialize_orchestrator_connectors,
+            raising=False,
+        )
+
+
+def _setup_connector_adapter_mock(monkeypatch, omni_module):
+    """Helper function to mock try_send_via_connector on the omni module.
+
+    This must be called AFTER importing omni module, to mock the function where it's actually used.
+    """
+
+    # Mock try_send_via_connector to always succeed
+    def _fake_try_send_via_connector(
+        connector,
+        stage_id,
+        next_stage_id,
+        req_id,
+        next_inputs,
+        sampling_params,
+        original_prompt,
+        next_stage_queue_submit_fn,
+        metrics,
+    ):
+        # Simulate successful send by calling the submit function
+        task = {
+            "request_id": req_id,
+            "engine_inputs": next_inputs,
+            "sampling_params": sampling_params,
+        }
+        next_stage_queue_submit_fn(task)
+        return True
+
+    # Mock directly on the omni module where it's used
+    monkeypatch.setattr(omni_module, "try_send_via_connector", _fake_try_send_via_connector)
+
+
 @pytest.fixture(autouse=True)
 def mock_get_config(monkeypatch):
     """Auto-mock get_config and related model loading functions to avoid model path validation."""
@@ -450,10 +572,19 @@ def mock_get_config(monkeypatch):
 def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config):
     """Test that stage configs are auto-loaded when stage_configs_path is None."""
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [
-            _FakeStageConfig(fake_stage_config),
-            _FakeStageConfig(fake_stage_config),
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        cfg0 = dict(fake_stage_config)
+        cfg0["stage_id"] = 0
+        cfg1 = dict(fake_stage_config)
+        cfg1["stage_id"] = 1
+        return None, [
+            _FakeStageConfig(cfg0),
+            _FakeStageConfig(cfg1),
         ]
 
     # Remove modules from cache BEFORE setting mocks
@@ -472,10 +603,11 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
-    # Mock load_stage_configs_from_model
+    # Mock load_and_resolve_stage_configs
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -491,12 +623,14 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
     import vllm_omni.entrypoints.omni as omni_module
 
     # Patch the imported function and class in the module
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, omni_module)
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
     # Verify: auto-loaded stage_configs and stage_list have consistent count
     assert isinstance(omni.stage_configs, list)
     assert len(omni.stage_configs) == 2
@@ -515,8 +649,15 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
 def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
     """Test that generate raises ValueError when sampling_params_list length doesn't match."""
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        cfg0 = dict(fake_stage_config)
+        cfg0["stage_id"] = 0
+        return None, [_FakeStageConfig(cfg0)]
 
     import sys
 
@@ -532,14 +673,10 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-        _fake_loader,
-        raising=False,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -551,12 +688,12 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
     with pytest.raises(ValueError):
         omni.generate(prompts=["hi"], sampling_params_list=[])
 
@@ -564,11 +701,18 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
 def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
     """Test multi-stage generation pipeline with queue polling."""
     stage_cfg0 = dict(fake_stage_config)
+    stage_cfg0["stage_id"] = 0
     stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["stage_id"] = 1
     stage_cfg1["processed_input"] = ["processed-for-stage-1"]
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
 
@@ -584,14 +728,10 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
-        _fake_loader,
-        raising=False,
-    )
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -603,8 +743,11 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, omni_module)
+    _setup_connector_adapter_mock(monkeypatch, omni_module)
 
     # Mock uuid.uuid4() to return a predictable value for request ID generation
     test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -613,7 +756,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
 
     # Generate the expected request ID format: "0_<uuid>"
     expected_request_id = f"0_{test_uuid}"
@@ -663,12 +806,19 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
 def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
     """Test that generate returns empty list when all stages have final_output=False."""
     stage_cfg0 = dict(fake_stage_config)
-    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg0["stage_id"] = 0
     stage_cfg0["final_output"] = False
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["stage_id"] = 1
     stage_cfg1["final_output"] = False
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
 
@@ -684,9 +834,10 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -698,8 +849,11 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, omni_module)
+    _setup_connector_adapter_mock(monkeypatch, omni_module)
 
     # Mock uuid.uuid4() to return a predictable value for request ID generation
     test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -708,7 +862,7 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
 
     # Generate the expected request ID format: "0_<uuid>"
     expected_request_id = f"0_{test_uuid}"
@@ -742,12 +896,19 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
 def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_config):
     """Test that generate uses default sampling params when sampling_params_list is None."""
     stage_cfg0 = dict(fake_stage_config)
-    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg0["stage_id"] = 0
     stage_cfg0["final_output"] = False
+    stage_cfg1 = dict(fake_stage_config)
+    stage_cfg1["stage_id"] = 1
     stage_cfg1["final_output"] = False
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)]
 
     import sys
 
@@ -763,9 +924,10 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -777,8 +939,11 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
+    # Apply connector and adapter mocks after importing omni module
+    _setup_connector_mocks(monkeypatch, omni_module)
+    _setup_connector_adapter_mock(monkeypatch, omni_module)
 
     # Mock uuid.uuid4() to return a predictable value for request ID generation
     test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -787,7 +952,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
 
     # Generate the expected request ID format: "0_<uuid>"
     expected_request_id = f"0_{test_uuid}"
@@ -814,8 +979,15 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
 def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
     """Test that _wait_for_stages_ready handles timeout correctly."""
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        cfg0 = dict(fake_stage_config)
+        cfg0["stage_id"] = 0
+        return None, [_FakeStageConfig(cfg0)]
 
     import sys
 
@@ -831,9 +1003,10 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -856,13 +1029,13 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(cfg, **kwargs))
 
     from vllm_omni.entrypoints.omni import Omni
 
     # Use very short timeout
-    omni = Omni(model="any", init_timeout=0.01)
+    omni = Omni(model=MODEL, init_timeout=0.01)
     # Verify that no stages are ready
     assert len(omni._stages_ready) == 0
 
@@ -870,8 +1043,15 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
 def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     """Test that generate handles error messages from stages correctly."""
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        cfg0 = dict(fake_stage_config)
+        cfg0["stage_id"] = 0
+        return None, [_FakeStageConfig(cfg0)]
 
     import sys
 
@@ -887,9 +1067,10 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -901,7 +1082,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
 
     # Mock uuid.uuid4() to return a predictable value for request ID generation
@@ -911,7 +1092,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
 
     # Generate the expected request ID format: "0_<uuid>"
     expected_request_id = f"0_{test_uuid}"
@@ -945,8 +1126,15 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
 def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
     """Test that close() sends shutdown signal to all input queues."""
 
-    def _fake_loader(model: str, base_engine_args=None):
-        return [_FakeStageConfig(fake_stage_config)]
+    def _fake_loader(
+        model: str,
+        stage_configs_path: str | None = None,
+        base_engine_args: dict | None = None,
+        default_stage_cfg_factory=None,
+    ):
+        cfg0 = dict(fake_stage_config)
+        cfg0["stage_id"] = 0
+        return None, [_FakeStageConfig(cfg0)]
 
     import sys
 
@@ -962,9 +1150,10 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
     _setup_multiprocessing_mocks(monkeypatch)
     _setup_ipc_mocks(monkeypatch)
     _setup_log_mocks(monkeypatch)
+    _setup_connector_mocks(monkeypatch)
 
     monkeypatch.setattr(
-        "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
+        "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
         _fake_loader,
         raising=False,
     )
@@ -976,12 +1165,12 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
 
     import vllm_omni.entrypoints.omni as omni_module
 
-    monkeypatch.setattr(omni_module, "load_stage_configs_from_model", _fake_loader)
+    monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
     monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
 
     from vllm_omni.entrypoints.omni import Omni
 
-    omni = Omni(model="any", init_timeout=1)
+    omni = Omni(model=MODEL, init_timeout=1)
 
     # Call close
     omni.close()

tests/entrypoints/test_utils.py

@@ -0,0 +1,271 @@
+"""Unit tests for vllm_omni.entrypoints.utils module."""
+
+from collections import Counter
+from dataclasses import dataclass
+from unittest.mock import patch
+
+import pytest
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.engine.arg_utils import OmniEngineArgs
+from vllm_omni.entrypoints.utils import (
+    _convert_dataclasses_to_dict,
+    _filter_dict_like_object,
+    filter_dataclass_kwargs,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestFilterDictLikeObject:
+    """Test suite for _filter_dict_like_object function."""
+
+    def test_simple_dict(self):
+        """Test filtering a simple dictionary with no callables."""
+        input_dict = {"key1": "value1", "key2": 42, "key3": [1, 2, 3]}
+        result = _filter_dict_like_object(input_dict)
+
+        assert result == input_dict
+        assert isinstance(result, dict)
+
+    def test_dict_with_nested_values(self):
+        """Test filtering dict with nested dict and list values."""
+        input_dict = {
+            "level1": {
+                "level2": {"key": "value"},
+                "list": [1, 2, 3],
+            },
+            "simple": "string",
+        }
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Nested dicts and lists should be recursively processed
+        assert result["simple"] == "string"
+        assert isinstance(result["level1"], dict)
+
+    def test_dict_with_dataclass_values(self):
+        """Test filtering dict containing dataclass values."""
+
+        @dataclass
+        class TestDataclass:
+            field1: str
+            field2: int
+
+        obj = TestDataclass(field1="test", field2=42)
+        input_dict = {"data": obj, "normal": "value"}
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Dataclass should be converted to dict by recursive _convert_dataclasses_to_dict
+        assert "data" in result
+        assert "normal" in result
+        assert result["normal"] == "value"
+
+    def test_dict_with_counter_values(self):
+        """Test filtering dict containing Counter objects."""
+        counter_obj = Counter({"a": 1, "b": 2})
+        input_dict = {"counter": counter_obj, "normal": "value"}
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Counter should be converted to regular dict
+        assert "counter" in result
+        assert "normal" in result
+        assert result["normal"] == "value"
+
+    def test_empty_dict(self):
+        """Test filtering an empty dictionary."""
+        result = _filter_dict_like_object({})
+        assert result == {}
+        assert isinstance(result, dict)
+
+    def test_dict_with_set_values(self):
+        """Test filtering dict with set values."""
+        input_dict = {"set_key": {1, 2, 3}, "normal": "value"}
+
+        result = _filter_dict_like_object(input_dict)
+
+        assert "set_key" in result
+        assert "normal" in result
+        # Set should be converted to list by _convert_dataclasses_to_dict
+        assert result["normal"] == "value"
+
+    def test_dict_with_none_values(self):
+        """Test filtering dict with None values."""
+        input_dict = {"key1": None, "key2": "value", "key3": 0}
+
+        result = _filter_dict_like_object(input_dict)
+
+        assert result == input_dict
+
+    def test_dict_with_mixed_types(self):
+        """Test filtering dict with mixed value types."""
+        input_dict = {
+            "string": "hello",
+            "int": 42,
+            "float": 3.14,
+            "bool": True,
+            "none": None,
+            "list": [1, 2, 3],
+            "tuple": (1, 2, 3),
+            "set": {1, 2, 3},
+            "dict": {"nested": "value"},
+        }
+
+        result = _filter_dict_like_object(input_dict)
+
+        assert "string" in result
+        assert "int" in result
+        assert "float" in result
+        assert "bool" in result
+        assert "none" in result
+        assert "list" in result
+        assert "tuple" in result
+        assert "set" in result
+        assert "dict" in result
+
+    def test_dict_preserves_key_types(self):
+        """Test that dict key types are preserved."""
+        input_dict = {
+            "string_key": "value1",
+            42: "value2",
+            (1, 2): "value3",  # tuple as key
+        }
+
+        result = _filter_dict_like_object(input_dict)
+
+        # Keys should remain the same
+        assert "string_key" in result
+        assert 42 in result
+        assert (1, 2) in result
+
+    def test_dict_with_recursive_structure(self):
+        """Test filtering dict with recursive/complex nested structure."""
+        input_dict = {
+            "level1": {
+                "level2": {
+                    "level3": {"key": "value"},
+                    "callable": lambda x: x,
+                }
+            },
+            "normal": "value",
+        }
+
+        with patch("vllm_omni.entrypoints.utils.logger"):
+            result = _filter_dict_like_object(input_dict)
+
+            # Normal key should exist
+            assert "normal" in result
+            # Level1 should exist
+            assert "level1" in result
+
+    def test_integration_with_convert_dataclasses(self):
+        """Test that _filter_dict_like_object integrates properly with _convert_dataclasses_to_dict."""
+
+        @dataclass
+        class Config:
+            name: str
+            count: int
+
+        input_dict = {
+            "config": Config(name="test", count=5),
+            "func": lambda x: x,
+            "normal": "value",
+        }
+
+        with patch("vllm_omni.entrypoints.utils.logger"):
+            result = _filter_dict_like_object(input_dict)
+
+            # Callable should be filtered
+            assert "func" not in result
+            # Config should be converted to dict
+            assert "config" in result
+            assert "normal" in result
+
+
+class TestConvertDataclassesToDict:
+    """Test suite for _convert_dataclasses_to_dict function."""
+
+    def test_uses_filter_dict_like_object(self):
+        """Test that _convert_dataclasses_to_dict uses _filter_dict_like_object for dicts."""
+        input_dict = {
+            "normal": "value",
+            "callable": lambda x: x,
+        }
+
+        with patch("vllm_omni.entrypoints.utils.logger"):
+            result = _convert_dataclasses_to_dict(input_dict)
+
+            # Callable should be filtered out by _filter_dict_like_object
+            assert "normal" in result
+            assert "callable" not in result
+
+
+class TestFilterDataclassKwargs:
+    """Test basic functionality of filter_dataclass_kwargs."""
+
+    def test_simple_filtering(self):
+        """Test basic dataclass kwargs filtering."""
+
+        @dataclass
+        class SimpleConfig:
+            name: str
+            count: int
+
+        kwargs = {"name": "test", "count": 42, "invalid": "should_be_removed"}
+        result = filter_dataclass_kwargs(SimpleConfig, kwargs)
+
+        assert "name" in result
+        assert "count" in result
+        assert "invalid" not in result
+
+    def test_invalid_dataclass_raises_error(self):
+        """Test that non-dataclass raises ValueError."""
+        with pytest.raises(ValueError, match="is not a dataclass"):
+            filter_dataclass_kwargs(dict, {})
+
+    def test_invalid_kwargs_type_raises_error(self):
+        """Test that non-dict kwargs raises ValueError."""
+
+        @dataclass
+        class SimpleConfig:
+            name: str
+
+        with pytest.raises(ValueError, match="kwargs must be a dictionary"):
+            filter_dataclass_kwargs(SimpleConfig, "invalid")
+
+    def test_filters_omni_engine_args_unknown_fields(self):
+        """Test that OmniEngineArgs kwargs are filtered to valid fields only."""
+        kwargs = {
+            "model": "dummy",
+            "stage_id": 1,
+            "engine_output_type": "image",
+            "unknown_field": "drop_me",
+        }
+
+        result = filter_dataclass_kwargs(OmniEngineArgs, kwargs)
+
+        assert "model" in result
+        assert "stage_id" in result
+        assert "engine_output_type" in result
+        assert "unknown_field" not in result
+
+    def test_filters_omni_diffusion_config_union_dataclass(self):
+        """Test that OmniDiffusionConfig filters nested dataclass in Union fields."""
+        kwargs = {
+            "model": "dummy",
+            "cache_config": {
+                "rel_l1_thresh": 0.3,
+                "extra_param": "should_drop",
+            },
+            "unknown_top": "drop_me",
+        }
+
+        result = filter_dataclass_kwargs(OmniDiffusionConfig, kwargs)
+
+        assert "model" in result
+        assert "cache_config" in result
+        assert "unknown_top" not in result
+        assert result["cache_config"]["rel_l1_thresh"] == 0.3
+        assert "extra_param" not in result["cache_config"]

vllm_omni/entrypoints/async_omni.py

@@ -38,7 +38,7 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)
 
 
-def _weak_close_cleanup_async(stage_list, stage_in_queues, ray_pg, output_handler):
+def _weak_close_cleanup_async(stage_list, stage_in_queues, stage_out_queues, ray_pg, output_handler, zmq_ctx=None):
     """Weak reference cleanup function for AsyncOmni instances."""
     if stage_list:
         for q in stage_in_queues:
@@ -46,6 +46,13 @@ def _weak_close_cleanup_async(stage_list, stage_in_queues, ray_pg, output_handle
                 q.put_nowait(SHUTDOWN_TASK)
             except Exception as e:
                 logger.warning(f"Failed to send shutdown signal to stage input queue: {e}")
+            close_fn = getattr(q, "close", None)
+            if callable(close_fn):
+                close_fn()
+        for q in stage_out_queues:
+            close_fn = getattr(q, "close", None)
+            if callable(close_fn):
+                close_fn()
         for stage in stage_list:
             try:
                 stage.stop_stage_worker()
@@ -55,6 +62,8 @@ def _weak_close_cleanup_async(stage_list, stage_in_queues, ray_pg, output_handle
     # Cancel output handler
     if output_handler is not None:
         output_handler.cancel()
+    if zmq_ctx is not None:
+        zmq_ctx.term()
 
 
 class AsyncOmni(OmniBase):
@@ -108,8 +117,10 @@ class AsyncOmni(OmniBase):
             _weak_close_cleanup_async,
             self.stage_list,
             self._stage_in_queues,
+            self._stage_out_queues,
             self._ray_pg,
             self.output_handler,
+            self._zmq_ctx,
         )
 
     def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]:

vllm_omni/entrypoints/cli/serve.py

@@ -7,19 +7,37 @@ diffusion models (e.g., Qwen-Image) through the same CLI interface.
 
 import argparse
 import os
+import signal
+from typing import Any
 
+import msgspec.msgpack
 import uvloop
+import zmq
 from vllm.entrypoints.cli.types import CLISubcommand
 from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
 from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
 from vllm.logger import init_logger
 from vllm.utils.argparse_utils import FlexibleArgumentParser
+from vllm.utils.network_utils import make_zmq_socket
+from vllm.v1.utils import get_engine_client_zmq_addr
 
+from vllm_omni.distributed.omni_connectors import (
+    get_connectors_config_for_stage,
+    load_omni_transfer_config,
+)
+from vllm_omni.distributed.omni_connectors.utils.initialization import (
+    resolve_omni_kv_config_for_stage,
+)
 from vllm_omni.entrypoints.cli.logo import log_logo
+from vllm_omni.entrypoints.omni import OmniBase, omni_snapshot_download
+from vllm_omni.entrypoints.omni_stage import OmniStage
 from vllm_omni.entrypoints.openai.api_server import omni_run_server
+from vllm_omni.entrypoints.utils import inject_omni_kv_config
 
 logger = init_logger(__name__)
 
+HANDSHAKE_TIMEOUT_MINS = 5
+
 DESCRIPTION = """Launch a local OpenAI-compatible API server to serve Omni models
 via HTTP. Supports both multi-stage LLM models and diffusion models.
 
@@ -55,9 +73,15 @@ class OmniServeCommand(CLISubcommand):
         if hasattr(args, "model_tag") and args.model_tag is not None:
             args.model = args.model_tag
 
-        uvloop.run(omni_run_server(args))
+        if args.headless:
+            run_headless(args)
+        else:
+            uvloop.run(omni_run_server(args))
 
     def validate(self, args: argparse.Namespace) -> None:
+        if args.stage_id is not None and (args.omni_master_address is None or args.omni_master_port is None):
+            raise ValueError("--stage-id requires both --omni-master-address and --omni-master-port to be set")
+
         # Skip validation for diffusion models as they have different requirements
         from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model
 
@@ -94,6 +118,12 @@ class OmniServeCommand(CLISubcommand):
             default=None,
             help="Path to the stage configs file. If not specified, the stage configs will be loaded from the model.",
         )
+        omni_config_group.add_argument(
+            "--stage-id",
+            type=int,
+            default=None,
+            help="Select and launch a single stage by stage_id.",
+        )
         omni_config_group.add_argument(
             "--stage-init-timeout",
             type=int,
@@ -142,6 +172,18 @@ class OmniServeCommand(CLISubcommand):
             default=None,
             help="The address of the Ray cluster to connect to.",
         )
+        omni_config_group.add_argument(
+            "--omni-master-address",
+            "-oma",
+            type=str,
+            help="Hostname or IP address of the Omni orchestrator (master).",
+        )
+        omni_config_group.add_argument(
+            "--omni-master-port",
+            "-omp",
+            type=int,
+            help="Port of the Omni orchestrator (master).",
+        )
 
         # Diffusion model specific arguments
         omni_config_group.add_argument(
@@ -251,5 +293,125 @@ class OmniServeCommand(CLISubcommand):
         return serve_parser
 
 
+def _create_default_diffusion_stage_cfg(args: argparse.Namespace) -> list[dict[str, Any]]:
+    omni_base = OmniBase.__new__(OmniBase)
+    return omni_base._create_default_diffusion_stage_cfg(vars(args))
+
+
+def run_headless(args: argparse.Namespace) -> None:
+    if args.api_server_count is not None and args.api_server_count > 1:
+        raise ValueError("api_server_count can't be set in headless mode")
+    if args.worker_backend != "multi_process":
+        raise ValueError("headless mode requires worker_backend=multi_process")
+
+    model = omni_snapshot_download(args.model)
+
+    omni_base = OmniBase.__new__(OmniBase)
+    args_dict = vars(args).copy()
+    args_dict["model"] = model
+    config_path, stage_configs = omni_base._resolve_stage_configs(model, args_dict)
+
+    single_stage_id = args.stage_id
+    if single_stage_id is None:
+        if len(stage_configs) != 1:
+            raise ValueError("--stage-id is required in headless mode for multi-stage configs")
+        single_stage_id = getattr(stage_configs[0], "stage_id", 0)
+
+    stage_config = None
+    for cfg in stage_configs:
+        if getattr(cfg, "stage_id", None) == single_stage_id:
+            stage_config = cfg
+            break
+    if stage_config is None:
+        raise ValueError(f"No stage matches stage_id={single_stage_id}.")
+
+    # TODO(wuhang): Support connectors config by cli
+    transfer_config = load_omni_transfer_config(config_path, default_shm_threshold=args.shm_threshold_bytes)
+    connectors_config = get_connectors_config_for_stage(transfer_config, single_stage_id)
+
+    omni_master_address = args.omni_master_address
+    omni_master_port = args.omni_master_port
+
+    # Perform handshake with orchestrator to get dynamically allocated endpoints
+    with zmq.Context() as zmq_ctx:
+        handshake_endpoint = get_engine_client_zmq_addr(
+            local_only=False, host=omni_master_address, port=omni_master_port
+        )
+
+        with make_zmq_socket(zmq_ctx, handshake_endpoint, zmq.REQ, bind=False, linger=5000) as handshake_socket:
+            # TODO(wuhang): Define protocol in python dataclass.
+            handshake_msg = {"type": "handshake", "stage_id": single_stage_id}
+            handshake_socket.send(msgspec.msgpack.encode(handshake_msg))
+
+            # Wait for response with timeout
+            if not handshake_socket.poll(timeout=HANDSHAKE_TIMEOUT_MINS * 60_000):
+                raise RuntimeError(
+                    f"Handshake timeout ({HANDSHAKE_TIMEOUT_MINS} minutes) for stage-{single_stage_id} "
+                    f"at {handshake_endpoint}"
+                )
+
+            try:
+                response = msgspec.msgpack.decode(handshake_socket.recv())
+            except msgspec.DecodeError as exc:
+                raise RuntimeError(
+                    f"Handshake decode failed for stage-{single_stage_id} at {handshake_endpoint}: {exc}"
+                ) from exc
+            except Exception as exc:  # pragma: no cover - unexpected decode errors
+                raise RuntimeError(
+                    f"Unexpected error decoding handshake for stage-{single_stage_id} at {handshake_endpoint}: {exc}"
+                ) from exc
+
+            if not response["ok"]:
+                error_msg = response["error"]
+                raise RuntimeError(f"Handshake failed for stage-{single_stage_id}: {error_msg}")
+
+            in_endpoint, out_endpoint = response["in_endpoint"], response["out_endpoint"]
+
+            logger.info(
+                f"[Headless] Stage-{single_stage_id} received endpoints via handshake: "
+                f"in={in_endpoint}, out={out_endpoint}"
+            )
+
+    shutdown_requested = False
+
+    def signal_handler(signum, frame):
+        nonlocal shutdown_requested
+        if shutdown_requested:
+            return
+        shutdown_requested = True
+        raise SystemExit
+
+    signal.signal(signal.SIGTERM, signal_handler)
+    signal.signal(signal.SIGINT, signal_handler)
+
+    stage = OmniStage(stage_config, stage_init_timeout=args.stage_init_timeout)
+    stage.attach_queues(in_endpoint, out_endpoint)
+
+    # Inject YAML-resolved connector config into omni_kv_config for in-engine usage.
+    try:
+        omni_conn_cfg, omni_from, omni_to = resolve_omni_kv_config_for_stage(transfer_config, single_stage_id)
+        if omni_conn_cfg:
+            inject_omni_kv_config(stage, omni_conn_cfg, omni_from, omni_to)  # type: ignore
+    except Exception as e:
+        logger.debug("[Headless] Failed to inject omni connector config into stage-%s: %s", single_stage_id, e)
+
+    old_env = os.environ.get("VLLM_LOGGING_PREFIX")
+    os.environ["VLLM_LOGGING_PREFIX"] = f"[Stage-{single_stage_id}] {'' if old_env is None else old_env}"
+    try:
+        stage.init_stage_worker(
+            model,
+            is_async=True,
+            shm_threshold_bytes=int(args.shm_threshold_bytes),
+            batch_timeout=int(args.batch_timeout),
+            connectors_config=connectors_config,
+            worker_backend="multi_process",
+            ignore_runtime_config=True,
+        )
+        if stage._proc is not None:
+            stage._proc.join()
+    finally:
+        stage.stop_stage_worker()
+
+
 def cmd_init() -> list[CLISubcommand]:
     return [OmniServeCommand()]

vllm_omni/entrypoints/omni.py

@@ -3,6 +3,7 @@
 import json
 import multiprocessing as mp
 import os
+import threading
 import time
 import uuid
 import weakref
@@ -11,10 +12,14 @@ from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Any, Literal, overload
 
 import huggingface_hub
+import msgspec.msgpack
+import zmq
 from omegaconf import OmegaConf
 from tqdm.auto import tqdm
 from vllm import SamplingParams
 from vllm.logger import init_logger
+from vllm.utils.network_utils import make_zmq_socket
+from vllm.v1.utils import get_engine_client_zmq_addr
 
 from vllm_omni.distributed.omni_connectors import (
     get_stage_connector_config,
@@ -35,10 +40,9 @@ from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     get_final_stage_id_for_e2e,
     inject_omni_kv_config,
-    load_stage_configs_from_model,
-    load_stage_configs_from_yaml,
-    resolve_model_config_path,
+    load_and_resolve_stage_configs,
 )
+from vllm_omni.entrypoints.zmq_utils import ZmqQueue
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams
 from vllm_omni.metrics import OrchestratorAggregator, StageRequestStats
 from vllm_omni.model_executor.model_loader.weight_utils import (
@@ -49,7 +53,16 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)
 
 
-def _weak_close_cleanup(stage_list, stage_in_queues, ray_pg):
+def _weak_close_cleanup(
+    stage_list,
+    stage_in_queues,
+    stage_out_queues,
+    ray_pg,
+    zmq_ctx=None,
+    handshake_stop: threading.Event | None = None,
+    zmq_handshake_socket: zmq.Socket | None = None,
+    handshake_thread: threading.Thread | None = None,
+):
     """Weak reference cleanup function for OmniBase instances."""
     if stage_list:
         for q in stage_in_queues:
@@ -57,6 +70,13 @@ def _weak_close_cleanup(stage_list, stage_in_queues, ray_pg):
                 q.put_nowait(SHUTDOWN_TASK)
             except Exception as e:
                 logger.warning(f"Failed to send shutdown signal to stage input queue: {e}")
+            close_fn = getattr(q, "close", None)
+            if callable(close_fn):
+                close_fn()
+        for q in stage_out_queues:
+            close_fn = getattr(q, "close", None)
+            if callable(close_fn):
+                close_fn()
         for stage in stage_list:
             try:
                 stage.stop_stage_worker()
@@ -64,6 +84,20 @@ def _weak_close_cleanup(stage_list, stage_in_queues, ray_pg):
                 logger.warning(f"Failed to stop stage worker: {e}")
     try_close_ray(ray_pg)
 
+    # Gracefully shutdown handshake server thread
+    if handshake_stop is not None:
+        handshake_stop.set()
+    if handshake_thread is not None:
+        handshake_thread.join(timeout=2.0)
+        if handshake_thread.is_alive():
+            logger.warning("Handshake server thread did not terminate gracefully within timeout")
+
+    # Close ZMQ resources after thread has exited
+    if zmq_handshake_socket is not None:
+        zmq_handshake_socket.close(0)
+    if zmq_ctx is not None:
+        zmq_ctx.term()
+
 
 def _dummy_snapshot_download(model_id):
     return model_id
@@ -125,12 +159,21 @@ class OmniBase:
 
         # Stage management attributes
         self.stage_list: list[OmniStage] = []
-        self._stage_in_queues: list[mp.Queue] = []
-        self._stage_out_queues: list[mp.Queue] = []
+        self._stage_in_queues: list[Any] = []
+        self._stage_out_queues: list[Any] = []
         self._stages_ready: set[int] = set()
         self._ray_pg = None
         self._queue_cls = None
         self._ctx = None
+        self._zmq_ctx: zmq.Context | None = None
+        self._zmq_master_address: str | None = None
+        self._zmq_master_port: int | None = None
+        self._zmq_handshake_socket: zmq.Socket | None = None
+        self._handshake_thread: threading.Thread | None = None
+        self._handshake_stop: threading.Event | None = None
+        self._handshake_endpoints: dict[int, tuple[str, str]] = {}
+        self._handshake_seen: set[int] = set()  # Track which stage IDs have completed ZMQ handshake
+        self._single_stage_id: int | None = None  # Optional: deploy only a specific stage ID
 
         # Initialize stages - each stage will create appropriate instance based on stage_type
         # Stage workers will automatically create OmniLLM or OmniDiffusion instances
@@ -206,35 +249,24 @@ class OmniBase:
         default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion"
         return default_stage_cfg
 
-    def _initialize_stages(self, model: str, kwargs: dict[str, Any]) -> None:
-        """Initialize stage list management."""
-        stage_init_timeout = kwargs.get("stage_init_timeout", 20)
-        shm_threshold_bytes = kwargs.get("shm_threshold_bytes", 65536)
-        init_timeout = kwargs.get("init_timeout", 300)
-        worker_backend = kwargs.get("worker_backend", "multi_process")
-        ray_address = kwargs.get("ray_address", None)
-        batch_timeout = kwargs.get("batch_timeout", 10)
+    def _resolve_stage_configs(self, model: str, kwargs: dict[str, Any]) -> tuple[str, list[Any]]:
+        """Resolve stage configs and inject defaults shared by orchestrator/headless."""
+        # TODO(wuhang):
+        # Remove kwargs as parameters in the future.
+        # Use dataclass directly for engine args.
+
         stage_configs_path = kwargs.get("stage_configs_path", None)
-        log_stats = kwargs.get("log_stats", False)
-
-        ### base engine args
-        tokenizer = kwargs.get("tokenizer", None)
-
-        base_engine_args = {"tokenizer": tokenizer} if tokenizer is not None else None
 
         # Load stage configurations from YAML
-        if stage_configs_path is None:
-            self.config_path = resolve_model_config_path(model)
-            self.stage_configs = load_stage_configs_from_model(model, base_engine_args=base_engine_args)
-            if not self.stage_configs:
-                default_stage_cfg = self._create_default_diffusion_stage_cfg(kwargs)
-                self.stage_configs = OmegaConf.create(default_stage_cfg)
-        else:
-            self.config_path = stage_configs_path
-            self.stage_configs = load_stage_configs_from_yaml(stage_configs_path, base_engine_args=base_engine_args)
+        config_path, stage_configs = load_and_resolve_stage_configs(
+            model,
+            stage_configs_path,
+            kwargs,
+            default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs),
+        )
 
         # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config.
-        for cfg in self.stage_configs:
+        for cfg in stage_configs:
             try:
                 if getattr(cfg, "stage_type", None) != "diffusion":
                     continue
@@ -253,6 +285,27 @@ class OmniBase:
             except Exception as e:
                 logger.warning("Failed to inject LoRA config for stage: %s", e)
 
+        return config_path, stage_configs
+
+    def _initialize_stages(self, model: str, kwargs: dict[str, Any]) -> None:
+        """Initialize stage list management."""
+        stage_init_timeout = kwargs.get("stage_init_timeout", 20)
+        shm_threshold_bytes = kwargs.get("shm_threshold_bytes", 65536)
+        init_timeout = kwargs.get("init_timeout", 300)
+        worker_backend = kwargs.get("worker_backend", "multi_process")
+        ray_address = kwargs.get("ray_address", None)
+        batch_timeout = kwargs.get("batch_timeout", 10)
+        log_stats = kwargs.get("log_stats", False)
+        self._single_stage_id = kwargs.get("stage_id", None)
+        self._zmq_master_address = kwargs.get("omni_master_address", None)
+        if self._zmq_master_address is None:
+            self._zmq_master_address = "127.0.0.1"
+            logger.info("No omni_master_address provided, defaulting to localhost (127.0.0.1)")
+        self._zmq_master_port = kwargs.get("omni_master_port", None)
+
+        # Resolve stage configs shared by orchestrator/headless paths.
+        self.config_path, self.stage_configs = self._resolve_stage_configs(model, kwargs)
+
         # Initialize connectors
         self.omni_transfer_config, self.connectors = initialize_orchestrator_connectors(
             self.config_path, worker_backend=worker_backend, shm_threshold_bytes=shm_threshold_bytes
@@ -281,7 +334,7 @@ class OmniBase:
         self.stage_list = [st for _, st in results]
         self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list]
         self.output_modalities = [st.final_output_type for st in self.stage_list]
-        logger.debug(f"[{self._name}] Loaded {len(self.stage_list)} stages")
+        logger.info(f"[{self._name}] Loaded {len(self.stage_list)} stages")
 
         if self.worker_backend == "ray":
             self._queue_cls = get_ray_queue_class()
@@ -307,10 +360,34 @@ class OmniBase:
             self._ray_pg = create_placement_group(
                 number_of_stages=len(self.stage_list), address=self.ray_address, strategy="PACK"
             )
+        else:
+            # Initialize ZMQ context
+            if self._zmq_ctx is None:
+                self._zmq_ctx = zmq.Context()
+
+            # Allocate endpoints for each stage
+            total_stages = len(self.stage_configs)
+            self._handshake_endpoints = {}
+            for sid in range(total_stages):
+                in_endpoint = get_engine_client_zmq_addr(local_only=False, host=self._zmq_master_address)
+                out_endpoint = get_engine_client_zmq_addr(local_only=False, host=self._zmq_master_address)
+                self._handshake_endpoints[sid] = (in_endpoint, out_endpoint)
+                logger.debug(
+                    f"[{self._name}] Allocated endpoints for stage-{sid}: in={in_endpoint}, out={out_endpoint}"
+                )
+
+            # Start handshake server
+            self.start_handshake_server()
 
         for stage_id, stage in enumerate[OmniStage](self.stage_list):
-            in_q = self._queue_cls()
-            out_q = self._queue_cls()
+            if self.worker_backend == "ray":
+                in_q = self._queue_cls()
+                out_q = self._queue_cls()
+            else:
+                in_endpoint, out_endpoint = self._handshake_endpoints[stage_id]
+                in_q = ZmqQueue(self._zmq_ctx, zmq.PUSH, bind=in_endpoint)
+                out_q = ZmqQueue(self._zmq_ctx, zmq.PULL, bind=out_endpoint)
+
             self._stage_in_queues.append(in_q)
             self._stage_out_queues.append(out_q)
             stage.attach_queues(in_q, out_q)
@@ -332,6 +409,12 @@ class OmniBase:
             except Exception as e:
                 logger.debug("[Omni] Failed to inject omni connector config into stage-%s: %s", stage_id, e)
 
+            if self._single_stage_id is not None and stage_id != int(self._single_stage_id):
+                logger.info(
+                    f"[{self._name}] Skipping initialization of stage-{stage_id} worker due to single_stage_id setting"
+                )
+                continue
+
             stage.init_stage_worker(
                 model,
                 is_async=self.is_async,
@@ -341,6 +424,7 @@ class OmniBase:
                 connectors_config=stage_connectors_config,
                 worker_backend=self.worker_backend,
                 ray_placement_group=self._ray_pg,
+                ignore_runtime_config=True if self._single_stage_id is not None else False,
             )
 
             logger.debug(f"[{self._name}] Stage-{stage_id} process started")
@@ -351,6 +435,9 @@ class OmniBase:
 
     def _wait_for_stages_ready(self, timeout: int = 120) -> None:
         """Wait for all stages to report readiness with optimized polling."""
+        if self._single_stage_id is not None and self.worker_backend != "ray":
+            timeout = self._wait_for_handshakes(timeout)
+
         num_stages = len(self.stage_list)
         deadline = time.time() + max(0, int(timeout))
 
@@ -501,6 +588,129 @@ class OmniBase:
         if hasattr(self, "_weak_finalizer"):
             self._weak_finalizer()
 
+    def _process_handshake_message(self, msg: Any) -> dict[str, Any]:
+        """Process incoming handshake message and generate response.
+
+        Args:
+            msg: Decoded message from client
+
+        Returns:
+            Response dictionary with ok status and either endpoints or error
+        """
+        if not isinstance(msg, dict) or msg.get("type") != "handshake":
+            return {"ok": False, "error": "invalid handshake payload"}
+
+        try:
+            stage_id = int(msg.get("stage_id"))
+        except (TypeError, ValueError) as e:
+            return {"ok": False, "error": f"invalid stage_id: {e}"}
+
+        endpoints = self._handshake_endpoints.get(stage_id)
+        if endpoints is None:
+            return {"ok": False, "error": f"unknown stage_id: {stage_id}"}
+
+        # Mark stage as seen and prepare success response
+        self._handshake_seen.add(stage_id)
+        in_endpoint, out_endpoint = endpoints
+
+        logger.info(
+            "[%s] Handshake received from stage-%s",
+            self._name,
+            stage_id,
+        )
+
+        return {
+            "ok": True,
+            "in_endpoint": in_endpoint,
+            "out_endpoint": out_endpoint,
+        }
+
+    def _run_handshake_server_loop(self) -> None:
+        """Main loop for handshake server - polls for messages and responds."""
+        poller = zmq.Poller()
+        poller.register(self._zmq_handshake_socket, zmq.POLLIN)
+
+        try:
+            while not self._handshake_stop.is_set():
+                events = poller.poll(1000)
+                has_message = any(sock == self._zmq_handshake_socket and event == zmq.POLLIN for sock, event in events)
+                if not has_message:
+                    continue
+
+                msg = msgspec.msgpack.decode(self._zmq_handshake_socket.recv())
+                response = msgspec.msgpack.encode(self._process_handshake_message(msg))
+                self._zmq_handshake_socket.send(response)
+        finally:
+            poller.unregister(self._zmq_handshake_socket)
+
+    def start_handshake_server(self) -> None:
+        """Start the ZMQ handshake server.
+
+        The handshake server allows distributed stages to discover their
+        queue endpoints by querying the orchestrator with their stage_id.
+        Skips starting if the server is already running or ZMQ is not initialized.
+        """
+        # Skip if already running or ZMQ not initialized
+        if self._handshake_thread is not None or self._zmq_ctx is None:
+            return
+
+        # Skip if master address/port not configured
+        if not self._zmq_master_address or self._zmq_master_port is None:
+            return
+
+        # Create server endpoint and socket
+        endpoint = get_engine_client_zmq_addr(
+            local_only=False, host=self._zmq_master_address, port=int(self._zmq_master_port)
+        )
+
+        self._handshake_stop = threading.Event()
+        self._zmq_handshake_socket = make_zmq_socket(self._zmq_ctx, endpoint, zmq.REP, bind=True, linger=5000)
+
+        # Start server thread
+        self._handshake_thread = threading.Thread(
+            target=self._run_handshake_server_loop, daemon=True, name="zmq-handshake-server"
+        )
+        self._handshake_thread.start()
+
+    def _wait_for_handshakes(self, timeout: int = 120) -> int:
+        """Wait for handshakes from all expected stages.
+
+        Args:
+            timeout: Timeout in seconds for waiting for handshakes. Default is 120s.
+
+        Returns:
+            Remaining timeout in seconds after waiting for handshakes.
+        """
+        total_stages = len(self.stage_configs)
+        expected = set(range(total_stages)) - {int(self._single_stage_id)}
+        if not expected:
+            return timeout
+
+        deadline = time.time() + max(0, int(timeout))
+        logger.info(f"[{self._name}] Waiting for handshakes from stages: {expected} (timeout: {timeout}s)")
+
+        # NOTE: _handshake_seen may be updated from the handshake server thread.
+        # It is intentionally used here without additional locking because:
+        #   - _handshake_seen only ever grows (stages are added but never removed), and
+        #   - we only check membership and set inclusion relative to `expected`.
+        # Under these monotonic semantics and the CPython GIL, concurrent reads/writes
+        # are safe for this usage and cannot violate correctness: we may observe a
+        # slightly stale view, but the loop condition remains valid and eventually
+        # becomes true once all expected stages have handshaked or the timeout elapses.
+        while not expected.issubset(self._handshake_seen) and time.time() < deadline:
+            time.sleep(1.0)
+
+        remaining_timeout = max(0, int(deadline - time.time()))
+
+        if not expected.issubset(self._handshake_seen):
+            missing = sorted(expected - self._handshake_seen)
+            logger.warning(
+                f"[{self._name}] Handshake timeout: {len(self._handshake_seen)}/{len(expected)} "
+                f"stages completed handshake. Missing stages: {missing}"
+            )
+
+        return remaining_timeout
+
     @property
     def _name(self) -> str:
         return "OmniBase"
@@ -546,7 +756,12 @@ class Omni(OmniBase):
             _weak_close_cleanup,
             self.stage_list,
             self._stage_in_queues,
+            self._stage_out_queues,
             self._ray_pg,
+            self._zmq_ctx,
+            self._handshake_stop,
+            self._zmq_handshake_socket,
+            self._handshake_thread,
         )
 
     @overload

vllm_omni/entrypoints/omni_llm.py

@@ -24,6 +24,7 @@ from vllm_omni.engine.arg_utils import OmniEngineArgs
 from vllm_omni.engine.input_processor import OmniInputProcessor, reinject_omni_fields
 from vllm_omni.engine.output_processor import MultimodalOutputProcessor
 from vllm_omni.entrypoints.utils import (
+    filter_dataclass_kwargs,
     load_stage_configs_from_model,
     load_stage_configs_from_yaml,
     resolve_model_config_path,
@@ -153,7 +154,7 @@ class OmniLLM(LLM):
             compilation_config=compilation_config_instance,
             structured_outputs_config=structured_outputs_instance,
             omni_kv_config=omni_kv_config,
-            **kwargs,
+            **filter_dataclass_kwargs(OmniEngineArgs, kwargs),
         )
 
         # Create the Engine (autoselects V0 vs V1)

vllm_omni/entrypoints/omni_stage.py

@@ -36,7 +36,7 @@ from vllm_omni.distributed.omni_connectors import build_stage_connectors
 from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector
 from vllm_omni.distributed.omni_connectors.connectors.base import OmniConnectorBase
 from vllm_omni.distributed.ray_utils.utils import kill_ray_actor, start_ray_actor
-from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
+from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs, OmniEngineArgs
 from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
 from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
 from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
@@ -49,7 +49,11 @@ from vllm_omni.entrypoints.stage_utils import (
     maybe_dump_to_shm,
     set_stage_devices,
 )
-from vllm_omni.entrypoints.utils import detect_pid_host
+from vllm_omni.entrypoints.utils import detect_pid_host, filter_dataclass_kwargs
+from vllm_omni.entrypoints.zmq_utils import (
+    ZmqQueue,
+    create_zmq_queue,
+)
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType, OmniSamplingParams, OmniTokensPrompt
 from vllm_omni.metrics import count_tokens_from_outputs
 from vllm_omni.outputs import OmniRequestOutput
@@ -204,7 +208,8 @@ def _resolve_worker_cls(engine_args: dict[str, Any]) -> None:
     worker_type = engine_args.pop("worker_type", None)
     if not worker_type:
         return
-    if engine_args.get("worker_cls"):
+    worker_cls = engine_args.get("worker_cls")
+    if worker_cls is not None and worker_cls != "auto":
         return
     from vllm_omni.platforms import current_omni_platform
 
@@ -260,6 +265,12 @@ class OmniStage:
         self.is_comprehension = getattr(stage_config, "is_comprehension", False)
         # Support for different stage types: "llm" (default) or "diffusion"
         self.stage_type: Literal["llm", "diffusion"] = getattr(stage_config, "stage_type", "llm")
+        if (
+            "stage_id" in stage_config.engine_args
+            and stage_config.engine_args.stage_id != self.stage_id
+            and self.stage_id is not None
+        ):
+            stage_config.engine_args.stage_id = self.stage_id
         if hasattr(stage_config, "custom_process_input_func"):
             # Import the module specified in the config (already a full module path)
             module_path, func_name = stage_config.custom_process_input_func.rsplit(".", 1)
@@ -282,8 +293,8 @@ class OmniStage:
         except TypeError as error:
             raise TypeError(f"Invalid default_sampling_params for stage {self.stage_id}: {error}") from error
         # Runtime orchestration state (added)
-        self._in_q: mp.Queue | None = None
-        self._out_q: mp.Queue | None = None
+        self._in_q: mp.queues.Queue | ZmqQueue | str | None = None
+        self._out_q: mp.queues.Queue | ZmqQueue | str | None = None
         self._proc: mp.Process | None = None
         self._shm_threshold_bytes: int = 65536
         self._stage_init_timeout: int = stage_init_timeout
@@ -345,12 +356,16 @@ class OmniStage:
         self.engine_outputs = engine_outputs
 
     # ----------------- New Orchestration APIs -----------------
-    def attach_queues(self, in_q: mp.Queue, out_q: mp.Queue) -> None:
+    def attach_queues(
+        self,
+        in_q: mp.queues.Queue | ZmqQueue | str | None,
+        out_q: mp.queues.Queue | ZmqQueue | str | None,
+    ) -> None:
         """Attach input and output queues for IPC communication.
 
         Args:
-            in_q: Input queue for receiving tasks from orchestrator
-            out_q: Output queue for sending results to orchestrator
+            in_q: Input queue for receiving tasks from orchestrator (queue object or endpoint string)
+            out_q: Output queue for sending results to orchestrator (queue object or endpoint string)
         """
         self._in_q = in_q
         self._out_q = out_q
@@ -397,6 +412,7 @@ class OmniStage:
         batch_timeout: int = 10,
         connectors_config: dict | None = None,
         worker_backend: str = "multi_process",
+        ignore_runtime_config: bool = False,
         **kwargs: Any,
     ) -> None:
         """Initialize and start the stage worker process.
@@ -412,6 +428,7 @@ class OmniStage:
             batch_timeout: Timeout in seconds for batching requests
             connectors_config: Configuration for stage connectors
             worker_backend: Backend type ("multi_process" or "ray")
+            ignore_runtime_config: Whether to ignore runtime configuration (default: False)
             **kwargs: Additional arguments (e.g. ray_placement_group)
 
         Raises:
@@ -429,7 +446,10 @@ class OmniStage:
         ctx = ctx or mp.get_context("spawn")
         # Prepare lightweight dict config for worker
         engine_args = _to_dict(self.engine_args)
-        runtime_cfg = _to_dict(getattr(self.stage_config, "runtime", {}))
+        if ignore_runtime_config:
+            runtime_cfg = {}
+        else:
+            runtime_cfg = _to_dict(getattr(self.stage_config, "runtime", {}))
         stage_payload: dict[str, Any] = {
             "stage_id": self.stage_id,
             "engine_args": engine_args,
@@ -438,6 +458,8 @@ class OmniStage:
             "connectors_config": connectors_config or {},
             "stage_type": self.stage_type,
             "engine_input_source": self.engine_input_source,
+            "final_output": self.final_output,
+            "final_output_type": self.final_output_type,
         }
         try:
             old_env = os.environ.get("VLLM_LOGGING_PREFIX")
@@ -449,9 +471,10 @@ class OmniStage:
                         _stage_worker_async_entry,
                         ray_placement_group,
                         self.stage_id,
-                        self,
                         model=model,
                         stage_payload=stage_payload,
+                        in_q=self._in_q,
+                        out_q=self._out_q,
                         batch_timeout=batch_timeout,
                         stage_init_timeout=self._stage_init_timeout,
                     )
@@ -472,9 +495,10 @@ class OmniStage:
                     self._proc = ctx.Process(
                         target=_stage_worker_async_entry,
                         args=(
-                            self,
                             model,
                             stage_payload,
+                            self._in_q.endpoint if isinstance(self._in_q, ZmqQueue) else self._in_q,
+                            self._out_q.endpoint if isinstance(self._out_q, ZmqQueue) else self._out_q,
                             batch_timeout,
                             self._stage_init_timeout,
                         ),
@@ -485,8 +509,8 @@ class OmniStage:
                         args=(
                             model,
                             stage_payload,
-                            self._in_q,
-                            self._out_q,
+                            self._in_q.endpoint if isinstance(self._in_q, ZmqQueue) else self._in_q,
+                            self._out_q.endpoint if isinstance(self._out_q, ZmqQueue) else self._out_q,
                             batch_timeout,
                             self._stage_init_timeout,
                         ),
@@ -510,6 +534,13 @@ class OmniStage:
                 self._in_q.put_nowait(SHUTDOWN_TASK)
             except Exception as e:
                 logger.warning("Failed to send shutdown to in_q: %s", e)
+            close_fn = getattr(self._in_q, "close", None)
+            if callable(close_fn):
+                close_fn()
+        if self._out_q is not None:
+            close_fn = getattr(self._out_q, "close", None)
+            if callable(close_fn):
+                close_fn()
 
         if hasattr(self, "_ray_actor") and self._ray_actor:
             kill_ray_actor(self._ray_actor)
@@ -632,8 +663,8 @@ class OmniStage:
 def _stage_worker(
     model: str,
     stage_payload: dict[str, Any],
-    in_q: mp.Queue,
-    out_q: mp.Queue,
+    in_q: mp.queues.Queue | ZmqQueue | str,
+    out_q: mp.queues.Queue | ZmqQueue | str,
     batch_timeout: int = 10,
     stage_init_timeout: int = 300,
 ) -> None:
@@ -644,6 +675,8 @@ def _stage_worker(
     import os as _os
     import time as _time
 
+    import zmq
+
     from vllm_omni.plugins import load_omni_general_plugins
 
     load_omni_general_plugins()
@@ -670,6 +703,21 @@ def _stage_worker(
     if stage_type != "diffusion":
         _resolve_worker_cls(engine_args)
 
+    # Resolve ZMQ queue endpoints if needed
+    zmq_ctx = None
+    if isinstance(in_q, str) or isinstance(out_q, str):
+        zmq_ctx = zmq.Context()
+        if isinstance(in_q, str):
+            in_q = create_zmq_queue(zmq_ctx, in_q, zmq.PULL)
+        if isinstance(out_q, str):
+            out_q = create_zmq_queue(zmq_ctx, out_q, zmq.PUSH)
+        # When using ZMQ (cross-node IPC), disable SHM so data is sent inline.
+        shm_threshold_bytes = sys.maxsize
+        logger.info(
+            "[Stage-%s] ZMQ transport detected; disabling SHM IPC (shm_threshold_bytes set to maxsize)",
+            stage_id,
+        )
+
     # Aggregates for running average
     _agg_total_tokens = 0
     _agg_total_gen_time_ms = 0.0
@@ -701,6 +749,7 @@ def _stage_worker(
             engine_args["stage_connector_spec"] = stage_connector_spec
             engine_args["stage_id"] = stage_id
         if stage_type == "diffusion":
+            engine_args = filter_dataclass_kwargs(OmniDiffusionConfig, engine_args)
             engine_args.pop("model_stage", None)
             engine_args.pop("model", None)
             stage_engine = OmniDiffusion(
@@ -710,6 +759,8 @@ def _stage_worker(
                 **engine_args,
             )
         else:
+            engine_args = filter_dataclass_kwargs(OmniEngineArgs, engine_args)
+            engine_args.pop("model", None)
             # Default to LLM engine
             stage_engine = OmniLLM(model=model, **engine_args)
 
@@ -1004,19 +1055,21 @@ def _stage_worker(
 
 
 def _stage_worker_async_entry(
-    omni_stage: OmniStage,
     model: str,
     stage_payload: dict[str, Any],
+    in_q: mp.queues.Queue | ZmqQueue | str,
+    out_q: mp.queues.Queue | ZmqQueue | str,
     batch_timeout: int = 10,
     stage_init_timeout: int = 300,
 ) -> None:
-    asyncio.run(_stage_worker_async(omni_stage, model, stage_payload, batch_timeout, stage_init_timeout))
+    asyncio.run(_stage_worker_async(model, stage_payload, in_q, out_q, batch_timeout, stage_init_timeout))
 
 
 async def _stage_worker_async(
-    omni_stage: OmniStage,
     model: str,
     stage_payload: dict[str, Any],
+    in_q: mp.queues.Queue | ZmqQueue | str,
+    out_q: mp.queues.Queue | ZmqQueue | str,
     batch_timeout: int = 10,
     stage_init_timeout: int = 300,
 ) -> None:
@@ -1026,6 +1079,8 @@ async def _stage_worker_async(
     import os as _os
     import time as _time
 
+    import zmq
+
     from vllm_omni.plugins import load_omni_general_plugins
 
     load_omni_general_plugins()
@@ -1045,12 +1100,26 @@ async def _stage_worker_async(
     shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536))
     connectors_config = stage_payload.get("connectors_config", {})
     stage_type = stage_payload.get("stage_type", "llm")
+    final_output = stage_payload.get("final_output", False)
+    final_output_type = stage_payload.get("final_output_type", None)
 
     if stage_type != "diffusion":
         _resolve_worker_cls(engine_args)
 
-    in_q = omni_stage._in_q
-    out_q = omni_stage._out_q
+    # Resolve ZMQ queue endpoints if needed
+    zmq_ctx = None
+    if isinstance(in_q, str) or isinstance(out_q, str):
+        zmq_ctx = zmq.Context()
+        if isinstance(in_q, str):
+            in_q = create_zmq_queue(zmq_ctx, in_q, zmq.PULL)
+        if isinstance(out_q, str):
+            out_q = create_zmq_queue(zmq_ctx, out_q, zmq.PUSH)
+        # When using ZMQ (cross-node IPC), disable SHM so data is sent inline.
+        shm_threshold_bytes = sys.maxsize
+        logger.info(
+            "[Stage-%s] ZMQ transport detected; disabling SHM IPC (shm_threshold_bytes set to maxsize)",
+            stage_id,
+        )
 
     # Aggregates for running average
     _agg_total_tokens = 0
@@ -1099,6 +1168,7 @@ async def _stage_worker_async(
             engine_args["stage_id"] = stage_id
         if stage_type == "diffusion":
             # For diffusion, we need to extract diffusion-specific config
+            engine_args = filter_dataclass_kwargs(OmniDiffusionConfig, engine_args)
             od_config = _build_od_config(engine_args, model)
 
             # Inject omni config for worker to access stage info
@@ -1115,6 +1185,8 @@ async def _stage_worker_async(
             )
             vllm_config = None  # Diffusion doesn't use vllm_config
         else:
+            engine_args = filter_dataclass_kwargs(AsyncOmniEngineArgs, engine_args)
+            engine_args.pop("model", None)
             omni_engine_args = AsyncOmniEngineArgs(model=model, **engine_args)
             usage_context = UsageContext.OPENAI_API_SERVER
             vllm_config = omni_engine_args.create_engine_config(usage_context=usage_context)
@@ -1126,14 +1198,13 @@ async def _stage_worker_async(
                     engine_args.get("disable_log_stats", False) or getattr(omni_engine_args, "disable_log_stats", False)
                 ),
             )
-    omni_stage.set_async_engine(stage_engine)
-    if hasattr(omni_stage.async_engine, "log_stats") and omni_stage.async_engine.log_stats:
+    if hasattr(stage_engine, "log_stats") and stage_engine.log_stats:
 
         async def _force_log():
             try:
                 while True:
                     await asyncio.sleep(10.0)
-                    await omni_stage.async_engine.do_log_stats()
+                    await stage_engine.do_log_stats()
             except asyncio.CancelledError:
                 pass
 
@@ -1330,7 +1401,7 @@ async def _stage_worker_async(
             batch_request_ids, batch_request_outputs, _gen_ms_list, batch_metrics
         ):
             try:
-                r_outputs = [output_strip(output, omni_stage)]
+                r_outputs = [output_strip(output, final_output, final_output_type)]
                 use_shm, payload = maybe_dump_to_shm(r_outputs, shm_threshold_bytes)
                 if use_shm:
                     out_q.put(
@@ -1416,7 +1487,7 @@ def make_stage_stats(_agg_total_tokens: int, _agg_total_gen_time_ms: float):
     return StageStats(total_token=_agg_total_tokens, total_gen_time_ms=_agg_total_gen_time_ms)
 
 
-def output_strip(r_output: RequestOutput | OmniRequestOutput, omni_stage: OmniStage):
+def output_strip(r_output: RequestOutput | OmniRequestOutput, final_output: bool, final_output_type: str | None):
     """
     Strip unnecessary multimodal outputs from stages results,
     in order to:
@@ -1425,7 +1496,7 @@ def output_strip(r_output: RequestOutput | OmniRequestOutput, omni_stage: OmniSt
     """
 
     # check multimodal data is required by stage output config.
-    if omni_stage.final_output and omni_stage.final_output_type != "text":
+    if final_output and final_output_type != "text":
         return r_output
 
     # If the request has already finished, should not be altered.

vllm_omni/entrypoints/utils.py

@@ -1,8 +1,9 @@
 import os
+import types
 from collections import Counter
-from dataclasses import asdict, is_dataclass
+from dataclasses import asdict, fields, is_dataclass
 from pathlib import Path
-from typing import Any
+from typing import Any, get_args, get_origin
 
 from omegaconf import OmegaConf
 from vllm.logger import init_logger
@@ -68,6 +69,52 @@ def _try_get_class_name_from_diffusers_config(model: str) -> str | None:
     return None
 
 
+def _filter_dict_like_object(obj: dict | Any) -> dict:
+    """Filter dict-like object by removing callables and recursively converting values.
+
+    Converts dict-like objects to regular dicts while filtering out callable values
+    that are incompatible with OmegaConf. Recursively processes values through
+    _convert_dataclasses_to_dict for nested object conversion.
+
+    Args:
+        obj: Dict or dict-like object to filter
+
+    Returns:
+        Regular dict with callables filtered out and values recursively converted
+
+    Raises:
+        TypeError: If obj doesn't support .items() method
+        ValueError: If dict conversion fails unexpectedly
+    """
+
+    def _is_callable_value(value: Any) -> bool:
+        if callable(value):
+            return True
+        return isinstance(
+            value,
+            (
+                types.FunctionType,
+                types.MethodType,
+                types.BuiltinFunctionType,
+                types.BuiltinMethodType,
+            ),
+        )
+
+    result = {}
+    filtered_keys = []
+    for k, v in obj.items():
+        if _is_callable_value(v):
+            filtered_keys.append(str(k))
+        else:
+            result[k] = _convert_dataclasses_to_dict(v)
+    if filtered_keys:
+        logger.warning(
+            f"Filtered out {len(filtered_keys)} callable object(s) from base_engine_args "
+            f"that are not compatible with OmegaConf: {filtered_keys}. "
+        )
+    return result
+
+
 def _convert_dataclasses_to_dict(obj: Any) -> Any:
     """Recursively convert non-serializable objects to OmegaConf-compatible types.
 
@@ -75,6 +122,7 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
     - Dataclass objects with Literal type annotations (e.g., StructuredOutputsConfig)
     - Counter objects (from collections or vllm.utils)
     - Set objects
+    - Callable objects (functions, methods, etc.)
     - Other non-primitive types
     """
     # IMPORTANT: Check Counter BEFORE dict, since Counter is a subclass of dict
@@ -99,17 +147,21 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any:
         result = asdict(obj)
         # Recursively process the result to convert any Counter objects
         return _convert_dataclasses_to_dict(result)
-    # Handle dictionaries (recurse into values)
+    # Handle dictionaries (recurse into values) and filter out callables(cause error in OmegaConf.create)
     # Note: This must come AFTER Counter check since Counter is a dict subclass
     if isinstance(obj, dict):
-        return {k: _convert_dataclasses_to_dict(v) for k, v in obj.items()}
+        return _filter_dict_like_object(obj)
+    # Handle callable objects (functions, methods, etc.) - skip them
+    # Note: This comes after dict/list checks to avoid misclassifying dict-like objects
+    if callable(obj):
+        return None
     # Handle lists and tuples (recurse into items)
     if isinstance(obj, (list, tuple)):
-        return type(obj)(_convert_dataclasses_to_dict(item) for item in obj)
+        return type(obj)(_convert_dataclasses_to_dict(item) for item in obj if not callable(item))
     # Try to convert any dict-like object (has keys/values methods) to dict
     if hasattr(obj, "keys") and hasattr(obj, "values") and not isinstance(obj, (str, bytes)):
         try:
-            return {k: _convert_dataclasses_to_dict(v) for k, v in obj.items()}
+            return _filter_dict_like_object(obj)
         except (TypeError, ValueError, AttributeError):
             # If conversion fails, return as-is
             return obj
@@ -236,6 +288,40 @@ def load_stage_configs_from_yaml(config_path: str, base_engine_args: dict | None
     return stage_args
 
 
+def load_and_resolve_stage_configs(
+    model: str,
+    stage_configs_path: str | None,
+    kwargs: dict | None,
+    default_stage_cfg_factory: Any = None,
+) -> tuple[str, list]:
+    """Load stage configurations from model or YAML file with fallback to defaults.
+
+    Args:
+        model: Model name or path
+        stage_configs_path: Optional path to YAML file containing stage configurations
+        kwargs: Engine arguments to merge with stage configs
+        default_stage_cfg_factory: Optional callable that takes no args and returns
+            default stage config list when no configs are found
+
+    Returns:
+        Tuple of (config_path, stage_configs)
+    """
+    if stage_configs_path is None:
+        config_path = resolve_model_config_path(model)
+        stage_configs = load_stage_configs_from_model(model, base_engine_args=kwargs)
+        if not stage_configs:
+            if default_stage_cfg_factory is not None:
+                default_stage_cfg = default_stage_cfg_factory()
+                stage_configs = OmegaConf.create(default_stage_cfg)
+            else:
+                stage_configs = []
+    else:
+        config_path = stage_configs_path
+        stage_configs = load_stage_configs_from_yaml(stage_configs_path, base_engine_args=kwargs)
+
+    return config_path, stage_configs
+
+
 def get_final_stage_id_for_e2e(
     output_modalities: list[str] | None, default_modalities: list[str], stage_list: list
 ) -> int:
@@ -282,6 +368,117 @@ def get_final_stage_id_for_e2e(
     return final_stage_id_for_e2e
 
 
+def filter_dataclass_kwargs(cls: Any, kwargs: dict) -> dict:
+    """Filter kwargs to only include fields defined in the dataclass.
+
+    Args:
+        cls: Dataclass type
+        kwargs: Keyword arguments to filter
+
+    Returns:
+        Filtered keyword arguments containing only valid dataclass fields
+    """
+    if not is_dataclass(cls):
+        raise ValueError(f"{cls} is not a dataclass")
+    if not isinstance(kwargs, dict):
+        raise ValueError("kwargs must be a dictionary")
+
+    def _filter_value(value: Any, annotation: Any) -> Any:
+        """Recursively filter nested dict/list values based on dataclass annotations."""
+        if annotation is None:
+            return value
+
+        origin = get_origin(annotation)
+        if origin is None:
+            if isinstance(annotation, type) and is_dataclass(annotation) and isinstance(value, dict):
+                return filter_dataclass_kwargs(annotation, value)
+            return value
+
+        if origin in (list, tuple, set):
+            args = get_args(annotation)
+            inner = args[0] if args else None
+            if isinstance(value, (list, tuple, set)):
+                return type(value)(_filter_value(v, inner) for v in value)
+            return value
+
+        if origin is dict:
+            args = get_args(annotation)
+            val_type = args[1] if len(args) > 1 else None
+            if isinstance(value, dict):
+                return {k: _filter_value(v, val_type) for k, v in value.items()}
+            return value
+
+        if origin is types.UnionType or origin is getattr(types, "UnionType", None):
+            for arg in get_args(annotation):
+                if isinstance(arg, type) and is_dataclass(arg) and isinstance(value, dict):
+                    return filter_dataclass_kwargs(arg, value)
+                # Try container-style filtering for union members
+                filtered = _filter_value(value, arg)
+                if filtered is not value:
+                    return filtered
+            return value
+
+        return value
+
+    valid_fields = {f.name: f for f in fields(cls) if f.init}
+    filtered_kwargs = {}
+    for k, v in kwargs.items():
+        if k not in valid_fields:
+            continue
+        field = valid_fields[k]
+        filtered_kwargs[k] = _filter_value(v, field.type)
+
+    return filtered_kwargs
+
+
+# TODO(wuhang): Remove after PR #1115.
+def build_base_engine_args(source: Any) -> dict[str, Any] | None:
+    """Build base engine args with tokenizer and parallel configuration.
+
+    Automatically detects whether source is a dict-like object or namespace object.
+
+    Args:
+        source: Source object (args namespace or kwargs dict) containing configuration.
+
+    Returns:
+        Dictionary containing tokenizer and parallel configuration overrides,
+        or None if no configuration is present.
+    """
+    # Auto-detect source type: dict-like objects have 'get' method
+    is_dict_like = hasattr(source, "get") and callable(getattr(source, "get"))
+
+    # Extract tokenizer
+    if is_dict_like:
+        tokenizer = source.get("tokenizer", None)
+    else:
+        tokenizer = getattr(source, "tokenizer", None)
+
+    base_engine_args = {"tokenizer": tokenizer} if tokenizer is not None else None
+
+    # Extract parallel configuration
+    parallel_keys = [
+        "tensor_parallel_size",
+        "pipeline_parallel_size",
+        "data_parallel_size",
+        "data_parallel_size_local",
+        "data_parallel_backend",
+        "distributed_executor_backend",
+    ]
+
+    if is_dict_like:
+        parallel_overrides = {k: source[k] for k in parallel_keys if k in source and source[k] is not None}
+    else:
+        parallel_overrides = {
+            k: getattr(source, k) for k in parallel_keys if hasattr(source, k) and getattr(source, k) is not None
+        }
+
+    if parallel_overrides:
+        base_engine_args = base_engine_args or {}
+        base_engine_args.update(parallel_overrides)
+
+    return base_engine_args
+
+
 # The following code detects if the process is running in a container and if
 # PID host is available. If so, we can use process-scoped memory tracking;
 # otherwise we need sequential init locks.

vllm_omni/entrypoints/zmq_utils.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""ZMQ-based queue utilities for Omni IPC."""
+
+from __future__ import annotations
+
+import queue
+from typing import Any
+
+import zmq
+from vllm.utils.network_utils import make_zmq_socket
+
+
+class ZmqQueue:
+    """Queue-like wrapper on a ZMQ socket."""
+
+    def __init__(
+        self,
+        ctx: zmq.Context,
+        socket_type: int,
+        *,
+        bind: str | None = None,
+        connect: str | None = None,
+        recv_timeout_ms: int | None = None,
+        send_timeout_ms: int | None = None,
+    ) -> None:
+        # Determine path and bind mode
+        path = bind if bind is not None else connect
+        if path is None:
+            raise ValueError("Either bind or connect must be specified")
+        bind_mode = bind is not None
+
+        self._socket = make_zmq_socket(ctx, path, socket_type, bind=bind_mode, linger=5000)
+
+        # Reusable poller for efficient polling operations
+        self._poller = zmq.Poller()
+        self._poller.register(self._socket, zmq.POLLIN)
+
+        # Store default timeout settings
+        self._default_recv_timeout = recv_timeout_ms
+        self._default_send_timeout = send_timeout_ms
+
+        # Apply timeout settings if specified
+        if recv_timeout_ms is not None:
+            self._socket.rcvtimeo = recv_timeout_ms
+        if send_timeout_ms is not None:
+            self._socket.sndtimeo = send_timeout_ms
+
+        self.endpoint = path
+
+    def put(self, obj: Any) -> None:
+        """Send an object to the queue. Blocks until sent or timeout."""
+        try:
+            self._socket.send_pyobj(obj)
+        except zmq.Again as e:
+            raise queue.Full() from e
+
+    def put_nowait(self, obj: Any) -> None:
+        """Send an object to the queue without blocking."""
+        try:
+            self._socket.send_pyobj(obj, flags=zmq.NOBLOCK)
+        except zmq.Again as e:
+            raise queue.Full() from e
+
+    def get(self, timeout: float | None = None) -> Any:
+        """Receive an object from the queue with optional timeout in seconds."""
+        if timeout is None:
+            return self._socket.recv_pyobj()
+
+        # Use the reusable poller for timeout handling
+        events = dict(self._poller.poll(int(timeout * 1000)))
+        if events.get(self._socket) == zmq.POLLIN:
+            return self._socket.recv_pyobj()
+        raise queue.Empty()
+
+    def get_nowait(self) -> Any:
+        """Receive an object from the queue without blocking."""
+        try:
+            return self._socket.recv_pyobj(flags=zmq.NOBLOCK)
+        except zmq.Again as e:
+            raise queue.Empty() from e
+
+    def empty(self) -> bool:
+        """Check if the queue is empty without blocking."""
+        events = dict(self._poller.poll(0))
+        return events.get(self._socket) != zmq.POLLIN
+
+    def close(self) -> None:
+        self._socket.close(0)
+
+
+def create_zmq_queue(ctx: zmq.Context, endpoint: str, socket_type: int) -> ZmqQueue:
+    """Create a ZmqQueue from an endpoint string and socket type."""
+    return ZmqQueue(ctx, socket_type, connect=endpoint)

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -15,7 +15,6 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
     Qwen3OmniMoeThinkerConfig,
 )
 from vllm.config import VllmConfig
-from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
 from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
@@ -355,13 +354,14 @@ class Qwen3OmniMoeForConditionalGeneration(
 
         # ========== Stage 3: Code2Wav ==========
         elif self.model_stage == "code2wav":
+            seq_token_counts: list[int] | None = kwargs.get("seq_token_counts")
+
             # Extract codec codes from input
             if input_ids.shape[0] % 16 == 0:
-                ubatch_slices = get_forward_context().ubatch_slices
-                if ubatch_slices is not None:
-                    max_seq_len = max(ubatch_slices) // 16
-                    batch_size = len(ubatch_slices)
-                    split_codes = torch.split(input_ids, ubatch_slices, dim=0)
+                if seq_token_counts is not None:
+                    max_seq_len = max(seq_token_counts) // 16
+                    batch_size = len(seq_token_counts)
+                    split_codes = torch.split(input_ids, seq_token_counts, dim=0)
                     codes = torch.zeros((batch_size, 16, max_seq_len), device=input_ids.device, dtype=input_ids.dtype)
                     for idx, code in enumerate(split_codes):
                         seq_len = code.shape[0] // 16
@@ -386,7 +386,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 codes = input_ids_flatten.reshape(1, 16, -1)
 
             # Generate audio from codec codes
-            audio_tensors = self.generate_audio(codes, voice_type)
+            audio_tensors = self.generate_audio(codes, voice_type, seq_token_counts)
 
             return audio_tensors
 
@@ -458,16 +458,22 @@ class Qwen3OmniMoeForConditionalGeneration(
 
     # ==================== Audio Generation ====================
 
-    def generate_audio(self, code: torch.Tensor, voice_type: str) -> list[torch.Tensor]:
+    def generate_audio(
+        self,
+        code: torch.Tensor,
+        voice_type: str,
+        seq_token_counts: list[int] | None = None,
+    ) -> list[torch.Tensor]:
         """
         Generate audio waveform from codec codes.
 
         Args:
-            code: [8, T] - 8-layer RVQ codec codes
+            code: [batch, num_quantizers, T] - RVQ codec codes
             voice_type: Voice type (not used in Qwen3, kept for compatibility)
+            seq_token_counts: Token count for each request in batch
 
         Returns:
-            audio_tensor: [1, waveform_len] - Audio waveform
+            list of audio waveforms
         """
         code2wav_dev = self._module_device(self.code2wav)
 
@@ -491,6 +497,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 talker_codes,
                 chunk_size=25,
                 left_context_size=25,
+                seq_token_counts=seq_token_counts,
             )
         else:
             # Use chunked decode for memory efficiency
@@ -498,6 +505,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 talker_codes,
                 chunk_size=300,
                 left_context_size=25,
+                seq_token_counts=seq_token_counts,
             )
 
         return audio_tensors

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py

@@ -22,7 +22,6 @@ from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
     SnakeBeta,
 )
 from vllm.config import VllmConfig  # type: ignore
-from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger  # type: ignore
 from vllm.model_executor.models.utils import (  # type: ignore
     AutoWeightsLoader,
@@ -163,6 +162,7 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
         codes: torch.Tensor,
         chunk_size: int = 300,
         left_context_size: int = 25,
+        seq_token_counts: list[int] | None = None,
     ) -> list[torch.Tensor]:
         """
         Decode long sequences in chunks to avoid OOM.
@@ -173,6 +173,7 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
             codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes
             chunk_size: Number of codec frames per chunk
             left_context_size: Number of overlapping frames for context
+            seq_token_counts: Token count for each request in batch
 
         Returns:
             list[torch.Tensor]: Complete waveform decoded from the input
@@ -197,12 +198,10 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
 
             start_index = end_index
 
-        ubatch_slices = get_forward_context().ubatch_slices
-        if ubatch_slices is not None:
-            code_seq_lens = [seq_len // self.config.num_quantizers for seq_len in ubatch_slices]
+        if seq_token_counts is not None:
+            code_seq_lens = [seq_len // self.config.num_quantizers for seq_len in seq_token_counts]
         else:
             # Fallback: assume all batch elements share the same sequence length.
-            # Create one entry per batch so that each element is processed.
             code_seq_lens = [codes.shape[-1]] * codes.shape[0]
         batch_wav = torch.cat(wavs, dim=-1)
         wavs = []
@@ -216,6 +215,7 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
         codes: torch.Tensor,
         chunk_size: int = 25,
         left_context_size: int = 25,
+        seq_token_counts: list[int] | None = None,
     ) -> list[torch.Tensor]:
         """
         Decode long sequences in chunks to avoid OOM.
@@ -226,21 +226,19 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
             codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes
             chunk_size: Number of codec frames per chunk
             left_context_size: Number of overlapping frames for context
+            seq_token_counts: Token count for each request in batch
 
         Returns:
             list[torch.Tensor]: Complete waveform decoded from the input
                 codes. For ``batch_size == 1``, this is a list containing a
                 single tensor with shape ``[1, waveform_len]``.
         """
-        # Decode chunk
         wavs = []
         batch_wav = self(codes)
-        ubatch_slices = get_forward_context().ubatch_slices
-        if ubatch_slices is not None:
-            code_seq_lens = [seq_len // self.config.num_quantizers for seq_len in ubatch_slices]
+        if seq_token_counts is not None:
+            code_seq_lens = [n // self.config.num_quantizers for n in seq_token_counts]
         else:
             # Fallback: assume all batch elements share the same sequence length.
-            # Create one entry per batch so that each element is processed.
             code_seq_lens = [codes.shape[-1]] * codes.shape[0]
         for idx, code_seq_len in enumerate(code_seq_lens):
             # TODO: need to optimize algorithms, current only support

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

@@ -9,6 +9,7 @@ import torch
 import torch.nn as nn
 from transformers.utils.hub import cached_file
 from vllm.config import VllmConfig
+from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.logger import init_logger
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput
@@ -116,6 +117,28 @@ class Qwen3TTSCode2Wav(nn.Module):
     def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None:
         return None
 
+    def _split_request_ids(self, ids: torch.Tensor, seq_token_counts: list[int] | None = None) -> list[torch.Tensor]:
+        """Split concatenated input_ids into per-request segments.
+
+        Uses seq_token_counts (injected by the runner via model_kwargs) when
+        available, falling back to forward-context ubatch_slices when
+        micro-batching is active. Returns [ids] for single-request batches.
+        """
+        if seq_token_counts is not None and len(seq_token_counts) > 1:
+            boundaries = [0]
+            for count in seq_token_counts:
+                boundaries.append(boundaries[-1] + count)
+            n = ids.numel()
+            return [ids[boundaries[i] : min(boundaries[i + 1], n)] for i in range(len(seq_token_counts))]
+        if is_forward_context_available():
+            slices = get_forward_context().ubatch_slices
+            if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices):
+                boundaries = [0]
+                for s in slices:
+                    boundaries.append(boundaries[-1] + s)
+                return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)]
+        return [ids]
+
     @torch.no_grad()
     def forward(
         self,
@@ -124,98 +147,129 @@ class Qwen3TTSCode2Wav(nn.Module):
         intermediate_tensors: Any = None,
         inputs_embeds: torch.Tensor | None = None,
         **kwargs: Any,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> OmniOutput:
         """Decode codec codes into audio waveform.
 
-        input_ids layout: [codec_context_frames, *flat_codes]
+        input_ids layout per request: [codec_context_frames, *flat_codes]
         where flat_codes is codebook-major [q*F].
+
+        When batched, uses forward context ubatch_slices to split the
+        concatenated input_ids and decode via a single batched forward pass.
         """
-        tok = self._ensure_speech_tokenizer_loaded()
+        self._ensure_speech_tokenizer_loaded()
         assert self._num_quantizers is not None
+        assert self._decode_upsample_rate is not None
         assert self._output_sample_rate is not None
 
-        sr_val = self._output_sample_rate
-        empty_ret = (
-            torch.zeros((0,), dtype=torch.float32),
-            torch.tensor(sr_val, dtype=torch.int32),
-        )
-
-        if input_ids is None:
-            return empty_ret
-
+        tok = self._speech_tokenizer
         q = int(self._num_quantizers)
-        ids = input_ids.reshape(-1).to(dtype=torch.long)
-        n_tokens = ids.numel()
+        upsample = int(self._decode_upsample_rate)
+        sr_val = int(self._output_sample_rate)
+        sr_tensor = torch.tensor(sr_val, dtype=torch.int32)
+        empty = torch.zeros((0,), dtype=torch.float32)
 
-        if n_tokens == 0:
-            return empty_ret
-
-        # input_ids[0] = codec_context_frames (prepended by stage_input_processor).
-        ctx_frames = int(ids[0].item())
-        ids = ids[1:]
-        n_tokens = ids.numel()
-
-        if n_tokens == 0:
-            return empty_ret
-
-        # Warmup / dummy_run: not divisible by num_quantizers.
-        if n_tokens % q != 0:
-            logger.warning(
-                "Code2Wav input_ids length %d not divisible by num_quantizers %d, "
-                "likely a warmup run; returning empty audio.",
-                n_tokens,
-                q,
+        if input_ids is None or input_ids.numel() == 0:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": [empty], "sr": [sr_tensor]},
             )
-            return empty_ret
 
-        total_frames = n_tokens // q
+        ids = input_ids.reshape(-1).to(dtype=torch.long)
+        request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts"))
 
-        # Reshape codebook-major flat [q*F] -> [q, F] -> [F, q] for SpeechTokenizer.
-        codes_fq = ids.reshape(q, total_frames).transpose(0, 1).contiguous()
+        # Parse each request: extract ctx_frames, validate, reshape codes.
+        # input_ids layout per request: [codec_context_frames, *flat_codes]
+        # where flat_codes is codebook-major [q*F].
+        parsed = []  # (ctx_frames, actual_frames)
+        valid_codes = []
+        valid_indices = []
+        for i, req_ids in enumerate(request_ids_list):
+            if req_ids.numel() < 2:
+                parsed.append((0, 0))
+                continue
+            ctx_frames = int(req_ids[0].item())
+            flat = req_ids[1:]
+            n = flat.numel()
+            # Warmup / dummy_run: not divisible by num_quantizers.
+            if n == 0 or n % q != 0:
+                if n > 0:
+                    logger.warning(
+                        "Code2Wav input_ids length %d not divisible by num_quantizers %d, "
+                        "likely a warmup run; returning empty audio.",
+                        n,
+                        q,
+                    )
+                parsed.append((0, 0))
+                continue
+            frames = n // q
+            # Reshape codebook-major flat [q*F] -> [q, F] -> [F, q] for SpeechTokenizer.
+            codes_fq = flat.reshape(q, frames).transpose(0, 1).contiguous()
+            parsed.append((ctx_frames, frames))
+            valid_codes.append({"audio_codes": codes_fq})
+            valid_indices.append(i)
 
-        if not self._logged_codec_stats and total_frames > 1:
+        num_req = len(request_ids_list)
+        if not valid_codes:
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={
+                    "model_outputs": [empty] * num_req,
+                    "sr": [sr_tensor] * num_req,
+                },
+            )
+
+        if not self._logged_codec_stats:
             self._logged_codec_stats = True
             try:
-                uniq = int(torch.unique(codes_fq).numel())
-                cmin = int(codes_fq.min().item())
-                cmax = int(codes_fq.max().item())
-                head = codes_fq[: min(2, total_frames), : min(8, q)].cpu().tolist()
+                c = valid_codes[0]["audio_codes"]
                 logger.info(
-                    "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] head=%s",
-                    total_frames,
+                    "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] head=%s batch=%d",
+                    c.shape[0],
                     q,
-                    uniq,
-                    cmin,
-                    cmax,
-                    head,
+                    int(torch.unique(c).numel()),
+                    int(c.min().item()),
+                    int(c.max().item()),
+                    c[: min(2, c.shape[0]), : min(8, q)].cpu().tolist(),
+                    len(valid_codes),
                 )
             except Exception:
                 pass
 
-        wavs, sr = tok.decode({"audio_codes": codes_fq})
-        if not wavs:
-            raise ValueError("SpeechTokenizer code2wav produced empty waveform list.")
-        audio_np = wavs[0].astype(np.float32, copy=False)
+        # Batched decode: single forward pass through SpeechTokenizer.
+        wavs, _ = tok.decode(valid_codes)
+        if len(wavs) != len(valid_codes):
+            raise RuntimeError(f"Code2Wav returned {len(wavs)} waveforms for {len(valid_codes)} requests")
 
-        # Trim left-context waveform samples (streaming sliding window).
-        if ctx_frames > 0:
-            upsample = self._decode_upsample_rate
-            if upsample is None or upsample <= 0:
-                raise ValueError(f"Invalid decode upsample rate: {upsample}")
-            cut = ctx_frames * upsample
-            if cut < audio_np.shape[0]:
-                audio_np = audio_np[cut:]
-            else:
-                logger.warning(
-                    "Context trim %d >= decoded length %d; returning empty audio.",
-                    cut,
-                    audio_np.shape[0],
-                )
-                return empty_ret
+        # Build per-request outputs, trimming padding and left-context.
+        audios = [empty] * num_req
+        srs = [sr_tensor] * num_req
 
-        audio_tensor = torch.from_numpy(audio_np).to(dtype=torch.float32).reshape(-1)
-        sr_tensor = torch.tensor(int(sr), dtype=torch.int32)
-        return audio_tensor, sr_tensor
+        for j, idx in enumerate(valid_indices):
+            ctx_frames, actual_frames = parsed[idx]
+            audio_np = wavs[j].astype(np.float32, copy=False)
+            # Trim decoder padding (output may be longer due to batch padding).
+            expected_len = actual_frames * upsample
+            if audio_np.shape[0] > expected_len:
+                audio_np = audio_np[:expected_len]
+            # Trim left-context waveform samples (streaming sliding window).
+            if ctx_frames > 0:
+                cut = ctx_frames * upsample
+                if cut < audio_np.shape[0]:
+                    audio_np = audio_np[cut:]
+                else:
+                    logger.warning(
+                        "Context trim %d >= decoded length %d; returning empty audio.",
+                        cut,
+                        audio_np.shape[0],
+                    )
+                    continue
+            if audio_np.shape[0] > 0:
+                audios[idx] = torch.from_numpy(audio_np).to(dtype=torch.float32).reshape(-1)
+
+        return OmniOutput(
+            text_hidden_states=None,
+            multimodal_outputs={"model_outputs": audios, "sr": srs},
+        )
 
     def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
         if isinstance(model_outputs, OmniOutput):

vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml

@@ -0,0 +1,105 @@
+# Same as qwen3_tts.yaml with batched talker and code2wav.
+# Stage 0: max_batch_size 4, stage 1: max_batch_size 4.
+# max_batch_size must be a power of two to align with CUDA graph capture sizes
+# (stage 0) and must match --batch-size in end2end.py / benchmark scripts.
+async_chunk: true
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 4
+    engine_args:
+      model_stage: qwen3_tts
+      model_arch: Qwen3TTSTalkerForConditionalGeneration
+      # Force stage-specific registered architecture.
+      hf_overrides:
+        architectures: [Qwen3TTSTalkerForConditionalGeneration]
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      enforce_eager: false
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: latent
+      gpu_memory_utilization: 0.3
+      distributed_executor_backend: "mp"
+      max_num_batched_tokens: 512
+      max_model_len: 4096
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
+    # Use named connector to apply runtime.connectors.extra.
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      seed: 42
+      detokenize: false
+      repetition_penalty: 1.05
+      stop_token_ids: [2150]
+
+  - stage_id: 1
+    stage_type: llm
+    runtime:
+      devices: "0"
+      max_batch_size: 4
+    engine_args:
+      model_stage: code2wav
+      model_arch: Qwen3TTSCode2Wav
+      # Force stage-specific registered architecture.
+      hf_overrides:
+        architectures: [Qwen3TTSCode2Wav]
+      worker_type: generation
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      enforce_eager: true
+      trust_remote_code: true
+      async_scheduling: false
+      enable_prefix_caching: false
+      engine_output_type: audio
+      gpu_memory_utilization: 0.2
+      distributed_executor_backend: "mp"
+      # Must be divisible by num_code_groups and cover (left_context + chunk).
+      max_num_batched_tokens: 8192
+      # async_chunk appends windows per step; max_model_len must cover accumulated stream.
+      max_model_len: 32768
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+    # Distributed connector configuration
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      detokenize: true
+      repetition_penalty: 1.0
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+        # Frame-aligned codec streaming transport.
+        codec_streaming: true
+        # Connector polling / timeout (unit: loop count, sleep interval in seconds).
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        # Align with Omni: small chunks with sufficient context overlap.
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

vllm_omni/worker/gpu_ar_model_runner.py

@@ -192,15 +192,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs),
             )
 
-            logger.debug(
-                "Running batch with cudagraph_mode: %s, batch_descriptor: %s, "
-                "should_ubatch: %s, num_tokens_across_dp: %s",
-                cudagraph_mode,
-                batch_desc,
-                should_ubatch,
-                num_tokens_across_dp,
-            )
-
             num_tokens_padded = batch_desc.num_tokens
             num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs
             ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices(
@@ -211,12 +202,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 self.parallel_config.num_ubatches,
             )
 
-            logger.debug(
-                "ubatch_slices: %s, ubatch_slices_padded: %s",
-                ubatch_slices,
-                ubatch_slices_padded,
-            )
-
             pad_attn = cudagraph_mode == CUDAGraphMode.FULL
 
             use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
@@ -313,15 +298,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
                 aux_hidden_states = None
 
             hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output)
-            if multimodal_outputs is not None:
-                keys_or_type = (
-                    list(multimodal_outputs.keys())
-                    if isinstance(multimodal_outputs, dict)
-                    else type(multimodal_outputs)
-                )
-                logger.debug(f"[AR] execute_model: multimodal_outputs keys = {keys_or_type}")
-            else:
-                logger.debug("[AR] execute_model: multimodal_outputs is None")
 
             if not self.broadcast_pp_output:
                 # Common case.

vllm_omni/worker/gpu_generation_model_runner.py

@@ -250,6 +250,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 intermediate_tensors,
             )
 
+            # [Omni] Pass token counts per request for code2wav output slicing
+            model_kwargs["seq_token_counts"] = tokens
+
         # Set cudagraph mode to none if calc_kv_scales is true.
         # KV scales calculation involves dynamic operations that are incompatible
         # with CUDA graph capture.
@@ -258,10 +261,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             # Mark KV scales as calculated after the first forward pass
             self.calculate_kv_scales = False
 
-        if ubatch_slices_padded is None:
-            # reuse ubatch_slices_padded for code2wav batching
-            ubatch_slices_padded = tokens
-
         # Run the model.
         # Use persistent buffers for CUDA graphs.
         # When spec decode is enabled, delay clearing connector metadata

vllm_omni/worker/gpu_model_runner.py

@@ -1148,12 +1148,6 @@ class OmniGPUModelRunner(GPUModelRunner):
         """Inject omni-specific kwargs into forward and cache model output"""
         model_kwargs_extra = self._build_model_kwargs_extra()
 
-        runtime_info = model_kwargs_extra.get("runtime_additional_information", [])
-        if runtime_info:
-            for i, info in enumerate(runtime_info):
-                if info:
-                    logger.debug(f"[OMNI] req[{i}] runtime_additional_information keys: {list(info.keys())}")
-
         model_output = super()._model_forward(
             input_ids=input_ids,
             positions=positions,