Revert "[Doc][Frontend][Model][VoxCPM2] Support instructions and per-request …"

This reverts commit 1069a144a8.

.buildkite/pipeline.yml

@@ -1,8 +1,9 @@
 # Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
-# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs if nightly-test
-# / main NIGHTLY so upload-nightly is not skipped together with test-ready/test-merge.
-#
-# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite).
+# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly
+# exception in upload_pipeline_with_skip_ci.sh (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
+# can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build.
+
+# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
 steps:
   - label: ":github: Resolve skip-ci & upload pipeline"
     key: upload-ci-pipeline
@@ -61,3 +62,15 @@ steps:
       - buildkite-agent pipeline upload .buildkite/test-nightly.yml
     agents:
       queue: "cpu_queue_premerge"
+
+  # L5 Test — main+WEEKLY=1 (scheduled), or PR with weekly-test label
+  - label: "Upload Weekly Pipeline"
+    depends_on: image-build
+    key: upload-weekly-pipeline
+    if: >-
+      (build.branch == "main" && build.env("WEEKLY") == "1") ||
+      (build.branch != "main" && build.pull_request.labels includes "weekly-test")
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-weekly.yml
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/scripts/hardware_ci/run_npu_test.sh

@@ -141,5 +141,5 @@ docker run \
     "${image_name}" \
     bash -c '
     set -e
-    VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
 '

.buildkite/test-amd-merge.yml

@@ -111,8 +111,8 @@ steps:
   commands:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
+    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
+    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni_expansion.py -m "advanced_model" --run-level "core_model"
 
 - label: "Omni Model Test Qwen3-Omni"
   agent_pool: mi325_2

.buildkite/test-amd-ready.yaml

@@ -112,19 +112,6 @@ steps:
               pytest -s -v tests/engine/test_async_omni_engine_abort.py
       '
 
-
-# NOTE: This test is not running any thing. It is skipped and deselected.
-# Currently it is = 1 skipped, 1 deselected, 17 warnings in 0.03s ======
-# - label: "Omni Model Test Qwen2-5-Omni"
-#   agent_pool: mi325_2
-#   depends_on: amd-build
-#   mirror_hardwares: [amdproduction]
-#   grade: Blocking
-#   commands:
-#     - export VLLM_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
-
 # - label: "Omni Model Test Qwen3-Omni"
 #   agent_pool: mi325_2
 #   depends_on: amd-build

.buildkite/test-amd.yaml

@@ -88,7 +88,7 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
     - pytest -s -v tests/engine/test_async_omni_engine_abort.py
 
 - label: "Omni Model Test Qwen3-Omni"

.buildkite/test-merge.yml

@@ -172,25 +172,6 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Omni Model Test"
-    timeout_in_minutes: 30
-    depends_on: upload-merge-pipeline
-    commands:
-      - export VLLM_LOGGING_LEVEL=DEBUG
-      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
   - label: "Qwen3-TTS CustomVoice E2E Test"
     depends_on: upload-merge-pipeline
     commands:
@@ -235,6 +216,28 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
+  # - label: "MOSS-TTS-Nano E2E Test"
+  #   depends_on: upload-merge-pipeline
+  #   commands:
+  #     - |
+  #       timeout 20m bash -c '
+  #         export VLLM_LOGGING_LEVEL=DEBUG
+  #         export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+  #         pytest -s -v tests/e2e/online_serving/test_moss_tts_nano.py tests/e2e/offline_inference/test_moss_tts_nano.py -m "advanced_model" --run-level "advanced_model"
+  #       '
+  #   agents:
+  #     queue: "gpu_1_queue"
+  #   plugins:
+  #     - docker#v5.2.0:
+  #         image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #         always-pull: true
+  #         propagate-environment: true
+  #         shm-size: "8gb"
+  #         environment:
+  #           - "HF_HOME=/fsx/hf_cache"
+  #         volumes:
+  #           - "/fsx/hf_cache:/fsx/hf_cache"
+
   - label: "Omni Model Test with H100"
     timeout_in_minutes: 30
     depends_on: upload-merge-pipeline
@@ -363,7 +366,7 @@ steps:
           set -e
           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_text2img.py -m "advanced_model" --run-level "advanced_model"
           pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
         '
     agents:

.buildkite/test-nightly.yml

@@ -10,7 +10,7 @@ steps:
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test"
     steps:
-      - label: ":full_moon: Omni · Function Test"
+      - label: ":full_moon: Omni · Function Test with H100"
         timeout_in_minutes: 90
         commands:
           - pytest -s -v tests/e2e/ -m "full_model and H100 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy
@@ -48,6 +48,25 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: Omni · Function Test with L4"
+        timeout_in_minutes: 90
+        commands:
+          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+          - pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
+        agents:
+          queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+        plugins:
+          - docker#v5.2.0:
+              image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+              always-pull: true
+              propagate-environment: true
+              shm-size: "8gb"
+              environment:
+                - "HF_HOME=/fsx/hf_cache"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: ":full_moon: Omni · Doc Test with L4"
         timeout_in_minutes: 90
         commands:
@@ -205,7 +224,7 @@ steps:
         timeout_in_minutes: 90
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          - pytest -s -v tests/e2e/ -m "full_model and L4 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy
+          - pytest -s -v tests/e2e/ -m "full_model and L4 and omni" -k "not test_qwen2_5" --run-level "full_model" --ignore=tests/e2e/accuracy
         agents:
           queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
         plugins:

.buildkite/test-weekly.yml

@@ -0,0 +1,85 @@
+env:
+  VLLM_WORKER_MULTIPROC_METHOD: spawn
+  HF_HUB_DOWNLOAD_TIMEOUT: 300
+  HF_HUB_ETAG_TIMEOUT: 60
+
+steps:
+  - label: "Reliability Test - qwen3-omni"
+    timeout_in_minutes: 180
+    depends_on: upload-weekly-pipeline
+    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
+    commands:
+      - pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
+    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: 2
+                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
+
+  - label: "Reliability Test - wan22"
+    timeout_in_minutes: 180
+    depends_on: upload-weekly-pipeline
+    if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
+    commands:
+      - pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
+    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: 2
+                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

.claude/skills/add-diffusion-model/SKILL.md

@@ -544,6 +544,15 @@ Pipelines without `SupportsModuleOffload` fall back to scanning
 well-known attribute names (`transformer`, `text_encoder`, `vae`,
 etc.), which fails for non-standard names.
 
+### Step 12: Performance Profiling
+
+After verifying correctness and implementing parallelism/caching, profile the model's performance to identify bottlenecks and ensure optimal execution.
+
+See the [Profiling Single-Stage Diffusion](../../../docs/contributing/profiling.md#3-profiling-single-stage-diffusion) guide for detailed instructions on:
+1. Using the PyTorch profiler (`profiler: "torch"`) to capture detailed CPU/CUDA traces.
+2. Using Nsight Systems (`nsys`) with `profiler: "cuda"` for low-overhead CUDA traces.
+3. Controlling profiling via `omni.start_profile()` and `omni.stop_profile()`.
+
 ---
 
 ## Iterative Development Tips

benchmarks/glm_image/README.md

@@ -0,0 +1,157 @@
+# GLM-Image Benchmarks
+
+Benchmark GLM-Image T2I (text-to-image) and I2I (image-to-image) performance across three backends: HuggingFace baseline, vLLM-Omni offline, and vLLM-Omni online serving.
+
+## Benchmarks
+
+| Benchmark | Script | Description |
+|-----------|--------|-------------|
+| HuggingFace Baseline | `huggingface/inference.py` | Single-GPU transformers + diffusers pipeline |
+| vLLM-Omni Offline | `vllm-omni/inference.py` | Offline inference with continuous batching |
+| vLLM-Omni Online | `benchmark_glm_image.py` | Online serving via `/v1/chat/completions` |
+
+## HuggingFace Baseline
+
+Single-request sequential inference using the reference HuggingFace pipeline.
+
+```bash
+# T2I
+CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
+    --model-path /path/to/GLM-Image --mode t2i --num-prompts 10
+
+# I2I
+CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
+    --model-path /path/to/GLM-Image --mode i2i --num-prompts 10
+```
+
+### Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--model-path` | `zai-org/GLM-Image` | Model path |
+| `--mode` | `t2i` | `t2i` or `i2i` |
+| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
+| `--num-prompts` | `10` | Number of images to generate |
+| `--width` / `--height` | `1024` | Output image size |
+| `--num-inference-steps` | `50` | Diffusion denoising steps |
+| `--output-dir` | `benchmarks/glm_image/huggingface/outputs` | Output directory |
+| `--output-file` | - | JSON file for metrics |
+
+## vLLM-Omni Offline
+
+Multi-GPU offline inference with pipeline parallelism and continuous batching.
+
+```bash
+# T2I
+CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
+    --model-path /path/to/GLM-Image --mode t2i --num-prompts 10
+
+# I2I
+CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
+    --model-path /path/to/GLM-Image --mode i2i --num-prompts 10
+```
+
+### Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--model-path` | `zai-org/GLM-Image` | Model path |
+| `--deploy-config` | - | Deploy config YAML |
+| `--mode` | `t2i` | `t2i` or `i2i` |
+| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
+| `--num-prompts` | `10` | Number of images to generate |
+| `--width` / `--height` | `1024` | Output image size |
+| `--num-inference-steps` | `50` | Diffusion denoising steps |
+| `--output-dir` | `benchmarks/glm_image/vllm-omni/outputs` | Output directory |
+| `--output-file` | - | JSON file for metrics |
+| `--stage-init-timeout` | `600` | Stage initialization timeout (s) |
+
+### Latency Computation
+
+In offline mode all requests are submitted simultaneously and processed with continuous batching. The per-request latency is computed by summing the actual per-stage times (with `stage_0_gen_ms` diffed against the previous request to remove accumulated queue/scheduling wait).
+
+## vLLM-Omni Online Serving
+
+### Start the server
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/GLM-Image \
+    --omni --port 8091 --host 0.0.0.0 \
+    --served-model-name glm-image
+```
+
+### Run the benchmark
+
+```bash
+# T2I
+python benchmarks/glm_image/benchmark_glm_image.py \
+    --mode t2i --num-prompts 10 --model glm-image
+
+# I2I
+python benchmarks/glm_image/benchmark_glm_image.py \
+    --mode i2i --num-prompts 10 --model glm-image
+
+# Custom dataset
+python benchmarks/glm_image/benchmark_glm_image.py \
+    --mode i2i --dataset custom \
+    --dataset-path prompts.json --num-prompts 5
+```
+
+### Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--mode` | `t2i` | `t2i` or `i2i` |
+| `--dataset` | `prompt` | `prompt`, `random`, or `custom` |
+| `--dataset-path` | - | JSON file path (required for `custom`) |
+| `--num-prompts` | `10` | Number of benchmark requests |
+| `--max-concurrency` | `1` | Max concurrent requests |
+| `--request-rate` | `inf` | Requests per second (Poisson arrival) |
+| `--warmup-requests` | `1` | Warmup requests before measurement |
+| `--width` / `--height` | `1024` | Output image size |
+| `--num-inference-steps` | `50` | Diffusion denoising steps |
+| `--seed` | - | Random seed |
+| `--model` | `default` | Model name (must match `--served-model-name`) |
+| `--host` | `localhost` | Server host |
+| `--port` | `8091` | Server port |
+| `--output-file` | - | JSON output file for metrics |
+| `--num-input-images` | `1` | Number of input images for random I2I |
+
+## Dataset
+
+The default dataset is hosted on [HuggingFace](https://huggingface.co/datasets/JaredforReal/glm-image-bench) (`prompt.json`). It is automatically downloaded and cached to `prompt/prompt.json` on first run. No manual setup needed.
+
+Each entry contains:
+
+- `t2i_prompt`: Text prompt for text-to-image generation
+- `i2i_prompt`: Text prompt for image-to-image editing
+- `image_url`: Source image URL for I2I (downloaded and cached on first use)
+
+Custom datasets use the same JSON format and can be provided via `--dataset-path`.
+
+## Pipeline Timings
+
+All three benchmarks report per-stage pipeline timings (in milliseconds):
+
+| Key | Description |
+|-----|-------------|
+| `preprocess_ms` | Input preprocessing (tokenization, multimodal encoding) |
+| `stage_0_gen_ms` | AR (autoregressive) model generation time |
+| `ar2diffusion_ms` | AR output to diffusion input conversion |
+| `stage_1_gen_ms` | Diffusion model denoising time |
+| `queue_wait_ms` | Queue wait time before processing |
+
+The stages are ordered by execution: `preprocess → stage_0 (AR) → ar2diffusion → stage_1 (Diffusion)`.
+
+## Sample Results
+
+Tested on 2x GPU with 10 prompts, 1024x1024, 50 denoising steps:
+
+| Backend | Mode | Latency Mean (s) | Throughput (img/s) |
+|---------|------|-------------------|--------------------|
+| HuggingFace | T2I | 72.6 | 0.014 |
+| HuggingFace | I2I | 70.9 | 0.014 |
+| vLLM-Omni Offline | T2I | 35.0 | 0.044 |
+| vLLM-Omni Offline | I2I | 31.0 | 0.053 |
+| vLLM-Omni Online | T2I | 38.8 | 0.026 |
+| vLLM-Omni Online | I2I | 34.7 | 0.029 |

benchmarks/glm_image/benchmark_glm_image.py

@@ -0,0 +1,464 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Online serving benchmark for GLM-Image (T2I and I2I modes).
+
+Sends requests to the /v1/chat/completions endpoint and reports end-to-end
+latency, throughput, and per-stage durations (when the server is started with
+--enable-diffusion-pipeline-profiler and/or --enable-ar-profiler).
+
+Supports three dataset types:
+  - prompt:   Use prompt.json (default). T2I uses t2i_prompt, I2I uses i2i_prompt
+              and sends source images from image_url.
+  - random:   Generate synthetic prompts (and random images for I2I).
+  - custom:   Load from a user-specified JSON file.
+
+Usage:
+    # T2I with prompt.json (default)
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode t2i --num-prompts 10
+
+    # I2I with prompt.json (downloads source images automatically)
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode i2i --num-prompts 10
+
+    # Random dataset
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode t2i --dataset random --num-prompts 20
+
+    # Custom dataset
+    python benchmarks/glm_image/benchmark_glm_image.py \
+        --mode i2i --dataset custom \
+        --dataset-path my_prompts.json --num-prompts 5
+"""
+
+import argparse
+import asyncio
+import base64
+import json
+import os
+import sys
+import tempfile
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import aiohttp
+import numpy as np
+import requests as sync_requests
+from PIL import Image
+from tqdm.asyncio import tqdm
+
+# Import backends from the diffusion benchmark (add parent dirs to path)
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "diffusion"))
+from backends import RequestFuncOutput
+
+BENCHMARK_DIR = Path(__file__).resolve().parent
+DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
+IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
+
+DATASET_REPO = "JaredforReal/glm-image-bench"
+DATASET_FILE = "prompt.json"
+
+
+def _ensure_prompt_json(dataset_path: str | None) -> str:
+    """Return path to prompt.json, downloading from HuggingFace if needed."""
+    if dataset_path:
+        return dataset_path
+    local = DEFAULT_PROMPT_JSON
+    if local.exists():
+        return str(local)
+    print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
+    try:
+        from huggingface_hub import hf_hub_download
+
+        downloaded = hf_hub_download(
+            repo_id=DATASET_REPO,
+            filename=DATASET_FILE,
+            repo_type="dataset",
+        )
+        local.parent.mkdir(parents=True, exist_ok=True)
+        import shutil
+
+        shutil.copy2(downloaded, local)
+        print(f"Saved to {local}")
+    except ImportError:
+        url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
+        import urllib.request
+
+        local.parent.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, local)
+        print(f"Saved to {local}")
+    return str(local)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class GLMImageRequest:
+    prompt: str
+    image_path: str | None = None  # Only for I2I mode
+
+
+def download_image(url: str) -> str:
+    """Download an image to cache and return the local path."""
+    IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
+    fname = url.rsplit("/", 1)[-1]
+    local_path = IMAGE_CACHE_DIR / fname
+    if local_path.exists():
+        return str(local_path)
+    resp = sync_requests.get(url, timeout=30)
+    resp.raise_for_status()
+    local_path.write_bytes(resp.content)
+    return str(local_path)
+
+
+def encode_image_as_data_url(path: str) -> str:
+    """Encode a local image file as a base64 data URL."""
+    with open(path, "rb") as f:
+        encoded = base64.b64encode(f.read()).decode("utf-8")
+    ext = Path(path).suffix.lower()
+    mime = {"png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg"}.get(ext, "image/png")
+    return f"data:{mime};base64,{encoded}"
+
+
+# ---------------------------------------------------------------------------
+# Datasets
+# ---------------------------------------------------------------------------
+
+
+class PromptDataset:
+    """Load from prompt.json. T2I uses t2i_prompt, I2I uses i2i_prompt + image_url."""
+
+    def __init__(self, args: argparse.Namespace):
+        path = _ensure_prompt_json(args.dataset_path)
+        with open(path, encoding="utf-8") as f:
+            raw = json.load(f)
+
+        prompt_key = "t2i_prompt" if args.mode == "t2i" else "i2i_prompt"
+        self.items: list[GLMImageRequest] = []
+
+        for entry in raw:
+            prompt = entry.get(prompt_key, "").strip()
+            if not prompt:
+                continue
+            image_path = None
+            if args.mode == "i2i":
+                url = entry.get("image_url", "")
+                if url:
+                    image_path = download_image(url)
+            self.items.append(GLMImageRequest(prompt=prompt, image_path=image_path))
+
+        if args.num_prompts and len(self.items) > args.num_prompts:
+            self.items = self.items[: args.num_prompts]
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, idx: int) -> GLMImageRequest:
+        return self.items[idx]
+
+    def get_requests(self) -> list[GLMImageRequest]:
+        return list(self.items)
+
+
+class RandomDataset:
+    """Generate synthetic prompts (and optional random images for I2I)."""
+
+    def __init__(self, args: argparse.Namespace):
+        self.args = args
+        self.num_prompts = args.num_prompts
+        self._random_image_paths: list[str] | None = None
+        if args.mode == "i2i":
+            self._random_image_paths = self._generate_random_images()
+
+    def _generate_random_images(self) -> list[str]:
+        paths: list[str] = []
+        for i in range(self.args.num_input_images):
+            img = Image.new("RGB", (512, 512), (128 + i * 30 % 128, 64, 192))
+            path = os.path.join(tempfile.gettempdir(), f"glm_image_bench_input_{i}.png")
+            img.save(path)
+            paths.append(path)
+        return paths
+
+    def __len__(self) -> int:
+        return self.num_prompts
+
+    def __getitem__(self, idx: int) -> GLMImageRequest:
+        image_path = None
+        if self._random_image_paths is not None:
+            image_path = self._random_image_paths[idx % len(self._random_image_paths)]
+        return GLMImageRequest(
+            prompt=f"A beautiful scene with vivid colors and intricate details, prompt {idx}",
+            image_path=image_path,
+        )
+
+    def get_requests(self) -> list[GLMImageRequest]:
+        return [self[i] for i in range(len(self))]
+
+
+class CustomDataset:
+    """Load from a user-specified JSON file.
+
+    Expected format:
+    [
+        {"prompt": "A cat sitting on a windowsill"},
+        {"prompt": "Make it look like winter", "image_path": "/path/to/img.png"}
+    ]
+    """
+
+    def __init__(self, args: argparse.Namespace):
+        if not args.dataset_path:
+            raise ValueError("--dataset-path is required for custom dataset")
+        with open(args.dataset_path, encoding="utf-8") as f:
+            raw = json.load(f)
+        self.items: list[GLMImageRequest] = []
+        for item in raw:
+            self.items.append(
+                GLMImageRequest(
+                    prompt=item.get("prompt", ""),
+                    image_path=item.get("image_path"),
+                )
+            )
+        if args.num_prompts and len(self.items) > args.num_prompts:
+            self.items = self.items[: args.num_prompts]
+
+    def __len__(self) -> int:
+        return len(self.items)
+
+    def __getitem__(self, idx: int) -> GLMImageRequest:
+        return self.items[idx]
+
+    def get_requests(self) -> list[GLMImageRequest]:
+        return list(self.items)
+
+
+# ---------------------------------------------------------------------------
+# Async request for GLM-Image (chat completions with image support)
+# ---------------------------------------------------------------------------
+
+
+async def async_glm_image_request(
+    req: GLMImageRequest,
+    api_url: str,
+    model: str,
+    session: aiohttp.ClientSession,
+    pbar: Any,
+    args: argparse.Namespace,
+) -> RequestFuncOutput:
+    """Send a single T2I or I2I request via chat completions endpoint."""
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    # Build messages
+    if req.image_path and args.mode == "i2i":
+        data_url = encode_image_as_data_url(req.image_path)
+        content = [
+            {"type": "text", "text": req.prompt},
+            {"type": "image_url", "image_url": {"url": data_url}},
+        ]
+    else:
+        content = req.prompt
+
+    messages = [{"role": "user", "content": content}]
+
+    extra_body: dict[str, Any] = {}
+    if args.height:
+        extra_body["height"] = args.height
+    if args.width:
+        extra_body["width"] = args.width
+    if args.num_inference_steps:
+        extra_body["num_inference_steps"] = args.num_inference_steps
+    if args.seed is not None:
+        extra_body["seed"] = args.seed
+
+    payload: dict[str, Any] = {
+        "model": model,
+        "messages": messages,
+    }
+    if extra_body:
+        payload["extra_body"] = extra_body
+
+    try:
+        async with session.post(api_url, json=payload) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+                try:
+                    choices = resp_json.get("choices", [])
+                    if choices and isinstance(choices, list):
+                        msg = choices[0].get("message", {})
+                        if isinstance(msg, dict):
+                            resp_content = msg.get("content", [])
+                            if resp_content and isinstance(resp_content, list) and len(resp_content) > 0:
+                                first_item = resp_content[0]
+                                if isinstance(first_item, dict):
+                                    output.stage_durations = first_item.get("stage_durations") or {}
+                                    output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
+                except (IndexError, TypeError, AttributeError):
+                    pass
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+
+async def iter_requests(n: int, request_rate: float) -> Any:
+    import random as _random
+
+    for i in range(n):
+        if request_rate != float("inf") and i > 0:
+            await asyncio.sleep(_random.expovariate(request_rate))
+        yield i
+
+
+def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float) -> dict[str, Any]:
+    success = [o for o in outputs if o.success]
+    errors = [o for o in outputs if not o.success]
+    latencies = [o.latency for o in success]
+    peak_mems = [o.peak_memory_mb for o in success if o.peak_memory_mb > 0]
+
+    stage_duration_lists: dict[str, list[float]] = {}
+    for o in success:
+        for stage, dur in (o.stage_durations or {}).items():
+            stage_duration_lists.setdefault(stage, []).append(dur)
+
+    return {
+        "duration": total_duration,
+        "completed_requests": len(success),
+        "failed_requests": len(errors),
+        "throughput_qps": len(success) / total_duration if total_duration > 0 else 0,
+        "latency_mean": float(np.mean(latencies)) if latencies else 0,
+        "latency_median": float(np.median(latencies)) if latencies else 0,
+        "latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
+        "latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
+        "peak_memory_mb_max": max(peak_mems) if peak_mems else 0,
+        "stage_durations_mean": {s: float(np.mean(v)) for s, v in stage_duration_lists.items()},
+        "stage_durations_p50": {s: float(np.percentile(v, 50)) for s, v in stage_duration_lists.items()},
+    }
+
+
+async def benchmark(args: argparse.Namespace) -> None:
+    api_url = f"http://{args.host}:{args.port}/v1/chat/completions"
+
+    # Load dataset
+    if args.dataset == "prompt":
+        dataset = PromptDataset(args)
+    elif args.dataset == "random":
+        dataset = RandomDataset(args)
+    elif args.dataset == "custom":
+        dataset = CustomDataset(args)
+    else:
+        raise ValueError(f"Unknown dataset: {args.dataset}")
+
+    glm_requests = dataset.get_requests()
+    print(f"Prepared {len(glm_requests)} requests (mode={args.mode}, dataset={args.dataset})")
+
+    semaphore = asyncio.Semaphore(args.max_concurrency) if args.max_concurrency else None
+
+    async def limited_request(idx: int, req: GLMImageRequest, session: aiohttp.ClientSession, pbar: Any):
+        if semaphore:
+            async with semaphore:
+                return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
+        return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
+
+    async with aiohttp.ClientSession() as session:
+        # Warmup
+        if args.warmup_requests and glm_requests:
+            print(f"Running {args.warmup_requests} warmup request(s)...")
+            for i in range(args.warmup_requests):
+                await limited_request(i, glm_requests[i % len(glm_requests)], session, None)
+
+        # Main benchmark
+        pbar = tqdm(total=len(glm_requests), disable=args.disable_tqdm)
+        start_time = time.perf_counter()
+        tasks = []
+        async for idx in iter_requests(len(glm_requests), args.request_rate):
+            tasks.append(asyncio.create_task(limited_request(idx, glm_requests[idx], session, pbar)))
+        outputs = await asyncio.gather(*tasks)
+        total_duration = time.perf_counter() - start_time
+        pbar.close()
+
+    # Metrics
+    metrics = calculate_metrics(outputs, total_duration)
+    metrics["mode"] = args.mode
+    metrics["model"] = args.model
+    metrics["dataset"] = args.dataset
+
+    print(f"\n{' GLM-Image Online Benchmark Result ':=^60}")
+    print(f"{'Mode:':<40} {args.mode}")
+    print(f"{'Model:':<40} {args.model}")
+    print(f"{'Dataset:':<40} {args.dataset}")
+    print("-" * 50)
+    print(f"{'Benchmark duration (s):':<40} {metrics['duration']:.2f}")
+    print(f"{'Request rate:':<40} {args.request_rate}")
+    print(f"{'Max concurrency:':<40} {args.max_concurrency}")
+    print(f"{'Successful requests:':<40} {metrics['completed_requests']}/{len(glm_requests)}")
+    print("-" * 50)
+    print(f"{'Throughput (req/s):':<40} {metrics['throughput_qps']:.2f}")
+    print(f"{'Latency Mean (s):':<40} {metrics['latency_mean']:.4f}")
+    print(f"{'Latency Median (s):':<40} {metrics['latency_median']:.4f}")
+    print(f"{'Latency P95 (s):':<40} {metrics['latency_p95']:.4f}")
+    print(f"{'Latency P99 (s):':<40} {metrics['latency_p99']:.4f}")
+
+    if metrics["peak_memory_mb_max"] > 0:
+        print("-" * 50)
+        print(f"{'Peak Memory Max (MB):':<40} {metrics['peak_memory_mb_max']:.2f}")
+
+    if metrics["stage_durations_mean"]:
+        print("-" * 50)
+        print("Stage Durations Mean:")
+        for stage, val in sorted(metrics["stage_durations_mean"].items()):
+            unit = "ms" if stage.endswith("_ms") else "s"
+            print(f"  {stage + ':':<38} {val:.4f} ({unit})")
+
+    print("=" * 60)
+
+    if args.output_file:
+        with open(args.output_file, "w") as f:
+            json.dump(metrics, f, indent=2)
+        print(f"Metrics saved to {args.output_file}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Benchmark GLM-Image T2I/I2I online serving.")
+    parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
+    parser.add_argument("--dataset", type=str, default="prompt", choices=["prompt", "random", "custom"])
+    parser.add_argument("--dataset-path", type=str, default=None)
+    parser.add_argument("--num-prompts", type=int, default=10)
+    parser.add_argument("--max-concurrency", type=int, default=1)
+    parser.add_argument("--request-rate", type=float, default=float("inf"))
+    parser.add_argument("--warmup-requests", type=int, default=1)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--num-inference-steps", type=int, default=50)
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--model", type=str, default="default")
+    parser.add_argument("--host", type=str, default="localhost")
+    parser.add_argument("--port", type=int, default=8091)
+    parser.add_argument("--output-file", type=str, default=None)
+    parser.add_argument("--disable-tqdm", action="store_true")
+    parser.add_argument("--num-input-images", type=int, default=1, help="For random I2I dataset.")
+    args = parser.parse_args()
+    asyncio.run(benchmark(args))
+
+
+if __name__ == "__main__":
+    main()

benchmarks/glm_image/huggingface/inference.py

@@ -0,0 +1,291 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+HuggingFace (transformers + diffusers) baseline benchmark for GLM-Image.
+
+Supports T2I and I2I modes with the prompt.json dataset.
+Downloads source images for I2I from image_url on first run and caches locally.
+
+Usage:
+    # T2I mode (text-to-image, no source images needed)
+    python benchmarks/glm_image/huggingface/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode t2i --num-prompts 10
+
+    # I2I mode (image-to-image, downloads source images)
+    python benchmarks/glm_image/huggingface/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode i2i --num-prompts 10
+
+    # With custom prompt.json
+    python benchmarks/glm_image/huggingface/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode i2i --dataset-path prompts.json --num-prompts 5
+"""
+
+import argparse
+import json
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import requests
+import torch
+from PIL import Image
+
+BENCHMARK_DIR = Path(__file__).resolve().parent.parent
+DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
+IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
+
+DATASET_REPO = "JaredforReal/glm-image-bench"
+DATASET_FILE = "prompt.json"
+
+
+def _ensure_prompt_json(dataset_path: str | None) -> str:
+    """Return path to prompt.json, downloading from HuggingFace if needed."""
+    if dataset_path:
+        return dataset_path
+    local = DEFAULT_PROMPT_JSON
+    if local.exists():
+        return str(local)
+    print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
+    try:
+        from huggingface_hub import hf_hub_download
+
+        downloaded = hf_hub_download(
+            repo_id=DATASET_REPO,
+            filename=DATASET_FILE,
+            repo_type="dataset",
+        )
+        local.parent.mkdir(parents=True, exist_ok=True)
+        import shutil
+
+        shutil.copy2(downloaded, local)
+        print(f"Saved to {local}")
+    except ImportError:
+        url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
+        import urllib.request
+
+        local.parent.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, local)
+        print(f"Saved to {local}")
+    return str(local)
+
+
+HEIGHT = 1024
+WIDTH = 1024
+SEED = 42
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 1.5
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+
+def load_dataset(
+    dataset_path: str | None,
+    mode: str,
+    num_prompts: int,
+) -> list[dict]:
+    """Load prompts from prompt.json and prepare per-request data."""
+    path = _ensure_prompt_json(dataset_path)
+    with open(path, encoding="utf-8") as f:
+        raw = json.load(f)
+
+    items = []
+    for entry in raw:
+        if mode == "t2i":
+            prompt_key = "t2i_prompt"
+        else:
+            prompt_key = "i2i_prompt"
+
+        prompt_text = entry.get(prompt_key, "").strip()
+        if not prompt_text:
+            continue
+
+        item = {"prompt": prompt_text}
+        if mode == "i2i":
+            item["image_url"] = entry.get("image_url", "")
+        items.append(item)
+
+    if num_prompts and len(items) > num_prompts:
+        items = items[:num_prompts]
+    return items
+
+
+def download_image(url: str, cache_dir: Path) -> str:
+    """Download an image to cache_dir and return the local path."""
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    fname = url.rsplit("/", 1)[-1]
+    local_path = cache_dir / fname
+    if local_path.exists():
+        return str(local_path)
+    print(f"  Downloading {url} ...")
+    resp = requests.get(url, timeout=30)
+    resp.raise_for_status()
+    local_path.write_bytes(resp.content)
+    return str(local_path)
+
+
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+
+def benchmark(args: argparse.Namespace) -> None:
+    from diffusers.pipelines.glm_image import GlmImagePipeline
+
+    print("=" * 60)
+    print("GLM-Image HuggingFace Baseline Benchmark")
+    print(f"Mode: {args.mode}  |  Model: {args.model_path}")
+    print(f"Size: {args.height}x{args.width}  |  Steps: {args.num_inference_steps}")
+    print("=" * 60)
+
+    # Load dataset
+    items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
+    if not items:
+        print("No prompts loaded. Exiting.")
+        return
+    print(f"Loaded {len(items)} prompts for {args.mode} mode")
+
+    # Download I2I source images
+    if args.mode == "i2i":
+        print("Preparing source images...")
+        for item in items:
+            url = item.get("image_url", "")
+            if url:
+                item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
+            else:
+                item["image_path"] = None
+
+    # Load pipeline
+    print(f"\nLoading pipeline from {args.model_path} ...")
+    t0 = time.perf_counter()
+    pipe = GlmImagePipeline.from_pretrained(
+        args.model_path,
+        torch_dtype=torch.bfloat16,
+        device_map="cuda",
+    )
+    init_time = time.perf_counter() - t0
+    print(f"Pipeline loaded in {init_time:.2f}s")
+
+    # Create output dir
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Run benchmark
+    generator = torch.Generator(device="cuda").manual_seed(args.seed)
+    latencies = []
+    success = 0
+    failed = 0
+
+    print(f"\nRunning {len(items)} requests sequentially...")
+    print("-" * 60)
+
+    for i, item in enumerate(items):
+        prompt = item["prompt"]
+        gen_kwargs: dict = {
+            "prompt": prompt,
+            "height": args.height,
+            "width": args.width,
+            "num_inference_steps": args.num_inference_steps,
+            "guidance_scale": args.guidance_scale,
+            "generator": generator,
+        }
+
+        if args.mode == "i2i":
+            img_path = item.get("image_path")
+            if img_path and os.path.exists(img_path):
+                gen_kwargs["image"] = [Image.open(img_path).convert("RGB")]
+            else:
+                print(f"  [{i + 1}] SKIP: no source image")
+                failed += 1
+                continue
+
+        t_start = time.perf_counter()
+        try:
+            result = pipe(**gen_kwargs)
+            image = result.images[0]
+            elapsed = time.perf_counter() - t_start
+            latencies.append(elapsed)
+            success += 1
+
+            out_path = os.path.join(args.output_dir, f"{i:04d}.png")
+            image.save(out_path)
+            print(f"  [{i + 1}/{len(items)}] {elapsed:.3f}s -> {out_path}")
+        except Exception as e:
+            elapsed = time.perf_counter() - t_start
+            failed += 1
+            print(f"  [{i + 1}/{len(items)}] FAILED ({elapsed:.3f}s): {e}")
+
+    # Report
+    total_gen_time = sum(latencies) if latencies else 0
+    print("\n" + "=" * 60)
+    print("HuggingFace Baseline Results")
+    print("=" * 60)
+    print(f"{'Mode:':<40} {args.mode}")
+    print(f"{'Model:':<40} {args.model_path}")
+    print(f"{'Image size:':<40} {args.height}x{args.width}")
+    print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
+    print("-" * 50)
+    print(f"{'Pipeline init time (s):':<40} {init_time:.2f}")
+    print(f"{'Successful:':<40} {success}/{len(items)}")
+    print(f"{'Failed:':<40} {failed}")
+    print("-" * 50)
+    if latencies:
+        arr = np.array(latencies)
+        print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
+        print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
+        print(f"{'Latency Mean (s):':<40} {arr.mean():.4f}")
+        print(f"{'Latency Median (s):':<40} {np.median(arr):.4f}")
+        print(f"{'Latency P95 (s):':<40} {np.percentile(arr, 95):.4f}")
+        print(f"{'Latency P99 (s):':<40} {np.percentile(arr, 99):.4f}")
+
+    print(f"\n{'Output dir:':<40} {args.output_dir}")
+    print("=" * 60)
+
+    # Save metrics JSON
+    metrics = {
+        "backend": "huggingface",
+        "mode": args.mode,
+        "model": args.model_path,
+        "height": args.height,
+        "width": args.width,
+        "num_inference_steps": args.num_inference_steps,
+        "init_time_s": init_time,
+        "completed_requests": success,
+        "failed_requests": failed,
+        "total_gen_time_s": total_gen_time,
+        "throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
+        "latency_mean": float(np.mean(latencies)) if latencies else 0,
+        "latency_median": float(np.median(latencies)) if latencies else 0,
+        "latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
+        "latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
+    }
+    if args.output_file:
+        with open(args.output_file, "w") as f:
+            json.dump(metrics, f, indent=2)
+        print(f"Metrics saved to {args.output_file}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="GLM-Image HuggingFace baseline benchmark")
+    parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
+    parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
+    parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
+    parser.add_argument("--num-prompts", type=int, default=10)
+    parser.add_argument("--height", type=int, default=HEIGHT)
+    parser.add_argument("--width", type=int, default=WIDTH)
+    parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
+    parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
+    parser.add_argument("--seed", type=int, default=SEED)
+    parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/huggingface/outputs")
+    parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
+    args = parser.parse_args()
+    benchmark(args)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/glm_image/vllm-omni/inference.py

@@ -0,0 +1,505 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+vLLM-Omni offline benchmark for GLM-Image.
+
+Supports T2I and I2I modes with the prompt.json dataset.
+Downloads source images for I2I from image_url on first run and caches locally.
+
+Usage:
+    # T2I mode
+    python benchmarks/glm_image/vllm-omni/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode t2i --num-prompts 10
+
+    # I2I mode (downloads source images)
+    python benchmarks/glm_image/vllm-omni/inference.py \
+        --model-path zai-org/GLM-Image \
+        --mode i2i --num-prompts 10
+"""
+
+import argparse
+import json
+import math
+import os
+import time
+from pathlib import Path
+
+import numpy as np
+import requests
+from PIL import Image
+from vllm import SamplingParams
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+BENCHMARK_DIR = Path(__file__).resolve().parent.parent
+DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
+IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
+DEFAULT_DEPLOY_CONFIG = "vllm_omni/deploy/glm_image.yaml"
+
+DATASET_REPO = "JaredforReal/glm-image-bench"
+DATASET_FILE = "prompt.json"
+
+
+def _ensure_prompt_json(dataset_path: str | None) -> str:
+    """Return path to prompt.json, downloading from HuggingFace if needed."""
+    if dataset_path:
+        return dataset_path
+    local = DEFAULT_PROMPT_JSON
+    if local.exists():
+        return str(local)
+    print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
+    try:
+        from huggingface_hub import hf_hub_download
+
+        downloaded = hf_hub_download(
+            repo_id=DATASET_REPO,
+            filename=DATASET_FILE,
+            repo_type="dataset",
+        )
+        local.parent.mkdir(parents=True, exist_ok=True)
+        import shutil
+
+        shutil.copy2(downloaded, local)
+        print(f"Saved to {local}")
+    except ImportError:
+        url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
+        import urllib.request
+
+        local.parent.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(url, local)
+        print(f"Saved to {local}")
+    return str(local)
+
+
+SEED = 42
+HEIGHT = 1024
+WIDTH = 1024
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 1.5
+
+GLM_IMAGE_EOS_TOKEN_ID = 16385
+GLM_IMAGE_VISION_VOCAB_SIZE = 16512
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+
+def load_dataset(
+    dataset_path: str | None,
+    mode: str,
+    num_prompts: int,
+) -> list[dict]:
+    path = _ensure_prompt_json(dataset_path)
+    with open(path, encoding="utf-8") as f:
+        raw = json.load(f)
+
+    items = []
+    for entry in raw:
+        prompt_key = "t2i_prompt" if mode == "t2i" else "i2i_prompt"
+        prompt_text = entry.get(prompt_key, "").strip()
+        if not prompt_text:
+            continue
+
+        item = {"prompt": prompt_text}
+        if mode == "i2i":
+            item["image_url"] = entry.get("image_url", "")
+        items.append(item)
+
+    if num_prompts and len(items) > num_prompts:
+        items = items[:num_prompts]
+    return items
+
+
+def download_image(url: str, cache_dir: Path) -> str:
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    fname = url.rsplit("/", 1)[-1]
+    local_path = cache_dir / fname
+    if local_path.exists():
+        return str(local_path)
+    print(f"  Downloading {url} ...")
+    resp = requests.get(url, timeout=30)
+    resp.raise_for_status()
+    local_path.write_bytes(resp.content)
+    return str(local_path)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def compute_max_tokens(height: int, width: int, is_i2i: bool = False) -> int:
+    factor = 32
+    token_h = height // factor
+    token_w = width // factor
+    large_tokens = token_h * token_w
+
+    # Small preview tokens (half resolution in each dimension)
+
+    ratio = token_h / token_w if token_w > 0 else 1.0
+    small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2)))
+    small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2)))
+    small_tokens = small_token_h * small_token_w
+
+    # Mode-dependent totals:
+    # - t2i: small + large + EOS
+    # - i2i: large + EOS
+    if is_i2i:
+        return large_tokens + 1
+    return small_tokens + large_tokens + 1
+
+
+def build_prompt_t2i(prompt: str, height: int, width: int, **gen_kw) -> dict:
+    return {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        "mm_processor_kwargs": {"target_h": height, "target_w": width},
+        **gen_kw,
+    }
+
+
+def build_prompt_i2i(prompt: str, image_path: str, height: int, width: int, **gen_kw) -> dict:
+    return {
+        "prompt": prompt,
+        "height": height,
+        "width": width,
+        "mm_processor_kwargs": {"target_h": height, "target_w": width},
+        "multi_modal_data": {"image": Image.open(image_path).convert("RGB")},
+        **gen_kw,
+    }
+
+
+def resolve_deploy_config(args: argparse.Namespace) -> str:
+    if args.deploy_config:
+        return args.deploy_config
+    if os.path.exists(DEFAULT_DEPLOY_CONFIG):
+        return DEFAULT_DEPLOY_CONFIG
+    fallback = Path(__file__).resolve().parents[3] / DEFAULT_DEPLOY_CONFIG
+    if fallback.exists():
+        return str(fallback)
+    raise FileNotFoundError("Deploy config not found. Specify --deploy-config.")
+
+
+# ---------------------------------------------------------------------------
+# Benchmark
+# ---------------------------------------------------------------------------
+
+
+def benchmark(args: argparse.Namespace) -> None:
+    is_i2i = args.mode == "i2i"
+
+    print("=" * 60)
+    print("GLM-Image vLLM-Omni Benchmark")
+    print(f"Mode: {args.mode}  |  Model: {args.model_path}")
+    print(f"Size: {args.height}x{args.width}  |  Steps: {args.num_inference_steps}")
+    print("=" * 60)
+
+    # Load dataset
+    items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
+    if not items:
+        print("No prompts loaded. Exiting.")
+        return
+    print(f"Loaded {len(items)} prompts for {args.mode} mode")
+
+    # Download I2I source images
+    if is_i2i:
+        print("Preparing source images...")
+        for item in items:
+            url = item.get("image_url", "")
+            if url:
+                item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
+            else:
+                item["image_path"] = None
+
+    # Init Omni
+    deploy_config = resolve_deploy_config(args)
+    print(f"\nInitializing vLLM-Omni (deploy config: {deploy_config}) ...")
+    t0 = time.perf_counter()
+
+    omni = Omni(
+        model=args.model_path,
+        deploy_config=deploy_config,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+        enable_ar_profiler=args.enable_ar_profiler,
+    )
+
+    init_time = time.perf_counter() - t0
+    print(f"Initialized in {init_time:.2f}s")
+
+    # Sampling params
+    max_tokens = compute_max_tokens(args.height, args.width, is_i2i=is_i2i)
+    ar_params = SamplingParams(
+        temperature=0.9,
+        top_p=0.75,
+        top_k=GLM_IMAGE_VISION_VOCAB_SIZE,
+        max_tokens=max_tokens,
+        stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],
+        seed=args.seed,
+        detokenize=False,
+        extra_args={"target_h": args.height, "target_w": args.width},
+    )
+    diff_params = OmniDiffusionSamplingParams(
+        num_inference_steps=args.num_inference_steps,
+        guidance_scale=args.guidance_scale,
+        height=args.height,
+        width=args.width,
+        seed=args.seed,
+    )
+    sampling_params_list = [ar_params, diff_params]
+
+    # Build all prompts
+    gen_kw = {
+        "seed": args.seed,
+        "num_inference_steps": args.num_inference_steps,
+        "guidance_scale": args.guidance_scale,
+    }
+    all_prompts = []
+    for item in items:
+        if is_i2i:
+            img_path = item.get("image_path")
+            if not img_path or not os.path.exists(img_path):
+                continue
+            all_prompts.append(build_prompt_i2i(item["prompt"], img_path, args.height, args.width, **gen_kw))
+        else:
+            all_prompts.append(build_prompt_t2i(item["prompt"], args.height, args.width, **gen_kw))
+
+    valid = len(all_prompts)
+    print(f"Valid prompts: {valid}")
+
+    # Create output dir
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    # Warmup: run 1 request to prime caches, CUDA graphs, etc.
+    if all_prompts:
+        print("Running warmup request...")
+        try:
+            warmup_prompt = [all_prompts[0]]
+            omni.generate(warmup_prompt, sampling_params_list, py_generator=False)
+            print("Warmup done.\n")
+        except Exception as e:
+            print(f"Warmup failed (continuing): {e}")
+
+    # Run
+    print(f"\nRunning {valid} requests...")
+    print("-" * 60)
+
+    latencies = []
+    all_stage_durations: list[dict[str, float]] = []
+    success = 0
+    failed = 0
+    wall_start = time.perf_counter()
+
+    try:
+        output_idx = 0
+        for stage_outputs in omni.generate(all_prompts, sampling_params_list, py_generator=True):
+            if stage_outputs.final_output_type == "image":
+                request_output = stage_outputs.request_output
+                request_id = getattr(request_output, "request_id", "")
+
+                images = getattr(request_output, "images", [])
+                if not images and hasattr(request_output, "multimodal_output"):
+                    mm = request_output.multimodal_output
+                    if isinstance(mm, dict):
+                        images = mm.get("images", [])
+
+                elapsed = time.perf_counter() - wall_start
+                if images:
+                    for img in images:
+                        if isinstance(img, Image.Image):
+                            out_path = os.path.join(args.output_dir, f"{output_idx:04d}.png")
+                            img.save(out_path)
+                    success += 1
+                    latencies.append(elapsed)
+                    stage_durations = getattr(stage_outputs, "stage_durations", {})
+                    if stage_durations:
+                        all_stage_durations.append(stage_durations)
+                    # Show wall-clock elapsed and pipeline breakdown if available
+                    preprocess_str = ""
+                    if "preprocess_ms" in stage_durations:
+                        preprocess_str = f" preprocess={stage_durations['preprocess_ms'] / 1000.0:.2f}s"
+                    print(f"  [{success}/{valid}] id={request_id[:8]} {elapsed:.2f}s{preprocess_str}")
+                    output_idx += 1
+                else:
+                    failed += 1
+    except Exception as e:
+        print(f"Error: {e}")
+        failed = valid - success
+
+    total_gen_time = time.perf_counter() - wall_start
+
+    # Diff stage_0_gen_ms with previous request to remove accumulated wait time.
+    # stage_0_gen_ms is measured from submit_ts (same for all requests submitted
+    # at once), so it accumulates queue/scheduling overhead across requests.
+    # Other stages and pipeline timings are per-request already.
+    _TIMING_ORDER = [
+        "preprocess_ms",
+        "stage_0_gen_ms",
+        "ar2diffusion_ms",
+        "stage_1_gen_ms",
+        "queue_wait_ms",
+    ]
+
+    per_request_actual: list[dict[str, float]] = []
+    prev_stage_0_ms = 0.0
+    for sd in all_stage_durations:
+        actual = dict(sd)
+        s0 = sd.get("stage_0_gen_ms", 0.0)
+        actual["stage_0_gen_ms"] = s0 - prev_stage_0_ms
+        prev_stage_0_ms = s0
+        per_request_actual.append(actual)
+
+    per_request_e2e_ms: list[float] = []
+    for actual in per_request_actual:
+        e2e_ms = sum(v for k, v in actual.items() if k in _TIMING_ORDER)
+        if e2e_ms > 0:
+            per_request_e2e_ms.append(e2e_ms)
+
+    # Report
+    print("\n" + "=" * 60)
+    print("vLLM-Omni Benchmark Results")
+    print("=" * 60)
+    print(f"{'Mode:':<40} {args.mode}")
+    print(f"{'Model:':<40} {args.model_path}")
+    print(f"{'Image size:':<40} {args.height}x{args.width}")
+    print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
+    print("-" * 50)
+    print(f"{'Init time (s):':<40} {init_time:.2f}")
+    print(f"{'Successful:':<40} {success}/{valid}")
+    print(f"{'Failed:':<40} {failed}")
+    print("-" * 50)
+
+    if per_request_e2e_ms:
+        per_request_s = np.array(per_request_e2e_ms) / 1000.0
+        print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
+        print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
+        print(f"{'Latency Mean (s):':<40} {per_request_s.mean():.4f}")
+        print(f"{'Latency Median (s):':<40} {np.median(per_request_s):.4f}")
+        print(f"{'Latency P95 (s):':<40} {np.percentile(per_request_s, 95):.4f}")
+        print(f"{'Latency P99 (s):':<40} {np.percentile(per_request_s, 99):.4f}")
+        print(f"{'Latency Min (s):':<40} {per_request_s.min():.4f}")
+        print(f"{'Latency Max (s):':<40} {per_request_s.max():.4f}")
+    elif latencies:
+        per_request = np.diff([0.0] + list(latencies))
+        print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
+        print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
+        print(f"{'Latency Mean (s) [wall-clock]:':<40} {per_request.mean():.4f}")
+        print(f"{'Latency Median (s) [wall-clock]:':<40} {np.median(per_request):.4f}")
+        print(f"{'Latency P95 (s) [wall-clock]:':<40} {np.percentile(per_request, 95):.4f}")
+        print(f"{'Latency P99 (s) [wall-clock]:':<40} {np.percentile(per_request, 99):.4f}")
+        print(f"{'Latency Min (s) [wall-clock]:':<40} {per_request.min():.4f}")
+        print(f"{'Latency Max (s) [wall-clock]:':<40} {per_request.max():.4f}")
+
+    if per_request_actual:
+        print("-" * 50)
+        print("Pipeline Timings Mean:")
+        for key in _TIMING_ORDER:
+            vals = [d.get(key, 0.0) for d in per_request_actual]
+            if any(v != 0 for v in vals):
+                unit = "ms" if key.endswith("_ms") else "s"
+                print(f"  {key + ':':<38} {np.mean(vals):.4f} ({unit})")
+        # Show any extra keys not in the ordered list
+        ordered_set = set(_TIMING_ORDER)
+        extra_keys = sorted(k for k in per_request_actual[0].keys() if k not in ordered_set)
+        for key in extra_keys:
+            vals = [d.get(key, 0.0) for d in per_request_actual]
+            if any(v != 0 for v in vals):
+                unit = "ms" if key.endswith("_ms") else "s"
+                print(f"  {key + ':':<38} {np.mean(vals):.4f} ({unit})")
+
+    print(f"\n{'Output dir:':<40} {args.output_dir}")
+    print("=" * 60)
+
+    # Metrics JSON
+    metrics = {
+        "backend": "vllm-omni",
+        "mode": args.mode,
+        "model": args.model_path,
+        "height": args.height,
+        "width": args.width,
+        "num_inference_steps": args.num_inference_steps,
+        "init_time_s": init_time,
+        "completed_requests": success,
+        "failed_requests": failed,
+        "total_gen_time_s": total_gen_time,
+        "throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
+    }
+    if per_request_e2e_ms:
+        per_request_s = np.array(per_request_e2e_ms) / 1000.0
+        metrics["latency_mean"] = float(per_request_s.mean())
+        metrics["latency_median"] = float(np.median(per_request_s))
+        metrics["latency_p95"] = float(np.percentile(per_request_s, 95))
+        metrics["latency_p99"] = float(np.percentile(per_request_s, 99))
+    elif latencies:
+        per_request = np.diff([0.0] + list(latencies))
+        metrics["latency_mean"] = float(per_request.mean())
+        metrics["latency_median"] = float(np.median(per_request))
+        metrics["latency_p95"] = float(np.percentile(per_request, 95))
+        metrics["latency_p99"] = float(np.percentile(per_request, 99))
+    else:
+        metrics["latency_mean"] = 0
+        metrics["latency_median"] = 0
+        metrics["latency_p95"] = 0
+        metrics["latency_p99"] = 0
+    if per_request_actual:
+        all_keys = list(_TIMING_ORDER) + sorted(k for k in per_request_actual[0].keys() if k not in set(_TIMING_ORDER))
+        stage_metrics = {}
+        for key in all_keys:
+            vals = [d.get(key, 0.0) for d in per_request_actual]
+            stage_metrics[key] = {
+                "mean": float(np.mean(vals)),
+                "median": float(np.median(vals)),
+                "p95": float(np.percentile(vals, 95)),
+            }
+        metrics["stage_durations"] = stage_metrics
+    if args.output_file:
+        with open(args.output_file, "w") as f:
+            json.dump(metrics, f, indent=2)
+        print(f"Metrics saved to {args.output_file}")
+
+    omni.close()
+    print("Done!")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="GLM-Image vLLM-Omni offline benchmark")
+    parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
+    parser.add_argument("--deploy-config", type=str, default=None, help="Deploy config YAML")
+    parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
+    parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
+    parser.add_argument("--num-prompts", type=int, default=10)
+    parser.add_argument("--height", type=int, default=HEIGHT)
+    parser.add_argument("--width", type=int, default=WIDTH)
+    parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
+    parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
+    parser.add_argument("--seed", type=int, default=SEED)
+    parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/vllm-omni/outputs")
+    parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
+    parser.add_argument("--stage-init-timeout", type=int, default=600)
+    parser.add_argument(
+        "--enable-diffusion-pipeline-profiler",
+        action="store_true",
+        help="Enable diffusion pipeline profiler for stage-level timing",
+    )
+    parser.add_argument(
+        "--enable-ar-profiler",
+        action="store_true",
+        help="Enable AR stage profiler to include AR timing in stage_durations",
+    )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        help="Enable detailed per-request pipeline stats logging",
+    )
+    args = parser.parse_args()
+    benchmark(args)
+
+
+if __name__ == "__main__":
+    main()

docs/.nav.yml

@@ -22,6 +22,7 @@ nav:
       - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
+      - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
       - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
       - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md

docs/cli/serve.md

@@ -12,9 +12,11 @@ The stage-based CLI is designed for deployments that require launching each pipe
   `vllm_omni/model_executor/stage_configs/`, the `--stage-configs-path` parameter remains mandatory.
 
 Example: Initializing Stage 0 (Orchestrator and API Server):
+The commands below show a common device mapping where Stage 0 uses GPU 0 and
+worker stages use GPU 1 via `CUDA_VISIBLE_DEVICES`.
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --port 8091 \
     --stage-id 0 \
     --omni-master-address 127.0.0.1 \
@@ -24,7 +26,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 Example: Initializing a Headless Worker Stage (Stage 1):
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -46,7 +48,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
 the stage-based CLI permits the direct initialization of Stage 1 with explicit parameters:
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --gpu-memory-utilization 0.5 \

docs/configuration/stage_configs.md

@@ -102,15 +102,17 @@ The stage-based CLI paradigm facilitates the execution of discrete pipeline stag
 
 For migrated architectures, the system automatically resolves and loads the bundled deployment YAML. Consequently, the primary execution path
 does **not** necessitate the explicit definition of `--deploy-config`:
+the example below uses `CUDA_VISIBLE_DEVICES=0` for Stage 0 and
+`CUDA_VISIBLE_DEVICES=1` for Stage 1.
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --port 8091 \
     --stage-id 0 \
     --omni-master-address 127.0.0.1 \
     --omni-master-port 26000
 
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -132,7 +134,7 @@ Conversely, in the context of the **stage-based CLI** paradigm, given that each
 can be defined uniformly via explicit CLI flags on the corresponding instantiation command, rendering composite `--stage-overrides` JSON strings unnecessary:
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --gpu-memory-utilization 0.5 \
@@ -189,7 +191,7 @@ Within the stage-based CLI paradigm, equivalent configuration parameters can inh
 as command-line arguments to the designated single-stage process instantiation:
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 0 \
     --max-num-seqs 8 \
     --omni-master-address 127.0.0.1 \

docs/contributing/ci/CI_5levels.md

@@ -108,7 +108,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
         /tests/dfx/stability/tests/test_qwen3_omni.json<br>
         /tests/dfx/stability/tests/test_wan22.json<br>
         <strong>Reliability:</strong><br>
-        tests/e2e/reliability/test_{model_name}.py
+        tests/dfx/reliability/test_reliability_{model_key}.py<br>
+        (e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>)
       </td>
       <td>
         <a href="#chapter-4-l5-level-testing-stability-and-reliability-testing">Chapter 4</a><br>
@@ -576,12 +577,17 @@ L5 level testing focuses on the performance of model services under ***long-runn
 ### 4.2 Testing Content and Scope
 
 -   ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
--   ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
+-   ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni or `multiprocessing.spawn` for Wan2.2 video workers), **large multimodal chat** or **`/v1/videos`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `<details>` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill).
 
 ### 4.3 Test Directory and Execution Files
 
 -   ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family)
--   ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
+-   ***Reliability Test Suite*** (`tests/dfx/reliability/`):
+    -   `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
+    -   `test_reliability_wan22.py` — Wan2.2 T2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
+    -   `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
+    -   `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
+    -   `README.md` — Short local run commands for this directory.
 
 ### 4.4 Execution Method and Example
 
@@ -662,8 +668,57 @@ All other optional parameters follow the same rules as the in Chapter 3.4.
 
 </details>
 
+<details>
+<summary> Reliability test suite (<code>tests/dfx/reliability</code>)</summary>
+
+#### Purpose and relationship to stability
+
+Reliability tests are **short fault-injection** integration runs (L5 **(b)** in `tests/dfx/reliability/README.md`). They complement **stability** JSON-driven long runs: instead of hours of steady traffic, they **perturb** the server (GPU OOM sidecar, fatal signals on selected processes) and check **failure mode** and **latency bounds** (e.g. chat or `/v1/videos` must not hang under concurrent fault-time load).
+
+#### Directory layout
+
+| Path | Responsibility |
+| ---- | -------------- |
+| `helpers.py` | Shared helpers used by current reliability suites: raw `POST`/`GET` probes (`/v1/chat/completions`, `/health`), OpenAI error parsing (`extract_openai_error_contract_from_bytes`), GPU OOM sidecar lifecycle (`inject_gpu_oom`, `stop_gpu_oom_hogs`), and process-kill injector builder (`make_process_kill_fault_injector`). |
+| `conftest.py` | Pytest fixtures: indirect `fault_injector`, `omni_server_after_fault` / `omni_server_after_fault_function` (run injector after server is ready, then yield server). |
+| `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
+| `test_reliability_wan22.py` | Wan2.2 T2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
+| `README.md` | Minimal run / collect examples. |
+
+#### Parametrization and markers
+
+- Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
+- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni paths often require **2** cards; Wan2.2 video paths **1** card).
+- **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
+
+#### CI trigger
+
+Weekly Buildkite (`.buildkite/test-weekly.yml`) runs, for example:
+
+```bash
+pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
+pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
+```
+
+#### Local commands
+
+```bash
+pytest --collect-only tests/dfx/reliability
+pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
+pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
+```
+
+#### Adding a new model suite
+
+1. Add `test_reliability_<model>.py` under `tests/dfx/reliability/`.
+2. Define **`RELIABILITY_SCENARIOS`** and pass them through **`create_reliability_omni_server_params()`** with the correct deploy or e2e stage-config directory (same pattern as existing files).
+3. Reuse **`helpers`** for OOM / kill / raw HTTP; prefer **`assert_fault_exception()`** and **`resolve_oom_device_spec()`** from `tests/dfx/conftest.py` for consistent device selection vs stage YAML.
+4. Register **`slow`** (and **`hardware_test`** if needed); extend **`.buildkite/test-weekly.yml`** when the suite should run in weekly L5.
+
+</details>
+
 -   -   ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
-    -   ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
+    -   ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow` and/or `pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow` (add `test_reliability_<suite>.py` for new models)
 
 ## Summary
 

docs/models/supported_models.md

@@ -20,6 +20,7 @@ th {
 | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
+| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
 | `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |

docs/user_guide/diffusion_features.md

@@ -108,7 +108,7 @@ The following tables show which models support each feature:
 | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
 | **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅  | ❌ | ❌ | ❌ |
-| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |

docs/user_guide/examples/offline_inference/glm_image.md

@@ -1,154 +1,87 @@
-# GLM-Image Multistage End-to-End Inference
+# GLM-Image Offline Inference
 
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/glm_image>.
-
-
-This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
+GLM-Image is a 2-stage image generation model (AR + Diffusion) supported by vLLM-Omni's
+declarative config system. The pipeline topology and stage structure are declared in
+`vllm_omni/model_executor/models/glm_image/pipeline.py`; deployment knobs live in
+`vllm_omni/deploy/glm_image.yaml`.
 
 ## Architecture
 
-GLM-Image uses a 2-stage pipeline:
-
 ```
-┌─────────────────────────────────────────────────────────────┐
-│                     GLM-Image Pipeline                       │
-├─────────────────────────────────────────────────────────────┤
-│                                                             │
-│  Stage 0 (AR Model)              Stage 1 (Diffusion)        │
-│  ┌─────────────────┐            ┌─────────────────────┐     │
-│  │ vLLM-optimized  │            │  GlmImagePipeline   │     │
-│  │ GlmImageFor     │  prior     │  ┌───────────────┐  │     │
-│  │ Conditional     │──tokens───►│  │ DiT Denoiser  │  │     │
-│  │ Generation      │            │  └───────────────┘  │     │
-│  │ (9B AR model)   │            │         │          │     │
-│  └─────────────────┘            │         ▼          │     │
-│         ▲                       │  ┌───────────────┐  │     │
-│         │                       │  │  VAE Decode   │──┼──► Image
-│    Text/Image                   │  └───────────────┘  │     │
-│      Input                      └─────────────────────┘     │
-│                                                             │
-└─────────────────────────────────────────────────────────────┘
+Stage 0 (AR Model)                Stage 1 (Diffusion)
+┌───────────────────┐            ┌─────────────────────┐
+│ vLLM-optimized    │  prior     │  GlmImagePipeline   │
+│ GlmImageFor       │──tokens──►│  ┌───────────────┐  │
+│ Conditional       │            │  │ DiT Denoiser  │  │
+│ Generation        │            │  └───────┬───────┘  │
+│ (9B AR model)     │            │          ▼          │
+└───────────────────┘            │  ┌───────────────┐  │
+        ▲                        │  │  VAE Decode   │──┼──► Image
+        │                        │  └───────────────┘  │
+   Text / Image                  └─────────────────────┘
+     Input
 ```
 
-## Features
+## Text-to-Image
 
-- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
-- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
-- **Text-to-Image**: Generate images from text descriptions
-- **Image-to-Image**: Edit existing images with text prompts
+```python
+from vllm_omni.entrypoints.omni import Omni
 
-## Usage
-
-### Text-to-Image
-
-```bash
-python end2end.py \
-    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
-    --prompt "A beautiful sunset over the ocean with sailing boats" \
-    --height 1024 \
-    --width 1024 \
-    --output output_t2i.png
+if __name__ == "__main__":
+    omni = Omni(model="zai-org/GLM-Image")
+    outputs = omni.generate(
+        "A photorealistic mountain landscape at sunset",
+        sampling_params={
+            "height": 1024,
+            "width": 1024,
+            "num_inference_steps": 50,
+            "guidance_scale": 1.5,
+            "seed": 42,
+        },
+    )
+    outputs[0].request_output.images[0].save("output.png")
 ```
 
-### Image-to-Image (Image Editing)
+## Image-to-Image (Image Editing)
 
-```bash
-python end2end.py \
-    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
-    --prompt "Transform this scene into a winter wonderland" \
-    --image input.png \
-    --output output_i2i.png
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(model="zai-org/GLM-Image")
+    outputs = omni.generate(
+        {
+            "prompt": "Convert this image to watercolor style",
+            "multi_modal_data": {
+                "image": "input.png",
+            },
+        },
+        sampling_params={
+            "height": 1024,
+            "width": 1024,
+            "num_inference_steps": 50,
+            "guidance_scale": 1.5,
+            "seed": 42,
+        },
+    )
+    outputs[0].request_output.images[0].save("output.png")
 ```
 
-### With Custom Parameters
+## Generation Parameters
 
-```bash
-python end2end.py \
-    --model-path /path/to/glm-image \
-    --config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
-    --prompt "A photorealistic cat sitting on a window sill" \
-    --height 1024 \
-    --width 1024 \
-    --num-inference-steps 50 \
-    --guidance-scale 1.5 \
-    --seed 42 \
-    --output output.png
-```
+| Parameter             | Type  | Default | Description                         |
+| --------------------- | ----- | ------- | ----------------------------------- |
+| `height`              | int   | 1024    | Image height in pixels              |
+| `width`               | int   | 1024    | Image width in pixels               |
+| `num_inference_steps` | int   | 50      | Number of diffusion denoising steps |
+| `guidance_scale`      | float | 1.5     | Classifier-free guidance scale      |
+| `seed`                | int   | None    | Optional random seed                |
+| `negative_prompt`     | str   | None    | Negative prompt                     |
 
-## Shell Scripts
+## VRAM Requirements
 
-### Run Text-to-Image
-
-```bash
-./run_t2i.sh
-```
-
-### Run Image-to-Image
-
-```bash
-./run_i2i.sh --image /path/to/input.png
-```
-
-## Stage Configuration
-
-The stage config (`glm_image.yaml`) defines:
-
-- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
-
-  - Model: `GlmImageForConditionalGeneration`
-  - Output: `token_ids` (prior tokens)
-
-- **Stage 1 (Diffusion)**: Uses diffusion engine
-  - Model: `GlmImagePipeline`
-  - Output: Generated image
-
-See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
-
-## Comparison with Single-Stage
-
-| Aspect      | Single-Stage (transformers) | Multistage (vLLM)   |
-| ----------- | --------------------------- | ------------------- |
-| AR Model    | transformers native         | vLLM PagedAttention |
-| Memory      | Higher (no KV cache opt)    | Lower (optimized)   |
-| Throughput  | Lower                       | Higher              |
-| Flexibility | Single GPU                  | Multi-GPU support   |
-
-## Troubleshooting
-
-### OOM Error
-
-Try reducing memory usage:
-
-```bash
-# In glm_image.yaml, adjust:
-gpu_memory_utilization: 0.5  # Reduce from 0.6
-```
-
-### Slow Initialization
-
-The first run loads model weights. Subsequent runs are faster:
-
-```bash
---stage-init-timeout 900  # Increase timeout for slow storage
-```
-
-## Requirements
-
-- vLLM-Omni with GLM-Image support
-- CUDA-capable GPU (recommended: H100/A100 with 80GB)
-- GLM-Image model weights
-
-## Example materials
-
-??? abstract "end2end.py"
-    ``````py
-    --8<-- "examples/offline_inference/glm_image/end2end.py"
-    ``````
-??? abstract "run_i2i.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/glm_image/run_i2i.sh"
-    ``````
-??? abstract "run_t2i.sh"
-    ``````sh
-    --8<-- "examples/offline_inference/glm_image/run_t2i.sh"
-    ``````
+| Stage             | VRAM                   |
+| :---------------- | :--------------------- |
+| Stage-0 (AR)      | **~18 GiB + KV Cache** |
+| Stage-1 (DiT+VAE) | **~20 GiB**            |
+| Total             | **~38 GiB + KV Cache** |

docs/user_guide/examples/offline_inference/internvla_a1.md

@@ -0,0 +1,205 @@
+# InternVLA-A1
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/internvla_a1>.
+
+
+## Setup
+This example is adapted from: https://github.com/InternRobotics/InternVLA-A1/blob/master/tests/policies/internvla_a1_3b/open_loop_genie1_real.ipynb
+
+This example runs the single-path vLLM offline inference workflow for InternVLA-A1 open-loop action prediction.
+
+Before running the script, export the required local paths:
+
+```bash
+# hf Jia-Zeng/InternVLA-A1-3B-FineTuned-Place_Markpen
+export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen
+# hf download InternRobotics/InternData-A1 real_lerobotv30/genie1/Genie1-Place_Markpen.tar.gz --repo-type dataset --local-dir /path/to/Genie1-Place_Markpen
+export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen
+export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct
+# hf tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors
+export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor
+```
+
+The shell entrypoint also accepts these variables as defaults, so you can keep the command line itself short.
+
+`INTERNVLA_A1_COSMOS_DIR` is expected to contain:
+
+- `encoder.safetensors`
+- `decoder.safetensors`
+
+Reference Hugging Face repo: `tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors`
+
+## Run examples
+
+Get into the example folder:
+
+```bash
+cd examples/offline_inference/internvla_a1
+```
+
+### Run one sample
+
+This runs a single sample through the maintained vLLM registry path.
+
+```bash
+bash run.sh --num-samples 1 --num-episodes 0
+```
+
+### Run open-loop evaluation against GT
+
+This runs one episode, writes evaluation metrics, and saves prediction-vs-GT plots when `matplotlib` is available.
+
+```bash
+bash run.sh --num-episodes 1
+```
+
+Outputs are written under:
+
+```bash
+outputs/internvla_a1/vllm_infer/
+```
+
+Typical files:
+
+- `summary.json`: top-level run summary
+- `registry/log.json`: per-episode GT comparison metrics
+- `registry/plots/*.jpg`: prediction-vs-GT figures
+
+### Eager float32 validation
+
+Use eager attention and float32 when you want the most stable numerical comparison baseline.
+
+```bash
+bash run.sh \
+  --num-episodes 1 \
+  --attn-implementation eager \
+  --dtype float32
+```
+
+### Optional runtime switches
+
+- `--dtype {bfloat16,float32}`: choose inference dtype
+- `--attn-implementation {eager,sdpa}`: switch attention backend
+- `--enable-regional-compile`: enable regional `torch.compile`
+- `--enable-warmup`: run pipeline warmup in initialization
+- `--skip-plots`: skip plot generation even if `matplotlib` is installed
+
+### Collect results and performance logs
+
+Use the helper below when you want to collect one-sample output, forward latency, GT comparison, plots, timing logs, and optional pytest output into a single result directory:
+
+```bash
+bash collect_results.sh
+```
+
+The script writes a timestamped directory under:
+
+```bash
+outputs/internvla_a1/collected_results/
+```
+
+Typical result files:
+
+- `env_summary.txt`: environment and path summary
+- `sample_run.log`: one-sample inference log
+- `forward_benchmark/forward_latency.json`: isolated `pipeline.forward` latency summary
+- `eval_run.log`: open-loop GT evaluation log
+- `eval_outputs/summary.json`: top-level output summary
+- `eval_outputs/registry/log.json`: GT comparison metrics
+- `eval_outputs/registry/plots/*.jpg`: prediction-vs-GT figures
+- `*_time.txt`: timing output
+- `*_gpu.csv`: sampled GPU usage
+
+### Benchmark forward latency
+
+Use the dedicated mode below when you want the isolated `pipeline.forward` latency instead of end-to-end script time:
+
+```bash
+python end2end.py \
+  --model-dir "$INTERNVLA_A1_MODEL_DIR" \
+  --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
+  --benchmark-forward \
+  --dtype bfloat16 \
+  --attn-implementation eager \
+  --warmup-iters 3 \
+  --benchmark-iters 10 \
+  --output-dir outputs/internvla_a1/forward_benchmark
+```
+
+The output JSON is written to `outputs/internvla_a1/forward_benchmark/forward_latency.json` and contains:
+
+- `cold_start_ms`
+- `warmup_summary`
+- `benchmark_summary`
+- `benchmark_samples_ms`
+
+### Reference results
+
+Reference run collected on `1x NVIDIA H200`, `bfloat16`, `eager`:
+
+- one-sample end-to-end run: `38s`
+- one-episode GT evaluation run: `45s`
+- `average_mse = 1.7173260857816786e-05`
+- `average_mae = 0.0011860118247568607`
+- `average_mse_joint = 7.42028441891307e-06`
+- `average_mae_joint = 0.0010777723509818316`
+- `average_mse_gripper = 8.544408774469048e-05`
+- `average_mae_gripper = 0.0019436875591054559`
+
+The reference plot filename is:
+
+- `eval_outputs/registry/plots/vllm_registry_open_loop_ep0.jpg`
+
+## Measure latency and VRAM
+
+You can collect end-to-end runtime and peak host memory with:
+
+```bash
+/usr/bin/time -v bash run.sh --num-episodes 1
+```
+
+For GPU memory snapshots, use:
+
+```bash
+nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv
+```
+
+Run `nvidia-smi` before and during inference to record the peak device memory for your PR report.
+
+## FAQ
+
+If `matplotlib` is missing, evaluation still runs and only plot generation is skipped. To enable plots:
+
+```bash
+pip install matplotlib
+```
+
+## Example materials
+
+- `end2end.py`: main offline inference and GT evaluation entrypoint
+- `run.sh`: shell wrapper with local path env vars
+- `collect_results.sh`: helper to gather result summaries and performance logs into one directory
+- `internvla_a1_common.py`: dataset, evaluation, and plotting helpers
+
+## Embedded source listings
+
+??? abstract "collect_results.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/internvla_a1/collect_results.sh"
+    ``````
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/internvla_a1/end2end.py"
+    ``````
+??? abstract "internvla_a1_common.py"
+    ``````py
+    --8<-- "examples/offline_inference/internvla_a1/internvla_a1_common.py"
+    ``````
+??? abstract "run.sh"
+    ``````sh
+    --8<-- "examples/offline_inference/internvla_a1/run.sh"
+    ``````
+??? abstract "standalone/outputs/qwena1/standalone_infer/log.json"
+    ``````json
+    --8<-- "examples/offline_inference/internvla_a1/standalone/outputs/qwena1/standalone_infer/log.json"
+    ``````

docs/user_guide/examples/online_serving/glm_image.md

@@ -1,99 +1,96 @@
 # GLM-Image Online Serving
 
-Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/glm_image>.
+GLM-Image is a 2-stage image generation model (AR + Diffusion) supported by vLLM-Omni's
+declarative config system. The pipeline topology and stage structure are declared in
+`vllm_omni/model_executor/models/glm_image/pipeline.py`; deployment knobs (GPU placement,
+memory, sampling params) live in `vllm_omni/deploy/glm_image.yaml`.
 
-
-This example demonstrates how to deploy GLM-Image for online image generation using vLLM-Omni.
-
-## 🛠️ Installation
-
-Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
-
-## Run examples (GLM-Image)
-
-**Note**: These examples work with the default configuration on **2× NVIDIA A100 (80GB)** or equivalent. Stage 0 (AR) and Stage 1 (Diffusion) each use one GPU by default. For single-GPU setups, modify the stage configuration to share the same device.
-
-### Launch the Server
+## Start Server
 
 ```bash
-# Use default configuration
 vllm serve zai-org/GLM-Image --omni --port 8091
 ```
 
-Or use the convenience script:
+The config system auto-detects the pipeline from the model's `model_index.json` — no
+manual `--stage-configs-path` or `--deploy-config` needed.
+
+By default, stage 0 (AR) runs on GPU 0 and stage 1 (Diffusion) on GPU 1. To colocate
+both stages on a single GPU, override per stage:
 
 ```bash
-cd examples/online_serving/glm_image
-bash run_server.sh
+vllm serve zai-org/GLM-Image --omni --port 8091 \
+    --stage-0-devices 0 --stage-1-devices 0
 ```
 
-If you have a custom stage configs file:
-
-```bash
-vllm serve zai-org/GLM-Image --omni --port 8091 --stage-configs-path /path/to/glm_image.yaml
-```
-
-### Send Requests
-
-Get into the glm_image folder:
-
-```bash
-cd examples/online_serving/glm_image
-```
-
-Send request via Python:
-
-```bash
-python openai_chat_client.py --prompt "A cute cat sitting on a window sill"
-```
-
-The Python client supports the following command-line arguments:
-
-- `--prompt` (or `-p`): Text prompt for generation (default: `A beautiful sunset over the ocean with sailing boats`)
-- `--output` (or `-o`): Output file path (default: `glm_image_output.png`)
-- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
-- `--image` (or `-i`): Input image path (for image-to-image editing)
-- `--height`: Image height in pixels (default: 1024)
-- `--width`: Image width in pixels (default: 1024)
-- `--steps`: Number of inference steps (default: 50)
-- `--guidance-scale`: Classifier-free guidance scale (default: 1.5)
-- `--seed`: Random seed (default: 42)
-- `--negative`: Negative prompt
-
-## Modality Control
-
-GLM-Image supports **text-to-image** and **image-to-image** modes.
-
-The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use the default configuration file: [`glm_image.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/glm_image.yaml)
-
-| Mode           | Input        | Output | Description                        |
-| -------------- | ------------ | ------ | ---------------------------------- |
-| Text-to-Image  | Text         | Image  | Generate images from text prompts  |
-| Image-to-Image | Image + Text | Image  | Edit images with text instructions |
+## API Calls
 
 ### Text-to-Image
 
 ```bash
-python openai_chat_client.py \
-    --prompt "A photorealistic mountain landscape at sunset" \
-    --height 1024 \
-    --width 1024 \
-    --output landscape.png
-
-# Or use the curl script:
-bash run_curl_text_to_image.sh "A futuristic city skyline at night"
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {"role": "user", "content": "A photorealistic mountain landscape at sunset"}
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 50,
+      "guidance_scale": 1.5,
+      "seed": 42
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
 ```
 
 ### Image-to-Image (Image Editing)
 
 ```bash
-python openai_chat_client.py \
-    --prompt "Convert this image to watercolor style" \
-    --image input.png \
-    --output watercolor.png
+curl -s http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [
+      {
+        "role": "user",
+        "content": [
+          {"type": "text", "text": "Convert this image to watercolor style"},
+          {"type": "image_url", "image_url": {"url": "data:image/png;base64,$(base64 -w0 input.png)}"}
+        ]
+      }
+    ],
+    "extra_body": {
+      "height": 1024,
+      "width": 1024,
+      "num_inference_steps": 50,
+      "guidance_scale": 1.5
+    }
+  }' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
+```
 
-# Or use the curl script:
-bash run_curl_image_edit.sh input.png "Convert to watercolor style"
+### Using the OpenAI Python SDK
+
+```python
+from openai import OpenAI
+import base64
+
+client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
+
+response = client.chat.completions.create(
+    model="zai-org/GLM-Image",
+    messages=[{"role": "user", "content": "A beautiful sunset over the ocean"}],
+    extra_body={
+        "height": 1024,
+        "width": 1024,
+        "num_inference_steps": 50,
+        "guidance_scale": 1.5,
+        "seed": 42,
+    },
+)
+
+img_url = response.choices[0].message.content[0].image_url.url
+_, b64_data = img_url.split(",", 1)
+with open("output.png", "wb") as f:
+    f.write(base64.b64decode(b64_data))
 ```
 
 For general-purpose request methods (curl, OpenAI SDK, Python `requests`), see
@@ -116,7 +113,7 @@ dimensions and count, use `size` and `n` rather than `height` or `width`.
 | `width`               | int   | 1024    | Image width in pixels               |
 | `num_inference_steps` | int   | 50      | Number of diffusion denoising steps |
 | `guidance_scale`      | float | 1.5     | Classifier-free guidance scale      |
-| `seed`                | int   | None    | Optional random seed; `/v1/images/*` generates one server-side if omitted |
+| `seed`                | int   | None    | Optional random seed                |
 | `negative_prompt`     | str   | None    | Negative prompt                     |
 
 ## Response Format
@@ -150,13 +147,12 @@ dimensions and count, use `size` and `n` rather than `height` or `width`.
 ## Extract Image
 
 ```bash
-# From a saved JSON response
 cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
 ```
 
 ## Architecture
 
-GLM-Image uses a 2-stage multistage pipeline:
+GLM-Image uses a 2-stage pipeline:
 
 ```
 Stage 0 (AR Model)                Stage 1 (Diffusion)
@@ -181,41 +177,13 @@ Stage 0 (AR Model)                Stage 1 (Diffusion)
 | Stage-1 (DiT+VAE) | **~20 GiB**            |
 | Total             | **~38 GiB + KV Cache** |
 
-## File Description
-
-| File                        | Description                           |
-| --------------------------- | ------------------------------------- |
-| `run_server.sh`             | Server startup script                 |
-| `run_curl_text_to_image.sh` | Text-to-image curl example            |
-| `run_curl_image_edit.sh`    | Image-to-image (editing) curl example |
-| `openai_chat_client.py`     | Python client (t2i + i2i)             |
-
 ## FAQ
 
-- If you encounter OOM errors, adjust `gpu_memory_utilization` in the stage config:
+- If you encounter OOM errors, adjust `gpu_memory_utilization` in the deploy config:
 
 ```yaml
-# In glm_image.yaml, reduce from default 0.6:
+# In vllm_omni/deploy/glm_image.yaml, reduce from default 0.6:
 gpu_memory_utilization: 0.5
 ```
 
 - The first request may be slow due to model warmup. Subsequent requests will be faster.
-
-## Example materials
-
-??? abstract "openai_chat_client.py"
-    ``````py
-    --8<-- "examples/online_serving/glm_image/openai_chat_client.py"
-    ``````
-??? abstract "run_curl_image_edit.sh"
-    ``````sh
-    --8<-- "examples/online_serving/glm_image/run_curl_image_edit.sh"
-    ``````
-??? abstract "run_curl_text_to_image.sh"
-    ``````sh
-    --8<-- "examples/online_serving/glm_image/run_curl_text_to_image.sh"
-    ``````
-??? abstract "run_server.sh"
-    ``````sh
-    --8<-- "examples/online_serving/glm_image/run_server.sh"
-    ``````

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -28,11 +28,13 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
 ### Launch individual stages (stage-based CLI)
 
 Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.
+The example below pins Stage 0 to GPU 0 and Stage 1/2 to GPU 1 via
+`CUDA_VISIBLE_DEVICES`.
 
 **1. Stage 0 (Thinker + API server)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --port 8091 \
     --stage-id 0 \
     --omni-master-address 127.0.0.1 \
@@ -42,7 +44,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 **2. Stage 1 (Talker)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -52,7 +54,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 **3. Stage 2 (Code2Wav)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 2 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -75,7 +77,7 @@ that kind of change. Since each command launches one stage, just pass the knob
 directly on that stage command:
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --gpu-memory-utilization 0.5 \

examples/offline_inference/internvla_a1/README.md

@@ -0,0 +1,30 @@
+# InternVLA-A1
+
+Full usage and result-reporting guidance lives in [docs/user_guide/examples/offline_inference/internvla_a1.md](../../../docs/user_guide/examples/offline_inference/internvla_a1.md).
+
+Quick start:
+
+```bash
+export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen
+export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen
+export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct
+# hf tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors
+export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor
+
+bash run.sh --num-samples 1 --num-episodes 0
+bash run.sh --num-episodes 1
+bash collect_results.sh
+```
+
+Expected files under `INTERNVLA_A1_COSMOS_DIR`:
+
+- `encoder.safetensors`
+- `decoder.safetensors`
+
+Reference Hugging Face repo: `tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors`
+
+Key entrypoints:
+
+- `run.sh`: wrapper for offline inference and GT evaluation
+- `collect_results.sh`: collect sample output, latency, metrics, plots, and logs
+- `end2end.py`: underlying Python entrypoint

examples/offline_inference/internvla_a1/collect_results.sh

@@ -0,0 +1,220 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$ROOT_DIR/../../.." && pwd)"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+RESULT_ROOT="${INTERNVLA_A1_RESULT_DIR:-$REPO_ROOT/outputs/internvla_a1/collected_results/$TIMESTAMP}"
+
+: "${INTERNVLA_A1_MODEL_DIR:?Please export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen}"
+: "${INTERNVLA_A1_DATASET_DIR:?Please export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen}"
+: "${INTERNVLA_A1_PROCESSOR_DIR:?Please export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct}"
+: "${INTERNVLA_A1_COSMOS_DIR:?Please export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor}"
+
+EVAL_OUTPUT_DIR="$RESULT_ROOT/eval_outputs"
+GPU_MONITOR_INTERVAL="${GPU_MONITOR_INTERVAL:-1}"
+
+mkdir -p "$RESULT_ROOT"
+
+export INTERNVLA_A1_MODEL_DIR
+export INTERNVLA_A1_DATASET_DIR
+export INTERNVLA_A1_PROCESSOR_DIR
+export INTERNVLA_A1_COSMOS_DIR
+
+write_env_summary() {
+  cat >"$RESULT_ROOT/env_summary.txt" <<EOF
+timestamp=$TIMESTAMP
+repo_root=$REPO_ROOT
+result_root=$RESULT_ROOT
+python=$(command -v python || true)
+python_version=$(python --version 2>&1 || true)
+model_dir=$INTERNVLA_A1_MODEL_DIR
+dataset_dir=$INTERNVLA_A1_DATASET_DIR
+processor_dir=$INTERNVLA_A1_PROCESSOR_DIR
+cosmos_dir=$INTERNVLA_A1_COSMOS_DIR
+pwd=$(pwd)
+EOF
+}
+
+capture_gpu_snapshot() {
+  local output_file="$1"
+  if command -v nvidia-smi >/dev/null 2>&1; then
+    nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu --format=csv >"$output_file" 2>&1 || true
+  else
+    echo "nvidia-smi not found" >"$output_file"
+  fi
+}
+
+start_gpu_monitor() {
+  local output_file="$1"
+  if ! command -v nvidia-smi >/dev/null 2>&1; then
+    echo "nvidia-smi not found" >"$output_file"
+    return 1
+  fi
+
+  (
+    echo "timestamp,index,name,memory.total [MiB],memory.used [MiB],utilization.gpu [%]"
+    while true; do
+      local now
+      now="$(date '+%Y-%m-%d %H:%M:%S')"
+      nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu --format=csv,noheader,nounits \
+        | awk -F',' -v ts="$now" '{gsub(/^[ \t]+|[ \t]+$/, "", $0); print ts "," $0}'
+      sleep "$GPU_MONITOR_INTERVAL"
+    done
+  ) >"$output_file" 2>/dev/null &
+  echo $!
+}
+
+stop_gpu_monitor() {
+  local monitor_pid="${1:-}"
+  if [[ -n "$monitor_pid" ]] && kill -0 "$monitor_pid" 2>/dev/null; then
+    kill "$monitor_pid" 2>/dev/null || true
+    wait "$monitor_pid" 2>/dev/null || true
+  fi
+}
+
+copy_eval_outputs() {
+  local source_dir="$1"
+  local target_dir="$2"
+  local source_real=""
+  local target_real=""
+  source_real="$(realpath "$source_dir" 2>/dev/null || echo "$source_dir")"
+  target_real="$(realpath "$target_dir" 2>/dev/null || echo "$target_dir")"
+  if [[ "$source_real" == "$target_real" ]]; then
+    return 0
+  fi
+  mkdir -p "$target_dir"
+  if [[ -d "$source_dir" ]]; then
+    cp -r "$source_dir"/. "$target_dir"/
+  fi
+}
+
+run_with_artifacts() {
+  local name="$1"
+  shift
+
+  local log_file="$RESULT_ROOT/${name}.log"
+  local time_file="$RESULT_ROOT/${name}_time.txt"
+  local gpu_file="$RESULT_ROOT/${name}_gpu.csv"
+  local status_file="$RESULT_ROOT/${name}_status.txt"
+  local monitor_pid=""
+
+  monitor_pid="$(start_gpu_monitor "$gpu_file" || true)"
+  set +e
+  local exit_code=0
+  local start_ts=0
+  local end_ts=0
+  start_ts="$(date +%s)"
+  if [[ -x /usr/bin/time ]]; then
+    /usr/bin/time -v -o "$time_file" "$@" >"$log_file" 2>&1
+    exit_code=$?
+    if [[ ! -s "$time_file" ]]; then
+      end_ts="$(date +%s)"
+      {
+        echo "timing_mode=usr_bin_time_empty_fallback"
+        echo "start_ts=$start_ts"
+        echo "end_ts=$end_ts"
+        echo "elapsed_seconds=$((end_ts - start_ts))"
+      } >"$time_file"
+    fi
+  else
+    "$@" >"$log_file" 2>&1
+    exit_code=$?
+    end_ts="$(date +%s)"
+    {
+      echo "timing_mode=shell_date"
+      echo "start_ts=$start_ts"
+      echo "end_ts=$end_ts"
+      echo "elapsed_seconds=$((end_ts - start_ts))"
+    } >"$time_file"
+  fi
+  set -e
+  stop_gpu_monitor "$monitor_pid"
+
+  echo "exit_code=$exit_code" >"$status_file"
+  if [[ $exit_code -ne 0 ]]; then
+    echo "[error] command failed for $name, see $log_file"
+    return $exit_code
+  fi
+}
+
+write_manifest() {
+  cat >"$RESULT_ROOT/README.txt" <<EOF
+InternVLA-A1 collected results
+
+Key files:
+- env_summary.txt: environment and path summary
+- sample_run.log / sample_run_time.txt: one-sample functional run
+- forward_benchmark.log / forward_benchmark_time.txt: pure pipeline.forward latency benchmark
+- eval_run.log / eval_run_time.txt: GT evaluation run
+- pytest_e2e.log / pytest_e2e_time.txt: offline e2e pytest result
+- gpu_info_before.csv / gpu_info_after.csv: point-in-time GPU snapshots
+- *_gpu.csv: sampled GPU usage during each run
+- eval_outputs/: copied output directory from the GT evaluation run
+
+Important outputs:
+- forward_benchmark/forward_latency.json
+- eval_outputs/summary.json
+- eval_outputs/registry/log.json
+- eval_outputs/registry/plots/
+EOF
+}
+
+write_skip_artifact() {
+  local name="$1"
+  local reason="$2"
+  echo "$reason" >"$RESULT_ROOT/${name}.log"
+  echo "timing_mode=skipped" >"$RESULT_ROOT/${name}_time.txt"
+  echo "exit_code=0" >"$RESULT_ROOT/${name}_status.txt"
+  echo "skipped_reason=$reason" >>"$RESULT_ROOT/${name}_status.txt"
+  if [[ ! -f "$RESULT_ROOT/${name}_gpu.csv" ]]; then
+    echo "skipped,$reason" >"$RESULT_ROOT/${name}_gpu.csv"
+  fi
+}
+
+write_env_summary
+write_manifest
+capture_gpu_snapshot "$RESULT_ROOT/gpu_info_before.csv"
+
+run_with_artifacts \
+  "sample_run" \
+  bash "$ROOT_DIR/run.sh" \
+  --output-dir "$RESULT_ROOT/sample_outputs" \
+  --num-samples 1 \
+  --num-episodes 0
+
+run_with_artifacts \
+  "forward_benchmark" \
+  python "$ROOT_DIR/end2end.py" \
+  --model-dir "$INTERNVLA_A1_MODEL_DIR" \
+  --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
+  --benchmark-forward \
+  --dtype bfloat16 \
+  --attn-implementation eager \
+  --warmup-iters 3 \
+  --benchmark-iters 10 \
+  --output-dir "$RESULT_ROOT/forward_benchmark"
+
+run_with_artifacts \
+  "eval_run" \
+  bash "$ROOT_DIR/run.sh" \
+  --output-dir "$EVAL_OUTPUT_DIR" \
+  --num-episodes 1
+
+copy_eval_outputs "$EVAL_OUTPUT_DIR" "$RESULT_ROOT/eval_outputs"
+
+if python - <<'PY' >/dev/null 2>&1
+import importlib.util
+raise SystemExit(0 if importlib.util.find_spec("pytest") else 1)
+PY
+then
+  run_with_artifacts \
+    "pytest_e2e" \
+    python -m pytest -sv tests/e2e/offline_inference/test_internvla_a1.py -m advanced_model
+else
+  write_skip_artifact "pytest_e2e" "pytest is not installed in the current python environment"
+fi
+
+capture_gpu_snapshot "$RESULT_ROOT/gpu_info_after.csv"
+
+echo "Results written to: $RESULT_ROOT"

examples/offline_inference/internvla_a1/end2end.py

@@ -0,0 +1,327 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import statistics
+import sys
+import time
+from pathlib import Path
+
+import torch
+
+_WORKSPACE_ROOT = Path(__file__).resolve().parent
+_REPO_ROOT = _WORKSPACE_ROOT.parents[2]
+if str(_REPO_ROOT) not in sys.path:
+    # Keep the shared helper import clean by doing script-only path bootstrap here.
+    sys.path.insert(0, str(_REPO_ROOT))
+
+from internvla_a1_common import (  # noqa: E402
+    A2DOpenLoopDataset,
+    collate_open_loop_samples,
+    make_shared_noise,
+    run_open_loop_evaluation,
+    select_indices,
+    tensor_dtype,
+    tensor_sha256,
+)
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig  # noqa: E402
+from vllm_omni.diffusion.models.internvla_a1 import (  # noqa: E402
+    InternVLAA1Config,
+    InternVLAA1TrainMetadata,
+)
+from vllm_omni.diffusion.models.internvla_a1.config import OBS_STATE  # noqa: E402
+from vllm_omni.diffusion.registry import initialize_model  # noqa: E402
+from vllm_omni.diffusion.request import OmniDiffusionRequest  # noqa: E402
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams  # noqa: E402
+
+
+def _required_path_arg(env_name: str, cli_value: str | None) -> str:
+    value = cli_value or os.getenv(env_name)
+    if not value:
+        raise ValueError(f"Missing required path: set --{env_name.lower().replace('_', '-')} or {env_name}.")
+    return value
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run vLLM InternVLA-A1 inference on a few samples.")
+    parser.add_argument("--model-dir")
+    parser.add_argument("--dataset-dir")
+    parser.add_argument("--num-samples", type=int, default=1)
+    parser.add_argument("--num-episodes", type=int, default=0)
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16")
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--compile-model", action="store_true")
+    parser.add_argument("--attn-implementation", default="eager")
+    parser.add_argument("--enable-regional-compile", action="store_true")
+    parser.add_argument("--enable-warmup", action="store_true")
+    parser.add_argument("--strict-load", action="store_true")
+    parser.add_argument("--output-dir", default="outputs/internvla_a1/vllm_infer")
+    parser.add_argument("--skip-plots", action="store_true")
+    parser.add_argument("--benchmark-forward", action="store_true")
+    parser.add_argument("--warmup-iters", type=int, default=3)
+    parser.add_argument("--benchmark-iters", type=int, default=10)
+    return parser.parse_args()
+
+
+def build_od_config(args: argparse.Namespace, processor_model_name: str) -> OmniDiffusionConfig:
+    return OmniDiffusionConfig(
+        model=str(Path(args.model_dir).resolve()),
+        model_class_name="InternVLAA1Pipeline",
+        dtype=tensor_dtype(args.dtype),
+        custom_pipeline_args={
+            "device": args.device,
+            "dtype": args.dtype,
+            "compile_model": args.compile_model,
+            "attn_implementation": args.attn_implementation,
+            "enable_regional_compile": args.enable_regional_compile,
+            "enable_warmup": args.enable_warmup,
+            "strict_load": args.strict_load,
+            "processor_model_name": processor_model_name,
+        },
+    )
+
+
+def build_dataset(args: argparse.Namespace) -> tuple[A2DOpenLoopDataset, InternVLAA1Config, InternVLAA1TrainMetadata]:
+    model_dir = Path(args.model_dir)
+    config = InternVLAA1Config.from_pretrained(model_dir)
+    config.device = args.device
+    config.dtype = args.dtype
+    config.compile_model = args.compile_model
+    config.attn_implementation = args.attn_implementation
+    config.enable_regional_compile = args.enable_regional_compile
+
+    train_meta = InternVLAA1TrainMetadata.from_pretrained(model_dir)
+    with open(model_dir / "stats.json", encoding="utf-8") as f:
+        train_stats = json.load(f)["a2d"]
+
+    dataset = A2DOpenLoopDataset(
+        args.dataset_dir,
+        config=config,
+        train_stats=train_stats,
+    )
+    return dataset, config, train_meta
+
+
+def run_one_path(
+    pipeline,
+    dataset: A2DOpenLoopDataset,
+    config: InternVLAA1Config,
+    args: argparse.Namespace,
+    indices: list[int],
+) -> list[dict[str, object]]:
+    results: list[dict[str, object]] = []
+    for index in indices:
+        sample = dataset.get_sample(index)
+        batch_inputs, _ = collate_open_loop_samples([sample], device=args.device, dtype=tensor_dtype(args.dtype))
+        noise = make_shared_noise(
+            args.seed,
+            index,
+            (
+                batch_inputs[OBS_STATE].shape[0],
+                config.chunk_size,
+                config.max_action_dim,
+            ),
+            args.device,
+        )
+        pred = run_pipeline_forward(pipeline, batch_inputs, noise)
+        pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+        results.append(
+            {
+                "path": "registry",
+                "index": index,
+                "episode_index": sample.episode_index,
+                "task": sample.task,
+                "seed": args.seed,
+                "shape": list(pred.shape),
+                "mean": float(pred.mean().item()),
+                "std": float(pred.std().item()),
+                "action_sha256": tensor_sha256(pred),
+                "first_action_prefix": pred[0, 0, :8].tolist(),
+            }
+        )
+    return results
+
+
+def run_pipeline_forward(
+    pipeline,
+    batch_inputs: dict[str, torch.Tensor],
+    noise: torch.Tensor,
+) -> torch.Tensor:
+    output = pipeline.forward(
+        OmniDiffusionRequest(
+            prompts=[""],
+            sampling_params=OmniDiffusionSamplingParams(
+                extra_args={
+                    "batch_inputs": batch_inputs,
+                    "noise": noise,
+                    "decode_image": False,
+                }
+            ),
+        )
+    )
+    if output.error:
+        raise RuntimeError(output.error)
+    if output.output is None:
+        raise RuntimeError("InternVLAA1Pipeline.forward returned no output tensor.")
+    return output.output
+
+
+def _synchronize(device: str) -> None:
+    if device.startswith("cuda") and torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+
+def _latency_summary(values_ms: list[float]) -> dict[str, float]:
+    sorted_values = sorted(values_ms)
+    p50_index = min(len(sorted_values) - 1, round(0.50 * (len(sorted_values) - 1)))
+    p90_index = min(len(sorted_values) - 1, round(0.90 * (len(sorted_values) - 1)))
+    return {
+        "mean_ms": float(statistics.mean(sorted_values)),
+        "stdev_ms": float(statistics.pstdev(sorted_values)) if len(sorted_values) > 1 else 0.0,
+        "min_ms": float(sorted_values[0]),
+        "max_ms": float(sorted_values[-1]),
+        "p50_ms": float(sorted_values[p50_index]),
+        "p90_ms": float(sorted_values[p90_index]),
+    }
+
+
+def benchmark_forward(
+    pipeline,
+    dataset: A2DOpenLoopDataset,
+    config: InternVLAA1Config,
+    args: argparse.Namespace,
+    index: int,
+    output_dir: Path,
+) -> dict[str, object]:
+    sample = dataset.get_sample(index)
+    batch_inputs, _ = collate_open_loop_samples([sample], device=args.device, dtype=tensor_dtype(args.dtype))
+    noise = make_shared_noise(
+        args.seed,
+        index,
+        (
+            batch_inputs[OBS_STATE].shape[0],
+            config.chunk_size,
+            config.max_action_dim,
+        ),
+        args.device,
+    )
+
+    _synchronize(args.device)
+    cold_start_begin = time.perf_counter()
+    pred = run_pipeline_forward(pipeline, batch_inputs, noise)
+    _synchronize(args.device)
+    cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0
+
+    warmup_ms: list[float] = []
+    for _ in range(args.warmup_iters):
+        _synchronize(args.device)
+        begin = time.perf_counter()
+        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _synchronize(args.device)
+        warmup_ms.append((time.perf_counter() - begin) * 1000.0)
+
+    benchmark_ms: list[float] = []
+    for _ in range(args.benchmark_iters):
+        _synchronize(args.device)
+        begin = time.perf_counter()
+        _ = run_pipeline_forward(pipeline, batch_inputs, noise)
+        _synchronize(args.device)
+        benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
+
+    pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+    summary = {
+        "mode": "forward_latency",
+        "model_dir": str(Path(args.model_dir).resolve()),
+        "dataset_dir": str(Path(args.dataset_dir).resolve()),
+        "sample_index": index,
+        "episode_index": sample.episode_index,
+        "task": sample.task,
+        "device": args.device,
+        "dtype": args.dtype,
+        "attn_implementation": args.attn_implementation,
+        "enable_regional_compile": args.enable_regional_compile,
+        "warmup_iters": args.warmup_iters,
+        "benchmark_iters": args.benchmark_iters,
+        "output_shape": list(pred.shape),
+        "cold_start_ms": cold_start_ms,
+        "warmup_summary": _latency_summary(warmup_ms) if warmup_ms else {},
+        "benchmark_summary": _latency_summary(benchmark_ms) if benchmark_ms else {},
+        "benchmark_samples_ms": benchmark_ms,
+    }
+    with open(output_dir / "forward_latency.json", "w", encoding="utf-8") as f:
+        json.dump(summary, f, ensure_ascii=False, indent=2)
+    print(json.dumps(summary, ensure_ascii=False, indent=2))
+    return summary
+
+
+def main() -> None:
+    args = parse_args()
+    args.model_dir = _required_path_arg("INTERNVLA_A1_MODEL_DIR", args.model_dir)
+    args.dataset_dir = _required_path_arg("INTERNVLA_A1_DATASET_DIR", args.dataset_dir)
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    dataset, config, train_meta = build_dataset(args)
+    indices = select_indices(dataset, args.num_samples)
+    od_config = build_od_config(args, train_meta.processor_model_name)
+
+    eval_summaries: dict[str, object] = {}
+    pipeline = initialize_model(od_config)
+    if args.benchmark_forward:
+        benchmark_forward(
+            pipeline=pipeline,
+            dataset=dataset,
+            config=config,
+            args=args,
+            index=indices[0],
+            output_dir=output_dir,
+        )
+        return
+    results = run_one_path(
+        pipeline=pipeline,
+        dataset=dataset,
+        config=config,
+        args=args,
+        indices=indices,
+    )
+    if args.num_episodes > 0:
+        eval_summaries["registry"] = run_open_loop_evaluation(
+            mode="vllm_registry",
+            policy=pipeline,
+            config=config,
+            dataset=dataset,
+            train_meta=train_meta,
+            collate_samples=collate_open_loop_samples,
+            run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise),
+            num_episodes=args.num_episodes,
+            seed=args.seed,
+            device=args.device,
+            dtype=tensor_dtype(args.dtype),
+            output_dir=output_dir / "registry",
+            skip_plots=args.skip_plots,
+        )
+
+    summary = {
+        "mode": "registry",
+        "model_dir": str(Path(args.model_dir).resolve()),
+        "dataset_dir": str(Path(args.dataset_dir).resolve()),
+        "device": args.device,
+        "dtype": args.dtype,
+        "attn_implementation": args.attn_implementation,
+        "enable_regional_compile": args.enable_regional_compile,
+        "seed": args.seed,
+        "indices": indices,
+        "results": results,
+        "output_dir": str(output_dir.resolve()),
+        "eval_summaries": eval_summaries,
+    }
+    with open(output_dir / "summary.json", "w", encoding="utf-8") as f:
+        json.dump(summary, f, ensure_ascii=False, indent=2)
+    print(json.dumps(summary, ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/internvla_a1/internvla_a1_common.py

@@ -0,0 +1,471 @@
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+import torch.nn.functional as F
+import torchvision
+
+from vllm_omni.diffusion.models.internvla_a1.config import (  # noqa: E402
+    OBS_IMAGES,
+    OBS_STATE,
+    OBS_TASK,
+    InternVLAA1Config,
+)
+
+
+def tensor_dtype(name: str) -> torch.dtype:
+    return torch.bfloat16 if name == "bfloat16" else torch.float32
+
+
+def select_indices(dataset, num_samples: int) -> list[int]:
+    indices: list[int] = []
+    for _, episode_indices in dataset.episode_start_indices(max_episodes=num_samples):
+        if episode_indices:
+            indices.append(episode_indices[0])
+    return indices[:num_samples]
+
+
+def make_shared_noise(seed: int, sample_index: int, shape: tuple[int, ...], device: str) -> torch.Tensor:
+    generator = torch.Generator(device="cpu")
+    generator.manual_seed(seed + sample_index)
+    noise = torch.randn(shape, generator=generator, dtype=torch.float32)
+    return noise.to(device=device)
+
+
+def tensor_sha256(tensor: torch.Tensor) -> str:
+    array = tensor.detach().contiguous().cpu().numpy()
+    return hashlib.sha256(array.tobytes()).hexdigest()
+
+
+def _clamp_index(index: int, start: int, end: int) -> int:
+    return max(start, min(end - 1, index))
+
+
+def _load_parquet_rows(path: Path) -> list[dict[str, Any]]:
+    return pq.read_table(path).to_pylist()
+
+
+def _load_json(path: Path) -> dict[str, Any]:
+    with open(path, encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _stack_stats(stats: dict[str, Any], keys: list[str]) -> dict[str, torch.Tensor]:
+    result = {}
+    for stat_name in ("mean", "std"):
+        values = []
+        for key in keys:
+            values.extend(stats[key][stat_name])
+        result[stat_name] = torch.tensor(values, dtype=torch.float32)
+    return result
+
+
+def normalize_vector(values: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
+    denom = torch.where(stats["std"] == 0, torch.ones_like(stats["std"]), stats["std"])
+    return (values - stats["mean"]) / denom
+
+
+def unnormalize_vector(values: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
+    return values * stats["std"] + stats["mean"]
+
+
+class TorchvisionVideoReaderCache:
+    def __init__(self, backend: str = "pyav") -> None:
+        self.backend = backend
+        self._readers: dict[str, Any] = {}
+        torchvision.set_video_backend(backend)
+
+    def get(self, path: str) -> Any:
+        reader = self._readers.get(path)
+        if reader is None:
+            reader = torchvision.io.VideoReader(path, "video")
+            self._readers[path] = reader
+        return reader
+
+    def decode_frames(self, path: str, timestamps: list[float], tolerance_s: float = 1e-4) -> torch.Tensor:
+        reader = self.get(path)
+        first_ts = min(timestamps)
+        last_ts = max(timestamps)
+        reader.seek(first_ts, keyframes_only=self.backend == "pyav")
+
+        loaded_frames: list[torch.Tensor] = []
+        loaded_ts: list[float] = []
+        for frame in reader:
+            current_ts = float(frame["pts"])
+            loaded_frames.append(frame["data"])
+            loaded_ts.append(current_ts)
+            if current_ts >= last_ts:
+                break
+
+        query_ts = torch.tensor(timestamps, dtype=torch.float32)
+        loaded_ts_tensor = torch.tensor(loaded_ts, dtype=torch.float32)
+        distances = torch.cdist(query_ts[:, None], loaded_ts_tensor[:, None], p=1)
+        min_dist, argmin = distances.min(dim=1)
+        if not torch.all(min_dist < tolerance_s):
+            raise RuntimeError(
+                f"Video timestamps are outside tolerance: query={query_ts.tolist()} "
+                f"loaded={loaded_ts_tensor.tolist()} path={path}"
+            )
+        return torch.stack([loaded_frames[i] for i in argmin]).float() / 255.0
+
+
+@dataclass
+class A2DOpenLoopSample:
+    index: int
+    episode_index: int
+    task: str
+    state_raw: torch.Tensor
+    action_raw: torch.Tensor
+    inputs: dict[str, torch.Tensor]
+
+
+class A2DOpenLoopDataset:
+    image_keys = [
+        "observation.images.head",
+        "observation.images.hand_left",
+        "observation.images.hand_right",
+    ]
+    state_keys = [
+        "observation.states.joint.position",
+        "observation.states.effector.position",
+    ]
+    action_keys = [
+        "actions.joint.position",
+        "actions.effector.position",
+    ]
+
+    def __init__(
+        self,
+        dataset_root: str | Path,
+        *,
+        config: InternVLAA1Config,
+        train_stats: dict[str, Any],
+        image_offsets: tuple[int, int] = (-15, 0),
+        tolerance_s: float = 1e-4,
+    ) -> None:
+        self.root = Path(dataset_root)
+        self.config = config
+        self.info = _load_json(self.root / "meta" / "info.json")
+        self.dataset_stats = _load_json(self.root / "meta" / "stats.json")
+        self.data_rows = _load_parquet_rows(self.root / "data" / "chunk-000" / "file-000.parquet")
+        self.episode_rows = _load_parquet_rows(self.root / "meta" / "episodes" / "chunk-000" / "file-000.parquet")
+        self.task_rows = _load_parquet_rows(self.root / "meta" / "tasks.parquet")
+
+        self.state_stats = _stack_stats(train_stats, self.state_keys)
+        self.action_stats = _stack_stats(train_stats, self.action_keys)
+        self.image_offsets = image_offsets
+        self.tolerance_s = tolerance_s
+        self.video_reader = TorchvisionVideoReaderCache(backend="pyav")
+
+    @property
+    def num_episodes(self) -> int:
+        return len(self.episode_rows)
+
+    @property
+    def physical_action_dim(self) -> int:
+        return 16
+
+    def episode_start_indices(self, max_episodes: int | None = None) -> list[tuple[int, list[int]]]:
+        rows = self.episode_rows if max_episodes is None else self.episode_rows[:max_episodes]
+        result = []
+        for ep in rows:
+            start = int(ep["dataset_from_index"])
+            end = int(ep["dataset_to_index"])
+            result.append((int(ep["episode_index"]), list(range(start, end, self.config.chunk_size))))
+        return result
+
+    def _task_text(self, task_index: int) -> str:
+        return self.task_rows[task_index]["__index_level_0__"]
+
+    def _episode_for_index(self, row: dict[str, Any]) -> dict[str, Any]:
+        return self.episode_rows[int(row["episode_index"])]
+
+    def _state_vector(self, row: dict[str, Any]) -> torch.Tensor:
+        return torch.tensor(row[self.state_keys[0]] + row[self.state_keys[1]], dtype=torch.float32)
+
+    def _action_vector(self, row: dict[str, Any]) -> torch.Tensor:
+        return torch.tensor(row[self.action_keys[0]] + row[self.action_keys[1]], dtype=torch.float32)
+
+    def _query_rows(self, idx: int, deltas: list[int]) -> list[dict[str, Any]]:
+        row = self.data_rows[idx]
+        episode = self._episode_for_index(row)
+        start = int(episode["dataset_from_index"])
+        end = int(episode["dataset_to_index"])
+        return [self.data_rows[_clamp_index(idx + delta, start, end)] for delta in deltas]
+
+    def _decode_camera_history(
+        self, episode: dict[str, Any], camera_key: str, rows: list[dict[str, Any]]
+    ) -> torch.Tensor:
+        timestamps = [float(r["timestamp"]) for r in rows]
+        shifted = [float(episode[f"videos/{camera_key}/from_timestamp"]) + ts for ts in timestamps]
+        chunk_idx = int(episode[f"videos/{camera_key}/chunk_index"])
+        file_idx = int(episode[f"videos/{camera_key}/file_index"])
+        path = self.root / self.info["video_path"].format(
+            video_key=camera_key,
+            chunk_index=chunk_idx,
+            file_index=file_idx,
+        )
+        frames = self.video_reader.decode_frames(str(path), shifted, tolerance_s=self.tolerance_s)
+        return frames
+
+    def get_sample(self, idx: int) -> A2DOpenLoopSample:
+        row = self.data_rows[idx]
+        episode = self._episode_for_index(row)
+        image_rows = self._query_rows(idx, list((-15, 0)))
+        action_rows = self._query_rows(idx, list(range(self.config.chunk_size)))
+        camera_images = [self._decode_camera_history(episode, camera_key, image_rows) for camera_key in self.image_keys]
+        state_raw = self._state_vector(row)
+        state_norm = normalize_vector(state_raw, self.state_stats)
+        action_raw = torch.stack([self._action_vector(action_row) for action_row in action_rows], dim=0)
+        task = self._task_text(int(row["task_index"]))
+        inputs = {
+            OBS_STATE: state_norm,
+            OBS_TASK: task,
+            f"{OBS_IMAGES}.image0": camera_images[0],
+            f"{OBS_IMAGES}.image1": camera_images[1],
+            f"{OBS_IMAGES}.image2": camera_images[2],
+            f"{OBS_IMAGES}.image0_mask": torch.tensor(True),
+            f"{OBS_IMAGES}.image1_mask": torch.tensor(True),
+            f"{OBS_IMAGES}.image2_mask": torch.tensor(True),
+        }
+        return A2DOpenLoopSample(
+            index=idx,
+            episode_index=int(row["episode_index"]),
+            task=task,
+            state_raw=state_raw,
+            action_raw=action_raw,
+            inputs=inputs,
+        )
+
+
+def collate_open_loop_samples(
+    samples: list[A2DOpenLoopSample],
+    *,
+    device: str,
+    dtype: torch.dtype,
+) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
+    first = samples[0]
+    batch_inputs: dict[str, torch.Tensor] = {}
+    for key in first.inputs:
+        values = [sample.inputs[key] for sample in samples]
+        if isinstance(values[0], torch.Tensor):
+            tensor = torch.stack(values, dim=0)
+            if tensor.dtype in (torch.int64, torch.bool):
+                batch_inputs[key] = tensor.to(device=device)
+            else:
+                batch_inputs[key] = tensor.to(device=device, dtype=dtype)
+        else:
+            batch_inputs[key] = values
+
+    metadata = {
+        "indices": [sample.index for sample in samples],
+        "episode_indices": [sample.episode_index for sample in samples],
+        "tasks": [sample.task for sample in samples],
+        "state_raw": torch.stack([sample.state_raw for sample in samples], dim=0),
+        "action_raw": torch.stack([sample.action_raw for sample in samples], dim=0),
+    }
+    return batch_inputs, metadata
+
+
+def plot_prediction_series(
+    *,
+    series: dict[str, np.ndarray],
+    output_path: Path,
+    title: str,
+) -> None:
+    import matplotlib.pyplot as plt
+
+    first_series = next(iter(series.values()))
+    num_dims = int(first_series.shape[1])
+    num_cols = 2
+    num_rows = math.ceil(num_dims / num_cols)
+    fig, axs = plt.subplots(num_rows, num_cols, figsize=(16, max(6, 3 * num_rows)))
+    axs = np.array(axs).reshape(-1)
+    x_values = np.arange(first_series.shape[0])
+    styles = [
+        ("Ground Truth", "blue", "-", 1.5),
+        ("VLLM", "green", "-.", 1.4),
+        ("Predicted", "red", "--", 1.4),
+    ]
+    style_map = {name: (color, linestyle, linewidth) for name, color, linestyle, linewidth in styles}
+
+    for dim in range(num_dims):
+        ax = axs[dim]
+        for name, values in series.items():
+            color, linestyle, linewidth = style_map.get(name, (None, "-", 1.2))
+            ax.plot(
+                x_values,
+                values[:, dim],
+                label=name,
+                color=color,
+                linestyle=linestyle,
+                linewidth=linewidth,
+            )
+        ax.set_title(f"Dimension {dim + 1}")
+        ax.set_xlabel("Time Step / Sample Index")
+        ax.set_ylabel(f"Value Dim {dim + 1}")
+        ax.legend(loc="upper right")
+        ax.grid(True, linestyle="--", alpha=0.7)
+
+    for dim in range(num_dims, len(axs)):
+        axs[dim].axis("off")
+
+    fig.suptitle(title, fontsize=16)
+    plt.tight_layout(rect=(0, 0, 1, 0.97))
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path)
+    plt.close(fig)
+
+
+def summarize_prediction_metrics(
+    *,
+    gt_tensor: torch.Tensor,
+    pred_tensor: torch.Tensor,
+    joint_dims: int = 14,
+) -> dict[str, float]:
+    return {
+        "mse": float(F.mse_loss(gt_tensor, pred_tensor, reduction="mean").item()),
+        "mae": float(F.l1_loss(gt_tensor, pred_tensor, reduction="mean").item()),
+        "mse_joint": float(F.mse_loss(gt_tensor[:, :joint_dims], pred_tensor[:, :joint_dims], reduction="mean").item()),
+        "mae_joint": float(F.l1_loss(gt_tensor[:, :joint_dims], pred_tensor[:, :joint_dims], reduction="mean").item()),
+        "mse_gripper": float(
+            F.mse_loss(gt_tensor[:, joint_dims:], pred_tensor[:, joint_dims:], reduction="mean").item()
+        ),
+        "mae_gripper": float(
+            F.l1_loss(gt_tensor[:, joint_dims:], pred_tensor[:, joint_dims:], reduction="mean").item()
+        ),
+    }
+
+
+def summarize_metric_list(values: list[float]) -> float | None:
+    return float(np.mean(values)) if values else None
+
+
+def run_open_loop_evaluation(
+    *,
+    mode: str,
+    policy,
+    config,
+    dataset,
+    train_meta,
+    collate_samples: Callable[..., tuple[dict[str, torch.Tensor], dict[str, Any]]],
+    run_sample_actions: Callable[[Any, dict[str, torch.Tensor], torch.Tensor], torch.Tensor],
+    num_episodes: int,
+    seed: int,
+    device: str,
+    dtype: torch.dtype,
+    output_dir: str | Path,
+    skip_plots: bool,
+) -> dict[str, Any]:
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    plots_dir = output_dir / "plots"
+    plots_dir.mkdir(parents=True, exist_ok=True)
+
+    mse_values: list[float] = []
+    mae_values: list[float] = []
+    mse_joint_values: list[float] = []
+    mae_joint_values: list[float] = []
+    mse_gripper_values: list[float] = []
+    mae_gripper_values: list[float] = []
+    per_episode: list[dict[str, Any]] = []
+
+    plotting_available = True
+    if not skip_plots:
+        try:
+            import matplotlib.pyplot as _plt  # noqa: F401
+        except ModuleNotFoundError:
+            plotting_available = False
+            print("[warning] matplotlib is not installed; plots will be skipped.")
+
+    episode_specs = dataset.episode_start_indices(max_episodes=num_episodes)
+    for episode_index, indices in episode_specs:
+        print(f"[{mode}] episode: {episode_index}")
+        pred_chunks = []
+        gt_chunks = []
+        visited_indices = []
+        task = None
+
+        for index in indices:
+            sample = dataset.get_sample(index)
+            task = sample.task
+            batch_inputs, meta = collate_samples([sample], device=device, dtype=dtype)
+            noise = make_shared_noise(
+                seed,
+                index,
+                (batch_inputs[OBS_STATE].shape[0], config.chunk_size, config.max_action_dim),
+                device,
+            )
+            with torch.no_grad():
+                pred = run_sample_actions(policy, batch_inputs, noise)
+
+            pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+            pred_phys = unnormalize_vector(pred, dataset.action_stats)
+            if train_meta.action_mode == "delta":
+                pred_phys[:, :, :14] += meta["state_raw"][:, None, :14]
+
+            gt_phys = meta["action_raw"][:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
+            pred_chunks.append(pred_phys[0])
+            gt_chunks.append(gt_phys[0])
+            visited_indices.append(index)
+
+        pred_tensor = torch.cat(pred_chunks, dim=0)
+        gt_tensor = torch.cat(gt_chunks, dim=0)
+        metrics = summarize_prediction_metrics(gt_tensor=gt_tensor, pred_tensor=pred_tensor)
+
+        mse_values.append(metrics["mse"])
+        mae_values.append(metrics["mae"])
+        mse_joint_values.append(metrics["mse_joint"])
+        mae_joint_values.append(metrics["mae_joint"])
+        mse_gripper_values.append(metrics["mse_gripper"])
+        mae_gripper_values.append(metrics["mae_gripper"])
+
+        if not skip_plots and plotting_available:
+            plot_prediction_series(
+                series={
+                    "Ground Truth": gt_tensor.numpy(),
+                    "VLLM": pred_tensor.numpy(),
+                },
+                output_path=plots_dir / f"{mode}_open_loop_ep{episode_index}.jpg",
+                title=f"{mode} Ground Truth vs VLLM",
+            )
+
+        episode_log = {
+            "episode_id": int(episode_index),
+            "task": task,
+            "visited_indices": visited_indices,
+            "num_pred_steps": int(pred_tensor.shape[0]),
+            **metrics,
+        }
+        per_episode.append(episode_log)
+        print(json.dumps({"mode": mode, **episode_log}, ensure_ascii=False))
+
+    summary = {
+        "mode": mode,
+        "num_episodes": len(per_episode),
+        "mse": mse_values,
+        "mae": mae_values,
+        "average_mse": summarize_metric_list(mse_values),
+        "average_mae": summarize_metric_list(mae_values),
+        "mse_joint": mse_joint_values,
+        "mae_joint": mae_joint_values,
+        "average_mse_joint": summarize_metric_list(mse_joint_values),
+        "average_mae_joint": summarize_metric_list(mae_joint_values),
+        "mse_gripper": mse_gripper_values,
+        "mae_gripper": mae_gripper_values,
+        "average_mse_gripper": summarize_metric_list(mse_gripper_values),
+        "average_mae_gripper": summarize_metric_list(mae_gripper_values),
+        "episodes": per_episode,
+    }
+    with open(output_dir / "log.json", "w", encoding="utf-8") as f:
+        json.dump(summary, f, indent=2, ensure_ascii=False)
+    return summary

examples/offline_inference/internvla_a1/run.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+REPO_ROOT="$(cd "$ROOT_DIR/../../.." && pwd)"
+ : "${INTERNVLA_A1_MODEL_DIR:?Please export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen}"
+ : "${INTERNVLA_A1_DATASET_DIR:?Please export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen}"
+ : "${INTERNVLA_A1_PROCESSOR_DIR:?Please export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct}"
+ : "${INTERNVLA_A1_COSMOS_DIR:?Please export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor}"
+INTERNVLA_A1_OUTPUT_DIR="${INTERNVLA_A1_OUTPUT_DIR:-$REPO_ROOT/outputs/internvla_a1/vllm_infer}"
+
+export INTERNVLA_A1_MODEL_DIR
+export INTERNVLA_A1_DATASET_DIR
+export INTERNVLA_A1_PROCESSOR_DIR
+export INTERNVLA_A1_COSMOS_DIR
+
+python "$ROOT_DIR/end2end.py" \
+  --model-dir "$INTERNVLA_A1_MODEL_DIR" \
+  --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
+  --output-dir "$INTERNVLA_A1_OUTPUT_DIR" \
+  --num-episodes "${INTERNVLA_A1_NUM_EPISODES:-1}" \
+  --attn-implementation eager \
+  "$@"

examples/offline_inference/moss_tts_nano/README.md

@@ -0,0 +1,101 @@
+# MOSS-TTS-Nano Offline Inference
+
+## Overview
+
+Single-stage offline TTS pipeline using the 0.1B MOSS-TTS-Nano AR LM and MOSS-Audio-Tokenizer-Nano codec. Outputs 48 kHz stereo WAV.
+
+## Quick Start
+
+```bash
+python end2end.py --text "Hello, this is MOSS-TTS-Nano."
+```
+
+The first run downloads `OpenMOSS-Team/MOSS-TTS-Nano` and `OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano` from Hugging Face.
+
+## Usage
+
+```
+python end2end.py [OPTIONS]
+
+Options:
+  --text TEXT               Text to synthesize (default: "Hello, this is MOSS-TTS-Nano speaking.")
+  --voice VOICE             Built-in voice preset (default: Junhao). See voice table below.
+  --mode MODE               voice_clone (default) or continuation
+  --prompt-audio PATH       Reference WAV/MP3 for custom voice cloning
+  --prompt-text TEXT        Reference transcript (continuation mode)
+  --max-new-frames N        Max AR frames, default 375 (~14 s audio)
+  --seed INT                Random seed for reproducibility
+  --audio-temperature F     Audio sampling temperature (default: 0.8)
+  --audio-top-k N           Audio top-k sampling (default: 25)
+  --audio-top-p F           Audio top-p sampling (default: 0.95)
+  --text-temperature F      Text layer temperature (default: 1.0)
+  --batch                   Run a built-in batch of diverse samples (ZH/EN/FR)
+  --output-dir DIR          Directory for WAV outputs (default: /tmp/moss_tts_nano_output)
+  --deploy-config PATH      Override deploy YAML (defaults to vllm_omni/deploy/moss_tts_nano.yaml)
+  --stage-init-timeout INT  Timeout in seconds for stage init (default: 120)
+```
+
+## Examples
+
+```bash
+# Built-in Chinese voice
+python end2end.py --text "你好，这是MOSS-TTS-Nano的语音合成演示。" --voice Junhao
+
+# Built-in English voice
+python end2end.py --text "Hello from MOSS-TTS-Nano." --voice Ava
+
+# Custom voice clone
+python end2end.py \
+    --text "Hello, this is a cloned voice." \
+    --prompt-audio /path/to/reference.wav \
+    --prompt-text "Exact transcript of the reference audio."
+
+# Batch synthesis (ZH/EN/FR)
+python end2end.py --batch --output-dir /tmp/batch_output
+
+# Reproducible output
+python end2end.py --text "Deterministic test." --seed 42
+```
+
+## Built-in Voice Presets
+
+| Voice | Language |
+|-------|----------|
+| `Junhao` | ZH |
+| `Zhiming` | ZH |
+| `Weiguo` | ZH |
+| `Xiaoyu` | ZH |
+| `Yuewen` | ZH |
+| `Lingyu` | ZH |
+| `Ava` | EN |
+| `Bella` | EN |
+| `Adam` | EN |
+| `Nathan` | EN |
+| `Sakura` | JA |
+| `Yui` | JA |
+| `Aoi` | JA |
+| `Hina` | JA |
+| `Mei` | JA |
+
+## Deploy Config
+
+Runtime knobs live in `vllm_omni/deploy/moss_tts_nano.yaml` (auto-loaded;
+override with `--deploy-config PATH`). Key stage-level settings:
+
+```yaml
+stages:
+  - stage_id: 0
+    gpu_memory_utilization: 0.3   # ~2 GB VRAM; increase for faster init
+    max_num_seqs: 4               # concurrent requests
+    max_model_len: 4096
+```
+
+## Output Format
+
+WAV files, 48 kHz, stereo (2-channel). The codec interleaves stereo as `[L, R, L, R, ...]` in the flat tensor returned by the model.
+
+## Troubleshooting
+
+- **`libnvrtc.so.13: cannot open shared object file`**: torchaudio 2.10+ torchcodec backend requires NVRTC. The model patches `torchaudio.load/save` automatically at load time to fall back to soundfile.
+- **`flash_attn not installed`**: The model falls back to `sdpa` attention automatically.
+- **Empty audio**: Check that `--text` is non-empty and the model loaded successfully (look for "MOSS-TTS-Nano LM loaded" in logs).

examples/offline_inference/moss_tts_nano/end2end.py

@@ -0,0 +1,180 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline inference example for MOSS-TTS-Nano via vLLM-Omni.
+
+Single-stage pipeline: the 0.1B AR LM and MOSS-Audio-Tokenizer-Nano codec
+both run inside one generation stage.  Output is 48 kHz stereo WAV.
+
+Supports:
+  - Built-in voice presets (--voice)
+  - Custom voice cloning via reference audio (--prompt-audio)
+  - Batch synthesis of multiple texts
+  - Reproducible output via --seed
+
+Usage examples:
+  # Built-in voice
+  python end2end.py --text "Hello from MOSS-TTS-Nano."
+
+  # Voice clone with reference audio
+  python end2end.py --text "Hello!" --prompt-audio /path/to/ref.wav
+
+  # Batch with different voices
+  python end2end.py --batch --output-dir /tmp/moss_output
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import soundfile as sf
+import torch
+from vllm import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+# Prevent multiprocessing from re-importing CUDA in the wrong context.
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+from vllm_omni import Omni  # noqa: E402
+
+MODEL = "OpenMOSS-Team/MOSS-TTS-Nano"
+
+BATCH_SAMPLES = [
+    {"text": "Hello, this is a test of MOSS-TTS-Nano.", "voice": "Ava", "label": "en_hello"},
+    {"text": "你好，这是 MOSS-TTS-Nano 的语音合成测试。", "voice": "Junhao", "label": "zh_hello"},
+    {"text": "Bonjour, ceci est un test de synthèse vocale.", "voice": "Bella", "label": "fr_hello"},
+]
+
+
+def build_request(
+    text: str,
+    voice: str = "Junhao",
+    mode: str = "voice_clone",
+    prompt_audio_path: str | None = None,
+    prompt_text: str | None = None,
+    max_new_frames: int = 375,
+    seed: int | None = None,
+    audio_temperature: float = 0.8,
+    audio_top_k: int = 25,
+    audio_top_p: float = 0.95,
+    text_temperature: float = 1.0,
+) -> dict:
+    """Build an Omni request payload for MOSS-TTS-Nano."""
+    additional: dict = {
+        "text": [text],
+        "voice": [voice],
+        "mode": [mode],
+        "max_new_frames": [max_new_frames],
+        "audio_temperature": [audio_temperature],
+        "audio_top_k": [audio_top_k],
+        "audio_top_p": [audio_top_p],
+        "text_temperature": [text_temperature],
+    }
+    if prompt_audio_path:
+        additional["prompt_audio_path"] = [str(prompt_audio_path)]
+    if prompt_text:
+        additional["prompt_text"] = [prompt_text]
+    if seed is not None:
+        additional["seed"] = [seed]
+
+    return {
+        "prompt": "<|im_start|>assistant\n",  # minimal placeholder prompt
+        "additional_information": additional,
+    }
+
+
+def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = 48000) -> None:
+    audio_np = waveform.float().numpy()
+    # Reshape stereo: inference_stream yields interleaved [samples*2]; reshape to [samples, 2]
+    if audio_np.ndim == 1 and audio_np.shape[0] % 2 == 0:
+        audio_np = audio_np.reshape(-1, 2)
+    sf.write(path, audio_np, sample_rate)
+    print(f"  Saved {path} ({audio_np.shape}, {sample_rate} Hz)")
+
+
+def main(args) -> None:
+    omni = Omni(
+        model=MODEL,
+        deploy_config=args.deploy_config,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    sampling_params = SamplingParams(
+        temperature=1.0,
+        top_p=1.0,
+        top_k=50,
+        max_tokens=4096,
+        seed=args.seed if args.seed is not None else 42,
+        detokenize=False,
+    )
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    if args.batch:
+        print(f"Running batch synthesis ({len(BATCH_SAMPLES)} samples)...")
+        inputs = [build_request(s["text"], voice=s["voice"], seed=args.seed) for s in BATCH_SAMPLES]
+        params_list = [sampling_params] * len(inputs)
+    else:
+        print(f"Synthesizing: {args.text!r}")
+        inputs = build_request(
+            text=args.text,
+            voice=args.voice,
+            mode=args.mode,
+            prompt_audio_path=args.prompt_audio,
+            prompt_text=args.prompt_text,
+            max_new_frames=args.max_new_frames,
+            seed=args.seed,
+            audio_temperature=args.audio_temperature,
+            audio_top_k=args.audio_top_k,
+            audio_top_p=args.audio_top_p,
+            text_temperature=args.text_temperature,
+        )
+        params_list = sampling_params
+
+    for stage_outputs in omni.generate(inputs, params_list):
+        for i, req_output in enumerate(stage_outputs.request_output):
+            for j, out in enumerate(req_output.outputs):
+                mm = out.multimodal_output
+                if mm is None:
+                    print(f"  [req {i}] No audio output.")
+                    continue
+                audio = mm.get("audio")
+                sr_tensor = mm.get("sr")
+                if audio is None:
+                    print(f"  [req {i}] No waveform in multimodal_output.")
+                    continue
+                sr = int(sr_tensor.item()) if sr_tensor is not None else 48000
+                label = BATCH_SAMPLES[i]["label"] if args.batch else f"output_{i}_{j}"
+                out_path = str(output_dir / f"{label}.wav")
+                save_audio(audio.cpu(), out_path, sr)
+
+    print("Done.")
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(description="MOSS-TTS-Nano offline inference")
+    parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.")
+    parser.add_argument("--voice", default="Junhao", help="Built-in voice preset name.")
+    parser.add_argument("--mode", default="voice_clone", choices=["voice_clone", "continuation"])
+    parser.add_argument("--prompt-audio", default=None, help="Path to reference audio for voice cloning.")
+    parser.add_argument("--prompt-text", default=None, help="Reference transcript (continuation mode).")
+    parser.add_argument("--max-new-frames", type=int, default=375, help="Max AR frames (~14s at default).")
+    parser.add_argument("--seed", type=int, default=None, help="Random seed.")
+    parser.add_argument("--audio-temperature", type=float, default=0.8)
+    parser.add_argument("--audio-top-k", type=int, default=25)
+    parser.add_argument("--audio-top-p", type=float, default=0.95)
+    parser.add_argument("--text-temperature", type=float, default=1.0)
+    parser.add_argument("--batch", action="store_true", help="Run built-in batch of diverse samples.")
+    parser.add_argument("--output-dir", default="/tmp/moss_tts_nano_output", help="Directory for WAV outputs.")
+    parser.add_argument(
+        "--deploy-config",
+        default=None,
+        help="Path to a deploy YAML; leave unset to auto-load vllm_omni/deploy/moss_tts_nano.yaml.",
+    )
+    parser.add_argument("--stage-init-timeout", type=int, default=120)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    main(parse_args())

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -242,8 +242,7 @@ async def run_single_request(
                 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)
+                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:
@@ -293,7 +292,7 @@ async def run_single_request(
     if text_parts:
         text_file = os.path.join(output_dir, f"{request_id}.txt")
         with open(text_file, "w", encoding="utf-8") as f:
-            f.write("\n".join(text_parts))
+            f.write("".join(text_parts))
         result["saved_files"].append(text_file)
         print(
             f"[Request {request_id}] Text saved to {text_file} "
@@ -393,6 +392,10 @@ async def run_all(args):
 
         # Use default sampling params from stage config (they are pre-configured
         # in the YAML for each stage).
+        #
+        # NOTE: Since we do not set the sampling params directly, .generate in
+        # will automatically set the output kind to delta, since this is what
+        # makes sense for most multimodal use-cases.
         sampling_params_list = None
 
         output_dir = args.output_dir

examples/online_serving/moss_tts_nano/README.md

@@ -0,0 +1,158 @@
+# MOSS-TTS-Nano
+
+## Model checkpoint
+
+| Model | Description |
+|-------|-------------|
+| `OpenMOSS-Team/MOSS-TTS-Nano` | 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec, 48 kHz stereo, ZH/EN/JA |
+
+## Gradio Demo
+
+An interactive Gradio demo is available with multilingual voice presets, custom voice cloning, and streaming support.
+
+```bash
+# Option 1: Launch server + Gradio together
+./run_gradio_demo.sh
+
+# Option 2: If server is already running
+python gradio_demo.py --api-base http://localhost:8091
+```
+
+Then open http://localhost:7860 in your browser.
+
+Features:
+
+- 15 built-in voice presets (6 ZH · 4 EN · 5 JA)
+- Custom voice cloning from uploaded reference audio
+- Streaming mode (progressive PCM playback)
+
+## Launch the Server
+
+```bash
+vllm serve OpenMOSS-Team/MOSS-TTS-Nano --omni --port 8091
+```
+
+The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no
+`--stage-configs-path`, `--trust-remote-code`, or `--enforce-eager` flags
+are needed.
+
+Or use the convenience script:
+
+```bash
+./run_server.sh
+```
+
+## Send TTS Request
+
+### Using curl
+
+```bash
+# Built-in voice preset
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "你好，这是语音合成测试。",
+        "voice": "Junhao",
+        "response_format": "wav"
+    }' --output output.wav
+
+# English preset
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, this is MOSS-TTS-Nano.",
+        "voice": "Ava",
+        "response_format": "wav"
+    }' --output output_en.wav
+```
+
+### Custom Voice Cloning
+
+Provide a reference audio (base64 data URL) and its transcript:
+
+```bash
+REF_AUDIO=$(base64 -w 0 /path/to/reference.wav)
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d "{
+        \"input\": \"Hello, this is a cloned voice.\",
+        \"voice\": \"Junhao\",
+        \"ref_audio\": \"data:audio/wav;base64,${REF_AUDIO}\",
+        \"ref_text\": \"Exact transcript of the reference audio.\"
+    }" --output cloned.wav
+```
+
+### Using Python
+
+```python
+import httpx
+
+response = httpx.post(
+    "http://localhost:8091/v1/audio/speech",
+    json={
+        "input": "你好，这是语音合成测试。",
+        "voice": "Junhao",
+        "response_format": "wav",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+### Streaming
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, streaming output from MOSS-TTS-Nano.",
+        "voice": "Ava",
+        "stream": true,
+        "response_format": "pcm"
+    }' --no-buffer | play -t raw -r 48000 -e signed -b 16 -c 2 -
+```
+
+**Note:** MOSS-TTS-Nano outputs at 48 kHz stereo (2-channel).
+
+## API Parameters
+
+MOSS-TTS-Nano uses the standard `/v1/audio/speech` endpoint.
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text to synthesize (ZH / EN / JA) |
+| `voice` | string | `"Junhao"` | Built-in voice name (see table below) |
+| `response_format` | string | `"wav"` | Audio format: wav, mp3, flac, pcm |
+| `ref_audio` | string | null | Base64 data URL for custom voice cloning |
+| `ref_text` | string | null | Transcript of reference audio (required when `ref_audio` is set) |
+| `stream` | bool | false | Stream raw PCM chunks |
+| `max_new_tokens` | int | 4096 | Maximum tokens to generate |
+
+## Built-in Voice Presets
+
+| Voice | Language | Description |
+|-------|----------|-------------|
+| `Junhao` | ZH | Male, standard Mandarin |
+| `Zhiming` | ZH | Male |
+| `Weiguo` | ZH | Male |
+| `Xiaoyu` | ZH | Female |
+| `Yuewen` | ZH | Female |
+| `Lingyu` | ZH | Female |
+| `Ava` | EN | Female, American English |
+| `Bella` | EN | Female |
+| `Adam` | EN | Male |
+| `Nathan` | EN | Male |
+| `Sakura` | JA | Female |
+| `Yui` | JA | Female |
+| `Aoi` | JA | Female |
+| `Hina` | JA | Female |
+| `Mei` | JA | Female |
+
+## Troubleshooting
+
+1. **`libnvrtc.so.13: cannot open shared object file`**: torchaudio 2.10+ defaults to torchcodec which requires NVRTC. vLLM-Omni patches this automatically at model load time to use soundfile instead.
+2. **Connection refused**: Ensure the server is running on the correct port.
+3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings.
+4. **Out of memory**: The default `gpu_memory_utilization=0.3` is conservative. Increase it in the stage config if you have more VRAM available.

examples/online_serving/moss_tts_nano/gradio_demo.py

@@ -0,0 +1,690 @@
+"""Gradio demo for MOSS-TTS-Nano with gapless streaming audio playback.
+
+Uses a custom AudioWorklet-based player (adapted from the Qwen3-TTS demo)
+for gap-free streaming.  Audio is streamed from the vLLM-Omni server
+through a same-origin proxy and played via the Web Audio API's AudioWorklet.
+
+MOSS-TTS-Nano outputs 48 kHz stereo (2-channel interleaved L,R,L,R,...).
+The AudioWorklet downmixes stereo to mono during playback.
+
+Also supports non-streaming mode (full audio download) via gr.Audio.
+
+Usage:
+    # Start the server first (see run_server.sh), then:
+    python gradio_demo.py --api-base http://localhost:8091
+
+    # Or use run_gradio_demo.sh to start both server and demo together.
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import io
+import json
+import logging
+
+try:
+    import gradio as gr
+except ImportError:
+    raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
+
+import httpx
+import numpy as np
+import soundfile as sf
+from fastapi import FastAPI, Request
+from fastapi.responses import Response, StreamingResponse
+
+logger = logging.getLogger(__name__)
+
+# MOSS-TTS-Nano outputs 48 kHz stereo (2-channel interleaved).
+PCM_SAMPLE_RATE = 48000
+PCM_CHANNELS = 2
+DEFAULT_API_BASE = "http://localhost:8091"
+
+# Built-in voice presets (matches the MOSS-TTS-Nano GitHub repo assets/)
+VOICE_PRESETS = {
+    # Chinese
+    "Junhao (ZH)": "Junhao",
+    "Zhiming (ZH)": "Zhiming",
+    "Weiguo (ZH)": "Weiguo",
+    "Xiaoyu (ZH)": "Xiaoyu",
+    "Yuewen (ZH)": "Yuewen",
+    "Lingyu (ZH)": "Lingyu",
+    # English
+    "Ava (EN)": "Ava",
+    "Bella (EN)": "Bella",
+    "Adam (EN)": "Adam",
+    "Nathan (EN)": "Nathan",
+    # Japanese
+    "Sakura (JA)": "Sakura",
+    "Yui (JA)": "Yui",
+    "Aoi (JA)": "Aoi",
+    "Hina (JA)": "Hina",
+    "Mei (JA)": "Mei",
+}
+PRESET_LABELS = list(VOICE_PRESETS.keys())
+DEFAULT_PRESET = "Junhao (ZH)"
+
+# ── AudioWorklet processor (loaded in browser via Blob URL) ──────────
+# Downmixes stereo-interleaved int16 PCM to mono float32 for playback.
+WORKLET_JS = r"""
+class TTSPlaybackProcessor extends AudioWorkletProcessor {
+    constructor() {
+        super();
+        this.queue = [];
+        this.buf = null;
+        this.pos = 0;
+        this.playing = false;
+        this.played = 0;
+        this.port.onmessage = (e) => {
+            if (e.data && e.data.type === 'clear') {
+                this.queue = []; this.buf = null; this.pos = 0; this.played = 0;
+                if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped'}); }
+                return;
+            }
+            // Input: Int16Array of stereo-interleaved samples [L0,R0,L1,R1,...]
+            // Convert to mono float32 by averaging L+R channels.
+            const pcm = e.data;
+            const monoLen = Math.floor(pcm.length / 2);
+            const mono = new Float32Array(monoLen);
+            for (let i = 0; i < monoLen; i++) {
+                mono[i] = (pcm[i * 2] + pcm[i * 2 + 1]) / 65536.0;
+            }
+            this.queue.push(mono);
+        };
+    }
+    process(inputs, outputs) {
+        const out = outputs[0][0];
+        for (let i = 0; i < out.length; i++) {
+            if (!this.buf || this.pos >= this.buf.length) {
+                if (this.queue.length > 0) {
+                    this.buf = this.queue.shift(); this.pos = 0;
+                } else {
+                    for (let j = i; j < out.length; j++) out[j] = 0;
+                    if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped', played:this.played}); }
+                    return true;
+                }
+            }
+            out[i] = this.buf[this.pos++];
+            this.played++;
+        }
+        if (!this.playing) { this.playing = true; this.port.postMessage({type:'started'}); }
+        return true;
+    }
+}
+registerProcessor('tts-playback-processor', TTSPlaybackProcessor);
+"""
+
+# ── Player HTML (container with metric cards) ────────────────────────
+PLAYER_HTML = """
+<div id="tts-player">
+  <div style="display:flex; align-items:center; gap:10px;">
+    <div id="tts-status-dot" style="width:10px;height:10px;border-radius:50%;background:#ccc;flex-shrink:0;"></div>
+    <span id="tts-status" style="font-weight:600;font-size:1.05em;">Ready</span>
+    <button id="tts-stop-btn" onclick="window.ttsStop()"
+      style="display:none; margin-left:auto; padding:5px 16px; border-radius:6px; border:1px solid #EF5552;
+             background:#fff; color:#EF5552; cursor:pointer; font-size:0.85em; transition:all 0.15s;">Stop</button>
+  </div>
+  <div id="tts-metrics" style="display:none; grid-template-columns:repeat(4,1fr); gap:10px; margin-top:12px;">
+    <div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
+      <div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">TTFP</div>
+      <div id="tts-m-ttfp" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
+    </div>
+    <div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
+      <div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">RTF</div>
+      <div id="tts-m-rtf" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
+    </div>
+    <div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
+      <div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">Audio</div>
+      <div id="tts-m-dur" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
+    </div>
+    <div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
+      <div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">Speed</div>
+      <div id="tts-m-speed" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
+    </div>
+  </div>
+  <div id="tts-rtf-bar-wrap" style="display:none; background:#e8ecf1; border-radius:4px; height:20px; overflow:hidden; position:relative; margin-top:10px;">
+    <div id="tts-rtf-bar" style="height:100%; border-radius:4px; transition:width 0.3s ease, background 0.3s ease; width:0%;"></div>
+    <span id="tts-rtf-label" style="position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);font-size:0.75em;font-weight:600;color:#444;"></span>
+  </div>
+  <div id="tts-elapsed" style="display:none; margin-top:6px; font-size:0.8em; color:#999; text-align:right;"></div>
+</div>
+"""
+
+
+def _build_player_js(sample_rate: int) -> str:
+    """Build the JavaScript that powers the AudioWorklet player."""
+    return f"""
+    <script>
+    const SR = {sample_rate};
+    const WC = {json.dumps(WORKLET_JS)};
+    let ctx = null, node = null, abort = null, gen = false, st = {{}};
+
+    async function init() {{
+        if (ctx) return;
+        ctx = new AudioContext({{ sampleRate: SR }});
+        const b = new Blob([WC], {{ type: 'application/javascript' }});
+        const u = URL.createObjectURL(b);
+        await ctx.audioWorklet.addModule(u);
+        URL.revokeObjectURL(u);
+        node = new AudioWorkletNode(ctx, 'tts-playback-processor');
+        node.connect(ctx.destination);
+        node.port.onmessage = (e) => {{
+            if (e.data.type === 'started') setStatus('Playing...', '#2E7D32');
+            else if (e.data.type === 'stopped' && !gen) {{
+                setStatus('Done', '#2E7D32'); showStats(true);
+                const btn = document.getElementById('tts-stop-btn');
+                if (btn) btn.style.display = 'none';
+            }}
+        }};
+    }}
+
+    function setStatus(text, color) {{
+        const s = document.getElementById('tts-status');
+        const d = document.getElementById('tts-status-dot');
+        if (s) s.textContent = text;
+        if (d) d.style.background = color || '#ccc';
+    }}
+
+    function showStats(fin) {{
+        if (!st.t0) return;
+        const elapsed = (fin && st.streamEnd ? (st.streamEnd - st.t0) : (performance.now() - st.t0)) / 1000;
+        // st.samples counts stereo int16 samples; mono frames = samples / 2
+        const dur = (st.samples / 2) / SR;
+        const mTtfp = document.getElementById('tts-m-ttfp');
+        const mRtf = document.getElementById('tts-m-rtf');
+        const mDur = document.getElementById('tts-m-dur');
+        const mSpeed = document.getElementById('tts-m-speed');
+        const bar = document.getElementById('tts-rtf-bar');
+        const barLabel = document.getElementById('tts-rtf-label');
+        const elapsedEl = document.getElementById('tts-elapsed');
+
+        if (mTtfp && st.ttfp != null) mTtfp.textContent = st.ttfp.toFixed(0) + 'ms';
+        if (mDur) mDur.textContent = dur.toFixed(1) + 's';
+
+        if (dur > 0 && elapsed > 0) {{
+            const rtf = elapsed / dur;
+            const speed = 1 / rtf;
+            if (mRtf) {{
+                mRtf.textContent = rtf.toFixed(2) + 'x';
+                mRtf.style.color = rtf < 1 ? '#2E7D32' : rtf < 1.5 ? '#e8a317' : '#EF5552';
+            }}
+            if (mSpeed) {{
+                mSpeed.textContent = speed.toFixed(1) + 'x';
+                mSpeed.style.color = speed > 1 ? '#2E7D32' : speed > 0.7 ? '#e8a317' : '#EF5552';
+            }}
+            if (bar) {{
+                const pct = Math.min(speed / 3 * 100, 100);
+                bar.style.width = pct + '%';
+                bar.style.background = speed > 1 ? 'linear-gradient(90deg,#66BB6A,#2E7D32)' : speed > 0.7 ? 'linear-gradient(90deg,#e8a317,#f0b866)' : 'linear-gradient(90deg,#EF5552,#f87171)';
+            }}
+            if (barLabel) barLabel.textContent = speed.toFixed(1) + 'x realtime';
+        }}
+        if (elapsedEl) {{
+            elapsedEl.style.display = 'block';
+            elapsedEl.textContent = fin ? 'Completed in ' + elapsed.toFixed(1) + 's  (' + st.chunks + ' chunks)' : elapsed.toFixed(1) + 's elapsed  (' + st.chunks + ' chunks)';
+        }}
+    }}
+
+    window.ttsStop = function() {{
+        if (abort) abort.abort();
+        if (node) node.port.postMessage({{ type: 'clear' }});
+        gen = false;
+        setStatus('Stopped', '#999');
+        const btn = document.getElementById('tts-stop-btn');
+        if (btn) btn.style.display = 'none';
+    }};
+
+    window.ttsGenerate = async function(payload) {{
+        try {{ await init(); if (ctx.state === 'suspended') await ctx.resume(); }}
+        catch (e) {{ const s = document.getElementById('tts-status'); if (s) s.textContent = 'Audio init error: ' + e.message; return; }}
+
+        if (abort) abort.abort();
+        node.port.postMessage({{ type: 'clear' }});
+        await new Promise(r => setTimeout(r, 50));
+        node.port.postMessage({{ type: 'clear' }});
+
+        gen = true;
+        st = {{ t0: null, chunks: 0, samples: 0, ttfp: null }};
+        setStatus('Connecting...', '#2E7D32');
+        const bEl = document.getElementById('tts-stop-btn');
+        if (bEl) bEl.style.display = 'inline-block';
+        const mp = document.getElementById('tts-metrics');
+        if (mp) {{ mp.style.display = 'grid'; ['tts-m-ttfp','tts-m-rtf','tts-m-dur','tts-m-speed'].forEach(id => {{ const e = document.getElementById(id); if(e) {{ e.textContent = '\\u2014'; e.style.color = '#333'; }} }}); }}
+        const bw = document.getElementById('tts-rtf-bar-wrap');
+        if (bw) bw.style.display = 'block';
+        const bar = document.getElementById('tts-rtf-bar');
+        if (bar) bar.style.width = '0%';
+        const bl = document.getElementById('tts-rtf-label');
+        if (bl) bl.textContent = '';
+        const ee = document.getElementById('tts-elapsed');
+        if (ee) {{ ee.style.display = 'none'; ee.textContent = ''; }}
+        abort = new AbortController();
+
+        try {{
+            st.t0 = performance.now();
+            const r = await fetch('/proxy/v1/audio/speech', {{
+                method: 'POST',
+                headers: {{ 'Content-Type': 'application/json' }},
+                body: JSON.stringify(payload),
+                signal: abort.signal,
+            }});
+            if (!r.ok) {{ const t = await r.text(); throw new Error('Server ' + r.status + ': ' + t.slice(0, 200)); }}
+            setStatus('Streaming...', '#43A047');
+
+            const reader = r.body.getReader();
+            let left = new Uint8Array(0);
+            while (true) {{
+                const {{ done, value }} = await reader.read();
+                if (done) break;
+                let raw;
+                if (left.length > 0) {{
+                    raw = new Uint8Array(left.length + value.length);
+                    raw.set(left); raw.set(value, left.length);
+                }} else {{ raw = value; }}
+                const usable = raw.length - (raw.length % 2);
+                left = usable < raw.length ? raw.slice(usable) : new Uint8Array(0);
+                if (usable > 0) {{
+                    const ab = new ArrayBuffer(usable);
+                    new Uint8Array(ab).set(raw.subarray(0, usable));
+                    const pcm = new Int16Array(ab);
+                    node.port.postMessage(pcm);
+                    st.chunks++;
+                    st.samples += pcm.length;
+                    if (st.ttfp == null) st.ttfp = performance.now() - st.t0;
+                    showStats(false);
+                }}
+            }}
+        }} catch (e) {{
+            if (e.name !== 'AbortError') {{
+                setStatus('Error: ' + e.message, '#EF5552');
+                console.error('TTS error:', e);
+            }}
+        }} finally {{
+            st.streamEnd = performance.now();
+            showStats(true);
+            gen = false;
+            if (st.samples > 0) setStatus('Finishing playback...', '#2E7D32');
+            else {{
+                setStatus('No audio received', '#999');
+                if (bEl) bEl.style.display = 'none';
+            }}
+        }}
+    }};
+    </script>
+"""
+
+
+def encode_audio_to_base64(audio_data: tuple) -> str:
+    """Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL."""
+    sample_rate, audio_np = audio_data
+    if audio_np.dtype != np.int16:
+        if audio_np.dtype in (np.float32, np.float64):
+            audio_np = np.clip(audio_np, -1.0, 1.0)
+            audio_np = (audio_np * 32767).astype(np.int16)
+        else:
+            audio_np = audio_np.astype(np.int16)
+    buf = io.BytesIO()
+    sf.write(buf, audio_np, sample_rate, format="WAV")
+    wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:audio/wav;base64,{wav_b64}"
+
+
+def build_payload(
+    text: str,
+    voice_preset: str,
+    ref_audio: tuple | None,
+    ref_text: str,
+    response_format: str = "pcm",
+    stream: bool = True,
+) -> dict:
+    """Build the /v1/audio/speech request payload."""
+    if not text or not text.strip():
+        raise gr.Error("Please enter text to synthesize.")
+
+    voice_name = VOICE_PRESETS.get(voice_preset, "Junhao")
+
+    payload: dict = {
+        "input": text.strip(),
+        "voice": voice_name,
+        "response_format": "pcm" if stream else response_format,
+        "stream": stream,
+    }
+
+    has_ref = ref_audio is not None
+    if has_ref:
+        payload["ref_audio"] = encode_audio_to_base64(ref_audio)
+        ref_text_stripped = ref_text.strip() if ref_text else ""
+        if ref_text_stripped:
+            payload["ref_text"] = ref_text_stripped
+        else:
+            raise gr.Error("Voice cloning requires a transcript of the reference audio.")
+
+    return payload
+
+
+def generate_speech(
+    api_base: str,
+    text: str,
+    voice_preset: str,
+    ref_audio: tuple | None,
+    ref_text: str,
+    response_format: str,
+) -> tuple:
+    """Non-streaming: call /v1/audio/speech and return full audio as (sr, np_array)."""
+    payload = build_payload(text, voice_preset, ref_audio, ref_text, response_format, stream=False)
+
+    try:
+        with httpx.Client(timeout=300.0) as client:
+            resp = client.post(
+                f"{api_base}/v1/audio/speech",
+                json=payload,
+                headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
+            )
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.")
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
+
+    if resp.status_code != 200:
+        raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+
+    try:
+        if response_format == "pcm":
+            # PCM is stereo interleaved; downmix to mono for Gradio.
+            samples = np.frombuffer(resp.content, dtype=np.int16)
+            n = len(samples) - (len(samples) % 2)
+            left = samples[:n:2].astype(np.float32)
+            right = samples[1:n:2].astype(np.float32)
+            mono = (left + right) / 65536.0
+            return (PCM_SAMPLE_RATE, mono)
+        # WAV/MP3/FLAC: soundfile handles header parsing.
+        audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
+        if audio_np.ndim > 1:
+            audio_np = audio_np.mean(axis=1)
+        return (int(sample_rate), audio_np.astype(np.float32))
+    except Exception as e:
+        raise gr.Error(f"Failed to decode audio: {e}")
+
+
+def create_app(api_base: str):
+    """Create the FastAPI app with streaming proxy + Gradio UI."""
+    fastapi_app = FastAPI()
+
+    # ── Streaming proxy (same-origin, no CORS issues) ────────────
+    @fastapi_app.post("/proxy/v1/audio/speech")
+    async def proxy_speech(request: Request):
+        body = await request.json()
+        logger.info(
+            "Proxy request: input=%r voice=%s stream=%s",
+            body.get("input", "")[:50],
+            body.get("voice"),
+            body.get("stream"),
+        )
+        try:
+            client = httpx.AsyncClient(timeout=300)
+            resp = await client.send(
+                client.build_request(
+                    "POST",
+                    f"{api_base}/v1/audio/speech",
+                    json=body,
+                    headers={"Authorization": "Bearer EMPTY", "Content-Type": "application/json"},
+                ),
+                stream=True,
+            )
+        except Exception as exc:
+            logger.exception("Proxy connection error")
+            await client.aclose()
+            return Response(content=str(exc), status_code=502)
+
+        if resp.status_code != 200:
+            content = await resp.aread()
+            logger.error("Proxy upstream error %d: %s", resp.status_code, content[:200])
+            await resp.aclose()
+            await client.aclose()
+            return Response(content=content, status_code=resp.status_code)
+
+        async def relay():
+            total = 0
+            try:
+                async for chunk in resp.aiter_bytes():
+                    total += len(chunk)
+                    yield chunk
+            except Exception:
+                logger.exception("Proxy relay error after %d bytes", total)
+            finally:
+                logger.info("Proxy relay done: %d bytes", total)
+                await resp.aclose()
+                await client.aclose()
+
+        return StreamingResponse(relay(), media_type="application/octet-stream")
+
+    # ── Gradio UI ────────────────────────────────────────────────
+    css = """
+    #generate-btn button { width: 100%; }
+    #streaming-player { border: 1px solid var(--border-color-primary) !important; border-radius: var(--block-radius) !important; padding: var(--block-padding) !important; }
+    """
+
+    theme = gr.themes.Default(
+        primary_hue=gr.themes.Color(
+            c50="#E8F5E9",
+            c100="#C8E6C9",
+            c200="#A5D6A7",
+            c300="#81C784",
+            c400="#66BB6A",
+            c500="#4CAF50",
+            c600="#43A047",
+            c700="#388E3C",
+            c800="#2E7D32",
+            c900="#1B5E20",
+            c950="#0D3B10",
+        ),
+    )
+
+    with gr.Blocks(title="MOSS-TTS-Nano Demo") as demo:
+        gr.HTML(f"""
+        <div style="display:flex; align-items:center; gap:16px; margin-bottom:8px;">
+          <img src="https://raw.githubusercontent.com/vllm-project/vllm-omni/main/docs/source/logos/vllm-omni-logo.png"
+               alt="vLLM-Omni" style="height:42px;">
+          <div>
+            <h1 style="margin:0; font-size:1.5em;">MOSS-TTS-Nano Streaming Demo</h1>
+            <span style="font-size:0.85em; color:#666;">
+              0.1B AR LM + MOSS-Audio-Tokenizer-Nano &middot; 48 kHz &middot; ZH / EN / JA
+              &nbsp;&middot;&nbsp; Served by
+              <a href="https://github.com/vllm-project/vllm-omni" target="_blank"
+              style="color:#43A047; text-decoration:none; font-weight:600;">vLLM-Omni</a>
+              &nbsp;&middot;&nbsp; <code style="background:#eef2f7; padding:2px 6px; border-radius:4px; font-size:0.9em;">{api_base}</code>
+            </span>
+          </div>
+        </div>
+        """)
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                text_input = gr.Textbox(
+                    label="Text to Synthesize",
+                    placeholder="Enter text in Chinese, English, or Japanese...",
+                    lines=4,
+                )
+
+                voice_preset = gr.Dropdown(
+                    choices=PRESET_LABELS,
+                    value=DEFAULT_PRESET,
+                    label="Voice Preset",
+                    info="Built-in multilingual voices. ZH = Chinese, EN = English, JA = Japanese.",
+                )
+
+                with gr.Accordion("Custom Voice Clone (optional)", open=False):
+                    gr.Markdown(
+                        "Upload a reference audio clip (10-30 s) and provide its transcript "
+                        "to clone that voice instead of using the built-in preset."
+                    )
+                    ref_audio = gr.Audio(
+                        label="Reference Audio",
+                        type="numpy",
+                        sources=["upload", "microphone"],
+                    )
+                    ref_text = gr.Textbox(
+                        label="Reference Audio Transcript (required for cloning)",
+                        placeholder="Exact transcript of the reference audio...",
+                        lines=2,
+                    )
+
+                with gr.Row():
+                    response_format = gr.Dropdown(
+                        choices=["wav", "mp3", "flac", "pcm"],
+                        value="pcm",
+                        label="Audio Format",
+                        interactive=False,
+                        scale=1,
+                    )
+                    stream_checkbox = gr.Checkbox(
+                        label="Stream (gapless)",
+                        value=True,
+                        info="AudioWorklet streaming (recommended)",
+                        scale=1,
+                    )
+
+                with gr.Row():
+                    generate_btn = gr.Button(
+                        "Generate Speech",
+                        variant="primary",
+                        size="lg",
+                        elem_id="generate-btn",
+                        scale=3,
+                    )
+                    reset_btn = gr.Button(
+                        "Reset",
+                        variant="secondary",
+                        size="lg",
+                        scale=1,
+                    )
+
+            with gr.Column(scale=2):
+                player_html = gr.HTML(
+                    value=PLAYER_HTML,
+                    visible=True,
+                    label="streaming player",
+                    elem_id="streaming-player",
+                )
+                audio_output = gr.Audio(
+                    label="Generated Audio",
+                    interactive=False,
+                    autoplay=True,
+                    visible=False,
+                )
+                gr.Examples(
+                    examples=[
+                        ["Hello, this is MOSS-TTS-Nano speaking.", "Ava (EN)"],
+                        ["The quick brown fox jumps over the lazy dog.", "Adam (EN)"],
+                        ["MOSS-TTS-Nanoは、軽量なテキスト読み上げモデルです。", "Sakura (JA)"],
+                    ],
+                    inputs=[text_input, voice_preset],
+                    label="Examples",
+                )
+                gr.HTML("""
+                <div style="text-align:center; padding:8px 0; margin-top:4px;">
+                  <a href="https://github.com/vllm-project/vllm-omni" target="_blank">
+                    <img src="https://raw.githubusercontent.com/vllm-project/vllm-omni/main/docs/source/logos/vllm-omni-logo.png"
+                         alt="vLLM-Omni" style="height:28px; opacity:0.7;">
+                  </a>
+                </div>
+                """)
+
+        # Hidden textbox to pass payload from Python -> JavaScript
+        hidden_payload = gr.Textbox(visible=False, elem_id="tts-payload")
+
+        # ── Event wiring ─────────────────────────────────────────
+        def on_stream_change(stream: bool):
+            if stream:
+                return (
+                    gr.update(value="pcm", interactive=False),
+                    gr.update(visible=True),  # player
+                    gr.update(visible=False),  # audio
+                )
+            return (
+                gr.update(value="wav", interactive=True),
+                gr.update(visible=False),
+                gr.update(visible=True),
+            )
+
+        stream_checkbox.change(
+            fn=on_stream_change,
+            inputs=[stream_checkbox],
+            outputs=[response_format, player_html, audio_output],
+        )
+
+        def on_reset():
+            return (
+                "",  # text
+                None,  # audio_output
+                "",  # hidden_payload
+                PLAYER_HTML,  # reset player
+            )
+
+        reset_btn.click(
+            fn=on_reset,
+            outputs=[text_input, audio_output, hidden_payload, player_html],
+            js="() => { if (window.ttsStop) window.ttsStop(); }",
+        )
+
+        all_inputs = [text_input, voice_preset, ref_audio, ref_text, response_format]
+
+        def on_generate(stream_enabled, *args):
+            import time as _time
+
+            text, voice, ref_a, ref_t, _fmt = args
+            if stream_enabled:
+                payload = build_payload(text, voice, ref_a, ref_t, stream=True)
+                payload["_nonce"] = int(_time.time() * 1000)
+                return json.dumps(payload), gr.update()
+            else:
+                audio = generate_speech(api_base, text, voice, ref_a, ref_t, _fmt)
+                return "", audio
+
+        generate_btn.click(
+            fn=on_generate,
+            inputs=[stream_checkbox] + all_inputs,
+            outputs=[hidden_payload, audio_output],
+        ).then(
+            fn=lambda p: p,
+            inputs=[hidden_payload],
+            outputs=[hidden_payload],
+            js="(p) => { if (p && p.trim()) { const d = JSON.parse(p); delete d._nonce; window.ttsGenerate(d); } return p; }",
+        )
+
+        demo.queue()
+
+    return gr.mount_gradio_app(
+        fastapi_app,
+        demo,
+        path="/",
+        css=css,
+        theme=theme,
+        head=_build_player_js(PCM_SAMPLE_RATE),
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Gradio demo for MOSS-TTS-Nano with gapless AudioWorklet streaming.",
+    )
+    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})")
+    parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)")
+    parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)")
+    parser.add_argument("--share", action="store_true", help="Share publicly via Gradio tunnel")
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO)
+    print(f"Connecting to vLLM-Omni server at: {args.api_base}")
+
+    app = create_app(args.api_base)
+
+    import uvicorn
+
+    uvicorn.run(app, host=args.host, port=args.port)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/moss_tts_nano/run_gradio_demo.sh

@@ -0,0 +1,43 @@
+#!/bin/bash
+# Launch MOSS-TTS-Nano server + Gradio demo together.
+#
+# Usage:
+#   ./run_gradio_demo.sh
+#   CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh
+
+set -e
+
+MODEL="${MODEL:-OpenMOSS-Team/MOSS-TTS-Nano}"
+PORT="${PORT:-8091}"
+GRADIO_PORT="${GRADIO_PORT:-7860}"
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+
+echo "Starting MOSS-TTS-Nano server (port $PORT)..."
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+vllm serve "$MODEL" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --omni &
+SERVER_PID=$!
+
+cleanup() {
+    echo "Stopping server (PID $SERVER_PID)..."
+    kill $SERVER_PID 2>/dev/null
+    wait $SERVER_PID 2>/dev/null
+}
+trap cleanup EXIT
+
+# Wait for server to be ready.
+echo "Waiting for server to start..."
+for i in $(seq 1 120); do
+    if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
+        echo "Server ready."
+        break
+    fi
+    sleep 2
+done
+
+echo "Starting Gradio demo (port $GRADIO_PORT)..."
+python "$SCRIPT_DIR/gradio_demo.py" \
+    --api-base "http://localhost:$PORT" \
+    --port "$GRADIO_PORT"

examples/online_serving/moss_tts_nano/run_server.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+# Launch vLLM-Omni server for MOSS-TTS-Nano
+#
+# Usage:
+#   ./run_server.sh
+#   CUDA_VISIBLE_DEVICES=0 PORT=8091 ./run_server.sh
+
+set -e
+
+MODEL="${MODEL:-OpenMOSS-Team/MOSS-TTS-Nano}"
+PORT="${PORT:-8091}"
+
+echo "Starting MOSS-TTS-Nano server with model: $MODEL"
+
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+vllm serve "$MODEL" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --omni

examples/online_serving/qwen3_omni/README.md

@@ -31,11 +31,13 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
 ### Launch individual stages (stage-based CLI)
 
 Use the stage-based CLI when you want to run one stage per process.
+The example below pins Stage 0 to GPU 0 and Stage 1/2 to GPU 1 via
+`CUDA_VISIBLE_DEVICES`.
 
 **1. Stage 0 (Thinker + API server)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --port 8091 \
     --stage-id 0 \
     --omni-master-address 127.0.0.1 \
@@ -45,7 +47,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 **2. Stage 1 (Talker)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -55,7 +57,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
 **3. Stage 2 (Code2Wav)**
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 2 \
     --headless \
     --omni-master-address 127.0.0.1 \
@@ -78,7 +80,7 @@ for this category of optimization. Given that each instantiation strictly initia
 parameter flags can be systematically assigned directly onto that specific stage's command sequence:
 
 ```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
     --stage-id 1 \
     --headless \
     --gpu-memory-utilization 0.5 \

recipes/Qwen/Qwen3-Omni.md

@@ -1,76 +1,141 @@
-# Qwen3-Omni for multimodal chat on 1x A100 80GB
+# Qwen3-Omni
 
 ## Summary
 
 - Vendor: Qwen
 - Model: `Qwen/Qwen3-Omni-30B-A3B-Instruct`
-- Task: Multimodal chat with text, image, audio, or video input
+- Task: Multimodal chat with text/image/audio/video input
 - Mode: Online serving with the OpenAI-compatible API
 - Maintainer: Community
 
 ## When to use this recipe
 
-Use this recipe when you want a known-good starting point for serving
-`Qwen/Qwen3-Omni-30B-A3B-Instruct` with vLLM-Omni on a single 80 GB A100 and
-validate the deployment with the existing multimodal client examples in this
-repository.
+Use this recipe as a practical baseline for running
+`Qwen/Qwen3-Omni-30B-A3B-Instruct` with the same serving paths already
+exercised by repository examples and tests.
 
 ## References
 
-- Upstream or canonical docs:
+- User guide:
   [`docs/user_guide/examples/online_serving/qwen3_omni.md`](../../docs/user_guide/examples/online_serving/qwen3_omni.md)
-- Related example under `examples/`:
+- Example guide:
   [`examples/online_serving/qwen3_omni/README.md`](../../examples/online_serving/qwen3_omni/README.md)
-- Related issue or discussion:
-  [RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
 
-## Hardware Support
-
-This recipe currently documents one tested-style reference configuration for
-CUDA GPU serving. Add more sections for other hardware as community validation
-lands.
-
-## GPU
-
-### 1x A100 80GB
-
-#### Environment
+## Environment
 
 - OS: Linux
 - Python: 3.10+
-- Driver / runtime: NVIDIA CUDA environment with an A100 80 GB GPU
-- vLLM version: Match the repository requirements for your checkout
-- vLLM-Omni version or commit: Use the commit you are deploying from
+- vLLM / vLLM-Omni: use versions from your current checkout, >=0.18.0
 
-#### Command
+## Start server (single command)
 
-Start the server from the repository root:
+From repository root:
 
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
-Async chunking is enabled by default in the bundled deployment config. For
-common runtime tuning, prefer CLI overrides instead of editing or passing a
-custom YAML file:
+Notes:
 
-```bash
-# Disable async chunking for /v1/realtime sessions
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
-  --no-async-chunk
-```
+- `--omni` is required.
+- The default deploy config `vllm_omni/deploy/qwen3_omni_moe.yaml` is loaded
+  automatically by model registry.
+- `async_chunk` is enabled by default in this deploy config.
+- Platform deltas under `platforms:` (NPU/ROCm/XPU) are merged automatically on
+  matching runtimes.
 
-Use a custom deploy config only for advanced cases such as custom topology,
-connector wiring, or a larger overlay of stage defaults:
+For advanced customization, pass an overlay YAML:
 
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
   --deploy-config /path/to/your_qwen3_omni_overrides.yaml
 ```
 
+#### Runtime tuning
+
+Prefer CLI overrides for day-to-day tuning:
+
+```bash
+# Disable async chunking when using /v1/realtime
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+  --no-async-chunk
+
+# Example per-stage tuning in unified launch
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+  --stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
+
+# Tune max_num_seqs per stage (single process launch)
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
+  --stage-overrides '{
+    "0": {"max_num_seqs": 8},
+    "1": {"max_num_seqs": 4},
+    "2": {"max_num_seqs": 4}
+  }'
+```
+
+#### Stage-based launch (one stage per process)
+
+Use three terminals (one per stage). Start with the default commands below, then
+add `--max-num-seqs` only if you need explicit per-stage concurrency control.
+
+Default stage-based commands:
+
+```bash
+# Stage 0: Thinker + API server
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --port 8091 \
+  --stage-id 0 \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+
+# Stage 1: Talker
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --stage-id 1 \
+  --headless \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+
+# Stage 2: Code2Wav
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --stage-id 2 \
+  --headless \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+```
+
+Optional: explicit per-stage `max_num_seqs` tuning:
+
+```bash
+# Stage 0
+CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --port 8091 \
+  --stage-id 0 \
+  --max-num-seqs 8 \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+
+# Stage 1
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --stage-id 1 \
+  --headless \
+  --max-num-seqs 4 \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+
+# Stage 2
+CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
+  --stage-id 2 \
+  --headless \
+  --max-num-seqs 4 \
+  --omni-master-address 127.0.0.1 \
+  --omni-master-port 26000 &
+```
+
+If you use custom deploy YAML, add `--deploy-config` to each stage command.
+
 #### Verification
 
-Run one of the existing example clients after the server is ready:
+After server startup, run a multimodal example client:
 
 ```bash
 python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
@@ -80,7 +145,7 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene
   --host localhost
 ```
 
-For a quick API smoke test, request text-only output:
+Quick API smoke test (text-only output):
 
 ```bash
 curl http://localhost:8091/v1/chat/completions \
@@ -92,8 +157,89 @@ curl http://localhost:8091/v1/chat/completions \
   }'
 ```
 
+`{"modalities":["text","audio"]}` means the model returns both text and audio in
+the same response. Use it when you want transcription/content text and TTS audio
+together.
+
+Quick API smoke test (text + audio output):
+
+```bash
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+    "messages": [{"role": "user", "content": "Describe vLLM in brief."}],
+    "modalities": ["text", "audio"]
+  }'
+```
+
+Realtime WebSocket check (`/v1/realtime`) requires async chunk disabled:
+
+```bash
+python examples/online_serving/qwen3_omni/openai_realtime_client.py \
+  --url ws://localhost:8091/v1/realtime \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --input-wav /path/to/input_16k_mono.wav \
+  --output-wav realtime_output.wav
+```
+
+#### Benchmark with `vllm bench`
+
+After the server is up, you can run online serving benchmarks with
+`vllm bench serve --omni`.
+
+Text-focused random workload:
+
+```bash
+vllm bench serve \
+  --omni \
+  --host localhost \
+  --port 8091 \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --endpoint /v1/chat/completions \
+  --backend openai-chat-omni \
+  --dataset-name random \
+  --num-prompts 40 \
+  --max-concurrency 4 \
+  --random-input-len 2500 \
+  --random-output-len 900 \
+  --ignore-eos \
+  --extra-body '{"modalities":["text"]}' \
+  --percentile-metrics ttft,tpot,itl,e2el
+```
+
+If you want benchmark requests to return both text and audio, switch
+`--extra-body` to:
+
+```bash
+--extra-body '{"modalities":["text","audio"]}'
+```
+
+Synthetic multimodal workload (`random-mm`):
+
+```bash
+vllm bench serve \
+  --omni \
+  --host localhost \
+  --port 8091 \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --endpoint /v1/chat/completions \
+  --backend openai-chat-omni \
+  --dataset-name random-mm \
+  --num-prompts 40 \
+  --max-concurrency 4 \
+  --random-input-len 2500 \
+  --random-output-len 900 \
+  --random-mm-base-items-per-request 2 \
+  --random-mm-limit-mm-per-prompt '{"image":1,"video":1,"audio":1}' \
+  --random-mm-bucket-config '{"(32, 32, 1)": 0.5, "(0, 1, 1)": 0.5}' \
+  --ignore-eos \
+  --extra-body '{"modalities":["text"]}' \
+  --percentile-metrics ttft,tpot,itl,e2el
+```
+
 #### Notes
 
-- Memory usage: Size depends on runtime options and output modalities; leave headroom for multimodal workloads. Prefer CLI overrides such as `--gpu-memory-utilization` for routine tuning.
-- Key flags: `--omni` is required; async chunking is enabled by default; use `--no-async-chunk` for realtime sessions and `--deploy-config` only for advanced custom deployments.
-- Known limitations: The `/v1/realtime` WebSocket flow is currently unsupported while async chunking is enabled. This starter recipe is intentionally narrow and focuses on the single-GPU online-serving path already documented in the repo examples.
+- `/v1/realtime` is unsupported while `async_chunk` is enabled.
+- The default deploy uses `SharedMemoryConnector`; this is for single-host
+  stage wiring.

requirements/musa.txt

@@ -1,6 +1,6 @@
 -r common.txt
 # MUSA platform dependencies
-torchada>=0.1.50
+torchada>=0.1.52
 onnxruntime>=1.23.2
 mate>=0.2.0
 flash_attn_3>=0.1.4

tests/benchmarks/test_seed_tts_dataset_variants.py

@@ -59,6 +59,7 @@ def mock_tokenizer(mocker):
     tokenizer.all_special_ids = []
     tokenizer.all_special_tokens = []
     tokenizer.vocab_size = 1
+    tokenizer.__len__.return_value = 1
     return tokenizer
 
 

tests/config/test_pipeline_registry.py

@@ -6,11 +6,7 @@ from __future__ import annotations
 
 import pytest
 
-from vllm_omni.config.pipeline_registry import (
-    _DIFFUSION_PIPELINES,
-    _OMNI_PIPELINES,
-    _VLLM_OMNI_PIPELINES,
-)
+from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
 from vllm_omni.config.stage_config import (
     _PIPELINE_REGISTRY,
     PipelineConfig,
@@ -23,17 +19,9 @@ from vllm_omni.config.stage_config import (
 class TestCentralRegistryDeclarations:
     """Every in-tree pipeline must be declared exactly once in the central registry."""
 
-    def test_union_contains_all_omni(self):
+    def test_omni_entries_visible_in_registry(self):
         for key in _OMNI_PIPELINES:
-            assert key in _VLLM_OMNI_PIPELINES
-
-    def test_union_contains_all_diffusion(self):
-        for key in _DIFFUSION_PIPELINES:
-            assert key in _VLLM_OMNI_PIPELINES
-
-    def test_no_duplicate_model_type_between_omni_and_diffusion(self):
-        overlap = set(_OMNI_PIPELINES) & set(_DIFFUSION_PIPELINES)
-        assert not overlap, f"Duplicate model_types across omni/diffusion: {overlap}"
+            assert key in _PIPELINE_REGISTRY
 
     def test_expected_omni_pipelines_present(self):
         # Guard against accidental removal during future refactors.

tests/dfx/conftest.py

@@ -1,13 +1,18 @@
 import json
 import os
+import re
 import subprocess
+import time
 from datetime import datetime
 from pathlib import Path
 from typing import Any
 
 import pytest
 
+from tests.dfx.reliability.helpers import list_remote_process_pids_by_pattern, post_chat_completions_raw
+from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import modify_stage_config
+from vllm_omni.platforms import current_omni_platform
 
 
 def load_configs(config_path: str) -> list[dict[str, Any]]:
@@ -26,18 +31,15 @@ def load_configs(config_path: str) -> list[dict[str, Any]]:
         raise RuntimeError(f"Failed to load configuration file: {str(e)}")
 
 
-def modify_stage(default_path, updates, deletes):
-    kwargs = {}
+def modify_stage(default_path: str, updates: dict[str, Any] | None, deletes: dict[str, Any] | None) -> str:
+    kwargs: dict[str, Any] = {}
     if updates is not None:
         kwargs["updates"] = updates
     if deletes is not None:
         kwargs["deletes"] = deletes
     if kwargs:
-        path = modify_stage_config(default_path, **kwargs)
-    else:
-        path = default_path
-
-    return path
+        return modify_stage_config(default_path, **kwargs)
+    return default_path
 
 
 def _build_serve_args(serve_args: Any) -> list[str]:
@@ -111,6 +113,122 @@ def create_unique_server_params(
     return unique_params
 
 
+def configs_with_platform_stage_configs(configs: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Resolve stage config for XPU vs CUDA/ROCm."""
+    out: list[dict[str, Any]] = []
+    for config in configs:
+        config_copy = json.loads(json.dumps(config))
+        if current_omni_platform.is_xpu():
+            config_copy["server_params"]["stage_config_name"] = "xpu/qwen3_omni_ci.yaml"
+        out.append(config_copy)
+    return out
+
+
+def extract_server_args_by_test_name(configs: list[dict[str, Any]]) -> dict[str, list[str] | None]:
+    mapping: dict[str, list[str] | None] = {}
+    for cfg in configs:
+        test_name = str(cfg.get("test_name"))
+        server_params = cfg.get("server_params") or {}
+        raw_args = server_params.get("server_args")
+        mapping[test_name] = [str(item) for item in raw_args] if isinstance(raw_args, list) else None
+    return mapping
+
+
+def create_reliability_omni_server_params(
+    configs: list[dict[str, Any]], stage_configs_dir: Path
+) -> list[OmniServerParams]:
+    adjusted_configs = configs_with_platform_stage_configs(configs)
+    unique_params = create_unique_server_params(adjusted_configs, stage_configs_dir)
+    server_args_by_name = extract_server_args_by_test_name(adjusted_configs)
+    return [
+        OmniServerParams(
+            model=model,
+            stage_config_path=stage_config_path,
+            server_args=server_args_by_name.get(test_name),
+        )
+        for test_name, model, stage_config_path, _stage_overrides_json, _extra_cli_args in unique_params
+    ]
+
+
+def supports_video_generation(model_name: str) -> bool:
+    lower = model_name.lower()
+    return any(key in lower for key in ("wan", "video", "i2v", "t2v"))
+
+
+def supports_chat_generation(model_name: str) -> bool:
+    return not supports_video_generation(model_name)
+
+
+def parse_stage_devices(stage_config_path: str) -> str:
+    text = Path(stage_config_path).read_text(encoding="utf-8")
+    raw_devices: list[str] = re.findall(r"^\s*devices:\s*\"?([0-9,\s]+)\"?\s*$", text, flags=re.MULTILINE)
+    devices: set[int] = set()
+    for item in raw_devices:
+        for token in item.split(","):
+            token = token.strip()
+            if token:
+                devices.add(int(token))
+    if not devices:
+        raise ValueError(f"No runtime.devices found in stage config: {stage_config_path}")
+    return ",".join(str(x) for x in sorted(devices))
+
+
+def resolve_oom_device_spec(config: dict[str, Any], stage_config_path: str | None) -> str:
+    explicit = config.get("device")
+    if explicit is not None:
+        return str(explicit)
+    if not stage_config_path:
+        return "0"
+    return parse_stage_devices(stage_config_path)
+
+
+def assert_fault_exception(exc: Exception, error_keywords: tuple[str, ...]) -> None:
+    text = str(exc).lower()
+    assert any(key in text for key in error_keywords), f"unexpected error under fault injection: {exc}"
+
+
+def wait_chat_request_ready(host: str, port: int, model: str, timeout_sec: int = 180) -> None:
+    """Poll a minimal chat request until success."""
+    deadline = time.time() + timeout_sec
+    payload = json.dumps(
+        {
+            "model": model,
+            "messages": [{"role": "user", "content": "Say hello in one short sentence."}],
+            "stream": False,
+            "modalities": ["text"],
+        }
+    )
+    last_error: str | None = None
+    while time.time() < deadline:
+        try:
+            status, body = post_chat_completions_raw(host, port, payload)
+            if status == 200:
+                return
+            last_error = f"http={status} body={body[:200]!r}"
+        except Exception as exc:  # noqa: BLE001
+            last_error = str(exc)
+        time.sleep(2)
+    raise TimeoutError(f"runtime-teardown warmup request did not succeed within {timeout_sec}s: {last_error}")
+
+
+def assert_no_extra_worker_processes(
+    baseline_pids: set[int],
+    worker_pattern: str,
+    timeout_sec: int = 60,
+) -> None:
+    """Ensure no extra worker PIDs remain after teardown."""
+    deadline = time.time() + timeout_sec
+    while time.time() < deadline:
+        current = set(list_remote_process_pids_by_pattern(worker_pattern))
+        extra = current - baseline_pids
+        if not extra:
+            return
+        time.sleep(2)
+    current = set(list_remote_process_pids_by_pattern(worker_pattern))
+    extra = sorted(current - baseline_pids)
+    assert not extra, f"orphan worker processes remain after container teardown: {extra}"
+
+
 def create_test_parameter_mapping(configs: list[dict[str, Any]]) -> dict[str, dict]:
     mapping = {}
     for config in configs:

tests/dfx/reliability/README.md

@@ -0,0 +1,14 @@
+# L5(b) Reliability (`tests/dfx/reliability`)
+
+This directory contains reliability fault-injection tests for key models.
+The detailed RFC document is maintained outside this repository (local/internal only).
+
+- Fault helpers: `conftest.py`, `helpers.py`
+- Qwen3 tests: `test_reliability_qwen3_omni.py`
+- Wan2.2 tests: `test_reliability_wan22.py`
+
+```bash
+pytest --collect-only tests/dfx/reliability
+pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
+pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
+```

tests/dfx/reliability/conftest.py

@@ -0,0 +1,29 @@
+"""Pytest fixtures for reliability tests."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from tests.dfx.reliability.helpers import FaultInjector
+
+
+@pytest.fixture
+def fault_injector(request: pytest.FixtureRequest) -> FaultInjector:
+    """Indirect only: ``request.param`` must be a ``FaultInjector`` callable."""
+    return request.param
+
+
+@pytest.fixture
+def omni_server_after_fault(omni_server: Any, fault_injector: FaultInjector):
+    """After ``omni_server`` is up, run ``fault_injector(omni_server)``, then yield the server."""
+    fault_injector(omni_server)
+    yield omni_server
+
+
+@pytest.fixture
+def omni_server_after_fault_function(omni_server_function: Any, fault_injector: FaultInjector):
+    """Inject fault after function-scoped server startup, then yield server."""
+    fault_injector(omni_server_function)
+    yield omni_server_function

tests/dfx/reliability/helpers.py

@@ -0,0 +1,567 @@
+"""Shared reliability fault-injection helpers.
+
+This module keeps fault injection callable from tests directly:
+- GPU OOM (CUDA sidecar memory hog)
+- process kill by pattern and signal
+- post-ready hooks via ``fault_injector`` / ``omni_server_after_fault`` fixtures
+"""
+
+from __future__ import annotations
+
+import http.client
+import json
+import logging
+import os
+import select
+import shlex
+import signal
+import subprocess
+import sys
+import time
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import psutil
+import pytest
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OomHandle:
+    """Handle for a started CUDA memory hog subprocess."""
+
+    proc: subprocess.Popen | None
+    device: int
+    target_mem_ratio: float
+    start_ts: float
+
+
+def post_chat_completions_raw(
+    host: str,
+    port: int,
+    body: bytes | str,
+    *,
+    content_type: str = "application/json",
+    timeout_sec: int = 120,
+) -> tuple[int, bytes]:
+    """POST /v1/chat/completions with raw bytes; returns (status, response_body)."""
+    conn = http.client.HTTPConnection(host, port, timeout=timeout_sec)
+    try:
+        headers = {"Content-Type": content_type}
+        payload = body.encode("utf-8") if isinstance(body, str) else body
+        conn.request("POST", "/v1/chat/completions", body=payload, headers=headers)
+        resp = conn.getresponse()
+        data = resp.read()
+        return resp.status, data
+    finally:
+        conn.close()
+
+
+def get_health_raw(host: str, port: int, *, timeout_sec: int = 20) -> tuple[int, bytes]:
+    """GET /health with stdlib HTTP client; returns (status, body)."""
+    conn = http.client.HTTPConnection(host, port, timeout=timeout_sec)
+    try:
+        conn.request("GET", "/health")
+        resp = conn.getresponse()
+        return resp.status, resp.read()
+    finally:
+        conn.close()
+
+
+def post_json_raw_http_client(
+    host: str,
+    port: int,
+    path: str,
+    payload: dict[str, Any],
+    *,
+    timeout_sec: int = 30,
+) -> tuple[int, bytes]:
+    """POST JSON to one endpoint with stdlib HTTP client; returns (status, body)."""
+    conn = http.client.HTTPConnection(host, port, timeout=timeout_sec)
+    try:
+        body = json.dumps(payload).encode("utf-8")
+        conn.request("POST", path, body=body, headers={"Content-Type": "application/json"})
+        resp = conn.getresponse()
+        return resp.status, resp.read()
+    finally:
+        conn.close()
+
+
+def post_json_raw(
+    host: str,
+    port: int,
+    path: str,
+    payload: dict[str, Any],
+    *,
+    timeout_sec: int = 30,
+) -> tuple[int, bytes]:
+    """POST JSON to one endpoint; returns (status, body)."""
+    return (
+        post_chat_completions_raw(
+            host,
+            port,
+            json.dumps(payload),
+            content_type="application/json",
+            timeout_sec=timeout_sec,
+        )
+        if path == "/v1/chat/completions"
+        else post_json_raw_http_client(
+            host,
+            port,
+            path,
+            payload,
+            timeout_sec=timeout_sec,
+        )
+    )
+
+
+def extract_openai_error_contract_from_bytes(response_body: bytes) -> dict[str, Any] | None:
+    """Best-effort parse OpenAI-style error object from raw response bytes."""
+    try:
+        payload = json.loads(response_body.decode("utf-8", errors="replace"))
+    except Exception:  # noqa: BLE001
+        return None
+    return extract_openai_error_contract_from_payload(payload)
+
+
+def extract_openai_error_contract_from_payload(payload: Any) -> dict[str, Any] | None:
+    """Best-effort parse OpenAI-style error object from decoded JSON payload."""
+    if not isinstance(payload, dict):
+        return None
+    error_obj = payload.get("error")
+    if not isinstance(error_obj, dict):
+        return None
+    if not isinstance(error_obj.get("message"), str):
+        return None
+    return error_obj
+
+
+def _build_sidecar_cmd(device: int, target_mem_ratio: float, hold_seconds: int, strict: bool) -> list[str]:
+    sidecar = r"""
+import sys
+import time
+import torch
+
+device = int(sys.argv[1])
+target_ratio = float(sys.argv[2])
+hold_seconds = int(sys.argv[3])
+strict = sys.argv[4] == "1"
+
+torch.cuda.init()
+torch.cuda.set_device(device)
+props = torch.cuda.get_device_properties(device)
+free_before, total_bytes = torch.cuda.mem_get_info(device)
+target_bytes = int(free_before * target_ratio)
+chunk_bytes = 256 * 1024 * 1024
+chunks = []
+allocated = 0
+
+while allocated < target_bytes:
+    req_bytes = min(chunk_bytes, target_bytes - allocated)
+    req_elems = max(1, req_bytes // 2)  # float16 -> 2 bytes
+    try:
+        chunk = torch.empty((req_elems,), dtype=torch.float16, device=f"cuda:{device}")
+        chunks.append(chunk)
+        allocated += chunk.numel() * 2
+    except RuntimeError:
+        break
+
+# In strict mode, keep filling with smaller chunks until allocator rejects.
+# This minimizes residual free memory and makes fault-path assertions steadier.
+if strict:
+    tail_chunk_bytes = [64 * 1024 * 1024, 16 * 1024 * 1024, 4 * 1024 * 1024, 1 * 1024 * 1024]
+    for tail_bytes in tail_chunk_bytes:
+        while True:
+            req_elems = max(1, tail_bytes // 2)
+            try:
+                chunk = torch.empty((req_elems,), dtype=torch.float16, device=f"cuda:{device}")
+                chunks.append(chunk)
+                allocated += chunk.numel() * 2
+            except RuntimeError:
+                break
+
+achieved_ratio = allocated / max(1, props.total_memory)
+achieved_free_ratio = allocated / max(1, int(free_before))
+free_after, _ = torch.cuda.mem_get_info(device)
+if strict and allocated < target_bytes:
+    print(
+        "ERROR:"
+        f"achieved_free_ratio={achieved_free_ratio:.4f};"
+        f"achieved_total_ratio={achieved_ratio:.4f};"
+        f"free_before={int(free_before)};"
+        f"free_after={int(free_after)};"
+        f"target_bytes={target_bytes};"
+        f"allocated={allocated}",
+        flush=True,
+    )
+    sys.exit(2)
+
+print(
+    "READY:"
+    f"achieved_free_ratio={achieved_free_ratio:.4f};"
+    f"achieved_total_ratio={achieved_ratio:.4f};"
+    f"free_before={int(free_before)};"
+    f"free_after={int(free_after)};"
+    f"target_bytes={target_bytes};"
+    f"allocated={allocated}",
+    flush=True,
+)
+if hold_seconds <= 0:
+    while True:
+        time.sleep(3600)
+time.sleep(hold_seconds)
+print("DONE", flush=True)
+"""
+    return [
+        sys.executable,
+        "-u",
+        "-c",
+        sidecar,
+        str(device),
+        str(target_mem_ratio),
+        str(hold_seconds),
+        "1" if strict else "0",
+    ]
+
+
+def start_gpu_oom_hog(
+    *,
+    device: int = 0,
+    target_mem_ratio: float = 0.95,
+    hold_seconds: int = 60,
+    startup_timeout_sec: int = 20,
+    strict: bool = True,
+    poll_interval_sec: float = 0.2,
+) -> OomHandle:
+    """Start a CUDA sidecar process that occupies GPU memory to trigger OOM.
+
+    Note:
+        ``target_mem_ratio`` is evaluated against free memory at injection start
+        (not total GPU memory), i.e. success gate is ``allocated / free_before``.
+        ``hold_seconds <= 0`` means keeping OOM pressure until the sidecar is
+        explicitly stopped via ``stop_gpu_oom_hog(s)``.
+    """
+    if os.name == "nt":
+        raise RuntimeError("CUDA OOM sidecar is intended for Linux CI/runtime.")
+    if not (0.0 <= target_mem_ratio < 1.0):
+        raise ValueError("target_mem_ratio should be in [0.0, 1.0).")
+
+    # Explicit opt-out for debugging: keep API shape stable while disabling injection.
+    if target_mem_ratio == 0.0:
+        print(f"[oom-sidecar][gpu={device}] DISABLED: target_mem_ratio=0.0 (no OOM injection)", flush=True)
+        return OomHandle(
+            proc=None,
+            device=device,
+            target_mem_ratio=target_mem_ratio,
+            start_ts=time.time(),
+        )
+
+    cmd = _build_sidecar_cmd(device, target_mem_ratio, hold_seconds, strict)
+    proc = subprocess.Popen(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+    )
+    assert proc.stdout is not None
+
+    deadline = time.time() + startup_timeout_sec
+    logs: list[str] = []
+    while time.time() < deadline:
+        ready, _, _ = select.select([proc.stdout], [], [], poll_interval_sec)
+        if ready:
+            line = proc.stdout.readline().strip()
+            if line:
+                logs.append(line)
+                print(f"[oom-sidecar][gpu={device}] {line}", flush=True)
+                if line.startswith("READY:"):
+                    return OomHandle(
+                        proc=proc,
+                        device=device,
+                        target_mem_ratio=target_mem_ratio,
+                        start_ts=time.time(),
+                    )
+                if line.startswith("ERROR:"):
+                    proc.terminate()
+                    raise RuntimeError(f"OOM sidecar failed to reach target: {line}")
+        if proc.poll() is not None:
+            break
+
+    proc.terminate()
+    if logs:
+        print(f"[oom-sidecar][gpu={device}] startup logs: {' | '.join(logs)}", flush=True)
+    raise TimeoutError(f"OOM sidecar startup timeout. logs={logs}")
+
+
+def stop_gpu_oom_hog(handle: OomHandle, *, timeout_sec: int = 5) -> None:
+    """Stop and cleanup CUDA OOM sidecar."""
+    proc = handle.proc
+    if proc is None:
+        return
+    if proc.poll() is not None:
+        return
+    proc.terminate()
+    try:
+        proc.wait(timeout=timeout_sec)
+    except subprocess.TimeoutExpired:
+        proc.kill()
+        proc.wait(timeout=timeout_sec)
+
+
+def inject_gpu_oom(
+    *,
+    device: int | str | list[int] = 0,
+    target_mem_ratio: float = 0.95,
+    hold_seconds: int = 60,
+    startup_timeout_sec: int = 20,
+    strict: bool = True,
+) -> OomHandle | list[OomHandle]:
+    """Convenience wrapper to start CUDA OOM sidecar(s).
+
+    Args:
+        device: One device id (``0``), comma-separated string (``"0,1,2"``),
+            or a list of device ids (``[0, 1, 2]``).
+        hold_seconds: OOM hold time in seconds; ``<=0`` keeps pressure until
+            ``stop_gpu_oom_hogs`` is called.
+    """
+    if isinstance(device, int):
+        devices = [device]
+    elif isinstance(device, str):
+        devices = [int(x.strip()) for x in device.split(",") if x.strip()]
+    else:
+        devices = [int(x) for x in device]
+    if not devices:
+        raise ValueError("device must not be empty.")
+
+    handles = [
+        start_gpu_oom_hog(
+            device=dev,
+            target_mem_ratio=target_mem_ratio,
+            hold_seconds=hold_seconds,
+            startup_timeout_sec=startup_timeout_sec,
+            strict=strict,
+        )
+        for dev in devices
+    ]
+    if len(handles) == 1:
+        return handles[0]
+    return handles
+
+
+def stop_gpu_oom_hogs(handles: OomHandle | list[OomHandle], *, timeout_sec: int = 5) -> None:
+    """Stop one or multiple OOM sidecars."""
+    if isinstance(handles, OomHandle):
+        stop_gpu_oom_hog(handles, timeout_sec=timeout_sec)
+        return
+    for handle in handles:
+        stop_gpu_oom_hog(handle, timeout_sec=timeout_sec)
+
+
+def _runtime_teardown_ssh_target() -> str:
+    target = os.getenv("RUNTIME_TEARDOWN_SSH_TARGET", "").strip()
+    # Default to root@127.0.0.1 for same-host SSH control path.
+    return target or "root@127.0.0.1"
+
+
+def _runtime_teardown_ssh_cmd(remote_cmd: str, *, step: str | None = None) -> subprocess.CompletedProcess[str]:
+    ssh_target = _runtime_teardown_ssh_target()
+    default_reuse_opts = "-o ControlMaster=auto -o ControlPersist=10m -o ControlPath=/tmp/vllm-rt-ssh-%r@%h:%p"
+    raw_opts = os.getenv("RUNTIME_TEARDOWN_SSH_OPTS", "").strip()
+    ssh_opts = shlex.split(raw_opts or default_reuse_opts)
+    timeout_sec = int(os.getenv("RUNTIME_TEARDOWN_SSH_TIMEOUT_SEC", "600"))
+    step_prefix = f"[runtime-teardown][ssh]{f'[{step}]' if step else ''}"
+    print(f"{step_prefix} target={ssh_target} running remote command...", flush=True)
+    # IMPORTANT: SSH joins remote argv into one shell command string. If we pass
+    # ["bash", "-c", remote_cmd] as separate argv items, remote shell parsing can
+    # make `-c` consume only the first word (e.g. "docker"), causing docker help.
+    # Wrap the whole command as one quoted string for remote bash -c.
+    remote_invocation = f"bash --noprofile --norc -c {shlex.quote(remote_cmd)}"
+    try:
+        out = subprocess.run(
+            ["ssh", *ssh_opts, ssh_target, remote_invocation],
+            check=False,
+            capture_output=True,
+            text=True,
+            timeout=timeout_sec,
+        )
+    except subprocess.TimeoutExpired as exc:
+        raise RuntimeError(
+            f"{step_prefix} timed out after {timeout_sec}s. Increase RUNTIME_TEARDOWN_SSH_TIMEOUT_SEC if needed."
+        ) from exc
+    print(f"{step_prefix} exit_code={out.returncode}", flush=True)
+    return out
+
+
+def list_remote_process_pids_by_pattern(pattern: str) -> list[int]:
+    """Return matched PIDs from remote host ``pgrep -f <pattern>`` via SSH."""
+    cmd = f"pgrep -f {shlex.quote(pattern)} || true"
+    out = _runtime_teardown_ssh_cmd(cmd, step="pgrep")
+    if out.returncode not in (0, 1):
+        raise RuntimeError(f"remote pgrep failed for pattern={pattern!r}: {out.stderr.strip()}")
+    return [int(item) for item in out.stdout.split() if item.strip().isdigit()]
+
+
+def inject_process_kill(
+    *,
+    grep_pattern: str,
+    signal_name: str = "SIGTERM",
+    limit: int | None = None,
+    allow_zero_match: bool = False,
+    execute_kill: bool = True,
+) -> list[int]:
+    """Kill processes matching pattern with selected signal."""
+    if os.name == "nt":
+        raise RuntimeError("process-kill helper currently supports POSIX platforms only.")
+    if not grep_pattern.strip():
+        raise ValueError("grep_pattern must not be empty.")
+
+    sig = getattr(signal, signal_name, None)
+    if sig is None:
+        raise ValueError(f"Unsupported signal_name: {signal_name}")
+
+    out = subprocess.run(
+        ["pgrep", "-f", grep_pattern],
+        check=False,
+        capture_output=True,
+        text=True,
+    )
+    pids = [int(item) for item in out.stdout.split() if item.strip().isdigit()]
+    if limit is not None:
+        pids = pids[:limit]
+
+    if not pids and not allow_zero_match:
+        raise RuntimeError(f"No process matched pattern: {grep_pattern}")
+
+    if execute_kill:
+        for pid in pids:
+            os.kill(pid, sig)
+    return pids
+
+
+def _safe_proc_info(pid: int) -> tuple[str, str]:
+    """Best-effort process name/cmdline lookup for debug logging."""
+    try:
+        proc = psutil.Process(pid)
+        name = proc.name()
+        cmdline = " ".join(proc.cmdline()) or "<empty-cmdline>"
+        return name, cmdline
+    except Exception:  # noqa: BLE001
+        return "<unknown>", "<unavailable>"
+
+
+def _list_server_process_tree(server: Any) -> list[int]:
+    """Return [root, descendants...] PIDs for the current test server instance."""
+    root_proc = getattr(server, "proc", None)
+    if root_proc is None or getattr(root_proc, "pid", None) is None:
+        return []
+
+    root_pid = int(root_proc.pid)
+    try:
+        root = psutil.Process(root_pid)
+    except Exception:  # noqa: BLE001
+        return [root_pid]
+
+    descendants = [child.pid for child in root.children(recursive=True)]
+    return [root_pid, *descendants]
+
+
+def _log_server_process_tree(server: Any) -> None:
+    """Print server process tree for debugging fault injection targets."""
+    pids = _list_server_process_tree(server)
+    if not pids:
+        logger.warning("[reliability][process-kill] current server has no visible process tree")
+        return
+    for pid in pids:
+        name, cmdline = _safe_proc_info(pid)
+        print(
+            f"[reliability][process-kill] current_server_proc pid={pid} name={name} cmdline={cmdline}",
+            flush=True,
+        )
+
+
+FaultInjector = Callable[[Any], None]
+"""Callable invoked with the live ``OmniServer`` after it is ready (see ``omni_server_after_fault``)."""
+
+
+def make_process_kill_fault_injector(
+    *,
+    grep_patterns: str | Sequence[str],
+    signal_name: str = "SIGKILL",
+    limit: int = 1,
+    post_kill_wait_seconds: float = 0.0,
+) -> FaultInjector:
+    """Build a post-ready injector that kills processes matched by ``pgrep -f``.
+
+    Tries each pattern in order until at least one PID is killed. If none match,
+    the returned callable issues ``pytest.skip`` (same behavior as the previous
+    inline reliability test).
+
+    Args:
+        grep_patterns: One pattern or an ordered list of patterns.
+        signal_name: Passed to :func:`inject_process_kill` (e.g. ``SIGKILL``).
+        limit: Maximum PIDs to kill per pattern (default ``1``).
+        post_kill_wait_seconds: Optional wait time after kill before test request starts.
+    """
+    patterns: tuple[str, ...] = (grep_patterns,) if isinstance(grep_patterns, str) else tuple(grep_patterns)
+
+    def _inject(server: Any) -> None:
+        _log_server_process_tree(server)
+        server_tree = set(_list_server_process_tree(server))
+        if not server_tree:
+            logger.warning(
+                "[reliability][process-kill] no server process tree found; fallback to global pgrep matching"
+            )
+        for pattern in patterns:
+            print(
+                f"[reliability][process-kill] trying pattern={pattern} signal={signal_name} limit={limit}",
+                flush=True,
+            )
+            pids = inject_process_kill(
+                grep_pattern=pattern,
+                signal_name=signal_name,
+                limit=limit,
+                allow_zero_match=True,
+                execute_kill=False,
+            )
+            filtered = [pid for pid in pids if not server_tree or pid in server_tree]
+            if pids and not filtered:
+                logger.warning(
+                    "[reliability][process-kill] pattern=%s matched non-server pids=%s, skip them",
+                    pattern,
+                    pids,
+                )
+                continue
+            if filtered:
+                sig = getattr(signal, signal_name, None)
+                if sig is None:
+                    raise ValueError(f"Unsupported signal_name: {signal_name}")
+                for pid in filtered:
+                    name, cmdline = _safe_proc_info(pid)
+                    print(
+                        f"[reliability][process-kill] killing pid={pid} name={name} signal={signal_name} cmdline={cmdline}",
+                        flush=True,
+                    )
+                    os.kill(pid, sig)
+                print(
+                    f"[reliability][process-kill] matched pattern={pattern} killed_pids={filtered} killed_count={len(filtered)}",
+                    flush=True,
+                )
+                if post_kill_wait_seconds > 0:
+                    print(
+                        f"[reliability][process-kill] waiting {post_kill_wait_seconds:.2f}s after kill",
+                        flush=True,
+                    )
+                    time.sleep(post_kill_wait_seconds)
+                return
+        logger.warning(
+            "[reliability][process-kill] no process matched patterns=%s signal=%s limit=%s",
+            patterns,
+            signal_name,
+            limit,
+        )
+        pytest.skip("no matching runtime process found for kill injection")
+
+    return _inject

tests/dfx/reliability/test_reliability_qwen3_omni.py

@@ -0,0 +1,682 @@
+"""
+Qwen3-Omni reliability integration tests.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import errno
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any, Protocol
+
+import pytest
+import torch
+
+from tests.dfx.conftest import (
+    assert_fault_exception,
+    create_reliability_omni_server_params,
+    resolve_oom_device_spec,
+)
+from tests.dfx.reliability.helpers import (
+    extract_openai_error_contract_from_bytes,
+    get_health_raw,
+    inject_gpu_oom,
+    make_process_kill_fault_injector,
+    post_chat_completions_raw,
+    post_json_raw,
+    stop_gpu_oom_hogs,
+)
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_audio, generate_synthetic_image, generate_synthetic_video
+from tests.helpers.runtime import dummy_messages_from_mix_data
+
+RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
+    {
+        "test_name": "qwen3_omni_reliability_async_chunk",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "stage_config_name": "qwen3_omni_moe.yaml",
+            "server_args": ["--async-chunk"],
+        },
+    },
+    {
+        "test_name": "qwen3_omni_reliability_default",
+        "server_params": {
+            "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+            "stage_config_name": "qwen3_omni_moe.yaml",
+            "server_args": ["--no-async-chunk"],
+        },
+    },
+]
+
+DEPLOY_CONFIGS_DIR = Path(__file__).resolve().parent.parent.parent.parent / "vllm_omni" / "deploy"
+
+
+def _default_oom_device_spec() -> str:
+    """Use currently visible CUDA ordinals to avoid invalid device index in sidecar."""
+    count = torch.cuda.device_count()
+    if count <= 0:
+        return "0"
+    return ",".join(str(i) for i in range(count))
+
+
+OOM_INJECTION_CONFIG = {
+    "device": _default_oom_device_spec(),
+    "target_mem_ratio": 0.92,
+    "hold_seconds": 0,
+    "startup_timeout_sec": 20,
+    "strict": False,
+}
+# Post-fault recovery probe: keep the API process alive on the same GPU.
+# ``strict=True`` sidecars keep allocating until CUDA refuses, which often
+# evicts/kills the server; here we only take a fraction of *initial* free
+# memory and stop, then rely on a heavy multimodal forward to hit OOM.
+OOM_RECOVER_INJECTION_CONFIG = {
+    "device": _default_oom_device_spec(),
+    # ``strict=False`` stops at this fraction of *initial* free memory per GPU (see helpers sidecar).
+    # On multi-GPU hosts a mild ratio plus a small e2e-style mix can still leave several GB free
+    # on a data-parallel device and the request succeeds; 0.92+ pairs better with the heavy
+    # probe payload below. Raise toward 0.94 only if fault phase still never fails; avoid strict=True
+    # here so the API process is not squeezed to zero slack.
+    "target_mem_ratio": 0.8,
+    "hold_seconds": 0,
+    "startup_timeout_sec": 20,
+    "strict": False,
+}
+FAULT_ERROR_KEYWORDS = (
+    "the request failed",
+    "oom",
+    "out of memory",
+    "cuda",
+    "orchestrator",
+    "timeout",
+    "connection",
+    "500",
+    "503",
+)
+RUNTIME_WORKER_PATTERN = "VLLM::Worker"
+
+
+class _HasServeArgs(Protocol):
+    model: str
+    serve_args: list[str]
+
+
+def _get_system_prompt() -> dict:
+    return {
+        "role": "system",
+        "content": [
+            {
+                "type": "text",
+                "text": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, "
+                    "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+                ),
+            }
+        ],
+    }
+
+
+def _get_mix_prompt() -> str:
+    return "What is recited in the audio? What is in this image? Describe the video briefly."
+
+
+def _mix_chat_completions_probe_payload(omni_server: _HasServeArgs) -> dict[str, Any]:
+    """Heavy multimodal ``/v1/chat/completions`` body for OOM recover tests.
+
+    Uses the same synthetic sizes as ``test_reliability_fault_gpu_oom_chat_large_payload_failure``
+    (large vision tensors + long text), not the light expansion mix (224² / 5s audio): under
+    ``strict=False`` sidecars, a small mix often still completes on a GPU that retains multi-GB
+    free after the hog stops at ``target_mem_ratio``.
+    """
+    video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(1280, 720, 161)['base64']}"
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+    audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(20, 1)['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=_get_system_prompt(),
+        video_data_url=video_data_url,
+        image_data_url=image_data_url,
+        audio_data_url=audio_data_url,
+        content_text="What is recited in the audio? What is in this image? What is in this video? " * 200,
+    )
+    return {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "modalities": ["text", "audio"],
+    }
+
+
+def _fault_keywords_match_response(*, status: int, body: bytes) -> bool:
+    """True when HTTP status or body text looks like an expected fault / OOM class outcome."""
+    if status >= 400:
+        return True
+    text = body.decode("utf-8", errors="replace").lower()
+    # Avoid matching bare ``"500"`` / ``"503"`` in successful streamed bodies (base64/SSE noise).
+    body_fault_hints = (
+        "the request failed",
+        "oom",
+        "out of memory",
+        "cuda",
+        "orchestrator",
+        "timeout",
+        "connection refused",
+        "connection reset",
+    )
+    return any(key in text for key in body_fault_hints)
+
+
+def _stage_config_path_from_omni_server(omni_server: _HasServeArgs) -> str | None:
+    args: list[str] = omni_server.serve_args
+    for i, arg in enumerate(args):
+        if arg == "--stage-configs-path" and i + 1 < len(args):
+            return args[i + 1]
+        if arg.startswith("--stage-configs-path="):
+            return arg.split("=", 1)[1]
+    return None
+
+
+def _looks_like_server_unreachable(exc: BaseException) -> bool:
+    """True when /health cannot be reached because nothing is listening (process exited)."""
+    if isinstance(exc, (ConnectionRefusedError, BrokenPipeError, ConnectionResetError)):
+        return True
+    errno_val = getattr(exc, "errno", None)
+    if isinstance(exc, OSError) and errno_val is not None:
+        return errno_val in (
+            errno.ECONNREFUSED,
+            errno.ECONNRESET,
+            errno.EPIPE,
+        )
+    msg = str(exc).lower()
+    return "connection refused" in msg or "actively refused" in msg
+
+
+QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLOY_CONFIGS_DIR)
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
+    omni_server_function,
+) -> None:
+    """Black-box: text chat vs omni audio output should expose a consistent error contract under OOM.
+
+    Speech-style output is requested via ``/v1/chat/completions`` with ``modalities`` that include ``audio``.
+    """
+    device_spec = resolve_oom_device_spec(
+        OOM_INJECTION_CONFIG,
+        _stage_config_path_from_omni_server(omni_server_function),
+    )
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    mix_base = _mix_chat_completions_probe_payload(omni_server_function)
+    oom_contract_timeout_sec = 120
+    try:
+        chat_status, chat_body = post_json_raw(
+            host,
+            port,
+            "/v1/chat/completions",
+            {
+                "model": mix_base["model"],
+                "messages": mix_base["messages"],
+                "stream": False,
+                "modalities": ["text"],
+            },
+            timeout_sec=oom_contract_timeout_sec,
+        )
+        speech_status, speech_body = post_json_raw(
+            host,
+            port,
+            "/v1/chat/completions",
+            {
+                "model": mix_base["model"],
+                "messages": mix_base["messages"],
+                "stream": False,
+                "modalities": ["text", "audio"],
+            },
+            timeout_sec=oom_contract_timeout_sec,
+        )
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+    # Under OOM pressure both chat and speech should surface runtime-class (5xx)
+    # failures instead of request-validation-class (4xx) errors.
+    chat_error = extract_openai_error_contract_from_bytes(chat_body)
+    speech_error = extract_openai_error_contract_from_bytes(speech_body)
+    print(chat_status, chat_error, speech_status, speech_error)
+
+    assert chat_status >= 500, f"expected chat error under OOM, got status={chat_status}"
+    assert speech_status >= 500, f"expected speech runtime error under OOM, got status={speech_status}"
+
+    assert chat_error is not None, f"chat error payload not OpenAI-compatible: {chat_body[:300]!r}"
+    assert speech_error is not None, f"speech error payload not OpenAI-compatible: {speech_body[:300]!r}"
+    assert "code" in chat_error, f"chat error lacks code field: {chat_error!r}"
+    assert "code" in speech_error, f"speech error lacks code field: {speech_error!r}"
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None:
+    device_spec = resolve_oom_device_spec(
+        OOM_INJECTION_CONFIG,
+        _stage_config_path_from_omni_server(omni_server_function),
+    )
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    try:
+        video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(1280, 720, 161)['base64']}"
+        image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+        audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(20, 1)['base64']}"
+        messages = dummy_messages_from_mix_data(
+            system_prompt=_get_system_prompt(),
+            video_data_url=video_data_url,
+            image_data_url=image_data_url,
+            audio_data_url=audio_data_url,
+            content_text=f"{_get_mix_prompt()} " * 200,
+        )
+        request_config = {
+            "model": omni_server_function.model,
+            "messages": messages,
+            "stream": True,
+            "key_words": {"audio": ["test"]},
+        }
+        try:
+            openai_client_function.send_omni_request(request_config, request_num=1)
+        except Exception as exc:
+            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+        else:
+            pytest.fail("expected large chat payload request failure during GPU OOM injection")
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None:
+    device_spec = resolve_oom_device_spec(
+        OOM_INJECTION_CONFIG,
+        _stage_config_path_from_omni_server(omni_server_function),
+    )
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    try:
+        video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(1280, 720, 161)['base64']}"
+        image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+        audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(20, 1)['base64']}"
+        messages = dummy_messages_from_mix_data(
+            system_prompt=_get_system_prompt(),
+            video_data_url=video_data_url,
+            image_data_url=image_data_url,
+            audio_data_url=audio_data_url,
+            content_text=f"{_get_mix_prompt()} " * 200,
+        )
+        request_config = {
+            "model": omni_server_function.model,
+            "messages": messages,
+            "stream": True,
+            "modalities": ["text", "audio"],
+            "key_words": {"audio": ["test"]},
+        }
+        try:
+            openai_client_function.send_omni_request(request_config, request_num=4)
+        except Exception as exc:
+            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+        else:
+            pytest.fail("expected concurrent request failure under OOM injection")
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="VLLM::Worker",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_request_failure(
+    omni_server_after_fault_function, openai_client_function
+) -> None:
+    messages = dummy_messages_from_mix_data(
+        system_prompt=_get_system_prompt(),
+        content_text="What is the capital of China? Answer in 20 words.",
+    )
+    request_config = {
+        "model": omni_server_after_fault_function.model,
+        "messages": messages,
+        "stream": False,
+        "modalities": ["text"],
+        "key_words": {"text": ["beijing"]},
+    }
+    try:
+        openai_client_function.send_omni_request(request_config, request_num=1)
+    except Exception as exc:
+        assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+    else:
+        pytest.fail("expected request failure after process-kill injection")
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="VLLM::Worker",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_health_fast_fail_and_concurrent(
+    omni_server_after_fault_function,
+) -> None:
+    """Black-box: after worker SIGKILL, /health→503, chat fails fast, concurrent chat does not hang."""
+    host = omni_server_after_fault_function.host
+    port = omni_server_after_fault_function.port
+    model = omni_server_after_fault_function.model
+
+    deadline = time.monotonic() + 20.0
+    last_observation = ""
+    saw_503 = False
+    health_final_status: int | None = None
+    health_final_body = b""
+    while time.monotonic() < deadline:
+        try:
+            status, body = get_health_raw(host, port, timeout_sec=5)
+            last_observation = f"http={status}, body={body[:200]!r}"
+            health_final_status, health_final_body = status, body
+            if status == 503:
+                saw_503 = True
+                break
+        except Exception as exc:  # noqa: BLE001
+            last_observation = f"exception={exc!r}"
+        time.sleep(0.5)
+    assert saw_503, (
+        f"[process_kill health] expected /health to become 503 after fault injection, got {last_observation}"
+    )
+
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": "Say hello in one short sentence."}],
+        "stream": False,
+        "modalities": ["text"],
+    }
+    ff_status: int | None = None
+    ff_body = b""
+    ff_exc: BaseException | None = None
+    start = time.monotonic()
+    try:
+        ff_status, ff_body = post_json_raw(host, port, "/v1/chat/completions", payload, timeout_sec=20)
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[process_kill fast_fail] request did not fail fast after fault: {elapsed:.2f}s"
+        assert ff_status >= 500, (
+            f"[process_kill fast_fail] expected server-side failure after fault, "
+            f"got status={ff_status}, body={ff_body[:200]!r}"
+        )
+    except Exception as exc:
+        ff_exc = exc
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[process_kill fast_fail] request exception was too slow after fault: {elapsed:.2f}s"
+
+    payload_json = json.dumps(
+        {
+            "model": model,
+            "messages": [{"role": "user", "content": "What is the capital of China? Answer in one word."}],
+            "stream": False,
+            "modalities": ["text"],
+        }
+    )
+    start = time.monotonic()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as executor:
+        futures = [
+            executor.submit(
+                post_chat_completions_raw,
+                host,
+                port,
+                payload_json,
+                timeout_sec=20,
+            )
+            for _ in range(4)
+        ]
+        done, pending = concurrent.futures.wait(
+            futures,
+            timeout=30,
+            return_when=concurrent.futures.ALL_COMPLETED,
+        )
+
+    elapsed = time.monotonic() - start
+    assert not pending, f"[process_kill concurrent] some fault-time requests hung: pending={len(pending)}"
+    assert elapsed < 30, f"[process_kill concurrent] fault-time request convergence is too slow: {elapsed:.2f}s"
+
+    fault_observed = False
+    conc_debug: list[Any] = []
+    for future in done:
+        try:
+            status, body = future.result()
+            conc_debug.append((status, body[:200]))
+            if status >= 500:
+                fault_observed = True
+        except Exception as exc:
+            conc_debug.append(repr(exc))
+            fault_observed = True
+    # DEBUG: remove before merge
+    print(
+        health_final_status,
+        health_final_body[:200],
+        ff_status,
+        ff_body[:200],
+        ff_exc,
+        conc_debug,
+    )
+    assert fault_observed, (
+        "[process_kill concurrent] expected at least one request to fail after process-kill fault injection"
+    )
+
+
+@pytest.mark.slow
+@pytest.mark.skip(reason="issue#2327")
+@hardware_test(res={"cuda": "H100"}, num_cards=2)
+@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_state_converges_after_fault_removed(
+    omni_server_function,
+) -> None:
+    """Black-box: under bounded GPU pressure, a heavy mix chat fails; after hog stops, the same request succeeds.
+
+    Uses a milder sidecar than ``OOM_INJECTION_CONFIG`` (``strict=False`` and a capped
+    ``target_mem_ratio``) so the serving process stays alive: pressure should starve
+    the multimodal forward pass, not exhaust the whole device allocator.
+
+    Fault and recovery phases use raw ``/v1/chat/completions`` HTTP so tests can assert
+    on status codes and JSON error payloads (or fault keywords in streamed bodies).
+    """
+    device_spec = resolve_oom_device_spec(
+        OOM_RECOVER_INJECTION_CONFIG,
+        _stage_config_path_from_omni_server(omni_server_function),
+    )
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_RECOVER_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_RECOVER_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_RECOVER_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_RECOVER_INJECTION_CONFIG["strict"],
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    mix_payload = _mix_chat_completions_probe_payload(omni_server_function)
+    mix_chat_timeout_sec = 180
+    print(
+        "[oom-recover] injection ready "
+        f"device={device_spec} target_mem_ratio={OOM_RECOVER_INJECTION_CONFIG['target_mem_ratio']} "
+        f"strict={OOM_RECOVER_INJECTION_CONFIG['strict']} hold_seconds={OOM_RECOVER_INJECTION_CONFIG['hold_seconds']}"
+    )
+    print(f"[oom-recover] target endpoint http://{host}:{port}/v1/chat/completions")
+
+    failure_observed = False
+    fault_status: int | None = None
+    fault_body = b""
+    try:
+        try:
+            print(f"[oom-recover] fault-phase request start timeout={mix_chat_timeout_sec}s")
+            fault_status, fault_body = post_json_raw(
+                host,
+                port,
+                "/v1/chat/completions",
+                mix_payload,
+                timeout_sec=mix_chat_timeout_sec,
+            )
+            print(f"[oom-recover] fault-phase request done status={fault_status} body_prefix={fault_body[:200]!r}")
+        except Exception as exc:
+            failure_observed = True
+            print(f"[oom-recover] fault-phase request raised {type(exc).__name__}: {exc!r}")
+            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+        else:
+            assert fault_status is not None
+            failure_observed = fault_status != 200 or _fault_keywords_match_response(
+                status=fault_status,
+                body=fault_body,
+            )
+            assert failure_observed, (
+                "expected mix multimodal request failure while OOM pressure is active; "
+                f"status={fault_status}, body_prefix={fault_body[:500]!r}"
+            )
+            if fault_status >= 400:
+                err = extract_openai_error_contract_from_bytes(fault_body)
+                if err is None:
+                    assert _fault_keywords_match_response(status=fault_status, body=fault_body), (
+                        "non-2xx without OpenAI-style error object should still mention fault hints in body; "
+                        f"status={fault_status}, body_prefix={fault_body[:500]!r}"
+                    )
+                else:
+                    assert isinstance(err.get("message"), str) and err["message"].strip(), (
+                        f"structured error should carry a non-empty message: {err!r}"
+                    )
+            else:
+                assert _fault_keywords_match_response(status=fault_status, body=fault_body), (
+                    "HTTP 200 under pressure should still surface fault hints in the response body; "
+                    f"body_prefix={fault_body[:500]!r}"
+                )
+    finally:
+        print("[oom-recover] stopping OOM sidecar(s)")
+        stop_gpu_oom_hogs(handle)
+        print("[oom-recover] OOM sidecar(s) stopped")
+
+    recovery_deadline = time.monotonic() + 90.0
+    terminal_health: int | None = None
+    unreachable_streak = 0
+    last_health_exc: BaseException | None = None
+    while time.monotonic() < recovery_deadline:
+        try:
+            status, _ = get_health_raw(host, port, timeout_sec=5)
+            unreachable_streak = 0
+            if status in (200, 503):
+                terminal_health = status
+                print(f"[oom-recover] terminal health reached status={terminal_health}")
+                break
+        except Exception as exc:
+            last_health_exc = exc
+            if _looks_like_server_unreachable(exc):
+                unreachable_streak += 1
+                if unreachable_streak >= 5:
+                    pytest.fail(
+                        "after OOM sidecar stopped, /health is unreachable (connection refused / reset). "
+                        "The APIServer process likely exited (e.g. orchestrator thread crash under OOM); "
+                        "this test expects the server to stay up for post-fault health polling. "
+                        f"last_exc={last_health_exc!r}"
+                    )
+            else:
+                unreachable_streak = 0
+        time.sleep(1.0)
+    else:
+        pytest.fail(
+            "server did not converge to a terminal health state after OOM pressure was removed; "
+            f"last_health_exc={last_health_exc!r}"
+        )
+
+    probe_payload = {
+        "model": omni_server_function.model,
+        "messages": [{"role": "user", "content": "What is the capital of China? Answer in one word."}],
+        "stream": False,
+        "modalities": ["text"],
+    }
+    start = time.monotonic()
+    post_mix_status: int | None = None
+    post_mix_body = b""
+    request_status: int | None = None
+    try:
+        assert terminal_health is not None
+        if terminal_health == 200:
+            print(f"[oom-recover] recovery-phase mix request start timeout={mix_chat_timeout_sec}s")
+            post_mix_status, post_mix_body = post_json_raw(
+                host,
+                port,
+                "/v1/chat/completions",
+                mix_payload,
+                timeout_sec=mix_chat_timeout_sec,
+            )
+            print(
+                "[oom-recover] recovery-phase mix request done "
+                f"status={post_mix_status} body_prefix={post_mix_body[:200]!r}"
+            )
+        else:
+            print("[oom-recover] terminal health=503, send lightweight probe request")
+            request_status, _ = post_json_raw(host, port, "/v1/chat/completions", probe_payload, timeout_sec=20)
+            print(f"[oom-recover] lightweight probe done status={request_status}")
+    except Exception:
+        post_mix_status = None
+        request_status = None
+    elapsed = time.monotonic() - start
+    assert elapsed < 200, f"post-fault request should not hang after OOM removal: {elapsed:.2f}s"
+
+    if terminal_health == 200:
+        assert post_mix_status == 200, (
+            "health recovered but repeated mix multimodal request did not return HTTP 200; "
+            f"status={post_mix_status}, body_prefix={post_mix_body[:600]!r}"
+        )
+        recover_err = extract_openai_error_contract_from_bytes(post_mix_body)
+        assert recover_err is None, f"unexpected error object after recovery: {recover_err!r}"
+    else:
+        assert request_status is None or request_status >= 500, (
+            "unhealthy terminal state should fail fast on requests, "
+            f"got health={terminal_health}, request_status={request_status}"
+        )

tests/dfx/reliability/test_reliability_wan22.py

@@ -0,0 +1,411 @@
+"""
+Wan2.2 reliability integration tests.
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+import os
+import time
+from pathlib import Path
+from typing import Any
+
+import pytest
+import requests
+
+from tests.dfx.conftest import (
+    assert_fault_exception,
+    create_reliability_omni_server_params,
+    resolve_oom_device_spec,
+    supports_video_generation,
+)
+from tests.dfx.reliability.helpers import (
+    get_health_raw,
+    inject_gpu_oom,
+    make_process_kill_fault_injector,
+    stop_gpu_oom_hogs,
+)
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import generate_synthetic_image
+
+RELIABILITY_SCENARIOS: list[dict[str, Any]] = [
+    {
+        "test_name": "wan22_t2v_reliability_default",
+        "server_params": {
+            "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+            "server_args": [
+                "--num-gpus",
+                "1",
+                "--boundary-ratio",
+                "0.875",
+                "--flow-shift",
+                "5.0",
+                "--disable-log-stats",
+            ],
+        },
+    }
+]
+
+E2E_STAGE_CONFIGS_DIR = Path(__file__).resolve().parent.parent / "e2e" / "stage_configs"
+OOM_INJECTION_CONFIG = {
+    "target_mem_ratio": 0.95,
+    "hold_seconds": 0,
+    "startup_timeout_sec": 20,
+    "strict": True,
+}
+FAULT_ERROR_KEYWORDS = (
+    "oom",
+    "out of memory",
+    "cuda",
+    "job failed",
+    "unknown error",
+    "internal server error",
+    "500 server error",
+)
+PROCESS_KILL_ERROR_KEYWORDS = (
+    "timeout",
+    "did not complete within",
+    "connection",
+    "engine",
+    "orchestrator",
+    "dead",
+    "internal",
+    "500",
+    "503",
+)
+
+WAN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, E2E_STAGE_CONFIGS_DIR)
+DIFFUSION_VIDEO_PARAMS = [param for param in WAN_PARAMS if supports_video_generation(param.model)]
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_fault_gpu_oom_video_large_request_failure(omni_server_function, openai_client_function) -> None:
+    stage_config_path = getattr(omni_server_function, "stage_config_path", None)
+    device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    try:
+        image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+        request_config = {
+            "form_data": {
+                "prompt": "Generate a realistic road-driving video with camera motion.",
+                "width": 512,
+                "height": 512,
+                "fps": 8,
+                "num_frames": 8,
+                "guidance_scale": 1.0,
+                "flow_shift": 5.0,
+                "num_inference_steps": 4,
+                "seed": 42,
+            },
+            "image_reference": image_data_url,
+            "stream": False,
+        }
+        try:
+            openai_client_function.send_video_diffusion_request(request_config, request_num=1)
+        except Exception as exc:
+            assert_fault_exception(exc, FAULT_ERROR_KEYWORDS)
+        else:
+            pytest.fail("expected large /v1/videos request failure during GPU OOM injection")
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="multiprocessing.spawn",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_video_request_failure(
+    omni_server_after_fault_function,
+    openai_client_function,
+) -> None:
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+    request_config = {
+        "form_data": {
+            "prompt": "Generate a realistic road-driving video with camera motion.",
+            "width": 512,
+            "height": 512,
+            "fps": 8,
+            "num_frames": 8,
+            "guidance_scale": 1.0,
+            "flow_shift": 5.0,
+            "num_inference_steps": 4,
+            "seed": 42,
+        },
+        "image_reference": image_data_url,
+        "stream": False,
+    }
+    try:
+        openai_client_function.send_video_diffusion_request(request_config, request_num=1)
+    except Exception as exc:
+        assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
+    else:
+        pytest.fail("expected /v1/videos request failure after process-kill injection")
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(os.name == "nt", reason="process-kill injection helper is POSIX-only")
+@pytest.mark.parametrize(
+    "fault_injector",
+    [
+        pytest.param(
+            make_process_kill_fault_injector(
+                grep_patterns="multiprocessing.spawn",
+                signal_name="SIGKILL",
+                limit=1,
+                post_kill_wait_seconds=2.0,
+            ),
+            id="runtime_process_chain",
+        ),
+    ],
+    indirect=True,
+)
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_fault_process_kill_video_health_fast_fail_and_concurrent(
+    omni_server_after_fault_function,
+    openai_client_function,
+) -> None:
+    """Black-box: after process kill, /v1/videos fails fast and concurrent calls don't hang.
+
+    /health→503 is checked last: optional ``pytest.skip`` then ``assert`` (drop the ``if`` to harden).
+    """
+    host = omni_server_after_fault_function.host
+    port = omni_server_after_fault_function.port
+    url = f"http://{host}:{port}/v1/videos"
+    image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(1280, 720)['base64']}"
+    request_config = {
+        "form_data": {
+            "prompt": "Generate a realistic road-driving video with camera motion.",
+            "width": 512,
+            "height": 512,
+            "fps": 8,
+            "num_frames": 8,
+            "guidance_scale": 1.0,
+            "flow_shift": 5.0,
+            "num_inference_steps": 4,
+            "seed": 42,
+        },
+        "image_reference": image_data_url,
+        "stream": False,
+    }
+
+    # Phase-1: one new /v1/videos request should fail fast.
+    payload = {
+        "prompt": "fast-fail check",
+        "width": "512",
+        "height": "512",
+        "fps": "8",
+        "num_frames": "8",
+        "num_inference_steps": "4",
+    }
+    start = time.monotonic()
+    try:
+        response = requests.post(
+            url,
+            data=payload,
+            headers={"Accept": "application/json"},
+            timeout=20,
+        )
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[process_kill fast_fail] /v1/videos did not fail fast after fault: {elapsed:.2f}s"
+        assert response.status_code >= 500, (
+            "[process_kill fast_fail] expected server-side error after fatal fault, "
+            f"got status={response.status_code}, body={response.text[:200]!r}"
+        )
+    except Exception:
+        elapsed = time.monotonic() - start
+        assert elapsed < 15, f"[process_kill fast_fail] /v1/videos exception was too slow after fault: {elapsed:.2f}s"
+
+    # Phase-2: concurrent requests should converge and at least one should fail.
+    start = time.monotonic()
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+        futures = [
+            executor.submit(openai_client_function.send_video_diffusion_request, request_config, 1) for _ in range(3)
+        ]
+        done, pending = concurrent.futures.wait(
+            futures,
+            timeout=40,
+            return_when=concurrent.futures.ALL_COMPLETED,
+        )
+
+    elapsed = time.monotonic() - start
+    assert not pending, f"[process_kill concurrent] some fault-time video requests hung: pending={len(pending)}"
+    assert elapsed < 40, f"[process_kill concurrent] fault-time video request convergence is too slow: {elapsed:.2f}s"
+
+    fault_observed = False
+    conc_debug: list[Any] = []
+    for future in done:
+        try:
+            future.result()
+            conc_debug.append("ok")
+        except Exception as exc:
+            fault_observed = True
+            conc_debug.append(repr(exc))
+            assert_fault_exception(exc, PROCESS_KILL_ERROR_KEYWORDS)
+    assert fault_observed, (
+        "[process_kill concurrent] expected at least one /v1/videos request to fail after fault; "
+        f"conc_debug={conc_debug}"
+    )
+
+    # Phase-3: /health→503 after fault. Optional skip today; remove the ``if`` later to harden.
+    deadline = time.monotonic() + 20.0
+    last_observation = ""
+    final_health_status: int | None = None
+    while time.monotonic() < deadline:
+        try:
+            status, body = get_health_raw(host, port, timeout_sec=5)
+            last_observation = f"http={status}, body={body[:200]!r}"
+            final_health_status = status
+            if status == 503:
+                break
+        except Exception as exc:  # noqa: BLE001
+            last_observation = f"exception={exc!r}"
+        time.sleep(0.5)
+    if final_health_status != 503:
+        pytest.skip("issue#3050")
+    assert final_health_status == 503, (
+        "[process_kill health] expected /health 503 after fatal fault, "
+        f"got status={final_health_status}, last_observation={last_observation}"
+    )
+
+
+@pytest.mark.slow
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.skip(reason="issue#2327")
+@pytest.mark.parametrize("omni_server_function", DIFFUSION_VIDEO_PARAMS, indirect=True)
+def test_reliability_video_oom_recovers_after_fault_removed(
+    omni_server_function,
+) -> None:
+    """Black-box: after removing OOM pressure, service converges to a terminal state and requests do not hang."""
+    stage_config_path = getattr(omni_server_function, "stage_config_path", None)
+    device_spec = resolve_oom_device_spec(OOM_INJECTION_CONFIG, stage_config_path)
+    handle = inject_gpu_oom(
+        device=device_spec,
+        target_mem_ratio=OOM_INJECTION_CONFIG["target_mem_ratio"],
+        hold_seconds=OOM_INJECTION_CONFIG["hold_seconds"],
+        startup_timeout_sec=OOM_INJECTION_CONFIG["startup_timeout_sec"],
+        strict=OOM_INJECTION_CONFIG["strict"],
+    )
+    host = omni_server_function.host
+    port = omni_server_function.port
+    create_url = f"http://{host}:{port}/v1/videos"
+
+    failure_observed = False
+    try:
+        for _ in range(3):
+            try:
+                response = requests.post(
+                    create_url,
+                    data={
+                        "prompt": "oom recover probe",
+                        "width": "512",
+                        "height": "512",
+                        "fps": "8",
+                        "num_frames": "8",
+                        "num_inference_steps": "4",
+                    },
+                    headers={"Accept": "application/json"},
+                    timeout=25,
+                )
+                if response.status_code >= 500:
+                    failure_observed = True
+                    break
+            except Exception:
+                failure_observed = True
+                break
+            time.sleep(1.0)
+    finally:
+        stop_gpu_oom_hogs(handle)
+
+    assert failure_observed, "expected at least one video request failure while OOM pressure is active"
+
+    recovery_deadline = time.monotonic() + 90.0
+    terminal_health: int | None = None
+    unreachable_streak = 0
+    last_health_exc: BaseException | None = None
+    while time.monotonic() < recovery_deadline:
+        try:
+            status, _ = get_health_raw(host, port, timeout_sec=5)
+            unreachable_streak = 0
+            if status in (200, 503):
+                terminal_health = status
+                break
+        except Exception as exc:  # noqa: BLE001
+            last_health_exc = exc
+            msg = str(exc).lower()
+            if "connection refused" in msg or "actively refused" in msg:
+                unreachable_streak += 1
+                if unreachable_streak >= 5:
+                    pytest.fail(
+                        "after OOM sidecar stopped, /health is unreachable repeatedly; "
+                        "the APIServer process likely exited unexpectedly "
+                        f"(last_exc={last_health_exc!r})"
+                    )
+            else:
+                unreachable_streak = 0
+        time.sleep(1.0)
+    else:
+        pytest.fail(
+            "wan22 server did not converge to terminal health (200/503) after OOM pressure was removed; "
+            f"last_health_exc={last_health_exc!r}"
+        )
+
+    assert terminal_health is not None
+    start = time.monotonic()
+    recovery_resp: requests.Response | None = None
+    try:
+        recovery_resp = requests.post(
+            create_url,
+            data={
+                "prompt": "post-recovery admission check",
+                "width": "512",
+                "height": "512",
+                "fps": "8",
+                "num_frames": "8",
+                "num_inference_steps": "4",
+            },
+            headers={"Accept": "application/json"},
+            timeout=30,
+        )
+        recovery_status = recovery_resp.status_code
+        recovery_text_prefix = recovery_resp.text[:300]
+    except Exception as exc:  # noqa: BLE001
+        recovery_status = None
+        recovery_text_prefix = repr(exc)
+    elapsed = time.monotonic() - start
+    assert elapsed < 30, f"post-fault /v1/videos should not hang after OOM removal: {elapsed:.2f}s"
+
+    if terminal_health == 200:
+        assert recovery_status == 200, (
+            "health recovered but /v1/videos admission did not succeed; "
+            f"status={recovery_status}, body={recovery_text_prefix!r}"
+        )
+        assert recovery_resp is not None
+        recovery_payload = recovery_resp.json()
+        assert "id" in recovery_payload, f"post-recovery create payload missing id: {recovery_payload!r}"
+    else:
+        assert recovery_status is None or recovery_status >= 500, (
+            "unhealthy terminal state should fail fast on /v1/videos, "
+            f"got health={terminal_health}, request_status={recovery_status}, body={recovery_text_prefix!r}"
+        )

tests/diffusion/cache/test_teacache_extractors.py

@@ -13,6 +13,8 @@ used by TeaCache. Each model's extractor can be tested by:
 
 Currently implemented:
 - TestFlux2KleinExtractor: Flux2Klein model extractor
+- TestFlux2Extractor: Flux2 model extractor
+- TestFluxExtractor: Flux model extractor
 """
 
 from abc import ABC, abstractmethod
@@ -22,7 +24,12 @@ import pytest
 import torch
 
 from tests.helpers.mark import hardware_test
-from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_context, extract_flux2_klein_context
+from vllm_omni.diffusion.cache.teacache.extractors import (
+    extract_flux2_context,
+    extract_flux2_klein_context,
+    extract_flux_context,
+)
+from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformer2DModel
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
@@ -234,7 +241,7 @@ class TestFlux2Extractor(BaseExtractorTest):
         in_channels (128) to inner_dim (6144), so modulated_input should match
         the projected shape, not the input shape.
         """
-        context = extract_flux2_klein_context(flux2_module, **sample_inputs)
+        context = extract_flux2_context(flux2_module, **sample_inputs)
 
         batch_size, img_seq_len, _ = sample_inputs["hidden_states"].shape
         inner_dim = flux2_module.inner_dim
@@ -277,3 +284,139 @@ class TestFlux2Extractor(BaseExtractorTest):
                 img_ids=torch.randint(0, 64, (1, 1024, 4)),
                 txt_ids=torch.randint(0, 64, (1, 512, 4)),
             )
+
+
+@pytest.mark.cpu
+class TestFluxExtractor(BaseExtractorTest):
+    """Test extract_flux_context function."""
+
+    @pytest.fixture(autouse=True)
+    def cpu_vllm_config(self):
+        """Force CPU custom-op dispatch for this test class."""
+        from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
+
+        with set_current_vllm_config(VllmConfig(device_config=DeviceConfig(device="cpu"))):
+            yield
+
+    @pytest.fixture(autouse=True)
+    def mock_flux_attention_backend(self):
+        """Use the SDPA backend so FLUX can be instantiated in CPU tests."""
+        from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
+
+        with patch("vllm_omni.diffusion.attention.layer.get_attn_backend", return_value=SDPABackend):
+            yield
+
+    def get_extractor(self):
+        return extract_flux_context
+
+    @pytest.fixture
+    def flux_module(self):
+        """Create a minimal FluxTransformer2DModel for testing."""
+        return FluxTransformer2DModel(
+            num_layers=2,
+            num_single_layers=2,
+            num_attention_heads=2,
+            attention_head_dim=16,
+            joint_attention_dim=32,
+            pooled_projection_dim=16,
+            axes_dims_rope=(4, 4, 8),
+        )
+
+    @pytest.fixture
+    def flux_module_without_guidance(self):
+        """Create a minimal non-guidance-distilled FLUX transformer."""
+        return FluxTransformer2DModel(
+            num_layers=2,
+            num_single_layers=2,
+            num_attention_heads=2,
+            attention_head_dim=16,
+            joint_attention_dim=32,
+            pooled_projection_dim=16,
+            guidance_embeds=False,
+            axes_dims_rope=(4, 4, 8),
+        )
+
+    def get_module(self, flux_module):
+        return flux_module
+
+    @pytest.fixture
+    def sample_inputs(self):
+        """Create sample input tensors for Flux."""
+        batch_size = 1
+        img_seq_len = 16
+        txt_seq_len = 8
+        in_channels = 64  # Flux default in_channels
+        txt_dim = 32
+        pooled_dim = 16
+
+        return {
+            "hidden_states": torch.randn(batch_size, img_seq_len, in_channels),
+            "encoder_hidden_states": torch.randn(batch_size, txt_seq_len, txt_dim),
+            "pooled_projections": torch.randn(batch_size, pooled_dim),
+            "timestep": torch.tensor([500]),
+            "img_ids": torch.randint(0, 64, (batch_size, img_seq_len, 3)),
+            "txt_ids": torch.randint(0, 64, (batch_size, txt_seq_len, 3)),
+            "guidance": torch.tensor([3.5]),
+        }
+
+    def get_sample_inputs(self, sample_inputs):
+        return sample_inputs
+
+    def test_modulated_input_shape(self, flux_module, sample_inputs):
+        """Test that modulated_input has the projected FLUX inner dimension."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+
+        batch_size, img_seq_len, _ = sample_inputs["hidden_states"].shape
+        assert context.modulated_input.shape == (batch_size, img_seq_len, flux_module.inner_dim)
+
+    def test_run_transformer_blocks_callable(self, flux_module, sample_inputs):
+        """Test that run_transformer_blocks is callable."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+        assert callable(context.run_transformer_blocks)
+
+    def test_postprocess_callable(self, flux_module, sample_inputs):
+        """Test that postprocess is callable."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+        assert callable(context.postprocess)
+
+    def test_postprocess_output_shape(self, flux_module, sample_inputs):
+        """Test that postprocess projects back to the input channel width."""
+        context = extract_flux_context(flux_module, **sample_inputs)
+        output = context.postprocess(context.hidden_states)
+
+        assert output.sample.shape == sample_inputs["hidden_states"].shape
+
+    def test_postprocess_return_tuple_when_return_dict_false(self, flux_module, sample_inputs):
+        """Test that postprocess honors return_dict=False."""
+        context = extract_flux_context(flux_module, **sample_inputs, return_dict=False)
+        output = context.postprocess(context.hidden_states)
+
+        assert isinstance(output, tuple)
+        assert len(output) == 1
+        assert output[0].shape == sample_inputs["hidden_states"].shape
+
+    def test_without_guidance(self, flux_module_without_guidance, sample_inputs):
+        """Test context extraction works for FLUX variants without guidance embeddings."""
+        inputs = sample_inputs.copy()
+        inputs["guidance"] = None
+
+        context = extract_flux_context(flux_module_without_guidance, **inputs)
+
+        assert context is not None
+        assert context.temb is not None
+
+    def test_invalid_module_raises_error(self):
+        """Test that invalid module without transformer_blocks raises ValueError."""
+        invalid_module = Mock()
+        invalid_module.transformer_blocks = []
+
+        with pytest.raises(ValueError, match="Module must have transformer_blocks"):
+            extract_flux_context(
+                invalid_module,
+                hidden_states=torch.randn(1, 16, 64),
+                encoder_hidden_states=torch.randn(1, 8, 32),
+                pooled_projections=torch.randn(1, 16),
+                timestep=torch.tensor([500]),
+                img_ids=torch.randint(0, 64, (1, 16, 3)),
+                txt_ids=torch.randint(0, 64, (1, 8, 3)),
+            )

tests/diffusion/test_diffusion_engine_metrics.py

@@ -0,0 +1,105 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Source-level regression tests for DiffusionEngine.
+
+These tests verify naming conventions and patterns by inspecting source code
+at the function level using AST. They are intentionally coupled to the source
+layout and should be updated whenever the metrics construction code is
+refactored.
+"""
+
+from __future__ import annotations
+
+import ast
+import os
+
+_ENGINE_PATH = os.path.normpath(
+    os.path.join(
+        os.path.dirname(__file__),
+        os.pardir,
+        os.pardir,
+        "vllm_omni",
+        "diffusion",
+        "diffusion_engine.py",
+    )
+)
+
+
+def _read_engine_source() -> str:
+    with open(_ENGINE_PATH) as f:
+        return f.read()
+
+
+def _get_function_source(source: str, class_name: str | None, func_name: str) -> str:
+    """Extract the source of a specific function/method using AST.
+
+    Args:
+        source: Full file source code.
+        class_name: Enclosing class name, or None for module-level functions.
+        func_name: Function/method name.
+
+    Returns:
+        Source code of the function body.
+    """
+    tree = ast.parse(source)
+    if class_name is not None:
+        for node in ast.walk(tree):
+            if isinstance(node, ast.ClassDef) and node.name == class_name:
+                for item in node.body:
+                    if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == func_name:
+                        result = ast.get_source_segment(source, item)
+                        assert result is not None, f"{class_name}.{func_name} source not found"
+                        return result
+    else:
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == func_name:
+                result = ast.get_source_segment(source, node)
+                assert result is not None, f"{func_name} source not found"
+                return result
+    raise AssertionError(f"Function {class_name + '.' if class_name else ''}{func_name} not found in source")
+
+
+class TestMetricKeys:
+    """Verify metric naming conventions in DiffusionEngine.step()."""
+
+    def test_no_duplicate_preprocess_key(self) -> None:
+        """step() should not contain 'preprocessing_time_ms' (duplicate of 'preprocess_time_ms')."""
+        source = _read_engine_source()
+        step_source = _get_function_source(source, "DiffusionEngine", "step")
+        assert "preprocessing_time_ms" not in step_source, (
+            "Found duplicate key 'preprocessing_time_ms' in step() — should only use 'preprocess_time_ms'"
+        )
+
+    def test_metric_key_naming_consistency(self) -> None:
+        """In step(): exec_time should use exec_total_time, total_time should use step_total_ms."""
+        source = _read_engine_source()
+        step_source = _get_function_source(source, "DiffusionEngine", "step")
+        lines = step_source.split("\n")
+
+        found_exec = False
+        found_total = False
+        for line in lines:
+            if '"diffusion_engine_exec_time_ms"' in line:
+                found_exec = True
+                assert "exec_total_time" in line, (
+                    "diffusion_engine_exec_time_ms should measure executor time only (exec_total_time)"
+                )
+            if '"diffusion_engine_total_time_ms"' in line:
+                found_total = True
+                assert "step_total_ms" in line, (
+                    "diffusion_engine_total_time_ms should measure full step time (step_total_ms)"
+                )
+        assert found_exec, "diffusion_engine_exec_time_ms key not found in step()"
+        assert found_total, "diffusion_engine_total_time_ms key not found in step()"
+
+
+class TestDummyRunAllocation:
+    """Verify _dummy_run generates exact-sized audio arrays."""
+
+    def test_no_oversized_allocation(self) -> None:
+        """_dummy_run should not allocate more audio than needed."""
+        source = _read_engine_source()
+        dummy_source = _get_function_source(source, "DiffusionEngine", "_dummy_run")
+        assert "audio_sr * audio_duration_sec" not in dummy_source, (
+            "_dummy_run should generate exact-sized audio, not allocate and slice"
+        )

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -12,6 +12,7 @@ from pytest_mock import MockerFixture
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.request import RequestStatus
 
+from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
@@ -111,7 +112,11 @@ def test_load_poll(build_adapter):
     request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
 
     adapter.load_async(request)
-    payload = {"code_predictor_codes": [[1]], "hidden_states": torch.tensor([[2.0]]), "finished": True}
+    payload: OmniPayload = {
+        "codes": {"audio": [[1]]},
+        "hidden_states": {"output": torch.tensor([[2.0]])},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
     connector.get.return_value = (payload, 16)
     adapter._poll_single_request(request)
 
@@ -154,12 +159,49 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon
 def test_update_request_payload(build_adapter):
     adapter, _ = build_adapter()
 
-    adapter._update_request_payload("ext", {"h": torch.tensor([[1.0]]), "codes": [1], "finished": False})
-    merged = adapter._update_request_payload("ext", {"h": torch.tensor([[2.0]]), "codes": [2], "finished": True})
+    first: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[1.0]])},
+        "codes": {"audio": [1]},
+        "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+    }
+    adapter._update_request_payload("ext", first)
+    second: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[2.0]])},
+        "codes": {"audio": [2]},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
+    merged = adapter._update_request_payload("ext", second)
 
-    assert torch.equal(merged["h"], torch.tensor([[1.0], [2.0]]))
-    assert merged["codes"] == [1, 2]
-    assert merged["finished"] is True
+    assert torch.equal(merged["hidden_states"]["output"], torch.tensor([[1.0], [2.0]]))
+    assert merged["codes"]["audio"] == [1, 2]
+    assert merged["meta"]["finished"].item() is True
+
+
+def test_load_poll_ar_request_additional_information_concats_tensors(build_adapter):
+    adapter, connector = build_adapter(stage_id=2, model_mode="ar")
+    request = _req("req-merged", RequestStatus.WAITING, external_req_id="ext-merged")
+
+    adapter.request_ids_mapping["req-merged"] = "ext-merged"
+    adapter.request_payload["ext-merged"] = {
+        "hidden_states": {"output": torch.tensor([[1.0]])},
+        "ids": {"prompt": [11, 12]},
+        "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+    }
+    payload: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[2.0]])},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
+    connector.get.return_value = (payload, 8)
+
+    adapter._poll_single_request(request)
+
+    assert torch.equal(
+        request.additional_information["hidden_states"]["output"],
+        torch.tensor([[1.0], [2.0]]),
+    )
+    # Keys absent from the new chunk are dropped (matches main's behavior).
+    assert "ids" not in request.additional_information
+    assert request.additional_information["meta"]["finished"].item() is True
 
 
 def test_process_and_restore_queues(build_adapter):
@@ -321,7 +363,10 @@ def test_cleanup_after_poll_flow(build_adapter):
     adapter.load_async(request)
 
     adapter.request_ids_mapping["req-flow"] = "ext-flow"
-    payload = {"hidden_states": torch.tensor([[1.0]]), "finished": True}
+    payload: OmniPayload = {
+        "hidden_states": {"output": torch.tensor([[1.0]])},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
     connector.get.return_value = (payload, 8)
     adapter._poll_single_request(request)
 

tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py

@@ -17,6 +17,7 @@ import time
 import pytest
 import torch
 
+from tests.helpers.mark import hardware_test
 from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
     ManagedBuffer,
     MooncakeTransferEngineConnector,
@@ -24,7 +25,7 @@ from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_c
 )
 
 # All tests in this file require Mooncake TransferEngine and an RDMA environment.
-pytestmark = [pytest.mark.parallel, pytest.mark.gpu]
+pytestmark = [pytest.mark.parallel, pytest.mark.core_model]
 
 # ---------------------------------------------------------------------------
 # Shared helpers
@@ -194,7 +195,7 @@ def _md5(tensor: torch.Tensor) -> str:
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 class TestBasicConnector:
     """Verify connector initialization, put, cleanup, and health check."""
 
@@ -254,7 +255,7 @@ class TestBasicConnector:
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 class TestEndToEnd:
     """E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload, mixed types."""
 
@@ -447,7 +448,7 @@ class TestEndToEnd:
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 class TestLifecycle:
     """Close, context manager, double-close safety."""
 
@@ -475,9 +476,7 @@ class TestLifecycle:
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.cuda
-@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 class TestGPUPool:
     """GPU memory pool: initialization, put (CPU/GPU tensor), E2E transfer."""
 
@@ -521,12 +520,11 @@ class TestGPUPool:
 
 
 # ---------------------------------------------------------------------------
-# 5. Stress / Correctness tests (marked slow, skipped in quick CI)
+# 5. Stress / Correctness tests, slow
 # ---------------------------------------------------------------------------
 
 
-@pytest.mark.slow
-@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
 class TestStressCorrectness:
     """
     Slow but high-value regression tests: concurrent put+get with data

tests/e2e/offline_inference/test_bagel_text2img.py

@@ -304,11 +304,13 @@ 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"})
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
 def test_bagel_text2img_mooncake_connector(run_level):
     """Test Bagel text2img with Mooncake connector for inter-stage communication."""
     if not _is_mooncake_master_available():
-        pytest.skip("mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)")
+        raise RuntimeError(
+            "mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)"
+        )
     MOONCAKE_HOST = "127.0.0.1"
     MOONCAKE_RPC_PORT = _find_free_port()
     MOONCAKE_HTTP_PORT = _find_free_port()

tests/e2e/offline_inference/test_internvla_a1.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Offline E2E test for InternVLA-A1 open-loop inference."""
+
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+EXAMPLE_SCRIPT = REPO_ROOT / "examples" / "offline_inference" / "internvla_a1" / "end2end.py"
+
+
+def _required_env(name: str) -> str:
+    value = os.getenv(name)
+    if not value:
+        pytest.skip(f"{name} is not set")
+    return value
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.diffusion
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
+def test_internvla_a1_offline_open_loop(run_level: str) -> None:
+    if run_level != "advanced_model":
+        pytest.skip("InternVLA-A1 offline evaluation requires real local checkpoints and dataset.")
+
+    model_dir = _required_env("INTERNVLA_A1_MODEL_DIR")
+    dataset_dir = _required_env("INTERNVLA_A1_DATASET_DIR")
+    processor_dir = _required_env("INTERNVLA_A1_PROCESSOR_DIR")
+    cosmos_dir = _required_env("INTERNVLA_A1_COSMOS_DIR")
+
+    env = os.environ.copy()
+    env["INTERNVLA_A1_PROCESSOR_DIR"] = processor_dir
+    env["INTERNVLA_A1_COSMOS_DIR"] = cosmos_dir
+
+    with tempfile.TemporaryDirectory(prefix="internvla_a1_e2e_") as tmpdir:
+        output_dir = Path(tmpdir) / "outputs"
+        cmd = [
+            sys.executable,
+            str(EXAMPLE_SCRIPT),
+            "--model-dir",
+            model_dir,
+            "--dataset-dir",
+            dataset_dir,
+            "--output-dir",
+            str(output_dir),
+            "--num-samples",
+            "1",
+            "--num-episodes",
+            "1",
+            "--dtype",
+            "float32",
+            "--attn-implementation",
+            "eager",
+            "--skip-plots",
+        ]
+        subprocess.run(cmd, check=True, env=env)
+
+        summary_path = output_dir / "summary.json"
+        assert summary_path.exists(), f"Missing summary file: {summary_path}"
+
+        with open(summary_path, encoding="utf-8") as f:
+            summary = json.load(f)
+
+        assert summary["mode"] == "registry"
+        assert summary["attn_implementation"] == "eager"
+        assert summary["dtype"] == "float32"
+        assert len(summary["results"]) == 1
+        assert "registry" in summary["eval_summaries"]
+
+        eval_summary = summary["eval_summaries"]["registry"]
+        assert eval_summary["num_episodes"] == 1
+        assert eval_summary["average_mse"] is not None
+        assert eval_summary["average_mae"] is not None
+
+        log_path = output_dir / "registry" / "log.json"
+        assert log_path.exists(), f"Missing evaluation log: {log_path}"

tests/e2e/offline_inference/test_moss_tts_nano.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E offline inference tests for MOSS-TTS-Nano single-stage pipeline."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from vllm import SamplingParams
+
+from tests.helpers.mark import hardware_test
+from vllm_omni import Omni
+
+MODEL_NAME = "OpenMOSS-Team/MOSS-TTS-Nano"
+SAMPLE_RATE = 48000
+
+DEFAULT_SAMPLING = SamplingParams(
+    temperature=1.0,
+    top_p=1.0,
+    top_k=50,
+    max_tokens=4096,
+    seed=42,
+    detokenize=False,
+)
+
+
+def _build_request(
+    text: str,
+    voice: str = "Junhao",
+    mode: str = "voice_clone",
+    max_new_frames: int = 100,  # short for tests
+    seed: int = 42,
+) -> dict:
+    return {
+        "prompt": "<|im_start|>assistant\n",
+        "additional_information": {
+            "text": [text],
+            "voice": [voice],
+            "mode": [mode],
+            "max_new_frames": [max_new_frames],
+            "seed": [seed],
+        },
+    }
+
+
+def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
+    """Run a single request and return (waveform, sample_rate)."""
+    for stage_outputs in omni.generate(request, DEFAULT_SAMPLING):
+        for req_output in stage_outputs.request_output:
+            mm = req_output.outputs[0].multimodal_output
+            assert mm is not None, "Expected multimodal_output to be non-None"
+            audio = mm.get("audio")
+            sr = mm.get("sr")
+            assert audio is not None, "Expected 'audio' key in multimodal_output"
+            assert isinstance(audio, torch.Tensor), f"audio should be Tensor, got {type(audio)}"
+            return audio.cpu(), int(sr.item()) if sr is not None else SAMPLE_RATE
+    raise AssertionError("No stage outputs received")
+
+
+@pytest.fixture(scope="module")
+def omni_engine():
+    """Module-scoped Omni engine to avoid re-loading the model for each test."""
+    return Omni(model=MODEL_NAME, stage_init_timeout=180)
+
+
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"})
+def test_moss_tts_nano_english(omni_engine):
+    """English TTS produces non-empty 48 kHz stereo audio."""
+    req = _build_request("Hello, this is a test of MOSS-TTS-Nano.", voice="Ava")
+    audio, sr = _collect_audio(omni_engine, req)
+
+    assert sr == SAMPLE_RATE, f"Expected sample_rate={SAMPLE_RATE}, got {sr}"
+    assert audio.numel() > 0, "Audio tensor should not be empty"
+    assert not torch.all(audio == 0), "Audio should not be all-zeros (silence)"
+
+
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"})
+def test_moss_tts_nano_chinese(omni_engine):
+    """Chinese TTS produces non-empty audio."""
+    req = _build_request("你好，这是语音合成测试。", voice="Junhao")
+    audio, sr = _collect_audio(omni_engine, req)
+
+    assert sr == SAMPLE_RATE
+    assert audio.numel() > 0
+    assert not torch.all(audio == 0)
+
+
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"})
+def test_moss_tts_nano_deterministic(omni_engine):
+    """Same seed produces identical waveforms."""
+    req = _build_request("Reproducible output test.", voice="Adam", seed=123)
+    audio1, _ = _collect_audio(omni_engine, req)
+    audio2, _ = _collect_audio(omni_engine, req)
+
+    assert audio1.shape == audio2.shape, "Waveform shapes should match with same seed"
+    assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms should match with same seed"
+
+
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"})
+def test_moss_tts_nano_batch(omni_engine):
+    """Batch of two requests returns audio for each."""
+    requests = [
+        _build_request("First request.", voice="Ava"),
+        _build_request("Second request.", voice="Bella"),
+    ]
+    results = []
+    for stage_outputs in omni_engine.generate(requests, [DEFAULT_SAMPLING] * 2):
+        for req_output in stage_outputs.request_output:
+            mm = req_output.outputs[0].multimodal_output
+            assert mm is not None
+            results.append(mm["audio"].cpu())
+
+    assert len(results) == 2, f"Expected 2 outputs, got {len(results)}"
+    for i, audio in enumerate(results):
+        assert audio.numel() > 0, f"Audio {i} is empty"
+
+
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"})
+def test_moss_tts_nano_voice_presets(omni_engine):
+    """Multiple built-in voice presets all produce valid audio."""
+    voices = ["Junhao", "Ava", "Bella"]
+    for voice in voices:
+        req = _build_request(f"Testing voice {voice}.", voice=voice)
+        audio, sr = _collect_audio(omni_engine, req)
+        assert audio.numel() > 0, f"Voice {voice} produced empty audio"
+        assert sr == SAMPLE_RATE

tests/e2e/online_serving/test_moss_tts_nano.py

@@ -0,0 +1,117 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+E2E Online tests for MOSS-TTS-Nano model with text input and audio output.
+
+These tests verify the /v1/audio/speech endpoint works correctly with
+actual model inference, not mocks.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "OpenMOSS-Team/MOSS-TTS-Nano"
+
+
+def get_prompt(prompt_type="text"):
+    """Text prompt for text-to-audio tests."""
+    prompts = {
+        "text": "Hello, this is a test of MOSS-TTS-Nano text to speech synthesis.",
+        "chinese": "你好，这是MOSS-TTS-Nano的语音合成测试。",
+    }
+    return prompts.get(prompt_type, prompts["text"])
+
+
+tts_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("moss_tts_nano.yaml"),
+            server_args=["--disable-log-stats"],
+        ),
+        id="moss_tts_nano",
+    )
+]
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_001(omni_server, openai_client) -> None:
+    """
+    Test text input processing and audio output via /v1/audio/speech.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: audio (48 kHz)
+    Input Setting: stream=False, voice=Junhao
+    Datasets: single request
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": False,
+        "response_format": "wav",
+        "voice": "Junhao",
+    }
+
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_002(omni_server, openai_client) -> None:
+    """
+    Test streaming text-to-audio via /v1/audio/speech.
+    Deploy Setting: default yaml
+    Input Modal: text
+    Output Modal: audio (48 kHz, PCM stream)
+    Input Setting: stream=True, voice=Ava
+    Datasets: single request
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt(),
+        "stream": True,
+        "response_format": "pcm",
+        "voice": "Ava",
+    }
+
+    openai_client.send_audio_speech_request(request_config)
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
+def test_text_to_audio_003(omni_server, openai_client) -> None:
+    """
+    Test Chinese text-to-audio via /v1/audio/speech.
+    Deploy Setting: default yaml
+    Input Modal: text (Chinese)
+    Output Modal: audio (48 kHz)
+    Input Setting: stream=False, voice=Junhao
+    Datasets: single request
+    """
+    request_config = {
+        "model": omni_server.model,
+        "input": get_prompt("chinese"),
+        "stream": False,
+        "response_format": "wav",
+        "voice": "Junhao",
+    }
+
+    openai_client.send_audio_speech_request(request_config)

tests/e2e/online_serving/test_qwen2_5_omni_expansion.py

@@ -56,11 +56,7 @@ def get_max_batch_size(size_type="few"):
     return batch_sizes.get(size_type, 5)
 
 
-@pytest.mark.skip(
-    reason="qwen2.5-omni model cannot recognize the synthetic video and audio data, need to update the test data to enable this test case"
-)
 @pytest.mark.advanced_model
-@pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)

tests/e2e/online_serving/test_qwen3_omni_expansion.py

@@ -469,7 +469,6 @@ def test_speaker_001(omni_server, openai_client) -> None:
 
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
-@pytest.mark.skip(reason="Known issue: occasional inaccuracy in voice recognition.")
 def test_speaker_002(omni_server, openai_client) -> None:
     """
     Input Modal: text only (one-word answer constraint).

tests/e2e/online_serving/test_voxtral_tts.py

@@ -82,6 +82,30 @@ class TestVoxtralTTSFixedVoice:
             f"Audio content too small ({len(response.content)} bytes), expected at least {MIN_AUDIO_BYTES} bytes"
         )
 
+    @pytest.mark.core_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_speech_english_streaming(self, omni_server) -> None:
+        """Test basic streaming English TTS generation."""
+        url = f"http://{omni_server.host}:{omni_server.port}/v1/audio/speech"
+        payload = {
+            "input": "Hello, how are you?",
+            "voice": "casual_female",
+            "language": "English",
+            "stream": True,
+            "response_format": "pcm",
+        }
+
+        with httpx.Client(timeout=120.0) as client:
+            with client.stream("POST", url, json=payload) as response:
+                assert response.status_code == 200
+                assert response.headers.get("content-type") == "audio/pcm"
+                total = sum(len(c) for c in response.iter_bytes())
+
+        assert total > MIN_AUDIO_BYTES, (
+            f"Streamed audio too small ({total} bytes), expected at least {MIN_AUDIO_BYTES} bytes"
+        )
+
     @pytest.mark.advanced_model
     @pytest.mark.omni
     @hardware_test(res={"cuda": "H100"}, num_cards=1)

tests/engine/test_output_processor.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard."""
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from vllm.outputs import PoolingRequestOutput
+from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine import FinishReason
+
+from vllm_omni.engine.output_modality import OutputModalityNames
+from vllm_omni.engine.output_processor import OmniRequestState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Audio is explicitly listed as a drainable modality
+AUDIO = OutputModalityNames.AUDIO
+
+# Latent is explicitly not drainable, but the choice doesn't matter here as
+# long as isn't listed as drainable. I.e., could also be arbitrary keys for
+# the purposes of these tests
+LATENT = OutputModalityNames.LATENT
+
+# NOTE: detokenizer and logprobs aren't really used here, but we mock them since
+# some of the utils called in vLLM superclass assert require them to be None.
+_DETOK = MagicMock(
+    output_token_ids=[0],
+    get_next_output_text=MagicMock(return_value=""),
+    num_output_tokens=MagicMock(return_value=1),
+)
+_LOGPROBS = MagicMock(logprobs=None, cumulative_logprob=None, prompt_logprobs=None)
+
+_DEFAULT_STATE_KWARGS = dict(
+    request_id="r",
+    external_req_id="r",
+    parent_req=None,
+    request_index=0,
+    lora_request=None,
+    prompt=None,
+    prompt_token_ids=[0],
+    prompt_embeds=None,
+    logprobs_processor=_LOGPROBS,
+    detokenizer=_DETOK,
+    max_tokens_param=None,
+    arrival_time=0.0,
+    queue=None,
+    log_stats=False,
+    stream_interval=1,
+)
+
+
+def _make_state(output_kind: RequestOutputKind):
+    return OmniRequestState(**_DEFAULT_STATE_KWARGS, output_kind=output_kind)
+
+
+def test_init_empty_dict():
+    """Ensure mm_accumulated is initially empty."""
+    assert _make_state(RequestOutputKind.CUMULATIVE).mm_accumulated == {}
+    assert _make_state(RequestOutputKind.DELTA).mm_accumulated == {}
+
+
+def test_delta_drains_output_modality_per_step():
+    """DELTA drains the mm_type key (output modality) but preserves hidden-state keys."""
+    s = _make_state(RequestOutputKind.DELTA)
+    audio1, audio2, hs1, hs2 = [torch.ones(num_elem) for num_elem in range(1, 5)]
+
+    # Add audio and hidden state tensors
+    s.add_multimodal_tensor(audio1, mm_type=AUDIO)  # should be drained
+    s.add_multimodal_tensor(hs1, mm_type=LATENT)  # shouldn't be drained
+
+    out1 = s._new_completion_output([1], None, None)
+    out1_audio = out1.multimodal_output[AUDIO]
+    out1_hidden = out1.multimodal_output[LATENT]
+    assert isinstance(out1_audio, torch.Tensor)
+    assert torch.equal(out1.multimodal_output[AUDIO], audio1)
+    assert isinstance(out1_hidden, torch.Tensor)
+    assert torch.equal(out1.multimodal_output[LATENT], hs1)
+
+    # After emission, hidden states should remain, but audio is drained
+    assert set(s.mm_accumulated.keys()) == {LATENT}
+
+    s.add_multimodal_tensor(audio2, AUDIO)
+    s.add_multimodal_tensor(hs2, mm_type=LATENT)
+    out2 = s._new_completion_output([2], None, None)
+    out2_audio = out2.multimodal_output[AUDIO]
+    out2_hidden = out2.multimodal_output[LATENT]
+    assert isinstance(out2_audio, torch.Tensor)
+    assert torch.equal(out2_audio, audio2)
+    # Since hidden isn't drained, it's grown to a list
+    assert isinstance(out2_hidden, list) and len(out2_hidden) == 2
+    assert torch.equal(hs1, out2_hidden[0])
+    assert torch.equal(hs2, out2_hidden[1])
+
+
+def test_cumulative_emits_consolidated_audio_each_step():
+    """Ensure cumulative accumulates and consolidates modality keys every step."""
+    s = _make_state(RequestOutputKind.CUMULATIVE)
+    # NOTE: audio is usually emitted as (1, size) chunks; we need to be sure
+    # to not change the tensor dimension when we consolidate
+    audio1 = torch.ones(1, 500)
+    s.add_multimodal_tensor(audio1, mm_type=AUDIO)
+    req_out = s.make_request_output([1], None, None, None)
+    assert req_out is not None
+    cons_audio = req_out.outputs[0].multimodal_output[AUDIO]
+    # Single chunk keeps original shape [1, 500]
+    assert isinstance(cons_audio, torch.Tensor) and cons_audio.shape == audio1.shape
+
+    audio2 = torch.ones(1, 300)
+    s.add_multimodal_tensor(audio2, mm_type=AUDIO)
+    req_out = s.make_request_output([2], None, None, None)
+    assert req_out is not None
+    cons_audio = req_out.outputs[0].multimodal_output[AUDIO]
+    # After consolidation, audio chunks are concatenated on last axis,
+    # preserving the [1, N] channel dimension
+    total_audio_len = audio1.shape[-1] + audio2.shape[-1]
+    assert isinstance(audio2, torch.Tensor) and cons_audio.shape == (1, total_audio_len)
+
+    assert "audio" in s.mm_accumulated
+
+
+def test_finish_consolidates_hidden_states():
+    """Ensure consolidation merges hidden-state tensor lists on finish."""
+    s = _make_state(RequestOutputKind.CUMULATIVE)
+    s.add_multimodal_tensor(torch.ones(5, 4), mm_type=LATENT)
+    s.add_multimodal_tensor(torch.ones(3, 4), mm_type=LATENT)
+
+    result = s.make_request_output([1], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+
+    hs = result.outputs[0].multimodal_output[LATENT]
+    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 8
+
+
+def test_finish_consolidation_for_hs_delta():
+    """Ensure finish doesn't drop the accumulated hidden states."""
+    s = _make_state(RequestOutputKind.DELTA)
+    # hidden state accumulation (nothing drained)
+    s.add_multimodal_tensor({"foo": torch.ones(5, 4)}, mm_type=LATENT)
+    result = s.make_request_output([0], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+    hs = result.outputs[0].multimodal_output["foo"]
+    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 5
+
+    # Since we don't drain the hidden states, if we add 3 elements, we should get 8
+    s.add_multimodal_tensor({"foo": torch.ones(3, 4)}, mm_type=LATENT)
+    result = s.make_request_output([0], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+    hs = result.outputs[0].multimodal_output["foo"]
+    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 8
+    assert "foo" in s.mm_accumulated
+
+
+def test_finish_consolidation_drains_mm_delta():
+    """Ensure making the request output drains modality deltas (e.g., audio)."""
+    s = _make_state(RequestOutputKind.DELTA)
+    # multimodal data accumulation (drained)
+    s.add_multimodal_tensor({AUDIO: torch.ones(5, 4)}, mm_type=AUDIO)
+    result = s.make_request_output([0], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+    hs = result.outputs[0].multimodal_output[AUDIO]
+    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 5
+
+    # Since we did drain the hidden states, we no longer get the 5 back
+    s.add_multimodal_tensor({AUDIO: torch.ones(3, 4)}, mm_type=AUDIO)
+    result = s.make_request_output([0], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+    hs = result.outputs[0].multimodal_output[AUDIO]
+    assert isinstance(hs, torch.Tensor) and hs.shape[0] == 3
+    assert AUDIO not in s.mm_accumulated  # drained
+
+
+@pytest.mark.parametrize("mm_type", [AUDIO, "hidden"])
+def test_final_only_consolidates_drainable_keys(mm_type):
+    """FINAL_ONLY never drains per-step, so modality keys and hidden state
+    keys both accumulate and are consolidated on finish."""
+    s = _make_state(RequestOutputKind.FINAL_ONLY)
+
+    # NOTE: Currently there is brittlness in the tensor stacking, so we just
+    # test a 1D tensor here. The intention is just to ensure audio /hidden
+    # behave the same.
+    s.add_multimodal_tensor(torch.ones(500), mm_type=mm_type)
+    # Non-finish step returns None without calling _new_completion_output
+    assert s.make_request_output([1], None, None, None) is None
+    assert mm_type in s.mm_accumulated
+
+    s.add_multimodal_tensor(torch.ones(300), mm_type=mm_type)
+    result = s.make_request_output([2], None, FinishReason.STOP, None)
+    assert result is not None and not isinstance(result, PoolingRequestOutput)
+
+    audio = result.outputs[0].multimodal_output[mm_type]
+    assert isinstance(audio, torch.Tensor)
+    assert audio.shape == (800,)
+
+
+def test_cumulative_token_ids_always_set():
+    """cumulative_token_ids is set for all output kinds."""
+    for kind in (RequestOutputKind.DELTA, RequestOutputKind.CUMULATIVE, RequestOutputKind.FINAL_ONLY):
+        s = _make_state(kind)
+        out = s._new_completion_output([42], None, None)
+        assert hasattr(out, "cumulative_token_ids")
+        # The mock detokenizer has output_token_ids=[0]
+        assert list(out.cumulative_token_ids) == [0]
+
+
+def test_cumulative_token_ids_is_a_copy():
+    """cumulative_token_ids must be a snapshot, not a live reference."""
+    s = _make_state(RequestOutputKind.DELTA)
+    out = s._new_completion_output([42], None, None)
+    assert out.cumulative_token_ids is not _DETOK.output_token_ids

tests/entrypoints/openai_api/test_image_server.py

@@ -28,7 +28,6 @@ from vllm_omni.entrypoints.openai.image_api_utils import (
 )
 from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.model_executor.stage_input_processors.glm_image import compute_max_tokens
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -216,7 +215,11 @@ def async_omni_test_client():
             ]
             default_sampling_params_list = [
                 SamplingParams(temperature=0.1),
-                OmniDiffusionSamplingParams(),
+                OmniDiffusionSamplingParams(
+                    num_inference_steps=4,
+                    guidance_scale=7.5,
+                    generator_device="cpu",
+                ),
             ]
             self.engine = SimpleNamespace(
                 stage_configs=stage_configs,
@@ -576,7 +579,11 @@ def test_multistage_images_async_omni_construction(async_omni_test_client):
 
 
 def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
-    """GLM-Image multistage requests must compute stage-0 max_tokens from size."""
+    """GLM-Image multistage: stage-0 gets target_h/w from requested size.
+
+    max_tokens comes from the deploy YAML default (upper-bound ceiling),
+    NOT computed dynamically from height/width.
+    """
 
     class FakeAsyncOmniClass(AsyncOmni):
         def __init__(self):
@@ -584,8 +591,9 @@ def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
                 SimpleNamespace(stage_type="llm", is_comprehension=True, model_arch="GlmImageForConditionalGeneration"),
                 SimpleNamespace(stage_type="diffusion", is_comprehension=False, model_arch="GlmImagePipeline"),
             ]
+            # YAML default max_tokens for GLM-Image AR stage (upper bound for 2048x2048 t2i)
             default_sampling_params_list = [
-                SamplingParams(temperature=0.1, seed=42),
+                SamplingParams(temperature=0.1, seed=42, max_tokens=4353),
                 OmniDiffusionSamplingParams(height=1024, width=1024),
             ]
             self.engine = SimpleNamespace(
@@ -646,7 +654,8 @@ def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
     captured = engine.captured_sampling_params_list
     assert captured is not None
     assert len(captured) == 2
-    assert captured[0].max_tokens == compute_max_tokens(1024, 1024, is_i2i=False)
+    # max_tokens comes from YAML default, not computed dynamically
+    assert captured[0].max_tokens == 4353
     assert captured[0].extra_args["target_h"] == 1024
     assert captured[0].extra_args["target_w"] == 1024
     assert captured[1].height == 1024
@@ -1161,9 +1170,10 @@ def test_image_edit_ignores_mock_like_multimodal_limit(async_omni_test_client):
     assert response.status_code == 200
     captured_prompt = engine.captured_prompt
     assert captured_prompt is not None
-    processed_images = captured_prompt["multi_modal_data"]["image"]
-    assert len(processed_images) == 1
-    assert processed_images[0].size == (16, 16)
+    # Multi-stage path uses "img2img" key for single reference image
+    processed_images = captured_prompt["multi_modal_data"]["img2img"]
+    assert isinstance(processed_images, Image.Image)
+    assert processed_images.size == (16, 16)
 
 
 def test_image_edit_parameter_pass(async_omni_test_client):

tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py

@@ -67,6 +67,10 @@ def test_build_multistage_generation_inputs_applies_stage_specific_overrides(ser
     assert len(sampling_params_list) == 3
     assert sampling_params_list[0].temperature == 0.2
     assert sampling_params_list[0].seed == 0
+    assert sampling_params_list[0].extra_args == {"target_h": 768, "target_w": 1024}
+    assert sampling_params_list[1] is not gen_params
+    assert sampling_params_list[2] is not gen_params
+    assert sampling_params_list[1] is not sampling_params_list[2]
     assert sampling_params_list[1].height == 768
     assert sampling_params_list[1].width == 1024
     assert sampling_params_list[1].seed == 0
@@ -75,8 +79,15 @@ def test_build_multistage_generation_inputs_applies_stage_specific_overrides(ser
     assert sampling_params_list[1].num_outputs_per_prompt == 2
     assert sampling_params_list[1].true_cfg_scale == 5.0
     assert sampling_params_list[1].lora_request.name == "adapter-a"
+    assert sampling_params_list[1].lora_scale == 0.6
     assert sampling_params_list[2].height == 768
     assert sampling_params_list[2].width == 1024
+    assert sampling_params_list[2].seed == 0
     assert sampling_params_list[2].num_inference_steps == 28
+    assert sampling_params_list[2].lora_request.name == "adapter-a"
+    assert sampling_params_list[2].lora_scale == 0.6
+    assert gen_params.lora_request is None
     assert engine.default_sampling_params_list[1].height is None
+    assert engine.default_sampling_params_list[1].lora_request is None
     assert engine.default_sampling_params_list[2].resolution == 640
+    assert engine.default_sampling_params_list[2].lora_request is None

tests/entrypoints/openai_api/test_serving_chat_sampling_params.py

@@ -38,7 +38,7 @@ def default_comprehension_params():
         temperature=0.4,
         top_p=0.9,
         top_k=1,
-        max_tokens=2048,
+        max_tokens=4353,
         seed=42,
         repetition_penalty=1.05,
     )
@@ -146,7 +146,7 @@ def test_preserves_yaml_defaults_when_no_request_params(serving_chat, mock_reque
     assert comprehension_params.temperature == 0.4
     assert comprehension_params.top_p == 0.9
     assert comprehension_params.top_k == 1  # YAML custom param preserved
-    assert comprehension_params.max_tokens == 2048
+    assert comprehension_params.max_tokens == 4353
     assert comprehension_params.seed == 42
     assert comprehension_params.repetition_penalty == 1.05  # YAML custom param preserved
 
@@ -190,7 +190,7 @@ def test_max_tokens_uses_yaml_default_when_not_specified(serving_chat, mock_requ
     """Test that max_tokens falls back to YAML default when not in request."""
     result = serving_chat._build_sampling_params_list_from_request(mock_request)
 
-    assert result[0].max_tokens == 2048
+    assert result[0].max_tokens == 4353
 
 
 def test_request_seed_overrides_yaml_default(serving_chat, mock_request):
@@ -588,12 +588,17 @@ class TestResolveHeightWidth:
 
 
 # =============================================================================
-# Tests for _apply_request_overrides with GLM-Image (max_tokens computation)
+# Tests for _apply_request_overrides with GLM-Image (target_h/w injection)
 # =============================================================================
 
 
 class TestApplyRequestOverridesGLMImage:
-    """Test dynamic max_tokens computation for GLM-Image AR stage."""
+    """Test target_h/w injection for GLM-Image AR stage.
+
+    max_tokens is NOT computed dynamically — it comes from the deploy YAML
+    default (e.g. 4353). _apply_request_overrides only injects target_h/w
+    into extra_args so the model can build M-RoPE position grids.
+    """
 
     @pytest.fixture
     def glm_serving_chat(self, mock_engine_client, mocker: MockerFixture):
@@ -601,7 +606,6 @@ class TestApplyRequestOverridesGLMImage:
 
         instance = object.__new__(OmniOpenAIServingChat)
         instance.engine_client = mock_engine_client
-        # Mock the image extraction to return no reference images (t2i by default)
         instance._extract_diffusion_prompt_and_images_from_messages = mocker.MagicMock(return_value=("a cat", []))
         return instance
 
@@ -623,40 +627,41 @@ class TestApplyRequestOverridesGLMImage:
         req.model_fields_set = set()
         return req
 
-    def test_t2i_computes_max_tokens(self, glm_serving_chat, glm_request, default_comprehension_params):
-        """t2i mode: max_tokens computed from height/width, no reference images."""
+    def test_t2i_injects_target_h_w(self, glm_serving_chat, glm_request, default_comprehension_params):
+        """t2i mode: target_h/w injected into extra_args, max_tokens unchanged."""
         result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request)
-        # t2i 1024x1024 = 256 + 1024 + 1 = 1281
-        assert result.max_tokens == 1281
         assert result.extra_args["target_h"] == 1024
         assert result.extra_args["target_w"] == 1024
+        # max_tokens stays at YAML default (not dynamically computed)
+        assert result.max_tokens == 4353
 
-    def test_i2i_computes_fewer_tokens(
+    def test_i2i_injects_target_h_w(
         self, glm_serving_chat, glm_request, default_comprehension_params, mocker: MockerFixture
     ):
-        """i2i mode: max_tokens should be smaller than t2i for same dimensions."""
-        # Make it detect reference images
+        """i2i mode: target_h/w injected, max_tokens unchanged."""
         glm_serving_chat._extract_diffusion_prompt_and_images_from_messages = mocker.MagicMock(
             return_value=("edit this", ["fake_image"])
         )
-
         result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request)
-        # i2i 1024x1024 = 1024 + 1 = 1025
-        assert result.max_tokens == 1025
+        assert result.extra_args["target_h"] == 1024
+        assert result.extra_args["target_w"] == 1024
+        # max_tokens stays at YAML default regardless of t2i/i2i
+        assert result.max_tokens == 4353
 
-    def test_dynamic_max_tokens_overrides_user_value(self, glm_serving_chat, glm_request, default_comprehension_params):
-        """When height/width are provided, dynamic computation overrides user max_tokens."""
+    def test_user_max_tokens_preserved(self, glm_serving_chat, glm_request, default_comprehension_params):
+        """User-provided max_tokens is respected (not overridden by dynamic computation)."""
         glm_request.max_tokens = 500
         glm_request.model_fields_set = {"max_tokens"}
 
         result = glm_serving_chat._apply_request_overrides(default_comprehension_params, glm_request)
-        # Dynamic computation from height/width always wins when present
-        assert result.max_tokens == 1281
+        assert result.max_tokens == 500
+        assert result.extra_args["target_h"] == 1024
+        assert result.extra_args["target_w"] == 1024
 
     def test_no_height_width_preserves_default(
         self, glm_serving_chat, mocker: MockerFixture, default_comprehension_params
     ):
-        """When no height/width in extra_body, keep YAML default max_tokens."""
+        """When no height/width in extra_body, keep YAML default max_tokens, no target_h/w."""
         req = mocker.MagicMock()
         req.temperature = None
         req.top_p = None
@@ -673,7 +678,9 @@ class TestApplyRequestOverridesGLMImage:
         req.model_fields_set = set()
 
         result = glm_serving_chat._apply_request_overrides(default_comprehension_params, req)
-        assert result.max_tokens == 2048  # YAML default
+        assert result.max_tokens == 4353  # YAML default
+        # No target_h/w injected when dimensions not provided
+        assert not result.extra_args or "target_h" not in (result.extra_args or {})
 
     def test_size_string_parsed_for_glm_image(
         self, glm_serving_chat, mocker: MockerFixture, default_comprehension_params
@@ -695,5 +702,7 @@ class TestApplyRequestOverridesGLMImage:
         req.model_fields_set = set()
 
         result = glm_serving_chat._apply_request_overrides(default_comprehension_params, req)
-        # 512x512 t2i = 256 + 256 + 1 = 513
-        assert result.max_tokens == 513
+        assert result.extra_args["target_h"] == 512
+        assert result.extra_args["target_w"] == 512
+        # max_tokens stays at YAML default (not dynamically computed)
+        assert result.max_tokens == 4353

tests/entrypoints/openai_api/test_stage_params.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+from vllm import SamplingParams
+
+from vllm_omni.entrypoints.openai.stage_params import (
+    build_stage_sampling_params_list,
+    get_default_sampling_params_list,
+    resolve_stage_sampling_params,
+)
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+
+def test_resolve_stage_sampling_params_clones_stage_default():
+    default = SamplingParams(temperature=0.2, seed=11)
+
+    resolved = resolve_stage_sampling_params(
+        SimpleNamespace(stage_type="llm"),
+        0,
+        [default],
+    )
+
+    assert resolved is not default
+    assert resolved.temperature == 0.2
+    assert resolved.seed == 11
+
+    resolved.seed = 99
+    assert default.seed == 11
+
+
+def test_resolve_stage_sampling_params_uses_diffusion_fallback_when_default_missing():
+    request_params = OmniDiffusionSamplingParams(height=768, width=1024, seed=7)
+
+    resolved = resolve_stage_sampling_params(
+        SimpleNamespace(stage_type="diffusion"),
+        1,
+        [],
+        diffusion_params=request_params,
+    )
+
+    assert resolved is not request_params
+    assert resolved.height == 768
+    assert resolved.width == 1024
+    assert resolved.seed == 7
+
+
+def test_build_stage_sampling_params_list_can_replace_diffusion_defaults():
+    request_params = OmniDiffusionSamplingParams(height=512, width=512)
+    diffusion_default = OmniDiffusionSamplingParams(height=1024, width=1024)
+    llm_default = SamplingParams(temperature=0.1)
+    stages = [
+        SimpleNamespace(stage_type="llm"),
+        SimpleNamespace(stage_type="diffusion"),
+        SimpleNamespace(stage_type="diffusion"),
+    ]
+
+    resolved = build_stage_sampling_params_list(
+        stages,
+        [llm_default, diffusion_default],
+        diffusion_params=request_params,
+        replace_diffusion_params=True,
+    )
+
+    assert resolved[0] is not llm_default
+    assert resolved[0].temperature == 0.1
+    assert resolved[1] is not request_params
+    assert resolved[2] is not request_params
+    assert resolved[1] is not resolved[2]
+    assert resolved[1].height == 512
+    assert resolved[2].height == 512
+    assert resolved[1].width == 512
+    assert resolved[2].width == 512
+
+
+def test_get_default_sampling_params_list_reads_engine_defaults_only():
+    assert get_default_sampling_params_list(SimpleNamespace()) == []
+    assert get_default_sampling_params_list(SimpleNamespace(default_sampling_params_list=None)) == []
+
+    defaults = [SamplingParams(max_tokens=8)]
+    resolved = get_default_sampling_params_list(SimpleNamespace(default_sampling_params_list=defaults))
+    assert resolved == defaults
+    assert resolved is not defaults

tests/entrypoints/test_async_omni.py

@@ -0,0 +1,134 @@
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+from vllm.sampling_params import RequestOutputKind, SamplingParams
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+async def _noop(**kw):
+    pass
+
+
+def get_fake_add_request(submitted_request_ids):
+    async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
+        del prompt, sampling_params_list, final_stage_id, kwargs
+        submitted_request_ids.append(request_id)
+
+    return fake_add_request_async
+
+
+def get_fake_abort(aborted_request_batches):
+    async def fake_abort_async(request_ids):
+        aborted_request_batches.append(list(request_ids))
+
+    return fake_abort_async
+
+
+async def fake_process_results(request_id, metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts):
+    del metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts
+    if request_id.startswith("cancel-"):
+        await asyncio.Future()
+        return
+    yield SimpleNamespace(
+        stage_id=0,
+        request_output=SimpleNamespace(outputs=[]),
+        finished=True,
+    )
+
+
+def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) -> AsyncOmni:
+    omni = object.__new__(AsyncOmni)
+    omni._pause_cond = asyncio.Condition()
+    omni._paused = False
+    omni.engine = SimpleNamespace(
+        num_stages=1,
+        add_request_async=fake_add_request,
+        abort_async=fake_abort_request,
+    )
+    omni.log_stats = False
+    omni.request_states = {}
+    omni._final_output_handler = lambda: None
+    omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params
+    omni._compute_final_stage_id = lambda output_modalities: 0
+    omni._process_orchestrator_results = fake_process_results
+    omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)
+    return omni
+
+
+def test_generate_accepts_request_after_repeated_cancellations():
+    async def run_test():
+        submitted_request_ids = []
+        aborted_request_batches = []
+
+        async def collect_outputs(request_id):
+            outputs = []
+            async for output in AsyncOmni.generate(
+                omni,
+                prompt={"prompt": "prompt"},
+                request_id=request_id,
+                sampling_params_list=[SimpleNamespace()],
+                output_modalities=["image"],
+            ):
+                outputs.append(output)
+            return outputs
+
+        omni = get_async_omni_instance(
+            fake_add_request=get_fake_add_request(submitted_request_ids),
+            fake_abort_request=get_fake_abort(aborted_request_batches),
+        )
+
+        assert len(await collect_outputs("baseline")) == 1
+
+        for idx in range(3):
+            task = asyncio.create_task(collect_outputs(f"cancel-{idx}"))
+            await asyncio.sleep(0)
+            task.cancel()
+            with pytest.raises(asyncio.CancelledError):
+                await task
+
+        assert len(await collect_outputs("after-cancel")) == 1
+        assert submitted_request_ids == [
+            "baseline",
+            "cancel-0",
+            "cancel-1",
+            "cancel-2",
+            "after-cancel",
+        ]
+        assert aborted_request_batches == [
+            ["cancel-0"],
+            ["cancel-1"],
+            ["cancel-2"],
+        ]
+
+    asyncio.run(run_test())
+
+
+@pytest.mark.parametrize(
+    "output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY, RequestOutputKind.CUMULATIVE]
+)
+def test_output_kind_is_preserved_with_explicit_sampling_params(output_kind):
+    """Ensure we don't change the output kind in async generate if params are provided directly."""
+
+    captured_params = []
+
+    async def capturing_add_request(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
+        del prompt, final_stage_id, kwargs
+        captured_params.extend(sampling_params_list)
+
+    async def run():
+        omni = get_async_omni_instance(fake_add_request=capturing_add_request)
+        sp = SamplingParams(output_kind=output_kind)
+        async for _ in omni.generate(
+            prompt={"prompt": "test"},
+            request_id="test-req",
+            sampling_params_list=[sp],
+            output_modalities=["text"],
+        ):
+            pass
+
+    asyncio.run(run())
+    assert captured_params[0].output_kind == output_kind

tests/entrypoints/test_async_omni_abort.py

@@ -1,85 +0,0 @@
-import asyncio
-from types import SimpleNamespace
-
-import pytest
-
-from vllm_omni.entrypoints.async_omni import AsyncOmni
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-def test_generate_accepts_request_after_repeated_cancellations():
-    async def run_test():
-        submitted_request_ids = []
-        aborted_request_batches = []
-
-        async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs):
-            del prompt, sampling_params_list, final_stage_id, kwargs
-            submitted_request_ids.append(request_id)
-
-        async def fake_abort_async(request_ids):
-            aborted_request_batches.append(list(request_ids))
-
-        async def fake_process_results(request_id, metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts):
-            del metrics, final_stage_id_for_e2e, req_start_ts, wall_start_ts
-            if request_id.startswith("cancel-"):
-                await asyncio.Future()
-                return
-            yield SimpleNamespace(
-                stage_id=0,
-                request_output=SimpleNamespace(outputs=[]),
-                finished=True,
-            )
-
-        async def collect_outputs(request_id):
-            outputs = []
-            async for output in AsyncOmni.generate(
-                omni,
-                prompt={"prompt": "prompt"},
-                request_id=request_id,
-                sampling_params_list=[SimpleNamespace()],
-                output_modalities=["image"],
-            ):
-                outputs.append(output)
-            return outputs
-
-        omni = object.__new__(AsyncOmni)
-        omni._pause_cond = asyncio.Condition()
-        omni._paused = False
-        omni.engine = SimpleNamespace(
-            num_stages=1,
-            add_request_async=fake_add_request_async,
-            abort_async=fake_abort_async,
-        )
-        omni.log_stats = False
-        omni.request_states = {}
-        omni._final_output_handler = lambda: None
-        omni.resolve_sampling_params_list = lambda params: params
-        omni._compute_final_stage_id = lambda output_modalities: 0
-        omni._process_orchestrator_results = fake_process_results
-        omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)
-
-        assert len(await collect_outputs("baseline")) == 1
-
-        for idx in range(3):
-            task = asyncio.create_task(collect_outputs(f"cancel-{idx}"))
-            await asyncio.sleep(0)
-            task.cancel()
-            with pytest.raises(asyncio.CancelledError):
-                await task
-
-        assert len(await collect_outputs("after-cancel")) == 1
-        assert submitted_request_ids == [
-            "baseline",
-            "cancel-0",
-            "cancel-1",
-            "cancel-2",
-            "after-cancel",
-        ]
-        assert aborted_request_batches == [
-            ["cancel-0"],
-            ["cancel-1"],
-            ["cancel-2"],
-        ]
-
-    asyncio.run(run_test())

tests/entrypoints/test_utils.py

@@ -7,6 +7,7 @@ from dataclasses import dataclass
 import pytest
 import torch
 from pytest_mock import MockerFixture
+from vllm.sampling_params import RequestOutputKind, SamplingParams
 
 from vllm_omni.config.yaml_util import create_config
 from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -15,6 +16,7 @@ from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.entrypoints.utils import (
     _convert_dataclasses_to_dict,
     _filter_dict_like_object,
+    coerce_param_message_types,
     filter_dataclass_kwargs,
     load_and_resolve_stage_configs,
     load_stage_configs_from_yaml,
@@ -424,3 +426,44 @@ class TestLoadStageConfigsFromYaml:
 
         assert stages[0]["engine_args"]["nested"]["base"] == 1
         assert stages[0]["engine_args"]["nested"]["override"] == 2
+
+
+class TestCumulativeStreamingCoercion:
+    @pytest.mark.parametrize("skip_clone", [True, False])
+    def test_cumulative_default_becomes_delta_if_stream(self, skip_clone):
+        """Ensure cumulative messages are coercible to delta if streaming."""
+        sp = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
+        sp.skip_clone = skip_clone
+        result = coerce_param_message_types([sp], is_streaming=True)[0]
+        assert isinstance(result, SamplingParams)
+        assert result.output_kind == RequestOutputKind.DELTA
+        assert (skip_clone and sp is result) or (not skip_clone and sp is not result)
+
+    @pytest.mark.parametrize("skip_clone", [True, False])
+    def test_cumulative_default_becomes_final_only_if_not_stream(self, skip_clone):
+        """Ensure cumulative messages are coercible to final only if not streaming."""
+        sp = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
+        sp.skip_clone = skip_clone
+        result = coerce_param_message_types([sp], is_streaming=False)[0]
+        assert isinstance(result, SamplingParams)
+        assert result.output_kind == RequestOutputKind.FINAL_ONLY
+        assert (skip_clone and sp is result) or (not skip_clone and sp is not result)
+
+    @pytest.mark.parametrize("is_streaming", [True, False])
+    @pytest.mark.parametrize("output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])
+    def test_non_cumulative_are_coerced(self, output_kind, is_streaming):
+        """Ensure non-cumulative params are coerced to the target type."""
+        sp = SamplingParams(output_kind=output_kind)
+        expected = RequestOutputKind.DELTA if is_streaming else RequestOutputKind.FINAL_ONLY
+        result = coerce_param_message_types([sp], is_streaming=is_streaming)[0]
+        assert isinstance(result, SamplingParams)
+        assert result.output_kind == expected
+
+    def test_coercion_applies_to_all_stages(self):
+        """Ensure all stages are coerced to DELTA for streaming."""
+        sp0 = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
+        sp1 = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
+        result = coerce_param_message_types([sp0, sp1], is_streaming=True)
+        assert all([isinstance(r, SamplingParams) for r in result])
+        assert result[0].output_kind == RequestOutputKind.DELTA
+        assert result[1].output_kind == RequestOutputKind.DELTA

tests/examples/offline_inference/test_text_to_image.py

@@ -27,9 +27,27 @@ def _skip_readme_snippet(language: str, code: str, h2_title: str) -> tuple[bool,
 
 README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH, skipif=_skip_readme_snippet)
 
+# Track: https://github.com/vllm-project/vllm-omni/issues/3123 — CLI example subprocess exit 1 in CI; re-enable when fixed.
+_SKIP_T2I_SNIPPET_IDS = frozenset(
+    {
+        "quick_start_002",
+        "more_cli_examples_001",
+        "more_cli_examples_002",
+        "more_cli_examples_003",
+        "more_cli_examples_007",
+        "more_cli_examples_008",
+    }
+)
+_SKIP_T2I_REASON = (
+    "Skipped: text_to_image README examples failing in CI (subprocess exit 1). "
+    "See https://github.com/vllm-project/vllm-omni/issues/3123"
+)
+
 
 @pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
 def test_text_to_image(snippet: ReadmeSnippet, example_runner: ExampleRunner):
+    if snippet.test_id in _SKIP_T2I_SNIPPET_IDS:
+        pytest.skip(_SKIP_T2I_REASON)
     should_skip, reason = snippet.skip
     if should_skip:
         pytest.skip(reason)

tests/helpers/assertions.py

@@ -318,7 +318,7 @@ def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str:
         top = outputs[0]
         label = str(top.get("label", "")).lower()
         conf = float(top.get("score", 0.0))
-        if conf < 0.5:
+        if conf < 0.6:
             gender = "unknown"
         elif ("female" in label) or ("жен" in label):
             gender = "female"

tests/helpers/fixtures/runtime.py

@@ -10,17 +10,27 @@ from __future__ import annotations
 
 import threading
 from collections.abc import Generator
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import pytest
-import yaml
 
-from tests.helpers.runtime import OmniServer
-from tests.helpers.stage_config import modify_stage_config
+if TYPE_CHECKING:
+    from tests.helpers.runtime import OmniServer
 
 omni_fixture_lock = threading.Lock()
 
 
+@pytest.fixture(scope="function")
+def omni_server_function(
+    request: pytest.FixtureRequest,
+    run_level: str,
+    model_prefix: str,
+) -> Generator[OmniServer, Any, None]:
+    from tests.helpers.runtime import run_omni_server
+
+    yield from run_omni_server(request, run_level, model_prefix, omni_fixture_lock)
+
+
 @pytest.fixture(scope="module")
 def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: str) -> Generator[OmniServer, Any, None]:
     """Start vLLM-Omni through the standard or stage-CLI launcher.
@@ -29,81 +39,9 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st
     The ``use_stage_cli`` flag on ``OmniServerParams`` routes the setup through the
     stage-CLI harness while still reusing the same fixture grouping semantics.
     """
-    with omni_fixture_lock:
-        from tests.helpers.runtime import OmniServer, OmniServerParams, OmniServerStageCli
+    from tests.helpers.runtime import run_omni_server
 
-        params: OmniServerParams = request.param
-        model = model_prefix + params.model
-        port = params.port
-        stage_config_path = params.stage_config_path
-        if run_level in {"advanced_model", "full_model"} and stage_config_path is not None:
-            with open(stage_config_path, encoding="utf-8") as f:
-                cfg = yaml.safe_load(f) or {}
-            # Strip ``load_format: dummy`` (CI overlay default) so advanced_model
-            # tests use real weights. New schema (``stages:``) writes the field
-            # flat at stage level; legacy schema (``stage_args:``) nests it as
-            # ``engine_args.load_format``. Handle both.
-            new_schema_stages = cfg.get("stages")
-            stage_key = "stages" if new_schema_stages is not None else "stage_args"
-            delete_path = "load_format" if new_schema_stages is not None else "engine_args.load_format"
-            stage_entries = cfg.get(stage_key, [])
-            stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage]
-            stage_config_path = modify_stage_config(
-                stage_config_path,
-                deletes={stage_key: {stage_id: [delete_path] for stage_id in stage_ids}},
-            )
-
-        server_args = params.server_args or []
-        if params.use_omni and params.stage_init_timeout is not None:
-            server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)]
-        else:
-            server_args = [*server_args, "--stage-init-timeout", "600"]
-        if params.init_timeout is not None:
-            server_args = [*server_args, "--init-timeout", str(params.init_timeout)]
-        else:
-            server_args = [*server_args, "--init-timeout", "900"]
-        if params.use_stage_cli:
-            if not params.use_omni:
-                raise ValueError("omni_server with use_stage_cli=True requires use_omni=True")
-            if stage_config_path is None:
-                raise ValueError("omni_server with use_stage_cli=True requires a stage_config_path")
-            server_args += ["--stage-configs-path", stage_config_path]
-
-            with OmniServerStageCli(
-                model,
-                stage_config_path,
-                server_args,
-                port=port,
-                env_dict=params.env_dict,
-            ) as server:
-                print("OmniServer started successfully")
-                yield server
-                print("OmniServer stopping...")
-        else:
-            if stage_config_path is not None:
-                server_args += ["--stage-configs-path", stage_config_path]
-
-            with (
-                OmniServer(
-                    model,
-                    server_args,
-                    port=port,
-                    env_dict=params.env_dict,
-                    use_omni=params.use_omni,
-                )
-                if port
-                else OmniServer(
-                    model,
-                    server_args,
-                    env_dict=params.env_dict,
-                    use_omni=params.use_omni,
-                )
-            ) as server:
-                print("OmniServer started successfully")
-                yield server
-                print("OmniServer stopping...")
-
-        print("OmniServer stopped")
+    yield from run_omni_server(request, run_level, model_prefix, omni_fixture_lock)
 
 
 @pytest.fixture
@@ -115,6 +53,15 @@ def openai_client(request: pytest.FixtureRequest, run_level: str):
     return OpenAIClientHandler(host=server.host, port=server.port, api_key="EMPTY", run_level=run_level)
 
 
+@pytest.fixture
+def openai_client_function(request: pytest.FixtureRequest, run_level: str):
+    """Resolve ``omni_server_function`` lazily for function-scoped reliability tests."""
+    from tests.helpers.runtime import OpenAIClientHandler
+
+    server = request.getfixturevalue("omni_server_function")
+    return OpenAIClientHandler(host=server.host, port=server.port, api_key="EMPTY", run_level=run_level)
+
+
 @pytest.fixture(scope="module")
 def omni_runner(request: pytest.FixtureRequest, model_prefix: str):
     from tests.helpers.runtime import OmniRunner

tests/helpers/runtime.py

@@ -9,7 +9,9 @@ import socket
 import subprocess
 import sys
 import tempfile
+import threading
 import time
+from collections.abc import Generator
 from dataclasses import dataclass
 from io import BytesIO
 from pathlib import Path
@@ -120,6 +122,89 @@ class OmniServerParams(NamedTuple):
     stage_init_timeout: int | None = None  # None: fixture supplies default (600 s)
 
 
+def run_omni_server(
+    request: Any,
+    run_level: str,
+    model_prefix: str,
+    omni_fixture_lock: threading.Lock,
+) -> Generator["OmniServer", Any, None]:
+    from tests.helpers.stage_config import modify_stage_config
+
+    with omni_fixture_lock:
+        params: OmniServerParams = request.param
+        model = model_prefix + params.model
+        port = params.port
+        stage_config_path = params.stage_config_path
+        if run_level in {"advanced_model", "full_model"} and stage_config_path is not None:
+            with open(stage_config_path, encoding="utf-8") as f:
+                cfg = yaml.safe_load(f) or {}
+            # Strip ``load_format: dummy`` (CI overlay default) so advanced_model
+            # tests use real weights. New schema (``stages:``) writes the field
+            # flat at stage level; legacy schema (``stage_args:``) nests it as
+            # ``engine_args.load_format``. Handle both.
+            new_schema_stages = cfg.get("stages")
+            stage_key = "stages" if new_schema_stages is not None else "stage_args"
+            delete_path = "load_format" if new_schema_stages is not None else "engine_args.load_format"
+            stage_entries = cfg.get(stage_key, [])
+            stage_ids = [stage["stage_id"] for stage in stage_entries if "stage_id" in stage]
+            stage_config_path = modify_stage_config(
+                stage_config_path,
+                deletes={stage_key: {stage_id: [delete_path] for stage_id in stage_ids}},
+            )
+
+        server_args = params.server_args or []
+        if params.use_omni and params.stage_init_timeout is not None:
+            server_args = [*server_args, "--stage-init-timeout", str(params.stage_init_timeout)]
+        else:
+            server_args = [*server_args, "--stage-init-timeout", "600"]
+        if params.init_timeout is not None:
+            server_args = [*server_args, "--init-timeout", str(params.init_timeout)]
+        else:
+            server_args = [*server_args, "--init-timeout", "900"]
+        if params.use_stage_cli:
+            if not params.use_omni:
+                raise ValueError("omni_server with use_stage_cli=True requires use_omni=True")
+            if stage_config_path is None:
+                raise ValueError("omni_server with use_stage_cli=True requires a stage_config_path")
+            server_args += ["--stage-configs-path", stage_config_path]
+
+            with OmniServerStageCli(
+                model,
+                stage_config_path,
+                server_args,
+                port=port,
+                env_dict=params.env_dict,
+            ) as server:
+                print("OmniServer started successfully")
+                yield server
+                print("OmniServer stopping...")
+        else:
+            if stage_config_path is not None:
+                server_args += ["--stage-configs-path", stage_config_path]
+
+            with (
+                OmniServer(
+                    model,
+                    server_args,
+                    port=port,
+                    env_dict=params.env_dict,
+                    use_omni=params.use_omni,
+                )
+                if port
+                else OmniServer(
+                    model,
+                    server_args,
+                    env_dict=params.env_dict,
+                    use_omni=params.use_omni,
+                )
+            ) as server:
+                print("OmniServer started successfully")
+                yield server
+                print("OmniServer stopping...")
+
+        print("OmniServer stopped")
+
+
 class OmniServer:
     """Omniserver for vLLM-Omni tests."""
 

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -44,7 +44,7 @@ def test_forward_trims_context_on_exact_frame_boundaries():
 
     out = model.forward(
         input_ids=torch.arange(12, dtype=torch.long),
-        runtime_additional_information=[{"left_context_size": 2}],
+        runtime_additional_information=[{"meta": {"left_context_size": 2}}],
     )
 
     audio = out.multimodal_outputs["model_outputs"][0]
@@ -57,7 +57,7 @@ def test_forward_trims_trailing_padding_without_context():
 
     out = model.forward(
         input_ids=torch.arange(12, dtype=torch.long),
-        runtime_additional_information=[{"left_context_size": 0}],
+        runtime_additional_information=[{"meta": {"left_context_size": 0}}],
     )
 
     audio = out.multimodal_outputs["model_outputs"][0]

tests/model_executor/models/test_fish_speech_regressions.py

@@ -45,7 +45,7 @@ def test_dac_decoder_mixed_batch_empty_request_does_not_misalign_indices():
 
     out = decoder.forward(
         input_ids=torch.arange(20, dtype=torch.long),
-        runtime_additional_information=[{}, {"left_context_size": 1}],
+        runtime_additional_information=[{}, {"meta": {"left_context_size": 1}}],
     )
 
     audios = out.multimodal_outputs["model_outputs"]

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -13,7 +13,7 @@ def _source_output(request_id: str, prompt_ids: list[int], out_ids: list[int], m
     return SimpleNamespace(
         request_id=request_id,
         prompt_token_ids=prompt_ids,
-        outputs=[SimpleNamespace(token_ids=out_ids, multimodal_output=mm)],
+        outputs=[SimpleNamespace(token_ids=out_ids, cumulative_token_ids=out_ids, multimodal_output=mm)],
     )
 
 
@@ -58,8 +58,8 @@ def test_text2flow_supports_batched_source_outputs():
     assert len(outputs) == 2
     assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
     assert outputs[1]["prompt_token_ids"] == [4, 5]
-    assert outputs[0]["additional_information"]["prefix_ids"] == [10, 11]
-    assert outputs[1]["additional_information"]["prefix_ids"] == [20, 21]
+    assert outputs[0]["additional_information"]["ids"]["prompt"] == [10, 11]
+    assert outputs[1]["additional_information"]["ids"]["prompt"] == [20, 21]
 
 
 def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
@@ -83,8 +83,8 @@ def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
     )
 
     assert payload is not None
-    assert payload["finished"].item() is True
-    assert payload["code_predictor_codes"] == [1, 2, 3]
+    assert payload["meta"]["finished"].item() is True
+    assert payload["codes"]["audio"] == [1, 2, 3]
     assert payload["token_offset"] == 0
     assert payload["left_context_size"] == 0
     assert payload["req_id"] == ["rid-0"]
@@ -111,8 +111,8 @@ def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes
     )
 
     assert payload is not None
-    assert payload["code_predictor_codes"] == []
-    assert payload["finished"].item() is True
+    assert payload["codes"]["audio"] == []
+    assert payload["meta"]["finished"].item() is True
 
 
 def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
@@ -138,7 +138,7 @@ def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
     )
 
     assert payload1 is not None
-    assert payload1["code_predictor_codes"] == [1, 2]
+    assert payload1["codes"]["audio"] == [1, 2]
     assert payload1["token_offset"] == 0
     assert payload2 is None
 
@@ -168,9 +168,9 @@ def test_talker2code2wav_async_chunk_waits_for_prelookahead_and_emits_cumulative
 
     assert payload_pending is None
     assert payload_ready is not None
-    assert payload_ready["code_predictor_codes"] == [1, 2, 3]
+    assert payload_ready["codes"]["audio"] == [1, 2, 3]
     assert payload_ready["token_offset"] == 0
-    assert payload_ready["finished"].item() is False
+    assert payload_ready["meta"]["finished"].item() is False
 
 
 def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
@@ -197,12 +197,12 @@ def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
     )
 
     assert payload_stream is not None
-    assert payload_stream["finished"].item() is False
-    assert payload_stream["code_predictor_codes"] == [3, 4, 5]
+    assert payload_stream["meta"]["finished"].item() is False
+    assert payload_stream["codes"]["audio"] == [3, 4, 5]
     assert payload_stream["token_offset"] == 0
     assert payload_final is not None
-    assert payload_final["finished"].item() is True
-    assert payload_final["code_predictor_codes"] == [3, 4, 5, 6]
+    assert payload_final["meta"]["finished"].item() is True
+    assert payload_final["codes"]["audio"] == [3, 4, 5, 6]
     assert payload_final["token_offset"] == 2
 
 
@@ -233,7 +233,7 @@ def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
 
     assert payload_pending is None
     assert payload_ready is not None
-    assert payload_ready["code_predictor_codes"] == [8, 9, 10, 11]
+    assert payload_ready["codes"]["audio"] == [8, 9, 10, 11]
     assert payload_ready["token_offset"] == 0
 
 
@@ -260,8 +260,8 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(
     )
 
     assert payload_stream is not None
-    assert payload_stream["finished"].item() is False
-    assert payload_stream["code_predictor_codes"] == [3, 4]
+    assert payload_stream["meta"]["finished"].item() is False
+    assert payload_stream["codes"]["audio"] == [3, 4]
     assert payload_final is not None
-    assert payload_final["finished"].item() is True
-    assert payload_final["code_predictor_codes"] == []
+    assert payload_final["meta"]["finished"].item() is True
+    assert payload_final["codes"]["audio"] == []

tests/model_executor/stage_input_processors/test_glm_image.py

@@ -27,7 +27,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 def _source_output(token_ids: list[int], mm_output: dict | None = None):
     """Create a minimal AR output mock."""
     return SimpleNamespace(
-        outputs=[SimpleNamespace(token_ids=token_ids)],
+        outputs=[SimpleNamespace(token_ids=token_ids, cumulative_token_ids=token_ids)],
         multimodal_output=mm_output,
     )
 
@@ -299,7 +299,7 @@ class TestAr2Diffusion:
     def test_i2i_with_mm_output(self):
         """Test image-to-image with prior_token_image_ids from AR model."""
         token_ids = list(range(1024)) + [16385]
-        mm_output = {"prior_token_image_ids": torch.tensor([1, 2, 3])}
+        mm_output = {"ids": {"prior_image": torch.tensor([1, 2, 3])}}
         stage_list = [_stage_with_outputs([_source_output(token_ids, mm_output)])]
 
         from PIL import Image

tests/model_executor/stage_input_processors/test_mimo_audio_flush_remaining_codes.py

@@ -13,23 +13,23 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 def _sentinel():
-    return {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
+    return {"codes": {"audio": []}, "meta": {"finished": torch.tensor(True, dtype=torch.bool)}}
 
 
 def test_flush_remaining_codes_when_no_codes_accumulated_missing_request_id():
     """No entry for request_id: treat as empty, return finished sentinel with empty codes."""
     tm = SimpleNamespace(code_prompt_token_ids={})
     out = _flush_remaining_codes(tm, "missing", chunk_size=3, left_context_size=3)
-    assert out["code_predictor_codes"] == _sentinel()["code_predictor_codes"]
-    assert out["finished"].equal(_sentinel()["finished"])
+    assert out["codes"]["audio"] == _sentinel()["codes"]["audio"]
+    assert out["meta"]["finished"].item() is True
 
 
 def test_flush_remaining_codes_when_no_codes_accumulated_empty_list():
     """Explicit empty accumulation list returns the same sentinel."""
     tm = SimpleNamespace(code_prompt_token_ids={"r": []})
     out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
-    assert out["code_predictor_codes"] == []
-    assert out["finished"].item() is True
+    assert out["codes"]["audio"] == []
+    assert out["meta"]["finished"].item() is True
 
 
 def test_flush_remaining_codes_partial_chunk_remaining():
@@ -41,8 +41,8 @@ def test_flush_remaining_codes_partial_chunk_remaining():
         code_prompt_token_ids={"r": [[1], [2], [3], [4], [5], [6], [7]]},
     )
     out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
-    assert out["finished"].item() is True
-    assert out["code_predictor_codes"] == [4, 5, 6, 7]
+    assert out["meta"]["finished"].item() is True
+    assert out["codes"]["audio"] == [4, 5, 6, 7]
 
 
 def test_flush_remaining_codes_when_length_is_exact_multiple_of_chunk_size():
@@ -52,7 +52,7 @@ def test_flush_remaining_codes_when_length_is_exact_multiple_of_chunk_size():
     )
     out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
     # context_length = chunk_size = 3, end_index = min(6, 6) -> all 6
-    assert out["code_predictor_codes"] == [1, 2, 3, 4, 5, 6]
+    assert out["codes"]["audio"] == [1, 2, 3, 4, 5, 6]
 
 
 @pytest.mark.parametrize(
@@ -74,5 +74,5 @@ def test_flush_remaining_codes_context_window_end_index(
     tm = SimpleNamespace(code_prompt_token_ids={"r": accumulated})
     out = _flush_remaining_codes(tm, "r", chunk_size=chunk_size, left_context_size=left_context)
     expected_flat = list(range(length - expected_end_index, length))
-    assert out["code_predictor_codes"] == expected_flat
-    assert out["finished"].item() is True
+    assert out["codes"]["audio"] == expected_flat
+    assert out["meta"]["finished"].item() is True

tests/model_executor/stage_input_processors/test_mimo_audio_llm2code2wav.py

@@ -19,7 +19,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 def _make_stage_list(codec_codes: torch.Tensor, request_id: str = "req-0"):
     """Build a minimal stage_list[0] with a single talker output carrying codec_codes."""
-    output = SimpleNamespace(multimodal_output={"code_predictor_codes": codec_codes})
+    output = SimpleNamespace(multimodal_output={"codes": {"audio": codec_codes}})
     talker_output = SimpleNamespace(outputs=[output], request_id=request_id)
     stage0 = SimpleNamespace(engine_outputs=[talker_output])
     return [stage0]

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -55,7 +55,7 @@ def _call(tm, rid, *, n_frames, finished=False, req_ic=None):
     tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)]
     return talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.zeros((0,))},
+        pooling_output={"codes": {"audio": torch.zeros((0,))}},
         request=_req(rid, finished=finished, initial_codec_chunk_frames=req_ic),
         is_finished=finished,
     )
@@ -65,7 +65,7 @@ def test_empty_returns_none():
     tm = _tm()
     p = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.zeros((0,))},
+        pooling_output={"codes": {"audio": torch.zeros((0,))}},
         request=_req("r", finished=False),
     )
     assert p is None
@@ -79,7 +79,8 @@ def test_eof_marker_when_finished_empty():
         request=_req("r", finished=True),
         is_finished=True,
     )
-    assert p == {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
+    assert p["codes"] == {"audio": []}
+    assert p["meta"]["finished"].item() is True
 
 
 def test_flush_on_finish():
@@ -92,8 +93,8 @@ def test_flush_on_finish():
         is_finished=True,
     )
     assert p is not None
-    assert p["finished"] is True
-    assert len(p["code_predictor_codes"]) == _Q * 24
+    assert p["meta"]["finished"].item() is True
+    assert len(p["codes"]["audio"]) == _Q * 24
 
 
 _CASES = [
@@ -158,8 +159,8 @@ def test_streaming_phases(config, n_frames, finished, expected):
     else:
         exp_ctx, exp_window = expected
         assert payload is not None
-        assert payload["left_context_size"] == exp_ctx
-        assert len(payload["code_predictor_codes"]) == _Q * exp_window
+        assert payload["meta"]["left_context_size"] == exp_ctx
+        assert len(payload["codes"]["audio"]) == _Q * exp_window
 
 
 def test_dynamic_ic_adapts_to_load():
@@ -169,14 +170,14 @@ def test_dynamic_ic_adapts_to_load():
     # Low load (1/8) -> IC=2 -> emit at 2
     p1 = _call(tm, "r", n_frames=2)
     assert p1 is not None
-    assert len(p1["code_predictor_codes"]) == _Q * 2
+    assert len(p1["codes"]["audio"]) == _Q * 2
 
     # High load: add 4 others -> active=5/8 -> IC=8 -> emit at 8
     for i in range(4):
         tm.code_prompt_token_ids[f"other-{i}"] = [[0]]
     p2 = _call(tm, "r", n_frames=8)
     assert p2 is not None
-    assert len(p2["code_predictor_codes"]) == _Q * 8
+    assert len(p2["codes"]["audio"]) == _Q * 8
 
     # Requests past initial phase still count in load factor
     tm2 = _tm(max_num_seqs=4)
@@ -185,7 +186,7 @@ def test_dynamic_ic_adapts_to_load():
     # active=4/4=1.0 -> IC=16
     p3 = _call(tm2, "new", n_frames=16)
     assert p3 is not None
-    assert len(p3["code_predictor_codes"]) == _Q * 16
+    assert len(p3["codes"]["audio"]) == _Q * 16
 
 
 def test_ic_load_change_mid_request():
@@ -206,6 +207,7 @@ def test_ic_load_change_mid_request():
     assert _call(tm, "r", n_frames=27) is None
     p3 = _call(tm, "r", n_frames=49)
     assert p3 is not None
+    assert p3["meta"]["left_context_size"] == 24
 
     # A *new* request under high load gets IC=16 (not IC=2).
     # Frame 2 would emit under IC=2 but must hold under IC=16.
@@ -253,14 +255,14 @@ def test_first_streaming_chunk_prepends_ref_code_context():
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.zeros((0,)), "ref_code": ref_code},
+        pooling_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )
 
     assert payload is not None
-    assert payload["left_context_size"] == 2
-    assert len(payload["code_predictor_codes"]) == _Q * 12
+    assert payload["meta"]["left_context_size"] == 2
+    assert len(payload["codes"]["audio"]) == _Q * 12
 
 
 def test_ref_code_context_applies_to_all_streaming_chunks():
@@ -274,15 +276,15 @@ def test_ref_code_context_applies_to_all_streaming_chunks():
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.zeros((0,)), "ref_code": ref_code},
+        pooling_output={"codes": {"audio": torch.zeros((0,)), "ref": ref_code}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )
 
     assert payload is not None
     # 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)
+    assert payload["meta"]["left_context_size"] == 10 + 2
+    assert len(payload["codes"]["audio"]) == _Q * (20 + 2)
 
 
 def test_ref_code_context_can_be_buffered_before_first_emit():
@@ -292,7 +294,7 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
 
     first_payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]]), "ref_code": ref_code},
+        pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]]), "ref": ref_code}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )
@@ -302,22 +304,22 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
     for _ in range(8):
         talker2code2wav_async_chunk(
             transfer_manager=tm,
-            pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]])},
+            pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
             request=_req(rid, finished=False, initial_codec_chunk_frames=10),
             is_finished=False,
         )
 
     payload = talker2code2wav_async_chunk(
         transfer_manager=tm,
-        pooling_output={"audio_codes": torch.tensor([[1, 2, 3, 4]])},
+        pooling_output={"codes": {"audio": torch.tensor([[1, 2, 3, 4]])}},
         request=_req(rid, finished=False, initial_codec_chunk_frames=10),
         is_finished=False,
     )
 
     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 payload["meta"]["left_context_size"] == 2
+    assert len(payload["codes"]["audio"]) == _Q * 12
     assert rid in tm.request_payload
 
 
@@ -332,8 +334,9 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
         dtype=torch.long,
     )
     output = SimpleNamespace(
-        multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code},
+        multimodal_output={"codes": {"audio": audio_codes, "ref": ref_code}},
         token_ids=list(range(3)),
+        cumulative_token_ids=list(range(3)),
     )
     stage = SimpleNamespace(
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
@@ -343,7 +346,7 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
 
     assert len(prompts) == 1
     prompt = prompts[0]
-    assert prompt["additional_information"] == {"left_context_size": [2]}
+    assert prompt["additional_information"] == {"meta": {"left_context_size": [2]}}
     assert prompt["prompt_token_ids"] == [
         9,
         8,
@@ -377,8 +380,9 @@ def test_non_async_processor_filters_out_of_range_codec_values():
         dtype=torch.long,
     )
     output = SimpleNamespace(
-        multimodal_output={"audio_codes": audio_codes, "ref_code": ref_code},
+        multimodal_output={"codes": {"audio": audio_codes, "ref": ref_code}},
         token_ids=list(range(4)),
+        cumulative_token_ids=list(range(4)),
     )
     stage = SimpleNamespace(
         engine_outputs=[SimpleNamespace(outputs=[output], finished=True)],
@@ -390,4 +394,4 @@ def test_non_async_processor_filters_out_of_range_codec_values():
     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]}
+    assert prompt["additional_information"] == {"meta": {"left_context_size": [1]}}

tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py

@@ -126,8 +126,8 @@ def test_flush_tail_when_finished():
     )
 
     assert payload is not None
-    assert payload["finished"].item() is True
-    codes = payload["code_predictor_codes"]
+    assert payload["meta"]["finished"].item() is True
+    codes = payload["codes"]["audio"]
     # Format: [ctx_frames, context_length, ...flat_codes]
     assert len(codes) >= 2  # At least ctx_frames + context_length header
     ctx_frames = codes[0]
@@ -149,10 +149,8 @@ def test_eof_marker_when_finished_with_no_frames():
         request=request,
     )
 
-    assert payload == {
-        "code_predictor_codes": [],
-        "finished": torch.tensor(True, dtype=torch.bool),
-    }
+    assert payload["codes"] == {"audio": []}
+    assert payload["meta"]["finished"].item() is True
 
 
 def test_normal_chunk_emission():
@@ -176,7 +174,7 @@ def test_normal_chunk_emission():
 
     # A chunk should be emitted
     assert payload is not None
-    codes = payload["code_predictor_codes"]
+    codes = payload["codes"]["audio"]
     ctx_frames = codes[0]
     context_length = codes[1]
     assert ctx_frames == 20  # 25 - 5(chunk_size_at_begin)
@@ -203,7 +201,7 @@ def test_small_initial_chunks():
         )
 
     assert payload is not None
-    codes = payload["code_predictor_codes"]
+    codes = payload["codes"]["audio"]
     ctx_frames = codes[0]
     context_length = codes[1]
     assert ctx_frames == 0
@@ -252,7 +250,7 @@ def test_context_handling_format():
         )
 
     assert payload is not None
-    codes = payload["code_predictor_codes"]
+    codes = payload["codes"]["audio"]
     # First two elements are ctx_frames and context_length
     ctx_frames = codes[0]
     context_length = codes[1]

tests/test_config_factory.py

@@ -701,7 +701,7 @@ stages:
 
 
 class TestPipelineDiscovery:
-    """Tests for the central pipeline registry (``pipeline_registry._VLLM_OMNI_PIPELINES``)."""
+    """Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``)."""
 
     def test_registry_has_known_models(self):
         """Built-in pipelines are lazy-loaded from the central declaration

tests/test_data_entry_keys.py

@@ -0,0 +1,252 @@
+"""Tests for data_entry_keys: TypedDict payload structure, flatten/unflatten, serialize/deserialize."""
+
+import torch
+
+from vllm_omni.data_entry_keys import (
+    OmniPayload,
+    deserialize_payload,
+    flatten_payload,
+    serialize_payload,
+    unflatten_payload,
+)
+from vllm_omni.engine import AdditionalInformationPayload
+
+
+class TestOmniPayload:
+    def test_nested_payload_structure(self):
+        """Verify OmniPayload can be constructed with nested dicts."""
+        payload: OmniPayload = {
+            "hidden_states": {"output": torch.tensor([1.0])},
+            "embed": {"prefill": torch.tensor([2.0])},
+            "codes": {"audio": torch.tensor([3.0])},
+            "ids": {"all": [1, 2, 3]},
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+        }
+        assert torch.equal(payload["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(payload["embed"]["prefill"], torch.tensor([2.0]))
+        assert torch.equal(payload["codes"]["audio"], torch.tensor([3.0]))
+        assert payload["ids"]["all"] == [1, 2, 3]
+        assert payload["meta"]["finished"].item() is True
+
+    def test_partial_payload(self):
+        """OmniPayload fields are all optional (total=False)."""
+        payload: OmniPayload = {"meta": {"finished": torch.tensor(False, dtype=torch.bool)}}
+        assert payload["meta"]["finished"].item() is False
+
+    def test_empty_payload(self):
+        payload: OmniPayload = {}
+        assert len(payload) == 0
+
+
+class TestFlattenPayload:
+    def test_basic_nested_to_dotted(self):
+        nested = {
+            "codes": {"audio": torch.tensor([1.0])},
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 5},
+        }
+        flat = flatten_payload(nested)
+        assert torch.equal(flat["codes.audio"], torch.tensor([1.0]))
+        assert flat["meta.finished"].item() is True
+        assert flat["meta.left_context_size"] == 5
+        assert "codes" not in flat
+        assert "meta" not in flat
+
+    def test_top_level_keys_preserved(self):
+        nested = {
+            "latent": torch.tensor([9.0]),
+            "generated_len": 42,
+        }
+        flat = flatten_payload(nested)
+        assert torch.equal(flat["latent"], torch.tensor([9.0]))
+        assert flat["generated_len"] == 42
+
+    def test_hidden_states_layers_expanded(self):
+        nested = {
+            "hidden_states": {
+                "output": torch.tensor([1.0]),
+                "layers": {
+                    0: torch.tensor([2.0]),
+                    24: torch.tensor([3.0]),
+                },
+            },
+        }
+        flat = flatten_payload(nested)
+        assert torch.equal(flat["hidden_states.output"], torch.tensor([1.0]))
+        assert torch.equal(flat["hidden_states.layer_0"], torch.tensor([2.0]))
+        assert torch.equal(flat["hidden_states.layer_24"], torch.tensor([3.0]))
+        assert "hidden_states.layers" not in flat
+
+    def test_empty_payload(self):
+        assert flatten_payload({}) == {}
+
+    def test_mixed_nested_and_top_level(self):
+        nested: OmniPayload = {
+            "codes": {"audio": torch.tensor([1.0])},
+            "latent": torch.tensor([2.0]),
+            "meta": {"finished": torch.tensor(False, dtype=torch.bool)},
+        }
+        flat = flatten_payload(nested)
+        assert set(flat.keys()) == {"codes.audio", "latent", "meta.finished"}
+
+
+class TestUnflattenPayload:
+    def test_basic_dotted_to_nested(self):
+        flat = {
+            "codes.audio": torch.tensor([1.0]),
+            "meta.finished": torch.tensor(True, dtype=torch.bool),
+            "meta.left_context_size": 5,
+        }
+        nested = unflatten_payload(flat)
+        assert torch.equal(nested["codes"]["audio"], torch.tensor([1.0]))
+        assert nested["meta"]["finished"].item() is True
+        assert nested["meta"]["left_context_size"] == 5
+
+    def test_top_level_keys_preserved(self):
+        flat = {"latent": torch.tensor([9.0]), "generated_len": 42}
+        nested = unflatten_payload(flat)
+        assert torch.equal(nested["latent"], torch.tensor([9.0]))
+        assert nested["generated_len"] == 42
+
+    def test_hidden_states_layers_collected(self):
+        flat = {
+            "hidden_states.output": torch.tensor([1.0]),
+            "hidden_states.layer_0": torch.tensor([2.0]),
+            "hidden_states.layer_24": torch.tensor([3.0]),
+        }
+        nested = unflatten_payload(flat)
+        assert torch.equal(nested["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(nested["hidden_states"]["layers"][0], torch.tensor([2.0]))
+        assert torch.equal(nested["hidden_states"]["layers"][24], torch.tensor([3.0]))
+
+    def test_empty_payload(self):
+        assert unflatten_payload({}) == {}
+
+
+class TestFlattenUnflattenRoundTrip:
+    def test_round_trip_simple(self):
+        original: OmniPayload = {
+            "codes": {"audio": torch.tensor([1.0, 2.0])},
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 10},
+            "ids": {"prompt": [1, 2, 3]},
+            "latent": torch.tensor([5.0]),
+        }
+        restored = unflatten_payload(flatten_payload(original))
+        assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"])
+        assert restored["meta"]["finished"].item() is True
+        assert restored["meta"]["left_context_size"] == 10
+        assert restored["ids"]["prompt"] == [1, 2, 3]
+        assert torch.equal(restored["latent"], original["latent"])
+
+    def test_round_trip_with_layers(self):
+        original = {
+            "hidden_states": {
+                "output": torch.tensor([1.0]),
+                "layers": {0: torch.tensor([2.0]), 24: torch.tensor([3.0])},
+            },
+        }
+        restored = unflatten_payload(flatten_payload(original))
+        assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][0], torch.tensor([2.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][24], torch.tensor([3.0]))
+
+    def test_round_trip_all_categories(self):
+        original: OmniPayload = {
+            "hidden_states": {"output": torch.tensor([1.0]), "last": torch.tensor([2.0])},
+            "embed": {"prefill": torch.tensor([3.0]), "tts_bos": torch.tensor([4.0])},
+            "codes": {"audio": torch.tensor([5.0]), "ref": torch.tensor([6.0])},
+            "ids": {"all": [1, 2], "prompt": [3, 4]},
+            "meta": {"finished": torch.tensor(False, dtype=torch.bool), "ar_width": 8},
+        }
+        restored = unflatten_payload(flatten_payload(original))
+        assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(restored["hidden_states"]["last"], torch.tensor([2.0]))
+        assert torch.equal(restored["embed"]["prefill"], torch.tensor([3.0]))
+        assert torch.equal(restored["embed"]["tts_bos"], torch.tensor([4.0]))
+        assert torch.equal(restored["codes"]["audio"], torch.tensor([5.0]))
+        assert torch.equal(restored["codes"]["ref"], torch.tensor([6.0]))
+        assert restored["ids"]["all"] == [1, 2]
+        assert restored["ids"]["prompt"] == [3, 4]
+        assert restored["meta"]["finished"].item() is False
+        assert restored["meta"]["ar_width"] == 8
+
+
+class TestSerializeDeserializePayload:
+    def test_tensor_round_trip(self):
+        original: OmniPayload = {
+            "hidden_states": {"output": torch.tensor([[1.0, 2.0], [3.0, 4.0]])},
+        }
+        wire = serialize_payload(original)
+        assert isinstance(wire, AdditionalInformationPayload)
+        restored = deserialize_payload(wire)
+        assert torch.equal(restored["hidden_states"]["output"], original["hidden_states"]["output"])
+
+    def test_list_round_trip(self):
+        original: OmniPayload = {
+            "ids": {"prompt": [10, 20, 30]},
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert restored["ids"]["prompt"] == [10, 20, 30]
+
+    def test_finished_tensor_round_trip(self):
+        original: OmniPayload = {
+            "meta": {"finished": torch.tensor(True, dtype=torch.bool), "left_context_size": 5},
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert isinstance(restored["meta"]["finished"], torch.Tensor)
+        assert restored["meta"]["finished"].dtype == torch.bool
+        assert restored["meta"]["finished"].item() is True
+        assert restored["meta"]["left_context_size"] == 5
+
+    def test_mixed_types_round_trip(self):
+        original: OmniPayload = {
+            "hidden_states": {"output": torch.tensor([1.0, 2.0])},
+            "ids": {"all": [1, 2, 3]},
+            "meta": {"finished": torch.tensor(False, dtype=torch.bool), "ar_width": 4},
+            "codes": {"audio": torch.tensor([3.0])},
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert torch.equal(restored["hidden_states"]["output"], original["hidden_states"]["output"])
+        assert restored["ids"]["all"] == [1, 2, 3]
+        assert restored["meta"]["finished"].item() is False
+        assert restored["meta"]["ar_width"] == 4
+        assert torch.equal(restored["codes"]["audio"], original["codes"]["audio"])
+
+    def test_hidden_states_layers_round_trip(self):
+        original = {
+            "hidden_states": {
+                "output": torch.tensor([1.0]),
+                "layers": {0: torch.tensor([2.0]), 24: torch.tensor([3.0])},
+            },
+        }
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert torch.equal(restored["hidden_states"]["output"], torch.tensor([1.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][0], torch.tensor([2.0]))
+        assert torch.equal(restored["hidden_states"]["layers"][24], torch.tensor([3.0]))
+
+    def test_tensor_dtype_preserved(self):
+        # bfloat16 excluded: numpy() doesn't support it; callers must cast before serializing.
+        for dtype in [torch.float16, torch.float32, torch.int64, torch.int32, torch.bool]:
+            original: OmniPayload = {"codes": {"audio": torch.tensor([1], dtype=dtype)}}
+            wire = serialize_payload(original)
+            restored = deserialize_payload(wire)
+            assert restored["codes"]["audio"].dtype == dtype, f"dtype mismatch for {dtype}"
+
+    def test_tensor_shape_preserved(self):
+        t = torch.randn(3, 4, 5)
+        original: OmniPayload = {"hidden_states": {"output": t}}
+        wire = serialize_payload(original)
+        restored = deserialize_payload(wire)
+        assert restored["hidden_states"]["output"].shape == (3, 4, 5)
+        assert torch.allclose(restored["hidden_states"]["output"], t)
+
+    def test_empty_payload_returns_none(self):
+        assert serialize_payload({}) is None
+
+    def test_none_values_skipped(self):
+        original: OmniPayload = {"meta": {"finished": None}}
+        wire = serialize_payload(original)
+        assert wire is None

tests/worker/test_omni_connector_mixin.py

@@ -891,9 +891,8 @@ class TestTPAsyncChunkFanout(unittest.TestCase):
     def test_rank0_only_polls_connector_for_tp_async_chunk(self):
         host = self._make_host(rank=0)
         payload = {
-            "code_predictor_codes": [10, 11],
-            "left_context_size": 0,
-            "finished": torch.tensor(False),
+            "codes": {"audio": [10, 11]},
+            "meta": {"left_context_size": 0, "finished": torch.tensor(False)},
         }
         host._omni_connector.get.return_value = (payload, 123)
         tp_group = _FakeTPGroup(world_size=2, rank_in_group=0)
@@ -1165,9 +1164,8 @@ class TestAsyncPayloadLifecycle(unittest.TestCase):
 
         host._omni_connector.get.return_value = (
             {
-                "code_predictor_codes": [20, 21, 22],
-                "left_context_size": 0,
-                "finished": torch.tensor(False),
+                "codes": {"audio": [20, 21, 22]},
+                "meta": {"left_context_size": 0, "finished": torch.tensor(False)},
             },
             1,
         )

tests/worker/test_omni_gpu_model_runner.py

@@ -119,7 +119,7 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
     runner.text_step = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
 
     runner.talker_mtp = DummyTalkerMTP()
-    runner.model = SimpleNamespace(talker_mtp_output_key="code_predictor_codes")
+    runner.model = SimpleNamespace(talker_mtp_output_key=("codes", "audio"))
     runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace())
 
     # Provide a minimal implementation that returns the expected 4-tuple.
@@ -188,8 +188,8 @@ def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
     # Validate per-request additional_information_cpu was updated
     info_r1 = runner.requests["r1"].additional_information_cpu
     info_r2 = runner.requests["r2"].additional_information_cpu
-    assert int(info_r1["code_predictor_codes"][0, 0]) == 0
-    assert int(info_r2["code_predictor_codes"][0, 0]) == 1
+    assert int(info_r1["codes"]["audio"][0, 0]) == 0
+    assert int(info_r2["codes"]["audio"][0, 0]) == 1
 
 
 def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):

vllm_omni/__init__.py

@@ -28,6 +28,7 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
 
 # Register custom configs (AutoConfig, AutoTokenizer) as early as possible.
 from vllm_omni.transformers_utils import configs as _configs  # noqa: F401, E402
+from vllm_omni.transformers_utils import parsers as _parsers  # noqa: F401, E402
 
 from .config import OmniModelConfig
 

vllm_omni/config/pipeline_registry.py

@@ -15,7 +15,13 @@ they just no longer need to self-register via ``register_pipeline(...)``.
 Adding a new pipeline:
     1. Define the ``PipelineConfig`` instance as a module-level variable in
        ``vllm_omni/.../pipeline.py``.
-    2. Add one line to ``_OMNI_PIPELINES`` or ``_DIFFUSION_PIPELINES`` below.
+    2. Add one line to ``_OMNI_PIPELINES`` below.
+
+Single-stage diffusion models continue to use the
+``_create_default_diffusion_stage_cfg`` fallback in
+``async_omni_engine.py`` — they don't need a registry entry. The empty
+``_DIFFUSION_PIPELINES`` placeholder previously here (#2915) was removed
+once #2987 (which would have populated it) was deferred.
 
 ``register_pipeline(config)`` in ``stage_config`` is still supported for
 out-of-tree plugins and tests that create pipelines at runtime; those override
@@ -79,13 +85,8 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.fish_speech.pipeline",
         "FISH_SPEECH_PIPELINE",
     ),
-}
-
-# --- Single-stage diffusion pipelines (populated in PR 3/N) ---
-_DIFFUSION_PIPELINES: dict[str, tuple[str, str]] = {}
-
-# Union view used by ``_LazyPipelineRegistry``; don't mutate at runtime.
-_VLLM_OMNI_PIPELINES: dict[str, tuple[str, str]] = {
-    **_OMNI_PIPELINES,
-    **_DIFFUSION_PIPELINES,
+    "moss_tts_nano": (
+        "vllm_omni.model_executor.models.moss_tts_nano.pipeline",
+        "MOSS_TTS_NANO_PIPELINE",
+    ),
 }

vllm_omni/config/stage_config.py

@@ -274,9 +274,9 @@ class _LazyPipelineRegistry:
 
     def _get_lazy_map(self) -> dict[str, tuple[str, str]]:
         if self._lazy_map is None:
-            from vllm_omni.config.pipeline_registry import _VLLM_OMNI_PIPELINES
+            from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
 
-            self._lazy_map = _VLLM_OMNI_PIPELINES
+            self._lazy_map = _OMNI_PIPELINES
         return self._lazy_map
 
     def _load_lazy(self, model_type: str) -> PipelineConfig | None:
@@ -374,7 +374,7 @@ _PIPELINE_REGISTRY = _LazyPipelineRegistry()
 def register_pipeline(pipeline: PipelineConfig) -> None:
     """Register a pipeline config dynamically.
 
-    In-tree pipelines are declared in ``pipeline_registry._VLLM_OMNI_PIPELINES``
+    In-tree pipelines are declared in ``pipeline_registry._OMNI_PIPELINES``
     and loaded lazily; calling ``register_pipeline`` is only needed for
     out-of-tree plugins or tests that build a ``PipelineConfig`` at runtime.
     A dynamic registration overrides the central-registry entry with the same

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -314,7 +314,7 @@ class OmniSchedulingCoordinator:
                 runtime_seed = None
                 if "left_context_size" in metadata:
                     runtime_seed = {
-                        "left_context_size": metadata["left_context_size"],
+                        "meta": {"left_context_size": metadata["left_context_size"]},
                     }
                 request._omni_initial_model_buffer = runtime_seed
                 if new_ids:

vllm_omni/data_entry_keys.py

@@ -0,0 +1,277 @@
+"""Structured payload types for inter-stage communication.
+
+Adding a new model?
+~~~~~~~~~~~~~~~~~~~
+Every key you put into the inter-stage payload (``additional_information``,
+``multimodal_output``, ``pooling_output``) **must** use the nested
+``OmniPayload`` TypedDict structure.  For each category, every known
+qualifier is an explicit field so misspellings are caught statically.
+
+Categories
+    hidden_states  – intermediate / output hidden-state tensors
+    embed          – embedding tensors (prefill, decode, special tokens)
+    ids            – token-ID sequences
+    codes          – codec / audio code tensors
+    meta           – scalar metadata, control flags, shapes
+
+This module provides:
+- Structured ``TypedDict`` types for static type checking (``OmniPayload``)
+- ``serialize_payload`` / ``deserialize_payload`` for transport across
+  process boundaries via ``AdditionalInformationPayload``
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, TypedDict
+
+import numpy as np
+import torch
+
+if TYPE_CHECKING:
+    from vllm_omni.engine import AdditionalInformationEntry, AdditionalInformationPayload
+
+# ── Structured payload types ──
+# These are TypedDicts (plain dicts at runtime, zero overhead) that give
+# static type checking and IDE autocomplete for inter-stage payloads.
+# Every field is optional (total=False) because each stage only populates
+# the subset it needs.
+
+
+class HiddenStates(TypedDict, total=False):
+    output: torch.Tensor
+    trailing_text: torch.Tensor
+    last: torch.Tensor
+    layers: dict[int, torch.Tensor]
+
+
+class Embeddings(TypedDict, total=False):
+    prefill: torch.Tensor
+    decode: torch.Tensor
+    cached_decode: torch.Tensor
+    tts_bos: torch.Tensor
+    tts_eos: torch.Tensor
+    tts_pad: torch.Tensor
+    tts_pad_projected: torch.Tensor
+    voice: torch.Tensor
+    speech_feat: torch.Tensor
+    thinker_reply: torch.Tensor
+
+
+class Codes(TypedDict, total=False):
+    audio: torch.Tensor
+    ref: torch.Tensor
+
+
+class Ids(TypedDict, total=False):
+    all: list[int]
+    prompt: list[int]
+    output: list[int]
+    speech_token: list[int]
+    prior_image: list[int]
+
+
+class OmniPayloadMeta(TypedDict, total=False):
+    finished: torch.Tensor
+    left_context_size: int
+    override_keys: list[tuple[str, str]]
+    num_processed_tokens: int
+    next_stage_prompt_len: int
+    ar_width: int
+    eol_token_id: int
+    visual_token_start_id: int
+    visual_token_end_id: int
+    gen_token_mask: torch.Tensor
+    omni_task: list[str]
+    height: int
+    width: int
+    decode_flag: bool
+    codec_streaming: bool
+    ref_code_len: int
+    talker_prefill_offset: int
+
+
+class OmniPayload(TypedDict, total=False):
+    hidden_states: HiddenStates
+    embed: Embeddings
+    ids: Ids
+    codes: Codes
+    meta: OmniPayloadMeta
+    latent: torch.Tensor
+    generated_len: int
+    model_outputs: list[torch.Tensor]
+    mtp_inputs: tuple[torch.Tensor, torch.Tensor]
+    speaker: Any
+    language: Any
+    request_id: str
+
+
+# ── Keys whose values are nested dicts (TypedDict sub-categories) ──
+_NESTED_KEYS = frozenset({"hidden_states", "embed", "ids", "codes", "meta"})
+
+# Sub-TypedDict for each nested category, used by runtime validation.
+_NESTED_SCHEMAS: dict[str, type] = {
+    "hidden_states": HiddenStates,
+    "embed": Embeddings,
+    "ids": Ids,
+    "codes": Codes,
+    "meta": OmniPayloadMeta,
+}
+
+_ROOT_KEYS: frozenset[str] = frozenset(OmniPayload.__annotations__.keys())
+
+
+def assert_payload(payload: dict[str, Any], *, context: str = "payload") -> None:
+    """Validate ``payload`` matches the ``OmniPayload`` nested schema.
+
+    TypedDict is a static-only contract in Python; this helper closes the
+    loop at runtime by rejecting:
+      * non-dict payloads
+      * top-level keys not declared on ``OmniPayload``
+      * nested-category values that aren't dicts
+      * sub-keys not declared on the matching nested TypedDict
+
+    Call at producer/consumer boundaries when a schema violation should
+    crash the pipeline instead of silently degrading audio quality.
+    """
+    assert isinstance(payload, dict), f"{context}: expected dict, got {type(payload).__name__}"
+    extra_top = set(payload) - _ROOT_KEYS
+    assert not extra_top, f"{context}: unknown top-level keys {sorted(extra_top)!r}"
+    for nested_key, schema in _NESTED_SCHEMAS.items():
+        if nested_key not in payload:
+            continue
+        sub = payload[nested_key]
+        assert isinstance(sub, dict), f"{context}: payload[{nested_key!r}] must be dict, got {type(sub).__name__}"
+        known_sub = frozenset(schema.__annotations__.keys())
+        extra_sub = set(sub) - known_sub
+        assert not extra_sub, f"{context}: payload[{nested_key!r}] unknown sub-keys {sorted(extra_sub)!r}"
+
+
+def flatten_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    """Flatten a nested ``OmniPayload`` to dotted keys.
+
+    Nested sub-dicts under ``_NESTED_KEYS`` are expanded:
+    ``{"codes": {"audio": tensor}}`` → ``{"codes.audio": tensor}``.
+    ``hidden_states["layers"]`` is expanded to ``hidden_states.layer_N``.
+    Top-level values are kept as-is.
+    """
+    if not payload:
+        return {}
+    flat: dict[str, Any] = {}
+    for key, value in payload.items():
+        if key in _NESTED_KEYS and isinstance(value, dict):
+            for qual, val in value.items():
+                if qual == "layers" and key == "hidden_states" and isinstance(val, dict):
+                    for layer_idx, tensor in val.items():
+                        flat[f"hidden_states.layer_{layer_idx}"] = tensor
+                else:
+                    flat[f"{key}.{qual}"] = val
+        else:
+            flat[key] = value
+    return flat
+
+
+def unflatten_payload(flat: dict[str, Any]) -> dict[str, Any]:
+    """Unflatten dotted keys back to nested dicts.
+
+    Reverse of :func:`flatten_payload`.
+    ``hidden_states.layer_N`` keys are collected into ``hidden_states.layers``.
+    """
+    result: dict[str, Any] = {}
+    for key, value in flat.items():
+        if "." in key:
+            type_key, qualifier = key.split(".", 1)
+            sub = result.setdefault(type_key, {})
+            if type_key == "hidden_states" and qualifier.startswith("layer_"):
+                layers = sub.setdefault("layers", {})
+                layer_idx = int(qualifier[len("layer_") :])
+                layers[layer_idx] = value
+            else:
+                sub[qualifier] = value
+        else:
+            result[key] = value
+    return result
+
+
+# ── dtype helpers ──
+_DTYPE_TO_NAME: dict[torch.dtype, str] = {
+    torch.float32: "float32",
+    torch.float16: "float16",
+    torch.bfloat16: "bfloat16",
+    torch.float64: "float64",
+    torch.int64: "int64",
+    torch.int32: "int32",
+    torch.int16: "int16",
+    torch.int8: "int8",
+    torch.uint8: "uint8",
+    torch.bool: "bool",
+}
+
+
+def _dtype_to_name(dtype: torch.dtype) -> str:
+    return _DTYPE_TO_NAME.get(dtype, str(dtype).replace("torch.", ""))
+
+
+def _serialize_tensor(t: torch.Tensor) -> AdditionalInformationEntry:
+    from vllm_omni.engine import AdditionalInformationEntry
+
+    t_cpu = t.detach().to("cpu").contiguous()
+    return AdditionalInformationEntry(
+        tensor_data=t_cpu.numpy().tobytes(),
+        tensor_shape=list(t_cpu.shape),
+        tensor_dtype=_dtype_to_name(t_cpu.dtype),
+    )
+
+
+def _deserialize_tensor(entry: AdditionalInformationEntry) -> torch.Tensor:
+    dt = np.dtype(entry.tensor_dtype or "float32")
+    arr = np.frombuffer(entry.tensor_data, dtype=dt)  # type: ignore[arg-type]
+    arr = arr.reshape(entry.tensor_shape)
+    return torch.from_numpy(arr.copy())
+
+
+def serialize_payload(
+    payload: OmniPayload,
+) -> AdditionalInformationPayload | None:
+    """Serialize an ``OmniPayload`` for EngineCore transport.
+
+    Uses :func:`flatten_payload` to produce dotted keys, then converts
+    each value to an ``AdditionalInformationEntry``.
+    """
+    from vllm_omni.engine import (
+        AdditionalInformationEntry,
+        AdditionalInformationPayload,
+    )
+
+    flat = flatten_payload(payload)
+    entries: dict[str, AdditionalInformationEntry] = {}
+
+    for key, value in flat.items():
+        if isinstance(value, torch.Tensor):
+            entries[key] = _serialize_tensor(value)
+        elif isinstance(value, list):
+            entries[key] = AdditionalInformationEntry(list_data=value)
+        elif value is not None:
+            entries[key] = AdditionalInformationEntry(scalar_data=value)
+
+    return AdditionalInformationPayload(entries=entries) if entries else None
+
+
+def deserialize_payload(
+    wire: AdditionalInformationPayload,
+) -> OmniPayload:
+    """Deserialize an ``AdditionalInformationPayload`` back to ``OmniPayload``.
+
+    Decodes entries to tensors/lists, then uses :func:`unflatten_payload`
+    to reconstruct the nested structure.
+    """
+    flat: dict[str, Any] = {}
+
+    for key, entry in wire.entries.items():
+        if entry.tensor_data is not None:
+            flat[key] = _deserialize_tensor(entry)
+        elif entry.list_data is not None:
+            flat[key] = entry.list_data
+        elif entry.scalar_data is not None:
+            flat[key] = entry.scalar_data
+
+    return unflatten_payload(flat)  # type: ignore[return-value]

vllm_omni/deploy/glm_image.yaml

@@ -10,6 +10,10 @@ async_chunk: false
 stages:
   # Stage 0: AR Model (GlmImageForConditionalGeneration)
   # Generates prior_token_ids for conditioning the diffusion process.
+  # max_tokens is set to the theoretical maximum (2048x2048 t2i ≈ 4353
+  # tokens). The AR model stops naturally at EOS (16385) well before
+  # this ceiling at lower resolutions. Do NOT override max_tokens per
+  # request — the actual resolution is communicated via target_h/w.
   - stage_id: 0
     max_num_seqs: 1
     gpu_memory_utilization: 0.6
@@ -21,6 +25,7 @@ stages:
       top_p: 0.75
       top_k: 16512
       stop_token_ids: [16385]
+      max_tokens: 4353
       seed: 42
       detokenize: false
 

vllm_omni/deploy/moss_tts_nano.yaml

@@ -0,0 +1,31 @@
+# MOSS-TTS-Nano deploy: single-stage AR TTS (0.1B).
+# Verified on 1x L4 24GB (and sufficient for H20 / A100 / 3090).
+#
+# Fields omitted from a stage fall back to StageDeployConfig defaults (see
+# vllm_omni/config/stage_config.py). Notable choices:
+#   * enforce_eager=true — the model uses upstream inference_stream() with
+#     trust_remote_code; CUDA graph capture is not wired up.
+#   * gpu_memory_utilization=0.3 — the AR LM + codec fit in ~2 GiB.
+#   * max_num_seqs=4 — mirrors the legacy stage_configs contract; per-request
+#     streaming generators are stored in self._stream_gens.
+#
+# MOSS-TTS-Nano has no multimodal input preprocessing, so skip_mm_profiling
+# bypasses the dummy-MM pass at profile time.
+async_chunk: false
+trust_remote_code: true
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 4
+    gpu_memory_utilization: 0.3
+    enforce_eager: true
+    max_num_batched_tokens: 4096
+    max_model_len: 4096
+    devices: "0"
+    skip_mm_profiling: true
+    default_sampling_params:
+      temperature: 1.0
+      top_p: 1.0
+      top_k: 50
+      max_tokens: 4096
+      seed: 42

vllm_omni/diffusion/attention/backends/utils/fa.py

@@ -16,9 +16,12 @@ from functools import lru_cache
 
 import torch
 import torch.nn.functional as F
+from vllm.logger import init_logger
 
 from vllm_omni.platforms import current_omni_platform
 
+logger = init_logger(__name__)
+
 # Flash Attention function detection with fallback chain
 flash_attn_func = None
 flash_attn_varlen_func = None
@@ -81,9 +84,38 @@ HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not No
 
 
 @lru_cache(maxsize=1)
-def is_mate_available() -> bool:
-    """Check if MATE (MUSA AI Tensor Engine) is available."""
-    return current_omni_platform.is_musa() and flash_attn_func is not None or flash_attn_varlen_func is not None
+def is_flash_attn_installed() -> bool:
+    """Return whether a Flash Attention backend package is importable.
+
+    Shared by CUDA/ROCm/MUSA platforms.
+    """
+    try:
+        # Check for any FA backend: FA3 (fa3_fwd_interface, flash_attn_interface) or FA2 (flash_attn)
+        # Try FA3 from fa3-fwd PyPI package
+        try:
+            import fa3_fwd_interface  # noqa: F401
+
+            return True
+        except (ImportError, ModuleNotFoundError):
+            pass
+
+        # Try FA3 from flash-attention source build
+        try:
+            import flash_attn_interface  # noqa: F401
+
+            return True
+        except (ImportError, ModuleNotFoundError):
+            pass
+
+        # Try FA2 from flash-attn package
+        from flash_attn import __version__
+
+        if __version__ < "2.6.0":
+            raise ImportError("install flash_attn >= 2.6.0")
+        return True
+    except (ImportError, ModuleNotFoundError):
+        logger.warning("No Flash Attention backend found, using pytorch SDPA implementation")
+        return False
 
 
 def _index_first_axis(tensor, indices):

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -1069,6 +1069,130 @@ def extract_flux2_context(
     )
 
 
+def extract_flux_context(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    pooled_projections: torch.Tensor,
+    timestep: torch.LongTensor,
+    img_ids: torch.Tensor,
+    txt_ids: torch.Tensor,
+    guidance: torch.Tensor | None = None,
+    joint_attention_kwargs: dict[str, Any] | None = None,
+    **kwargs: Any,
+) -> CacheContext:
+    """
+    Extract cache context for FluxTransformer2DModel.
+
+    This mirrors the standard FLUX.1 transformer forward path while exposing
+    the first modulated image stream tensor as TeaCache's similarity signal.
+
+    Args:
+        module: FluxTransformer2DModel instance
+        hidden_states: Input image hidden states tensor
+        encoder_hidden_states: Text encoder outputs
+        pooled_projections: Pooled text conditioning
+        timestep: Current diffusion timestep
+        img_ids: Image position IDs for RoPE
+        txt_ids: Text position IDs for RoPE
+        guidance: Optional guidance scale for guidance-distilled variants
+        joint_attention_kwargs: Additional attention kwargs
+        **kwargs: Additional keyword arguments
+
+    Returns:
+        CacheContext with all information needed for generic caching
+    """
+    from diffusers.models.modeling_outputs import Transformer2DModelOutput
+
+    if not hasattr(module, "transformer_blocks") or len(module.transformer_blocks) == 0:
+        raise ValueError("Module must have transformer_blocks")
+
+    # ============================================================================
+    # PREPROCESSING (Flux1-specific)
+    # ============================================================================
+    hidden_states = module.x_embedder(hidden_states)
+    timestep = timestep.to(device=hidden_states.device, dtype=hidden_states.dtype) * 1000
+
+    if guidance is not None:
+        guidance = guidance.to(device=hidden_states.device, dtype=hidden_states.dtype) * 1000
+
+    temb = (
+        module.time_text_embed(timestep, pooled_projections)
+        if guidance is None
+        else module.time_text_embed(timestep, guidance, pooled_projections)
+    )
+    encoder_hidden_states = module.context_embedder(encoder_hidden_states)
+
+    if txt_ids.ndim == 3:
+        txt_ids = txt_ids[0]
+    if img_ids.ndim == 3:
+        img_ids = img_ids[0]
+
+    ids = torch.cat((txt_ids, img_ids), dim=0)
+    if current_omni_platform.is_npu():
+        freqs_cos, freqs_sin = module.pos_embed(ids.cpu())
+        image_rotary_emb = (freqs_cos.npu(), freqs_sin.npu())
+    else:
+        image_rotary_emb = module.pos_embed(ids)
+
+    # ============================================================================
+    # EXTRACT MODULATED INPUT (for cache decision)
+    # ============================================================================
+    first_block = module.transformer_blocks[0]
+    modulated_input, *_ = first_block.norm1(hidden_states, emb=temb)
+
+    # ============================================================================
+    # DEFINE TRANSFORMER EXECUTION (Flux1-specific)
+    # ============================================================================
+    def run_transformer_blocks() -> tuple[torch.Tensor, torch.Tensor]:
+        h = hidden_states
+        e = encoder_hidden_states
+
+        for block in module.transformer_blocks:
+            e, h = block(
+                hidden_states=h,
+                encoder_hidden_states=e,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+
+        for block in module.single_transformer_blocks:
+            e, h = block(
+                hidden_states=h,
+                encoder_hidden_states=e,
+                temb=temb,
+                image_rotary_emb=image_rotary_emb,
+                joint_attention_kwargs=joint_attention_kwargs,
+            )
+
+        return (h, e)
+
+    return_dict = kwargs.get("return_dict", True)
+
+    # ============================================================================
+    # DEFINE POSTPROCESSING
+    # ============================================================================
+    def postprocess(h: torch.Tensor) -> Any:
+        h = module.norm_out(h, temb)
+        output = module.proj_out(h)
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+    # ============================================================================
+    # RETURN CONTEXT
+    # ============================================================================
+    return CacheContext(
+        modulated_input=modulated_input,
+        hidden_states=hidden_states,
+        encoder_hidden_states=encoder_hidden_states,
+        temb=temb,
+        run_transformer_blocks=run_transformer_blocks,
+        postprocess=postprocess,
+    )
+
+
 # Registry for model-specific extractors
 # Key: Transformer class name
 # Value: extractor function with signature (module, *args, **kwargs) -> CacheContext
@@ -1083,8 +1207,8 @@ EXTRACTOR_REGISTRY: dict[str, Callable] = {
     "StableAudioDiTModel": extract_stable_audio_context,
     "Flux2Transformer2DModel": extract_flux2_context,
     "LongCatImageTransformer2DModel": extract_longcat_context,
+    "FluxTransformer2DModel": extract_flux_context,
     # Future models:
-    # "FluxTransformer2DModel": extract_flux_context,
     # "CogVideoXTransformer3DModel": extract_cogvideox_context,
 }
 

vllm_omni/diffusion/diffusion_engine.py

@@ -190,22 +190,21 @@ class DiffusionEngine:
 
         metrics = {
             "preprocess_time_ms": preprocess_time * 1000,
-            "diffusion_engine_exec_time_ms": (time.perf_counter() - diffusion_engine_start_time) * 1000,
-            "diffusion_engine_total_time_ms": exec_total_time * 1000,
+            "diffusion_engine_exec_time_ms": exec_total_time * 1000,
+            "diffusion_engine_total_time_ms": step_total_ms,
             "image_num": int(request.sampling_params.num_outputs_per_prompt),
             "resolution": int(request.sampling_params.resolution),
             "postprocess_time_ms": postprocess_time * 1000,
         }
-        if self.pre_process_func is not None:
-            metrics["preprocessing_time_ms"] = preprocess_time * 1000
 
         # Handle single request or multiple requests
+        is_audio_output = supports_audio_output(self.od_config.model_class_name)
         if len(request.prompts) == 1:
             # Single request: return single OmniRequestOutput
             prompt = request.prompts[0]
             request_id = request.request_ids[0] if request.request_ids else ""
 
-            if supports_audio_output(self.od_config.model_class_name):
+            if is_audio_output:
                 request_audio_payload = outputs[0] if len(outputs) == 1 else outputs
                 return [
                     OmniRequestOutput.from_diffusion(
@@ -265,7 +264,7 @@ class DiffusionEngine:
                 request_outputs = outputs[start_idx:end_idx] if output_idx < len(outputs) else []
                 output_idx = end_idx
 
-                if supports_audio_output(self.od_config.model_class_name):
+                if is_audio_output:
                     request_audio_payload = request_outputs[0] if len(request_outputs) == 1 else request_outputs
                     results.append(
                         OmniRequestOutput.from_diffusion(
@@ -416,9 +415,7 @@ class DiffusionEngine:
 
         if supports_audio_input(self.od_config.model_class_name):
             audio_sr = 16000
-            audio_duration_sec = 4
-            audio_array = np.random.randn(audio_sr * audio_duration_sec).astype(np.float32)
-            dummy_audio = audio_array[audio_sr * 1 : audio_sr * 3]
+            dummy_audio = np.random.randn(audio_sr * 2).astype(np.float32)
         else:
             dummy_audio = None
 

vllm_omni/diffusion/envs.py

@@ -5,9 +5,6 @@ import os
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.attention.backends.utils.fa import is_mate_available
 from vllm_omni.platforms import current_omni_platform
 
 if TYPE_CHECKING:
@@ -30,8 +27,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")),
 }
 
-logger = init_logger(__name__)
-
 
 class PackagesEnvChecker:
     """Singleton class for checking package availability."""
@@ -46,58 +41,17 @@ class PackagesEnvChecker:
 
     def initialize(self):
         packages_info = {}
-        packages_info["has_flash_attn"] = self._check_flash_attn(packages_info)
+        packages_info["has_flash_attn"] = self._check_flash_attn()
         self.packages_info = packages_info
 
-    def _check_flash_attn(self, packages_info) -> bool:
+    def _check_flash_attn(self) -> bool:
         """Check if flash attention is available and compatible."""
         platform = current_omni_platform
 
-        # MUSA uses MATE for flash attention
-        if platform.is_musa():
-            return is_mate_available()
-
-        # Flash attention requires CUDA-like platforms (CUDA or ROCm)
-        if not platform.is_cuda_alike():
-            return False
-
-        # Check if devices are available
         if platform.get_device_count() == 0:
             return False
 
-        try:
-            gpu_name = platform.get_device_name()
-            # Turing/Tesla/T4 GPUs don't support flash attention well
-            if "Turing" in gpu_name or "Tesla" in gpu_name or "T4" in gpu_name:
-                return False
-
-            # Check for any FA backend: FA3 (fa3_fwd_interface, flash_attn_interface) or FA2 (flash_attn)
-            # Try FA3 from fa3-fwd PyPI package
-            try:
-                import fa3_fwd_interface  # noqa: F401
-
-                return True
-            except (ImportError, ModuleNotFoundError):
-                pass
-
-            # Try FA3 from flash-attention source build
-            try:
-                import flash_attn_interface  # noqa: F401
-
-                return True
-            except (ImportError, ModuleNotFoundError):
-                pass
-
-            # Try FA2 from flash-attn package
-            from flash_attn import __version__
-
-            if __version__ < "2.6.0":
-                raise ImportError("install flash_attn >= 2.6.0")
-            return True
-        except (ImportError, ModuleNotFoundError):
-            if not packages_info.get("has_aiter", False):
-                logger.warning("No Flash Attention backend found, using pytorch SDPA implementation")
-            return False
+        return platform.has_flash_attn_package()
 
     def get_packages_info(self) -> dict:
         """Get the packages info dictionary."""

vllm_omni/diffusion/models/internvla_a1/__init__.py

@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""InternVLA-A1 diffusion model components."""
+
+from vllm_omni.diffusion.models.internvla_a1.config import (
+    InternVLAA1Config,
+    InternVLAA1TrainMetadata,
+)
+from vllm_omni.diffusion.models.internvla_a1.model_internvla_a1 import (
+    InternVLAA1,
+    InternVLAA1Policy,
+)
+from vllm_omni.diffusion.models.internvla_a1.pipeline_internvla_a1 import (
+    InternVLAA1Pipeline,
+    get_internvla_a1_post_process_func,
+)
+
+__all__ = [
+    "InternVLAA1",
+    "InternVLAA1Config",
+    "InternVLAA1Pipeline",
+    "InternVLAA1TrainMetadata",
+    "InternVLAA1Policy",
+    "get_internvla_a1_post_process_func",
+]

vllm_omni/diffusion/models/internvla_a1/adapter_qwen3_vl.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    ALL_ATTENTION_FUNCTIONS,
+    BaseModelOutputWithPast,
+    Cache,
+    DynamicCache,
+    FlashAttentionKwargs,
+    Qwen3VLConfig,
+    Qwen3VLPreTrainedModel,
+    Qwen3VLTextConfig,
+    Qwen3VLTextRotaryEmbedding,
+    Qwen3VLVisionModel,
+    Unpack,
+    apply_rotary_pos_emb,
+    check_model_inputs,
+    create_causal_mask,
+    deprecate_kwarg,
+    eager_attention_forward,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLForConditionalGeneration as HFQwen3VLForConditionalGeneration,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLModel as HFQwen3VLModel,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLTextAttention as HFQwen3VLTextAttention,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLTextDecoderLayer as HFQwen3VLTextDecoderLayer,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLTextModel as HFQwen3VLTextModel,
+)
+from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+    Qwen3VLTextRMSNorm as HFQwen3VLTextRMSNorm,
+)
+
+
+class Qwen3VLTextRMSNorm(HFQwen3VLTextRMSNorm):
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return (self.weight * hidden_states).to(input_dtype)
+
+
+class Qwen3VLTextAttention(HFQwen3VLTextAttention):
+    def __init__(self, config: Qwen3VLTextConfig, layer_idx: int):
+        super().__init__(config, layer_idx)
+        self.q_norm = Qwen3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = Qwen3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None,
+        past_key_values: Cache | None = None,
+        cache_position: torch.LongTensor | None = None,
+        use_cache: bool = False,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None:
+            if use_cache:
+                cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+                key_states, value_states = past_key_values.update(
+                    key_states, value_states, self.layer_idx, cache_kwargs
+                )
+            else:
+                key_states = torch.cat([past_key_values[self.layer_idx][0], key_states], dim=2)
+                value_states = torch.cat([past_key_values[self.layer_idx][1], value_states], dim=2)
+
+        attention_interface = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+
+class Qwen3VLTextDecoderLayer(HFQwen3VLTextDecoderLayer):
+    def __init__(self, config: Qwen3VLTextConfig, layer_idx: int):
+        super().__init__(config, layer_idx)
+        self.self_attn = Qwen3VLTextAttention(config=config, layer_idx=layer_idx)
+        self.input_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+
+class Qwen3VLTextModel(HFQwen3VLTextModel):
+    config: Qwen3VLTextConfig
+    _no_split_modules = ["Qwen3VLTextDecoderLayer"]
+    _repeated_blocks = ["Qwen3VLTextDecoderLayer"]
+
+    def __init__(self, config: Qwen3VLTextConfig):
+        Qwen3VLPreTrainedModel.__init__(self, config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [Qwen3VLTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = Qwen3VLTextRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        self.post_init()
+
+    @check_model_inputs
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        cache_position: torch.LongTensor | None = None,
+        visual_pos_masks: torch.Tensor | None = None,
+        deepstack_visual_embeds: list[torch.Tensor] | None = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple | BaseModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if use_cache and past_key_values is None and not torch.jit.is_tracing():
+            past_key_values = DynamicCache(config=self.config)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+        elif position_ids.ndim == 2:
+            position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+        if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+            text_position_ids = position_ids[0]
+            position_ids = position_ids[1:]
+        else:
+            text_position_ids = position_ids[0]
+
+        attention_mask = create_causal_mask(
+            config=self.config,
+            input_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            cache_position=cache_position,
+            past_key_values=past_key_values,
+            position_ids=text_position_ids,
+        )
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for layer_idx, decoder_layer in enumerate(self.layers):
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_ids=text_position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                position_embeddings=position_embeddings,
+                **kwargs,
+            )
+
+            if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)):
+                hidden_states = self._deepstack_process(
+                    hidden_states,
+                    visual_pos_masks,
+                    deepstack_visual_embeds[layer_idx],
+                )
+
+        hidden_states = self.norm(hidden_states)
+
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+class Qwen3VLModel(HFQwen3VLModel):
+    base_model_prefix = ""
+    _checkpoint_conversion_mapping = {}
+    accepts_loss_kwargs = False
+    config: Qwen3VLConfig
+    _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"]
+
+    def __init__(self, config):
+        Qwen3VLPreTrainedModel.__init__(self, config)
+        self.visual = Qwen3VLVisionModel._from_config(config.vision_config)
+        self.language_model = Qwen3VLTextModel._from_config(config.text_config)
+        self.rope_deltas = None
+        self.post_init()
+
+
+class Qwen3VLForConditionalGeneration(HFQwen3VLForConditionalGeneration):
+    _checkpoint_conversion_mapping = {}
+    _tied_weights_keys = ["lm_head.weight"]
+    accepts_loss_kwargs = False
+    config: Qwen3VLConfig
+
+    def __init__(self, config):
+        Qwen3VLPreTrainedModel.__init__(self, config)
+        self.model = Qwen3VLModel(config)
+        self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+        self.post_init()
+
+
+__all__ = [
+    "Qwen3VLVisionModel",
+    "Qwen3VLForConditionalGeneration",
+    "Qwen3VLModel",
+    "Qwen3VLPreTrainedModel",
+    "Qwen3VLTextModel",
+    "apply_rotary_pos_emb",
+    "eager_attention_forward",
+]

vllm_omni/diffusion/models/internvla_a1/config.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from dataclasses import fields as dataclass_fields
+from pathlib import Path
+from typing import Any
+
+from huggingface_hub.constants import HF_HOME
+
+ACTION = "action"
+OBS_STR = "observation"
+OBS_PREFIX = OBS_STR + "."
+OBS_STATE = OBS_STR + ".state"
+OBS_TASK = OBS_STR + ".task"
+OBS_IMAGE = OBS_STR + ".image"
+OBS_IMAGES = OBS_IMAGE + "s"
+OPENPI_ATTENTION_MASK_VALUE = -2.3819763e38
+
+DEFAULT_QWEN3_VL_MODEL = os.getenv("INTERNVLA_A1_PROCESSOR_DIR", "Qwen/Qwen3-VL-2B-Instruct")
+DEFAULT_COSMOS_REPO = "nvidia/Cosmos-Tokenizer-CI8x8"
+DEFAULT_COSMOS_DIR = Path(
+    os.getenv(
+        "INTERNVLA_A1_COSMOS_DIR",
+        str(Path(HF_HOME) / "hub" / "Cosmos-Tokenizer-CI8x8"),
+    )
+)
+
+
+@dataclass
+class InternVLAA1Config:
+    """Standalone-compatible InternVLA-A1 config with a few fake-smoke defaults."""
+
+    type: str = "internvla_a1"
+    qwen3_vl_variant: str = "qwen3_vl_28l"
+    action_expert_variant: str = "qwen3_28l"
+    dtype: str = "bfloat16"
+    device: str = "cuda"
+    chunk_size: int = 50
+    n_action_steps: int = 50
+    max_state_dim: int = 32
+    max_action_dim: int = 32
+    num_inference_steps: int = 10
+    time_sampling_beta_alpha: float = 1.5
+    time_sampling_beta_beta: float = 1.0
+    time_sampling_scale: float = 0.999
+    time_sampling_offset: float = 0.001
+    min_period: float = 4e-3
+    max_period: float = 4.0
+    image_resolution: tuple[int, int] = (224, 224)
+    empty_cameras: int = 0
+    gradient_checkpointing: bool = False
+    compile_model: bool = False
+    compile_mode: str = "max-autotune"
+    enable_regional_compile: bool = False
+    regional_compile_dynamic: bool = True
+    enable_suffix_static_context_optimization: bool = False
+    attn_implementation: str = "eager"
+    tokenizer_max_length: int = 48
+    freeze_vision_encoder: bool = False
+    train_expert_only: bool = False
+    train_vlm_only: bool = False
+    scale_factor: int = 8
+    lambda_gen: float = 0.01
+    input_features: dict[str, Any] = field(default_factory=dict)
+    output_features: dict[str, Any] = field(default_factory=dict)
+
+    # Fake/smoke-only knobs used by the lightweight pipeline path.
+    vocab_size: int = 256
+    hidden_size: int = 128
+    intermediate_size: int = 256
+    num_attention_heads: int = 4
+    num_hidden_layers: int = 2
+    num_cameras: int = 3
+    image_history: int = 2
+    pixel_feature_dim: int = 48
+
+    @classmethod
+    def from_pretrained(cls, checkpoint_dir: str | Path) -> InternVLAA1Config:
+        checkpoint_dir = Path(checkpoint_dir)
+        with open(checkpoint_dir / "config.json", encoding="utf-8") as f:
+            raw = json.load(f)
+        return cls.from_model_config(raw)
+
+    @classmethod
+    def from_model_config(cls, model_config: dict[str, Any] | None) -> InternVLAA1Config:
+        if not model_config:
+            return cls()
+
+        raw = dict(model_config)
+        raw["image_resolution"] = tuple(raw.get("image_resolution", [224, 224]))
+        allowed = {item.name for item in dataclass_fields(cls)}
+        filtered = {key: value for key, value in raw.items() if key in allowed}
+        return cls(**filtered)
+
+
+@dataclass
+class InternVLAA1TrainMetadata:
+    action_mode: str = "delta"
+    processor_model_name: str = DEFAULT_QWEN3_VL_MODEL
+
+    @classmethod
+    def from_pretrained(cls, checkpoint_dir: str | Path) -> InternVLAA1TrainMetadata:
+        checkpoint_dir = Path(checkpoint_dir)
+        path = checkpoint_dir / "train_config.json"
+        if not path.exists():
+            return cls()
+
+        with open(path, encoding="utf-8") as f:
+            raw = json.load(f)
+
+        dataset_cfg = raw.get("dataset", {})
+        processor_model_name = DEFAULT_QWEN3_VL_MODEL
+        for transform in dataset_cfg.get("data_transforms", {}).get("inputs", []):
+            if transform.get("type") == "internvla_a1_processor":
+                processor_model_name = transform.get(
+                    "pretrained_model_name_or_path",
+                    processor_model_name,
+                )
+                break
+
+        return cls(
+            action_mode=dataset_cfg.get("action_mode", "delta"),
+            processor_model_name=processor_model_name,
+        )

vllm_omni/diffusion/models/internvla_a1/cosmos_ci_torch.py

@@ -0,0 +1,416 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file vendors the minimal Continuous Image (CI) tokenizer implementation
+# needed by InternVLA-A1 from NVIDIA/Cosmos-Tokenizer. The code is intentionally
+# reduced to the CI image path only so the runtime does not depend on the
+# external `cosmos_tokenizer` package.
+
+from __future__ import annotations
+
+import math
+from collections import OrderedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+_WAVELETS = {
+    "haar": torch.tensor([0.7071067811865476, 0.7071067811865476]),
+    "rearrange": torch.tensor([1.0, 1.0]),
+}
+_PERSISTENT = False
+
+
+def nonlinearity(x: torch.Tensor) -> torch.Tensor:
+    return x * torch.sigmoid(x)
+
+
+def normalize(in_channels: int, num_groups: int = 32) -> nn.GroupNorm:
+    return nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+class IdentityDistribution(nn.Module):
+    def forward(self, parameters: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        zeros = torch.tensor([0.0], device=parameters.device, dtype=parameters.dtype)
+        return parameters, (zeros, zeros)
+
+
+class Patcher(nn.Module):
+    def __init__(self, patch_size: int = 1, patch_method: str = "haar") -> None:
+        super().__init__()
+        self.patch_size = patch_size
+        self.patch_method = patch_method
+        self.register_buffer("wavelets", _WAVELETS[patch_method], persistent=_PERSISTENT)
+        self.range = range(int(torch.log2(torch.tensor(self.patch_size)).item()))
+        self.register_buffer(
+            "_arange",
+            torch.arange(_WAVELETS[patch_method].shape[0]),
+            persistent=_PERSISTENT,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.patch_method == "haar":
+            return self._haar(x)
+        if self.patch_method == "rearrange":
+            return self._arrange(x)
+        raise ValueError(f"Unknown patch method: {self.patch_method}")
+
+    def _dwt(self, x: torch.Tensor, mode: str = "reflect", *, rescale: bool = False) -> torch.Tensor:
+        dtype = x.dtype
+        h = self.wavelets
+        n = h.shape[0]
+        g = x.shape[1]
+        hl = h.flip(0).reshape(1, 1, -1).repeat(g, 1, 1).to(dtype=dtype)
+        hh = (h * ((-1) ** self._arange)).reshape(1, 1, -1).repeat(g, 1, 1).to(dtype=dtype)
+
+        x = F.pad(x, pad=(n - 2, n - 1, n - 2, n - 1), mode=mode).to(dtype)
+        xl = F.conv2d(x, hl.unsqueeze(2), groups=g, stride=(1, 2))
+        xh = F.conv2d(x, hh.unsqueeze(2), groups=g, stride=(1, 2))
+        xll = F.conv2d(xl, hl.unsqueeze(3), groups=g, stride=(2, 1))
+        xlh = F.conv2d(xl, hh.unsqueeze(3), groups=g, stride=(2, 1))
+        xhl = F.conv2d(xh, hl.unsqueeze(3), groups=g, stride=(2, 1))
+        xhh = F.conv2d(xh, hh.unsqueeze(3), groups=g, stride=(2, 1))
+
+        out = torch.cat([xll, xlh, xhl, xhh], dim=1)
+        if rescale:
+            out = out / 2
+        return out
+
+    def _haar(self, x: torch.Tensor) -> torch.Tensor:
+        for _ in self.range:
+            x = self._dwt(x, rescale=True)
+        return x
+
+    def _arrange(self, x: torch.Tensor) -> torch.Tensor:
+        return rearrange(
+            x,
+            "b c (h p1) (w p2) -> b (c p1 p2) h w",
+            p1=self.patch_size,
+            p2=self.patch_size,
+        ).contiguous()
+
+
+class UnPatcher(nn.Module):
+    def __init__(self, patch_size: int = 1, patch_method: str = "haar") -> None:
+        super().__init__()
+        self.patch_size = patch_size
+        self.patch_method = patch_method
+        self.register_buffer("wavelets", _WAVELETS[patch_method], persistent=_PERSISTENT)
+        self.range = range(int(torch.log2(torch.tensor(self.patch_size)).item()))
+        self.register_buffer(
+            "_arange",
+            torch.arange(_WAVELETS[patch_method].shape[0]),
+            persistent=_PERSISTENT,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.patch_method == "haar":
+            return self._ihaar(x)
+        if self.patch_method == "rearrange":
+            return self._iarrange(x)
+        raise ValueError(f"Unknown patch method: {self.patch_method}")
+
+    def _idwt(self, x: torch.Tensor, *, rescale: bool = False) -> torch.Tensor:
+        dtype = x.dtype
+        h = self.wavelets
+        n = h.shape[0]
+        g = x.shape[1] // 4
+        hl = h.flip([0]).reshape(1, 1, -1).repeat([g, 1, 1]).to(dtype=dtype)
+        hh = (h * ((-1) ** self._arange)).reshape(1, 1, -1).repeat(g, 1, 1).to(dtype=dtype)
+
+        xll, xlh, xhl, xhh = torch.chunk(x.to(dtype), 4, dim=1)
+        yl = F.conv_transpose2d(xll, hl.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        yl += F.conv_transpose2d(xlh, hh.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        yh = F.conv_transpose2d(xhl, hl.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        yh += F.conv_transpose2d(xhh, hh.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0))
+        y = F.conv_transpose2d(yl, hl.unsqueeze(2), groups=g, stride=(1, 2), padding=(0, n - 2))
+        y += F.conv_transpose2d(yh, hh.unsqueeze(2), groups=g, stride=(1, 2), padding=(0, n - 2))
+
+        if rescale:
+            y = y * 2
+        return y
+
+    def _ihaar(self, x: torch.Tensor) -> torch.Tensor:
+        for _ in self.range:
+            x = self._idwt(x, rescale=True)
+        return x
+
+    def _iarrange(self, x: torch.Tensor) -> torch.Tensor:
+        return rearrange(
+            x,
+            "b (c p1 p2) h w -> b c (h p1) (w p2)",
+            p1=self.patch_size,
+            p2=self.patch_size,
+        )
+
+
+class Upsample(nn.Module):
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3)
+        return self.conv(x)
+
+
+class Downsample(nn.Module):
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(F.pad(x, (0, 1, 0, 1), mode="constant", value=0))
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, *, in_channels: int, out_channels: int | None = None, dropout: float, **_: object) -> None:
+        super().__init__()
+        out_channels = in_channels if out_channels is None else out_channels
+        self.norm1 = normalize(in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = normalize(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.nin_shortcut = (
+            nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
+            if in_channels != out_channels
+            else nn.Identity()
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.conv1(nonlinearity(self.norm1(x)))
+        h = self.conv2(self.dropout(nonlinearity(self.norm2(h))))
+        return self.nin_shortcut(x) + h
+
+
+class AttnBlock(nn.Module):
+    def __init__(self, in_channels: int) -> None:
+        super().__init__()
+        self.norm = normalize(in_channels)
+        self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+        self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        q = self.q(h)
+        k = self.k(h)
+        v = self.v(h)
+        b, c, height, width = q.shape
+        q = q.reshape(b, c, height * width).permute(0, 2, 1)
+        k = k.reshape(b, c, height * width)
+        attn = torch.bmm(q, k) * (int(c) ** -0.5)
+        attn = F.softmax(attn, dim=2)
+        v = v.reshape(b, c, height * width)
+        attn = attn.permute(0, 2, 1)
+        h = torch.bmm(v, attn).reshape(b, c, height, width)
+        return x + self.proj_out(h)
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        channels: int,
+        channels_mult: list[int],
+        num_res_blocks: int,
+        attn_resolutions: list[int],
+        dropout: float,
+        resolution: int,
+        z_channels: int,
+        spatial_compression: int,
+        **kwargs: object,
+    ) -> None:
+        super().__init__()
+        self.num_resolutions = len(channels_mult)
+        self.num_res_blocks = num_res_blocks
+
+        patch_size = int(kwargs.get("patch_size", 1))
+        self.patcher = Patcher(patch_size, str(kwargs.get("patch_method", "rearrange")))
+        in_channels = in_channels * patch_size * patch_size
+        self.num_downsamples = int(math.log2(spatial_compression)) - int(math.log2(patch_size))
+        assert self.num_downsamples <= self.num_resolutions
+
+        self.conv_in = nn.Conv2d(in_channels, channels, kernel_size=3, stride=1, padding=1)
+        curr_res = resolution // patch_size
+        in_ch_mult = (1,) + tuple(channels_mult)
+        self.down = nn.ModuleList()
+        block_in = channels
+        for i_level in range(self.num_resolutions):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_in = channels * in_ch_mult[i_level]
+            block_out = channels * channels_mult[i_level]
+            for _ in range(self.num_res_blocks):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            down = nn.Module()
+            down.block = block
+            down.attn = attn
+            if i_level < self.num_downsamples:
+                down.downsample = Downsample(block_in)
+                curr_res //= 2
+            self.down.append(down)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout)
+        self.norm_out = normalize(block_in)
+        self.conv_out = nn.Conv2d(block_in, z_channels, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.patcher(x)
+        hs = [self.conv_in(x)]
+        for i_level in range(self.num_resolutions):
+            for i_block in range(self.num_res_blocks):
+                h = self.down[i_level].block[i_block](hs[-1])
+                if len(self.down[i_level].attn) > 0:
+                    h = self.down[i_level].attn[i_block](h)
+                hs.append(h)
+            if i_level < self.num_downsamples:
+                hs.append(self.down[i_level].downsample(hs[-1]))
+
+        h = hs[-1]
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        h = nonlinearity(self.norm_out(h))
+        return self.conv_out(h)
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        out_channels: int,
+        channels: int,
+        channels_mult: list[int],
+        num_res_blocks: int,
+        attn_resolutions: list[int],
+        dropout: float,
+        resolution: int,
+        z_channels: int,
+        spatial_compression: int,
+        **kwargs: object,
+    ) -> None:
+        super().__init__()
+        self.num_resolutions = len(channels_mult)
+        self.num_res_blocks = num_res_blocks
+
+        patch_size = int(kwargs.get("patch_size", 1))
+        self.unpatcher = UnPatcher(patch_size, str(kwargs.get("patch_method", "rearrange")))
+        out_ch = out_channels * patch_size * patch_size
+        self.num_upsamples = int(math.log2(spatial_compression)) - int(math.log2(patch_size))
+        assert self.num_upsamples <= self.num_resolutions
+
+        block_in = channels * channels_mult[self.num_resolutions - 1]
+        curr_res = (resolution // patch_size) // 2 ** (self.num_resolutions - 1)
+        self.z_shape = (1, z_channels, curr_res, curr_res)
+        self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1)
+
+        self.mid = nn.Module()
+        self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout)
+        self.mid.attn_1 = AttnBlock(block_in)
+        self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout)
+
+        self.up = nn.ModuleList()
+        for i_level in reversed(range(self.num_resolutions)):
+            block = nn.ModuleList()
+            attn = nn.ModuleList()
+            block_out = channels * channels_mult[i_level]
+            for _ in range(self.num_res_blocks + 1):
+                block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, dropout=dropout))
+                block_in = block_out
+                if curr_res in attn_resolutions:
+                    attn.append(AttnBlock(block_in))
+            up = nn.Module()
+            up.block = block
+            up.attn = attn
+            if i_level >= (self.num_resolutions - self.num_upsamples):
+                up.upsample = Upsample(block_in)
+                curr_res *= 2
+            self.up.insert(0, up)
+
+        self.norm_out = normalize(block_in)
+        self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        h = self.conv_in(z)
+        h = self.mid.block_1(h)
+        h = self.mid.attn_1(h)
+        h = self.mid.block_2(h)
+        for i_level in reversed(range(self.num_resolutions)):
+            for i_block in range(self.num_res_blocks + 1):
+                h = self.up[i_level].block[i_block](h)
+                if len(self.up[i_level].attn) > 0:
+                    h = self.up[i_level].attn[i_block](h)
+            if i_level >= (self.num_resolutions - self.num_upsamples):
+                h = self.up[i_level].upsample(h)
+        h = self.conv_out(nonlinearity(self.norm_out(h)))
+        return self.unpatcher(h)
+
+
+class ContinuousImageTokenizer(nn.Module):
+    def __init__(self, z_channels: int, z_factor: int, latent_channels: int, **kwargs: object) -> None:
+        super().__init__()
+        self.latent_channels = latent_channels
+        self.encoder = Encoder(z_channels=z_factor * z_channels, **kwargs)
+        self.decoder = Decoder(z_channels=z_channels, **kwargs)
+        self.quant_conv = nn.Conv2d(z_factor * z_channels, z_factor * latent_channels, 1)
+        self.post_quant_conv = nn.Conv2d(latent_channels, z_channels, 1)
+        self.distribution = IdentityDistribution()
+
+    def encoder_module(self) -> nn.Sequential:
+        return nn.Sequential(
+            OrderedDict(
+                [
+                    ("encoder", self.encoder),
+                    ("quant_conv", self.quant_conv),
+                    ("distribution", self.distribution),
+                ]
+            )
+        )
+
+    def decoder_module(self) -> nn.Sequential:
+        return nn.Sequential(
+            OrderedDict(
+                [
+                    ("post_quant_conv", self.post_quant_conv),
+                    ("decoder", self.decoder),
+                ]
+            )
+        )
+
+    def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
+        h = self.encoder(x)
+        moments = self.quant_conv(h)
+        return self.distribution(moments)
+
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        return self.decoder(self.post_quant_conv(z))
+
+
+def build_cosmos_ci_torch_model(spatial_compression: int) -> ContinuousImageTokenizer:
+    return ContinuousImageTokenizer(
+        attn_resolutions=[32],
+        channels=128,
+        channels_mult=[2, 4, 4],
+        dropout=0.0,
+        in_channels=3,
+        spatial_compression=spatial_compression,
+        num_res_blocks=2,
+        out_channels=3,
+        resolution=1024,
+        patch_size=4,
+        patch_method="haar",
+        latent_channels=16,
+        z_channels=16,
+        z_factor=1,
+        name="CI",
+    )

vllm_omni/diffusion/models/internvla_a1/model_cosmos.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+from .cosmos_ci_torch import build_cosmos_ci_torch_model
+
+_SAFETENSORS_EXTENSION = ".safetensors"
+
+
+def is_safetensors_checkpoint(path: str | Path) -> bool:
+    return Path(path).suffix.lower() == _SAFETENSORS_EXTENSION
+
+
+def infer_cosmos_ci_spatial_compression(path: str | Path) -> int:
+    normalized = str(path).upper()
+    if "CI8X8" in normalized or "8X8" in normalized:
+        return 8
+    if "CI16X16" in normalized or "16X16" in normalized:
+        return 16
+    raise ValueError(f"Unable to infer Cosmos CI spatial compression from checkpoint path: {path}")
+
+
+def _load_safetensors_checkpoint(path: str | Path, device: str) -> dict[str, torch.Tensor]:
+    return load_file(str(path), device=device)
+
+
+def _load_native_cosmos_component(
+    checkpoint_filepath: str | Path,
+    *,
+    component: str,
+    device: str,
+) -> torch.nn.Module:
+    full_model = build_cosmos_ci_torch_model(
+        spatial_compression=infer_cosmos_ci_spatial_compression(checkpoint_filepath)
+    )
+    if component == "encoder":
+        model = full_model.encoder_module()
+    elif component == "decoder":
+        model = full_model.decoder_module()
+    else:
+        raise ValueError(f"Unsupported Cosmos component: {component}")
+
+    if not is_safetensors_checkpoint(checkpoint_filepath):
+        raise ValueError(f"Native Cosmos torch backend expects `.safetensors` checkpoints. Got: {checkpoint_filepath}")
+    state_dict = _load_safetensors_checkpoint(checkpoint_filepath, device)
+    model.load_state_dict(state_dict, strict=False)
+    return model.eval().to(device)
+
+
+def load_cosmos_component(
+    checkpoint_filepath: str | Path,
+    *,
+    component: str,
+    device: str = "cuda",
+) -> torch.nn.Module:
+    return _load_native_cosmos_component(
+        checkpoint_filepath,
+        component=component,
+        device=device,
+    )
+
+
+class ImageTokenizer(torch.nn.Module):
+    def __init__(
+        self,
+        checkpoint_enc: str,
+        checkpoint_dec: str,
+        device: str = "cuda",
+    ) -> None:
+        super().__init__()
+        self._enc_model = load_cosmos_component(
+            checkpoint_enc,
+            component="encoder",
+            device=device,
+        )
+        self._dec_model = load_cosmos_component(
+            checkpoint_dec,
+            component="decoder",
+            device=device,
+        )
+
+    def decode(self, input_latent: torch.Tensor) -> torch.Tensor:
+        original_dtype = input_latent.dtype
+        model_dtype = next(self.parameters()).dtype
+        output_tensor = self._dec_model(input_latent.to(model_dtype))
+        return output_tensor.to(original_dtype)
+
+    def encode(self, input_tensor: torch.Tensor) -> torch.Tensor:
+        original_dtype = input_tensor.dtype
+        model_dtype = next(self.parameters()).dtype
+        output_latent = self._enc_model(input_tensor.to(model_dtype))
+        if isinstance(output_latent, torch.Tensor):
+            return output_latent.to(original_dtype)
+        return output_latent[0].to(original_dtype)

vllm_omni/diffusion/models/internvla_a1/model_internvla_a1.py

@@ -0,0 +1,962 @@
+from __future__ import annotations
+
+import math
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal
+
+import torch
+import torch._dynamo as dynamo
+import torch.nn.functional as F
+from einops import rearrange
+from safetensors import safe_open
+from torch import Tensor, nn
+from transformers import Qwen3VLProcessor
+from transformers.models.auto import CONFIG_MAPPING
+from vllm.model_executor.models.utils import AutoWeightsLoader
+
+from .adapter_qwen3_vl import (
+    Qwen3VLForConditionalGeneration,
+    Qwen3VLTextModel,
+    apply_rotary_pos_emb,
+    eager_attention_forward,
+)
+from .config import (
+    DEFAULT_COSMOS_DIR,
+    DEFAULT_COSMOS_REPO,
+    DEFAULT_QWEN3_VL_MODEL,
+    OBS_IMAGES,
+    OBS_PREFIX,
+    OBS_STATE,
+    OBS_TASK,
+    OPENPI_ATTENTION_MASK_VALUE,
+    InternVLAA1Config,
+)
+from .model_cosmos import ImageTokenizer
+
+
+def get_safe_dtype(target_dtype: torch.dtype, device_type: str) -> torch.dtype:
+    if device_type == "mps" and target_dtype == torch.float64:
+        return torch.float32
+    if device_type == "cpu" and target_dtype == torch.bfloat16:
+        return torch.float32
+    return target_dtype
+
+
+def create_sinusoidal_pos_embedding(
+    time: torch.Tensor,
+    dimension: int,
+    min_period: float,
+    max_period: float,
+    device: torch.device,
+) -> Tensor:
+    if dimension % 2 != 0:
+        raise ValueError(f"dimension ({dimension}) must be divisible by 2")
+    dtype = get_safe_dtype(torch.float64, device.type)
+    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)
+    period = min_period * (max_period / min_period) ** fraction
+    scaling_factor = 1.0 / period * 2 * math.pi
+    sin_input = scaling_factor[None, :] * time[:, None]
+    return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)
+
+
+def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor:
+    cumsum = torch.cumsum(att_masks, dim=1)
+    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]
+    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]
+    return att_2d_masks & pad_2d_masks
+
+
+def pad_vector(vector: torch.Tensor, new_dim: int) -> torch.Tensor:
+    if vector.shape[-1] >= new_dim:
+        return vector
+    return F.pad(vector, (0, new_dim - vector.shape[-1]))
+
+
+def resize_with_pad(images: torch.Tensor, size: tuple[int, int]) -> torch.Tensor:
+    target_h, target_w = size
+    if images.ndim != 4:
+        raise ValueError(f"Expected [T, C, H, W], got {tuple(images.shape)}")
+    _, _, src_h, src_w = images.shape
+    scale = min(target_h / src_h, target_w / src_w)
+    resized_h = max(1, int(round(src_h * scale)))
+    resized_w = max(1, int(round(src_w * scale)))
+    resized = F.interpolate(images, size=(resized_h, resized_w), mode="bilinear", align_corners=False)
+    pad_h = target_h - resized_h
+    pad_w = target_w - resized_w
+    pad_top = pad_h // 2
+    pad_bottom = pad_h - pad_top
+    pad_left = pad_w // 2
+    pad_right = pad_w - pad_left
+    return F.pad(resized, (pad_left, pad_right, pad_top, pad_bottom))
+
+
+class Qwen3VLInputBuilder:
+    def __init__(
+        self,
+        processor_model_name: str = DEFAULT_QWEN3_VL_MODEL,
+        max_length: int = 48,
+        spatial_merge_size: int = 2,
+    ) -> None:
+        processor_model_name = os.getenv("INTERNVLA_A1_PROCESSOR_DIR", processor_model_name)
+        local_path = Path(processor_model_name).expanduser()
+        if local_path.exists():
+            self.processor = Qwen3VLProcessor.from_pretrained(str(local_path), local_files_only=True)
+        else:
+            if os.getenv("INTERNVLA_A1_PROCESSOR_DIR"):
+                raise FileNotFoundError(f"INTERNVLA_A1_PROCESSOR_DIR points to a missing path: {local_path}")
+            self.processor = Qwen3VLProcessor.from_pretrained(processor_model_name)
+        self.max_length = max_length
+        self.spatial_merge_size = spatial_merge_size
+
+    def build(self, camera_images: list[torch.Tensor], task: str) -> dict[str, torch.Tensor]:
+        input_ids: list[int] = []
+        attention_mask: list[int] = []
+        pixel_values: list[torch.Tensor] = []
+        image_grid_thw: list[torch.Tensor] = []
+
+        for image_history in camera_images:
+            current_image = image_history[1]
+            image_inputs = self.processor.image_processor(current_image, do_rescale=False)
+            pixel_values.append(image_inputs.pixel_values)
+            image_grid_thw.append(image_inputs.image_grid_thw)
+            num_img_token = int(torch.prod(image_inputs.image_grid_thw[0]).item()) // (self.spatial_merge_size**2)
+            input_ids.extend(
+                [self.processor.vision_start_token_id]
+                + [self.processor.image_token_id] * num_img_token
+                + [self.processor.vision_end_token_id]
+            )
+            attention_mask.extend([1] * (num_img_token + 2))
+
+        text_inputs = self.processor.tokenizer(
+            task,
+            max_length=self.max_length,
+            padding="max_length",
+            padding_side="right",
+            truncation=True,
+        )
+        input_ids.extend(text_inputs.input_ids)
+        attention_mask.extend(text_inputs.attention_mask)
+
+        return {
+            f"{OBS_PREFIX}pixel_values": torch.cat(pixel_values, dim=0),
+            f"{OBS_PREFIX}image_grid_thw": torch.cat(image_grid_thw, dim=0),
+            f"{OBS_PREFIX}input_ids": torch.tensor(input_ids, dtype=torch.long),
+            f"{OBS_PREFIX}attention_mask": torch.tensor(attention_mask, dtype=torch.long),
+        }
+
+
+@dataclass
+class SuffixStaticContext:
+    state_emb: torch.Tensor
+    full_att_2d_masks_4d: torch.Tensor
+    position_ids: torch.Tensor
+
+
+def compute_layer_complete(
+    layer_idx: int,
+    inputs_embeds: list[torch.Tensor],
+    attention_mask: torch.Tensor,
+    position_ids: torch.Tensor,
+    und_expert: nn.Module,
+    gen_expert: nn.Module,
+    act_expert: nn.Module,
+) -> list[torch.Tensor]:
+    models = [und_expert.language_model, gen_expert, act_expert]
+    query_states = []
+    key_states = []
+    value_states = []
+    for i, hidden_states in enumerate(inputs_embeds):
+        layer = models[i].layers[layer_idx]
+        hidden_states = layer.input_layernorm(hidden_states)
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)
+        if layer.self_attn.q_proj.weight.dtype == torch.bfloat16:
+            hidden_states = hidden_states.to(dtype=torch.bfloat16)
+        query_state = layer.self_attn.q_norm(layer.self_attn.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_state = layer.self_attn.k_norm(layer.self_attn.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        query_states.append(query_state)
+        key_states.append(key_state)
+        value_states.append(value_state)
+
+    query_states = torch.cat(query_states, dim=2)
+    key_states = torch.cat(key_states, dim=2)
+    value_states = torch.cat(value_states, dim=2)
+    dummy_tensor = torch.zeros(
+        query_states.shape[0],
+        query_states.shape[2],
+        query_states.shape[-1],
+        device=query_states.device,
+        dtype=query_states.dtype,
+    )
+    cos, sin = und_expert.model.language_model.rotary_emb(dummy_tensor, position_ids)
+    query_states, key_states = apply_rotary_pos_emb(
+        query_states,
+        key_states,
+        cos,
+        sin,
+        unsqueeze_dim=1,
+    )
+    scaling = und_expert.language_model.layers[layer_idx].self_attn.scaling
+    att_output, _ = eager_attention_forward(
+        und_expert.language_model.layers[layer_idx].self_attn,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        scaling,
+    )
+
+    head_dim = und_expert.language_model.layers[layer_idx].self_attn.head_dim
+    num_attention_heads = und_expert.language_model.layers[layer_idx].self_attn.config.num_attention_heads
+    batch_size = query_states.shape[0]
+    att_output = att_output.reshape(batch_size, -1, num_attention_heads * head_dim)
+
+    outputs_embeds = []
+    start_pos = 0
+    for i, hidden_states in enumerate(inputs_embeds):
+        layer = models[i].layers[layer_idx]
+        end_pos = start_pos + hidden_states.shape[1]
+        att_chunk = att_output[:, start_pos:end_pos]
+        if att_chunk.dtype != layer.self_attn.o_proj.weight.dtype:
+            att_chunk = att_chunk.to(layer.self_attn.o_proj.weight.dtype)
+        out_emb = layer.self_attn.o_proj(att_chunk)
+        out_emb = out_emb + hidden_states
+        residual = out_emb
+        out_emb = layer.post_attention_layernorm(out_emb)
+        if layer.mlp.up_proj.weight.dtype == torch.bfloat16:
+            out_emb = out_emb.to(dtype=torch.bfloat16)
+        out_emb = layer.mlp(out_emb)
+        out_emb = out_emb + residual
+        outputs_embeds.append(out_emb)
+        start_pos = end_pos
+    return outputs_embeds
+
+
+class QwenConfig:
+    def __init__(
+        self,
+        head_dim: int,
+        hidden_size: int,
+        intermediate_size: int,
+        num_attention_heads: int,
+        num_hidden_layers: int,
+        num_key_value_heads: int,
+    ) -> None:
+        self.head_dim = head_dim
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_attention_heads = num_attention_heads
+        self.num_hidden_layers = num_hidden_layers
+        self.num_key_value_heads = num_key_value_heads
+
+
+def get_qwen_config(variant: str) -> QwenConfig:
+    num_hidden_layers = int(variant.split("_")[-1][:-1])
+    if variant.startswith("qwen3_vl"):
+        return QwenConfig(128, 2048, 6144, 16, num_hidden_layers, 8)
+    if variant.startswith("qwen3"):
+        return QwenConfig(128, 1024, 3072, 16, num_hidden_layers, 8)
+    raise ValueError(f"Unknown variant: {variant}")
+
+
+def resolve_cosmos_checkpoint_paths(
+    *,
+    encoder_path: str | Path | None = None,
+    decoder_path: str | Path | None = None,
+) -> tuple[Path, Path]:
+    encoder_override = encoder_path or os.getenv("INTERNVLA_A1_COSMOS_ENCODER_PATH")
+    decoder_override = decoder_path or os.getenv("INTERNVLA_A1_COSMOS_DECODER_PATH")
+    default_encoder_name = "encoder.safetensors"
+    default_decoder_name = "decoder.safetensors"
+    checkpoint_enc = (
+        Path(encoder_override).expanduser()
+        if encoder_override is not None
+        else DEFAULT_COSMOS_DIR / default_encoder_name
+    )
+    checkpoint_dec = (
+        Path(decoder_override).expanduser()
+        if decoder_override is not None
+        else DEFAULT_COSMOS_DIR / default_decoder_name
+    )
+
+    missing = [str(path) for path in (checkpoint_enc, checkpoint_dec) if not path.exists()]
+    if missing:
+        missing_display = ", ".join(missing)
+        raise FileNotFoundError(
+            "InternVLA-A1 requires local Cosmos tokenizer checkpoints. "
+            f"Missing: {missing_display}. "
+            f"Download {DEFAULT_COSMOS_REPO} and set INTERNVLA_A1_COSMOS_DIR to that directory, "
+            "or point INTERNVLA_A1_COSMOS_ENCODER_PATH / INTERNVLA_A1_COSMOS_DECODER_PATH "
+            "to explicit checkpoint files."
+        )
+
+    return checkpoint_enc, checkpoint_dec
+
+
+class Qwen3VLWithExpertModel(nn.Module):
+    def __init__(
+        self,
+        vlm_config: QwenConfig,
+        action_expert_config: QwenConfig,
+        precision: Literal["bfloat16", "float32"] = "bfloat16",
+    ) -> None:
+        super().__init__()
+
+        vlm_config_hf = CONFIG_MAPPING["qwen3_vl"]()
+        vlm_config_hf.text_config.hidden_size = vlm_config.hidden_size
+        vlm_config_hf.text_config.intermediate_size = vlm_config.intermediate_size
+        vlm_config_hf.text_config.num_attention_heads = vlm_config.num_attention_heads
+        vlm_config_hf.text_config.head_dim = vlm_config.head_dim
+        vlm_config_hf.text_config.num_hidden_layers = vlm_config.num_hidden_layers
+        vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_key_value_heads
+        vlm_config_hf.text_config.max_position_embeddings = 262144
+        vlm_config_hf.text_config.rope_scaling = {
+            "mrope_interleaved": True,
+            "mrope_section": [24, 20, 20],
+            "rope_type": "default",
+        }
+        vlm_config_hf.text_config.tie_word_embeddings = True
+        vlm_config_hf.tie_word_embeddings = True
+        vlm_config_hf.vision_config.deepstack_visual_indexes = [5, 11, 17]
+        vlm_config_hf.vision_config.depth = 24
+        vlm_config_hf.vision_config.hidden_size = 1024
+        vlm_config_hf.vision_config.intermediate_size = 4096
+        vlm_config_hf.vision_config.out_hidden_size = 2048
+
+        self.und_expert = Qwen3VLForConditionalGeneration(config=vlm_config_hf)
+
+        gen_expert_config_hf = CONFIG_MAPPING["qwen3_vl_text"]()
+        gen_expert_config_hf.head_dim = action_expert_config.head_dim
+        gen_expert_config_hf.hidden_size = action_expert_config.hidden_size
+        gen_expert_config_hf.intermediate_size = action_expert_config.intermediate_size
+        gen_expert_config_hf.num_attention_heads = action_expert_config.num_attention_heads
+        gen_expert_config_hf.num_hidden_layers = action_expert_config.num_hidden_layers
+        gen_expert_config_hf.num_key_value_heads = action_expert_config.num_key_value_heads
+        gen_expert_config_hf.max_position_embeddings = self.und_expert.config.text_config.max_position_embeddings
+        gen_expert_config_hf.rope_scaling = self.und_expert.config.text_config.rope_scaling
+        self.gen_expert = Qwen3VLTextModel(config=gen_expert_config_hf)
+        self.gen_expert.embed_tokens = None
+        self.gen_expert.lm_head = None
+
+        act_expert_config_hf = CONFIG_MAPPING["qwen3_vl_text"]()
+        act_expert_config_hf.head_dim = action_expert_config.head_dim
+        act_expert_config_hf.hidden_size = action_expert_config.hidden_size
+        act_expert_config_hf.intermediate_size = action_expert_config.intermediate_size
+        act_expert_config_hf.num_attention_heads = action_expert_config.num_attention_heads
+        act_expert_config_hf.num_hidden_layers = action_expert_config.num_hidden_layers
+        act_expert_config_hf.num_key_value_heads = action_expert_config.num_key_value_heads
+        act_expert_config_hf.max_position_embeddings = self.und_expert.config.text_config.max_position_embeddings
+        act_expert_config_hf.rope_scaling = self.und_expert.config.text_config.rope_scaling
+        self.act_expert = Qwen3VLTextModel(config=act_expert_config_hf)
+        self.act_expert.embed_tokens = None
+        self.act_expert.lm_head = None
+
+        self.to_selected_precision(precision)
+
+    def to_selected_precision(self, precision: Literal["bfloat16", "float32"]) -> None:
+        if precision == "bfloat16":
+            self.to(dtype=torch.bfloat16)
+        elif precision == "float32":
+            # The fp32 path keeps the whole module in full precision, so no
+            # selective post-conversion fixups are required.
+            self.to(dtype=torch.float32)
+            return
+        else:
+            raise ValueError(f"Unsupported precision: {precision}")
+
+        # Match the original runtime numerics by keeping normalization layers
+        # in fp32 even when the rest of the module runs in bf16.
+        keep_fp32 = ["input_layernorm", "post_attention_layernorm", "model.norm"]
+        for name, param in self.named_parameters():
+            if any(key in name for key in keep_fp32):
+                param.data = param.data.to(dtype=torch.float32)
+
+    def forward(
+        self,
+        attention_mask: torch.Tensor | None,
+        position_ids: torch.LongTensor | None,
+        past_key_values: Any,
+        inputs_embeds: list[torch.Tensor | None],
+        use_cache: bool,
+    ) -> tuple[list[torch.Tensor | None], Any]:
+        if inputs_embeds[1] is None and inputs_embeds[2] is None:
+            prefix_output = self.und_expert.language_model.forward(
+                inputs_embeds=inputs_embeds[0],
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+            return [prefix_output.last_hidden_state, None, None], prefix_output.past_key_values
+
+        if inputs_embeds[0] is None and inputs_embeds[2] is None:
+            middle_output = self.gen_expert.forward(
+                inputs_embeds=inputs_embeds[1],
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+            return [None, middle_output.last_hidden_state, None], middle_output.past_key_values
+
+        if inputs_embeds[0] is None and inputs_embeds[1] is None:
+            suffix_output = self.act_expert.forward(
+                inputs_embeds=inputs_embeds[2],
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+            )
+            return [None, None, suffix_output.last_hidden_state], None
+
+        models = [self.und_expert.language_model, self.gen_expert, self.act_expert]
+        stacked_inputs = [inputs_embeds[0], inputs_embeds[1], inputs_embeds[2]]
+        for layer_idx in range(self.und_expert.config.text_config.num_hidden_layers):
+            stacked_inputs = compute_layer_complete(
+                layer_idx,
+                stacked_inputs,
+                attention_mask,
+                position_ids,
+                und_expert=self.und_expert,
+                gen_expert=self.gen_expert,
+                act_expert=self.act_expert,
+            )
+        outputs = [models[i].norm(stacked_inputs[i]) for i in range(3)]
+        return outputs, None
+
+
+class InternVLAA1(nn.Module):
+    def __init__(
+        self,
+        config: InternVLAA1Config,
+        *,
+        cosmos_encoder_path: str | Path | None = None,
+        cosmos_decoder_path: str | Path | None = None,
+    ) -> None:
+        super().__init__()
+        self.config = config
+
+        vlm_config = get_qwen_config(config.qwen3_vl_variant)
+        action_expert_config = get_qwen_config(config.action_expert_variant)
+        self.qwen3_vl_with_expert = Qwen3VLWithExpertModel(vlm_config, action_expert_config, precision=config.dtype)
+
+        cosmos_encoder_path, cosmos_decoder_path = resolve_cosmos_checkpoint_paths(
+            encoder_path=cosmos_encoder_path,
+            decoder_path=cosmos_decoder_path,
+        )
+        self.cosmos = ImageTokenizer(
+            checkpoint_enc=str(cosmos_encoder_path),
+            checkpoint_dec=str(cosmos_decoder_path),
+            device=config.device,
+        )
+
+        hidden_size = action_expert_config.hidden_size
+        vae_dim = 16
+        ds = config.scale_factor
+        self.cosmos_in_proj = nn.Conv2d(vae_dim, hidden_size, kernel_size=1, stride=1, padding=0)
+        self.downsample_conv = nn.Conv2d(hidden_size, hidden_size, kernel_size=ds, stride=ds, padding=0)
+        self.upsample_conv = nn.ConvTranspose2d(hidden_size, hidden_size, kernel_size=ds, stride=ds, padding=0)
+        self.cosmos_out_proj = nn.Linear(hidden_size, vae_dim)
+        self.cosmos_out_layer_norm = nn.LayerNorm(hidden_size)
+
+        self.action_in_proj = nn.Linear(config.max_action_dim, hidden_size)
+        self.action_out_proj = nn.Linear(hidden_size, config.max_action_dim)
+        self.state_proj = nn.Linear(config.max_state_dim, hidden_size)
+        self.action_time_mlp_in = nn.Linear(2 * hidden_size, hidden_size)
+        self.action_time_mlp_out = nn.Linear(hidden_size, hidden_size)
+
+        if config.compile_model:
+            torch.set_float32_matmul_precision("high")
+            self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode)
+
+        self.cosmos.eval()
+        for param in self.cosmos.parameters():
+            param.requires_grad = False
+
+        self.set_attention_implementation(config.attn_implementation)
+
+    def set_attention_implementation(self, attn_implementation: str) -> None:
+        self.config.attn_implementation = attn_implementation
+        self.qwen3_vl_with_expert.und_expert.config.text_config._attn_implementation = attn_implementation
+        self.qwen3_vl_with_expert.und_expert.language_model.config._attn_implementation = attn_implementation
+        self.qwen3_vl_with_expert.gen_expert.config._attn_implementation = attn_implementation
+        self.qwen3_vl_with_expert.act_expert.config._attn_implementation = attn_implementation
+
+    def _prepare_attention_masks_4d(
+        self,
+        att_2d_masks: torch.Tensor,
+        *,
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        att_2d_masks_4d = att_2d_masks[:, None, :, :]
+        bias_dtype = torch.float32 if self.config.attn_implementation == "eager" else dtype
+        valid_bias = torch.zeros((), dtype=bias_dtype, device=att_2d_masks.device)
+        invalid_bias = torch.tensor(
+            OPENPI_ATTENTION_MASK_VALUE,
+            dtype=bias_dtype,
+            device=att_2d_masks.device,
+        )
+        return torch.where(att_2d_masks_4d, valid_bias, invalid_bias)
+
+    def sample_noise(self, shape: tuple[int, ...], device: torch.device) -> torch.Tensor:
+        return torch.normal(0.0, 1.0, shape, dtype=torch.float32, device=device)
+
+    @dynamo.disable
+    def embed_prefix(
+        self,
+        pixel_values: torch.Tensor,
+        image_grid_thw: torch.Tensor,
+        lang_tokens: torch.Tensor,
+        lang_masks: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        image_token_id = self.qwen3_vl_with_expert.und_expert.config.image_token_id
+        pixel_values = pixel_values.view(-1, pixel_values.shape[-1])
+        image_grid_thw = image_grid_thw.view(-1, 3)
+        image_embs, _ = self.qwen3_vl_with_expert.und_expert.visual(pixel_values, image_grid_thw)
+
+        embs = self.qwen3_vl_with_expert.und_expert.get_input_embeddings()(lang_tokens)
+        batch_size, seq_len, hidden_dim = embs.shape
+        embs = embs.view(-1, hidden_dim)
+        lang_tokens_flat = lang_tokens.view(-1)
+        embs[lang_tokens_flat == image_token_id] = image_embs
+        embs = embs.view(batch_size, seq_len, hidden_dim)
+
+        pad_masks = lang_masks.to(torch.bool)
+        att_masks = torch.zeros_like(pad_masks, dtype=torch.bool, device=pad_masks.device)
+        return embs, pad_masks, att_masks
+
+    def get_cosmos_features(self, images: torch.Tensor) -> torch.Tensor:
+        shape = images.shape[:-3]
+        channels, height, width = images.shape[-3:]
+        images = images.reshape(-1, channels, height, width)
+        if (height, width) != (256, 256):
+            images = F.interpolate(images, size=(256, 256), mode="bilinear", align_corners=False)
+        images = images * 2 - 1
+        features = self.cosmos.encode(images)
+        channels, height, width = features.shape[-3:]
+        return features.view(*shape, channels, height, width)
+
+    def embed_middle(
+        self,
+        images: torch.Tensor,
+        img_masks: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        device = images.device
+        batch_size, n_view, timesteps = images.shape[:3]
+        features = self.get_cosmos_features(images)
+        features = rearrange(features, "b n t c h w -> (b n t) c h w")
+        features = self.cosmos_in_proj(features)
+        features = self.downsample_conv(features)
+        features = rearrange(features, "(b n t) c h w -> b n t c h w", b=batch_size, n=n_view, t=timesteps)
+
+        _, _, _, _, height, width = features.shape
+        embs = rearrange(features, "b n t c h w -> b (n t h w) c")
+        pad_masks = torch.zeros((batch_size, n_view, timesteps, height, width), dtype=torch.bool, device=device)
+        pad_masks[img_masks] = True
+        pad_masks = rearrange(pad_masks, "b n t h w -> b (n t h w)")
+        att_masks = torch.tensor([1] + [0] * (embs.shape[1] - 1), dtype=torch.bool, device=device)
+        att_masks = att_masks[None, :].expand(batch_size, att_masks.shape[0])
+        return embs, pad_masks, att_masks
+
+    def prepare_suffix_static_context(
+        self,
+        state: torch.Tensor,
+        prefix_pad_masks: torch.Tensor,
+        max_prefix_position_ids: torch.Tensor,
+    ) -> SuffixStaticContext:
+        if self.state_proj.weight.dtype == torch.float32:
+            state = state.to(torch.float32)
+
+        state_emb = self.state_proj(state)
+        device = state_emb.device
+        batch_size = state_emb.shape[0]
+        suffix_len = self.config.chunk_size + 1
+        prefix_len = prefix_pad_masks.shape[1]
+
+        suffix_pad_masks = torch.ones(batch_size, suffix_len, dtype=torch.bool, device=device)
+        suffix_att_masks = torch.tensor([1] + [1] + [0] * (self.config.chunk_size - 1), dtype=torch.bool, device=device)
+        suffix_att_masks = suffix_att_masks[None, :].expand(batch_size, suffix_len)
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
+        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+        full_att_2d_masks_4d = self._prepare_attention_masks_4d(
+            torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2),
+            dtype=state_emb.dtype,
+        )
+        position_ids = (
+            torch.arange(1, suffix_len + 1, device=max_prefix_position_ids.device)
+            .repeat(3, 1, 1)
+            .to(max_prefix_position_ids)
+            + max_prefix_position_ids
+        )
+        return SuffixStaticContext(
+            state_emb=state_emb,
+            full_att_2d_masks_4d=full_att_2d_masks_4d,
+            position_ids=position_ids,
+        )
+
+    def embed_suffix(
+        self,
+        state: torch.Tensor,
+        noisy_actions: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self.state_proj.weight.dtype == torch.float32:
+            state = state.to(torch.float32)
+
+        state_emb = self.state_proj(state)
+        batch_size = state_emb.shape[0]
+        device = state_emb.device
+        state_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=device)
+
+        time_emb = create_sinusoidal_pos_embedding(
+            timestep,
+            self.action_in_proj.out_features,
+            min_period=self.config.min_period,
+            max_period=self.config.max_period,
+            device=timestep.device,
+        ).to(dtype=timestep.dtype)
+        action_emb = self.action_in_proj(noisy_actions)
+        time_emb = time_emb[:, None, :].expand_as(action_emb)
+        action_time_emb = torch.cat([action_emb, time_emb], dim=2)
+        action_time_emb = self.action_time_mlp_out(F.silu(self.action_time_mlp_in(action_time_emb)))
+        action_mask = torch.ones(batch_size, action_time_emb.shape[1], dtype=torch.bool, device=device)
+
+        embs = torch.cat([state_emb[:, None, :], action_time_emb], dim=1)
+        pad_masks = torch.cat([state_mask, action_mask], dim=1)
+        att_masks = torch.tensor([1] + [1] + [0] * (self.config.chunk_size - 1), dtype=embs.dtype, device=device)
+        att_masks = att_masks[None, :].expand(batch_size, self.config.chunk_size + 1)
+        return embs, pad_masks, att_masks
+
+    def get_position_ids(
+        self,
+        lang_tokens: torch.Tensor,
+        image_grid_thw: torch.Tensor | None,
+        pad_masks: torch.Tensor,
+    ) -> tuple[torch.Tensor, Any]:
+        seq_len = lang_tokens.shape[1]
+        padded_lang_tokens = torch.ones_like(pad_masks).to(lang_tokens) * 777
+        padded_lang_tokens[:, :seq_len] = lang_tokens
+        attention_mask = pad_masks.to(lang_tokens)
+        if image_grid_thw is not None:
+            image_grid_thw = image_grid_thw.view(-1, 3)
+        return self.qwen3_vl_with_expert.und_expert.model.get_rope_index(
+            padded_lang_tokens,
+            image_grid_thw,
+            attention_mask=attention_mask,
+        )
+
+    @torch.no_grad()
+    def sample_actions(
+        self,
+        images: torch.Tensor,
+        img_masks: torch.Tensor,
+        pixel_values: torch.Tensor,
+        image_grid_thw: torch.Tensor,
+        lang_tokens: torch.Tensor,
+        lang_masks: torch.Tensor,
+        state: torch.Tensor,
+        *,
+        noise: torch.Tensor | None = None,
+        num_steps: int | None = None,
+        decode_image: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if num_steps is None:
+            num_steps = self.config.num_inference_steps
+        batch_size = state.shape[0]
+        device = state.device
+        dtype = state.dtype
+        if noise is None:
+            noise = self.sample_noise((batch_size, self.config.chunk_size, self.config.max_action_dim), device)
+
+        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(
+            pixel_values,
+            image_grid_thw,
+            lang_tokens,
+            lang_masks,
+        )
+        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)
+        prefix_position_ids, _ = self.get_position_ids(lang_tokens, image_grid_thw, prefix_pad_masks)
+        prefix_att_2d_masks_4d = self._prepare_attention_masks_4d(
+            prefix_att_2d_masks,
+            dtype=prefix_embs.dtype,
+        )
+        _, past_key_values = self.qwen3_vl_with_expert.forward(
+            attention_mask=prefix_att_2d_masks_4d,
+            position_ids=prefix_position_ids,
+            past_key_values=None,
+            inputs_embeds=[prefix_embs, None, None],
+            use_cache=True,
+        )
+        max_prefix_position_ids = prefix_position_ids.max(dim=-1, keepdim=True).values
+
+        middle_embs, middle_pad_masks, middle_att_masks = self.embed_middle(images[:, :, :2], img_masks)
+        middle_len = middle_pad_masks.shape[1]
+        prefix_len = prefix_pad_masks.shape[1]
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, middle_len, prefix_len)
+        middle_att_2d_masks = make_att_2d_masks(middle_pad_masks, middle_att_masks)
+        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, middle_att_2d_masks], dim=2)
+        middle_position_ids = (
+            torch.arange(1, middle_len + 1, device=max_prefix_position_ids.device)
+            .repeat(3, 1, 1)
+            .to(max_prefix_position_ids)
+            + max_prefix_position_ids
+        )
+        full_att_2d_masks_4d = self._prepare_attention_masks_4d(
+            full_att_2d_masks,
+            dtype=middle_embs.dtype,
+        )
+        (_, middle_out, _), past_key_values = self.qwen3_vl_with_expert.forward(
+            attention_mask=full_att_2d_masks_4d,
+            position_ids=middle_position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=[None, middle_embs, None],
+            use_cache=True,
+        )
+
+        max_position_ids = middle_position_ids.max(dim=-1, keepdim=True).values
+        curr_pad_masks = torch.cat([prefix_pad_masks, middle_pad_masks], dim=1)
+        dt = torch.tensor(-1.0 / num_steps, dtype=torch.float32, device=device)
+        x_t = noise
+        time = torch.tensor(1.0, dtype=torch.float32, device=device)
+        suffix_static = None
+        if self.config.enable_suffix_static_context_optimization:
+            suffix_static = self.prepare_suffix_static_context(state, curr_pad_masks, max_position_ids)
+        while time >= -dt / 2:
+            expanded_time = time.expand(batch_size)
+            if suffix_static is None:
+                v_t = self.denoise_step(
+                    state,
+                    curr_pad_masks,
+                    past_key_values,
+                    max_position_ids,
+                    x_t.to(dtype),
+                    expanded_time.to(dtype),
+                )
+            else:
+                v_t = self.denoise_step_optimized(
+                    suffix_static, past_key_values, x_t.to(dtype), expanded_time.to(dtype)
+                )
+            x_t = x_t + dt * v_t
+            time += dt
+
+        return x_t, None if not decode_image else middle_out
+
+    def denoise_step(
+        self,
+        state: torch.Tensor,
+        prefix_pad_masks: torch.Tensor,
+        past_key_values: Any,
+        max_prefix_position_ids: torch.Tensor,
+        x_t: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(state, x_t, timestep)
+        suffix_len = suffix_pad_masks.shape[1]
+        batch_size = prefix_pad_masks.shape[0]
+        prefix_len = prefix_pad_masks.shape[1]
+        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)
+        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)
+        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)
+        position_ids = (
+            torch.arange(1, suffix_len + 1, device=max_prefix_position_ids.device)
+            .repeat(3, 1, 1)
+            .to(max_prefix_position_ids)
+            + max_prefix_position_ids
+        )
+        full_att_2d_masks_4d = self._prepare_attention_masks_4d(
+            full_att_2d_masks,
+            dtype=suffix_embs.dtype,
+        )
+        outputs_embeds, _ = self.qwen3_vl_with_expert.forward(
+            attention_mask=full_att_2d_masks_4d,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=[None, None, suffix_embs],
+            use_cache=False,
+        )
+        suffix_out = outputs_embeds[2][:, -self.config.chunk_size :].to(dtype=torch.float32)
+        return self.action_out_proj(suffix_out)
+
+    def denoise_step_optimized(
+        self,
+        suffix_static: SuffixStaticContext,
+        past_key_values: Any,
+        x_t: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        state_emb = suffix_static.state_emb
+        time_emb = create_sinusoidal_pos_embedding(
+            timestep,
+            self.action_in_proj.out_features,
+            min_period=self.config.min_period,
+            max_period=self.config.max_period,
+            device=timestep.device,
+        ).to(dtype=timestep.dtype)
+        action_emb = self.action_in_proj(x_t)
+        time_emb = time_emb[:, None, :].expand_as(action_emb)
+        action_time_emb = torch.cat([action_emb, time_emb], dim=2)
+        action_time_emb = self.action_time_mlp_out(F.silu(self.action_time_mlp_in(action_time_emb)))
+        suffix_embs = torch.cat([state_emb[:, None, :], action_time_emb], dim=1)
+        outputs_embeds, _ = self.qwen3_vl_with_expert.forward(
+            attention_mask=suffix_static.full_att_2d_masks_4d,
+            position_ids=suffix_static.position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=[None, None, suffix_embs],
+            use_cache=False,
+        )
+        suffix_out = outputs_embeds[2][:, -self.config.chunk_size :].to(dtype=torch.float32)
+        return self.action_out_proj(suffix_out)
+
+
+class InternVLAA1Policy(nn.Module):
+    _AUTO_WEIGHTS_IGNORE_PREFIXES = ["model.cosmos."]
+
+    def __init__(
+        self,
+        config: InternVLAA1Config,
+        *,
+        processor_model_name: str = DEFAULT_QWEN3_VL_MODEL,
+    ) -> None:
+        super().__init__()
+        self.config = config
+        self.input_builder = Qwen3VLInputBuilder(
+            processor_model_name=processor_model_name,
+            max_length=config.tokenizer_max_length,
+        )
+        cosmos_encoder_path, cosmos_decoder_path = resolve_cosmos_checkpoint_paths()
+        self.model = InternVLAA1(
+            config,
+            cosmos_encoder_path=cosmos_encoder_path,
+            cosmos_decoder_path=cosmos_decoder_path,
+        )
+        self.model.to(config.device)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        checkpoint_dir: str | Path,
+        *,
+        config: InternVLAA1Config | None = None,
+        processor_model_name: str = DEFAULT_QWEN3_VL_MODEL,
+        strict: bool = False,
+    ) -> InternVLAA1Policy:
+        checkpoint_dir = Path(checkpoint_dir)
+        if config is None:
+            config = InternVLAA1Config.from_pretrained(checkpoint_dir)
+        instance = cls(config, processor_model_name=processor_model_name)
+        checkpoint_path = checkpoint_dir / "model.safetensors"
+        if not checkpoint_path.exists():
+            raise FileNotFoundError(f"InternVLAA1 checkpoint not found: {checkpoint_path}")
+
+        loader = AutoWeightsLoader(
+            instance,
+            ignore_unexpected_prefixes=cls._AUTO_WEIGHTS_IGNORE_PREFIXES,
+        )
+        with safe_open(str(checkpoint_path), framework="pt", device=str(config.device)) as f:
+            file_keys = list(f.keys())
+            loaded = loader.load_weights((name, f.get_tensor(name)) for name in file_keys)
+
+        if strict:
+            expected = {
+                name
+                for name, _ in instance.state_dict().items()
+                if not any(name.startswith(prefix) for prefix in cls._AUTO_WEIGHTS_IGNORE_PREFIXES)
+            }
+            missing = sorted(expected - loaded)
+            if missing:
+                preview = ", ".join(missing[:10])
+                suffix = " ..." if len(missing) > 10 else ""
+                raise RuntimeError(f"Missing weights after AutoWeightsLoader load: {preview}{suffix}")
+
+        instance.eval()
+        return instance
+
+    def to(self, *args, **kwargs):  # type: ignore[override]
+        super().to(*args, **kwargs)
+        self.model.cosmos.to(torch.bfloat16)
+        self.model.action_out_proj.to(torch.float32)
+        return self
+
+    def _prepare_resized_histories(self, batch: dict[str, Any]) -> list[torch.Tensor]:
+        resized_images: list[torch.Tensor] = []
+        for i in range(3):
+            image_history = batch[f"{OBS_IMAGES}.image{i}"]
+            if image_history.shape[0] != 1:
+                raise ValueError(f"InternVLAA1Policy only supports bs=1, got image batch {tuple(image_history.shape)}")
+            resized_images.append(resize_with_pad(image_history[0], self.config.image_resolution))
+        return resized_images
+
+    def _preprocess_images(
+        self,
+        batch: dict[str, torch.Tensor],
+        *,
+        resized_histories: list[torch.Tensor] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if resized_histories is None:
+            resized_histories = self._prepare_resized_histories(batch)
+        images = torch.stack(resized_histories, dim=0).unsqueeze(0)
+        img_masks = torch.stack([batch[f"{OBS_IMAGES}.image{i}_mask"] for i in range(3)], dim=1)
+        return images, img_masks
+
+    def prepare_state(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
+        return pad_vector(batch[OBS_STATE], self.config.max_state_dim)
+
+    def _get_task(self, batch: dict[str, Any]) -> str:
+        task = batch[OBS_TASK]
+        if isinstance(task, str):
+            return task
+        if isinstance(task, (list, tuple)):
+            if len(task) != 1:
+                raise ValueError(f"InternVLAA1Policy only supports bs=1, got {len(task)} tasks")
+            item = task[0]
+            if not isinstance(item, str):
+                raise TypeError(f"Expected task string, got {type(item)!r}")
+            return item
+        raise TypeError(f"Unsupported task payload type: {type(task)!r}")
+
+    def _prepare_qwen_prefix_inputs(
+        self,
+        resized_histories: list[torch.Tensor],
+        task: str,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        qwen_inputs = self.input_builder.build(resized_histories, task)
+        device = self.model.action_in_proj.weight.device
+        return (
+            qwen_inputs[f"{OBS_PREFIX}pixel_values"]
+            .unsqueeze(0)
+            .to(device=device, dtype=self.model.action_in_proj.weight.dtype),
+            qwen_inputs[f"{OBS_PREFIX}image_grid_thw"].unsqueeze(0).to(device=device),
+            qwen_inputs[f"{OBS_PREFIX}input_ids"].unsqueeze(0).to(device=device),
+            qwen_inputs[f"{OBS_PREFIX}attention_mask"].unsqueeze(0).to(device=device),
+        )
+
+    @torch.inference_mode()
+    def forward(
+        self,
+        batch: dict[str, Any],
+        *,
+        noise: torch.Tensor | None = None,
+        decode_image: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        task = self._get_task(batch)
+        resized_histories = self._prepare_resized_histories(batch)
+        pixel_values, image_grid_thw, lang_tokens, lang_masks = self._prepare_qwen_prefix_inputs(
+            resized_histories,
+            task,
+        )
+        state = self.prepare_state(batch)
+        images, img_masks = self._preprocess_images(batch, resized_histories=resized_histories)
+        return self.model.sample_actions(
+            images,
+            img_masks,
+            pixel_values,
+            image_grid_thw,
+            lang_tokens,
+            lang_masks,
+            state,
+            noise=noise,
+            decode_image=decode_image,
+        )