p

Signed-off-by: khluu <khluu000@gmail.com>

.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
 

.buildkite/test-merge.yml

@@ -12,6 +12,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -29,6 +30,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -46,6 +48,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -64,6 +67,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -81,6 +85,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -98,6 +103,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -116,6 +122,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -134,6 +141,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -152,6 +160,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -178,6 +187,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -197,6 +207,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -226,6 +237,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -261,6 +277,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -272,37 +293,50 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  # - label: "Bagel Text2Img Model Test with H100"
-  #   timeout_in_minutes: 30
-  #   depends_on: upload-merge-pipeline
-  #   commands:
-  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 1
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
+  - label: "Bagel Model Test with H100 (Real Weights)"
+    timeout_in_minutes: 60
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 55m bash -c '
+          set -e
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          export VLLM_IMAGE_FETCH_TIMEOUT=60
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate

.buildkite/test-nightly.yml

@@ -30,6 +30,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -58,6 +63,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -86,6 +92,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -116,6 +127,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -149,6 +161,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:

.buildkite/test-ready.yml

@@ -12,6 +12,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -28,6 +29,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -44,6 +46,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -60,6 +63,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -76,6 +80,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -93,6 +98,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -110,6 +116,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -137,6 +144,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -159,6 +167,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -221,6 +230,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -268,7 +278,7 @@ steps:
         timeout 30m bash -c '
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_TEST_CLEAN_GPU_MEMORY=1
-          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model"
         '
     agents:
       queue: "mithril-h100-pool"
@@ -288,6 +298,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -306,7 +321,7 @@ steps:
         timeout 30m bash -c '
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_TEST_CLEAN_GPU_MEMORY=1
-          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model"
         '
     agents:
       queue: "mithril-h100-pool"
@@ -326,6 +341,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -345,7 +365,7 @@ steps:
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_TEST_CLEAN_GPU_MEMORY=1
           export VLLM_IMAGE_FETCH_TIMEOUT=60
-          pytest -s -v tests/e2e/online_serving/test_bagel_online.py
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model"
         '
     agents:
       queue: "mithril-h100-pool"
@@ -365,6 +385,11 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:

.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.

README.md

@@ -17,7 +17,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 *Latest News* 🔥
 
