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

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

.buildkite/bootstrap-intel-omni.sh

@@ -0,0 +1,51 @@
+#!/bin/bash
+# vllm-omni Intel bootstrap
+# Uses static pipeline-intel.yaml for Intel XPU tests
+# Based on: bootstrap-amd-omni.sh
+
+set -euo pipefail
+
+if [[ -z "${DOCS_ONLY_DISABLE:-}" ]]; then
+    DOCS_ONLY_DISABLE=0
+fi
+
+upload_pipeline() {
+    echo "--- 🛠 Preparing Intel pipeline"
+    buildkite-agent pipeline upload .buildkite/pipeline-intel.yaml
+}
+
+get_diff() {
+    $(git add .)
+    echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
+}
+
+get_diff_main() {
+    $(git add .)
+    echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
+}
+
+file_diff=$(get_diff)
+if [[ $BUILDKITE_BRANCH == "main" ]]; then
+    file_diff=$(get_diff_main)
+fi
+
+# ----------------------------------------------------------------------
+# Early exit: skip pipeline if all changed files are under docs/
+# ----------------------------------------------------------------------
+if [[ "${DOCS_ONLY_DISABLE}" != "1" ]] && [[ -n "${file_diff:-}" ]]; then
+    docs_only=1
+    while IFS= read -r f; do
+        [[ -z "$f" ]] && continue
+        if [[ "$f" != docs/* ]]; then
+            docs_only=0
+            break
+        fi
+    done < <(printf '%s\n' "$file_diff" | tr ' ' '\n' | tr -d '\r')
+
+    if [[ "$docs_only" -eq 1 ]]; then
+        buildkite-agent annotate ":memo: CI skipped — docs only" --style "info"
+        exit 0
+    fi
+fi
+
+upload_pipeline

.buildkite/pipeline-intel.yaml

@@ -0,0 +1,16 @@
+# .buildkite/pipeline-intel.yaml
+steps:
+  - group: "INTEL Tests"
+    steps:
+      - label: "Intel: XPU Test"
+        command: "bash .buildkite/scripts/hardware_ci/run-xpu-test.sh"
+        agents:
+          queue: "intel-gpu-omni"
+        env:
+          DOCKER_BUILDKIT: "1"
+          # Buildkite will automatically replace this with the actual commit hash
+          VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
+          VLLM_VERSION: "v0.17.0"
+        priority: 100
+        timeout_in_minutes: 60
+        soft_fail: true

.buildkite/test-amd.yaml

@@ -101,6 +101,7 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+    - export VLLM_ROCM_USE_AITER=0
     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
 

.github/ISSUE_TEMPLATE/400-bug-report.yml

@@ -74,28 +74,21 @@ body:
       If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
 
       ```python
-      from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
+      from vllm_omni.entrypoints.omni import Omni
+      from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+      from vllm import SamplingParams
 
-      # Create stage configurations
-      ar_config = create_ar_stage_config(
-          stage_id=0,
-          model_path="Qwen/Qwen3-0.6B",
-          input_modalities=["text"],
-          output_modalities=["text"]
+      omni = Omni(
+          model="Qwen/Qwen-Image",
+          stage_configs_path="/path/to/stage_configs.yaml",
       )
 
-      dit_config = create_dit_stage_config(
-          stage_id=1,
-          model_path="stabilityai/stable-diffusion-2-1",
-          input_modalities=["text"],
-          output_modalities=["image"]
-      )
-
-      # Initialize OmniLLM
-      omni = OmniLLM([ar_config, dit_config])
-
-      # Generate
-      outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
+      prompts = [{"prompt": "A scenic watercolor painting of a lighthouse at sunset"}]
+      sampling_params_list = [
+          SamplingParams(max_tokens=1),
+          OmniDiffusionSamplingParams(num_outputs_per_prompt=1),
+      ]
+      outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
       ```
 
       If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.

benchmarks/qwen3-omni/README.md

@@ -56,9 +56,9 @@ What it does:
 - Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
 - Uses `benchmarks/build_dataset/top100.txt` and writes to:
   - Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
-    - `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
-    - `omni_llm_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
-    - `omni_llm_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
+    - `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
+    - `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
+    - `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
   - Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
 
 Key checks:

benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh

@@ -26,12 +26,12 @@ else
                       --log-stats \
                       --log-dir $log_dir
     echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
-    echo "  - omni_llm_pipeline_text                       run dir/base name"
-    echo "  - omni_llm_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
-    echo "  - omni_llm_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
-    echo "  - omni_llm_pipeline_text.stage0.log                per-stage detailed logs"
-    echo "  - omni_llm_pipeline_text.stage1.log"
-    echo "  - omni_llm_pipeline_text.stage2.log"
+    echo "  - omni_pipeline_text                       run dir/base name"
+    echo "  - omni_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
+    echo "  - omni_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
+    echo "  - omni_pipeline_text.stage0.log                per-stage detailed logs"
+    echo "  - omni_pipeline_text.stage1.log"
+    echo "  - omni_pipeline_text.stage2.log"
     echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
     echo "  - outputs/             Generated txt and wav files, there should be 100 text and wav files generated respectively"
 fi

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml

@@ -10,8 +10,6 @@ stage_args:
     engine_args:
       model_stage: qwen3_tts
       model_arch: Qwen3TTSTalkerForConditionalGeneration
-      hf_overrides:
-        architectures: [Qwen3TTSTalkerForConditionalGeneration]
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       enforce_eager: false
@@ -43,8 +41,6 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3TTSCode2Wav
-      hf_overrides:
-        architectures: [Qwen3TTSCode2Wav]
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml

@@ -11,8 +11,6 @@ stage_args:
     engine_args:
       model_stage: qwen3_tts
       model_arch: Qwen3TTSTalkerForConditionalGeneration
-      hf_overrides:
-        architectures: [Qwen3TTSTalkerForConditionalGeneration]
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       enforce_eager: false
@@ -44,8 +42,6 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3TTSCode2Wav
-      hf_overrides:
-        architectures: [Qwen3TTSCode2Wav]
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true

docs/api/README.md

@@ -6,19 +6,13 @@ Main entry points for vLLM-Omni inference and serving.
 
 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
 - [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
-- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
 - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
 - [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
 - [vllm_omni.entrypoints.omni.Omni][]
-- [vllm_omni.entrypoints.omni.OmniBase][]
-- [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
-- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
-- [vllm_omni.entrypoints.omni_stage.OmniStage][]
-- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
-- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
+- [vllm_omni.entrypoints.omni_base.OmniBase][]
 
 ## Inputs
 
@@ -48,9 +42,7 @@ Engine classes for offline and online inference.
 - [vllm_omni.engine.OmniEngineCoreOutputs][]
 - [vllm_omni.engine.OmniEngineCoreRequest][]
 - [vllm_omni.engine.PromptEmbedsPayload][]
-- [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
 - [vllm_omni.engine.arg_utils.OmniEngineArgs][]
-- [vllm_omni.engine.input_processor.OmniInputProcessor][]
 - [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
 - [vllm_omni.engine.output_processor.OmniRequestState][]
 

docs/configuration/stage_configs.md

@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
 For offline (Assume necessary dependencies have ben imported):
 ```python
 model_name = "Qwen/Qwen2.5-Omni-7B"
-omni_llm = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
+omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
 ```
 
 For online serving:
@@ -30,7 +30,7 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to
 
 Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
 ```python
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
 stage_args:
   - stage_id: 0 # mark the unique id for each stage
     runtime: # The disaggregated configuration

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -1,4 +1,4 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
 stage_args:
   - stage_id: 0
     runtime:

docs/contributing/README.md

@@ -107,7 +107,7 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 - `[CI/Build]` for build or continuous integration improvements.
 - `[Doc]` for documentation fixes and improvements.
 - `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
-- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
+- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `Omni`/`AsyncOmni`, etc.)
 - `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
 - `[Core]` for changes in the core vLLM-Omni logic (e.g., `OmniProcessor`, `OmniARScheduler`, etc.)
 - `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.

docs/contributing/ci/CI_5levels.md

@@ -168,11 +168,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
-│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
-│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
-│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
-│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
 │   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)

docs/contributing/ci/tests_style.md

@@ -24,8 +24,6 @@ End-to-end tests verify the complete functionality of a system or component. For
 
 - **`tests/e2e/online_serving/`**: Tests for online serving scenarios (e.g., API server tests)
 
-**Example:** The test file for `vllm_omni/entrypoints/omni_llm.py` should be located at `tests/entrypoints/test_omni_llm.py`.
-
 ## Test Directory Structure
 
 The ideal directory structure mirrors the source code organization. Legend: `✅` = test exists, `⬜` = suggested to add.
@@ -75,11 +73,6 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
-│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
-│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
-│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
-│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
-│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
 │   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
@@ -170,7 +163,7 @@ vllm_omni/                                    tests/
 
 ### Naming Conventions
 
-- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
+- **Unit Tests**: Use `test_<module_name>.py` format. Example: `stage_utils.py` → `test_stage_utils.py`
 
 - **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`
 

docs/contributing/model/adding_omni_model.md

@@ -330,54 +330,46 @@ Stage transitions are the mechanism by which outputs from one stage are converte
 
 ### Where Stage Transitions Are Called
 
-Stage transitions happen automatically in the orchestrator (`OmniLLM` class) during the generation loop. Here's the detailed flow:
+Stage transitions happen automatically in the runtime orchestrator. Here's the detailed flow:
 
-1. **Location**: `vllm_omni/entrypoints/omni_llm.py` in the `_run_generation()` method
+1. **Location**: `vllm_omni/engine/orchestrator.py` in `_forward_to_next_stage()`
 2. **Trigger**: When a stage completes processing and produces outputs
 3. **Execution Flow**:
    ```python
-   # In omni_llm.py, _run_generation() method (around line 345-460)
+   # In orchestrator.py
+   next_stage_id = stage_id + 1
+   next_client = self.stage_clients[next_stage_id]
+   params = req_state.sampling_params_list[next_stage_id]
 
-   # Main orchestrator loop polls each stage for completed requests
-   for stage_id, stage in enumerate(self.stage_list):
-       result = stage.try_collect()  # Get completed request
-       if result is None:
-           continue
+   # Save current stage outputs so stage_input_processors can consume them.
+   self.stage_clients[stage_id].set_engine_outputs([output])
 
-       # Store outputs from this stage
-       engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
-       stage.set_engine_outputs(engine_outputs)
+   # THIS IS WHERE STAGE TRANSITION HAPPENS
+   next_inputs = next_client.process_engine_inputs(
+       stage_list=self.stage_clients,
+       prompt=req_state.prompt,
+   )
 
-       # Check if there's a next stage to forward to
-       next_stage_id = stage_id + 1
-       if next_stage_id < len(self.stage_list):
-           next_stage: OmniStage = self.stage_list[next_stage_id]
-
-           # THIS IS WHERE STAGE TRANSITION HAPPENS
-           next_inputs = next_stage.process_engine_inputs(
-               self.stage_list,
-               [request_id_to_prompt[req_id]]
-           )
-
-           # Submit to next stage
-           task = {
-               "type": OmniStageTaskType.GENERATE,
-               "request_id": req_id,
-               "engine_inputs": next_inputs[0],
-               "sampling_params": sampling_params_list[next_stage_id],
-           }
-           next_stage.submit(task)
+   # Build and submit request(s) to the next stage.
+   for next_input in next_inputs:
+       request = build_engine_core_request_from_tokens(
+           request_id=req_id,
+           prompt=next_input,
+           params=params,
+           model_config=self.stage_vllm_configs[next_stage_id].model_config,
+       )
+       await next_client.add_request_async(request)
    ```
 
 ### How Stage Transitions Work
 
 The stage transition process follows these steps:
 
-1. **Stage Completion**: When a stage finishes processing a request, it stores outputs via `stage.set_engine_outputs(engine_outputs)`
+1. **Stage Completion**: When a stage finishes processing a request, the orchestrator stores outputs via `stage_client.set_engine_outputs(...)`
 
 2. **Transition Detection**: The orchestrator checks if there's a next stage and calls `process_engine_inputs()` on it
 
-3. **Input Processing**: The `process_engine_inputs()` method in `OmniStage` (`omni_stage.py`) handles the transition:
+3. **Input Processing**: The stage input processor configured in stage YAML (under `vllm_omni/model_executor/stage_input_processors/`) handles the transition:
    ```python
    def process_engine_inputs(
        self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None

docs/contributing/profiling.md

@@ -23,30 +23,32 @@ export VLLM_PROFILER_MAX_ITERS=1
 The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
 ```python
 # Profile all stages
-omni_llm.start_profile()
+omni.start_profile()
 
 # Only profile Stage 1
-omni_llm.start_profile(stages=[1])
+omni.start_profile(stages=[1])
 ```
 
 ```python
 # Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
-omni_llm.start_profile(stages=[0, 2])
+omni.start_profile(stages=[0, 2])
 ```
 
 **Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
 
 ```python
-from vllm_omni import omni_llm
+from vllm_omni.entrypoints.omni import Omni
+
+omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
 
 profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
 # 1. Start profiling if enabled
 if profiler_enabled:
-    omni_llm.start_profile(stages=[0])
+    omni.start_profile(stages=[0])
 
 # Initialize generator
-omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
 total_requests = len(prompts)
 processed_count = 0
@@ -57,21 +59,21 @@ for stage_outputs in omni_generator:
     # ... [Output processing logic for text/audio would go here] ...
 
     # Update count to track when to stop profiling
-    processed_count += len(stage_outputs.request_output)
+    processed_count += 1
 
     # 2. Check if all requests are done to stop the profiler safely
     if profiler_enabled and processed_count >= total_requests:
         print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
 
         # Stop the profiler while workers are still active
-        omni_llm.stop_profile()
+        omni.stop_profile()
 
         # Wait for traces to flush to disk
         print("[Info] Waiting 30s for workers to write trace files to disk...")
         time.sleep(30)
         print("[Info] Trace export wait time finished.")
 
-omni_llm.close()
+omni.close()
 ```
 
 

docs/design/architecture_overview.md

@@ -67,12 +67,12 @@ According to analysis for current popular open-source models, most of them have
 | Component         | Description                                                                                                                              |
 | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
 | **OmniRouter**    | provide an intelligent router for Omni-modality requests dispatch                                                                        |
-| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
+| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
 | **AR**            | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management                                 |
 | **Diffusion**     | natively implemented and optimized using acceleration components                                                                         |
 | **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages                    |
 
-Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.
+Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
 
 ## Main features
 
@@ -127,10 +127,10 @@ Taking **Qwen3-Omni** as an example:
 The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
 
 ```
-# Create an omni_lm with HF model name.
+# Create an omni runtime with HF model name.
 from vllm_omni.entrypoints.omni import Omni
 
-omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
+omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
 
 # Example prompts.
 om_inputs = {"prompt": prompt,
@@ -140,7 +140,7 @@ om_inputs = {"prompt": prompt,
              }}
 
 # Generate texts and audio from the multi-modality inputs.
-outputs = omni_lm.generate(om_inputs, sampling_params_list)
+outputs = omni.generate(om_inputs, sampling_params_list)
 ```
 
 ## Online Serving

docs/design/module/ar_module.md

@@ -85,9 +85,6 @@ classDiagram
     class InputProcessor {
         +process_inputs() EngineCoreRequest
     }
-    class OmniInputProcessor {
-        +process_inputs() OmniEngineCoreRequest
-    }
 
     class VLLMOutputProcessor {
         +process_outputs() OutputProcessorOutput
@@ -96,7 +93,6 @@ classDiagram
         +process_outputs() OutputProcessorOutput
         +_route_and_normalize()
     }
-    InputProcessor <|-- OmniInputProcessor
     VLLMOutputProcessor <|-- MultimodalOutputProcessor
 ```
 
@@ -105,7 +101,7 @@ classDiagram
 - **Scheduler**: `OmniARScheduler` extends `vllm.v1.core.sched.scheduler.Scheduler` to enrich scheduled requests with omni-specific payloads
 - **Worker**: `GPUARWorker` extends `vllm.v1.worker.gpu_worker.Worker` to initialize AR-specific model runners
 - **ModelRunner**: `GPUARModelRunner` extends `OmniGPUModelRunner` → `vllm.v1.worker.gpu_model_runner.GPUModelRunner` to expose hidden states and handle multimodal outputs
-- **InputProcessor**: `OmniInputProcessor` extends `vllm.v1.engine.input_processor.InputProcessor` to serialize prompt embeddings and additional information
+- **InputProcessor**: Stage-0 uses upstream `vllm.v1.engine.input_processor.InputProcessor`; `AsyncOmniEngine` then restores omni-specific payloads (for example `additional_information` and `prompt_embeds`) when building `OmniEngineCoreRequest`
 - **OutputProcessor**: `MultimodalOutputProcessor` extends `vllm.v1.engine.output_processor.OutputProcessor` to route and accumulate multimodal outputs
 
 ## 3. Scheduler Design
@@ -118,7 +114,7 @@ The following diagram illustrates the request flow through the AR module compone
 
 ```mermaid
 flowchart TD
-    A[OmniInputProcessor] -->|OmniEngineCoreRequest| B[OmniARScheduler]
+    A[InputProcessor stage-0 in AsyncOmniEngine] -->|EngineCoreRequest then upgraded to OmniEngineCoreRequest| B[OmniARScheduler]
     B -->|schedule: OmniNewRequestData| C[GPUARWorker]
     C -->|SchedulerOutput| D[GPUARModelRunner]
     D -->|execute_model: None| E[Model Forward Pass]
@@ -297,15 +293,17 @@ The input/output processing pipeline handles serialization, routing, and accumul
 ```mermaid
 sequenceDiagram
     participant Client
-    participant OmniInputProcessor
+    participant AsyncOmniEngine
+    participant InputProcessor
     participant Scheduler
     participant ModelRunner
     participant MultimodalOutputProcessor
     participant Client
 
-    Client->>OmniInputProcessor: prompt + prompt_embeds + additional_info
-    OmniInputProcessor->>OmniInputProcessor: Serialize tensors to payloads
-    OmniInputProcessor->>Scheduler: OmniEngineCoreRequest (with payloads)
+    Client->>AsyncOmniEngine: prompt + prompt_embeds + additional_info
+    AsyncOmniEngine->>InputProcessor: process_inputs()
+    InputProcessor->>Scheduler: EngineCoreRequest
+    AsyncOmniEngine->>AsyncOmniEngine: _upgrade_to_omni_request() + serialize_additional_information()
     Scheduler->>ModelRunner: OmniNewRequestData (with payloads)
     ModelRunner->>ModelRunner: Decode payloads → CPU tensors
     ModelRunner->>ModelRunner: Overlay prompt_embeds on inputs_embeds
@@ -318,43 +316,18 @@ sequenceDiagram
     MultimodalOutputProcessor->>Client: RequestOutput (with multimodal_output)
 ```
 
-### OmniInputProcessor
+### Stage-0 Input Processing
 
-`OmniInputProcessor` extends the base input processor to serialize prompt embeddings and additional information for inter-stage transfer.
-
-#### Payload Serialization
-
-Converts PyTorch tensors to serialized payloads:
+Stage-0 now uses upstream `InputProcessor` directly, and `AsyncOmniEngine` upgrades the request to `OmniEngineCoreRequest` while restoring omni-specific payloads.
 
 ```python
-def process_inputs(self, ...) -> OmniEngineCoreRequest:
-    # Serialize prompt_embeds
-    if "prompt_embeds" in decoder_inputs:
-        pe_cpu = decoder_inputs["prompt_embeds"].detach().to("cpu").contiguous()
-        prompt_embeds_payload = PromptEmbedsPayload(
-            data=pe_cpu.numpy().tobytes(),
-            shape=[seq_len, hidden_size],
-            dtype=dtype_str,
-        )
-
-    # Serialize additional_information
-    if "additional_information" in decoder_inputs:
-        entries = {}
-        for key, value in raw_info.items():
-            if isinstance(value, torch.Tensor):
-                entry = AdditionalInformationEntry(
-                    tensor_data=value.numpy().tobytes(),
-                    tensor_shape=list(value.shape),
-                    tensor_dtype=dtype_str,
-                )
-            entries[key] = entry
-        additional_information_payload = AdditionalInformationPayload(entries=entries)
-
-    return OmniEngineCoreRequest(
-        # ... standard fields ...
-        prompt_embeds=prompt_embeds_payload,
-        additional_information=additional_information_payload,
-    )
+request = self.input_processor.process_inputs(
+    request_id=request_id,
+    prompt=prompt,
+    params=params,
+    supported_tasks=self.supported_tasks,
+)
+request = _upgrade_to_omni_request(request, prompt)
 ```
 
 ### MultimodalOutputProcessor
@@ -400,13 +373,13 @@ The AR module of vLLM-Omni extends vLLM through strategic inheritance and minima
 ### Key Design Patterns
 
 1. **Inheritance over composition**: Extends vLLM classes to preserve compatibility with existing scheduling, batching, and execution mechanisms
-2. **Payload serialization**: Uses `PromptEmbedsPayload` and `AdditionalInformationPayload` for efficient inter-stage data transfer
+2. **Payload serialization**: Uses serialized `additional_information` payloads together with prompt-embedding handoff for efficient inter-stage transfer
 3. **Two-phase execution**: Maintains vLLM's execute/sample separation for AR models while supporting single-phase execution for generation models
 4. **Multimodal routing**: Routes outputs by `output_type` and accumulates tensors incrementally to support streaming
 
 ### Differences from vLLM
 
-- **Payload support**: Serialized prompt embeddings and additional information enable direct transfer between pipeline stages
+- **Payload support**: Serialized additional information and prompt embeddings enable direct transfer between pipeline stages
 - **Multimodal handling**: Extended input/output processors support images, audio, and other modalities alongside text
 - **Hidden state exposure**: AR model runners expose per-request hidden states via `pooler_output` for downstream consumption
 - **Generation scheduler**: Fast-path scheduling for basic heterogeneous architectures that complete in one step

docs/design/module/async_omni_architecture.md

@@ -0,0 +1,203 @@
+# AsyncOmni Architecture (Qwen3-Omni Example)
+
+## 1. System Architecture
+
+```text
+• ┌─────────────────────────────────────────────────────────────────────────────────┐
+  │                                    API Layer                                    │
+  │  ┌─────────────────────────────────────┐  ┌──────────────────────────────────┐  │
+  │  │ AsyncOmni (EngineClient)            │  │ Omni                             │  │
+  │  │ • generate() / abort() / shutdown() │  │ • generate()                     │  │
+  │  │ • _final_output_handler()           │  │                                  |  │
+  │  └─────────────────────────────────────┘  └──────────────────────────────────┘  │
+  ├─────────────────────────────────────────────────────────────────────────────────┤
+  │                              Engine Layer (Proxy)                               │
+  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
+  │  │ AsyncOmniEngine                                                           │  │
+  │  │ • _bootstrap_orchestrator() & _initialize_stages()                        │  │
+  │  │ • add_request() / add_request_async() -> input_processor.process_inputs() │  │
+  │  │ • try_get_output() / try_get_output_async()                               │  │
+  │  └───────────────────┬─────────────────────────────────▲─────────────────────┘  │
+  │         request_queue (janus.Queue)        output_queue (janus.Queue)           │
+  ├──────────────────────┼─────────────────────────────────┼────────────────────────┤
+  │                      ▼        Orchestration Layer      │                        │
+  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
+  │  │ Orchestrator [background thread]                                          │  │
+  │  │ • _request_handler()                                                      │  │
+  │  │     -  stage_client.add_request_async() & _prewarm_async_chunk_stages()   │  │
+  │  │ • _orchestration_output_handler()                                         │  │
+  │  │     -  _process_stage_outputs() -> output_processors[i].process_outputs() │  │
+  │  │     -  _route_output() & _forward_to_next_stage()                         │  │
+  │  └──────────┬─────────────────────────┬────────────────────────┬─────────────┘  │
+  ├─────────────┼─────────────────────────┼────────────────────────┼────────────────┤
+  │             │                 Communication Layer              │                │
+  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
+  │  │ StageEngineCoreClient │ │ StageEngineCoreClient │ │ StageDiffusionClient  │  │
+  │  │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │  │
+  │  │ • Msgpack codec       │ │ • Msgpack codec       │ │ • Msgpack codec       │  │
+  │  └──────────┬────────────┘ └──────────┬────────────┘ └──────────┬────────────┘  │
+  │             ▼ ZMQ IPC                 ▼ ZMQ IPC                 ▼ ZMQ IPC       │
+  ├─────────────────────────────────────────────────────────────────────────────────┤
+  │                                 Execution Layer                                 │
+  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
+  │  │ StageCoreProc         │ │ StageCoreProc         │ │ DiffusionEngine       │  │
+  │  │ [background process]  │ │ [background process]  │ │ [background process]  │  │
+  │  └───────────────────────┘ └───────────────────────┘ └───────────────────────┘  │
+  └─────────────────────────────────────────────────────────────────────────────────┘
+```
+
+## 2. Execution Flow (Arrow Steps, one generate request)
+
+```text
+[1] App
+    -> AsyncOmni.generate(prompt, request_id)
+
+[2] AsyncOmni
+    -> _final_output_handler()   (started on first request)
+    -> AsyncOmniEngine.add_request(stage_id=0, ...)
+
+[3] AsyncOmniEngine.add_request
+    -> (if stage-0 is llm and input is not EngineCoreRequest)
+       InputProcessor.process_inputs()
+       OutputProcessor[0].add_request()
+    -> request_queue.put(add_request_msg)
+
+[4] Orchestrator._request_handler
+    -> _handle_add_request(msg)
+    -> stage_clients[0].add_request_async(...)
+
+[5] Orchestrator._orchestration_loop (loop)
+    -> poll stage output
+       - llm stage: await get_output_async()
+       - diffusion stage: get_diffusion_output_async()
+    -> (llm stage) output_processors[i].process_outputs(...)
+    -> _route_output(...)
+    -> if finished and not final_stage and non-async-chunk:
+         _forward_to_next_stage(...)
+         -> next_stage.add_request_async(...)
+    -> output_queue.put(output)
+
+[6] AsyncOmni._final_output_loop (background coroutine)
+    -> AsyncOmniEngine.try_get_output_async()
+    -> route by request_id to ClientRequestState.queue
+
+[7] AsyncOmni._process_orchestrator_results
+    -> read from ClientRequestState.queue
+    -> _process_single_result(...)
+    -> yield OmniRequestOutput
+
+[8] Exit condition
+    -> receive result["finished"] == True
+    -> generate() ends
+```
+
+## 3. Runtime Sequence (one generate request)
+
+```mermaid
+sequenceDiagram
+    participant APP as App
+    participant AO as AsyncOmni
+    participant ENG as AsyncOmniEngine
+    participant ORCH as Orchestrator
+    participant S0 as Stage-0 Client
+    participant SN as Next Stage Client
+
+    APP->>AO: generate
+    AO->>AO: start output_handler once
+    AO->>ENG: add_request(stage_id=0, ...)
+    ENG->>ENG: input_processor.process_inputs()
+    ENG->>ORCH: request_queue.put(add_request)
+
+    ORCH->>ORCH: _handle_add_request
+    ORCH->>S0: add_request_async
+
+    loop poll route forward
+        ORCH->>S0: get_output_async / get_diffusion_output_async
+        ORCH->>ORCH: _route_output
+        alt need forward to next stage
+            ORCH->>SN: add_request_async
+        end
+        ORCH-->>ENG: output_queue.put
+    end
+
+    AO->>ENG: try_get_output_async
+    ENG-->>AO: message
+    AO-->>APP: yield OmniRequestOutput
+```
+
+## 4. Comparison
+
+Previous topology (reference):
+
+```text
+┌────────────────────────────────────────────────────────────────────────────┐
+│ Main Process                                                               │
+│  ┌──────────────────────┐   ┌────────────────────────────────────────────┐ │
+│  │ generate()           │   │ final_output_handler()                     │ │
+│  └──────────────────────┘   └────────────────────────────────────────────┘ │
+└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
+  mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)
+           ▼▲                        ▼▲                        ▼▲
+┌───────────────────────┐  ┌───────────────────────┐  ┌──────────────────────┐
+│ Worker Proc-0         │  │ Worker Proc-1         │  │ Worker Proc-2        │
+│ (Thinker LLM)         │  │ (Talker LLM)          │  │ (Vocoder)            │
+│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
+│  │_stage_worker   │   │  │  │_stage_worker   │   │  │  │_stage_worker   │  │
+│  │_async()        │   │  │  │_async()        │   │  │  │_async()        │  │
+│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
+│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
+│  │output_handler()│   │  │  │output_handler()│   │  │  │output_handler()│  │
+│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
+└──────────┬────────────┘  └──────────┬────────────┘  └──────────┬───────────┘
+       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ
+┌──────────────────────┐   ┌──────────────────────┐   ┌──────────────────────┐
+│ EngineCore Proc-0    │   │ EngineCore Proc-1    │   │ EngineCore Proc-2    │
+│ (Thinker)            │   │ (Talker)             │   │ (Vocoder)            │
+└──────────────────────┘   └──────────────────────┘   └──────────────────────┘
+```
+
+Current topology:
+
+```text
+┌────────────────────────────────────────────────────────────────────────────┐
+│ Main Process                                                               │
+│  ┌──────────────────────────────────────────────────────────────────────┐  │
+│  │ Main Thread                                                          │  │
+│  │  ┌──────────────────────┐   ┌─────────────────────────────────────┐  │  │
+│  │  │ generate()           │   │ final_output_handler()              │  │  │
+│  │  └──────────────────────┘   └─────────────────────────────────────┘  │  │
+│  └──────────────────────────────────────────────────────────────────────┘  │
+│         janus.Queue (request_queue) ▼  ▲ janus.Queue (output_queue)        │
+│  ┌──────────────────────────────────────────────────────────────────────┐  │
+│  │ Orchestrator Thread                                                  │  │
+│  │  ┌──────────────────────┐  ┌──────────────────────────────────────┐  │  │
+│  │  │ _request_handler()   │  │ _orchestration_output_handler()      │  │  │
+│  │  └──────────────────────┘  └──────────────────────────────────────┘  │  │
+│  │  ┌────────────────────────────────────────────────────────────────┐  │  │
+│  │  │ _orchestration_loop(): poll/process/route outputs for all stages│  │  │
+│  │  └────────────────────────────────────────────────────────────────┘  │  │
+│  └───────┬─────────────────────────┬─────────────────────────┬──────────┘  │
+└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
+       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ  
+  ┌──────────────────────┐  ┌──────────────────────┐  ┌──────────────────────┐
+  │ EngineCore Proc-0    │  │ EngineCore Proc-1    │  │ EngineCore Proc-2    │
+  │ (Thinker)            │  │ (Talker)             │  │ (Vocoder)            │
+  └──────────────────────┘  └──────────────────────┘  └──────────────────────┘
+```
+
+
+Test scripts:
+```bash
+# enter offline inference folder.
+cd examples/offline_inference/qwen2_5_omni
+python end2end.py --output-dir output_audio --query-type use_mixed_modalities
+
+cd ../qwen3_omni
+python end2end.py --output-dir output_audio --query-type text --async-chunk --enable-stats
+
+cd ../bagel
+python end2end.py --prompts "A cute cat"
+
+cd ../text_to_image
+python text_to_image.py --prompt "a cup of coffee on the table" --output output.png
+```

docs/features/custom_pipeline.md

@@ -97,7 +97,7 @@ outputs = omni.generate(
 )
 
 # Access custom trajectory data
-output = outputs[0].request_output[0]
+output = outputs[0].request_output
 print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
 print(f"Trajectory latents shape: {output.latents.shape}")
 ```

docs/getting_started/quickstart.md

@@ -42,7 +42,7 @@ if __name__ == "__main__":
     omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output[0].images
+    images = outputs[0].request_output.images
     images[0].save("coffee.png")
 ```
 
@@ -70,7 +70,7 @@ if __name__ == "__main__":
     ]
     omni_outputs = omni.generate(prompts)
     for i_prompt, prompt_output in enumerate(omni_outputs):
-        this_request_output = prompt_output.request_output[0]
+        this_request_output = prompt_output.request_output
         this_images = this_request_output.images
         for i_image, image in enumerate(this_images):
             image.save(f"p{i_prompt}-img{i_image}.jpg")

docs/user_guide/diffusion/quantization/gguf.md

@@ -22,10 +22,16 @@
 CLI (examples/offline_inference/text_to_image/text_to_image.py)
   |
   v
-Omni (vllm_omni/entrypoints/omni.py)
+Omni (vllm_omni/entrypoints/__init__.py)
   |
   v
-OmniStage (diffusion)
+AsyncOmniEngine
+  |
+  v
+Orchestrator
+  |
+  v
+StageDiffusionClient
   |
   v
 DiffusionWorker
@@ -64,7 +70,13 @@ _generate_with_async_omni
 AsyncOmni
   |
   v
-DiffusionEngine
+AsyncOmniEngine
+  |
+  v
+Orchestrator
+  |
+  v
+StageDiffusionClient
   |
   v
 OmniRequestOutput

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -19,7 +19,7 @@ if __name__ == "__main__":
     omni = Omni(model="Qwen/Qwen-Image")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output[0].images
+    images = outputs[0].request_output.images
     images[0].save("coffee.png")
 ```
 
@@ -37,7 +37,7 @@ if __name__ == "__main__":
     ]
     outputs = omni.generate(prompts)
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 !!! info
@@ -72,7 +72,7 @@ if __name__ == "__main__":
       }
     ])
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 ## Local CLI Usage

examples/offline_inference/bagel/end2end.py

@@ -34,7 +34,7 @@ def parse_args():
         help="Path to input image for img2img.",
     )
 
-    # OmniLLM init args
+    # Omni runtime init args
     parser.add_argument("--log-stats", action="store_true", default=False)
     parser.add_argument("--init-sleep-seconds", type=int, default=20)
     parser.add_argument("--batch-timeout", type=int, default=5)
@@ -163,25 +163,17 @@ def main():
 
     omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 
-    for i, req_output in enumerate(omni_outputs):
+    img_idx = 0
+    for req_output in omni_outputs:
         images = getattr(req_output, "images", None)
-        if not images and hasattr(req_output, "output"):
-            if isinstance(req_output.output, list):
-                images = req_output.output
-            else:
-                images = [req_output.output]
+        if not images:
+            continue
 
-        if images:
-            for j, img in enumerate(images):
-                img.save(f"output_{i}_{j}.png")
-
-        if hasattr(req_output, "request_output") and req_output.request_output:
-            for stage_out in req_output.request_output:
-                if hasattr(stage_out, "images") and stage_out.images:
-                    for k, img in enumerate(stage_out.images):
-                        save_path = f"output_{i}_stage_{getattr(stage_out, 'stage_id', '?')}_{k}.png"
-                        img.save(save_path)
-                        print(f"[Info] Saved stage output image to {save_path}")
+        for j, img in enumerate(images):
+            save_path = f"output_{img_idx}_{j}.png"
+            img.save(save_path)
+            print(f"[Info] Saved image to {save_path}")
+        img_idx += 1
 
     print(omni_outputs)
 

examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py

@@ -167,27 +167,26 @@ def run_e2e():
     print(f"Received {len(outputs)} outputs.")
     for i, output in enumerate(outputs):
         try:
-            ro_list = output.request_output or []
-            if not ro_list:
+            ro = output.request_output
+            if ro is None:
                 print("No request_output found.")
                 continue
 
-            for ro in ro_list:
-                # Multimodal output may be attached to RequestOutput or CompletionOutput.
-                mm = getattr(ro, "multimodal_output", None)
-                if not mm and ro.outputs:
-                    mm = getattr(ro.outputs[0], "multimodal_output", None)
+            # Multimodal output may be attached to RequestOutput or CompletionOutput.
+            mm = getattr(ro, "multimodal_output", None)
+            if not mm and ro.outputs:
+                mm = getattr(ro.outputs[0], "multimodal_output", None)
 
-                if mm:
-                    print(f"Multimodal output keys: {mm.keys()}")
-                    if "audio" in mm:
-                        audio_out = mm["audio"]
-                        print(f"Generated Audio Shape: {audio_out.shape}")
-                        out_path = f"output_{i}.wav"
-                        sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
-                        print(f"Saved audio to {out_path}")
-                else:
-                    print("No multimodal output found.")
+            if mm:
+                print(f"Multimodal output keys: {mm.keys()}")
+                if "audio" in mm:
+                    audio_out = mm["audio"]
+                    print(f"Generated Audio Shape: {audio_out.shape}")
+                    out_path = f"output_{i}.wav"
+                    sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
+                    print(f"Saved audio to {out_path}")
+            else:
+                print("No multimodal output found.")
         except Exception as e:
             print(f"Error inspecting output: {e}")
     omni.close()

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -209,7 +209,7 @@ async def main():
     if not outputs:
         raise ValueError("No output produced from omni.generate()")
 
-    first_out = outputs[0].request_output[0]
+    first_out = outputs[0].request_output
     req_out: OmniRequestOutput = first_out
 
     # Verify trajectory data (from custom pipeline)

examples/offline_inference/fish_speech/end2end.py

@@ -129,8 +129,14 @@ def main(args):
 
     t_start = time.perf_counter()
     for stage_outputs in omni.generate(inputs):
-        for output in stage_outputs.request_output:
-            _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
+        request_output = stage_outputs.request_output
+        if request_output is None or not request_output.outputs:
+            continue
+        _save_wav(
+            output_dir,
+            request_output.request_id,
+            request_output.outputs[0].multimodal_output,
+        )
     t_end = time.perf_counter()
     logger.info("Total inference time: %.1f ms", (t_end - t_start) * 1000)
 

examples/offline_inference/glm_image/end2end.py

@@ -310,36 +310,35 @@ def main(args: argparse.Namespace) -> None:
 
     output_count = 0
     for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True):
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "image":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
+            request_id = output.request_id
 
-                # Get generated images
-                images = output.images if hasattr(output, "images") else []
-                if not images and hasattr(output, "multimodal_output"):
-                    images = output.multimodal_output.get("images", [])
+            # Get generated images
+            images = output.images if hasattr(output, "images") else []
+            if not images and hasattr(output, "multimodal_output"):
+                images = output.multimodal_output.get("images", [])
 
-                # Save each generated image
-                for idx, img in enumerate(images):
-                    if args.num_prompts == 1 and len(images) == 1:
-                        output_path = args.output
-                    else:
-                        base, ext = os.path.splitext(args.output)
-                        output_path = f"{base}_{request_id}_{idx}{ext}"
+            # Save each generated image
+            for idx, img in enumerate(images):
+                if args.num_prompts == 1 and len(images) == 1:
+                    output_path = args.output
+                else:
+                    base, ext = os.path.splitext(args.output)
+                    output_path = f"{base}_{request_id}_{idx}{ext}"
 
-                    if isinstance(img, Image.Image):
-                        save_image(img, output_path)
-                    else:
-                        print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
+                if isinstance(img, Image.Image):
+                    save_image(img, output_path)
+                else:
+                    print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
 
-                    output_count += 1
+                output_count += 1
 
         elif stage_outputs.final_output_type == "text":
             # AR stage output (intermediate, for debugging)
             if args.verbose:
-                for output in stage_outputs.request_output:
-                    print(f"AR output for request {output.request_id}:")
-                    print(f"  Token count: {len(output.outputs[0].token_ids)}")
+                print(f"AR output for request {output.request_id}:")
+                print(f"  Token count: {len(output.outputs[0].token_ids)}")
 
     gen_time = time.time() - gen_start_time
     print(f"\nGeneration completed in {gen_time:.2f}s")

examples/offline_inference/helios/end2end.py

@@ -293,12 +293,11 @@ def main():
                 )
 
             if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
-                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
-                    inner_output = first_item.request_output[0]
-                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
-                        frames = inner_output.images[0] if inner_output.images else None
-                        if frames is None:
-                            raise ValueError("No video frames found in output.")
+                inner_output = first_item.request_output
+                if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
+                    frames = inner_output.images[0] if inner_output.images else None
+                    if frames is None:
+                        raise ValueError("No video frames found in output.")
             elif hasattr(first_item, "images") and first_item.images:
                 frames = first_item.images
             else:

examples/offline_inference/hunyuan_image3/image_to_text.py

@@ -9,7 +9,7 @@ from PIL import Image
 from vllm_omni.entrypoints.omni import Omni
 
 """
-The tencent/HunyuanImage-3.0-Instruct base model is built on the Hunyuan v1 architecture, specifically the tencent/Hunyuan-A13B-Instruct model. It utilizes two tokenizer delimiter templates:
+The tencent/HunyuanImage-3.0-Instruct base model uses the tencent/Hunyuan-A13B-Instruct backbone. It utilizes two tokenizer delimiter templates:
 
 1) Pretrained template (default for gen_text mode), which concatenates system, image
    tokens, and user question WITHOUT role delimiters:
@@ -70,8 +70,8 @@ def main(args: argparse.Namespace) -> None:
     prompts = [prompt_dict]
     omni_outputs = omni.generate(prompts=prompts)
 
-    prompt_text = omni_outputs[0].request_output[0].prompt
-    generated_text = omni_outputs[0].request_output[0].outputs[0].text
+    prompt_text = omni_outputs[0].request_output.prompt
+    generated_text = omni_outputs[0].request_output.outputs[0].text
     print(f"Prompt: {prompt_text}")
     print(f"Text: {generated_text}")
 

examples/offline_inference/image_to_image/image_edit.py

@@ -463,12 +463,12 @@ def main():
         raise ValueError("No output generated from omni.generate()")
 
     # Extract images from OmniRequestOutput
-    # omni.generate() returns list[OmniRequestOutput], extract images from request_output[0].images
+    # omni.generate() returns list[OmniRequestOutput], extract images from request_output.images
     first_output = outputs[0]
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output[0]
+    req_out = first_output.request_output
     if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
         raise ValueError("Invalid request_output structure or missing 'images' key")
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -338,8 +338,6 @@ def main():
             audio = frames.multimodal_output["audio"]
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
-            if isinstance(inner_output, list):
-                inner_output = inner_output[0] if inner_output else None
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -99,14 +99,12 @@ def main() -> None:
 
     lines: list[str] = []
     for stage_outputs in outputs:
-        req_outputs = getattr(stage_outputs, "request_output", stage_outputs)
-        req_outputs = req_outputs if isinstance(req_outputs, list) else [req_outputs]
-        for ro in req_outputs:
-            text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
-            lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
-            lines.append("answer:\n")
-            lines.append(text.strip() + "\n")
-            lines.append("\n")
+        ro = getattr(stage_outputs, "request_output", stage_outputs)
+        text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
+        lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
+        lines.append("answer:\n")
+        lines.append(text.strip() + "\n")
+        lines.append("\n")
 
     print("\n".join(lines))
 

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -148,19 +148,16 @@ def _collect_images(outputs: list) -> list[torch.Tensor]:
     """Extract all image tensors produced by the final (DiT) stage."""
     images: list[torch.Tensor] = []
     for out in outputs:
-        ro_list = getattr(out, "request_output", out)
-        if not isinstance(ro_list, list):
-            ro_list = [ro_list]
-        for ro_item in ro_list:
-            for completion in getattr(ro_item, "outputs", None) or []:
-                mm = getattr(completion, "multimodal_output", None)
-                if not isinstance(mm, dict) or "image" not in mm:
-                    raise RuntimeError(f"Missing image in multimodal output: {mm}")
-                payload = mm["image"]
-                for tensor in payload if isinstance(payload, list) else [payload]:
-                    if not isinstance(tensor, torch.Tensor):
-                        raise TypeError(f"Expected image tensor, got {type(tensor)}")
-                    images.append(tensor)
+        ro_item = getattr(out, "request_output", out)
+        for completion in getattr(ro_item, "outputs", None) or []:
+            mm = getattr(completion, "multimodal_output", None)
+            if not isinstance(mm, dict) or "image" not in mm:
+                raise RuntimeError(f"Missing image in multimodal output: {mm}")
+            payload = mm["image"]
+            for tensor in payload if isinstance(payload, list) else [payload]:
+                if not isinstance(tensor, torch.Tensor):
+                    raise TypeError(f"Expected image tensor, got {type(tensor)}")
+                images.append(tensor)
     return images
 
 

examples/offline_inference/mimo_audio/end2end.py

@@ -180,11 +180,11 @@ def main(args):
     # Get the query function and call it with appropriate parameters
     query_func = query_map[args.query_type]
 
-    omni_llm = Omni(
+    omni = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
         log_stats=args.enable_stats,
-        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
+        log_file=("omni_pipeline.log" if args.enable_stats else None),
         init_sleep_seconds=args.init_sleep_seconds,
         batch_timeout=args.batch_timeout,
         init_timeout=args.init_timeout,
@@ -290,7 +290,7 @@ def main(args):
     prompts = [copy.deepcopy(query_result) for _ in range(args.num_prompts)]
 
     print("prompts", prompts)
-    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
+    omni_outputs = omni.generate(prompts, sampling_params_list)
 
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     if args.query_type is not None:
@@ -298,45 +298,44 @@ def main(args):
     os.makedirs(output_dir, exist_ok=True)
 
     for stage_outputs in omni_outputs:
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                text_output = output.outputs[0].text
-                # Save aligned text file per request
-                prompt_text = output.prompt
-                out_txt = os.path.join(output_dir, f"{request_id}.txt")
-                lines = []
-                lines.append("Prompt:\n")
-                lines.append(str(prompt_text) + "\n")
-                lines.append("vllm_text_output:\n")
-                lines.append(str(text_output).strip() + "\n")
-                try:
-                    with open(out_txt, "w", encoding="utf-8") as f:
-                        print("lines", lines)
-                        f.writelines(lines)
-                except Exception as e:
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
+            request_id = output.request_id
+            text_output = output.outputs[0].text
+            # Save aligned text file per request
+            prompt_text = output.prompt
+            out_txt = os.path.join(output_dir, f"{request_id}.txt")
+            lines = []
+            lines.append("Prompt:\n")
+            lines.append(str(prompt_text) + "\n")
+            lines.append("vllm_text_output:\n")
+            lines.append(str(text_output).strip() + "\n")
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    print("lines", lines)
+                    f.writelines(lines)
+            except Exception as e:
+                print(f"[Warn] Failed writing text file {out_txt}: {e}")
+            print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
         elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output.get("audio")
+            request_id = output.request_id
+            audio_tensor = output.outputs[0].multimodal_output.get("audio")
 
-                if audio_tensor is None:
-                    continue
+            if audio_tensor is None:
+                continue
 
-                output_wav = os.path.join(output_dir, f"{request_id}.wav")
+            output_wav = os.path.join(output_dir, f"{request_id}.wav")
 
-                # Convert to numpy array and ensure correct format
-                audio_numpy = audio_tensor.float().detach().cpu().numpy()
+            # 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=24000, format="WAV")
-                print(f"Request ID: {request_id}, Audio saved to {output_wav}")
+            # Save audio file with explicit WAV format
+            sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+            print(f"Request ID: {request_id}, Audio saved to {output_wav}")
 
 
 def parse_args():

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -320,7 +320,7 @@ def main(args):
         query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
     else:
         query_result = query_func()
-    omni_llm = Omni(
+    omni = Omni(
         model=model_name,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
@@ -378,9 +378,11 @@ def main(args):
             prompt["modalities"] = output_modalities
 
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
-    if profiler_enabled:
-        omni_llm.start_profile(stages=[0])
-    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+    if profiler_enabled and hasattr(omni, "start_profile"):
+        omni.start_profile(stages=[0])
+    elif profiler_enabled:
+        print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.")
+    omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
@@ -389,43 +391,42 @@ def main(args):
     total_requests = len(prompts)
     processed_count = 0
     for stage_outputs in omni_generator:
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                text_output = output.outputs[0].text
-                # Save aligned text file per request
-                prompt_text = output.prompt
-                out_txt = os.path.join(output_dir, f"{request_id}.txt")
-                lines = []
-                lines.append("Prompt:\n")
-                lines.append(str(prompt_text) + "\n")
-                lines.append("vllm_text_output:\n")
-                lines.append(str(text_output).strip() + "\n")
-                try:
-                    with open(out_txt, "w", encoding="utf-8") as f:
-                        f.writelines(lines)
-                except Exception as e:
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+            request_id = output.request_id
+            text_output = output.outputs[0].text
+            # Save aligned text file per request
+            prompt_text = output.prompt
+            out_txt = os.path.join(output_dir, f"{request_id}.txt")
+            lines = []
+            lines.append("Prompt:\n")
+            lines.append(str(prompt_text) + "\n")
+            lines.append("vllm_text_output:\n")
+            lines.append(str(text_output).strip() + "\n")
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.writelines(lines)
+            except Exception as e:
+                print(f"[Warn] Failed writing text file {out_txt}: {e}")
+            print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output["audio"]
-                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
-                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+            request_id = output.request_id
+            audio_tensor = output.outputs[0].multimodal_output["audio"]
+            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+            sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
+            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
-        processed_count += len(stage_outputs.request_output)
-        if profiler_enabled and processed_count >= total_requests:
+        processed_count += 1
+        if profiler_enabled and hasattr(omni, "stop_profile") and processed_count >= total_requests:
             print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
             # Stop the profiler while workers are still alive
-            omni_llm.stop_profile()
+            omni.stop_profile()
 
             print("[Info] Waiting 30s for workers to write massive trace files to disk...")
             time.sleep(30)
             print("[Info] Trace export wait finished.")
 
-    omni_llm.close()
+    omni.close()
 
 
 def parse_args():

examples/offline_inference/qwen3_omni/end2end.py

@@ -294,7 +294,7 @@ def main(args):
     else:
         query_result = query_func()
 
-    omni_llm = Omni(
+    omni = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
@@ -354,8 +354,8 @@ def main(args):
 
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
     if profiler_enabled:
-        omni_llm.start_profile(stages=[0])
-    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+        omni.start_profile(stages=[0])
+    omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     os.makedirs(output_dir, exist_ok=True)
@@ -366,51 +366,50 @@ def main(args):
     print(f"query type: {args.query_type}")
 
     for stage_outputs in omni_generator:
+        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                text_output = output.outputs[0].text
-                # Save aligned text file per request
-                prompt_text = output.prompt
-                out_txt = os.path.join(output_dir, f"{request_id}.txt")
-                lines = []
-                lines.append("Prompt:\n")
-                lines.append(str(prompt_text) + "\n")
-                lines.append("vllm_text_output:\n")
-                lines.append(str(text_output).strip() + "\n")
-                try:
-                    with open(out_txt, "w", encoding="utf-8") as f:
-                        f.writelines(lines)
-                except Exception as e:
-                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
-                print(f"Request ID: {request_id}, Text saved to {out_txt}")
+            request_id = output.request_id
+            text_output = output.outputs[0].text
+            # Save aligned text file per request
+            prompt_text = output.prompt
+            out_txt = os.path.join(output_dir, f"{request_id}.txt")
+            lines = []
+            lines.append("Prompt:\n")
+            lines.append(str(prompt_text) + "\n")
+            lines.append("vllm_text_output:\n")
+            lines.append(str(text_output).strip() + "\n")
+            try:
+                with open(out_txt, "w", encoding="utf-8") as f:
+                    f.writelines(lines)
+            except Exception as e:
+                print(f"[Warn] Failed writing text file {out_txt}: {e}")
+            print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output["audio"]
-                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+            request_id = output.request_id
+            audio_tensor = output.outputs[0].multimodal_output["audio"]
+            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
 
-                # Convert to numpy array and ensure correct format
-                audio_numpy = audio_tensor.float().detach().cpu().numpy()
+            # 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=24000, 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=24000, format="WAV")
+            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
-        processed_count += len(stage_outputs.request_output)
+        processed_count += 1
         if profiler_enabled and processed_count >= total_requests:
             print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
             # Stop the profiler while workers are still alive
-            omni_llm.stop_profile()
+            omni.stop_profile()
 
             print("[Info] Waiting 30s for workers to write trace files to disk...")
             time.sleep(30)
             print("[Info] Trace export wait time finished.")
-    omni_llm.close()
+    omni.close()
 
 
 def parse_args():

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -231,52 +231,47 @@ async def run_single_request(
             sampling_params_list=sampling_params_list,
             output_modalities=output_modalities,
         ):
-            if not isinstance(omni_output.request_output, list):
-                outputs_list = [omni_output.request_output]
-            else:
-                outputs_list = omni_output.request_output
+            output = omni_output.request_output
+            if omni_output.final_output_type == "text":
+                if stage_0_first_output_ts is None:
+                    stage_0_first_output_ts = time.perf_counter()
+                text_output = output.outputs[0].text
+                if output.finished:
+                    text_parts.append(text_output)
+            elif omni_output.final_output_type == "audio":
+                mm_out = output.outputs[0].multimodal_output
+                if mm_out and "audio" in mm_out:
+                    if first_audio_ts is None:
+                        first_audio_ts = time.perf_counter()
+                    if audio_sr is None and "sr" in mm_out:
+                        sr_val = mm_out["sr"]
+                        audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
+                        samplerate = audio_sr
+                    audio_data = mm_out["audio"]
+                    if isinstance(audio_data, list):
+                        new_chunks = audio_data[audio_list_consumed:]
+                        audio_list_consumed = len(audio_data)
+                    elif isinstance(audio_data, torch.Tensor):
+                        new_chunks = [audio_data]
+                        audio_last_tensor = audio_data
+                    else:
+                        new_chunks = []
 
-            for output in outputs_list:
-                if omni_output.final_output_type == "text":
-                    if stage_0_first_output_ts is None:
-                        stage_0_first_output_ts = time.perf_counter()
-                    text_output = output.outputs[0].text
-                    if output.finished:
-                        text_parts.append(text_output)
-                elif omni_output.final_output_type == "audio":
-                    mm_out = output.outputs[0].multimodal_output
-                    if mm_out and "audio" in mm_out:
-                        if first_audio_ts is None:
-                            first_audio_ts = time.perf_counter()
-                        if audio_sr is None and "sr" in mm_out:
-                            sr_val = mm_out["sr"]
-                            audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
-                            samplerate = audio_sr
-                        audio_data = mm_out["audio"]
-                        if isinstance(audio_data, list):
-                            new_chunks = audio_data[audio_list_consumed:]
-                            audio_list_consumed = len(audio_data)
-                        elif isinstance(audio_data, torch.Tensor):
-                            new_chunks = [audio_data]
-                            audio_last_tensor = audio_data
-                        else:
-                            new_chunks = []
-
-                        if stream_audio_to_disk and new_chunks:
-                            if sf_writer is None:
-                                sf_writer = sf.SoundFile(
-                                    wav_file,
-                                    mode="w",
-                                    samplerate=samplerate,
-                                    channels=1,
-                                    subtype="FLOAT",
-                                )
-                            for chunk in new_chunks:
-                                chunk_np = chunk.float().detach().cpu().numpy().flatten()
-                                sf_writer.write(chunk_np)
-                                audio_samples_written += len(chunk_np)
-                        else:
-                            audio_chunks.extend(new_chunks)
+                    if stream_audio_to_disk and new_chunks:
+                        if sf_writer is None:
+                            sf_writer = sf.SoundFile(
+                                wav_file,
+                                mode="w",
+                                samplerate=samplerate,
+                                channels=1,
+                                subtype="FLOAT",
+                            )
+                        for chunk in new_chunks:
+                            chunk_np = chunk.float().detach().cpu().numpy().flatten()
+                            sf_writer.write(chunk_np)
+                            audio_samples_written += len(chunk_np)
+                    else:
+                        audio_chunks.extend(new_chunks)
     finally:
         if sf_writer is not None:
             sf_writer.close()

examples/offline_inference/qwen3_tts/end2end.py

@@ -51,16 +51,68 @@ def _estimate_prompt_len(
 
             tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
             cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
-            _cache[model_name] = (tok, getattr(cfg, "talker_config", None))
 
-        tok, tcfg = _cache[model_name]
+            # Load speech tokenizer (codec encoder) for exact ref_code_len.
+            speech_tok = None
+            try:
+                import os
+
+                from transformers.utils import cached_file
+
+                from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
+
+                st_cfg_path = cached_file(model_name, "speech_tokenizer/config.json")
+                if st_cfg_path:
+                    speech_tok = Qwen3TTSTokenizer.from_pretrained(
+                        os.path.dirname(st_cfg_path), torch_dtype=torch.bfloat16
+                    )
+                    logger.info("Loaded speech tokenizer for exact ref_code_len estimation")
+            except Exception as e:
+                logger.debug("Could not load speech tokenizer: %s", e)
+
+            _cache[model_name] = (tok, getattr(cfg, "talker_config", None), speech_tok)
+
+        tok, tcfg, speech_tok = _cache[model_name]
         task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
+
+        def _estimate_ref_code_len(ref_audio: object) -> int | None:
+            """Encode ref_audio with the actual codec to get exact frame count."""
+            if not isinstance(ref_audio, (str, list)):
+                return None
+            audio_path = ref_audio[0] if isinstance(ref_audio, list) else ref_audio
+            if not isinstance(audio_path, str) or not audio_path.strip():
+                return None
+            try:
+                from vllm.multimodal.media import MediaConnector
+
+                connector = MediaConnector(allowed_local_media_path="/")
+                audio, sr = connector.fetch_audio(audio_path)
+                import numpy as np
+
+                wav_np = np.asarray(audio, dtype=np.float32)
+
+                if speech_tok is not None:
+                    enc = speech_tok.encode(wav_np, sr=int(sr), return_dict=True)
+                    ref_code = getattr(enc, "audio_codes", None)
+                    if isinstance(ref_code, list):
+                        ref_code = ref_code[0] if ref_code else None
+                    if ref_code is not None and hasattr(ref_code, "shape"):
+                        shape = ref_code.shape
+                        return int(shape[0]) if len(shape) == 2 else int(shape[1]) if len(shape) == 3 else None
+
+                # Fallback: estimate from duration
+                codec_hz = getattr(tcfg, "codec_frame_rate", None) or 12
+                return int(len(audio) / sr * codec_hz)
+            except Exception:
+                return None
+
         return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
             additional_information=additional_information,
             task_type=task_type,
             tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
             codec_language_id=getattr(tcfg, "codec_language_id", None),
             spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
+            estimate_ref_code_len=_estimate_ref_code_len,
         )
     except Exception as exc:
         logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
@@ -325,8 +377,8 @@ def main(args):
     for batch_start in range(0, len(inputs), batch_size):
         batch = inputs[batch_start : batch_start + batch_size]
         for stage_outputs in omni.generate(batch):
-            for output in stage_outputs.request_output:
-                _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
+            output = stage_outputs.request_output
+            _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
 
 
 async def main_streaming(args):

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -177,7 +177,7 @@ def main():
     output = outputs[0]
     if not hasattr(output, "request_output") or not output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
-    request_output = output.request_output[0]
+    request_output = output.request_output
     if not hasattr(request_output, "multimodal_output"):
         raise ValueError("No multimodal_output found in request_output")
 

examples/offline_inference/text_to_image/README.md

@@ -16,7 +16,7 @@ if __name__ == "__main__":
     omni = Omni(model="Qwen/Qwen-Image")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output[0].images
+    images = outputs[0].request_output.images
     images[0].save("coffee.png")
 ```
 
@@ -34,7 +34,7 @@ if __name__ == "__main__":
     ]
     outputs = omni.generate(prompts)
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 !!! info
@@ -69,7 +69,7 @@ if __name__ == "__main__":
       }
     ])
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 ## Local CLI Usage

examples/offline_inference/text_to_image/gradio_demo.py

@@ -112,7 +112,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
         )
         images_outputs = []
         for output in outputs:
-            req_out = output.request_output[0]
+            req_out = output.request_output
             if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
                 raise ValueError("Invalid request_output structure or missing 'images' key")
             images = req_out.images

examples/offline_inference/text_to_image/text_to_image.py

@@ -14,7 +14,6 @@ from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
-from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
 
@@ -400,20 +399,18 @@ def main():
         else:
             print("[Profiler] No valid profiling data returned.")
 
-    # Extract images from OmniRequestOutput
-    # omni.generate() returns list[OmniRequestOutput], extract images from the first output
+    # omni.generate() returns list[OmniRequestOutput]
     if not outputs or len(outputs) == 0:
         raise ValueError("No output generated from omni.generate()")
     logger.info(f"Outputs: {outputs}")
 
-    # Extract images from request_output[0]['images']
     first_output = outputs[0]
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output[0]
-    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
-        raise ValueError("Invalid request_output structure or missing 'images' key")
+    req_out = first_output.request_output
+    if not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images'.")
 
     images = req_out.images
     if not images:

examples/offline_inference/text_to_video/text_to_video.py

@@ -240,8 +240,6 @@ def main():
             audio = frames.multimodal_output["audio"]
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
-            if isinstance(inner_output, list):
-                inner_output = inner_output[0] if inner_output else None
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]

requirements/common.txt

@@ -17,3 +17,4 @@ einops>=0.8.1
 prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0
+janus>=1.0.0

tests/comfyui/test_comfyui_integration.py

@@ -11,6 +11,7 @@ import time
 import traceback
 from collections.abc import Iterable, Sequence
 from enum import StrEnum, auto
+from types import SimpleNamespace
 from typing import Any, NamedTuple
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -31,6 +32,7 @@ from vllm import SamplingParams
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.utils.argparse_utils import FlexibleArgumentParser
 
+from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
 from vllm_omni.entrypoints.cli.serve import OmniServeCommand
 from vllm_omni.inputs.data import OmniSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
@@ -43,7 +45,7 @@ class ServerCase(NamedTuple):
 
     served_model: str
     stage_list: list
-    stage_configs: list[dict]
+    stage_configs: list[Any]
     outputs: list[OmniRequestOutput]
 
 
@@ -262,6 +264,40 @@ def _assert_model_param_values(received: OmniSamplingParams, expected: dict):
         )
 
 
+def _make_stage_config(
+    stage_type: str,
+    *,
+    is_comprehension: bool = False,
+    model_stage: str | None = None,
+):
+    engine_args = SimpleNamespace()
+    if model_stage is not None:
+        engine_args.model_stage = model_stage
+    return SimpleNamespace(
+        stage_type=stage_type,
+        is_comprehension=is_comprehension,
+        engine_args=engine_args,
+    )
+
+
+def _stage_type(stage: Any) -> str | None:
+    return getattr(stage, "stage_type", None)
+
+
+def _build_output_modalities(stage_configs: list[Any]) -> list[str]:
+    modalities: list[str] = []
+    for stage in stage_configs:
+        final_output = getattr(stage, "final_output", False)
+        final_output_type = getattr(stage, "final_output_type", None)
+        if final_output and isinstance(final_output_type, str):
+            modalities.append(final_output_type)
+    if modalities:
+        return modalities
+    if any(_stage_type(stage) == "diffusion" for stage in stage_configs):
+        return ["image"]
+    return ["text", "audio"]
+
+
 def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: SamplingCase, server_case: ServerCase):
     async def _mock_generate(*args, **kwargs):
         received_sampling_params_list: Sequence[OmniSamplingParams] | None = (
@@ -367,18 +403,28 @@ def mock_async_omni(server_case: ServerCase, sampling_case: SamplingCase):
             new=_mock_preprocess_chat,
         ),
     ):
-        mock_instance = AsyncMock()
+        mock_instance = AsyncMock(spec=RealAsyncOmni)
         mock_instance.generate = _build_mock_outputs(server_case.outputs, sampling_case, server_case)
 
         mock_instance.stage_list = server_case.stage_list
         mock_instance.stage_configs = server_case.stage_configs
+        mock_instance.output_modalities = _build_output_modalities(server_case.stage_configs)
         mock_instance.default_sampling_params_list = [
-            SamplingParams() if stage.get("stage_type") != "diffusion" else MagicMock()
+            SamplingParams() if _stage_type(stage) != "diffusion" else MagicMock()
             for stage in server_case.stage_configs
         ]
         mock_instance.errored = False
         mock_instance.dead_error = RuntimeError("Mock engine error")
-        mock_instance.model_config = MagicMock(max_model_len=4096, io_processor_plugin=None)
+        mock_instance.model_config = MagicMock(
+            max_model_len=4096,
+            io_processor_plugin=None,
+            allowed_local_media_path=None,
+            allowed_media_domains=None,
+        )
+        # Mimic Qwen3-TTS talker speaker config so CustomVoice validation passes.
+        mock_instance.model_config.hf_config = MagicMock()
+        mock_instance.model_config.hf_config.talker_config = MagicMock()
+        mock_instance.model_config.hf_config.talker_config.speaker_id = {"Vivian": 0}
         mock_instance.io_processor = MagicMock()
         mock_instance.input_processor = MagicMock()
         mock_instance.shutdown = MagicMock()
@@ -444,7 +490,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="Tongyi-MAI/Z-Image-Turbo",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_images_endpoint()],
             ),
             "Tongyi-MAI/Z-Image-Turbo",
@@ -455,7 +501,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="ByteDance-Seed/BAGEL-7B-MoT",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_chat_endpoint()],
             ),
             "ByteDance-Seed/BAGEL-7B-MoT",
@@ -466,7 +512,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="Qwen/Qwen-Image-Edit",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_images_endpoint()],
             ),
             "Qwen/Qwen-Image-Edit",
@@ -477,7 +523,7 @@ def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni
             ServerCase(
                 served_model="ByteDance-Seed/BAGEL-7B-MoT",
                 stage_list=["diffusion"],
-                stage_configs=[{"stage_type": "diffusion"}],
+                stage_configs=[_make_stage_config("diffusion")],
                 outputs=[_build_diffusion_image_output_for_chat_endpoint()],
             ),
             "ByteDance-Seed/BAGEL-7B-MoT",
@@ -542,9 +588,9 @@ async def test_image_generation_node(api_server: str, model: str, image_input: b
                     MagicMock(is_comprehension=False, model_stage="llm"),
                 ],
                 stage_configs=[
-                    {"stage_type": "llm"},
-                    {"stage_type": "llm"},
-                    {"stage_type": "llm"},
+                    _make_stage_config("llm", is_comprehension=True, model_stage="thinker"),
+                    _make_stage_config("llm", is_comprehension=False, model_stage="talker"),
+                    _make_stage_config("llm", is_comprehension=False, model_stage="code2wav"),
                 ],
                 outputs=[_build_audio_chat_output(), _build_text_output("Understanding response")],
             ),
@@ -599,7 +645,7 @@ async def test_understanding_node(api_server: str, sampling_case: SamplingCase):
             ServerCase(
                 served_model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
                 stage_list=["llm"],
-                stage_configs=[{"stage_type": "llm"}],
+                stage_configs=[_make_stage_config("llm", model_stage="qwen3_tts")],
                 outputs=[_build_audio_speech_output()],
             ),
             VLLMOmniTTS,
@@ -617,7 +663,7 @@ async def test_understanding_node(api_server: str, sampling_case: SamplingCase):
             ServerCase(
                 served_model="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
                 stage_list=["llm"],
-                stage_configs=[{"stage_type": "llm"}],
+                stage_configs=[_make_stage_config("llm", model_stage="qwen3_tts")],
                 outputs=[_build_audio_speech_output()],
             ),
             VLLMOmniVoiceClone,

tests/conftest.py

@@ -1733,7 +1733,9 @@ class OmniRunner:
         Returns:
             List of SamplingParams with default decoding for each stage
         """
-        return [st.default_sampling_params for st in self.omni.stage_list]
+        if not hasattr(self.omni, "default_sampling_params_list"):
+            raise AttributeError("Omni.default_sampling_params_list is not available")
+        return list(self.omni.default_sampling_params_list)
 
     def get_omni_inputs(
         self,
@@ -1991,9 +1993,9 @@ class OmniRunnerHandler:
             audio_content = None
             for stage_output in outputs:
                 if getattr(stage_output, "final_output_type", None) == "text":
-                    text_content = stage_output.request_output[0].outputs[0].text
+                    text_content = stage_output.request_output.outputs[0].text
                 if getattr(stage_output, "final_output_type", None) == "audio":
-                    audio_content = stage_output.request_output[0].outputs[0].multimodal_output["audio"]
+                    audio_content = stage_output.request_output.outputs[0].multimodal_output["audio"]
 
             result.audio_content = audio_content
             result.text_content = text_content

tests/diffusion/cache/test_cache_backends.py

@@ -16,6 +16,7 @@ from unittest.mock import Mock, patch
 import pytest
 
 from vllm_omni.diffusion.cache.cache_dit_backend import (
+    CUSTOM_DIT_ENABLERS,
     CacheDiTBackend,
 )
 from vllm_omni.diffusion.cache.selector import get_cache_backend
@@ -122,6 +123,52 @@ class TestCacheDiTBackend:
         assert "cache_config" in call_args[1]
         assert backend._last_num_inference_steps == 100
 
+    def test_hunyuan_custom_enabler_registered(self):
+        """Test HunyuanImage3 custom cache-dit enabler is registered."""
+        assert "HunyuanImage3Pipeline" in CUSTOM_DIT_ENABLERS
+
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.BlockAdapter")
+    @patch("vllm_omni.diffusion.cache.cache_dit_backend.cache_dit")
+    def test_enable_hunyuan_pipeline_uses_model_transformer(self, mock_cache_dit, mock_block_adapter):
+        """Test HunyuanImage3 custom enabler uses pipeline.model for cache enable/refresh."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline"
+        mock_pipeline.model = Mock()
+        mock_pipeline.model.layers = Mock()
+
+        mock_cache_dit.enable_cache = Mock()
+        mock_cache_dit.refresh_context = Mock()
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+        backend.enable(mock_pipeline)
+
+        assert backend.enabled is True
+        assert backend._refresh_func is not None
+        mock_block_adapter.assert_called_once()
+        adapter_kwargs = mock_block_adapter.call_args.kwargs
+        assert adapter_kwargs["transformer"] is mock_pipeline.model
+        assert adapter_kwargs["blocks"] is mock_pipeline.model.layers
+        assert adapter_kwargs["forward_pattern"] == adapter_kwargs["forward_pattern"].__class__.Pattern_4
+        assert len(adapter_kwargs["params_modifiers"]) == 1
+        mock_cache_dit.enable_cache.assert_called_once()
+
+        backend.refresh(mock_pipeline, num_inference_steps=12)
+        mock_cache_dit.refresh_context.assert_called_once()
+        call_args = mock_cache_dit.refresh_context.call_args
+        assert call_args[0][0] is mock_pipeline.model
+        assert call_args[1]["num_inference_steps"] == 12
+
+    def test_enable_hunyuan_pipeline_requires_model_layers(self):
+        """Test HunyuanImage3 enabler fails with a formatted pipeline class name."""
+        mock_pipeline = Mock()
+        mock_pipeline.__class__.__name__ = "HunyuanImage3Pipeline"
+        mock_pipeline.model = Mock(spec=[])
+
+        backend = CacheDiTBackend({"Fn_compute_blocks": 2})
+
+        with pytest.raises(ValueError, match="HunyuanImage3Pipeline"):
+            backend.enable(mock_pipeline)
+
 
 class TestTeaCacheBackend:
     """Test TeaCacheBackend implementation."""

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -70,6 +70,7 @@ def mock_dist():
         patch.object(dist, "get_world_size", return_value=2),
         patch.object(dist, "get_rank", return_value=0),
         patch.object(dist, "is_initialized", return_value=True),
+        patch.object(dist, "all_reduce", return_value=None),
         patch.object(dist, "gather", return_value=None),
         patch.object(dist, "broadcast", return_value=None),
     ):

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -54,6 +54,8 @@ def build_adapter(monkeypatch, mocker: MockerFixture):
             self._pending_save_reqs = deque()
             self._finished_save_reqs = set()
             self.stop_event = threading.Event()
+            self._recv_cond = threading.Condition()
+            self._save_cond = threading.Condition()
 
         monkeypatch.setattr(OmniTransferAdapterBase, "__init__", _fake_base_init)
         monkeypatch.setattr(

tests/e2e/offline_inference/stage_configs/npu/qwen2_5_omni_ci.yaml

@@ -1,4 +1,4 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni for multi-stage omni runtime.
 
 # This config is optimized for CI e2e tests.
 stage_args:

tests/e2e/offline_inference/test_bagel_img2img.py

@@ -98,9 +98,9 @@ def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
         if images := getattr(req_output, "images", None):
             return images[0]
         if hasattr(req_output, "request_output") and req_output.request_output:
-            for stage_out in req_output.request_output:
-                if hasattr(stage_out, "images") and stage_out.images:
-                    return stage_out.images[0]
+            stage_out = req_output.request_output
+            if hasattr(stage_out, "images") and stage_out.images:
+                return stage_out.images[0]
     return None
 
 

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -14,6 +14,9 @@ Equivalent to running:
 """
 
 import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
 import signal
 import socket
 import subprocess
@@ -96,9 +99,9 @@ def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
         if images := getattr(req_output, "images", None):
             return images[0]
         if hasattr(req_output, "request_output") and req_output.request_output:
-            for stage_out in req_output.request_output:
-                if hasattr(stage_out, "images") and stage_out.images:
-                    return stage_out.images[0]
+            stage_out = req_output.request_output
+            if hasattr(stage_out, "images") and stage_out.images:
+                return stage_out.images[0]
     return None
 
 

tests/e2e/offline_inference/test_cache_dit.py

@@ -72,13 +72,13 @@ def test_cache_dit(model_name: str):
                 num_outputs_per_prompt=1,  # Single output for speed
             ),
         )
-        # Extract images from request_output[0]['images']
+        # Extract images from request_output['images']
         first_output = outputs[0]
         assert first_output.final_output_type == "image"
         if not hasattr(first_output, "request_output") or not first_output.request_output:
             raise ValueError("No request_output found in OmniRequestOutput")
 
-        req_out = first_output.request_output[0]
+        req_out = first_output.request_output
         if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
             raise ValueError("Invalid request_output structure or missing 'images' key")
 

tests/e2e/offline_inference/test_diffusion_lora.py

@@ -37,7 +37,7 @@ def test_diffusion_model(model_name: str, tmp_path: Path):
         if not hasattr(first_output, "request_output") or not first_output.request_output:
             raise ValueError("No request_output found in OmniRequestOutput")
 
-        req_out = first_output.request_output[0]
+        req_out = first_output.request_output
         if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
             raise ValueError("Invalid request_output structure or missing 'images' key")
         return req_out.images

tests/e2e/offline_inference/test_expert_parallel.py

@@ -129,7 +129,7 @@ def _run_inference(
         elapsed_ms = (time.time() - start) * 1000
 
         return InferenceResult(
-            images=outputs[0].request_output[0].images,
+            images=outputs[0].images,
             elapsed_ms=elapsed_ms,
         )
     finally:

tests/e2e/offline_inference/test_sequence_parallel.py

@@ -130,7 +130,7 @@ def _run_inference(
         elapsed_ms = (time.time() - start) * 1000
 
         return InferenceResult(
-            images=outputs[0].request_output[0].images,
+            images=outputs[0].request_output.images,
             elapsed_ms=elapsed_ms,
         )
     finally:

tests/e2e/offline_inference/test_stable_audio_model.py

@@ -57,7 +57,7 @@ def test_stable_audio_model(model_name: str):
     assert first_output.final_output_type == "image"
     assert hasattr(first_output, "request_output") and first_output.request_output
 
-    req_out = first_output.request_output[0]
+    req_out = first_output.request_output
     assert isinstance(req_out, OmniRequestOutput)
     assert req_out.final_output_type == "audio"
     assert hasattr(req_out, "multimodal_output") and req_out.multimodal_output

tests/e2e/offline_inference/test_t2i_model.py

@@ -28,7 +28,7 @@ if current_omni_platform.is_npu():
     models = ["Tongyi-MAI/Z-Image-Turbo", "Qwen/Qwen-Image"]
 elif current_omni_platform.is_rocm():
     # TODO: When ROCm support is ready, remove this branch.
-    # vLLM V0.11.0 has issues running riverclouds/qwen_image_random
+    # Current upstream vLLM has issues running riverclouds/qwen_image_random
     # on ROCm
     models = ["Tongyi-MAI/Z-Image-Turbo"]
 
@@ -62,13 +62,13 @@ def test_diffusion_model(model_name: str, run_level):
                 num_outputs_per_prompt=2,
             ),
         )
-        # Extract images from request_output[0]['images']
+        # Extract images from request_output['images']
         first_output = outputs[0]
         assert first_output.final_output_type == "image"
         if not hasattr(first_output, "request_output") or not first_output.request_output:
             raise ValueError("No request_output found in OmniRequestOutput")
 
-        req_out = first_output.request_output[0]
+        req_out = first_output.request_output
         if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
             raise ValueError("Invalid request_output structure or missing 'images' key")
 

tests/e2e/offline_inference/test_t2v_model.py

@@ -51,7 +51,7 @@ def test_video_diffusion_model(model_name: str):
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output[0]
+    req_out = first_output.request_output
     if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
         raise ValueError("Invalid request_output structure or missing 'images' key")
 

tests/e2e/offline_inference/test_teacache.py

@@ -68,13 +68,13 @@ def test_teacache(model_name: str):
                 num_outputs_per_prompt=1,  # Single output for speed
             ),
         )
-        # Extract images from request_output[0]['images']
+        # Extract images from request_output['images']
         first_output = outputs[0]
         assert first_output.final_output_type == "image"
         if not hasattr(first_output, "request_output") or not first_output.request_output:
             raise ValueError("No request_output found in OmniRequestOutput")
 
-        req_out = first_output.request_output[0]
+        req_out = first_output.request_output
         if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
             raise ValueError("Invalid request_output structure or missing 'images' key")
 

tests/e2e/offline_inference/test_vae_decode_parallelism.py

@@ -72,6 +72,7 @@ def is_nextstep_model(model_name: str) -> bool:
 
 
 def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile, vae_patch_parallel_size=1):
+    m = None
     try:
         parallel_config = DiffusionParallelConfig(
             tensor_parallel_size=tp,
@@ -105,7 +106,7 @@ def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile,
         )
         end = time.perf_counter()
         first_output = outputs[0]
-        req_out = first_output.request_output[0]
+        req_out = first_output.request_output
         frames = req_out.images[0]
         if isinstance(frames, torch.Tensor):
             frames = frames.detach().cpu().numpy()
@@ -115,7 +116,8 @@ def model_run(model_configs, tp, out_height, out_width, out_frames, using_tile,
         cost = (end - start) * 1000
         return frames, cost
     finally:
-        m.close()
+        if m is not None:
+            m.close()
         cleanup_dist_env_and_memory()
 
 

tests/e2e/offline_inference/test_zimage_parallelism.py

@@ -68,7 +68,7 @@ def _extract_single_image(outputs) -> Image.Image:
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output[0]
+    req_out = first_output.request_output
     if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
         raise ValueError("Invalid request_output structure or missing 'images' key")
 

tests/e2e/online_serving/test_qwen3_tts_base.py

@@ -17,6 +17,7 @@ import tempfile
 from pathlib import Path
 
 import httpx
+import numpy as np
 import pytest
 
 from tests.conftest import (
@@ -99,6 +100,34 @@ def assert_not_silence(pcm_bytes: bytes):
 
 
 MIN_AUDIO_BYTES = 10000
+MIN_HNR_DB = 1.2  # Clean voice clone > 1.2 dB; distorted < 1.0 dB
+
+
+def compute_hnr_db(pcm_samples: np.ndarray, sr: int = 24000) -> float:
+    """Compute mean Harmonic-to-Noise Ratio (dB) for speech quality.
+
+    Clean cloned speech has HNR > 1.2 dB; distorted speech (e.g. lost
+    ref_code decoder context) drops below 1.0 dB.
+    """
+    frame_len = int(0.03 * sr)  # 30ms frames
+    hop = frame_len // 2
+    hnr_values: list[float] = []
+
+    for start in range(0, len(pcm_samples) - frame_len, hop):
+        frame = pcm_samples[start : start + frame_len]
+        if np.max(np.abs(frame)) < 0.01:
+            continue
+        ac = np.correlate(frame, frame, mode="full")[len(frame) - 1 :]
+        ac = ac / (ac[0] + 1e-10)
+        min_lag = int(sr / 400)
+        max_lag = min(int(sr / 80), len(ac))
+        if min_lag >= max_lag:
+            continue
+        peak = float(np.max(ac[min_lag:max_lag]))
+        if 0 < peak < 1:
+            hnr_values.append(10 * np.log10(peak / (1 - peak + 1e-10)))
+
+    return float(np.mean(hnr_values)) if hnr_values else 0.0
 
 
 class TestQwen3TTSBaseVoiceClone:
@@ -165,3 +194,30 @@ class TestQwen3TTSBaseVoiceClone:
         assert response.headers.get("content-type") == "audio/wav"
         assert verify_wav_audio(response.content), "Response is not valid WAV"
         assert len(response.content) > MIN_AUDIO_BYTES
+
+    @pytest.mark.core_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "L4"}, num_cards=4)
+    def test_base_voice_clone_no_distortion(self, omni_server) -> None:
+        """Audio must not be distorted (regression for ref_code context loss).
+
+        When the decoder loses ref_code speaker context on later streaming
+        chunks, HNR drops below 1.0 dB.  Clean voice clone should be > 1.2 dB.
+        See https://github.com/vllm-project/vllm-omni/issues/1944
+        """
+        response = make_base_speech_request(
+            host=omni_server.host,
+            port=omni_server.port,
+            response_format="pcm",
+        )
+
+        assert response.status_code == 200, f"Request failed: {response.text}"
+        assert len(response.content) > MIN_AUDIO_BYTES
+
+        pcm_samples = np.frombuffer(response.content, dtype=np.int16).astype(np.float32) / 32768.0
+        hnr = compute_hnr_db(pcm_samples)
+        print(f"Voice clone HNR: {hnr:.2f} dB (threshold: {MIN_HNR_DB} dB)")
+        assert hnr >= MIN_HNR_DB, (
+            f"Audio distortion detected: HNR={hnr:.2f} dB < {MIN_HNR_DB} dB. "
+            "Voice clone decoder may be losing ref_code speaker context on later chunks."
+        )

tests/e2e/stage_configs/qwen2_5_omni_ci.yaml

@@ -1,4 +1,4 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni for multi-stage omni runtime.
 
 # The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
 # This config is optimized for CI e2e tests.

tests/e2e/stage_configs/rocm/qwen2_5_omni_ci.yaml

@@ -1,4 +1,4 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni for multi-stage omni runtime.
 
 # The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
 # This config is optimized for CI e2e tests.

tests/e2e/stage_configs/xpu/qwen2_5_omni_ci.yaml

@@ -1,9 +1,9 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni for multi-stage omni runtime.
 
 # The following config is verified with 2 * Intel Arc Pro B60 XPU.
 stage_args:
   - stage_id: 0
-    stage_type: llm # Use llm stage type to launch OmniLLM
+    stage_type: llm # Use llm stage type for AR stages
     runtime:
       process: true # Run this stage in a separate process
       devices: "0" # Visible devices for this stage
@@ -35,7 +35,7 @@ stage_args:
       detokenize: True
       repetition_penalty: 1.1
   - stage_id: 1
-    stage_type: llm # Use llm stage type to launch OmniLLM
+    stage_type: llm # Use llm stage type for AR stages
     runtime:
       process: true
       devices: "1"
@@ -67,7 +67,7 @@ stage_args:
       stop_token_ids: [8294]
 
   - stage_id: 2
-    stage_type: llm # Use llm stage type to launch OmniLLM
+    stage_type: llm # Use llm stage type for AR stages
     runtime:
       process: true
       devices: "2"

tests/e2e/stage_configs/xpu/qwen3_omni_ci.yaml

@@ -6,7 +6,7 @@
 # The following config is verified with 8 * Intel Arc Pro B60 XPU.
 stage_args:
 - stage_id: 0
-  stage_type: llm # Use llm stage type to launch OmniLLM
+  stage_type: llm # Use llm stage type for AR stages
   runtime:
     devices: "0,1,2,3"
     max_batch_size: 1
@@ -41,7 +41,7 @@ stage_args:
     repetition_penalty: 1.05
 
 - stage_id: 1
-  stage_type: llm # Use llm stage type to launch OmniLLM
+  stage_type: llm # Use llm stage type for AR stages
   runtime:
     devices: "4"
     max_batch_size: 1
@@ -75,7 +75,7 @@ stage_args:
     stop_token_ids: [2150]
 
 - stage_id: 2
-  stage_type: llm # Use llm stage type to launch OmniLLM
+  stage_type: llm # Use llm stage type for AR stages
   runtime:
     devices: "5"
     max_batch_size: 1

tests/engine/test_async_omni_engine_input.py

@@ -0,0 +1,63 @@
+from unittest.mock import Mock
+
+import pytest
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine import EngineCoreRequest
+
+from vllm_omni.engine import OmniEngineCoreRequest
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_engine_core_request() -> EngineCoreRequest:
+    return EngineCoreRequest(
+        request_id="req-1",
+        prompt_token_ids=[1, 1, 1],
+        mm_features=None,
+        sampling_params=SamplingParams(max_tokens=8),
+        pooling_params=None,
+        arrival_time=0.0,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+    )
+
+
+def test_build_add_request_message_preserves_additional_information():
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [{"stage_type": "llm"}]
+    engine.supported_tasks = ("speech",)
+
+    input_processor = Mock()
+    input_processor.process_inputs.return_value = _make_engine_core_request()
+    engine.input_processor = input_processor
+
+    output_processor = Mock()
+    engine.output_processors = [output_processor]
+
+    prompt = {
+        "prompt_token_ids": [1, 1, 1],
+        "additional_information": {
+            "text": ["hello world"],
+            "speaker": ["vivian"],
+        },
+    }
+
+    msg = engine._build_add_request_message(
+        request_id="req-1",
+        prompt=prompt,
+        sampling_params_list=[params],
+        final_stage_id=0,
+        arrival_time=0.0,
+    )
+
+    request = msg["prompt"]
+    assert isinstance(request, OmniEngineCoreRequest)
+    assert request.external_req_id == "req-1"
+    assert request.additional_information is not None
+    assert request.additional_information.entries["text"].list_data == ["hello world"]
+    assert request.additional_information.entries["speaker"].list_data == ["vivian"]
+    output_processor.add_request.assert_called_once()

tests/entrypoints/openai_api/test_image_server.py

@@ -10,6 +10,7 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py.
 import base64
 import io
 from argparse import Namespace
+from types import SimpleNamespace
 
 import pytest
 from fastapi.testclient import TestClient
@@ -115,7 +116,10 @@ class FakeAsyncOmni:
     """Fake AsyncOmni that yields a single diffusion output."""
 
     def __init__(self):
-        self.stage_list = ["llm", "diffusion"]
+        self.stage_configs = [
+            SimpleNamespace(stage_type="llm"),
+            SimpleNamespace(stage_type="diffusion"),
+        ]
         self.default_sampling_params_list = [SamplingParams(temperature=0.1), OmniDiffusionSamplingParams()]
         self.captured_sampling_params_list = None
         self.captured_prompt = None
@@ -129,21 +133,25 @@ class FakeAsyncOmni:
 
 @pytest.fixture
 def mock_async_diffusion(mocker: MockerFixture):
-    """Mock AsyncOmniDiffusion instance that returns fake images"""
-    mock = mocker.Mock()
-    mock.is_running = True  # For health endpoint
-    mock.check_health = mocker.AsyncMock()  # For LLM mode health check
+    """Mock diffusion engine that matches the current async-generator API."""
 
-    async def generate(**kwargs):
-        # Return n PIL images wrapped in result object
-        n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
-        mock.captured_sampling_params_list = kwargs["sampling_params_list"]
-        mock.captured_prompt = kwargs["prompt"]
-        images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
-        return MockGenerationResult(images)
+    class MockAsyncDiffusion:
+        def __init__(self) -> None:
+            self.is_running = True
+            self.check_health = mocker.AsyncMock()
+            self.captured_sampling_params_list = None
+            self.captured_prompt = None
+            self.generate_calls = 0
 
-    mock.generate = mocker.AsyncMock(side_effect=generate)
-    return mock
+        async def generate(self, **kwargs):
+            self.generate_calls += 1
+            n = kwargs["sampling_params_list"][0].num_outputs_per_prompt
+            self.captured_sampling_params_list = kwargs["sampling_params_list"]
+            self.captured_prompt = kwargs["prompt"]
+            images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
+            yield MockGenerationResult(images)
+
+    return MockAsyncDiffusion()
 
 
 @pytest.fixture
@@ -159,7 +167,7 @@ def test_client(mock_async_diffusion):
     # Set up app state with diffusion engine
     app.state.engine_client = mock_async_diffusion
     app.state.diffusion_engine = mock_async_diffusion  # Also set for health endpoint
-    app.state.stage_configs = [{"stage_type": "diffusion"}]
+    app.state.stage_configs = [SimpleNamespace(stage_type="diffusion")]
     app.state.diffusion_model_name = "Qwen/Qwen-Image"  # For models endpoint
     app.state.args = Namespace(
         default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}',
@@ -180,7 +188,32 @@ def async_omni_test_client():
     app.include_router(router)
 
     app.state.engine_client = FakeAsyncOmni()
-    app.state.stage_configs = [{"stage_type": "llm"}, {"stage_type": "diffusion"}]
+    app.state.stage_configs = [
+        SimpleNamespace(stage_type="llm"),
+        SimpleNamespace(stage_type="diffusion"),
+    ]
+    app.state.args = Namespace(
+        default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
+        max_generated_image_size=4096,  # 64*64
+    )
+    return TestClient(app)
+
+
+@pytest.fixture
+def async_omni_stage_configs_only_client():
+    """Create test client with refactored AsyncOmni compatibility surface only."""
+    from fastapi import FastAPI
+
+    from vllm_omni.entrypoints.openai.api_server import router
+
+    app = FastAPI()
+    app.include_router(router)
+
+    engine = FakeAsyncOmni()
+    assert not hasattr(engine, "stage_list")
+    app.state.engine_client = engine
+    # Intentionally do not populate app.state.stage_configs. Refactored
+    # AsyncOmni exposes stage_configs on the engine instance.
     app.state.args = Namespace(
         default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}',
         max_generated_image_size=4096,  # 64*64
@@ -291,6 +324,45 @@ def test_generate_images_async_omni_sampling_params(async_omni_test_client):
     assert captured[1].seed == 7
 
 
+def test_generate_images_async_omni_stage_configs_only(async_omni_stage_configs_only_client):
+    """Regression: image generation accepts refactored AsyncOmni without stage_list."""
+    response = async_omni_stage_configs_only_client.post(
+        "/v1/images/generations",
+        json={
+            "prompt": "a castle",
+            "n": 1,
+            "size": "256x256",
+            "seed": 11,
+        },
+    )
+    assert response.status_code == 200
+    data = response.json()
+    assert len(data["data"]) == 1
+    engine = async_omni_stage_configs_only_client.app.state.engine_client
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+    assert captured[1].seed == 11
+
+
+def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only_client):
+    """Regression: image edits accepts refactored AsyncOmni without stage_list."""
+    img_bytes = make_test_image_bytes((16, 16))
+    response = async_omni_stage_configs_only_client.post(
+        "/v1/images/edits",
+        files=[("image", img_bytes)],
+        data={
+            "prompt": "edit me",
+            "size": "auto",
+        },
+    )
+    assert response.status_code == 200
+    engine = async_omni_stage_configs_only_client.app.state.engine_client
+    captured = engine.captured_sampling_params_list
+    assert captured is not None
+    assert len(captured) == 2
+
+
 def test_generate_multiple_images(test_client):
     """Test generating multiple images"""
     response = test_client.post(
@@ -538,13 +610,12 @@ def test_parameters_passed_through(test_client, mock_async_diffusion):
     )
     assert response.status_code == 200
 
-    # Ensure generate() was called exactly once
-    mock_async_diffusion.generate.assert_awaited_once()
-    call_kwargs = mock_async_diffusion.generate.call_args[1]["sampling_params_list"][0]
-    assert call_kwargs.num_inference_steps == 100
-    assert call_kwargs.guidance_scale == 7.5
-    assert call_kwargs.true_cfg_scale == 3.0
-    assert call_kwargs.seed == 42
+    assert mock_async_diffusion.generate_calls == 1
+    captured = mock_async_diffusion.captured_sampling_params_list[0]
+    assert captured.num_inference_steps == 100
+    assert captured.guidance_scale == 7.5
+    assert captured.true_cfg_scale == 3.0
+    assert captured.seed == 42
 
 
 def test_model_field_omitted_works(test_client):

tests/entrypoints/openai_api/test_serving_chat_sampling_params.py

@@ -63,9 +63,9 @@ def mock_engine_client(
     default_other_params,
     mocker: MockerFixture,
 ):
-    """Create mock engine client with stage_list and default_sampling_params_list."""
+    """Create mock engine client with stage_configs and default_sampling_params_list."""
     engine_client = mocker.MagicMock()
-    engine_client.stage_list = [mock_comprehension_stage, mock_other_stage]
+    engine_client.stage_configs = [mock_comprehension_stage, mock_other_stage]
     engine_client.default_sampling_params_list = [
         default_comprehension_params,
         default_other_params,
@@ -311,7 +311,7 @@ def test_get_comprehension_stage_index_finds_second_stage(mocker: MockerFixture)
     comprehension.is_comprehension = True
 
     instance.engine_client = mocker.MagicMock()
-    instance.engine_client.stage_list = [other, comprehension]
+    instance.engine_client.stage_configs = [other, comprehension]
 
     assert instance._get_comprehension_stage_index() == 1
 
@@ -328,7 +328,7 @@ def test_get_comprehension_stage_index_raises_when_not_found(mocker: MockerFixtu
     stage2.is_comprehension = False
 
     instance.engine_client = mocker.MagicMock()
-    instance.engine_client.stage_list = [stage1, stage2]
+    instance.engine_client.stage_configs = [stage1, stage2]
 
     with pytest.raises(ValueError, match="No comprehension stage"):
         instance._get_comprehension_stage_index()

tests/entrypoints/openai_api/test_serving_speech.py

@@ -487,7 +487,7 @@ class TestTTSMethods:
     def speech_server(self, mocker: MockerFixture):
         mock_engine_client = mocker.MagicMock()
         mock_engine_client.errored = False
-        mock_engine_client.stage_list = None
+        mock_engine_client.stage_configs = []
         mock_engine_client.tts_max_instructions_length = None
         mock_models = mocker.MagicMock()
         mock_models.is_base_model.return_value = True
@@ -499,7 +499,7 @@ class TestTTSMethods:
 
     def test_is_tts_detection_no_stage(self, speech_server):
         """Test TTS model detection when no TTS stage exists."""
-        # Fixture creates server with stage_list = None -> _is_tts should be False
+        # Fixture creates server with stage_configs = [] -> _is_tts should be False
         assert speech_server._is_tts is False
         assert speech_server._tts_stage is None
 
@@ -511,9 +511,9 @@ class TestTTSMethods:
 
         # Create a TTS stage
         mock_stage = mocker.MagicMock()
-        mock_stage.model_stage = "qwen3_tts"
+        mock_stage.engine_args.model_stage = "qwen3_tts"
         mock_stage.tts_args = {}
-        mock_engine_client.stage_list = [mock_stage]
+        mock_engine_client.stage_configs = [mock_stage]
 
         mock_models = mocker.MagicMock()
         mock_models.is_base_model.return_value = True
@@ -614,7 +614,7 @@ class TestTTSMethods:
         """Test _load_supported_speakers."""
         mock_engine_client = mocker.MagicMock()
         mock_engine_client.errored = False
-        mock_engine_client.stage_list = None
+        mock_engine_client.stage_configs = []
 
         # Mock talker_config with mixed-case speaker names
         mock_talker_config = mocker.MagicMock()
@@ -757,7 +757,7 @@ class TestTTSMethods:
         """Test CLI override (stored in engine_client) takes highest priority."""
         mock_engine_client = mocker.MagicMock()
         mock_engine_client.errored = False
-        mock_engine_client.stage_list = None
+        mock_engine_client.stage_configs = []
         # CLI override is stored in engine_client
         mock_engine_client.tts_max_instructions_length = 1000
         mock_models = mocker.MagicMock()
@@ -781,9 +781,9 @@ class TestTTSMethods:
 
         # Mock stage with tts_args
         mock_stage = mocker.MagicMock()
-        mock_stage.model_stage = "qwen3_tts"
+        mock_stage.engine_args.model_stage = "qwen3_tts"
         mock_stage.tts_args = {"max_instructions_length": 750}
-        mock_engine_client.stage_list = [mock_stage]
+        mock_engine_client.stage_configs = [mock_stage]
 
         server = OmniOpenAIServingSpeech(
             engine_client=mock_engine_client,
@@ -804,9 +804,9 @@ class TestTTSMethods:
 
         # Mock stage with tts_args
         mock_stage = mocker.MagicMock()
-        mock_stage.model_stage = "qwen3_tts"
+        mock_stage.engine_args.model_stage = "qwen3_tts"
         mock_stage.tts_args = {"max_instructions_length": 750}
-        mock_engine_client.stage_list = [mock_stage]
+        mock_engine_client.stage_configs = [mock_stage]
 
         server = OmniOpenAIServingSpeech(
             engine_client=mock_engine_client,
@@ -820,7 +820,7 @@ class TestTTSMethods:
         """Test instructions length validation uses cached _max_instructions_length."""
         mock_engine_client = mocker.MagicMock()
         mock_engine_client.errored = False
-        mock_engine_client.stage_list = None
+        mock_engine_client.stage_configs = []
         # CLI override with max length of 10 characters
         mock_engine_client.tts_max_instructions_length = 10
         mock_models = mocker.MagicMock()
@@ -1034,15 +1034,12 @@ class TestAsyncOmniSupportedTasks:
     @pytest.mark.asyncio
     async def test_tts_only_no_generate_task(self):
         """TTS-only models (audio output, no text) should not include 'generate'."""
-        from unittest.mock import MagicMock
+        from types import SimpleNamespace
 
         from vllm_omni.entrypoints.async_omni import AsyncOmni
 
         omni = AsyncOmni.__new__(AsyncOmni)
-        omni.output_modalities = [None, "audio"]
-        stage = MagicMock()
-        stage.is_comprehension = False
-        omni.stage_list = [stage]
+        omni.engine = SimpleNamespace(supported_tasks=("speech",))
         tasks = await omni.get_supported_tasks()
         assert "generate" not in tasks
         assert "speech" in tasks
@@ -1050,15 +1047,12 @@ class TestAsyncOmniSupportedTasks:
     @pytest.mark.asyncio
     async def test_omni_model_includes_generate(self):
         """Models with text output (e.g. Qwen3-Omni) should include 'generate'."""
-        from unittest.mock import MagicMock
+        from types import SimpleNamespace
 
         from vllm_omni.entrypoints.async_omni import AsyncOmni
 
         omni = AsyncOmni.__new__(AsyncOmni)
-        omni.output_modalities = ["text", None, "audio"]
-        stage = MagicMock()
-        stage.is_comprehension = True
-        omni.stage_list = [stage]
+        omni.engine = SimpleNamespace(supported_tasks=("generate", "speech"))
         tasks = await omni.get_supported_tasks()
         assert "generate" in tasks
 

tests/entrypoints/openai_api/test_video_server.py

@@ -44,7 +44,7 @@ class MockVideoResult:
 
 class FakeAsyncOmni:
     def __init__(self):
-        self.stage_list = ["diffusion"]
+        self.stage_configs = [SimpleNamespace(stage_type="diffusion")]
         self.captured_prompt = None
         self.captured_sampling_params_list = None
 

tests/entrypoints/test_async_omni.py

@@ -0,0 +1,65 @@
+from types import SimpleNamespace
+
+import pytest
+from vllm.entrypoints.openai.models.protocol import BaseModelPath
+from vllm.entrypoints.openai.models.serving import OpenAIServingModels
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.asyncio
+async def test_get_supported_tasks_returns_engine_supported_tasks():
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(supported_tasks=("generate", "speech"))
+
+    supported_tasks = await omni.get_supported_tasks()
+
+    assert supported_tasks == ("generate", "speech")
+
+
+def test_model_config_and_vllm_config_forward_from_comprehension_stage():
+    model_config = SimpleNamespace(model="Qwen/Qwen3-TTS")
+    vllm_config = SimpleNamespace(model_config=model_config)
+    renderer = SimpleNamespace(name="renderer")
+    input_processor = SimpleNamespace(renderer=renderer)
+    io_processor = SimpleNamespace(name="io-processor")
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        stage_clients=[SimpleNamespace(is_comprehension=False), SimpleNamespace(is_comprehension=True)],
+        stage_vllm_configs=[None, vllm_config],
+    )
+    omni.input_processor = input_processor
+    omni.io_processor = io_processor
+
+    assert omni.vllm_config is vllm_config
+    assert omni.model_config is model_config
+    assert omni.renderer is renderer
+    assert omni.input_processor is input_processor
+    assert omni.io_processor is io_processor
+
+
+def test_openai_serving_models_can_consume_async_omni_compat_attrs():
+    model_config = SimpleNamespace(model="Qwen/Qwen3-TTS", max_model_len=32768)
+    vllm_config = SimpleNamespace(model_config=model_config)
+    renderer = SimpleNamespace(name="renderer")
+    input_processor = SimpleNamespace(renderer=renderer)
+    io_processor = SimpleNamespace(name="io-processor")
+    omni = object.__new__(AsyncOmni)
+    omni.engine = SimpleNamespace(
+        stage_clients=[SimpleNamespace(is_comprehension=True)],
+        stage_vllm_configs=[vllm_config],
+    )
+    omni.input_processor = input_processor
+    omni.io_processor = io_processor
+
+    serving_models = OpenAIServingModels(
+        engine_client=omni,
+        base_model_paths=[BaseModelPath(name="tts-model", model_path="Qwen/Qwen3-TTS")],
+    )
+
+    assert serving_models.model_config is model_config
+    assert serving_models.renderer is renderer
+    assert serving_models.io_processor is io_processor
+    assert serving_models.input_processor is input_processor

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -3,87 +3,51 @@
 
 import pytest
 
-from vllm_omni.entrypoints import utils as utils_module
-from vllm_omni.entrypoints.async_omni import AsyncOmni
+from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
-MODEL = "riverclouds/qwen_image_random"
+
+def test_default_stage_config_includes_cache_backend():
+    """Ensure cache knobs survive the default diffusion-stage builder."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "cache_backend": "cache_dit",
+            "cache_config": '{"Fn_compute_blocks": 2}',
+            "vae_use_slicing": True,
+            "ulysses_degree": 2,
+        }
+    )[0]
+
+    engine_args = stage_cfg["engine_args"]
+    assert stage_cfg["stage_type"] == "diffusion"
+    assert engine_args["cache_backend"] == "cache_dit"
+    assert engine_args["cache_config"]["Fn_compute_blocks"] == 2
+    assert engine_args["vae_use_slicing"] is True
+    assert engine_args["parallel_config"].ulysses_degree == 2
+    assert engine_args["model_stage"] == "diffusion"
 
 
-def _noop_inline_engine(self, model, stage_config, kwargs):
-    self._inline_diffusion = False
-    self._inline_engine = None
+def test_default_cache_config_used_when_missing():
+    """Ensure default cache_config is synthesized when only backend is given."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "cache_backend": "cache_dit",
+        }
+    )[0]
 
-
-def test_default_stage_config_includes_cache_backend(monkeypatch):
-    """Ensure cache_backend/cache_config are preserved in default diffusion stage."""
-    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)
-    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
-
-    omni = AsyncOmni(
-        model=MODEL,
-        cache_backend="cache_dit",
-        cache_config='{"Fn_compute_blocks": 2}',
-        vae_use_slicing=True,
-        ulysses_degree=2,
-    )
-
-    stage_cfg = omni.stage_configs[0]
-    engine_args = stage_cfg.engine_args
-
-    assert engine_args.get("cache_backend") == "cache_dit"
-    cache_config = engine_args.get("cache_config")
-    assert cache_config["Fn_compute_blocks"] == 2
-    assert engine_args.get("vae_use_slicing") is True
-    parallel_config = engine_args.get("parallel_config")
-    if hasattr(parallel_config, "get"):
-        ulysses_degree = parallel_config.get("ulysses_degree")
-    else:
-        ulysses_degree = getattr(parallel_config, "ulysses_degree", None)
-    assert ulysses_degree == 2
-
-
-def test_default_cache_config_used_when_missing(monkeypatch):
-    """Ensure default cache_config is applied when cache_backend is set."""
-    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)
-    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
-
-    omni = AsyncOmni(
-        model=MODEL,
-        cache_backend="cache_dit",
-    )
-
-    engine_args = omni.stage_configs[0].engine_args
-    cache_config = engine_args.get("cache_config")
+    cache_config = stage_cfg["engine_args"]["cache_config"]
     assert cache_config is not None
     assert cache_config["Fn_compute_blocks"] == 1
 
 
-def test_default_stage_devices_from_sequence_parallel(monkeypatch):
-    """Ensure devices list reflects sequence parallel size when no parallel_config is provided."""
-    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)
-    monkeypatch.setattr(AsyncOmni, "_init_inline_diffusion_engine", _noop_inline_engine)
+def test_default_stage_devices_from_sequence_parallel():
+    """Ensure runtime devices reflect computed diffusion world size."""
+    stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
+        {
+            "ulysses_degree": 2,
+            "ring_degree": 2,
+        }
+    )[0]
 
-    omni = AsyncOmni(
-        model=MODEL,
-        ulysses_degree=2,
-        ring_degree=2,
-    )
-
-    stage_cfg = omni.stage_configs[0]
-    runtime = stage_cfg.runtime
-    if hasattr(runtime, "get"):
-        devices = runtime.get("devices")
-    else:
-        devices = getattr(runtime, "devices", None)
-    assert devices == "0,1,2,3"
+    assert stage_cfg["runtime"]["devices"] == "0,1,2,3"

tests/entrypoints/test_omni_input_preprocessor.py

@@ -1,63 +0,0 @@
-import pytest
-
-from vllm_omni.inputs.preprocess import OmniInputPreprocessor
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def _make_preprocessor(monkeypatch):
-    preprocessor = object.__new__(OmniInputPreprocessor)
-    monkeypatch.setattr(preprocessor, "_truncate_inputs", lambda tokens, tokenization_kwargs=None: tokens)
-    monkeypatch.setattr(
-        preprocessor,
-        "_process_multimodal",
-        lambda *args, **kwargs: {"prompt_token_ids": [1, 2, 3]},
-    )
-    monkeypatch.setattr(preprocessor, "_tokenize_prompt", lambda prompt_text, tokenization_kwargs=None: [9, 8, 7])
-    return preprocessor
-
-
-def test_process_tokens_keeps_additional_information(monkeypatch):
-    preprocessor = _make_preprocessor(monkeypatch)
-    parsed = {
-        "prompt_token_ids": [1, 2, 3],
-        "prompt_embeds": "embeds",
-        "additional_information": {"task": ["tts"], "lang": ["auto"]},
-    }
-
-    inputs = OmniInputPreprocessor._process_tokens(preprocessor, parsed)
-
-    assert inputs["prompt_token_ids"] == [1, 2, 3]
-    assert inputs["prompt_embeds"] == "embeds"
-    assert inputs["additional_information"] == {"task": ["tts"], "lang": ["auto"]}
-
-
-def test_process_text_keeps_additional_information(monkeypatch):
-    preprocessor = _make_preprocessor(monkeypatch)
-    parsed = {
-        "prompt": "hello",
-        "prompt_embeds": "embeds",
-        "additional_information": {"speaker": ["alice"]},
-    }
-
-    inputs = OmniInputPreprocessor._process_text(preprocessor, parsed)
-
-    assert inputs["prompt_token_ids"] == [9, 8, 7]
-    assert inputs["prompt_embeds"] == "embeds"
-    assert inputs["additional_information"] == {"speaker": ["alice"]}
-
-
-def test_process_text_multimodal_skips_empty_payloads(monkeypatch):
-    preprocessor = _make_preprocessor(monkeypatch)
-    parsed = {
-        "prompt": "hello",
-        "multi_modal_data": {"image": "fake"},
-        "prompt_embeds": None,
-        "additional_information": None,
-    }
-
-    inputs = OmniInputPreprocessor._process_text(preprocessor, parsed)
-
-    assert inputs["prompt_token_ids"] == [1, 2, 3]
-    assert "prompt_embeds" not in inputs
-    assert "additional_information" not in inputs

tests/entrypoints/test_omni_stage_diffusion_config.py

@@ -1,31 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import pytest
-
-from vllm_omni.entrypoints.omni_stage import _build_od_config
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_build_od_config_includes_diffusion_fields():
-    engine_args = {
-        "cache_backend": "cache_dit",
-        "cache_config": {"Fn_compute_blocks": 2},
-        "vae_use_slicing": True,
-    }
-    od_config = _build_od_config(engine_args, model="dummy-model")
-
-    assert od_config["model"] == "dummy-model"
-    assert od_config["cache_backend"] == "cache_dit"
-    assert od_config["cache_config"]["Fn_compute_blocks"] == 2
-    assert od_config["vae_use_slicing"] is True
-
-
-def test_build_od_config_respects_explicit_config():
-    engine_args = {
-        "od_config": {"cache_backend": "tea_cache"},
-        "cache_backend": "cache_dit",
-    }
-    od_config = _build_od_config(engine_args, model="dummy-model")
-    assert od_config == {"cache_backend": "tea_cache"}

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -227,12 +227,14 @@ def test_first_streaming_chunk_prepends_ref_code_context():
     assert len(payload["code_predictor_codes"]) == _Q * 12
 
 
-def test_ref_code_context_only_applies_to_first_streaming_chunk():
+def test_ref_code_context_applies_to_all_streaming_chunks():
+    """ref_code is prepended as decoder context on every chunk, not just the first."""
     tm = _tm()
     rid = "r-ref2"
     tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(20)]
     tm.put_req_chunk[rid] = 1
     ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+    tm.request_payload[rid] = ref_code
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
@@ -242,8 +244,9 @@ def test_ref_code_context_only_applies_to_first_streaming_chunk():
     )
 
     assert payload is not None
-    assert payload["left_context_size"] == 10
-    assert len(payload["code_predictor_codes"]) == _Q * 20
+    # ref_code (2 frames) prepended as left context on second chunk too
+    assert payload["left_context_size"] == 10 + 2
+    assert len(payload["code_predictor_codes"]) == _Q * (20 + 2)
 
 
 def test_ref_code_context_can_be_buffered_before_first_emit():
@@ -276,9 +279,10 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
     )
 
     assert payload is not None
+    # ref_code (2 frames) is kept (not popped) for subsequent chunks
     assert payload["left_context_size"] == 2
     assert len(payload["code_predictor_codes"]) == _Q * 12
-    assert rid not in tm.request_payload
+    assert rid in tm.request_payload
 
 
 def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
@@ -317,3 +321,27 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
         4,
         8,
     ]
+
+
+def test_non_async_processor_filters_out_of_range_codec_values():
+    """Frames with values >= codebook_size (e.g. stop_token_id=2150) are filtered."""
+    ref_code = torch.tensor([[9, 9, 9, 9]], dtype=torch.long)
+    audio_codes = torch.tensor(
+        [
+            [0, 0, 0, 0],  # zero-padded (filtered)
+            [1, 2, 3, 4],  # valid
+            [2150, 0, 0, 0],  # stop token (filtered)
+            [5, 6, 7, 8],  # valid
+        ],
+        dtype=torch.long,
+    )
+    output = SimpleNamespace(multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code})
+    stage = SimpleNamespace(engine_outputs=[SimpleNamespace(outputs=[output])])
+
+    prompts = talker2code2wav(stage_list=[stage], engine_input_source=[0])
+
+    assert len(prompts) == 1
+    prompt = prompts[0]
+    # Only ref_code (1 frame) + 2 valid frames = 3 frames * 4 quantizers = 12 codes
+    assert len(prompt["prompt_token_ids"]) == 4 * 3
+    assert prompt["additional_information"] == {"left_context_size": [1]}

tests/perf/scripts/run_benchmark.py

@@ -50,11 +50,14 @@ def create_unique_server_params(configs: list[dict[str, Any]]) -> list[tuple[str
     for config in configs:
         test_name = config["test_name"]
         model = config["server_params"]["model"]
-        stage_config_name = config["server_params"]["stage_config_name"]
-        stage_config_path = str(Path(__file__).parent.parent / "stage_configs" / stage_config_name)
-        delete = config["server_params"].get("delete", None)
-        update = config["server_params"].get("update", None)
-        stage_config_path = modify_stage(stage_config_path, update, delete)
+        stage_config_name = config["server_params"].get("stage_config_name")
+        if stage_config_name:
+            stage_config_path = str(Path(__file__).parent.parent / "stage_configs" / stage_config_name)
+            delete = config["server_params"].get("delete", None)
+            update = config["server_params"].get("update", None)
+            stage_config_path = modify_stage(stage_config_path, update, delete)
+        else:
+            stage_config_path = None
 
         server_param = (test_name, model, stage_config_path)
         if server_param not in seen:
@@ -98,7 +101,10 @@ def omni_server(request):
 
         print(f"Starting OmniServer with test: {test_name}, model: {model}")
 
-        with OmniServer(model, ["--stage-configs-path", stage_config_path, "--stage-init-timeout", "120"]) as server:
+        server_args = ["--stage-init-timeout", "120"]
+        if stage_config_path:
+            server_args = ["--stage-configs-path", stage_config_path] + server_args
+        with OmniServer(model, server_args) as server:
             server.test_name = test_name
             print("OmniServer started successfully")
             yield server

tests/perf/stage_configs/qwen3_omni.yaml

@@ -7,7 +7,7 @@
 async_chunk: false
 stage_args:
   - stage_id: 0
-    stage_type: llm  # Use llm stage type to launch OmniLLM
+    stage_type: llm  # Use llm stage type for AR stages
     runtime:
       devices: "0"
       max_batch_size: 64
@@ -38,7 +38,7 @@ stage_args:
       repetition_penalty: 1.05
 
   - stage_id: 1
-    stage_type: llm  # Use llm stage type to launch OmniLLM
+    stage_type: llm  # Use llm stage type for AR stages
     runtime:
       devices: "1"
       max_batch_size: 64
@@ -69,7 +69,7 @@ stage_args:
       stop_token_ids: [2150]
 
   - stage_id: 2
-    stage_type: llm  # Use llm stage type to launch OmniLLM
+    stage_type: llm  # Use llm stage type for AR stages
     runtime:
       devices: "1"
       max_batch_size: 64

tests/perf/stage_configs/qwen3_tts.yaml

@@ -7,14 +7,13 @@ async_chunk: true
 stage_args:
   - stage_id: 0
     stage_type: llm
+    is_comprehension: true
     runtime:
       devices: "0"
       max_batch_size: 4
     engine_args:
       model_stage: qwen3_tts
       model_arch: Qwen3TTSTalkerForConditionalGeneration
-      hf_overrides:
-        architectures: [Qwen3TTSTalkerForConditionalGeneration]
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       enforce_eager: false
@@ -46,8 +45,6 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3TTSCode2Wav
-      hf_overrides:
-        architectures: [Qwen3TTSCode2Wav]
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true

tests/perf/tests/test.json

@@ -152,8 +152,7 @@
     {
         "test_name": "test_qwen3_tts",
         "server_params": {
-            "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
-            "stage_config_name": "qwen3_tts.yaml"
+            "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
         },
         "benchmark_params": [
             {

tools/nightly/generate_nightly_perf_excel.py

@@ -22,6 +22,49 @@ LOGGER = logging.getLogger(__name__)
 
 GREY_BLOCK_FILL = PatternFill(start_color="D3D3D3", fill_type="solid")
 
+# Diffusion sheet columns (Qwen-Image diffusion benchmark).
+DIFFUSION_BENCHMARK_COLUMNS: tuple[str, ...] = (
+    "duration",
+    "completed_requests",
+    "failed_requests",
+    "throughput_qps",
+    "latency_mean",
+    "latency_median",
+    "latency_p50",
+    "latency_p99",
+    "peak_memory_mb_max",
+    "peak_memory_mb_mean",
+    "peak_memory_mb_median",
+    "slo_attainment_rate",
+)
+
+DIFFUSION_NUMERIC_FORMAT_COLUMNS: tuple[str, ...] = DIFFUSION_BENCHMARK_COLUMNS
+
+DIFFUSION_SUMMARY_COLUMNS: tuple[str, ...] = (
+    "date",
+    "test_name",
+    "model",
+    "backend",
+    "dataset",
+    "task",
+    "completed_requests",
+    "failed_requests",
+    "duration",
+    "throughput_qps",
+    "latency_mean",
+    "latency_median",
+    "latency_p50",
+    "latency_p99",
+    "peak_memory_mb_max",
+    "peak_memory_mb_mean",
+    "peak_memory_mb_median",
+    "slo_attainment_rate",
+    "commit_sha",
+    "build_id",
+    "build_url",
+    "source_file",
+)
+
 # Benchmark metric columns: grey the latest row's cell when value changed vs previous date.
 BENCHMARK_COLUMNS: tuple[str, ...] = (
     "num_prompts",
@@ -53,10 +96,29 @@ BENCHMARK_COLUMNS: tuple[str, ...] = (
 NUMERIC_FORMAT_COLUMNS: tuple[str, ...] = tuple(
     c for c in BENCHMARK_COLUMNS if c not in ("request_rate", "max_concurrency")
 )
+DATASET_NAME_ALLOWED = ("random", "random-mm")
 
 _COLUMNS_FILENAME = "nightly_perf_summary_columns.txt"
+_RESULT_JSON_PREFIX = "result_test_"
+_DIFFUSION_JSON_PREFIX = "diffusion_perf_"
 DEFAULT_INPUT_DIR = os.getenv("DEFAULT_INPUT_DIR") if os.getenv("DEFAULT_INPUT_DIR") else "tests"
 DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR") if os.getenv("DEFAULT_OUTPUT_DIR") else "tests"
+DEFAULT_DIFFUSION_INPUT_DIR = os.getenv("DIFFUSION_BENCHMARK_DIR")
+# Read omni/diffusion benchmarks from DEFAULT_INPUT_DIR, if DEFAULT_DIFFUSION_INPUT_DIR is not set.
+
+
+def _omni_group_key(record: dict[str, Any]) -> tuple[Any, ...]:
+    return (
+        record.get("model_id") or "",
+        record.get("test_name") or "",
+        record.get("dataset_name") or "",
+        record.get("max_concurrency") if record.get("max_concurrency") is not None else 0,
+        record.get("num_prompts") if record.get("num_prompts") is not None else 0,
+    )
+
+
+def _diffusion_group_key(record: dict[str, Any]) -> tuple[Any, ...]:
+    return (record.get("test_name") or "",)
 
 
 def _load_summary_columns(script_dir: str) -> list[str]:
@@ -68,6 +130,8 @@ def _load_summary_columns(script_dir: str) -> list[str]:
         "backend",
         "model_id",
         "tokenizer_id",
+        "test_name",
+        "dataset_name",
         "num_prompts",
         "request_rate",
         "burstiness",
@@ -106,6 +170,33 @@ def _load_summary_columns(script_dir: str) -> list[str]:
     return columns if columns else default
 
 
+def _ensure_omni_summary_columns(summary_columns: list[str]) -> list[str]:
+    """Ensure omni summary contains required columns, even when a custom columns file exists."""
+    required = ("test_name", "dataset_name", "source_file")
+    existing = set(summary_columns)
+    if all(c in existing for c in required):
+        return summary_columns
+
+    out = list(summary_columns)
+    insert_after = "tokenizer_id"
+    if "test_name" not in existing:
+        if insert_after in existing:
+            idx = out.index(insert_after) + 1
+            out.insert(idx, "test_name")
+        else:
+            out.append("test_name")
+    existing = set(out)
+    if "dataset_name" not in existing:
+        if "test_name" in out:
+            idx = out.index("test_name") + 1
+            out.insert(idx, "dataset_name")
+        else:
+            out.append("dataset_name")
+    if "source_file" not in set(out):
+        out.append("source_file")
+    return out
+
+
 def _vllm_omni_root() -> str:
     """Resolve vllm-omni repo root: directory that contains a 'tests' subdir (and usually 'tools')."""
     path = os.path.dirname(os.path.abspath(__file__))
@@ -122,6 +213,11 @@ def _default_input_dir() -> str:
     return os.path.join(root, DEFAULT_INPUT_DIR)
 
 
+def _default_diffusion_input_dir(input_dir: str) -> str:
+    """Default diffusion input dir: DIFFUSION_BENCHMARK_DIR if set, else fall back to omni input dir."""
+    return DEFAULT_DIFFUSION_INPUT_DIR if DEFAULT_DIFFUSION_INPUT_DIR else input_dir
+
+
 def _default_output_file() -> str:
     """Default: vllm-omni root / DEFAULT_OUTPUT_DIR / nightly_perf_<timestamp>.xlsx."""
     ts = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
@@ -139,12 +235,23 @@ def parse_args() -> argparse.Namespace:
         default=_default_input_dir(),
         help="Directory containing performance JSON files; default is <vllm-omni-root>/DEFAULT_INPUT_DIR.",
     )
+    parser.add_argument(
+        "--diffusion-input-dir",
+        type=str,
+        default=None,
+        help=(
+            "Directory containing diffusion_perf_*.json files; default is "
+            "DIFFUSION_BENCHMARK_DIR, fallback to --input-dir."
+        ),
+    )
     parser.add_argument(
         "--output-file",
         type=str,
         default=_default_output_file(),
-        help="Output path of the Excel report; \
-            default is <vllm-omni-root>/DEFAULT_OUTPUT_DIR/nightly_perf_<timestamp>.xlsx.",
+        help=(
+            "Output path of the Excel report; default is "
+            "<vllm-omni-root>/DEFAULT_OUTPUT_DIR/nightly_perf_<timestamp>.xlsx."
+        ),
     )
     parser.add_argument(
         "--commit-sha",
@@ -183,11 +290,54 @@ def _load_json_file(path: str) -> dict[str, Any] | None:
     return data
 
 
-def _iter_json_records(input_dir: str) -> Iterable[dict[str, Any]]:
-    """Iterate over JSON files in the input directory and yield normalized records.
-    commit_sha/build_id/build_url are not set here; they are applied later only to
-    rows with the latest date (see _apply_build_metadata_to_latest_only).
+def _parse_from_filename(filename: str) -> dict[str, Any]:
+    """Parse test-related metadata from a result JSON filename.
+
+    Expected pattern (after prefix/suffix stripped):
+    <test_name>_<dataset_name>_<max_concurrency>_<num_prompts>_<timestamp>
     """
+    name, ext = os.path.splitext(filename)
+    if ext != ".json" or not name.startswith(_RESULT_JSON_PREFIX):
+        return {}
+
+    core = name[len(_RESULT_JSON_PREFIX) :]
+    parts = core.split("_")
+    if len(parts) < 5:
+        LOGGER.warning("filename '%s' does not match expected pattern, skip parsing test metadata", filename)
+        return {}
+
+    timestamp = parts[-1]
+    num_prompts_str = parts[-2]
+    max_concurrency_str = parts[-3]
+    dataset_name = parts[-4]
+    test_name = "_".join(parts[:-4]) if parts[:-4] else ""
+
+    parsed: dict[str, Any] = {}
+
+    if len(timestamp) >= 15:
+        parsed["date"] = timestamp
+
+    if dataset_name in DATASET_NAME_ALLOWED:
+        parsed["dataset_name"] = dataset_name
+
+    try:
+        parsed["num_prompts"] = int(num_prompts_str)
+    except (TypeError, ValueError):
+        pass
+
+    try:
+        parsed["max_concurrency"] = int(max_concurrency_str)
+    except (TypeError, ValueError):
+        pass
+
+    if test_name:
+        parsed["test_name"] = test_name
+
+    return parsed
+
+
+def _iter_omni_json_records(input_dir: str) -> Iterable[dict[str, Any]]:
+    """Iterate over result_test_*.json files and yield normalized omni records."""
     if not os.path.isdir(input_dir):
         LOGGER.warning("input dir '%s' does not exist or is not a directory", input_dir)
         return
@@ -195,6 +345,8 @@ def _iter_json_records(input_dir: str) -> Iterable[dict[str, Any]]:
     for entry in sorted(os.listdir(input_dir)):
         if not entry.endswith(".json"):
             continue
+        if not entry.startswith(_RESULT_JSON_PREFIX):
+            continue
         full_path = os.path.join(input_dir, entry)
         if not os.path.isfile(full_path):
             continue
@@ -204,17 +356,89 @@ def _iter_json_records(input_dir: str) -> Iterable[dict[str, Any]]:
             continue
 
         record: dict[str, Any] = dict(data)
-        record.setdefault("date", datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S"))
+        filename_meta = _parse_from_filename(os.path.basename(full_path))
+
+        if "date" not in record or not record["date"]:
+            if "date" in filename_meta:
+                record["date"] = filename_meta["date"]
+            else:
+                record["date"] = datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
+
+        if "num_prompts" not in record or record["num_prompts"] is None:
+            if "num_prompts" in filename_meta:
+                record["num_prompts"] = filename_meta["num_prompts"]
+
+        if "max_concurrency" not in record or record["max_concurrency"] is None:
+            if "max_concurrency" in filename_meta:
+                record["max_concurrency"] = filename_meta["max_concurrency"]
+
+        if "test_name" not in record or not record.get("test_name"):
+            if "test_name" in filename_meta:
+                record["test_name"] = filename_meta["test_name"]
+
+        if "dataset_name" not in record or not record.get("dataset_name"):
+            if "dataset_name" in filename_meta:
+                record["dataset_name"] = filename_meta["dataset_name"]
+
+        record["source_file"] = os.path.basename(full_path)
+        yield record
+
+
+def _parse_diffusion_from_filename(filename: str) -> dict[str, Any]:
+    """Parse diffusion test_name/date from filename: diffusion_perf_<test_name>_<YYYYMMDD-HHMMSS>.json"""
+    name, ext = os.path.splitext(filename)
+    if ext != ".json" or not name.startswith(_DIFFUSION_JSON_PREFIX):
+        return {}
+    core = name[len(_DIFFUSION_JSON_PREFIX) :]
+    parts = core.split("_")
+    if len(parts) < 2:
+        return {}
+    timestamp = parts[-1]
+    test_name = "_".join(parts[:-1]) if parts[:-1] else ""
+    parsed: dict[str, Any] = {}
+    if len(timestamp) >= 15:
+        parsed["date"] = timestamp
+    if test_name:
+        parsed["test_name"] = test_name
+    return parsed
+
+
+def _iter_diffusion_json_records(input_dir: str) -> Iterable[dict[str, Any]]:
+    """Iterate over diffusion_perf_*.json files and yield normalized diffusion records."""
+    if not os.path.isdir(input_dir):
+        LOGGER.warning("diffusion input dir '%s' does not exist or is not a directory", input_dir)
+        return
+
+    for entry in sorted(os.listdir(input_dir)):
+        if not entry.endswith(".json"):
+            continue
+        if not entry.startswith(_DIFFUSION_JSON_PREFIX):
+            continue
+        full_path = os.path.join(input_dir, entry)
+        if not os.path.isfile(full_path):
+            continue
+
+        data = _load_json_file(full_path)
+        if data is None:
+            continue
+
+        record: dict[str, Any] = dict(data)
+        filename_meta = _parse_diffusion_from_filename(os.path.basename(full_path))
+        if "date" not in record or not record.get("date"):
+            record["date"] = filename_meta.get("date") or datetime.now(timezone.utc).strftime("%Y%m%d-%H%M%S")
+        if "test_name" not in record or not record.get("test_name"):
+            if "test_name" in filename_meta:
+                record["test_name"] = filename_meta["test_name"]
         record["source_file"] = os.path.basename(full_path)
         yield record
 
 
 def _collect_records(input_dir: str) -> list[dict[str, Any]]:
-    """Collect all JSON records into a list."""
-    records: list[dict[str, Any]] = []
-    for record in _iter_json_records(input_dir):
-        records.append(record)
-    return records
+    return list(_iter_omni_json_records(input_dir))
+
+
+def _collect_diffusion_records(diffusion_input_dir: str) -> list[dict[str, Any]]:
+    return list(_iter_diffusion_json_records(diffusion_input_dir))
 
 
 def _apply_build_metadata_to_latest_only(
@@ -241,9 +465,17 @@ def _apply_build_metadata_to_latest_only(
 
 
 def _sort_records_for_summary(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
-    """Sort so that same model_id is grouped, newest date first within each group."""
+    """Sort so that same test configuration is grouped, newest date first within each group."""
     by_date_desc = sorted(records, key=lambda r: (r.get("date") or ""), reverse=True)
-    return sorted(by_date_desc, key=lambda r: (r.get("model_id") or ""))
+    return sorted(
+        by_date_desc,
+        key=_omni_group_key,
+    )
+
+
+def _sort_diffusion_records_for_summary(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    by_date_desc = sorted(records, key=lambda r: (r.get("date") or ""), reverse=True)
+    return sorted(by_date_desc, key=_diffusion_group_key)
 
 
 def _values_differ(a: Any, b: Any) -> bool:
@@ -261,21 +493,22 @@ def _values_differ(a: Any, b: Any) -> bool:
     return a != b
 
 
-def _apply_benchmark_change_highlight(
+def _apply_change_highlight(
     ws,
-    summary_columns: Sequence[str],
+    columns: Sequence[str],
     records: Sequence[dict[str, Any]],
+    benchmark_columns: Sequence[str],
+    group_key_fn,
 ) -> None:
-    """Grey cells in the latest row of each model when a benchmark metric changed vs previous date."""
+    """Grey cells in the latest row of each group when a metric changed vs previous row in the same group."""
     if not records:
         return
-    col_to_index = {c: i + 1 for i, c in enumerate(summary_columns)}
-    # Walk by model_id blocks (records already sorted by model_id, date desc).
+    col_to_index = {c: i + 1 for i, c in enumerate(columns)}
     i = 0
     while i < len(records):
-        model_id = records[i].get("model_id")
+        group_key = group_key_fn(records[i])
         block_start = i
-        while i < len(records) and records[i].get("model_id") == model_id:
+        while i < len(records) and group_key_fn(records[i]) == group_key:
             i += 1
         block_end = i
         newest_idx = block_start
@@ -283,7 +516,7 @@ def _apply_benchmark_change_highlight(
         if prev_idx is None:
             continue
         excel_row = newest_idx + 2
-        for col in BENCHMARK_COLUMNS:
+        for col in benchmark_columns:
             if col not in col_to_index:
                 continue
             cur_val = records[newest_idx].get(col)
@@ -343,18 +576,20 @@ def _write_sheet(
         ws.append(row_values)
 
 
-def _format_benchmark_columns(
+def _format_numeric_columns(
     ws,
     columns: Sequence[str],
+    numeric_columns: Sequence[str],
     num_data_rows: int,
+    width: int,
 ) -> None:
-    """Set number format and column width for numeric benchmark columns so values display (no ####)."""
-    numeric_set = set(NUMERIC_FORMAT_COLUMNS)
+    """Set number format and column width for numeric columns so values display (no ####)."""
+    numeric_set = set(numeric_columns)
     for c, col_name in enumerate(columns):
         if col_name not in numeric_set:
             continue
         col_letter = get_column_letter(c + 1)
-        ws.column_dimensions[col_letter].width = 14
+        ws.column_dimensions[col_letter].width = width
         for r in range(2, 2 + num_data_rows):
             cell = ws.cell(row=r, column=c + 1)
             if cell.value is not None and isinstance(cell.value, (int, float)):
@@ -367,6 +602,35 @@ def _format_benchmark_columns(
                     pass
 
 
+def _set_column_width(ws, columns: Sequence[str], col_name: str, width: int) -> None:
+    """Set a single column width by column name if present."""
+    try:
+        idx = list(columns).index(col_name)
+    except ValueError:
+        return
+    ws.column_dimensions[get_column_letter(idx + 1)].width = width
+
+
+def _apply_column_widths(ws, columns: Sequence[str], widths: dict[str, int]) -> None:
+    for name, width in widths.items():
+        _set_column_width(ws, columns, name, width)
+
+
+_OMNI_SUMMARY_WIDTHS = {
+    "endpoint_type": 10,
+    "backend": 10,
+    "model_id": 20,
+    "tokenizer_id": 20,
+    "test_name": 20,
+    "dataset_name": 10,
+}
+
+_DIFFUSION_SUMMARY_WIDTHS = {
+    "test_name": 30,
+    "model": 15,
+}
+
+
 def _ensure_parent_dir(path: str) -> None:
     """Ensure that the parent directory of the output file exists."""
     parent = os.path.dirname(os.path.abspath(path))
@@ -377,34 +641,79 @@ def _ensure_parent_dir(path: str) -> None:
 
 def generate_excel_report(
     input_dir: str,
+    diffusion_input_dir: str,
     output_file: str,
     commit_sha: str | None,
     build_id: str | None,
     build_url: str | None,
 ) -> None:
-    """Main logic: load JSON records and generate an Excel report."""
+    """Main logic: load JSON records and generate an Excel report with multiple sheets."""
     script_dir = os.path.dirname(os.path.abspath(__file__))
-    summary_columns = _load_summary_columns(script_dir)
+    omni_summary_columns = _ensure_omni_summary_columns(_load_summary_columns(script_dir))
 
-    records = _collect_records(input_dir)
-    if not records:
-        LOGGER.warning("no valid json records found under '%s'", input_dir)
+    omni_records = _collect_records(input_dir)
+    diffusion_records = _collect_diffusion_records(diffusion_input_dir)
 
-    sorted_records = _sort_records_for_summary(records)
-    _apply_build_metadata_to_latest_only(sorted_records, commit_sha, build_id, build_url)
+    if not omni_records:
+        LOGGER.warning("no valid omni json records found under '%s'", input_dir)
+    if not diffusion_records:
+        LOGGER.warning("no valid diffusion json records found under '%s'", diffusion_input_dir)
+
+    omni_sorted = _sort_records_for_summary(omni_records)
+    diffusion_sorted = _sort_diffusion_records_for_summary(diffusion_records)
+
+    _apply_build_metadata_to_latest_only(omni_sorted, commit_sha, build_id, build_url)
+    _apply_build_metadata_to_latest_only(diffusion_sorted, commit_sha, build_id, build_url)
 
     wb = Workbook()
-    ws_summary = wb.active
-    ws_summary.title = "summary"
+    ws_omni_summary = wb.active
+    ws_omni_summary.title = "omni_summary"
+    _write_sheet(ws_omni_summary, omni_summary_columns, omni_sorted, numeric_columns=NUMERIC_FORMAT_COLUMNS)
+    _format_numeric_columns(
+        ws_omni_summary,
+        omni_summary_columns,
+        NUMERIC_FORMAT_COLUMNS,
+        len(omni_sorted),
+        width=14,
+    )
+    _apply_column_widths(ws_omni_summary, omni_summary_columns, _OMNI_SUMMARY_WIDTHS)
+    _apply_change_highlight(
+        ws_omni_summary,
+        omni_summary_columns,
+        omni_sorted,
+        BENCHMARK_COLUMNS,
+        _omni_group_key,
+    )
 
-    _write_sheet(ws_summary, summary_columns, sorted_records, numeric_columns=NUMERIC_FORMAT_COLUMNS)
-    _format_benchmark_columns(ws_summary, summary_columns, len(sorted_records))
-    _apply_benchmark_change_highlight(ws_summary, summary_columns, sorted_records)
+    ws_diff_summary = wb.create_sheet(title="diffusion_summary")
+    _write_sheet(
+        ws_diff_summary, DIFFUSION_SUMMARY_COLUMNS, diffusion_sorted, numeric_columns=DIFFUSION_NUMERIC_FORMAT_COLUMNS
+    )
+    _format_numeric_columns(
+        ws_diff_summary,
+        DIFFUSION_SUMMARY_COLUMNS,
+        DIFFUSION_NUMERIC_FORMAT_COLUMNS,
+        len(diffusion_sorted),
+        width=16,
+    )
+    _apply_column_widths(ws_diff_summary, DIFFUSION_SUMMARY_COLUMNS, _DIFFUSION_SUMMARY_WIDTHS)
+    _apply_change_highlight(
+        ws_diff_summary,
+        DIFFUSION_SUMMARY_COLUMNS,
+        diffusion_sorted,
+        DIFFUSION_BENCHMARK_COLUMNS,
+        _diffusion_group_key,
+    )
 
-    if sorted_records:
-        raw_columns = _build_raw_columns(sorted_records, summary_columns)
-        ws_raw = wb.create_sheet(title="raw")
-        _write_sheet(ws_raw, raw_columns, sorted_records)
+    if omni_sorted:
+        omni_raw_columns = _build_raw_columns(omni_sorted, omni_summary_columns)
+        ws_omni_raw = wb.create_sheet(title="omni_raw")
+        _write_sheet(ws_omni_raw, omni_raw_columns, omni_sorted)
+
+    if diffusion_sorted:
+        diffusion_raw_columns = _build_raw_columns(diffusion_sorted, DIFFUSION_SUMMARY_COLUMNS)
+        ws_diff_raw = wb.create_sheet(title="diffusion_raw")
+        _write_sheet(ws_diff_raw, diffusion_raw_columns, diffusion_sorted)
 
     _ensure_parent_dir(output_file)
     wb.save(output_file)
@@ -423,9 +732,11 @@ def main() -> None:
     commit_sha = args.commit_sha or os.getenv("BUILDKITE_COMMIT")
     build_id = args.build_id or os.getenv("BUILDKITE_BUILD_ID")
     build_url = args.build_url or os.getenv("BUILDKITE_BUILD_URL")
+    diffusion_input_dir = args.diffusion_input_dir or _default_diffusion_input_dir(args.input_dir)
 
     generate_excel_report(
         input_dir=args.input_dir,
+        diffusion_input_dir=diffusion_input_dir,
         output_file=args.output_file,
         commit_sha=commit_sha,
         build_id=build_id,

tools/pre_commit/check_pickle_imports.py

@@ -16,7 +16,6 @@ import regex as re
 #  alternatives like msgpack or pydantic that are already in use in vLLM. Only
 #  add to this list if absolutely necessary and after careful security review.
 ALLOWED_FILES = {
-    "vllm_omni/entrypoints/omni_llm.py",
     "tests/e2e/offline_inference/utils.py",
     "tests/utils.py",
     "vllm_omni/diffusion/distributed/group_coordinator.py",

vllm_omni/__init__.py

@@ -24,10 +24,7 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
 from vllm_omni.transformers_utils import configs as _configs  # noqa: F401, E402
 
 from .config import OmniModelConfig
-from .entrypoints.async_omni import AsyncOmni
-
-# Main entry points
-from .entrypoints.omni import Omni
+from .entrypoints import AsyncOmni, Omni
 
 from .version import __version__, __version_tuple__  # isort:skip
 

vllm_omni/config/stage_config.py

@@ -49,10 +49,11 @@ class StageType(str, Enum):
 
 @dataclass
 class StageConfig:
-    """Per-stage configuration — pipeline-structure fields only.
+    """Per-stage configuration from pipeline YAML.
 
-    Engine params (gpu_memory_utilization, tp_size, etc.) come from CLI,
-    NOT from this class.
+    Topology fields (stage_id, input_sources, etc.) define the DAG.
+    Engine and runtime defaults come from the YAML; CLI overrides take
+    precedence via ``runtime_overrides``.
     """
 
     # Identity
@@ -71,6 +72,14 @@ class StageConfig:
     hf_config_name: str | None = None
     is_comprehension: bool = False
 
+    # Per-stage engine args from pipeline YAML (defaults)
+    yaml_engine_args: dict[str, Any] = field(default_factory=dict)
+    # Per-stage runtime config from pipeline YAML (devices, max_batch_size)
+    yaml_runtime: dict[str, Any] = field(default_factory=dict)
+    # Pass-through fields from pipeline YAML (default_sampling_params,
+    # output_connectors, input_connectors, tts_args, etc.)
+    yaml_extras: dict[str, Any] = field(default_factory=dict)
+
     # Runtime overrides (populated from CLI, not from pipeline YAML)
     runtime_overrides: dict[str, Any] = field(default_factory=dict)
 
@@ -80,11 +89,11 @@ class StageConfig:
         Returns:
             OmegaConf DictConfig with stage configuration in legacy format.
         """
-        # Build engine_args dict with required fields
-        engine_args: dict[str, Any] = {
-            "model_stage": self.model_stage,
-        }
+        # Start with YAML engine_args defaults
+        engine_args: dict[str, Any] = dict(self.yaml_engine_args)
 
+        # Overlay topology-level fields
+        engine_args["model_stage"] = self.model_stage
         if self.worker_type:
             engine_args["worker_type"] = self.worker_type
         if self.scheduler_cls:
@@ -92,19 +101,22 @@ class StageConfig:
         if self.hf_config_name:
             engine_args["hf_config_name"] = self.hf_config_name
 
-        # Apply runtime overrides (CLI args)
+        # CLI overrides take precedence over YAML defaults
         for key, value in self.runtime_overrides.items():
             if key not in ("devices", "max_batch_size"):
                 engine_args[key] = value
 
-        # Build runtime config
-        runtime: dict[str, Any] = {
-            "process": True,
-            "max_batch_size": self.runtime_overrides.get("max_batch_size", 1),
-        }
+        # Build runtime config from YAML defaults + CLI overrides
+        runtime: dict[str, Any] = dict(self.yaml_runtime)
+        runtime.setdefault("process", True)
+        runtime.setdefault("max_batch_size", self.runtime_overrides.get("max_batch_size", 1))
         if "devices" in self.runtime_overrides:
             runtime["devices"] = self.runtime_overrides["devices"]
 
+        # Inject max_num_seqs from max_batch_size (legacy compat)
+        max_batch_size = int(runtime.get("max_batch_size", 1) or 1)
+        engine_args.setdefault("max_num_seqs", max_batch_size)
+
         # Build full config dict
         config_dict: dict[str, Any] = {
             "stage_id": self.stage_id,
@@ -120,6 +132,10 @@ class StageConfig:
         if self.custom_process_input_func:
             config_dict["custom_process_input_func"] = self.custom_process_input_func
 
+        # Pass through extra YAML fields (default_sampling_params,
+        # output_connectors, input_connectors, tts_args, etc.)
+        config_dict.update(self.yaml_extras)
+
         return create_config(config_dict)
 
 
@@ -249,6 +265,15 @@ class StageConfigFactory:
         if errors:
             logger.warning(f"Pipeline validation warnings for {model}: {errors}")
 
+        # Inject pipeline-wide async_chunk into ALL stages' engine_args.
+        # The legacy loader (load_stage_configs_from_yaml) sets async_chunk
+        # on every stage so that build_engine_args_dict() can inject the
+        # stage_connector_spec.  AsyncOmniEngine.__init__ also reads it
+        # from stage_configs[0].engine_args.async_chunk.
+        if pipeline.async_chunk:
+            for stage in pipeline.stages:
+                stage.yaml_engine_args.setdefault("async_chunk", True)
+
         # Apply CLI overrides
         result: list[StageConfig] = []
         for stage in pipeline.stages:
@@ -348,6 +373,25 @@ class StageConfigFactory:
 
         return cls._parse_pipeline_yaml(pipeline_path, model_type)
 
+    # Keys consumed as explicit StageConfig fields — everything else is
+    # passed through via yaml_extras.
+    _KNOWN_STAGE_KEYS: set[str] = {
+        "stage_id",
+        "model_stage",
+        "stage_type",
+        "input_sources",
+        "engine_input_source",
+        "custom_process_input_func",
+        "final_output",
+        "final_output_type",
+        "worker_type",
+        "scheduler_cls",
+        "hf_config_name",
+        "is_comprehension",
+        "engine_args",
+        "runtime",
+    }
+
     @classmethod
     def _parse_pipeline_yaml(cls, path: Path, model_type: str) -> ModelPipeline:
         """Parse a pipeline YAML file.
@@ -375,27 +419,65 @@ class StageConfigFactory:
                 input_sources = []
             input_sources = list(input_sources)
 
+            # Extract per-stage engine_args and runtime dicts
+            raw_ea = stage_data.get("engine_args", None)
+            yaml_engine_args = to_dict(raw_ea) if raw_ea is not None else {}
+            raw_rt = stage_data.get("runtime", None)
+            yaml_runtime = to_dict(raw_rt) if raw_rt is not None else {}
+
+            # Topology-level fields that also live inside engine_args in legacy
+            # YAMLs (worker_type, scheduler_cls, etc.) — read from both places.
+            worker_type = stage_data.get("worker_type", None) or yaml_engine_args.pop("worker_type", None)
+            scheduler_cls = stage_data.get("scheduler_cls", None) or yaml_engine_args.pop("scheduler_cls", None)
+            hf_config_name = stage_data.get("hf_config_name", None) or yaml_engine_args.pop("hf_config_name", None)
+            model_stage = getattr(stage_data, "model_stage", None) or yaml_engine_args.pop("model_stage", None)
+
+            # Collect pass-through fields (default_sampling_params,
+            # output_connectors, input_connectors, tts_args, etc.)
+            yaml_extras: dict[str, Any] = {}
+            for key in stage_data:
+                if key not in cls._KNOWN_STAGE_KEYS:
+                    val = stage_data[key]
+                    try:
+                        yaml_extras[key] = to_dict(val)
+                    except ValueError:
+                        yaml_extras[key] = val
+
             stage = StageConfig(
                 stage_id=stage_data.stage_id,
-                model_stage=stage_data.model_stage,
+                model_stage=model_stage or "",
                 stage_type=stage_type,
                 input_sources=input_sources,
                 custom_process_input_func=stage_data.get("custom_process_input_func", None),
                 final_output=stage_data.get("final_output", False),
                 final_output_type=stage_data.get("final_output_type", None),
-                worker_type=stage_data.get("worker_type", None),
-                scheduler_cls=stage_data.get("scheduler_cls", None),
-                hf_config_name=stage_data.get("hf_config_name", None),
+                worker_type=worker_type,
+                scheduler_cls=scheduler_cls,
+                hf_config_name=hf_config_name,
                 is_comprehension=stage_data.get("is_comprehension", False),
+                yaml_engine_args=yaml_engine_args,
+                yaml_runtime=yaml_runtime,
+                yaml_extras=yaml_extras,
             )
             stages.append(stage)
 
         # Get pipeline-wide flags
         async_chunk = config_data.get("async_chunk", False)
 
-        # Get optional connector config
-        connectors = to_dict(config_data.connectors) if hasattr(config_data, "connectors") else None
-        edges = to_dict(config_data.edges) if hasattr(config_data, "edges") else None
+        # Get optional connector config — check both top-level and nested
+        # under ``runtime`` (legacy stage_configs format).
+        connectors = None
+        edges = None
+        if hasattr(config_data, "connectors"):
+            connectors = to_dict(config_data.connectors)
+        if hasattr(config_data, "edges"):
+            edges = to_dict(config_data.edges)
+        if hasattr(config_data, "runtime") and config_data.runtime is not None:
+            top_runtime = config_data.runtime
+            if connectors is None and hasattr(top_runtime, "connectors"):
+                connectors = to_dict(top_runtime.connectors)
+            if edges is None and hasattr(top_runtime, "edges"):
+                edges = to_dict(top_runtime.edges)
 
         return ModelPipeline(
             model_type=getattr(config_data, "model_type", model_type),
@@ -421,9 +503,21 @@ class StageConfigFactory:
 
             hf_config = get_config(model, trust_remote_code=trust_remote_code)
             return hf_config.model_type, hf_config
+        except Exception:
+            pass
+
+        # Fallback: read config.json directly for custom model types that
+        # are not registered with transformers (e.g. qwen3_tts).
+        try:
+            from vllm.transformers_utils.config import get_hf_file_to_dict
+
+            config_dict = get_hf_file_to_dict("config.json", model, revision=None)
+            if config_dict and "model_type" in config_dict:
+                return config_dict["model_type"], None
         except Exception as e:
             logger.debug(f"Failed to auto-detect model type for {model}: {e}")
-            return None, None
+
+        return None, None
 
     # Keys that should never be forwarded as engine overrides (internal /
     # orchestrator-only knobs, complex objects, etc.).

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -606,6 +606,76 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No
     return refresh_cache_context
 
 
+def enable_cache_for_hunyuan_image3(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for HunyuanImage3 pipeline.
+
+    HunyuanImage3 stores its main transformer stack at ``pipeline.model`` with
+    decoder blocks in ``pipeline.model.layers``.
+    """
+    db_cache_config = _build_db_cache_config(cache_config)
+
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    transformer = getattr(pipeline, "model", None)
+    if transformer is None or not hasattr(transformer, "layers"):
+        raise ValueError(
+            f"HunyuanImage3 cache-dit enabler expects pipeline.model.layers, got pipeline={pipeline.__class__.__name__}"
+        )
+
+    logger.info(
+        f"Enabling cache-dit on HunyuanImage3 model: "
+        f"Fn={db_cache_config.Fn_compute_blocks}, "
+        f"Bn={db_cache_config.Bn_compute_blocks}, "
+        f"W={db_cache_config.max_warmup_steps}, "
+    )
+    # HunyuanImage3 decoder layers are single-stream `hidden_states` blocks
+    # (`HunyuanImage3DecoderLayer.forward(hidden_states, ...)`) while caller
+    # still reads outputs as tuple (`layer_outputs[0]`). Pattern_4 matches this:
+    # hidden-only input with hidden-first tuple output.
+    modifier = ParamsModifier(
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    cache_dit.enable_cache(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=transformer.layers,
+            forward_pattern=ForwardPattern.Pattern_4,
+            params_modifiers=[modifier],
+            check_forward_pattern=False,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        transformer = getattr(pipeline, "model", None)
+        if transformer is None:
+            raise ValueError("HunyuanImage3 cache-dit refresh expects pipeline.model to exist.")
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(transformer, num_inference_steps=num_inference_steps, verbose=verbose)
+        else:
+            cache_dit.refresh_context(
+                transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
 class BagelCachedContextManager(CachedContextManager):
     """
     Custom CachedContextManager for Bagel that safely handles NaiveCache objects
@@ -1064,6 +1134,7 @@ CUSTOM_DIT_ENABLERS.update(
         "Wan22Pipeline": enable_cache_for_wan22,
         "Wan22I2VPipeline": enable_cache_for_wan22,
         "Wan22TI2VPipeline": enable_cache_for_wan22,
+        "HunyuanImage3Pipeline": enable_cache_for_hunyuan_image3,
         "FluxPipeline": enable_cache_for_flux,
         "Flux2KleinPipeline": enable_cache_for_flux2_klein,
         "LongCatImagePipeline": enable_cache_for_longcat_image,

vllm_omni/diffusion/diffusion_engine.py

@@ -8,6 +8,7 @@ from typing import Any
 
 import numpy as np
 import PIL.Image
+import torch
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -106,8 +107,19 @@ class DiffusionEngine:
                 for i, prompt in enumerate(request.prompts)
             ]
 
+        # When CPU offload is enabled, move output to CPU before
+        # post-processing to avoid device OOM — model weights may still
+        # reside on the device and leave no headroom for intermediates.
+        output_data = output.output
+        if (
+            self.od_config.enable_cpu_offload
+            and isinstance(output_data, torch.Tensor)
+            and output_data.device.type != "cpu"
+        ):
+            output_data = output_data.cpu()
+
         postprocess_start_time = time.perf_counter()
-        outputs = self.post_process_func(output.output) if self.post_process_func is not None else output.output
+        outputs = self.post_process_func(output_data) if self.post_process_func is not None else output_data
         audio_payload = None
         if isinstance(outputs, dict):
             audio_payload = outputs.get("audio")

vllm_omni/diffusion/stage_diffusion_client.py

@@ -0,0 +1,146 @@
+"""Stage Diffusion Client for vLLM-Omni multi-stage runtime.
+
+Wraps AsyncOmniDiffusion to expose the same interface the Orchestrator
+expects from any stage client.
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.engine.stage_init_utils import StageMetadata
+from vllm_omni.entrypoints.async_omni_diffusion import AsyncOmniDiffusion
+from vllm_omni.outputs import OmniRequestOutput
+
+if TYPE_CHECKING:
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+
+logger = init_logger(__name__)
+
+
+class StageDiffusionClient:
+    """Wraps AsyncOmniDiffusion for use inside the Orchestrator.
+
+    Exposes the same attributes and async methods the Orchestrator
+    uses on StageEngineCoreClient, but routes execution through
+    DiffusionEngine instead of vLLM EngineCore.
+    """
+
+    stage_type: str = "diffusion"
+
+    def __init__(
+        self,
+        model: str,
+        od_config: OmniDiffusionConfig,
+        metadata: StageMetadata,
+    ) -> None:
+        self.stage_id = metadata.stage_id
+        self.final_output = metadata.final_output
+        self.final_output_type = metadata.final_output_type
+        self.default_sampling_params = metadata.default_sampling_params
+        self.custom_process_input_func = metadata.custom_process_input_func
+        self.engine_input_source = metadata.engine_input_source
+
+        self._engine = AsyncOmniDiffusion(model=model, od_config=od_config)
+        self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue()
+        self._tasks: dict[str, asyncio.Task] = {}
+
+        logger.info("[StageDiffusionClient] Stage-%s initialized", self.stage_id)
+
+    async def add_request_async(
+        self,
+        request_id: str,
+        prompt: OmniPromptType,
+        sampling_params: OmniDiffusionSamplingParams,
+    ) -> None:
+        task = asyncio.create_task(
+            self._run(request_id, prompt, sampling_params),
+            name=f"diffusion-{request_id}",
+        )
+        self._tasks[request_id] = task
+
+    async def _run(
+        self,
+        request_id: str,
+        prompt: OmniPromptType,
+        sampling_params: OmniDiffusionSamplingParams,
+    ) -> None:
+        try:
+            result = await self._engine.generate(prompt, sampling_params, request_id)
+            await self._output_queue.put(result)
+        except Exception as e:
+            logger.exception(
+                "[StageDiffusionClient] Stage-%s req=%s failed: %s",
+                self.stage_id,
+                request_id,
+                e,
+            )
+        finally:
+            self._tasks.pop(request_id, None)
+
+    def get_diffusion_output_async(self) -> OmniRequestOutput | None:
+        try:
+            return self._output_queue.get_nowait()
+        except asyncio.QueueEmpty:
+            return None
+
+    async def abort_requests_async(self, request_ids: list[str]) -> None:
+        for rid in request_ids:
+            task = self._tasks.pop(rid, None)
+            if task:
+                task.cancel()
+
+    async def collective_rpc_async(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple[Any, ...] = (),
+        kwargs: dict[str, Any] | None = None,
+    ) -> Any:
+        """Best-effort control RPC shim for diffusion stages.
+
+        TODO(AsyncOmni): add dedicated wrappers on AsyncOmniDiffusion for the
+        remaining control APIs instead of reaching into its underlying engine.
+        """
+        kwargs = kwargs or {}
+
+        if method in {"add_lora", "remove_lora", "list_loras", "pin_lora", "start_profile", "stop_profile"}:
+            target = getattr(self._engine, method, None)
+            if target is None:
+                return {
+                    "supported": False,
+                    "todo": True,
+                    "reason": f"AsyncOmniDiffusion.{method} is not implemented",
+                }
+            result = target(*args, **kwargs)
+            if timeout is not None:
+                return await asyncio.wait_for(result, timeout=timeout)
+            return await result
+
+        if method in {"sleep", "wake_up"}:
+            loop = asyncio.get_running_loop()
+            return await loop.run_in_executor(
+                self._engine._executor,
+                self._engine.engine.collective_rpc,
+                method,
+                timeout,
+                args,
+                kwargs,
+                None,
+            )
+
+        return {
+            "supported": False,
+            "todo": True,
+            "reason": f"Diffusion stage collective_rpc method {method} is not implemented yet",
+        }
+
+    def shutdown(self) -> None:
+        for task in self._tasks.values():
+            task.cancel()
+        self._tasks.clear()
+        self._engine.close()

vllm_omni/distributed/omni_connectors/adapter.py

@@ -1,13 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# temporary for compatibility with vllm_omni.entrypoints.omni_stage.py
-# and vllm_omni.entrypoints.omni_llm.py
-
 import time
 from collections.abc import Callable
 from typing import Any
 
-from vllm_omni.entrypoints.stage_utils import OmniStageTaskType
 from vllm_omni.metrics import OrchestratorAggregator
 
 from .utils.logging import get_connector_logger
@@ -66,7 +62,7 @@ def try_send_via_connector(
         if success:
             # Send lightweight notification via queue
             notify_payload = {
-                "type": OmniStageTaskType.GENERATE,
+                "type": "generate",
                 "request_id": req_id,
                 "sampling_params": sampling_params,
                 "from_connector": True,

vllm_omni/distributed/omni_connectors/transfer_adapter/base.py

@@ -2,7 +2,6 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import threading
-import time
 from collections import deque
 from typing import Any
 
@@ -33,6 +32,8 @@ class OmniTransferAdapterBase:
         self._finished_save_reqs = set()
 
         self.stop_event = threading.Event()
+        self._recv_cond = threading.Condition()
+        self._save_cond = threading.Condition()
 
         self.recv_thread = threading.Thread(target=self.recv_loop, daemon=True)
         self.recv_thread.start()
@@ -45,22 +46,37 @@ class OmniTransferAdapterBase:
         raise NotImplementedError
 
     def recv_loop(self):
-        """Loop to poll for incoming data."""
+        """Loop to poll for incoming data.
+
+        Process each pending request exactly once per pass.  When no request
+        made progress, back off 1 ms instead of tight-spinning on failed
+        shm_open syscalls (which can burn a full CPU core).
+        """
         while not self.stop_event.is_set():
-            # Iterate over a snapshot of pending requests
-            while self._pending_load_reqs:
+            n = len(self._pending_load_reqs)
+            any_success = False
+            for _ in range(n):
+                if not self._pending_load_reqs:
+                    break
                 request = self._pending_load_reqs.popleft()
                 request_id = request.request_id
                 self.request_ids_mapping[request_id] = request.external_req_id
                 try:
                     is_success = self._poll_single_request(request)
-                    if not is_success:
+                    if is_success:
+                        any_success = True
+                    else:
                         self._pending_load_reqs.append(request)
                 except Exception as e:
                     self._pending_load_reqs.append(request)
                     logger.warning(f"Error receiving data for {request_id}: {e}")
 
-            time.sleep(0.001)
+            # Timeout is the fallback for lock-free append/notify races.
+            with self._recv_cond:
+                if not self._pending_load_reqs and not self.stop_event.is_set():
+                    self._recv_cond.wait(timeout=0.1)
+                elif not any_success and not self.stop_event.is_set():
+                    self._recv_cond.wait(timeout=0.001)
 
     def save_loop(self):
         """Loop to send outgoing data."""
@@ -71,7 +87,10 @@ class OmniTransferAdapterBase:
                     self._send_single_request(task)
                 except Exception as e:
                     logger.warning(f"Error saving data for {task.get('request_id')}: {e}")
-            time.sleep(0.001)
+
+            with self._save_cond:
+                if not self._pending_save_reqs and not self.stop_event.is_set():
+                    self._save_cond.wait(timeout=0.1)
 
     def _poll_single_request(self, *args, **kwargs):
         """Poll connector for a single request task.
@@ -105,4 +124,13 @@ class OmniTransferAdapterBase:
 
     def shutdown(self):
         """Stop background loops and close the connector."""
-        raise NotImplementedError
+        self.stop_event.set()
+        with self._recv_cond:
+            self._recv_cond.notify_all()
+        with self._save_cond:
+            self._save_cond.notify_all()
+        if self.connector is not None:
+            try:
+                self.connector.close()
+            except Exception:
+                pass

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -95,6 +95,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         if not hasattr(request, "additional_information"):
             request.additional_information = None
         self._pending_load_reqs.append(request)
+        with self._recv_cond:
+            self._recv_cond.notify()
 
     def save_async(
         self,
@@ -116,6 +118,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             "is_finished": request.is_finished(),
         }
         self._pending_save_reqs.append(task)
+        with self._save_cond:
+            self._save_cond.notify()
 
     def _poll_single_request(self, request: Request):
         stage_id = self.connector.stage_id
@@ -302,6 +306,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         )
         while len(running_queue) > self.scheduler_max_num_seqs:
             request = running_queue.pop()
+            request.status = RequestStatus.PREEMPTED
             waiting_queue.prepend_requests([request])
 
     def restore_queues(self, waiting_queue: Any, running_queue: list[Request]) -> None:

vllm_omni/engine/arg_utils.py

@@ -1,8 +1,10 @@
+import argparse
+import dataclasses
 from dataclasses import dataclass, field
 from typing import Any
 
 import vllm.envs as envs
-from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
+from vllm.engine.arg_utils import EngineArgs
 from vllm.logger import init_logger
 from vllm.transformers_utils.gguf_utils import is_gguf
 
@@ -86,6 +88,12 @@ class OmniEngineArgs(EngineArgs):
         load_omni_general_plugins()
         super().__post_init__()
 
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace) -> "OmniEngineArgs":
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+        engine_args = cls(**{attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)})
+        return engine_args
+
     def _ensure_omni_models_registered(self):
         if hasattr(self, "_omni_models_registered"):
             return True
@@ -115,170 +123,8 @@ class OmniEngineArgs(EngineArgs):
         self._ensure_omni_models_registered()
 
         # Keep compatibility when async args are constructed from partial payloads.
-        limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {})
-        language_model_only = getattr(self, "language_model_only", False)
-        enable_mm_embeds = getattr(self, "enable_mm_embeds", False)
-        interleave_mm_strings = getattr(self, "interleave_mm_strings", False)
-        media_io_kwargs = getattr(self, "media_io_kwargs", {})
-        skip_mm_profiling = getattr(self, "skip_mm_profiling", False)
-        mm_processor_kwargs = getattr(self, "mm_processor_kwargs", None)
-        mm_processor_cache_gb = getattr(self, "mm_processor_cache_gb", 4)
-        mm_processor_cache_type = getattr(self, "mm_processor_cache_type", None)
-        mm_shm_cache_max_object_size_mb = getattr(self, "mm_shm_cache_max_object_size_mb", 128)
-        mm_encoder_only = getattr(self, "mm_encoder_only", False)
-        mm_encoder_tp_mode = getattr(self, "mm_encoder_tp_mode", "weights")
-        mm_encoder_attn_backend = getattr(self, "mm_encoder_attn_backend", None)
-        video_pruning_rate = getattr(self, "video_pruning_rate", 0.0)
-
-        # Build stage_connector_config from stage_connector_spec
-        stage_connector_config = {
-            "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"),
-            "extra": self.stage_connector_spec.get("extra", {}).copy(),
-        }
-        stage_connector_config["extra"]["stage_id"] = self.stage_id
-
-        # Create OmniModelConfig directly from engine args
-        # Note: We pass the actual init parameters matching vLLM's EngineArgs.create_model_config()
-        omni_config = OmniModelConfig(
-            # Base ModelConfig fields (matching vLLM's EngineArgs.create_model_config)
-            model=self.model,
-            model_weights=self.model_weights,
-            hf_config_path=self.hf_config_path,
-            runner=self.runner,
-            convert=self.convert,
-            tokenizer=self.tokenizer,
-            tokenizer_mode=self.tokenizer_mode,
-            trust_remote_code=self.trust_remote_code,
-            allowed_local_media_path=self.allowed_local_media_path,
-            allowed_media_domains=self.allowed_media_domains,
-            dtype=self.dtype,
-            seed=self.seed,
-            revision=self.revision,
-            code_revision=self.code_revision,
-            hf_token=self.hf_token,
-            hf_overrides=self.hf_overrides,
-            tokenizer_revision=self.tokenizer_revision,
-            max_model_len=self.max_model_len,
-            quantization=self.quantization,
-            allow_deprecated_quantization=self.allow_deprecated_quantization,
-            enforce_eager=self.enforce_eager,
-            enable_return_routed_experts=self.enable_return_routed_experts,
-            max_logprobs=self.max_logprobs,
-            logprobs_mode=self.logprobs_mode,
-            disable_sliding_window=self.disable_sliding_window,
-            disable_cascade_attn=self.disable_cascade_attn,
-            skip_tokenizer_init=self.skip_tokenizer_init,
-            enable_prompt_embeds=self.enable_prompt_embeds,
-            served_model_name=self.served_model_name,
-            language_model_only=language_model_only,
-            limit_mm_per_prompt=limit_mm_per_prompt,
-            enable_mm_embeds=enable_mm_embeds,
-            interleave_mm_strings=interleave_mm_strings,
-            media_io_kwargs=media_io_kwargs,
-            skip_mm_profiling=skip_mm_profiling,
-            config_format=self.config_format,
-            mm_processor_kwargs=mm_processor_kwargs,
-            mm_processor_cache_gb=mm_processor_cache_gb,
-            mm_processor_cache_type=mm_processor_cache_type,
-            mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
-            mm_encoder_only=mm_encoder_only,
-            mm_encoder_tp_mode=mm_encoder_tp_mode,
-            mm_encoder_attn_backend=mm_encoder_attn_backend,
-            pooler_config=self.pooler_config,
-            generation_config=self.generation_config,
-            override_generation_config=self.override_generation_config,
-            enable_sleep_mode=self.enable_sleep_mode,
-            model_impl=self.model_impl,
-            override_attention_dtype=self.override_attention_dtype,
-            logits_processors=self.logits_processors,
-            video_pruning_rate=video_pruning_rate,
-            io_processor_plugin=self.io_processor_plugin,
-            # Omni-specific fields
-            stage_id=self.stage_id,
-            async_chunk=self.async_chunk,
-            model_stage=self.model_stage,
-            model_arch=self.model_arch,
-            worker_type=self.worker_type,
-            engine_output_type=self.engine_output_type,
-            hf_config_name=self.hf_config_name,
-            custom_process_next_stage_input_func=self.custom_process_next_stage_input_func,
-            stage_connector_config=stage_connector_config,
-            omni_kv_config=self.omni_kv_config,
-            task_type=self.task_type,
-        )
-        omni_config.hf_config.architectures = omni_config.architectures
-
-        return omni_config
-
-
-@dataclass
-class AsyncOmniEngineArgs(AsyncEngineArgs):
-    """Async engine arguments for omni models, extending base AsyncEngineArgs.
-    Adds omni-specific configuration fields for multi-stage pipeline
-    processing and output type specification in async contexts.
-    Args:
-        stage_id: Identifier for the stage in a multi-stage pipeline (default: 0)
-        model_stage: Stage type identifier, e.g., "thinker" or "talker"
-            (default: "thinker")
-        model_arch: Model architecture name
-            (default: "Qwen2_5OmniForConditionalGeneration")
-        engine_output_type: Optional output type specification for the engine.
-            Used to route outputs to appropriate processors (e.g., "image",
-            "audio", "latents"). If None, output type is inferred.
-        stage_connector_spec: Extra configuration for stage connector
-        worker_type: Model Type, e.g., "ar" or "generation"
-        task_type: Default task type for TTS models (CustomVoice, VoiceDesign, or Base).
-            If not specified, will be inferred from model path.
-    """
-
-    stage_id: int = 0
-    model_stage: str = "thinker"
-    model_arch: str | None = None
-    engine_output_type: str | None = None
-    hf_config_name: str | None = None
-    custom_process_next_stage_input_func: str | None = None
-    stage_connector_spec: dict[str, Any] = field(default_factory=dict)
-    async_chunk: bool = False
-    omni_kv_config: dict | None = None
-    quantization_config: Any | None = None
-    worker_type: str | None = None
-    task_type: str | None = None
-
-    def __post_init__(self) -> None:
-        load_omni_general_plugins()
-        super().__post_init__()
-
-    def _ensure_omni_models_registered(self):
-        if hasattr(self, "_omni_models_registered"):
-            return True
-        register_omni_models_to_vllm()
-        self._omni_models_registered = True
-        return True
-
-    def create_model_config(self) -> OmniModelConfig:
-        """Create an OmniModelConfig from these engine arguments.
-        Returns:
-            OmniModelConfig instance with all configuration fields set
-        """
-        # GGUF files need a specific model loader path in vLLM.
-        if is_gguf(self.model):
-            self.quantization = self.load_format = "gguf"
-
-        if not envs.VLLM_ENABLE_V1_MULTIPROCESSING:
-            logger.warning(
-                "The global random seed is set to %d. Since "
-                "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may "
-                "affect the random state of the Python process that "
-                "launched vLLM.",
-                self.seed,
-            )
-
-        # register omni models to avoid model not found error
-        self._ensure_omni_models_registered()
-
-        # Keep compatibility when async args are constructed from partial payloads.
-        limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {})
         language_model_only = getattr(self, "language_model_only", False)
+        limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {})
         enable_mm_embeds = getattr(self, "enable_mm_embeds", False)
         interleave_mm_strings = getattr(self, "interleave_mm_strings", False)
         media_io_kwargs = getattr(self, "media_io_kwargs", {})

vllm_omni/engine/input_processor.py

@@ -1,301 +0,0 @@
-import time
-from collections.abc import Mapping
-from typing import Any
-
-import numpy as np
-import torch
-from vllm.config import VllmConfig
-from vllm.inputs import ProcessorInputs, PromptType
-from vllm.inputs.parse import split_enc_dec_inputs
-from vllm.logger import init_logger
-from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry
-from vllm.multimodal.inputs import MultiModalFeatureSpec
-from vllm.multimodal.utils import argsort_mm_positions
-from vllm.platforms import current_platform
-from vllm.pooling_params import PoolingParams
-from vllm.renderers import BaseRenderer
-from vllm.sampling_params import SamplingParams
-from vllm.tasks import SupportedTask
-from vllm.utils import length_from_prompt_token_ids_or_embeds
-from vllm.utils.jsontree import json_iter_leaves
-from vllm.v1.engine.input_processor import InputProcessor
-
-from vllm_omni.engine import (
-    AdditionalInformationEntry,
-    AdditionalInformationPayload,
-    OmniEngineCoreRequest,
-    PromptEmbedsPayload,
-)
-from vllm_omni.inputs.preprocess import OmniInputPreprocessor
-from vllm_omni.lora.request import LoRARequest
-
-logger = init_logger(__name__)
-
-_OMNI_EXTRA_KEYS = (
-    "additional_information",
-    "prompt_embeds",
-    "negative_prompt",
-    "negative_prompt_embeds",
-)
-
-
-def reinject_omni_fields(
-    results: list[ProcessorInputs],
-    original_prompts: list[dict],
-) -> None:
-    """Re-inject omni-specific fields that the upstream renderer discards.
-
-    The upstream renderer's ``process_for_engine`` creates new dicts that only
-    copy standard vLLM fields (prompt_token_ids, multi_modal_data, …).
-    Omni-specific fields such as ``additional_information`` and
-    ``prompt_embeds`` are silently dropped.  This helper copies them back from
-    the *original* parsed prompts into the renderer outputs so they survive
-    into ``OmniInputProcessor.process_inputs()``.
-    """
-    for result, orig in zip(results, original_prompts):
-        if not isinstance(orig, dict):
-            continue
-        for key in _OMNI_EXTRA_KEYS:
-            val = orig.get(key)
-            if val is not None and key not in result:
-                result[key] = val
-
-
-class OmniInputProcessor(InputProcessor):
-    """Processor for omni models, handling multimodal inputs and embeddings.
-
-    Extends the base vLLM Processor with support for processing prompt
-    embeddings and additional information payloads, enabling direct transfer
-    of pre-computed embeddings between pipeline stages.
-
-    Args:
-        vllm_config: Global vLLM configuration
-        mm_registry: Multi-modal registry for processing multimodal inputs
-    """
-
-    @staticmethod
-    def _dtype_to_name(dtype: torch.dtype) -> str:
-        """Convert torch dtype to string representation.
-
-        Args:
-            dtype: PyTorch dtype to convert
-
-        Returns:
-            String representation of the dtype (e.g., "float32", "int64")
-        """
-        mapping = {
-            torch.float32: "float32",
-            torch.float: "float32",
-            torch.float16: "float16",
-            torch.half: "float16",
-            torch.bfloat16: "bfloat16",
-            torch.float64: "float64",
-            torch.double: "float64",
-            torch.int64: "int64",
-            torch.long: "int64",
-            torch.int32: "int32",
-            torch.int: "int32",
-            torch.int16: "int16",
-            torch.short: "int16",
-            torch.int8: "int8",
-            torch.uint8: "uint8",
-            torch.bool: "bool",
-        }
-        return mapping.get(dtype, str(dtype).replace("torch.", ""))
-
-    def __init__(
-        self,
-        vllm_config: VllmConfig,
-        renderer: BaseRenderer | None = None,
-        *,
-        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
-    ):
-        super().__init__(vllm_config, renderer=renderer, mm_registry=mm_registry)
-        self.input_preprocessor = OmniInputPreprocessor(
-            vllm_config=vllm_config,
-            renderer=self.renderer,
-            mm_registry=mm_registry,
-        )
-
-    def process_inputs(
-        self,
-        request_id: str,
-        prompt: PromptType | ProcessorInputs,
-        params: SamplingParams | PoolingParams,
-        supported_tasks: tuple[SupportedTask, ...] = ("generate",),
-        arrival_time: float | None = None,
-        lora_request: LoRARequest | None = None,
-        tokenization_kwargs: dict[str, Any] | None = None,
-        trace_headers: Mapping[str, str] | None = None,
-        priority: int = 0,
-        data_parallel_rank: int | None = None,
-        resumable: bool = False,
-    ) -> OmniEngineCoreRequest:
-        """Process input prompt into an engine core request.
-
-        Converts a prompt (text, tokens, or multimodal) into an
-        OmniEngineCoreRequest that can be processed by the engine.
-        Handles prompt embeddings and additional information payloads
-        for direct transfer between stages.
-
-        Args:
-            request_id: Unique identifier for this request
-            prompt: Input prompt (text, token IDs, embeddings, or multimodal)
-            params: Sampling or pooling parameters for generation
-            supported_tasks: Tuple of supported tasks for validation
-            arrival_time: Optional arrival timestamp (defaults to current time)
-            lora_request: Optional LoRA adapter request
-            tokenization_kwargs: Optional additional tokenization arguments
-            trace_headers: Optional tracing headers for observability
-            priority: Request priority (higher values processed first)
-            data_parallel_rank: Optional data parallel rank for distributed
-                inference
-            resumable: Whether the request supports streaming input
-
-        Returns:
-            OmniEngineCoreRequest ready for the engine
-
-        Raises:
-            ValueError: If data_parallel_rank is out of range or prompt_embeds
-                has incorrect shape
-        """
-        self._validate_params(params, supported_tasks)
-        self._validate_lora(lora_request)
-
-        parallel_config = self.vllm_config.parallel_config
-        dp_size = parallel_config.data_parallel_size
-        dp_local_size = parallel_config.data_parallel_size_local
-        num_ranks = dp_local_size if parallel_config.local_engines_only else dp_size
-        if data_parallel_rank is not None and not (0 <= data_parallel_rank < num_ranks):
-            raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {num_ranks}).")
-
-        # Short-circuit for prompts already processed by the renderer
-        # (they carry a "type" key).  Raw prompts must still go through the
-        # omni preprocessor which preserves additional_information, etc.
-        if isinstance(prompt, dict) and "type" in prompt:
-            if arrival_time is None:
-                arrival_time = prompt.get("arrival_time", time.time())
-            processed_inputs: ProcessorInputs = prompt  # type: ignore[assignment]
-        else:
-            if arrival_time is None:
-                arrival_time = time.time()
-
-            processed_inputs = self.input_preprocessor.preprocess(
-                prompt,
-                tokenization_kwargs=tokenization_kwargs,
-            )
-
-        current_platform.validate_request(processed_inputs, params)
-
-        encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)
-        self._validate_model_inputs(encoder_inputs, decoder_inputs)
-
-        # Normalize decoder prompt access across TypedDict variants.
-        if decoder_inputs["type"] == "embeds":
-            prompt_token_ids = None
-            prompt_embeds = decoder_inputs["prompt_embeds"]
-        else:
-            prompt_token_ids = decoder_inputs["prompt_token_ids"]
-            prompt_embeds = decoder_inputs.get("prompt_embeds")
-
-        sampling_params = None
-        pooling_params = None
-        if isinstance(params, SamplingParams):
-            # TODO: can we avoid cloning here in multiproc case?
-            sampling_params = params.clone()
-            # If unset max tokens, then generate up to the max_model_len.
-            if sampling_params.max_tokens is None:
-                seq_len = length_from_prompt_token_ids_or_embeds(prompt_token_ids, prompt_embeds)
-                sampling_params.max_tokens = self.model_config.max_model_len - seq_len
-            sampling_params.update_from_generation_config(
-                self.generation_config_fields,
-                self.renderer.get_eos_token_id(),
-            )
-            if self.tokenizer is not None:
-                sampling_params.update_from_tokenizer(self.tokenizer)
-        else:
-            pooling_params = params.clone()
-
-        # Multimodal related.
-        mm_features: list[MultiModalFeatureSpec] | None = None
-
-        if decoder_inputs["type"] == "multimodal":
-            decoder_mm_inputs = decoder_inputs["mm_kwargs"]
-            decoder_mm_positions = decoder_inputs["mm_placeholders"]
-            decoder_mm_hashes = decoder_inputs["mm_hashes"]
-
-            if not all(isinstance(leaf, str) for leaf in json_iter_leaves(decoder_mm_hashes)):
-                raise ValueError(
-                    f"mm_hashes must contain only strings, got: {decoder_mm_hashes}. "
-                    "This is likely due to an incorrect custom implementation of "
-                    "MultiModalProcessor.apply method."
-                )
-
-            # Merge and flatten multimodal placeholders, hashes and inputs
-            # from dictionaries to lists, and sort them by each item's position
-            # in the input sequence.
-            sorted_mm_idxs = argsort_mm_positions(decoder_mm_positions)
-
-            mm_features = []
-            for modality, idx in sorted_mm_idxs:
-                base_mm_hash = decoder_mm_hashes[modality][idx]
-                mm_features.append(
-                    MultiModalFeatureSpec(
-                        data=decoder_mm_inputs[modality][idx],
-                        modality=modality,
-                        identifier=self._get_mm_identifier(base_mm_hash, lora_request),
-                        mm_position=decoder_mm_positions[modality][idx],
-                        mm_hash=base_mm_hash,
-                    )
-                )
-
-        # Compatibility: decode serialized prompt embeds if provided.
-        if isinstance(prompt_embeds, PromptEmbedsPayload):
-            prompt_embeds = self._decode_prompt_embeds(prompt_embeds)
-
-        additional_information_payload: AdditionalInformationPayload | None = None
-        raw_info: dict[str, Any] | AdditionalInformationPayload | None = decoder_inputs.get("additional_information")
-        if isinstance(raw_info, AdditionalInformationPayload):
-            additional_information_payload = raw_info
-        elif raw_info is not None:
-            entries: dict[str, AdditionalInformationEntry] = {}
-            for key, value in raw_info.items():
-                if isinstance(value, torch.Tensor):
-                    v_cpu = value.detach().to("cpu").contiguous()
-                    dtype_str = self._dtype_to_name(v_cpu.dtype)
-                    data_bytes = v_cpu.numpy().tobytes()
-                    entry = AdditionalInformationEntry(
-                        tensor_data=data_bytes,
-                        tensor_shape=[int(x) for x in list(v_cpu.shape)],
-                        tensor_dtype=dtype_str,
-                    )
-                elif isinstance(value, list):
-                    entry = AdditionalInformationEntry(list_data=value)
-                else:
-                    raise ValueError("additional_information values must be Tensor or list")
-                entries[key] = entry
-            additional_information_payload = AdditionalInformationPayload(entries=entries)
-
-        return OmniEngineCoreRequest(
-            request_id=request_id,
-            prompt_token_ids=prompt_token_ids,
-            mm_features=mm_features,
-            sampling_params=sampling_params,
-            pooling_params=pooling_params,
-            arrival_time=arrival_time,
-            lora_request=lora_request,
-            cache_salt=decoder_inputs.get("cache_salt"),
-            priority=priority,
-            data_parallel_rank=data_parallel_rank,
-            trace_headers=trace_headers,
-            prompt_embeds=prompt_embeds,
-            additional_information=additional_information_payload,
-            resumable=resumable,
-        )
-
-    @staticmethod
-    def _decode_prompt_embeds(payload: PromptEmbedsPayload) -> torch.Tensor:
-        dtype = getattr(np, payload.dtype)
-        arr = np.frombuffer(payload.data, dtype=dtype)
-        arr = arr.reshape(payload.shape)
-        return torch.from_numpy(arr)

vllm_omni/engine/orchestrator.py

@@ -0,0 +1,812 @@
+"""
+Orchestrator for vLLM-Omni multi-stage runtime.
+
+Runs inside a background thread with its own asyncio event loop.
+Owns all StageEngineCoreClient instances, input/output processors,
+and handles stage-to-stage transfer logic.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import copy
+import time as _time
+from dataclasses import dataclass, field
+from typing import Any
+
+import janus
+import torch
+from vllm.config import ModelConfig
+from vllm.logger import init_logger
+from vllm.outputs import RequestOutput
+from vllm.pooling_params import PoolingParams
+from vllm.sampling_params import SamplingParams
+from vllm.v1.engine import EngineCoreOutputs
+
+from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length
+from vllm_omni.engine import (
+    OmniEngineCoreRequest,
+)
+from vllm_omni.engine.serialization import serialize_additional_information
+from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
+from vllm_omni.metrics.stats import StageStats
+from vllm_omni.metrics.utils import count_tokens_from_outputs
+
+logger = init_logger(__name__)
+
+
+def build_engine_core_request_from_tokens(
+    request_id: str,
+    prompt: dict[str, Any],
+    params: SamplingParams | PoolingParams,
+    arrival_time: float | None = None,
+    model_config: ModelConfig | None = None,
+) -> OmniEngineCoreRequest:
+    """Build an OmniEngineCoreRequest directly from an OmniTokensPrompt.
+
+    Lightweight alternative to the full InputProcessor pipeline - skips
+    tokenization, multimodal preprocessing, LoRA validation, and platform
+    validation.  Intended for stage 1+ where the upstream stage has already
+    produced token IDs and optional embeddings.
+    """
+    if arrival_time is None:
+        arrival_time = _time.time()
+
+    prompt_token_ids = prompt["prompt_token_ids"]
+
+    # Clone params and set max_tokens if needed
+    sampling_params = None
+    pooling_params = None
+    if isinstance(params, SamplingParams):
+        sampling_params = params.clone()
+        if sampling_params.max_tokens is None and model_config is not None:
+            sampling_params.max_tokens = model_config.max_model_len - len(prompt_token_ids)
+    else:
+        pooling_params = params.clone()
+
+    prompt_embeds: torch.Tensor | None = prompt.get("prompt_embeds")
+
+    # Serialize additional_information if present
+    additional_info_payload = serialize_additional_information(
+        prompt.get("additional_information"),
+        log_prefix=f"build_engine_core_request_from_tokens req={request_id}",
+    )
+
+    return OmniEngineCoreRequest(
+        request_id=request_id,
+        prompt_token_ids=prompt_token_ids,
+        mm_features=None,
+        sampling_params=sampling_params,
+        pooling_params=pooling_params,
+        arrival_time=arrival_time,
+        lora_request=None,
+        cache_salt=None,
+        data_parallel_rank=None,
+        prompt_embeds=prompt_embeds,
+        additional_information=additional_info_payload,
+    )
+
+
+# ============================================================
+# Orchestrator internals (run inside the background thread)
+# ============================================================
+
+
+@dataclass
+class OrchestratorRequestState:
+    """Per-request bookkeeping inside the Orchestrator."""
+
+    request_id: str
+    prompt: Any = None
+    sampling_params_list: list[Any] = field(default_factory=list)
+    final_stage_id: int = -1
+
+    # Metrics: timestamp when request was submitted to each stage
+    stage_submit_ts: dict[int, float] = field(default_factory=dict)
+
+
+class Orchestrator:
+    """Runs inside a background thread's asyncio event loop.
+
+    Owns all StageEngineCoreClient instances, input/output processors,
+    and handles stage-to-stage transfer logic.
+    """
+
+    def __init__(
+        self,
+        request_async_queue: janus.AsyncQueue[dict[str, Any]],
+        output_async_queue: janus.AsyncQueue[dict[str, Any]],
+        rpc_async_queue: janus.AsyncQueue[dict[str, Any]],
+        stage_clients: list[Any],
+        output_processors: list[Any],
+        stage_vllm_configs: list[Any],
+        *,
+        async_chunk: bool = False,
+    ) -> None:
+        self.request_async_queue = request_async_queue
+        self.output_async_queue = output_async_queue
+        self.rpc_async_queue = rpc_async_queue
+
+        self.num_stages = len(stage_clients)
+        self.async_chunk = bool(async_chunk)
+
+        self.stage_clients: list[Any] = stage_clients
+        self.output_processors: list[Any] = output_processors
+        self.stage_vllm_configs: list[Any] = stage_vllm_configs
+
+        # Per-request state
+        self.request_states: dict[str, OrchestratorRequestState] = {}
+
+        # CFG companion tracking
+        self._companion_map: dict[str, dict[str, str]] = {}
+        self._companion_to_parent: dict[str, str] = {}
+        self._companion_ids: set[str] = set()
+        self._companion_done: dict[str, set[str]] = {}
+        self._deferred_parents: dict[str, dict[str, Any]] = {}
+
+        # Per-stage metrics accumulators.
+        self._batch_seq: list[int] = [0] * self.num_stages
+        self._agg_total_tokens: list[int] = [0] * self.num_stages
+        self._agg_total_gen_time_ms: list[float] = [0.0] * self.num_stages
+
+        # Shutdown coordination
+        self._shutdown_event = asyncio.Event()
+        self._stages_shutdown = False
+
+    async def run(self) -> None:
+        """Main entry point for the Orchestrator event loop."""
+        logger.info("[Orchestrator] Starting event loop")
+
+        request_task = asyncio.create_task(self._request_handler(), name="orchestrator-request-handler")
+        output_task = asyncio.create_task(
+            self._orchestration_output_handler(),
+            name="orchestrator-stage-output-handler",
+        )
+
+        try:
+            # Run both tasks concurrently; if either fails the other is cancelled.
+            await asyncio.gather(request_task, output_task)
+        except asyncio.CancelledError:
+            raise
+        except Exception:
+            logger.exception("[Orchestrator] Fatal error in orchestrator tasks")
+            raise
+        finally:
+            self._shutdown_event.set()
+            for t in (request_task, output_task):
+                if not t.done():
+                    t.cancel()
+            try:
+                await asyncio.gather(request_task, output_task, return_exceptions=True)
+            except Exception:
+                pass
+
+            self._shutdown_stages()
+
+            # Cancel any remaining tasks spawned by wait_for / gather so
+            # the event loop can close cleanly without "pending task" errors.
+            loop = asyncio.get_running_loop()
+            pending = [t for t in asyncio.all_tasks(loop) if t is not asyncio.current_task() and not t.done()]
+            for t in pending:
+                t.cancel()
+            if pending:
+                await asyncio.gather(*pending, return_exceptions=True)
+
+    async def _request_handler(self) -> None:
+        """Read messages from the main thread via request_async_queue."""
+        while True:
+            msg = await self.request_async_queue.get()
+            msg_type = msg.get("type")
+
+            if msg_type == "add_request":
+                await self._handle_add_request(msg)
+            elif msg_type == "add_companion_request":
+                await self._handle_add_companion(msg)
+            elif msg_type == "abort":
+                await self._handle_abort(msg)
+            elif msg_type == "collective_rpc":
+                await self._handle_collective_rpc(msg)
+            elif msg_type == "shutdown":
+                logger.info("[Orchestrator] Received shutdown signal")
+                self._shutdown_event.set()
+                self._shutdown_stages()
+                break
+            else:
+                logger.warning(f"[Orchestrator] Unknown message type: {msg_type}")
+
+    async def _orchestration_output_handler(self) -> None:
+        """Poll all stages, handle transfers, send final outputs to main."""
+        try:
+            await self._orchestration_loop()
+        except asyncio.CancelledError:
+            logger.debug("[Orchestrator] _orchestration_output_handler cancelled")
+            return
+
+    async def _orchestration_loop(self) -> None:
+        """Inner loop for _orchestration_output_handler (clean cancellation).
+
+        Control flow: poll raw → process through output processor → route.
+        """
+        while not self._shutdown_event.is_set():
+            idle = True
+            for stage_id in range(self.num_stages):
+                if self._shutdown_event.is_set():
+                    return
+
+                # 1) Diffusion stage: poll non-blocking queue
+                # TODO (Peiqi): the output of diffusion stage is OmniRequestOutput,
+                # which is different from EngineCoreOutputs (LLM stages). We may want to unify
+                # the output format in the future to simplify the processing logic in Orchestrator.
+                stage_client = self.stage_clients[stage_id]
+                if stage_client.stage_type == "diffusion":
+                    output = stage_client.get_diffusion_output_async()
+                    if output is not None:
+                        idle = False
+                        req_state = self.request_states.get(output.request_id)
+                        if req_state is not None:
+                            stage_metrics = self._build_stage_metrics(stage_id, output.request_id, [output], req_state)
+                            await self._route_output(stage_id, output, req_state, stage_metrics)
+                    continue
+
+                # 1) Poll raw outputs from the stage
+                try:
+                    raw_outputs = await asyncio.wait_for(self._poll_stage_raw(stage_id), timeout=0.001)
+                except asyncio.TimeoutError:
+                    continue
+                except asyncio.CancelledError:
+                    raise
+                except Exception:
+                    if self._shutdown_event.is_set():
+                        return
+                    logger.exception(
+                        "[Orchestrator] _poll_stage_raw failed for stage-%s",
+                        stage_id,
+                    )
+                    raise
+
+                if raw_outputs is None:
+                    continue
+                idle = False
+
+                # 2) Process raw outputs through the output processor
+                request_outputs = await self._process_stage_outputs(stage_id, raw_outputs)
+
+                # 3) Route each processed output
+                for output in request_outputs:
+                    req_state = self.request_states.get(output.request_id)
+                    if req_state is None:
+                        logger.warning(
+                            "[Orchestrator] Dropping output for unknown req %s at stage-%s (known reqs: %s)",
+                            output.request_id,
+                            stage_id,
+                            list(self.request_states.keys()),
+                        )
+                        continue
+                    stage_metrics = None
+                    if output.finished:
+                        stage_metrics = self._build_stage_metrics(
+                            stage_id,
+                            output.request_id,
+                            [output],
+                            req_state,
+                        )
+                    await self._route_output(stage_id, output, req_state, stage_metrics)
+
+            if idle:
+                await asyncio.sleep(0.001)
+            else:
+                await asyncio.sleep(0)
+
+    async def _route_output(
+        self,
+        stage_id: int,
+        output: Any,
+        req_state: OrchestratorRequestState,
+        stage_metrics: Any,
+    ) -> None:
+        """Route a processed output: send to main thread and/or forward to next stage."""
+        req_id = output.request_id
+        finished = output.finished
+        submit_ts = req_state.stage_submit_ts.get(stage_id)
+        stage_client = self.stage_clients[stage_id]
+
+        # CFG companion handling: companions don't produce user-visible output
+        # and don't forward to the next stage directly.
+        if finished and req_id in self._companion_ids:
+            parent_id = self._companion_to_parent.get(req_id)
+            if parent_id is not None:
+                self._companion_done.setdefault(parent_id, set()).add(req_id)
+                logger.debug(
+                    "[Orchestrator] CFG companion %s done (parent=%s)",
+                    req_id,
+                    parent_id,
+                )
+                # Check if parent is waiting and all companions are done
+                if parent_id in self._deferred_parents and self._all_companions_done(parent_id):
+                    deferred = self._deferred_parents.pop(parent_id)
+                    parent_state = self.request_states.get(parent_id)
+                    if parent_state is not None:
+                        await self._forward_to_next_stage(
+                            parent_id,
+                            deferred["stage_id"],
+                            deferred["output"],
+                            parent_state,
+                        )
+            self.request_states.pop(req_id, None)
+            return
+
+        if stage_client.final_output:
+            await self.output_async_queue.put(
+                {
+                    "type": "output",
+                    "request_id": req_id,
+                    "stage_id": stage_id,
+                    "engine_outputs": output,
+                    "metrics": stage_metrics,
+                    "finished": finished and stage_id == req_state.final_stage_id,
+                    "stage_submit_ts": submit_ts,
+                }
+            )
+        elif stage_metrics is not None:
+            await self.output_async_queue.put(
+                {
+                    "type": "stage_metrics",
+                    "request_id": req_id,
+                    "stage_id": stage_id,
+                    "metrics": stage_metrics,
+                    "stage_submit_ts": submit_ts,
+                }
+            )
+
+        if finished and stage_id < req_state.final_stage_id and not self.async_chunk:
+            # If this parent has CFG companions, defer forwarding until all done
+            if req_id in self._companion_map and not self._all_companions_done(req_id):
+                self._deferred_parents[req_id] = {
+                    "stage_id": stage_id,
+                    "output": output,
+                }
+                logger.debug(
+                    "[Orchestrator] Parent %s deferred, waiting for CFG companions",
+                    req_id,
+                )
+            else:
+                await self._forward_to_next_stage(req_id, stage_id, output, req_state)
+
+        if finished and stage_id == req_state.final_stage_id:
+            self._cleanup_companion_state(req_id)
+            self.request_states.pop(req_id, None)
+
+    def _cleanup_companion_state(self, parent_id: str) -> None:
+        """Remove all companion tracking state for a completed parent."""
+        role_map = self._companion_map.pop(parent_id, {})
+        for cid in role_map.values():
+            self._companion_ids.discard(cid)
+            self._companion_to_parent.pop(cid, None)
+        self._companion_done.pop(parent_id, None)
+        self._deferred_parents.pop(parent_id, None)
+
+    def _all_companions_done(self, parent_id: str) -> bool:
+        """Check whether all CFG companions for a parent request have finished."""
+        role_map = self._companion_map.get(parent_id, {})
+        if not role_map:
+            return True
+        done_set = self._companion_done.get(parent_id, set())
+        return all(cid in done_set for cid in role_map.values())
+
+    def _build_stage_metrics(
+        self,
+        stage_id: int,
+        req_id: str,
+        request_outputs: list[RequestOutput],
+        req_state: OrchestratorRequestState,
+    ) -> StageRequestMetrics:
+        """Build StageRequestMetrics for a finished request at a stage.
+
+        Reuses StageRequestMetrics so OrchestratorMetrics and downstream
+        metric handlers can consume a stable schema.
+        """
+        now = _time.time()
+        submit_ts = req_state.stage_submit_ts.get(stage_id, now)
+        stage_gen_time_ms = (now - submit_ts) * 1000.0
+
+        num_tokens_out = count_tokens_from_outputs(request_outputs)
+        num_tokens_in = 0
+        if stage_id == 0:
+            for ro in request_outputs:
+                ptids = getattr(ro, "prompt_token_ids", None)
+                if ptids is not None:
+                    num_tokens_in += len(ptids)
+
+        # Monotonic batch counter per stage.
+        self._batch_seq[stage_id] += 1
+        batch_id = self._batch_seq[stage_id]
+
+        # Accumulate for running-average stage_stats
+        self._agg_total_tokens[stage_id] += num_tokens_out
+        self._agg_total_gen_time_ms[stage_id] += stage_gen_time_ms
+
+        return StageRequestMetrics(
+            num_tokens_in=num_tokens_in,
+            num_tokens_out=num_tokens_out,
+            stage_gen_time_ms=stage_gen_time_ms,
+            batch_id=batch_id,
+            batch_size=1,
+            rx_decode_time_ms=0.0,
+            rx_transfer_bytes=0,
+            rx_in_flight_time_ms=0.0,
+            stage_stats=StageStats(
+                total_token=self._agg_total_tokens[stage_id],
+                total_gen_time_ms=self._agg_total_gen_time_ms[stage_id],
+            ),
+        )
+
+    async def _forward_to_next_stage(
+        self,
+        req_id: str,
+        stage_id: int,
+        output: Any,
+        req_state: OrchestratorRequestState,
+    ) -> None:
+        """Forward output from current stage to the next stage.
+
+        Handles the full pipeline: set outputs on current stage, compute
+        next-stage inputs, build lightweight requests, and submit them.
+        """
+        next_stage_id = stage_id + 1
+        next_client = self.stage_clients[next_stage_id]
+        params = req_state.sampling_params_list[next_stage_id]
+
+        if next_client.stage_type == "diffusion":
+            self.stage_clients[stage_id].set_engine_outputs([output])
+            if next_client.custom_process_input_func is not None:
+                diffusion_prompt = next_client.custom_process_input_func(
+                    self.stage_clients,
+                    next_client.engine_input_source,
+                    req_state.prompt,
+                    False,
+                )
+                if isinstance(diffusion_prompt, list):
+                    diffusion_prompt = diffusion_prompt[0]
+            else:
+                diffusion_prompt = req_state.prompt
+
+            # Attach CFG companion KV request IDs so the diffusion model
+            # runner can fetch companion KV caches alongside the primary one.
+            cfg_ids = self._companion_map.get(req_id)
+            if cfg_ids:
+                from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+                if isinstance(params, OmniDiffusionSamplingParams):
+                    params = copy.deepcopy(params)
+                    params.cfg_kv_request_ids = cfg_ids
+                    logger.info(
+                        "[Orchestrator] Attaching cfg_kv_request_ids=%s to req %s",
+                        cfg_ids,
+                        req_id,
+                    )
+
+            await next_client.add_request_async(req_id, diffusion_prompt, params)
+            req_state.stage_submit_ts[next_stage_id] = _time.time()
+            return
+
+        self.stage_clients[stage_id].set_engine_outputs([output])
+
+        # Process inputs for next stage
+        try:
+            next_inputs = next_client.process_engine_inputs(
+                stage_list=self.stage_clients,
+                prompt=req_state.prompt,
+            )
+        except Exception:
+            logger.exception(
+                "[Orchestrator] req=%s process_engine_inputs FAILED for stage-%s",
+                req_id,
+                next_stage_id,
+            )
+            raise
+
+        # Build and submit requests for each input
+        for next_input in next_inputs:
+            request = build_engine_core_request_from_tokens(
+                request_id=req_id,
+                prompt=next_input,
+                params=params,
+                model_config=self.stage_vllm_configs[next_stage_id].model_config,
+            )
+
+            # TODO: Here we directly use the req id to assign.
+            request.external_req_id = request.request_id
+
+            self.output_processors[next_stage_id].add_request(
+                request=request,
+                prompt=None,
+                parent_req=None,
+                request_index=0,
+                queue=None,
+            )
+
+            await next_client.add_request_async(request)
+
+        # Record submit timestamp for the next stage
+        req_state.stage_submit_ts[next_stage_id] = _time.time()
+
+    async def _poll_stage_raw(self, stage_id: int) -> EngineCoreOutputs | None:
+        """Pull raw EngineCoreOutputs from a stage client without processing.
+
+        Returns the raw outputs object, or None when there is nothing
+        to consume.
+        """
+        outputs = await self.stage_clients[stage_id].get_output_async()
+        if not outputs.outputs:
+            return None
+        return outputs
+
+    async def _process_stage_outputs(self, stage_id: int, raw_outputs: EngineCoreOutputs) -> list[RequestOutput]:
+        """Run the output processor on raw outputs, returning RequestOutputs.
+
+        Also handles abort forwarding and scheduler stats updates.
+        """
+        processor = self.output_processors[stage_id]
+
+        processed = processor.process_outputs(
+            raw_outputs.outputs,
+            raw_outputs.timestamp,
+            None,
+        )
+
+        if processed.reqs_to_abort:
+            await self.stage_clients[stage_id].abort_requests_async(processed.reqs_to_abort)
+
+        if raw_outputs.scheduler_stats is not None:
+            processor.update_scheduler_stats(raw_outputs.scheduler_stats)
+
+        return processed.request_outputs
+
+    async def _handle_add_request(self, msg: dict[str, Any]) -> None:
+        """Handle an add_request message from the main thread."""
+        stage_id = 0
+        request_id = msg["request_id"]
+        prompt = msg["prompt"]
+        original_prompt = msg.get("original_prompt", prompt)
+        sampling_params_list = msg["sampling_params_list"]
+        if not sampling_params_list:
+            raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.")
+        params = sampling_params_list[0]
+        final_stage_id = msg["final_stage_id"]
+
+        logger.info(
+            "[Orchestrator] _handle_add_request: stage=%s req=%s "
+            "prompt_type=%s original_prompt_type=%s final_stage=%s "
+            "num_sampling_params=%d",
+            stage_id,
+            request_id,
+            type(prompt).__name__,
+            type(original_prompt).__name__,
+            final_stage_id,
+            len(sampling_params_list),
+        )
+
+        # Track request state - use original_prompt so downstream stages
+        # (e.g. thinker2talker) can access the raw dict with multi_modal_data.
+        req_state = OrchestratorRequestState(
+            request_id=request_id,
+            prompt=original_prompt,
+            sampling_params_list=sampling_params_list,
+            final_stage_id=final_stage_id,
+        )
+        req_state.stage_submit_ts[stage_id] = _time.time()
+        self.request_states[request_id] = req_state
+
+        # Stage-0 prompt is already a fully-formed OmniEngineCoreRequest
+        # (pre-processed by AsyncOmniEngine.add_request, output processor
+        # already registered there) - submit directly.
+        request = prompt
+        stage_client = self.stage_clients[stage_id]
+        if stage_client.stage_type == "diffusion":
+            await stage_client.add_request_async(request_id, prompt, params)
+        else:
+            await stage_client.add_request_async(request)
+
+        if self.async_chunk and stage_id == 0 and final_stage_id > 0:
+            await self._prewarm_async_chunk_stages(request_id, request, req_state)
+
+    async def _prewarm_async_chunk_stages(
+        self,
+        request_id: str,
+        stage0_request: Any,
+        req_state: OrchestratorRequestState,
+    ) -> None:
+        """Pre-submit downstream stages for async-chunk mode.
+
+        In async-chunk mode, stages exchange data through connectors/chunk adapters,
+        so downstream stages should be armed once at request start instead of waiting
+        for stage-finished forwarding.
+        """
+        if req_state.final_stage_id <= 0:
+            return
+
+        prompt_token_ids = getattr(stage0_request, "prompt_token_ids", None)
+        if prompt_token_ids is None:
+            logger.warning(
+                "[Orchestrator] async_chunk prewarm skipped for req=%s: stage0 prompt_token_ids missing",
+                request_id,
+            )
+            return
+
+        # Pre-arm stage-1+ with placeholder prompt IDs.
+        try:
+            next_prompt_len = max(1, compute_talker_prompt_ids_length(prompt_token_ids))
+        except Exception:
+            next_prompt_len = max(1, len(prompt_token_ids))
+        original_prompt = req_state.prompt
+        if isinstance(original_prompt, dict):
+            base_input = copy.deepcopy(original_prompt)
+        else:
+            base_input = {}
+        base_input["prompt_token_ids"] = [0] * next_prompt_len
+        base_input["multi_modal_data"] = None
+        base_input["mm_processor_kwargs"] = None
+
+        for next_stage_id in range(1, req_state.final_stage_id + 1):
+            next_client = self.stage_clients[next_stage_id]
+            params = req_state.sampling_params_list[next_stage_id]
+
+            if next_client.stage_type == "diffusion":
+                await next_client.add_request_async(request_id, req_state.prompt, params)
+                req_state.stage_submit_ts[next_stage_id] = _time.time()
+                continue
+
+            request = build_engine_core_request_from_tokens(
+                request_id=request_id,
+                prompt=base_input,
+                params=params,
+                model_config=self.stage_vllm_configs[next_stage_id].model_config,
+            )
+            request.external_req_id = request.request_id
+
+            self.output_processors[next_stage_id].add_request(
+                request=request,
+                prompt=None,
+                parent_req=None,
+                request_index=0,
+                queue=None,
+            )
+            await next_client.add_request_async(request)
+            req_state.stage_submit_ts[next_stage_id] = _time.time()
+
+    async def _handle_add_companion(self, msg: dict[str, Any]) -> None:
+        """Handle an add_companion_request message: submit companion to stage 0."""
+        companion_id = msg["companion_id"]
+        parent_id = msg["parent_id"]
+        role = msg["role"]
+        companion_prompt = msg["prompt"]
+        sampling_params_list = msg["sampling_params_list"]
+
+        # Register companion mapping
+        if parent_id not in self._companion_map:
+            self._companion_map[parent_id] = {}
+        self._companion_map[parent_id][role] = companion_id
+        self._companion_ids.add(companion_id)
+        self._companion_to_parent[companion_id] = parent_id
+        self._companion_done.setdefault(parent_id, set())
+
+        companion_state = OrchestratorRequestState(
+            request_id=companion_id,
+            prompt=companion_prompt,
+            sampling_params_list=sampling_params_list,
+            final_stage_id=0,
+        )
+        companion_state.stage_submit_ts[0] = _time.time()
+        self.request_states[companion_id] = companion_state
+
+        request = companion_prompt  # Already a processed OmniEngineCoreRequest
+        stage_client = self.stage_clients[0]
+        await stage_client.add_request_async(request)
+
+        logger.info(
+            "[Orchestrator] CFG companion submitted: %s (role=%s, parent=%s)",
+            companion_id,
+            role,
+            parent_id,
+        )
+
+    async def _handle_abort(self, msg: dict[str, Any]) -> None:
+        """Handle an abort message from the main thread."""
+        request_ids = msg["request_ids"]
+        # Also abort any CFG companions for aborted parents
+        companion_ids_to_abort: list[str] = []
+        for req_id in request_ids:
+            role_map = self._companion_map.pop(req_id, {})
+            for cid in role_map.values():
+                companion_ids_to_abort.append(cid)
+                self._companion_ids.discard(cid)
+                self._companion_to_parent.pop(cid, None)
+                self.request_states.pop(cid, None)
+            self._companion_done.pop(req_id, None)
+            self._deferred_parents.pop(req_id, None)
+
+        all_ids_to_abort = list(request_ids) + companion_ids_to_abort
+        for stage_id in range(self.num_stages):
+            await self.stage_clients[stage_id].abort_requests_async(all_ids_to_abort)
+        for req_id in request_ids:
+            self.request_states.pop(req_id, None)
+        logger.info("[Orchestrator] Aborted request(s) %s", request_ids)
+
+    async def _handle_collective_rpc(self, msg: dict[str, Any]) -> None:
+        """Handle a control-plane RPC request from the main thread.
+
+        TODO(AsyncOmni): parallelize stage dispatch if control latency becomes
+        noticeable. The current sequential fanout keeps the first version simple
+        and deterministic.
+        """
+        rpc_id = msg["rpc_id"]
+        method = msg["method"]
+        timeout = msg.get("timeout")
+        args = tuple(msg.get("args", ()))
+        kwargs = dict(msg.get("kwargs") or {})
+        requested_stage_ids = msg.get("stage_ids")
+        stage_ids = list(range(self.num_stages)) if requested_stage_ids is None else list(requested_stage_ids)
+
+        results: list[Any] = []
+        for stage_id in stage_ids:
+            if stage_id < 0 or stage_id >= self.num_stages:
+                results.append(
+                    {
+                        "supported": False,
+                        "todo": True,
+                        "error": f"Invalid stage id {stage_id}",
+                    }
+                )
+                continue
+
+            stage_client = self.stage_clients[stage_id]
+            try:
+                if hasattr(stage_client, "collective_rpc_async"):
+                    stage_result = await stage_client.collective_rpc_async(
+                        method=method,
+                        timeout=timeout,
+                        args=args,
+                        kwargs=kwargs,
+                    )
+                else:
+                    stage_result = {
+                        "supported": False,
+                        "todo": True,
+                        "reason": (f"{stage_client.__class__.__name__}.collective_rpc_async is not implemented yet"),
+                    }
+            except Exception as exc:
+                logger.exception(
+                    "[Orchestrator] collective_rpc failed: stage=%s method=%s",
+                    stage_id,
+                    method,
+                )
+                stage_result = {
+                    "supported": False,
+                    "error": str(exc),
+                }
+
+            results.append(stage_result)
+
+        await self.rpc_async_queue.put(
+            {
+                "type": "collective_rpc_result",
+                "rpc_id": rpc_id,
+                "method": method,
+                "stage_ids": stage_ids,
+                "results": results,
+            }
+        )
+
+    def _shutdown_stages(self) -> None:
+        """Shutdown all stage clients."""
+        if self._stages_shutdown:
+            return
+
+        self._stages_shutdown = True
+        logger.info("[Orchestrator] Shutting down all stages")
+        for stage_id, stage_client in enumerate(self.stage_clients):
+            try:
+                stage_client.shutdown()
+                logger.info(f"[Orchestrator] Stage {stage_id} shut down")
+            except Exception as e:
+                logger.warning(f"[Orchestrator] Failed to shutdown stage {stage_id}: {e}")