-- [2026/03] Joined vLLM HK First Meetup! Video on [youtube](https://youtu.be/sgwNfsNnR9I)
+- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
 - [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
 - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.

benchmarks/diffusion/backends.py

@@ -35,6 +35,7 @@ class RequestFuncOutput:
     error: str = ""
     start_time: float = 0.0
     response_body: dict[str, Any] = field(default_factory=dict)
+    stage_durations: dict[str, float] = field(default_factory=dict)
     peak_memory_mb: float = 0.0
     slo_achieved: bool | None = None
 
@@ -55,6 +56,7 @@ async def async_request_chat_completions(
     input: RequestFuncInput,
     session: aiohttp.ClientSession,
     pbar: tqdm | None = None,
+    enable_diffusion_pipeline_profiler: bool = False,
 ) -> RequestFuncOutput:
     output = RequestFuncOutput()
     output.start_time = time.perf_counter()
@@ -108,6 +110,18 @@ async def async_request_chat_completions(
                 output.success = True
                 if "peak_memory_mb" in resp_json:
                     output.peak_memory_mb = resp_json["peak_memory_mb"]
+                try:
+                    choices = resp_json.get("choices", [])
+                    if choices and isinstance(choices, list):
+                        msg = choices[0].get("message", {})
+                        if isinstance(msg, dict):
+                            content = msg.get("content", [])
+                            if content and isinstance(content, list) and len(content) > 0:
+                                first_item = content[0]
+                                if isinstance(first_item, dict):
+                                    output.stage_durations = first_item.get("stage_durations")
+                except (IndexError, TypeError, AttributeError):
+                    pass
             else:
                 output.error = f"HTTP {response.status}: {await response.text()}"
                 output.success = False

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

docker/Dockerfile.npu

@@ -1,17 +1,17 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.14.0rc1
+ARG VLLM_ASCEND_TAG=v0.17.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm-ascend
-RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
-RUN pip install -v -e .
-
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN pip install -v -e .
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docker/Dockerfile.npu.a3

@@ -1,17 +1,17 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
+ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-WORKDIR /vllm-workspace/vllm-ascend
-RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
-RUN pip install -v -e .
-
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN pip install -v -e .
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

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/README.md

@@ -14,6 +14,10 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 
 - **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
 
+## Multi-Stage Recipes
+
+- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
+
 ## Optimization Features
 
 - **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss

docs/configuration/pd_disaggregation.md

@@ -0,0 +1,176 @@
+# Prefill-Decode (PD) Disaggregation
+
+PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
+stages so prompt processing and token generation can run on different workers.
+
+This is documented as a stage-config recipe instead of a bundled YAML because the
+deployment-specific values usually change per environment:
+
+- GPU placement
+- `tensor_parallel_size`
+- connector backend and connector ports
+- connector IPs or bootstrap addresses
+
+Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml)
+and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
+changes below.
+
+## Requirements
+
+- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
+- A KV connector supported by vLLM, such as `MooncakeConnector`
+- Matching `tensor_parallel_size` on the prefill and decode thinker stages
+
+## 1. Split the thinker into prefill and decode stages
+
+Replace the original thinker stage with two stages:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    stage_type: llm
+    is_prefill_only: true
+    runtime:
+      devices: "0"
+      max_batch_size: 16
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_producer"
+        kv_rank: 0
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25201
+    final_output: false
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+
+  - stage_id: 1
+    stage_type: llm
+    is_decode_only: true
+    runtime:
+      devices: "1"
+      max_batch_size: 64
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_type: ar
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.9
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      max_num_batched_tokens: 32768
+      hf_config_name: thinker_config
+      tensor_parallel_size: 1
+      kv_transfer_config:
+        kv_connector: "MooncakeConnector"
+        kv_role: "kv_consumer"
+        kv_rank: 1
+        kv_parallel_size: 2
+        kv_connector_extra_config:
+          mooncake_bootstrap_port: 25202
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 2048
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05
+```
+
+Notes:
+
+- `is_prefill_only: true` marks the thinker stage that only saves KV.
+- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
+- `kv_transfer_config` is required on both stages.
+- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
+  prefill side only processes the prompt and exports KV.
+
+## 2. Shift the downstream stages by one index
+
+After inserting the extra thinker stage, renumber the remaining stages:
+
+```yaml
+  - stage_id: 2
+    runtime:
+      devices: "2"
+    engine_input_source: [1]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
+
+  - stage_id: 3
+    runtime:
+      devices: "2"
+      max_batch_size: 1
+    engine_input_source: [2]
+    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
+```
+
+Compared with the default Qwen3-Omni config:
+
+- the talker becomes stage `2` instead of stage `1`
+- the code2wav stage becomes stage `3` instead of stage `2`
+- the talker now reads from decode stage `1`
+
+## 3. Add runtime edges for the four-stage pipeline
+
+```yaml
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1
+    - from: 1
+      to: 2
+      window_size: -1
+    - from: 2
+      to: 3
+      window_size: -1
+```
+
+## 4. Launch with your custom config
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+    --stage-configs-path /path/to/qwen3_omni_pd.yaml
+```
+
+## Operational Notes
+
+- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
+  pair. Keep prefill and decode at the same `tensor_parallel_size`.
+- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
+  separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
+- Choose connector ports and addresses that match your deployment. The values
+  shown above are examples only.

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_diffusion_model.md

@@ -815,6 +815,107 @@ class WanTransformer3DModel(nn.Module):
 
 ---
 
+### Diffusion Timing (Performance Profiling)
+When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
+vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
+
+!!! info
+      `DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
+
+This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
+
+**Enabling Diffusion Timing**
+
+
+Enable timing by setting:
+```
+vllm serve Qwen/Qwen-Image --omni --port 8091 --enable-diffusion-pipeline-profiler
+```
+You can optionally specify which modules to profile:
+```
+class YourPipeline(xxx, DiffusionPipelineProfilerMixin):
+    def __init__(self, xxx):
+        ...
+        self.setup_diffusion_pipeline_profiler(profiler_targets=["diffuse"], enable_diffusion_pipeline_profiler)
+```
+If not specified, the default targets are used:
+```
+["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
+```
+**Adding DiffusionPipelineProfilerMixin to a Pipeline**
+To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
+```python
+from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
+
+class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
+    # Optional: Specify custom timing targets
+    _PROFILER_TARGETS = ["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
+
+    def __init__(
+        self,
+        *,
+        od_config: OmniDiffusionConfig,
+        prefix: str = "",
+    ):
+        super().__init__()
+        self.od_config = od_config
+        self.parallel_config = od_config.parallel_config
+        # initialize pipeline components
+        ...
+
+        # initialize timing profiler
+        self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
+```
+The mixin dynamically wraps selected methods and records their execution time during inference.
+
+If you need to fetch the execution time of different modules, you will need to pass `self.stage_durations` to `DiffusionOutput`, as shown below:
+
+```diff
+-   return DiffusionOutput(output=img)
++   return DiffusionOutput(
+            output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
+        )
+```
+
+**Pipeline Design for Timing**
+The current diffusion timing utility is function-based, meaning it measures the execution time of individual methods.
+
+When implementing a new pipeline, avoid putting all logic inside a single function (e.g., forward). Instead, structure the pipeline in a modular way by separating key stages into independent methods, such as the diffusion loop.
+
+For example:
+```
+def forward(self, req: OmniDiffusionRequest):
+    prompt_embeds = self.encode_prompt(req)
+    latents = self.diffuse(prompt_embeds, req)
+    images = self.vae.decode(latents)
+    return DiffusionOutput(output=images)
+```
+This allows the timing utility to measure each stage (e.g., encode_prompt, diffuse, vae.decode) separately and helps identify performance bottlenecks more easily.
+
+
+**Default Profiled Modules**
+
+By default, the following pipeline modules are timed:
+```
+vae.encode
+vae.decode
+diffuse
+text_encoder.forward
+tokenizer.forward
+```
+
+**Example Output**
+
+When enabled, timing logs appear like this:
+```
+[DiffusionTiming] text_encoder.forward took 0.018s
+[DiffusionTiming] diffuse took 2.412s
+[DiffusionTiming] vae.decode took 0.063s
+```
+These measurements help identify bottlenecks during model adaptation and optimization
+
+
+
 ## Troubleshooting
 
 

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/installation/npu/npu.inc.md

@@ -33,18 +33,25 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
+cd /vllm-workspace/vllm
+git pull origin main
+git fetch origin --tags
+git checkout v0.16.0
+
 # Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
 # we have to pin vllm-ascend at the current commit.
 cd /vllm-workspace/vllm-ascend
+git pull origin main
 git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
 pip install -v -e .
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
 git clone -b v0.16.0 https://github.com/vllm-project/vllm-omni.git
-
 cd vllm-omni
-pip install -v -e .
+pip install -v -e . --no-build-isolation
+# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
+
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```
 
@@ -61,22 +68,22 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
 
 ```bash
-# Pin vLLM version to 0.16.0
+# Pin vLLM version to 0.17.0
 cd /vllm-workspace/vllm
 git pull origin main
 git fetch origin --tags
-git checkout v0.16.0
+git checkout v0.17.0
 VLLM_TARGET_DEVICE=empty pip install -v -e .
 
 # Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
-cd ../vllm-ascend
+cd /vllm-workspace/vllm-ascend
 git pull origin main
 git fetch origin --tags
-git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+git checkout v0.17.0
 pip install -v -e .
 
 # Install vLLM-Omni from the latest main branch
-cd ../vllm-omni
+cd /vllm-workspace/vllm-omni
 git clone https://github.com/vllm-project/vllm-omni.git
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .

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)
@@ -58,6 +58,11 @@ def parse_args():
         help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
     )
     parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     args = parser.parse_args()
     return args
@@ -103,6 +108,7 @@ def main():
             "shm_threshold_bytes": args.shm_threshold_bytes,
             "worker_backend": args.worker_backend,
             "ray_address": args.ray_address,
+            "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
         }
     )
 
@@ -163,25 +169,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

@@ -247,6 +247,7 @@ def main(args: argparse.Namespace) -> None:
         stage_configs_path=config_path,
         log_stats=args.enable_stats,
         stage_init_timeout=args.stage_init_timeout,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     init_time = time.time() - start_time
@@ -310,36 +311,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")
@@ -464,6 +464,11 @@ def parse_args() -> argparse.Namespace:
         default=False,
         help="Enable verbose output",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     return parser.parse_args()
 

examples/offline_inference/helios/end2end.py

@@ -173,6 +173,11 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of initial denoising steps using zero prediction (default: 1).",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     # Memory & parallelism
     parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
@@ -207,6 +212,7 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     # Validate I2V / V2V arguments
@@ -293,12 +299,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:
@@ -41,6 +41,11 @@ def parse_args() -> argparse.Namespace:
         required=True,
         help="Pretrain template prompt: <|startoftext|>{system}<img>{question}",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -52,7 +57,10 @@ def load_image(image_path: str) -> Image.Image:
 
 
 def main(args: argparse.Namespace) -> None:
-    omni = Omni(model=args.model)
+    omni = Omni(
+        model=args.model,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+    )
 
     prompt_dict = {
         "prompt": args.prompt,
@@ -70,8 +78,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

@@ -325,6 +325,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -389,6 +394,7 @@ def main():
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
     print("Pipeline loaded")
 
@@ -463,12 +469,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

@@ -176,6 +176,11 @@ def parse_args() -> argparse.Namespace:
             "Default 1 means pure sharding (no replication). "
         ),
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -281,6 +286,7 @@ def main():
         model_class_name=model_class_name,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     if profiler_enabled:
@@ -338,8 +344,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

@@ -43,6 +43,11 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--top-p", type=float, default=0.9)
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--trust-remote-code", action="store_true")
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -70,6 +75,7 @@ def main() -> None:
         model=args.model,
         stage_configs_path=args.stage_config,
         trust_remote_code=args.trust_remote_code,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
     try:
         sp = SamplingParams(
@@ -99,14 +105,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,15 +180,16 @@ 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,
         shm_threshold_bytes=args.shm_threshold_bytes,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     thinker_sampling_params = SamplingParams(
@@ -290,7 +291,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 +299,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():
@@ -434,6 +434,11 @@ def parse_args():
         default="../../../model_executor/stage_configs/mimo_audio.yaml",
         help="Path to a stage configs file.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     return parser.parse_args()
 

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -320,13 +320,14 @@ 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,
         batch_timeout=args.batch_timeout,
         init_timeout=args.init_timeout,
         shm_threshold_bytes=args.shm_threshold_bytes,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
     thinker_sampling_params = SamplingParams(
         temperature=0.0,  # Deterministic - no randomness
@@ -378,9 +379,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 +392,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():
@@ -539,6 +541,11 @@ def parse_args():
         default=False,
         help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 

examples/offline_inference/qwen3_omni/end2end.py

@@ -294,11 +294,12 @@ 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,
         stage_init_timeout=args.stage_init_timeout,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     thinker_sampling_params = SamplingParams(
@@ -354,8 +355,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 +367,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():
@@ -527,6 +527,11 @@ def parse_args():
         default=False,
         help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     return parser.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()
@@ -387,6 +382,7 @@ async def run_all(args):
             stage_configs_path=args.stage_configs_path,
             log_stats=args.log_stats,
             stage_init_timeout=args.stage_init_timeout,
+            enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         )
 
         # Use default sampling params from stage config (they are pre-configured
@@ -589,6 +585,11 @@ def parse_args():
         default=16000,
         help="Sampling rate for audio loading.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 

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)
@@ -319,14 +371,15 @@ def main(args):
         stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     batch_size = args.batch_size
     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):
@@ -340,6 +393,7 @@ async def main_streaming(args):
         stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     for i, prompt in enumerate(inputs):
@@ -481,6 +535,11 @@ def parse_args():
         default=1,
         help="Number of prompts per batch (default: 1, sequential).",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
 
     return parser.parse_args()
 

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -90,6 +90,11 @@ def parse_args() -> argparse.Namespace:
         default=44100,
         help="Sample rate for output audio (Stable Audio uses 44100 Hz).",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -133,7 +138,10 @@ def main():
     print(f"{'=' * 60}\n")
 
     # Initialize Omni with Stable Audio model
-    omni = Omni(model=args.model)
+    omni = Omni(
+        model=args.model,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+    )
 
     # Calculate audio end time
     audio_end_in_s = args.audio_start + args.audio_length
@@ -177,7 +185,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
 
 
@@ -218,6 +217,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="[NextStep-1.1 only] Apply layer normalization to sampled tokens.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -305,6 +309,7 @@ def main():
         "parallel_config": parallel_config,
         "enforce_eager": args.enforce_eager,
         "enable_cpu_offload": args.enable_cpu_offload,
+        "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
         **lora_args,
         **quant_kwargs,
     }
@@ -400,20 +405,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

@@ -132,6 +132,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable expert parallelism for MoE layers.",
     )
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler to display stage durations.",
+    )
     return parser.parse_args()
 
 
@@ -184,6 +189,7 @@ def main():
         enforce_eager=args.enforce_eager,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     if profiler_enabled:
@@ -240,8 +246,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

@@ -1236,15 +1236,14 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st
         port = params.port
         stage_config_path = params.stage_config_path
         if run_level == "advanced_model" and stage_config_path is not None:
+            # Dynamically detect stages from config to avoid KeyError
+            # for models with fewer stages (e.g., BAGEL has 2, Qwen3-Omni has 3)
+            with open(stage_config_path, encoding="utf-8") as f:
+                _cfg = yaml.safe_load(f) or {}
+            _stage_ids = [s["stage_id"] for s in _cfg.get("stage_args", []) if "stage_id" in s]
             stage_config_path = modify_stage_config(
                 stage_config_path,
-                deletes={
-                    "stage_args": {
-                        0: ["engine_args.load_format"],
-                        1: ["engine_args.load_format"],
-                        2: ["engine_args.load_format"],
-                    }
-                },
+                deletes={"stage_args": {sid: ["engine_args.load_format"] for sid in _stage_ids}},
             )
 
         server_args = params.server_args or []
@@ -1733,7 +1732,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,
@@ -1974,9 +1975,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/bagel_mooncake_ci.yaml

@@ -21,6 +21,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       tensor_parallel_size: 1
+      load_format: dummy
       omni_kv_config:
         need_send_cache: true
         kv_transfer_criteria:
@@ -54,6 +55,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       tensor_parallel_size: 1
+      load_format: dummy
       omni_kv_config:
         need_recv_cache: true
     engine_input_source: [0]

tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml

@@ -21,6 +21,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       tensor_parallel_size: 1
+      load_format: dummy
       omni_kv_config:
         need_send_cache: true
         kv_transfer_criteria:
@@ -53,6 +54,7 @@ stage_args:
       enable_prefix_caching: false
       max_num_batched_tokens: 32768
       tensor_parallel_size: 1
+      load_format: dummy
       omni_kv_config:
         need_recv_cache: true
     engine_input_source: [0]

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

@@ -22,6 +22,7 @@ import pytest
 from PIL import Image
 from vllm.assets.image import ImageAsset
 
+from tests.conftest import modify_stage_config
 from tests.utils import hardware_test
 from vllm_omni.entrypoints.omni import Omni
 
@@ -98,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
 
 
@@ -168,17 +169,39 @@ def _generate_bagel_img2img(
     return generated_image
 
 
+def _resolve_stage_config(config_path: str, run_level: str) -> str:
+    """Resolve stage config based on run level.
+
+    For advanced_model (real weights), strip load_format: dummy so the model
+    falls back to loading real weights from HuggingFace.
+    """
+    if run_level == "advanced_model":
+        return modify_stage_config(
+            config_path,
+            deletes={
+                "stage_args": {
+                    0: ["engine_args.load_format"],
+                    1: ["engine_args.load_format"],
+                }
+            },
+        )
+    return config_path
+
+
 @pytest.mark.core_model
+@pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"})
-def test_bagel_img2img_shared_memory_connector():
+def test_bagel_img2img_shared_memory_connector(run_level):
     """Test Bagel img2img with shared memory connector."""
     input_image = _load_input_image()
     config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
+    config_path = _resolve_stage_config(config_path, run_level)
     omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300)
 
     try:
         generated_image = _generate_bagel_img2img(omni, input_image)
-        _validate_pixels(generated_image)
+        if run_level == "advanced_model":
+            _validate_pixels(generated_image)
     finally:
         omni.close()

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
@@ -25,6 +28,7 @@ from typing import Any
 import pytest
 from PIL import Image
 
+from tests.conftest import modify_stage_config
 from tests.utils import hardware_test
 from vllm_omni.entrypoints.omni import Omni
 
@@ -96,9 +100,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
 
 
@@ -155,17 +159,39 @@ def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Ima
     return generated_image
 
 
+def _resolve_stage_config(config_path: str, run_level: str) -> str:
+    """Resolve stage config based on run level.
+
+    For advanced_model (real weights), strip load_format: dummy so the model
+    falls back to loading real weights from HuggingFace.
+    """
+    if run_level == "advanced_model":
+        return modify_stage_config(
+            config_path,
+            deletes={
+                "stage_args": {
+                    0: ["engine_args.load_format"],
+                    1: ["engine_args.load_format"],
+                }
+            },
+        )
+    return config_path
+
+
 @pytest.mark.core_model
+@pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"})
-def test_bagel_text2img_shared_memory_connector():
+def test_bagel_text2img_shared_memory_connector(run_level):
     """Test Bagel text2img with shared memory connector."""
     config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
+    config_path = _resolve_stage_config(config_path, run_level)
     omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300)
 
     try:
         generated_image = _generate_bagel_image(omni)
-        _validate_pixels(generated_image)
+        if run_level == "advanced_model":
+            _validate_pixels(generated_image)
     finally:
         omni.close()
 
@@ -248,9 +274,10 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
 
 
 @pytest.mark.core_model
+@pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"})
-def test_bagel_text2img_mooncake_connector():
+def test_bagel_text2img_mooncake_connector(run_level):
     """Test Bagel text2img with Mooncake connector for inter-stage communication."""
     MOONCAKE_HOST = "127.0.0.1"
     MOONCAKE_RPC_PORT = _find_free_port()
@@ -288,10 +315,12 @@ def test_bagel_text2img_mooncake_connector():
             http_port=MOONCAKE_HTTP_PORT,
         )
 
+        temp_config_file = _resolve_stage_config(temp_config_file, run_level)
         omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file, stage_init_timeout=300)
 
         generated_image = _generate_bagel_image(omni)
-        _validate_pixels(generated_image)
+        if run_level == "advanced_model":
+            _validate_pixels(generated_image)
 
     finally:
         if omni:

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

@@ -36,7 +36,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_bagel_expansion.py

@@ -0,0 +1,132 @@
+"""
+Comprehensive tests of diffusion features that are available in online serving mode
+and are supported by the following models:
+- ByteDance-Seed/BAGEL-7B-MoT
+Coverage:
+- TeaCache
+- Cache-DiT
+- CFG-Parallel
+- Tensor-Parallel
+
+assert_diffusion_response validates successful generation and the expected
+512x512 resolution.
+"""
+
+import pytest
+
+from tests.conftest import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+    dummy_messages_from_mix_data,
+)
+from tests.utils import hardware_marks
+
+PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution."
+NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark"
+
+SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
+PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+
+def _get_diffusion_feature_cases(model: str):
+    """Return L4 diffusion feature cases for Bagel.
+    TeaCache, Cache-DiT, CFG-Parallel, Tensor-Parallel.
+    """
+
+    return [
+        # TeaCache (single-card)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "tea_cache",
+                ],
+            ),
+            id="single_card_teacache",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+        # Cache-DiT (single-card)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                ],
+            ),
+            id="single_card_cache_dit",
+            marks=SINGLE_CARD_FEATURE_MARKS,
+        ),
+        # CFG-Parallel size 2 (2 GPUs, TeaCache backend)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "tea_cache",
+                    "--cfg-parallel-size",
+                    "2",
+                ],
+            ),
+            id="parallel_cfg_2",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+        # Tensor-Parallel size 2 (2 GPUs, Cache-DiT backend)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                server_args=[
+                    "--cache-backend",
+                    "cache_dit",
+                    "--tensor-parallel-size",
+                    "2",
+                ],
+            ),
+            id="parallel_tp_2",
+            marks=PARALLEL_FEATURE_MARKS,
+        ),
+    ]
+
+
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases("ByteDance-Seed/BAGEL-7B-MoT"),
+    indirect=True,
+)
+def test_bagel(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+):
+    """L4 diffusion feature coverage for Bagel on H100.
+
+    This test exercises:
+    - TeaCache
+    - Cache-DiT
+    - CFG-Parallel (size=2)
+    - Tensor-Parallel (size=2)
+
+    Validation is delegated to assert_diffusion_response in tests.conftest,
+    which checks output dimensions and basic correctness.
+    """
+
+    messages = dummy_messages_from_mix_data(content_text=PROMPT)
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 512,
+            "width": 512,
+            "num_inference_steps": 2,
+            # Enable CFG for models that use classifier-free guidance
+            "negative_prompt": NEGATIVE_PROMPT,
+            "true_cfg_scale": 4.0,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)

tests/e2e/online_serving/test_bagel_online.py

@@ -22,22 +22,18 @@ Equivalent to running:
 
 import base64
 import os
-import signal
-import socket
-import subprocess
-import sys
-import time
 from io import BytesIO
 from pathlib import Path
-from typing import Any
 
 import pytest
-import requests
-from PIL import Image
 from vllm.assets.image import ImageAsset
 
+from tests.conftest import OmniServerParams
 from tests.utils import hardware_test
 
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
 MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
 STAGE_CONFIGS_PATH = str(
     Path(__file__).parent.parent / "offline_inference" / "stage_configs" / "bagel_sharedmemory_ci.yaml"
@@ -46,202 +42,74 @@ STAGE_CONFIGS_PATH = str(
 TEXT2IMG_PROMPT = "A cute cat"
 IMG2IMG_PROMPT = "Change the grass color to red"
 
-
-class BagelOmniServer:
-    """Context manager to start/stop a vLLM-Omni server for Bagel model tests."""
-
-    def __init__(
-        self,
-        model: str = MODEL,
-        stage_configs_path: str = STAGE_CONFIGS_PATH,
-        env_dict: dict[str, str] | None = None,
-    ) -> None:
-        self.model = model
-        self.stage_configs_path = stage_configs_path
-        self.env_dict = env_dict
-        self.proc: subprocess.Popen | None = None
-        self.host = "127.0.0.1"
-        self.port = _find_free_port()
-
-    @property
-    def base_url(self) -> str:
-        return f"http://{self.host}:{self.port}"
-
-    def _start_server(self) -> None:
-        env = os.environ.copy()
-        if self.env_dict is not None:
-            env.update(self.env_dict)
-
-        cmd = [
-            sys.executable,
-            "-m",
-            "vllm_omni.entrypoints.cli.main",
-            "serve",
-            self.model,
-            "--omni",
-            "--host",
-            self.host,
-            "--port",
-            str(self.port),
-            "--stage-configs-path",
-            self.stage_configs_path,
-            "--stage-init-timeout",
-            "300",
-        ]
-
-        self.proc = subprocess.Popen(
-            cmd,
-            env=env,
-            start_new_session=True,
-        )
-
-        try:
-            if not _wait_for_port(self.host, self.port, timeout=600, proc=self.proc):
-                self.terminate()
-                raise RuntimeError(f"Server failed to start within 600 seconds on {self.host}:{self.port}")
-        except Exception:
-            self.terminate()
-            raise
-
-    def __enter__(self):
-        self._start_server()
-        return self
-
-    def terminate(self) -> None:
-        if self.proc:
-            try:
-                os.killpg(os.getpgid(self.proc.pid), signal.SIGTERM)
-            except ProcessLookupError:
-                pass
-            try:
-                self.proc.wait(timeout=30)
-            except subprocess.TimeoutExpired:
-                try:
-                    os.killpg(os.getpgid(self.proc.pid), signal.SIGKILL)
-                except ProcessLookupError:
-                    pass
-                self.proc.wait()
-            self.proc = None
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.terminate()
+# Create parameter combinations for model and stage config
+test_params = [
+    OmniServerParams(
+        model=MODEL,
+        stage_config_path=STAGE_CONFIGS_PATH,
+        server_args=["--stage-init-timeout", "300"],
+    ),
+]
 
 
-def _find_free_port() -> int:
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-        s.bind(("127.0.0.1", 0))
-        s.listen(1)
-        return s.getsockname()[1]
+def _build_text2img_messages(prompt: str) -> list[dict]:
+    """Build OpenAI-format messages for text2img generation."""
+    return [
+        {
+            "role": "user",
+            "content": [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}],
+        }
+    ]
 
 
-def _wait_for_port(host: str, port: int, timeout: int = 600, proc: subprocess.Popen | None = None) -> bool:
-    start = time.time()
-    while time.time() - start < timeout:
-        if proc is not None and proc.poll() is not None:
-            # Server process exited early
-            return False
-        try:
-            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-                sock.settimeout(1)
-                if sock.connect_ex((host, port)) == 0:
-                    return True
-        except Exception:
-            pass
-        time.sleep(2)
-    return False
+def _build_img2img_messages(prompt: str, image_b64: str) -> list[dict]:
+    """Build OpenAI-format messages for img2img generation."""
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"},
+                {
+                    "type": "image_url",
+                    "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"},
+                },
+            ],
+        }
+    ]
 
 
-def _send_chat_request(
-    server_url: str,
-    prompt: str,
-    *,
-    modality: str = "text2img",
-    image: Image.Image | None = None,
-    timeout: int = 300,
-) -> dict[str, Any]:
-    """Send a chat completion request matching the openai_chat_client.py format."""
-    content: list[dict[str, Any]] = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
-
-    if image is not None:
-        buffer = BytesIO()
-        image.save(buffer, format="JPEG")
-        b64_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
-        content.append(
-            {
-                "type": "image_url",
-                "image_url": {"url": f"data:image/jpeg;base64,{b64_data}"},
-            }
-        )
-
-    payload: dict[str, Any] = {
-        "messages": [{"role": "user", "content": content}],
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100"})
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_bagel_text2img_online(omni_server, openai_client) -> None:
+    """Test Bagel text2img via OpenAI-compatible chat completions API."""
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_text2img_messages(TEXT2IMG_PROMPT),
+        "modalities": ["image"],
     }
 
-    if modality in ("text2img", "img2img"):
-        payload["modalities"] = ["image"]
-
-    resp = requests.post(
-        f"{server_url}/v1/chat/completions",
-        headers={"Content-Type": "application/json"},
-        json=payload,
-        timeout=timeout,
-    )
-    resp.raise_for_status()
-    return resp.json()
-
-
-def _extract_image_from_response(data: dict[str, Any]) -> Image.Image | None:
-    """Extract the generated PIL Image from a chat completion response."""
-    for choice in data.get("choices", []):
-        content = choice.get("message", {}).get("content")
-        if isinstance(content, list) and content:
-            first_item = content[0]
-            if isinstance(first_item, dict) and "image_url" in first_item:
-                url = first_item["image_url"].get("url", "")
-                if url.startswith("data:image"):
-                    _, b64 = url.split(",", 1)
-                    return Image.open(BytesIO(base64.b64decode(b64)))
-    return None
+    openai_client.send_diffusion_request(request_config)
 
 
 @pytest.mark.core_model
+@pytest.mark.advanced_model
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100"})
-def test_bagel_text2img_online():
-    """Test Bagel text2img via OpenAI-compatible chat completions API."""
-    with BagelOmniServer() as server:
-        response_data = _send_chat_request(
-            server.base_url,
-            TEXT2IMG_PROMPT,
-            modality="text2img",
-        )
-
-        image = _extract_image_from_response(response_data)
-        assert image is not None, f"No image in response: {response_data}"
-        image.load()
-
-        w, h = image.size
-        assert w > 0 and h > 0, f"Invalid image size: {image.size}"
-
-
-@pytest.mark.core_model
-@pytest.mark.diffusion
-@hardware_test(res={"cuda": "H100"})
-def test_bagel_img2img_online():
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_bagel_img2img_online(omni_server, openai_client) -> None:
     """Test Bagel img2img via OpenAI-compatible chat completions API."""
     input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
+    buffer = BytesIO()
+    input_image.save(buffer, format="JPEG")
+    image_b64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
 
-    with BagelOmniServer() as server:
-        response_data = _send_chat_request(
-            server.base_url,
-            IMG2IMG_PROMPT,
-            modality="img2img",
-            image=input_image,
-        )
+    request_config = {
+        "model": omni_server.model,
+        "messages": _build_img2img_messages(IMG2IMG_PROMPT, image_b64),
+        "modalities": ["image"],
+    }
 
-        image = _extract_image_from_response(response_data)
-        assert image is not None, f"No image in response: {response_data}"
-        image.load()
-
-        w, h = image.size
-        assert w > 0 and h > 0, f"Invalid image size: {image.size}"
+    openai_client.send_diffusion_request(request_config)

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": [
             {