Fix README.md typo

Signed-off-by: Samit <285365963@qq.com>

.buildkite/test-nightly.yml

@@ -394,7 +394,7 @@ steps:
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
     steps:
-      - label: ":full_moon: TTS · Function Test"
+      - label: ":full_moon: TTS · Function Test with L4"
         timeout_in_minutes: 90
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
@@ -413,6 +413,45 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
+      # - label: ":full_moon: TTS · Function Test with H100"
+      #   timeout_in_minutes: 90
+      #   commands:
+      #     - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+      #     - pytest -s -v tests/e2e/ -m "full_model and H100 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
+      #   agents:
+      #     queue: "mithril-h100-pool"
+      #   plugins:
+      #     - kubernetes:
+      #         podSpec:
+      #           containers:
+      #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+      #               resources:
+      #                 limits:
+      #                   nvidia.com/gpu: 1
+      #               volumeMounts:
+      #                 - name: devshm
+      #                   mountPath: /dev/shm
+      #                 - name: hf-cache
+      #                   mountPath: /root/.cache/huggingface
+      #               env:
+      #                 - name: HF_HOME
+      #                   value: /root/.cache/huggingface
+      #                 - name: HF_TOKEN
+      #                   valueFrom:
+      #                     secretKeyRef:
+      #                       name: hf-token-secret
+      #                       key: token
+      #           nodeSelector:
+      #             node.kubernetes.io/instance-type: gpu-h100-sxm
+      #           volumes:
+      #             - name: devshm
+      #               emptyDir:
+      #                 medium: Memory
+      #             - name: hf-cache
+      #               hostPath:
+      #                 path: /mnt/hf-cache
+      #                 type: DirectoryOrCreate
+
       - label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
         timeout_in_minutes: 30
         commands:
@@ -762,6 +801,45 @@ steps:
                       type: DirectoryOrCreate
 
 
+      - label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100"
+        timeout_in_minutes: 120
+        commands:
+          - pytest -sv tests/e2e/online_serving/test_glm_image_expansion.py -m "full_model and diffusion and H100" --run-level "full_model"
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 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: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
         timeout_in_minutes: 60
         commands:
@@ -859,6 +937,46 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: ":full_moon: HunyuanImage3-DIT · Accuracy Test"
+        timeout_in_minutes: 180
+        commands:
+          - export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct"
+          - export HUNYUAN_IMAGE3_DEVICES="0,1,2,3"
+          - pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m full_model --run-level full_model
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 4
+                    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: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
         key: nightly-diffusion-x2iat-performance-qwen-image
         timeout_in_minutes: 180

.buildkite/test-ready.yml

@@ -273,6 +273,42 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: "Omni · Qwen3-Omni Realtime WebSocket"
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_streaming_audio_input_pcm_output_async_chunk_smoke -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
       - label: "Omni · VoxCPM2 Test"
         timeout_in_minutes: 20
         commands:
@@ -363,6 +399,29 @@ steps:
               volumes:
                 - "/fsx/hf_cache:/fsx/hf_cache"
 
+      - label: "TTS · Higgs-Audio-V3 Test"
+        commands:
+          - |
+            timeout 20m bash -c "
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              export VLLM_USE_DEEP_GEMM=0
+              export VLLM_MOE_USE_DEEP_GEMM=0
+              pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "gpu_4_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"
+                - "HF_TOKEN"
+              volumes:
+                - "/fsx/hf_cache:/fsx/hf_cache"
+
       - label: "Diffusion · Qwen Image Edit Test"
         commands:
           - |

.claude/skills/precheck-pr/SKILL.md

@@ -0,0 +1,97 @@
+---
+name: precheck-pr
+description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
+---
+
+# PR Pre-Check
+
+Self-review your branch before creating a PR against `vllm-project/vllm-omni`. Two modes: **quick** catches showstoppers, **full** does a thorough maintainer-grade review. Never posts to GitHub; the report is for the contributor's terminal only.
+
+## Mode Selection
+
+| Mode | When | Time |
+|------|------|------|
+| **Quick** | About to push, final sanity check | ~3 min |
+| **Full** | Ready for review, want maintainer-level scan | ~10 min |
+
+Default to quick if unsure. Run full before marking a PR "ready for review."
+
+## Workflow
+
+### Step 1: Detect Base Branch
+
+```bash
+BASE_SHA=$(git merge-base HEAD origin/main 2>/dev/null \
+         || git merge-base HEAD main 2>/dev/null \
+         || echo origin/main)
+echo "diffing against ${BASE_SHA}"
+git diff --name-only ${BASE_SHA}...HEAD
+```
+
+### Step 2: Validate PR Title
+
+Check the most recent commit message (or branch name if no commit yet) against the project convention. Valid prefixes:
+
+| Prefix | Applies to |
+|--------|-----------|
+| `[Bugfix]` | Bug fixes |
+| `[CI/Build]` | Build or CI improvements |
+| `[Doc]` | Documentation changes |
+| `[Model]` | New/improved models (include model name) |
+| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
+| `[Kernel]` | CUDA/kernel changes |
+| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
+| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
+| `[Misc]` | Other changes (use sparingly) |
+
+✗ if: missing prefix, wrong case (`[bugfix]`), or WIP/Draft in title.
+⚠ if: `[Model]` prefix without the model identifier (e.g., `[Model] Add new model` — should be `[Model] Add <ModelName> ...`).
+
+### Step 3: Categorize the PR
+
+| Diff contains | PR type |
+|---------------|---------|
+| New files under `vllm_omni/model_executor/models/<name>/` | **New Model** |
+| Changes to `vllm_omni/diffusion/` | **Diffusion Model** |
+| `[Bugfix]` prefix or single-file fix | **Bug Fix** |
+| Perf/benchmark/throughput claims in commit msg or diff | **Performance** |
+| Everything else | **General** |
+
+If multiple rows apply (e.g., a diffusion model is also a new model), union the checklists.
+
+### Step 4: Run Checklist
+
+Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type from [references/checklists.md](references/checklists.md). Each item produces ✓, ✗, or ⚠.
+
+### Step 5: Print Report
+
+```
+Pre-check report for <branch>
+
+  Mode: quick | full
+  Type: <new-model | diffusion-model | bug-fix | perf | general>
+
+  Dimension          Result
+  ─────────────────  ──────
+  PR title format    ✓
+  PR desc integrity  ✓
+  Registry/config    ✓
+  Dead code          ⚠ 2 warnings
+  Accuracy           ✓
+  Benchmark          ✗ missing software versions
+
+  Verdict: 1 blocking | 2 warnings | recommend fixing ✗ before PR
+```
+
+**Severity:**
+
+| Mark | Meaning |
+|------|---------|
+| ✗ | Blocking — fix before opening PR |
+| ⚠ | Warning — consider fixing |
+| ✓ | Pass |
+| — | Skipped (not applicable) |
+
+### Stop Here
+
+Do not post comments, open PRs, or modify files. The report is for the contributor's terminal only.

.claude/skills/precheck-pr/references/checklists.md

@@ -0,0 +1,135 @@
+# Pre-Check Checklists
+
+Each item is a concrete check to run against the diff. ✗ = fix before PR. ⚠ = consider fixing.
+
+---
+
+## All PRs: Title Format (check first)
+
+PR title must follow the project convention documented in `docs/contributing/README.md`:
+
+| Prefix | Applies to |
+|--------|-----------|
+| `[Bugfix]` | Bug fixes |
+| `[CI/Build]` | Build or CI improvements |
+| `[Doc]` | Documentation changes |
+| `[Model]` | New/improved models (include model name) |
+| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
+| `[Kernel]` | CUDA/kernel changes |
+| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
+| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
+| `[Misc]` | Other changes (use sparingly) |
+
+**Examples:**
+```
+✓ [Model] Add MiniCPM-o 4.5 support
+✓ [Bugfix] Fix Ovis image text encoder dtype
+✓ [Core] Qwen3-Omni performance optimization
+✓ [Hardware][Ascend] Fix VoxCPM2 on Ascend
+✗ Fix Ovis image text encoder dtype              ← missing prefix
+✗ [bugfix] Fix Ovis image text encoder           ← wrong case
+✗ [WIP] Add model                                ← WIP in title
+```
+
+- [ ] **Title has valid prefix** — ✗ if missing prefix, wrong case, or WIP/Draft in title. For `[Model]` PRs, the title MUST also include the model identifier (e.g., `[Model] Add <ModelName> ...`). ⚠ if missing the model name.
+
+---
+
+## Bug Fix PRs
+
+### Quick
+
+- [ ] **PR body has 4 sections:** what broke (error message), repro steps, root cause, fix description
+- [ ] **Fix is minimal:** diff touches only files directly related to the bug. No unrelated refactors or style changes.
+- [ ] **Repro is runnable:** the PR body includes a copy-paste command that reproduces the bug
+- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
+
+### Full (adds)
+
+- [ ] **Regression test exists:** `git diff --name-only` includes at least one `tests/` file
+- [ ] **Regression test reproduces the original error:** the test would fail on main, pass on this branch
+- [ ] **Fix matches root cause exactly:** no "fix the symptom + something else" — if the root cause is dtype, the fix is dtype, not dtype + reformatting
+- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added
+- [ ] **Upstream pattern match:** the fix follows the same pattern as existing code for similar cases (grep for analogous `from_pretrained` calls, etc.)
+- [ ] **Environment documented:** torch/transformers/vllm versions listed if the bug is version-dependent
+
+---
+
+## Performance PRs
+
+### Quick
+
+- [ ] **Before/after numbers in PR body:** at minimum, one metric (RTF, RPS, latency, VRAM) with main vs PR values
+- [ ] **Hardware specified:** GPU model, count, VRAM. "L4" or "H20" is sufficient. "tested on GPU" is not.
+- [ ] **Concurrency scaling:** if the optimization claims better throughput, numbers at 2+ concurrency levels
+- [ ] **No unexplained regressions:** if any metric goes down at any concurrency level, explain why
+- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
+
+### Full (adds)
+
+- [ ] **Software versions:** torch, CUDA, vllm, vllm-omni versions in PR body
+- [ ] **Warmup stated:** "1 warmup + 3 measured" or equivalent
+- [ ] **Stddev or range:** not just a single best run. "average of 2 rounds" is minimum.
+- [ ] **VRAM measured:** peak GPU memory before and after
+- [ ] **Benchmark script checked in:** under `tests/` or `examples/`, runnable by reviewer
+- [ ] **Exact command line in PR:** not pseudocode, not "run the benchmark"
+- [ ] **Pytest summary included:** `36 passed, 3 skipped, 17 warnings in 2897.66s` — proves the whole suite passes
+- [ ] **Config changes documented:** if `max_num_seqs`, `gpu_memory_utilization`, or deploy YAML changed, old and new values stated with reasoning
+
+---
+
+## New Model PRs
+
+### Quick
+
+- [ ] **PR body file list matches `git diff --name-only`:** every file path in the description exists in the diff
+- [ ] **Registry entries match claimed count:** count `_OMNI_MODELS` or `_DIFFUSION_MODELS` entries added vs claimed in body
+- [ ] **Pipeline config `custom_process_input_func` paths resolve:** each referenced module.function exists
+- [ ] **`__init__.py` exports match registry:** each registered class appears in `__init__.py` and `__all__`
+- [ ] **Shell script modes match README:** `run_curl.sh` supports every mode the README documents
+- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
+
+### Full (adds)
+
+- [ ] **Dead code scan (5 patterns):**
+  1. Dead `forward()` — defined but never called through the pipeline path
+  2. Dead factory/builder functions — instantiated nowhere
+  3. Dead wrappers — thin wrappers with no callers
+  4. Dead branch guards — `if version > X` where X is always satisfied
+  5. Unused parameters — `__init__` args never accessed
+- [ ] **Copy-paste detection:** no string constant defined in 3+ files, no cross-module validation duplication, no near-identical shape coercion fns
+- [ ] **Import hygiene:** no re-export-only imports, no module-level side effects without docs, no `import os` in function body when already at top. Function-body imports are acceptable when the import is only reachable from a closure or factory return value (i.e., `import` would otherwise execute at module load even though the user may never trigger that path). Either lift to top-of-file or annotate with `# noqa: PLC0415`.
+- [ ] **Accuracy:** at least one test compares output against reference; audio = valid WAV at correct sample rate; determinism verified if seed claimed. ⚠ if missing and not explicitly deferred to a follow-up PR or RFC section; — if explicitly deferred (the deferral target SHOULD be linked in the PR body).
+- [ ] **Performance:** hardware, software versions, warmup, and metrics (RTF, VRAM, RPS, latency) stated. ⚠ if missing and not explicitly deferred; — if explicitly deferred.
+- [ ] **Benchmark settings:** model config (TP, PP, max_model_len, enforce_eager, quant), runtime config (batch, input/output spec), environment versions
+- [ ] **Benchmark script checked in + exact command line + pytest summary**
+
+---
+
+## Diffusion Model PRs
+
+Diffusion models live under `vllm_omni/diffusion/`. In addition to the New Model checklist above, verify:
+
+- [ ] **Transformer adapter follows contract:** inherits from `nn.Module`, implements `load_weights()`, uses `vllm_omni.diffusion.attention.layer.Attention`
+- [ ] **Pipeline has both offline and online paths:** `forward()` for batch generation, streaming path for online serving
+- [ ] **Cache-DiT integration:** at least one acceleration (TeaCache, FBCache, etc.) is wired up. ✗ for new **top-level** pipelines. For **subclass** pipelines that inherit `__call__` from a parent where Cache-DiT is already absent, downgrade to ⚠ and recommend filing a follow-up issue against the parent pipeline family.
+- [ ] **Parallelism config:** TP/SP/USP/CFG-Parallel options exposed in pipeline config
+- [ ] **Docs table updated:** model added to `docs/models/supported_models.md` with correct metadata
+- [ ] **E2E test exists:** at least one test exercises the full generate path
+
+---
+
+## General PRs
+
+### Quick
+
+- [ ] **PR body matches diff:** description doesn't claim changes not in the diff
+- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
+- [ ] **CI gates passing:** DCO, pre-commit, build — check `gh pr view --json statusCheckRollup` if a PR already exists
+
+### Full (adds)
+
+- [ ] **No new dead code in changed files:** scan for unreachable paths, unused imports
+- [ ] **Test coverage:** if core code changed (`engine/`, `stages/`, `connectors/`), tests are added or existing tests cover the path
+- [ ] **Import hygiene:** no redundant imports added
+- [ ] **No unrelated changes:** diff scope matches PR description

.claude/skills/readme.md

@@ -25,7 +25,12 @@ include:
 - `add-tts-model`: covers integration of new TTS models and related serving
   workflows
 - `generate-release-note`: helps prepare release notes for repository changes
+- `precheck-pr`: self-check a branch before creating a PR — validates PR title
+  format, catches dead code, verifies accuracy/perf claims, and confirms merge
+  readiness
 - `review-pr`: provides a structured workflow for reviewing pull requests
+- `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the
+  latest vllm-ascend NPUModelRunner
 
 ## Maintenance Guidelines
 

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,4 +1,6 @@
 <!-- markdownlint-disable -->
+**Before submitting:** run the precheck-pr skill with code agent for a self-check against project conventions — catches dead code, missing benchmarks, and title format issues.
+
 PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
 
 ## Purpose
@@ -18,4 +20,5 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 - [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
 </details>
 
-**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
+**BEFORE SUBMITTING,** please read <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>
+(anything written below this line will be removed by GitHub Actions)

README.md

@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 ---
 
 *Latest News* 🔥
+- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
 - [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
 - [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
 - [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!

benchmarks/accuracy/common.py

@@ -151,6 +151,9 @@ class VllmOmniImageClient:
         seed: int | None = None,
         negative_prompt: str | None = None,
         output_compression: int | None = None,
+        bot_task: str | None = None,
+        sys_type: str | None = None,
+        system_prompt: str | None = None,
     ) -> Image.Image:
         if not isinstance(images, list):
             images = [images]
@@ -170,6 +173,12 @@ class VllmOmniImageClient:
             data["negative_prompt"] = negative_prompt
         if output_compression is not None:
             data["output_compression"] = str(output_compression)
+        if bot_task is not None:
+            data["bot_task"] = bot_task
+        if sys_type is not None:
+            data["sys_type"] = sys_type
+        if system_prompt is not None:
+            data["system_prompt"] = system_prompt
 
         files = [
             (

benchmarks/accuracy/image_to_image/gedit_bench.py

@@ -64,6 +64,47 @@ def resolve_model_name(*, model_name: str | None, model: str | None = None, outp
     raise ValueError("model-name is required when it cannot be inferred from model or output-root")
 
 
+def _coerce_score_int(value: Any) -> int:
+    if isinstance(value, bool):
+        raise TypeError(f"Expected numeric score, got bool: {value!r}")
+    if isinstance(value, (int, float)):
+        return int(value)
+    if isinstance(value, str):
+        stripped = value.strip()
+        if stripped.lstrip("-").isdigit():
+            return int(stripped)
+    raise TypeError(f"Expected numeric score, got {type(value).__name__}: {value!r}")
+
+
+def _expand_score_entry(entry: Any) -> list[int]:
+    if isinstance(entry, dict):
+        if "score" in entry:
+            return [_coerce_score_int(entry["score"])]
+        if "value" in entry:
+            return [_coerce_score_int(entry["value"])]
+        numeric_values = [_coerce_score_int(value) for value in entry.values()]
+        if numeric_values:
+            return numeric_values
+        raise TypeError(f"Score dict has no numeric values: {entry!r}")
+    return [_coerce_score_int(entry)]
+
+
+def _normalize_score_list(score: Any) -> list[int]:
+    if isinstance(score, dict):
+        entries: list[Any] = [score]
+    elif isinstance(score, list):
+        entries = score
+    else:
+        entries = [score]
+
+    normalized: list[int] = []
+    for entry in entries:
+        normalized.extend(_expand_score_entry(entry))
+    if not normalized:
+        raise ValueError(f"Empty score list after normalization: {score!r}")
+    return normalized
+
+
 def parse_score_payload(raw_text: str) -> dict[str, Any]:
     try:
         parsed = extract_json_object(raw_text)
@@ -73,10 +114,7 @@ def parse_score_payload(raw_text: str) -> dict[str, Any]:
             parsed = {"score": json.loads(stripped), "reasoning": ""}
         else:
             raise
-    score = parsed.get("score", [])
-    if not isinstance(score, list):
-        score = [score]
-    parsed["score"] = [int(value) for value in score]
+    parsed["score"] = _normalize_score_list(parsed.get("score", []))
     return parsed
 
 
@@ -391,6 +429,9 @@ class GEditBenchRunner:
         num_inference_steps: int = 20,
         guidance_scale: float | None = None,
         seed: int | None = 42,
+        bot_task: str | None = None,
+        sys_type: str | None = None,
+        system_prompt: str | None = None,
     ):
         self.dataset_ref = dataset_ref
         self.output_root = output_root
@@ -400,6 +441,9 @@ class GEditBenchRunner:
         self.num_inference_steps = num_inference_steps
         self.guidance_scale = guidance_scale
         self.seed = seed
+        self.bot_task = bot_task
+        self.sys_type = sys_type
+        self.system_prompt = system_prompt
         self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
 
     def generate(
@@ -478,6 +522,9 @@ class GEditBenchRunner:
             num_inference_steps=self.num_inference_steps,
             guidance_scale=self.guidance_scale,
             seed=self.seed,
+            bot_task=self.bot_task,
+            sys_type=self.sys_type,
+            system_prompt=self.system_prompt,
         )
         edited_image.save(output_path)
         return {
@@ -633,6 +680,25 @@ def build_parser() -> argparse.ArgumentParser:
     generate.add_argument("--num-inference-steps", type=int, default=20)
     generate.add_argument("--guidance-scale", type=float, default=None)
     generate.add_argument("--seed", type=int, default=42)
+    generate.add_argument(
+        "--bot-task",
+        type=str,
+        default=None,
+        choices=["think", "recaption", "think_recaption", "vanilla"],
+        help="HunyuanImage-3 prompting mode for /v1/images/edits (optional).",
+    )
+    generate.add_argument(
+        "--sys-type",
+        type=str,
+        default=None,
+        help="Hunyuan system-prompt type (e.g. en_unified). Defaults from --bot-task when omitted.",
+    )
+    generate.add_argument(
+        "--system-prompt",
+        type=str,
+        default=None,
+        help="Custom Hunyuan system prompt body (sys_type=custom).",
+    )
     generate.add_argument("--workers", type=int, default=1)
     generate.add_argument("--max-samples", type=int, default=None)
     generate.add_argument("--samples-per-group", type=int, default=None)
@@ -675,6 +741,9 @@ def main(argv: list[str] | None = None) -> int:
             num_inference_steps=args.num_inference_steps,
             guidance_scale=args.guidance_scale,
             seed=args.seed,
+            bot_task=args.bot_task,
+            sys_type=args.sys_type,
+            system_prompt=args.system_prompt,
         )
         records = runner.generate(
             model_name=model_name,

docs/api/README.md

@@ -92,7 +92,6 @@ Core scheduling and caching components.
 Configuration classes.
 
 - [vllm_omni.config.model.OmniModelConfig][]
-- [vllm_omni.config.stage_config.ModelPipeline][]
 - [vllm_omni.config.stage_config.StageConfig][]
 - [vllm_omni.config.stage_config.StageConfigFactory][]
 - [vllm_omni.config.stage_config.StageType][]

docs/cli/bench/serve.md

@@ -2,29 +2,29 @@
 The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
 
 ## Notes
-We currently only support using the "openai-chat-omni" backend.
+We currently only support using the "openai-chat-omni" and "openai-image-edits-omni" backend.
 
 ## Basic Parameter Description
 You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
-- `--omni`  
+- `--omni`
   Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
 
-- `--backend`  
+- `--backend`
   Specify the backend adapter as openai-chat-omni, using OpenAI Chat compatible API behavior as the protocol. Currently only openai-chat-omni is supported.
 
-- `--model`  
+- `--model`
   The model identifier to load, filled according to the models supported by vLLM-Omni.
 
-- `--endpoint`  
+- `--endpoint`
   The API endpoint exposed externally, to which clients send their requests.
 
-- `--dataset-name`  
+- `--dataset-name`
   The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio).
 
-- `--num-prompts`  
+- `--num-prompts`
   The total number of requests to send, an integer.
 
-- `--max-concurrency`  
+- `--max-concurrency`
   "Maximum number of concurrent requests. This can be used "
         "to help simulate an environment where a higher level component "
         "is enforcing a maximum number of concurrent requests. While the "
@@ -34,69 +34,73 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
         "actual request rate may be lower than specified with --request-rate, "
         "if the server is not processing requests fast enough to keep up."
 
-- `--request-rate`  
+- `--request-rate`
   "Number of requests per second. If this is inf, "
         "then all the requests are sent at time 0. "
         "Otherwise, we use Poisson process or gamma distribution "
         "to synthesize the request arrival times."
 
-- `--ignore-eos`  
+- `--ignore-eos`
   "Set ignore_eos flag when sending the benchmark request."
 
-- `--metric-percentiles`  
+- `--metric-percentiles`
   Comma-separated list of percentiles for selected metrics. "
         "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "
         "Default value is \"99\"."
         "Use \"--percentile-metrics\" to select metrics.
 
-- `--percentile-metrics`  
+- `--percentile-metrics`
         "Comma-separated list of selected metrics to report percentiles."
                     "This argument specifies the metrics to report percentiles."
-                    'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
+                    'Allowed metric names are "ttft", "tpot", "itl", "ttfc", "tpoc", "icl", '
+                    '"tpop", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
 
-- `--save-result`  
+- `-print-stage`
+Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. Disabled by default.
+
+- `--save-result`
 Specify to save benchmark results to a json file
 
-- `--save-detailed`  
+- `--save-detailed`
 "When saving the results, whether to include per request "
         "information such as response, error, ttfs, tpots, etc."
 
-- `--result-dir`  
+- `--result-dir`
  "Specify directory to save benchmark json results."
         "If not specified, results are saved in the current directory."
 
-- `--result-filename`  
+- `--result-filename`
 "Specify the filename to save benchmark json results."
         "If not specified, results will be saved in "
         "{label}-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
 
-- `--random-prefix-len`  
+- `--random-prefix-len`
   Number of fixed prefix tokens before the random context in a request.
   The total input length is the sum of random-prefix-len and a random
   context length sampled from [input_len * (1 - range_ratio),
   input_len * (1 + range_ratio)].Only the random and random-mm modes
   support this parameter.
 
-- `--random-input-len`  
+- `--random-input-len`
   Number of input tokens per request.Only the random and random-mm modes support this parameter.
 
-- `--random-output-len`  
+- `--random-output-len`
   Number of output tokens per request.Only the random and random-mm modes support this parameter.
 
-- `--random-range-ratio`  
+- `--random-range-ratio`
   Range ratio for sampling input/output length,
   used only for random sampling. Must be in the range [0, 1) to define
   a symmetric sampling range
   [length * (1 - range_ratio), length * (1 + range_ratio)].
   Only the random and random-mm modes support this parameter.
 
-- `--random-mm-base-items-per-request`  
+- `--random-mm-base-items-per-request`
   Base number of multimodal items per request for random-mm.
   Actual per-request count is sampled around this base using
   --random-mm-num-mm-items-range-ratio.
   Only the random-mm mode supports this parameter.
 
-- `--random-mm-limit-mm-per-prompt`  
+- `--random-mm-limit-mm-per-prompt`
   Per-modality hard caps for items attached per request, e.g.
   '{"image": 3, "video": 1, "audio": 1}'. The sampled per-request item
   count is clamped to the sum of these limits. When a modality
@@ -104,7 +108,7 @@ Specify to save benchmark results to a json file
   renormalized.
   Only the random-mm mode supports this parameter.
 
-- `--random-mm-num-mm-items-range-ratio`  
+- `--random-mm-num-mm-items-range-ratio`
   Range ratio r in [0, 1] for sampling items per request.
   We sample uniformly from the closed integer range
   [floor(n*(1-r)), ceil(n*(1+r))]
@@ -115,7 +119,7 @@ Specify to save benchmark results to a json file
   An error is raised if the computed min exceeds the max.
   Only the random-mm mode supports this parameter.
 
-- `--random-mm-bucket-config`  
+- `--random-mm-bucket-config`
   The bucket config is a dictionary mapping a multimodal item
   sampling configuration to a probability.
   Currently allows for 3 modalities: audio, images and videos.
@@ -276,7 +280,9 @@ Generate synthetic image、video、audio inputs alongside random text prompts to
 
 Notes:
 
-- Works only with online benchmark via the OpenAI backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
+- Works only with online benchmark via:
+  - the OpenAI chat backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
+  - the OpenAI edit image backend (`--backend openai-image-edits-omni`) and endpoint `/v1/images/edits`.
 
 Start the server (example):
 
@@ -357,3 +363,105 @@ How sampling works:
 This should be seen as an edge case, and if this behavior can be avoided by setting `--random-mm-limit-mm-per-prompt` to a large number. Note that this might result in errors due to engine config `--limit-mm-per-prompt`.
 - The resulting request contains synthetic image data in `multi_modal_data` (OpenAI Chat format). When `random-mm` is used with the OpenAI Chat backend, prompts remain text and MM content is attached via `multi_modal_data`.
 </details>
+
+### Multi-Stage Benchmark
+
+<details class="admonition abstract" markdown="1">
+<summary>Show more</summary>
+
+Start the server (change the port and model):
+
+```bash
+vllm serve --omni --port 29999 --model ~/Qwen3-Omni-30B-A3B-Instruct
+```
+
+Then run the benchmarking script (remember to add the flag `--print-stage`):
+```bash
+vllm bench serve --omni \
+  --dataset-name random \
+  --port 29999 \
+  --max-concurrency 1 \
+  --num-warmups 2 \
+  --model ~/Qwen3-Omni-30B-A3B-Instruct \
+  --endpoint /v1/chat/completions \
+  --backend openai-chat-omni \
+  --num-prompts 3 \
+  --random-input-len 2500 \
+  --ignore-eos \
+  --percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,ttfc,tpoc,icl \
+  --random-output-len 900 \
+  --extra_body '{"modalities": ["text","audio"]}' \
+  --print-stage
+```
+
+Besides the "Serving Benchmark Result", there will be "Stage Benchmark Result" below:
+```text
+============= Stage Benchmark Result =============
+=============== Stage 0 (thinker) ================
+-------------------Stage Timing-------------------
+Mean stage_gen_time (ms):                18243.16
+Median stage_gen_time (ms):              18220.83
+P99 stage_gen_time (ms):                 18296.75
+================== Text Result ===================
+Stage generated tokens:                  2700
+-----------Serving Time to First Token------------
+Mean Serving TTFT (ms):                  816.31
+Median Serving TTFT (ms):                818.94
+P99 Serving TTFT (ms):                   829.76
+-----Time per Output Token (excl. 1st token)------
+Mean TPOP (ms):                          19.41
+Median TPOP (ms):                        19.40
+P99 TPOP (ms):                           19.47
+---------------Inter-token Latency----------------
+Mean ITL (ms):                           19.41
+Median ITL (ms):                         19.12
+P99 ITL (ms):                            23.58
+================ Stage 1 (talker) ================
+-------------------Stage Timing-------------------
+Mean stage_gen_time (ms):                18849.64
+Median stage_gen_time (ms):              13641.15
+P99 stage_gen_time (ms):                 33655.58
+============= Internal Stream Result =============
+-----------Serving Time to First Chunk------------
+Mean Serving TTFC (ms):                  1029.14
+Median Serving TTFC (ms):                1027.98
+P99 Serving TTFC (ms):                   1032.59
+-----Time per Output Chunk (excl. 1st chunk)------
+Mean TPOP (ms):                          17.53
+Median TPOP (ms):                        18.70
+P99 TPOP (ms):                           18.88
+---------------Inter-chunk Latency----------------
+Mean ICL (ms):                           16.27
+Median ICL (ms):                         18.19
+P99 ICL (ms):                            25.86
+=============== Stage 2 (code2wav) ===============
+-------------------Stage Timing-------------------
+Mean stage_gen_time (ms):                18860.74
+Median stage_gen_time (ms):              13655.84
+P99 stage_gen_time (ms):                 33659.79
+================== Audio Result ==================
+Stage audio duration generated(s):       260.11
+Stage audio frames generated:            6242730
+-----------Serving Time to First Packet-----------
+Mean Serving AUDIO_TTFP (ms):            1366.16
+Median Serving AUDIO_TTFP (ms):          1367.20
+P99 Serving AUDIO_TTFP (ms):             1368.26
+-----------------Real Time Factor-----------------
+Mean AUDIO_RTF:                          0.24
+Median AUDIO_RTF:                        0.26
+P99 AUDIO_RTF:                           0.27
+```
+Explanation:
+- stage_gen_time: Time from submitting a request to a specific stage to that stage finishing generation.
+
+- Serving TTFC (Time to First Chunk): Time from the HTTP request being accepted by the serving frontend to
+  the stage producing its first non-empty output chunk. Keep the name TTFT in text stage, TTFP in audio stage
+  for easier mapping.
+
+- TPOP (Time per Output Chunk): Average time from the first output chunk to stage completion, divided by the
+  number of remaining output chunks. The TPOP abbreviation follows Qwen3.5-Omni Technical Report.
+
+- ICL (Inter-Chunk Latency): Time between two consecutive output chunks produced by the same stage. Keep the
+  name ITL in text stage for easier mapping.
+
+</details>

docs/contributing/README.md

@@ -116,6 +116,20 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 !!! note
     If the PR spans more than one category, please include all relevant prefixes.
 
+### Pre-Check Before Submitting
+
+Before submitting a PR, run the precheck-pr skill with code agent for a self-review against project conventions:
+
+```
+/precheck-pr
+```
+
+The skill offers two modes:
+- **Quick (~3 min):** catches showstoppers — PR title format, missing benchmark claims, rebase status
+- **Full (~10 min):** thorough maintainer-grade review — dead code scan, copy-paste detection, import hygiene
+
+The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the skill in `.claude/skills/precheck-pr/` for the full checklist.
+
 ### Local Test
 Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
 

docs/design/module/dit_module.md

@@ -305,8 +305,8 @@ class GPUWorker:
         )
 
         # Load model
-        model_loader = DiffusersPipelineLoader(load_config)
-        self.pipeline = model_loader.load_model(od_config, load_device=f"cuda:{rank}")
+        model_loader = DiffusersPipelineLoader(load_config, self.od_config)
+        self.pipeline = model_loader.load_model(load_device=f"cuda:{rank}")
 
         # Setup cache backend
         from vllm_omni.diffusion.cache.selector import get_cache_backend

docs/getting_started/installation/gpu/cuda.inc.md

@@ -45,7 +45,7 @@ If you do not need to modify source code of vLLM, you can directly install the s
 uv pip install vllm==0.22.0 --torch-backend=auto
 ```
 
-The 0.22.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.22.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -70,8 +70,8 @@ git checkout v0.22.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.22.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_35_x86_64.whl
+#For CUDA 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash

docs/getting_started/quickstart.md

@@ -22,7 +22,7 @@ source .venv/bin/activate
 uv pip install vllm==0.22.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721
+uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni

docs/models/supported_models.md

@@ -33,7 +33,7 @@ th {
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, T2V with sound | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
+| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
 | `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
@@ -50,6 +50,8 @@ th {
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
 | `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
+| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
+| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
 | `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
 | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
@@ -62,6 +64,10 @@ th {
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
+| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
+| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
+| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
+| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
 | `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
 | `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
 | `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
@@ -69,11 +75,13 @@ th {
 | `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
 | `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
 | `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
+| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
 | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
 | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
 | `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
 |`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
+| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
 | `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 |`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
 

docs/serving/speech_api.md

@@ -116,7 +116,7 @@ Content-Type: application/json
 | `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
 | `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
 
-**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
+**Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian.
 
 #### Voice Clone Parameters (Base task)
 

docs/serving/videos_api.md

@@ -67,8 +67,9 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
 
 | Parameter | Type | Default | Description |
 |-----------|------|---------|-------------|
-| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
-| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
+| `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests |
+| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` |
+| `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` |
 | `width` | integer | model default | Output video width |
 | `height` | integer | model default | Output video height |
 | `num_frames` | integer | model default | Number of generated frames |
@@ -80,6 +81,8 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
 | `flow_shift` | number | null | Scheduler flow-shift value |
 | `true_cfg_scale` | number | null | True CFG scale when supported by the model |
 | `seed` | integer | null | Random seed for reproducibility |
+| `generate_sound` | boolean | false | Request model-generated audio for video models that support sound generation |
+| `sound_duration` | number | null | Duration in seconds for generated audio; defaults to generated video duration |
 | `negative_prompt` | string | null | Text describing what to avoid in the generated video |
 | `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
 | `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
@@ -123,6 +126,41 @@ curl -s http://localhost:8091/v1/videos \
   -F "fps=16"
 ```
 
+### Video-to-Video
+
+For models that support video conditioning, upload the reference video with
+`input_reference`:
+
+```bash
+curl -s http://localhost:8091/v1/videos \
+  -F "prompt=continue this motion with consistent subjects and lighting" \
+  -F "input_reference=@input.mp4;type=video/mp4" \
+  -F "width=1280" \
+  -F "height=720" \
+  -F "num_frames=80" \
+  -F "fps=16"
+```
+
+You can also pass a JSON-safe video URL or `data:video/...;base64,...` payload
+through `video_reference`. Do not send `video_reference` together with
+`input_reference` or `image_reference`.
+
+```bash
+curl -s http://localhost:8091/v1/videos \
+  -F "prompt=continue this motion with consistent subjects and lighting" \
+  -F 'video_reference={"video_url":"https://example.com/input.mp4"}' \
+  -F "width=1280" \
+  -F "height=720" \
+  -F "num_frames=80" \
+  -F "fps=16"
+```
+
+JSON references currently support `image_url`/`video_url`; `file_id` references
+are not implemented yet. Models may expose additional V2V controls through
+`extra_params`. For example, Cosmos3 supports
+`condition_frame_indexes_vision` and `condition_video_keep` to select which
+decoded reference frames are used as clean conditioning.
+
 ### Synchronous Generation
 
 ```bash
@@ -146,7 +184,10 @@ export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
 
 ## Model-Specific Examples
 
-For complete text-to-video and image-to-video walkthroughs, see:
+For complete text-to-video, image-to-video, and model-specific video-to-video
+walkthroughs, see:
 
 - [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
 - [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
+- [Cosmos3 recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
+  for model-specific video-to-video examples and conditioning controls

docs/user_guide/diffusion_features.md

@@ -145,9 +145,10 @@ The following tables show which models support each feature:
 | Model                        | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
 | **Wan2.2**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ✅         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
+| **Wan2.2-S2V**               |     ❌     |     ✅      |           ❌           |       ❌        |         ✅         |         ❌         |   ❌    |             ✅             |  ✅ (encode/decode)   |       ❌        |        ❌         |
 | **Wan2.1-VACE**              |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |      ✅ (decode)      |       ❌        |        ❌         |
 | **LTX-2**                    |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
-| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |
+| **LTX-2.3**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **Helios**                   |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |
 | **HunyuanVideo-1.5 T2V I2V** |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ✅    |             ✅             |  ✅ (encode/decode)   |       ✅        |        ❌         |
 | **DreamID-Omni**             |     ❌     |     ❌      |           ❌           |       ✅        |         ❌         |         ❌         |   ✅    |             ✅             |          ❌           |       ❌        |        ❌         |

docs/user_guide/quantization/fp8.md

@@ -28,6 +28,50 @@ in deep DiT blocks.
 Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
 guide. FP8 on Ampere may use a weight-only path where available.
 
+### Faster FP8 GEMM on Blackwell (quack)
+
+On Blackwell (SM 100+), vLLM runs FP8 linears through the FlashInfer kernel, which
+applies the bias as a separate kernel after the GEMM. On the small GEMMs in video
+DiTs this bias add is a significant overhead. Installing the optional `quack` kernel
+lets vLLM-Omni fuse `alpha * (A @ B) + bias` into a single CuteDSL GEMM, recovering
+that overhead (e.g. HunyuanVideo-1.5 FP8 goes from slower-than-BF16 to faster).
+
+```bash
+# CUDA 12.9
+pip install vllm-omni[quack]
+
+# CUDA 13.x
+pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
+```
+
+It is enabled automatically once installed (no flag needed) and is **Blackwell-only**:
+on Hopper/Ada the CUTLASS FP8 kernel already fuses bias, so quack is not used there.
+Set `VLLM_OMNI_USE_QUACK_FP8=0` to force the FlashInfer path. If `quack-kernels` is
+not installed, FP8 still works — it just keeps the unfused FlashInfer path.
+
+#### Compile cache and warmup
+
+quack JIT-compiles its kernel once per distinct GEMM shape (tens of seconds, longer
+the first time across all autotuned configs). The compiled `.o` files are cached on
+disk and reused on later runs, so this is a one-time cost — **not** per request.
+
+vLLM-Omni points that cache at `~/.cache/vllm_omni/quack` (override with
+`QUACK_CACHE_DIR`) instead of quack's default under `/tmp`, so it survives restarts.
+In containers, set `QUACK_CACHE_DIR` to a mounted/persistent path — or bake it into
+the image — so the first cold start does not recompile. The engine's startup dummy
+run already exercises the kernels, so with a warm cache the first real request is fast.
+
+To pre-warm specific shapes (e.g. at image build time):
+
+```python
+from vllm_omni.quantization.quack_fp8 import warmup_quack_fp8
+# (M, K, N) per linear; M = number of tokens for your resolution/frame count
+warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
+```
+
+> The PyPI package is `quack-kernels` (imported as `quack`); plain `pip install
+> quack` is an unrelated statistics library. Requires CUDA 12.9+ and Python 3.12.
+
 ## Model Type Support
 
 ### Diffusion Model (Qwen-Image, Wan2.2)

examples/offline_inference/cosmos3/inputs/v2v.json

@@ -0,0 +1,14 @@
+{
+    "prompt": "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. The robotic arm is in the process of pouring a transparent liquid into the glass. To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. The lighting is bright, casting soft shadows on the table. The robotic arm's movements are smooth and controlled, demonstrating precision in its task. As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. The other robotic arm on the right side also stays stationary throughout the video. The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper.",
+    "negative_prompt": "",
+    "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4",
+    "height": 720,
+    "width": 1280,
+    "num_frames": 189,
+    "num_inference_steps": 35,
+    "guidance_scale": 6.0,
+    "flow_shift": 10.0,
+    "fps": 24,
+    "condition_frame_indexes_vision": [0, 1],
+    "condition_video_keep": "first"
+}

examples/offline_inference/speech_to_video/speech_to_video.py

@@ -16,9 +16,10 @@ Usage:
 """
 
 import argparse
-import os
+import json
 import time
 from pathlib import Path
+from typing import Any
 
 import torch
 
@@ -29,6 +30,16 @@ from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
 
+def parse_profiler_config(value: str) -> dict[str, Any]:
+    try:
+        config = json.loads(value)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
+    if not isinstance(config, dict):
+        raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
+    return config
+
+
 def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Generate a talking-head video from a reference image and audio (Wan2.2 S2V)."
@@ -87,6 +98,12 @@ def parse_args() -> argparse.Namespace:
         default=3.0,
         help="Scheduler flow shift (default: 3.0).",
     )
+    parser.add_argument(
+        "--boundary-ratio",
+        type=float,
+        default=None,
+        help="Boundary split ratio for low/high DiT stages (S2V). Default varies by model.",
+    )
     parser.add_argument(
         "--output",
         type=str,
@@ -178,30 +195,10 @@ def parse_args() -> argparse.Namespace:
         help="Enable cache-dit summary logging after diffusion forward passes.",
     )
     parser.add_argument(
-        "--profile-dir",
-        type=str,
+        "--profiler-config",
+        type=parse_profiler_config,
         default=None,
-        help="Enable torch profiler and save traces to this directory.",
-    )
-    parser.add_argument(
-        "--profile-record-shapes",
-        action="store_true",
-        help="Record tensor shapes in profiler (increases trace size).",
-    )
-    parser.add_argument(
-        "--profile-with-stack",
-        action="store_true",
-        help="Record stack traces in profiler (increases overhead).",
-    )
-    parser.add_argument(
-        "--profile-with-memory",
-        action="store_true",
-        help="Profile memory usage.",
-    )
-    parser.add_argument(
-        "--profile-with-flops",
-        action="store_true",
-        help="Estimate FLOPs for operations.",
+        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
     )
     return parser.parse_args()
 
@@ -239,23 +236,9 @@ def main():
         ulysses_degree=args.ulysses_degree,
     )
 
-    # Check if profiling is requested via CLI args or environment variable
-    profile_dir = args.profile_dir or os.getenv("VLLM_TORCH_PROFILER_DIR")
-    profiler_enabled = bool(profile_dir)
-    profiler_config = None
-    if profiler_enabled:
-        from vllm.config import ProfilerConfig
+    profiler_enabled = args.profiler_config is not None
 
-        profiler_config = ProfilerConfig(
-            profiler="torch",
-            torch_profiler_dir=profile_dir,
-            torch_profiler_record_shapes=args.profile_record_shapes,
-            torch_profiler_with_stack=args.profile_with_stack,
-            torch_profiler_with_memory=args.profile_with_memory,
-            torch_profiler_with_flops=args.profile_with_flops,
-        )
-
-    omni = Omni(
+    omni_kwargs = dict(
         model=args.model,
         model_class_name="WanS2VPipeline",
         flow_shift=args.flow_shift,
@@ -269,8 +252,12 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_cache_dit_summary=args.enable_cache_dit_summary,
-        profiler_config=profiler_config,
+        profiler_config=args.profiler_config,
     )
+    if args.boundary_ratio is not None:
+        omni_kwargs["boundary_ratio"] = args.boundary_ratio
+
+    omni = Omni(**omni_kwargs)
 
     # Print generation configuration
     print(f"\n{'=' * 60}")
@@ -325,13 +312,16 @@ def main():
     if profiler_enabled:
         print("\n[Profiler] Stopping profiler and collecting results...")
         profile_results = omni.stop_profile()
-        if profile_results and isinstance(profile_results, list):
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
             print("\n" + "=" * 60)
             print("PROFILING RESULTS:")
-            for i, result in enumerate(profile_results):
-                print(f"\nStage {i}:")
-                if result:
-                    print(f"  {result}")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
             print("=" * 60)
         else:
             print("[Profiler] No valid profiling data returned.")

examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline higgs-audio v3 inference example.
+
+Runs Stage 0 (Qwen3 talker) + Stage 1 (HiggsAudio codec) end-to-end
+through the vLLM-Omni engine without going through the HTTP server, and
+saves a 24 kHz mono WAV per prompt.
+
+Example (plain TTS):
+
+    python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
+        --texts "Hello world." \\
+        --output-dir results/higgs_v3_wavs
+
+Example (voice clone):
+
+    python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
+        --texts "Hello world." \\
+        --ref-audio path/to/reference.wav \\
+        --ref-text "Transcript of the reference clip." \\
+        --output-dir results/higgs_v3_wavs
+"""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
+
+import re
+import time
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+
+from vllm_omni import Omni
+
+SAMPLE_RATE = 24_000
+DEFAULT_TEXTS = (
+    "Hello world.",
+    "The quick brown fox jumps over the lazy dog.",
+    "Today is a beautiful day for a walk in the park.",
+)
+
+
+def parse_args():
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Offline higgs-audio v3 inference")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="bosonai/higgs-audio-v3-tts-4b",
+        help="Stage-0 talker model id or path.",
+    )
+    parser.add_argument(
+        "--texts",
+        type=str,
+        nargs="+",
+        default=list(DEFAULT_TEXTS),
+        help="One or more plain-text prompts to synthesize.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        default="results/higgs_v3_wavs",
+        help="Directory to write per-prompt WAV files.",
+    )
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        default=None,
+        help="Override the deploy config path.",
+    )
+    parser.add_argument(
+        "--max-new-tokens",
+        type=int,
+        default=2048,
+        help="Cap on Stage-0 codec frames per request.",
+    )
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Reference audio for voice clone (WAV/FLAC/MP3 path).",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcript of the reference audio. Optional but improves fidelity.",
+    )
+    return parser.parse_args()
+
+
+def _slugify(text: str) -> str:
+    slug = re.sub(r"\s+", "_", text.strip().lower())
+    slug = re.sub(r"[^a-z0-9_]+", "", slug)
+    return slug[:48] or "prompt"
+
+
+def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
+    audio = multimodal_output.get("model_outputs")
+    if audio is None:
+        audio = multimodal_output.get("audio")
+    if audio is None:
+        raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
+    if isinstance(audio, list):
+        valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
+        if not valid:
+            raise ValueError("audio list is empty")
+        return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
+    return torch.as_tensor(audio).float().cpu().reshape(-1)
+
+
+def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
+    arr = pcm.numpy()
+    if arr.dtype.kind == "f":
+        arr = np.clip(arr, -1.0, 1.0)
+        arr = (arr * 32767.0).astype(np.int16)
+    else:
+        arr = arr.astype(np.int16)
+    return arr
+
+
+def main():
+    args = parse_args()
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    from transformers import AutoTokenizer
+
+    from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
+        HiggsAudioV3TokenizerAdapter,
+        apply_delay_pattern,
+        encode_reference_audio,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
+    adapter = HiggsAudioV3TokenizerAdapter(tokenizer)
+
+    # Load and encode reference audio once (if voice cloning)
+    ref_codes_delayed: torch.Tensor | None = None
+    if args.ref_audio is not None:
+        ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
+        if ref_wav.ndim == 2:
+            ref_wav = ref_wav.mean(axis=1)
+        ref_codes_raw = encode_reference_audio(ref_wav, int(ref_sr))
+        ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
+        print(f"Reference   : {args.ref_audio}")
+        print(f"Ref codes   : {ref_codes_raw.shape[0]} frames -> {ref_codes_delayed.shape[0]} delayed")
+        if args.ref_text:
+            print(f"Ref text    : {args.ref_text}")
+
+    print(f"Model       : {args.model}")
+    print(f"Prompts     : {len(args.texts)}")
+    print(f"Output dir  : {output_dir}")
+    print(f"Voice clone : {'yes' if ref_codes_delayed is not None else 'no'}")
+
+    total_elapsed = 0.0
+    total_dur = 0.0
+    for text in args.texts:
+        if ref_codes_delayed is not None:
+            prompt_ids = adapter.build_prompt(
+                text,
+                num_ref_tokens=int(ref_codes_delayed.shape[0]),
+                reference_text=args.ref_text,
+            )
+            prompt = {
+                "prompt_token_ids": prompt_ids,
+                "additional_information": {
+                    "audio_input_ids": ref_codes_delayed.to(torch.long),
+                    "audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool),
+                },
+            }
+        else:
+            prompt_ids = adapter.build_prompt(text)
+            prompt = {"prompt_token_ids": prompt_ids}
+
+        t_start = time.perf_counter()
+        outputs = engine.generate([prompt])
+        elapsed = time.perf_counter() - t_start
+        total_elapsed += elapsed
+
+        mm = outputs[0].outputs[0].multimodal_output
+        pcm = _extract_pcm(mm)
+        slug = _slugify(text)
+        out_path = output_dir / f"{slug}.wav"
+        sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
+        dur = pcm.numel() / SAMPLE_RATE
+        total_dur += dur
+        print(f"  {slug:<50} dur={dur:5.2f}s  -> {out_path}")
+
+    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
+    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/text_to_speech/fish_speech/gradio_demo.py

@@ -51,13 +51,18 @@ def build_payload(
     response_format: str = "wav",
     stream: bool = False,
 ) -> dict:
-    """Build the /v1/audio/speech request payload."""
+    """Build the /v1/audio/speech request payload.
+
+    Fish Speech has no built-in speakers; omit ``voice`` so the model
+    generates with a random timbre (controllable via text tags like
+    ``[cool]``, ``[warm]``, etc.).  When ``ref_audio`` is supplied we
+    enter voice-cloning mode instead.
+    """
     if not text or not text.strip():
         raise gr.Error("Please enter text to synthesize.")
 
     payload: dict = {
         "input": text.strip(),
-        "voice": "default",
         "response_format": "pcm" if stream else response_format,
         "stream": stream,
     }
@@ -193,7 +198,8 @@ def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
 
                 with gr.Accordion("Voice Cloning (optional)", open=False):
                     gr.Markdown(
-                        "Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice."
+                        "Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice. "
+                        "Leave empty for random timbre; use text tags like `[cool]`, `[warm]`, `[happy]` to steer the output."
                     )
                     ref_audio = gr.Audio(
                         label="Reference Audio",

examples/online_serving/text_to_speech/fish_speech/speech_client.py

@@ -39,7 +39,6 @@ def run_tts(args) -> None:
     payload = {
         "model": args.model,
         "input": args.text,
-        "voice": "default",
         "response_format": args.response_format,
     }
 

examples/online_serving/text_to_speech/higgs_audio_v3/README.md

@@ -0,0 +1,40 @@
+# Higgs-Audio V3 Online Serving
+
+## Start the server
+
+```bash
+# Default: GPU 0, port 8095
+./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
+
+# Custom GPU / port
+PORT=8096 GPUS=0,1 ./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
+```
+
+## Plain text TTS
+
+```bash
+python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
+    --base-url http://localhost:8095 \
+    --output-dir /tmp/higgs_v3_batch \
+    --prompts "Hello world." "The quick brown fox jumps over the lazy dog."
+```
+
+## Voice clone
+
+```bash
+python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
+    --base-url http://localhost:8095 \
+    --output-dir /tmp/higgs_v3_clone \
+    --ref-audio path/to/reference.wav \
+    --ref-text "transcript of the reference clip" \
+    --prompts "Text to synthesize in the cloned voice."
+```
+
+## curl example
+
+```bash
+curl -X POST http://localhost:8095/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{"model": "higgs_audio_v3", "input": "Hello world."}' \
+    --output hello.wav
+```

examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Batch client for the higgs-audio v3 online server.
+
+Sends prompts to ``/v1/audio/speech`` and saves the returned WAV files.
+
+Usage (plain text -> speech):
+
+  python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
+      --base-url http://localhost:8095 \
+      --output-dir /tmp/higgs_v3_batch \
+      --prompts "Hello world." "The quick brown fox jumps over the lazy dog."
+
+Usage (voice clone):
+
+  python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
+      --base-url http://localhost:8095 \
+      --output-dir /tmp/higgs_v3_clone \
+      --ref-audio path/to/reference.wav \
+      --ref-text "the transcript of the reference clip" \
+      --prompts "Hello world."
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import sys
+from pathlib import Path
+
+DEFAULT_PROMPTS = (
+    "Hello world.",
+    "The quick brown fox jumps over the lazy dog.",
+    "Today is a beautiful day for a walk in the park.",
+    "Innovation distinguishes between a leader and a follower.",
+)
+
+
+def _slug(text: str) -> str:
+    import re
+
+    s = re.sub(r"\s+", "_", text.strip().lower())
+    return re.sub(r"[^a-z0-9_]+", "", s)[:32] or "prompt"
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument("--base-url", default="http://localhost:8095")
+    parser.add_argument("--model", default="higgs_audio_v3")
+    parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
+    parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_v3_batch"))
+    parser.add_argument("--format", choices=("wav", "pcm"), default="wav")
+    parser.add_argument("--max-new-tokens", type=int, default=2048)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--timeout-s", type=float, default=120.0)
+    parser.add_argument(
+        "--ref-audio",
+        type=Path,
+        default=None,
+        help="Reference clip for voice clone (WAV/FLAC/MP3 path). Pair with --ref-text.",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Transcript of the reference clip. Optional but improves fidelity.",
+    )
+    args = parser.parse_args()
+
+    ref_audio_data_url: str | None = None
+    if args.ref_audio is not None:
+        if not args.ref_audio.exists():
+            print(f"ref-audio file not found: {args.ref_audio}", file=sys.stderr)
+            return 2
+        mime = "audio/wav" if args.ref_audio.suffix.lower() == ".wav" else "audio/mpeg"
+        ref_b64 = base64.b64encode(args.ref_audio.read_bytes()).decode("ascii")
+        ref_audio_data_url = f"data:{mime};base64,{ref_b64}"
+
+    try:
+        import httpx
+    except ImportError:
+        print("this client needs `httpx`. Install with `pip install httpx`.", file=sys.stderr)
+        return 2
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    url = args.base_url.rstrip("/") + "/v1/audio/speech"
+    failures = 0
+    with httpx.Client(timeout=args.timeout_s) as client:
+        for prompt in args.prompts:
+            payload = {
+                "model": args.model,
+                "input": prompt,
+                "response_format": args.format,
+                "max_new_tokens": args.max_new_tokens,
+                "seed": args.seed,
+            }
+            if ref_audio_data_url is not None:
+                payload["ref_audio"] = ref_audio_data_url
+                if args.ref_text:
+                    payload["ref_text"] = args.ref_text
+            resp = client.post(url, json=payload)
+            if resp.status_code != 200:
+                print(f"[FAIL] {prompt!r} -> {resp.status_code}: {resp.text[:200]}", file=sys.stderr)
+                failures += 1
+                continue
+            suffix = ".wav" if args.format == "wav" else ".pcm"
+            out = args.output_dir / f"{_slug(prompt)}{suffix}"
+            out.write_bytes(resp.content)
+            print(f"[ ok ] {prompt!r} -> {out} ({len(resp.content)} bytes)")
+
+    return 1 if failures else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+# Launch vLLM-Omni server for higgs-audio v3.
+#
+# Supports plain text TTS and voice cloning via /v1/audio/speech.
+#
+# Usage:
+#   ./run_server.sh                 # default port 8095, GPU 0
+#   PORT=8096 GPUS=0,1 ./run_server.sh
+#   MODEL=/path/to/local/checkpoint ./run_server.sh
+
+set -e
+
+MODEL="${MODEL:-bosonai/higgs-audio-v3-tts-4b}"
+PORT="${PORT:-8095}"
+GPUS="${GPUS:-0}"
+GPU_UTIL="${GPU_UTIL:-0.6}"
+
+echo "Starting higgs-audio v3 server"
+echo "  MODEL=$MODEL"
+echo "  PORT=$PORT"
+echo "  CUDA_VISIBLE_DEVICES=$GPUS"
+
+CUDA_VISIBLE_DEVICES="$GPUS" \
+VLLM_USE_DEEP_GEMM=0 \
+VLLM_MOE_USE_DEEP_GEMM=0 \
+vllm-omni serve "$MODEL" \
+    --deploy-config vllm_omni/deploy/higgs_multimodal_qwen3.yaml \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --gpu-memory-utilization "$GPU_UTIL" \
+    --trust-remote-code \
+    --omni

examples/quantization/check_modelopt_fp8_export.py

@@ -0,0 +1,247 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Verify a ModelOpt FP8 diffusers checkpoint exported by
+quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py.
+
+Three checks:
+  A. transformer/config.json has a sane quantization_config block.
+  B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors.
+  C. transformer disk size is materially smaller than a BF16 baseline.
+
+Example:
+    python examples/quantization/check_modelopt_fp8_export.py \\
+        --output ./hv15-480p-modelopt-fp8
+
+    # Optional: compare disk size against a local or HF BF16 baseline.
+    python examples/quantization/check_modelopt_fp8_export.py \\
+        --output ./hv15-480p-modelopt-fp8 \\
+        --baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from collections import Counter
+from pathlib import Path
+
+SUPPORTED_ALGOS = ("FP8",)
+
+
+def _check_config(transformer_dir: Path) -> tuple[int, str | None]:
+    """Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn."""
+    cfg_path = transformer_dir / "config.json"
+    if not cfg_path.exists():
+        print(f"[FAIL] {cfg_path} missing.")
+        return 1, None
+
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+
+    qc = cfg.get("quantization_config")
+    if not isinstance(qc, dict):
+        print(f"[FAIL] No `quantization_config` block in {cfg_path}.")
+        return 1, None
+
+    print(f"[A] quantization_config from {cfg_path}:")
+    print(json.dumps(qc, indent=2))
+
+    quant_algo = qc.get("quant_algo")
+    issues = []
+    if qc.get("quant_method") != "modelopt":
+        issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')")
+    if quant_algo not in SUPPORTED_ALGOS:
+        issues.append(
+            f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)"
+        )
+
+    if issues:
+        print("[A] WARN — config looks incomplete:")
+        for issue in issues:
+            print(f"    - {issue}")
+        return 2, quant_algo
+    print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).")
+    return 0, quant_algo
+
+
+def _read_safetensors_header(path: Path) -> dict:
+    """Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors.
+
+    Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}.
+    Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ...
+    """
+    import struct
+
+    with open(path, "rb") as f:
+        header_len = struct.unpack("<Q", f.read(8))[0]
+        header = json.loads(f.read(header_len))
+    header.pop("__metadata__", None)
+    return header
+
+
+def _classify_weight_scale_granularity(weight_scale_shapes: list[list[int]]) -> str:
+    """Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes.
+
+    ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1
+    interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 —
+    and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block.
+    """
+    if not weight_scale_shapes:
+        return "no weight_scale tensors found"
+
+    def meaningful_dims(shape: list[int]) -> int:
+        return sum(1 for d in shape if d > 1)
+
+    per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0)
+    per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1)
+    per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2)
+    total = len(weight_scale_shapes)
+    if per_tensor == total:
+        return "per-tensor (all scalar scales)"
+    if per_channel == total:
+        return "per-channel (1 meaningful dim)"
+    if per_block == total:
+        return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)"
+    return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}"
+
+
+def _check_safetensors(transformer_dir: Path) -> int:
+    """Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header."""
+    files = sorted(transformer_dir.glob("*.safetensors"))
+    if not files:
+        print(f"[FAIL] No *.safetensors in {transformer_dir}.")
+        return 1
+
+    header_dtype_counts: Counter[str] = Counter()
+    sample_quant_weight_keys: list[str] = []
+    sample_scale_keys: list[str] = []
+    weight_scale_shapes: list[list[int]] = []
+    sample_weight_scale_entries: list[tuple[str, list[int]]] = []
+
+    for f in files:
+        try:
+            header = _read_safetensors_header(f)
+        except Exception as exc:
+            print(f"[B] WARN — could not parse header of {f}: {exc}")
+            continue
+        for k, info in header.items():
+            dtype = info.get("dtype", "?")
+            header_dtype_counts[dtype] += 1
+            # FP8 stores weights as F8_E4M3 directly.
+            if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5:
+                sample_quant_weight_keys.append(k)
+            if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5:
+                sample_scale_keys.append(k)
+            if k.endswith(".weight_scale"):
+                weight_scale_shapes.append(info.get("shape", []))
+                if len(sample_weight_scale_entries) < 5:
+                    sample_weight_scale_entries.append((k, info.get("shape", [])))
+
+    print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):")
+    for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]):
+        marker = "  <-- FP8" if dtype.startswith("F8") else ""
+        print(f"    {dtype:10s} {count:>6d}{marker}")
+
+    quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8"))
+    if quant_count == 0:
+        print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.")
+        return 1
+
+    print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.")
+    if sample_quant_weight_keys:
+        print(f"    sample quantized weight tensors: {sample_quant_weight_keys[:3]}")
+    if sample_scale_keys:
+        print(f"    sample scale tensors:            {sample_scale_keys[:3]}")
+    print("    (Note: torch's get_tensor() may return these as bf16 views on some versions —")
+    print("     irrelevant; vLLM's loader uses native FP8 ops.)")
+
+    # Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D).
+    print(f"\n    weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}")
+    for key, shape in sample_weight_scale_entries[:3]:
+        print(f"      {key}: shape {shape}")
+    return 0
+
+
+def _disk_size_gib(p: Path) -> float:
+    return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3)
+
+
+def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int:
+    """Returns 0 always (informational only)."""
+    quant_size = _disk_size_gib(transformer_dir)
+    print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB")
+
+    if baseline is None:
+        print("[C] SKIP — pass --baseline <path or HF id> to compare against BF16.")
+        return 0
+
+    baseline_path = Path(baseline)
+    if not baseline_path.exists():
+        # Try HF download.
+        try:
+            from huggingface_hub import snapshot_download
+        except ImportError:
+            print("[C] SKIP — huggingface_hub not installed and baseline not a local path.")
+            return 0
+        print(f"    Downloading baseline transformer from HF: {baseline}")
+        baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"]))
+
+    bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path
+    bf16_size = _disk_size_gib(bf16_dir)
+    if bf16_size == 0:
+        print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}")
+        return 0
+
+    # Expected reduction for FP8: ~50%.
+    min_reduction = 30
+    reduction = (1 - quant_size / bf16_size) * 100
+    print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})")
+    print(f"[C] Disk reduction: {reduction:.1f}%  (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)")
+    if reduction < min_reduction:
+        print(
+            f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization."
+        )
+    return 0
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.")
+    p.add_argument(
+        "--baseline",
+        default=None,
+        help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.",
+    )
+    args = p.parse_args()
+
+    out_root = Path(args.output).expanduser().resolve()
+    transformer_dir = out_root / "transformer"
+    if not transformer_dir.exists():
+        print(f"[FAIL] {transformer_dir} does not exist.")
+        sys.exit(1)
+
+    print(f"Checking: {out_root}\n")
+
+    fail = 0
+    config_status, _quant_algo = _check_config(transformer_dir)
+    fail |= config_status
+    fail |= _check_safetensors(transformer_dir)
+    _check_size_vs_baseline(transformer_dir, args.baseline)
+
+    print()
+    if fail == 0:
+        print("=" * 60)
+        print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.")
+    elif fail == 1:
+        print("=" * 60)
+        print("FAILURES detected — calibration may need to be re-run.")
+        sys.exit(1)
+    else:
+        print("=" * 60)
+        print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.")
+
+
+if __name__ == "__main__":
+    main()

examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py

@@ -0,0 +1,448 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint.
+
+Calibrates the DiT transformer using a small video prompt set and exports a
+diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
+The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
+(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
+
+Layers kept full precision match the #2728 / #2795 pattern: modulation,
+AdaLayerNorm, entry/exit projections, embeddings, the token refiner path,
+and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention
+empirically degrades under FP8 (see #2920 ablation).
+
+Example:
+    python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\
+        --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
+        --output ./hv15-480p-modelopt-fp8 \\
+        --overwrite
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import re
+import shutil
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+from diffusers import DiffusionPipeline
+
+DEFAULT_PROMPTS = [
+    "A dog running across a field of golden wheat.",
+    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
+    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
+    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
+    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
+    "A close-up of a blooming rose covered in morning dew, soft natural light.",
+    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
+    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
+]
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.")
+    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
+    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
+    p.add_argument("--height", type=int, default=480)
+    p.add_argument("--width", type=int, default=832)
+    p.add_argument(
+        "--num-frames",
+        type=int,
+        default=33,
+        help="Frames per calibration sample. 33 matches the typical short benchmark.",
+    )
+    p.add_argument("--guidance-scale", type=float, default=6.0)
+    p.add_argument(
+        "--calib-steps",
+        type=int,
+        default=10,
+        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
+    )
+    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument(
+        "--prompt",
+        action="append",
+        default=[],
+        help="Custom calibration prompt. Repeat to provide multiple.",
+    )
+    p.add_argument(
+        "--quantize-mha",
+        action="store_true",
+        help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.",
+    )
+    p.add_argument(
+        "--weight-block-size",
+        type=str,
+        default=None,
+        help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). "
+        "Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at "
+        "negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.",
+    )
+    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
+    return p
+
+
+def _parse_block_size(spec: str | None) -> list[int] | None:
+    if spec is None:
+        return None
+    parts = [int(x) for x in spec.split(",") if x.strip()]
+    if len(parts) != 2:
+        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
+    return parts
+
+
+def _require_modelopt() -> Any:
+    try:
+        import modelopt.torch.quantization as mtq
+    except ModuleNotFoundError as exc:
+        raise SystemExit(
+            "NVIDIA ModelOpt is not installed. Install with:\n"
+            "  pip install 'nvidia-modelopt[all]'\n"
+            f"Original error: {exc}"
+        ) from exc
+    return mtq
+
+
+def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
+    model_path = args.model
+    output_dir = Path(args.output).expanduser().resolve()
+    if output_dir.exists():
+        if not args.overwrite:
+            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
+        shutil.rmtree(output_dir)
+    return model_path, output_dir
+
+
+def _select_dtype(name: str) -> torch.dtype:
+    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
+
+
+def _build_prompts(args: argparse.Namespace) -> list[str]:
+    prompts = args.prompt or DEFAULT_PROMPTS
+    if args.calib_size <= 0:
+        raise SystemExit("--calib-size must be positive.")
+    if len(prompts) < args.calib_size:
+        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
+        prompts = (prompts * repeats)[: args.calib_size]
+    return prompts[: args.calib_size]
+
+
+# Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern).
+# - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding
+# - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation
+# - proj_out: final output projection
+# - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear
+def _filter_func_hv15(name: str) -> bool:
+    pattern = re.compile(
+        r"(proj_out.*|"
+        r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|"
+        r"time_embed|cond_type_embed|"
+        r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|"
+        r"token_refiner).*)"
+    )
+    return pattern.match(name) is not None
+
+
+def _mha_filter_func(name: str) -> bool:
+    pattern = re.compile(
+        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
+    )
+    return pattern.match(name) is not None
+
+
+def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
+    if not hasattr(mtq, "disable_quantizer"):
+        return
+    mtq.disable_quantizer(backbone, _filter_func_hv15)
+    if not quantize_mha:
+        mtq.disable_quantizer(backbone, _mha_filter_func)
+
+
+def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
+    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
+    if hasattr(pipe, "set_progress_bar_config"):
+        pipe.set_progress_bar_config(disable=True)
+    pipe.to("cuda")
+    return pipe
+
+
+def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
+    generator = torch.Generator(device="cuda")
+
+    # Try to set guidance on the pipeline's guider object up front (modern
+    # diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls
+    # back silently — calibration uses whatever default the pipeline ships with.
+    guider = getattr(pipe, "guider", None)
+    if guider is not None and hasattr(guider, "guidance_scale"):
+        try:
+            guider.guidance_scale = args.guidance_scale
+        except Exception:
+            pass
+
+    base_kwargs = dict(
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.calib_steps,
+        output_type="latent",
+    )
+
+    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
+        with torch.inference_mode():
+            for idx, prompt in enumerate(prompts):
+                generator.manual_seed(args.seed + idx)
+                # Try with guidance_scale first; fall back without on TypeError
+                # for pipelines (like HV-1.5) that take CFG via guider config.
+                try:
+                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
+                except TypeError as exc:
+                    if "guidance_scale" not in str(exc):
+                        raise
+                    pipe(prompt=prompt, generator=generator, **base_kwargs)
+
+    return forward_loop
+
+
+def _summarize_export(output_dir: Path) -> None:
+    cfg_path = output_dir / "transformer" / "config.json"
+    if not cfg_path.exists():
+        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
+        return
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+    qc = cfg.get("quantization_config")
+    if not isinstance(qc, dict):
+        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
+        return
+    print("Export summary:")
+    print(f"  quant_method: {qc.get('quant_method')}")
+    print(f"  quant_algo:   {qc.get('quant_algo')}")
+    producer = qc.get("producer")
+    if isinstance(producer, dict):
+        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
+    print(f"  config path:  {cfg_path}")
+
+
+def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
+    """Convert in-memory weights of quantized modules to actual FP8 storage.
+
+    `export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't
+    in ModelOpt's recognized-model registry), so we must call the per-weight
+    export helper ourselves. Same workaround as the HunyuanImage-3 calibration
+    helper.
+    """
+    from modelopt.torch.export.quant_utils import (
+        QUANTIZATION_NONE,
+        get_quantization_format,
+        quantizer_attr_names,
+        weight_attr_names,
+    )
+    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+
+    exported = 0
+    for name, module in backbone.named_modules():
+        try:
+            quantization_format = get_quantization_format(module)
+        except Exception as exc:
+            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
+            continue
+        if quantization_format == QUANTIZATION_NONE:
+            continue
+        for weight_name in weight_attr_names(module):
+            quantizer_attrs = quantizer_attr_names(weight_name)
+            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
+            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
+                continue
+            _export_quantized_weight(module, dtype, weight_name)
+            exported += 1
+    return exported
+
+
+def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
+    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).
+
+    Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When
+    `weight_block_size` is given, advertise block-wise weight quantization in
+    the saved metadata (so consumers know to expect multi-element scale tensors).
+    """
+    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
+    if weight_block_size is not None:
+        weights_cfg["strategy"] = "block"
+        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
+    return {
+        "config_groups": {
+            "group_0": {
+                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
+                "weights": weights_cfg,
+                "targets": ["Linear"],
+            }
+        },
+        "ignore": [
+            "context_embedder*",
+            "context_embedder_2*",
+            "cond_type_embed*",
+            "image_embedder*",
+            "norm1.linear*",
+            "norm1_context.linear*",
+            "norm2*",
+            "norm2_context*",
+            "norm_out*",
+            "proj_out*",
+            "time_embed*",
+            "token_refiner*",
+            "x_embedder*",
+        ],
+        "producer": {"name": "modelopt"},
+        "quant_algo": "FP8",
+        "quant_method": "modelopt",
+    }
+
+
+def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
+    """Inject quant_algo: FP8 + config_groups into transformer/config.json so
+    vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8."""
+    cfg_path = output_dir / "transformer" / "config.json"
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+
+    new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size)
+    existing = cfg.get("quantization_config")
+    if isinstance(existing, dict):
+        producer = existing.get("producer")
+        if isinstance(producer, dict):
+            new_qc["producer"] = producer
+
+    cfg["quantization_config"] = new_qc
+    with cfg_path.open("w", encoding="utf-8") as f:
+        json.dump(cfg, f, indent=2)
+
+
+def _save_pipeline_with_fp8_transformer(
+    pipe: DiffusionPipeline,
+    model_path: str,
+    output_dir: Path,
+    max_shard_size: str = "5GB",
+) -> None:
+    """Save the pipeline with the (now FP8) transformer.
+
+    Copies the source directory verbatim except for `transformer/`, then
+    saves the transformer with quantizers hidden so the state dict contains
+    only the FP8 weights + scale tensors.
+    """
+    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
+
+    src = Path(model_path)
+    if not src.exists():
+        from huggingface_hub import snapshot_download
+
+        src = Path(snapshot_download(model_path))
+
+    if output_dir.exists():
+        shutil.rmtree(output_dir)
+    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer"))
+
+    transformer_out = output_dir / "transformer"
+    # `hide_quantizers_from_state_dict` walks named_modules(); pass the actual
+    # nn.Module (transformer), not the diffusers Pipeline wrapper.
+    with hide_quantizers_from_state_dict(pipe.transformer):
+        pipe.transformer.save_pretrained(
+            str(transformer_out),
+            safe_serialization=True,
+            max_shard_size=max_shard_size,
+        )
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
+
+    mtq = _require_modelopt()
+    model_path, output_dir = _ensure_paths(args)
+    dtype = _select_dtype(args.dtype)
+    prompts = _build_prompts(args)
+
+    print("Quantization plan:")
+    weight_block_size = _parse_block_size(args.weight_block_size)
+
+    print(f"  input:           {args.model}")
+    print(f"  output:          {output_dir}")
+    print(f"  dtype:           {dtype}")
+    print(f"  height/width:    {args.height}x{args.width}")
+    print(f"  num_frames:      {args.num_frames}")
+    print(f"  calib_size:      {len(prompts)}")
+    print(f"  calib_steps:     {args.calib_steps}")
+    print(f"  quantize_mha:    {args.quantize_mha}")
+    print(
+        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
+    )
+
+    pipe = _load_pipeline(model_path, dtype)
+    backbone = pipe.transformer
+
+    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+    if weight_block_size is not None:
+        # Switch from per-tensor (default) to block-wise weight quantization.
+        # ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer.
+        quant_config["quant_cfg"]["*weight_quantizer"] = {
+            "num_bits": (4, 3),  # E4M3 (FP8 weights, same as default)
+            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
+        }
+        print(
+            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
+            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
+        )
+
+    forward_loop = _build_forward_loop(pipe, args, prompts)
+    quantized = mtq.quantize(backbone, quant_config, forward_loop)
+    if quantized is not None:
+        pipe.transformer = quantized
+        backbone = quantized
+
+    _disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha)
+
+    print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,")
+    print("so we have to call the per-weight export helper ourselves)...")
+    exported = _force_export_quantized_weights(backbone, dtype)
+    print(f"  -> {exported} weights converted to FP8 in memory")
+    if exported == 0:
+        raise SystemExit(
+            "No quantized weights were exported. Calibration may have skipped every layer "
+            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
+            "weight quantizers."
+        )
+
+    print("\nSaving pipeline with FP8 transformer...")
+    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
+    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
+    print(f"Saved to: {output_dir}")
+    _summarize_export(output_dir)
+
+    print("\nNext: validate the checkpoint with vllm-omni:")
+    print(
+        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
+        f"    --model {output_dir} \\\n"
+        "    --quantization fp8 \\\n"
+        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
+        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
+        "    --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n"
+        "    --output outputs/hv15_modelopt_fp8.mp4 \\\n"
+        "    --enforce-eager"
+    )
+    print(
+        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
+        "checkpoint's config.json has modelopt metadata.)"
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/quantization/quantize_wan2_2_modelopt_fp8.py

@@ -0,0 +1,471 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint.
+
+Calibrates the DiT transformer using a small video prompt set and exports a
+diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
+The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
+(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
+
+Layers kept full precision match the #2728 / #2795 pattern: condition embedder
+(time/text/image), patch embedding, modulation (scale_shift_table), final
+norm + proj_out, and sequence-parallel helpers. All attention + FFN linears
+are quantized — static calibration handles the numerics that online FP8
+couldn't (see #2920 ablation).
+
+Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that
+fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here.
+
+Example:
+    python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\
+        --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
+        --output ./wan22-ti2v-modelopt-fp8 \\
+        --overwrite
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import re
+import shutil
+import sys
+from pathlib import Path
+from typing import Any
+
+import torch
+from diffusers import DiffusionPipeline
+
+DEFAULT_PROMPTS = [
+    "A dog running across a field of golden wheat.",
+    "An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
+    "A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
+    "A crackling campfire at night under a starry sky, sparks rising into the dark.",
+    "An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
+    "A close-up of a blooming rose covered in morning dew, soft natural light.",
+    "A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
+    "A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
+]
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.")
+    p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
+    p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
+    p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).")
+    p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).")
+    p.add_argument(
+        "--num-frames",
+        type=int,
+        default=49,
+        help="Frames per calibration sample. 49 matches the typical short benchmark; "
+        "use 17 to reduce memory pressure during calibration.",
+    )
+    p.add_argument("--guidance-scale", type=float, default=5.0)
+    p.add_argument(
+        "--calib-steps",
+        type=int,
+        default=10,
+        help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
+    )
+    p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
+    p.add_argument("--seed", type=int, default=42)
+    p.add_argument(
+        "--prompt",
+        action="append",
+        default=[],
+        help="Custom calibration prompt. Repeat to provide multiple.",
+    )
+    p.add_argument(
+        "--quantize-mha",
+        action="store_true",
+        help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention "
+        "sequences amplified FP8 drift in the online ablation (see #2920).",
+    )
+    p.add_argument(
+        "--weight-block-size",
+        type=str,
+        default=None,
+        help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. "
+        "Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 "
+        "for the HV-1.5 investigation status before relying on this.",
+    )
+    p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
+    return p
+
+
+def _parse_block_size(spec: str | None) -> list[int] | None:
+    if spec is None:
+        return None
+    parts = [int(x) for x in spec.split(",") if x.strip()]
+    if len(parts) != 2:
+        raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
+    return parts
+
+
+def _require_modelopt() -> Any:
+    try:
+        import modelopt.torch.quantization as mtq
+    except ModuleNotFoundError as exc:
+        raise SystemExit(
+            "NVIDIA ModelOpt is not installed. Install with:\n"
+            "  pip install 'nvidia-modelopt[all]'\n"
+            f"Original error: {exc}"
+        ) from exc
+    return mtq
+
+
+def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
+    model_path = args.model
+    output_dir = Path(args.output).expanduser().resolve()
+    if output_dir.exists():
+        if not args.overwrite:
+            raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
+        shutil.rmtree(output_dir)
+    return model_path, output_dir
+
+
+def _select_dtype(name: str) -> torch.dtype:
+    return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
+
+
+def _build_prompts(args: argparse.Namespace) -> list[str]:
+    prompts = args.prompt or DEFAULT_PROMPTS
+    if args.calib_size <= 0:
+        raise SystemExit("--calib-size must be positive.")
+    if len(prompts) < args.calib_size:
+        repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
+        prompts = (prompts * repeats)[: args.calib_size]
+    return prompts[: args.calib_size]
+
+
+# Layers to KEEP at full precision. Wan2.2's module naming:
+# - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V)
+# - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip)
+# - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard)
+# - norm_out: AdaLayerNorm final
+# - proj_out: final nn.Linear
+# - timestep_proj_prepare / output_scale_shift_prepare: SP helpers
+def _filter_func_wan22(name: str) -> bool:
+    pattern = re.compile(
+        r"(proj_out.*|"
+        r".*(condition_embedder|patch_embedding|"
+        r"norm_out|scale_shift_table|"
+        r"timestep_proj_prepare|output_scale_shift_prepare).*)"
+    )
+    return pattern.match(name) is not None
+
+
+def _mha_filter_func(name: str) -> bool:
+    pattern = re.compile(
+        r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
+    )
+    return pattern.match(name) is not None
+
+
+def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
+    if not hasattr(mtq, "disable_quantizer"):
+        return
+    mtq.disable_quantizer(backbone, _filter_func_wan22)
+    if not quantize_mha:
+        mtq.disable_quantizer(backbone, _mha_filter_func)
+
+
+def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
+    pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
+    if hasattr(pipe, "set_progress_bar_config"):
+        pipe.set_progress_bar_config(disable=True)
+    pipe.to("cuda")
+    return pipe
+
+
+def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
+    generator = torch.Generator(device="cuda")
+
+    # Try setting guidance on the pipeline's guider if present (newer diffusers APIs).
+    guider = getattr(pipe, "guider", None)
+    if guider is not None and hasattr(guider, "guidance_scale"):
+        try:
+            guider.guidance_scale = args.guidance_scale
+        except Exception:
+            pass
+
+    base_kwargs = dict(
+        height=args.height,
+        width=args.width,
+        num_frames=args.num_frames,
+        num_inference_steps=args.calib_steps,
+        output_type="latent",
+    )
+
+    def forward_loop(*_unused_args, **_unused_kwargs) -> None:
+        with torch.inference_mode():
+            for idx, prompt in enumerate(prompts):
+                generator.manual_seed(args.seed + idx)
+                # Try with guidance_scale first; fall back without on TypeError
+                # for pipelines that take CFG via guider config only.
+                try:
+                    pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
+                except TypeError as exc:
+                    if "guidance_scale" not in str(exc):
+                        raise
+                    pipe(prompt=prompt, generator=generator, **base_kwargs)
+
+    return forward_loop
+
+
+def _summarize_export(output_dir: Path) -> None:
+    cfg_path = output_dir / "transformer" / "config.json"
+    if not cfg_path.exists():
+        print(f"[warn] {cfg_path} missing.", file=sys.stderr)
+        return
+    with cfg_path.open(encoding="utf-8") as f:
+        cfg = json.load(f)
+    qc = cfg.get("quantization_config")
+    if not isinstance(qc, dict):
+        print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
+        return
+    print("Export summary:")
+    print(f"  quant_method: {qc.get('quant_method')}")
+    print(f"  quant_algo:   {qc.get('quant_algo')}")
+    producer = qc.get("producer")
+    if isinstance(producer, dict):
+        print(f"  producer:     {producer.get('name')} {producer.get('version')}")
+    print(f"  config path:  {cfg_path}")
+
+
+def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
+    """Convert in-memory weights of quantized modules to actual FP8 storage.
+
+    `export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't
+    in ModelOpt's recognized-model registry), so we must call the per-weight
+    export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3
+    calibration helpers.
+    """
+    from modelopt.torch.export.quant_utils import (
+        QUANTIZATION_NONE,
+        get_quantization_format,
+        quantizer_attr_names,
+        weight_attr_names,
+    )
+    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+
+    exported = 0
+    for name, module in backbone.named_modules():
+        try:
+            quantization_format = get_quantization_format(module)
+        except Exception as exc:
+            print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
+            continue
+        if quantization_format == QUANTIZATION_NONE:
+            continue
+        for weight_name in weight_attr_names(module):
+            quantizer_attrs = quantizer_attr_names(weight_name)
+            weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
+            if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
+                continue
+            _export_quantized_weight(module, dtype, weight_name)
+            exported += 1
+    return exported
+
+
+def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
+    """Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913)."""
+    weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
+    if weight_block_size is not None:
+        weights_cfg["strategy"] = "block"
+        weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
+    return {
+        "config_groups": {
+            "group_0": {
+                "input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
+                "weights": weights_cfg,
+                "targets": ["Linear"],
+            }
+        },
+        "ignore": [
+            "condition_embedder*",
+            "norm_out*",
+            "output_scale_shift_prepare*",
+            "patch_embedding*",
+            "proj_out*",
+            "scale_shift_table*",
+            "timestep_proj_prepare*",
+        ],
+        "producer": {"name": "modelopt"},
+        "quant_algo": "FP8",
+        "quant_method": "modelopt",
+    }
+
+
+def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
+    """Inject quant_algo: FP8 + config_groups into each transformer's config.json
+    so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.
+
+    Patches both ``transformer`` and ``transformer_2`` (A14B) when present."""
+    for sub in ("transformer", "transformer_2"):
+        cfg_path = output_dir / sub / "config.json"
+        if not cfg_path.exists():
+            continue
+        with cfg_path.open(encoding="utf-8") as f:
+            cfg = json.load(f)
+
+        new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size)
+        existing = cfg.get("quantization_config")
+        if isinstance(existing, dict):
+            producer = existing.get("producer")
+            if isinstance(producer, dict):
+                new_qc["producer"] = producer
+
+        cfg["quantization_config"] = new_qc
+        with cfg_path.open("w", encoding="utf-8") as f:
+            json.dump(cfg, f, indent=2)
+
+
+def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]:
+    """Return [(attr_name, module)] for every DiT on the pipeline.
+
+    Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2``
+    (high-noise). TI2V-5B has only ``transformer``.
+    """
+    out = []
+    for attr in ("transformer", "transformer_2"):
+        mod = getattr(pipe, attr, None)
+        if mod is not None:
+            out.append((attr, mod))
+    return out
+
+
+def _save_pipeline_with_fp8_transformer(
+    pipe: DiffusionPipeline,
+    model_path: str,
+    output_dir: Path,
+    max_shard_size: str = "5GB",
+) -> None:
+    """Copy source dir verbatim minus transformer dirs, then save each quantized DiT.
+
+    Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well
+    as the single-DiT TI2V-5B.
+    """
+    from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
+
+    src = Path(model_path)
+    if not src.exists():
+        from huggingface_hub import snapshot_download
+
+        src = Path(snapshot_download(model_path))
+
+    if output_dir.exists():
+        shutil.rmtree(output_dir)
+    shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2"))
+
+    for attr, backbone in _list_transformers(pipe):
+        transformer_out = output_dir / attr
+        # Pass the nn.Module (transformer), not the Pipeline wrapper.
+        with hide_quantizers_from_state_dict(backbone):
+            backbone.save_pretrained(
+                str(transformer_out),
+                safe_serialization=True,
+                max_shard_size=max_shard_size,
+            )
+
+
+def main() -> None:
+    args = _build_parser().parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
+
+    mtq = _require_modelopt()
+    model_path, output_dir = _ensure_paths(args)
+    dtype = _select_dtype(args.dtype)
+    prompts = _build_prompts(args)
+    weight_block_size = _parse_block_size(args.weight_block_size)
+
+    print("Quantization plan:")
+    print(f"  input:           {args.model}")
+    print(f"  output:          {output_dir}")
+    print(f"  dtype:           {dtype}")
+    print(f"  height/width:    {args.height}x{args.width}")
+    print(f"  num_frames:      {args.num_frames}")
+    print(f"  calib_size:      {len(prompts)}")
+    print(f"  calib_steps:     {args.calib_steps}")
+    print(f"  quantize_mha:    {args.quantize_mha}")
+    print(
+        f"  weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
+    )
+
+    pipe = _load_pipeline(model_path, dtype)
+    transformers = _list_transformers(pipe)
+    if not transformers:
+        raise SystemExit("Pipeline has no transformer or transformer_2 attribute.")
+    print(f"  found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}")
+
+    quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
+    if weight_block_size is not None:
+        quant_config["quant_cfg"]["*weight_quantizer"] = {
+            "num_bits": (4, 3),
+            "block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
+        }
+        print(
+            f"  -> overriding weight quantizer with block_sizes={weight_block_size} "
+            f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
+        )
+
+    forward_loop = _build_forward_loop(pipe, args, prompts)
+
+    # Quantize each DiT. The forward_loop runs the full pipeline, so quantizers
+    # on every attached transformer (A14B: transformer + transformer_2) get
+    # calibrated during the same passes.
+    for attr, backbone in transformers:
+        print(f"\nQuantizing {attr}...")
+        quantized = mtq.quantize(backbone, quant_config, forward_loop)
+        if quantized is not None:
+            setattr(pipe, attr, quantized)
+        _disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha)
+
+    # Export weights to FP8 only AFTER all transformers are calibrated. Exporting
+    # inside the loop would convert transformer's weights to real FP8, and the
+    # next transformer's forward_loop (full pipeline) would then fake-quantize an
+    # already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn".
+    print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,")
+    print("so we have to call the per-weight export helper ourselves)...")
+    total_exported = 0
+    for attr, backbone in _list_transformers(pipe):
+        exported = _force_export_quantized_weights(backbone, dtype)
+        print(f"  -> {exported} weights converted to FP8 in {attr}")
+        total_exported += exported
+
+    if total_exported == 0:
+        raise SystemExit(
+            "No quantized weights were exported. Calibration may have skipped every layer "
+            "(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
+            "weight quantizers."
+        )
+
+    print("\nSaving pipeline with FP8 transformer(s)...")
+    _save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
+    _patch_quant_config(output_dir, weight_block_size=weight_block_size)
+    print(f"Saved to: {output_dir}")
+    _summarize_export(output_dir)
+
+    print("\nNext: validate the checkpoint with vllm-omni:")
+    print(
+        "  python examples/offline_inference/text_to_video/text_to_video.py \\\n"
+        f"    --model {output_dir} \\\n"
+        "    --quantization fp8 \\\n"
+        "    --prompt 'A dog running across a field of golden wheat.' \\\n"
+        f"    --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
+        "    --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n"
+        "    --output outputs/wan22_modelopt_fp8.mp4"
+    )
+    print(
+        "\n  (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
+        "checkpoint's config.json has modelopt metadata.)"
+    )
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -64,7 +64,9 @@ dev = [
     "lpips>=0.1.4", # for quantization fp8
     "orjson>=3.11.7", # to make v1/embeddings API fast
     "voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests
-    "torchdiffeq>=0.2.5" # for covo_audio Test
+    "torchdiffeq>=0.2.5", # for covo_audio Test
+    "FlagEmbedding", # for hunyuan_image3 tests
+    "kernels<0.15.0" # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
 ]
 
 demo = [
@@ -88,6 +90,12 @@ minicpmo = [
     "stepaudio2-minicpmo",
 ]
 
+# Optional Blackwell-only fused-bias FP8 GEMM; package is `quack-kernels` (imported
+# as `quack`), not the unrelated `quack`. See docs/user_guide/quantization/fp8.md.
+quack = [
+    "quack-kernels>=0.3.11",
+]
+
 docs = [
     "mkdocs>=1.5.0",
     "mkdocs-api-autonav",

recipes/BosonAI/Higgs-Audio-V3-TTS.md

@@ -0,0 +1,110 @@
+# Higgs-Audio V3 TTS
+
+> Multilingual text-to-speech with voice cloning on 1×H100
+
+## Summary
+
+- Vendor: Boson AI
+- Model: `bosonai/higgs-audio-v3-tts-4b`
+- Task: Text-to-speech synthesis with optional voice cloning (100+ languages)
+- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API; offline batch inference
+- Maintainer: @yuekaiz
+
+## When to use this recipe
+
+Use this recipe to serve `bosonai/higgs-audio-v3-tts-4b` for high-quality
+multilingual TTS. The model generates 24 kHz speech, supports zero-shot voice
+cloning from a reference clip, and handles 100+ languages with inline control
+tokens for emotion, style, and prosody. The architecture is a ~4B Qwen3
+backbone with fused multi-codebook embedding/head (8 codebooks × 1026 vocab,
+MusicGen-style delay pattern).
+
+## References
+
+- Model card: [bosonai/higgs-audio-v3-tts-4b](https://huggingface.co/bosonai/higgs-audio-v3-tts-4b)
+- Offline example: [`examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py`](../../examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py)
+- Online example: [`examples/online_serving/text_to_speech/higgs_audio_v3/`](../../examples/online_serving/text_to_speech/higgs_audio_v3/)
+- Benchmark results: see Performance section below
+
+## Hardware Support
+
+## GPU
+
+### 1×H100 80GB
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12+
+- CUDA: 12.x
+- vLLM version: 0.22.0
+- vLLM-Omni version or commit: `36e048fd` (branch `higgs-v3`)
+
+#### Command
+
+**Online serving:**
+
+```bash
+vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
+    --host 0.0.0.0 --port 8095 \
+    --trust-remote-code --omni
+```
+
+The default deploy config `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` is
+loaded automatically by model registry (HF `model_type=higgs_multimodal_qwen3`).
+
+**Offline batch inference:**
+
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
+    --texts "Hello world." "The quick brown fox jumps over the lazy dog." \
+    --output-dir results/higgs_v3_wavs
+```
+
+**Offline voice clone:**
+
+```bash
+python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
+    --texts "Text to synthesize in the cloned voice." \
+    --ref-audio path/to/reference.wav \
+    --ref-text "Transcript of the reference clip." \
+    --output-dir results/higgs_v3_clone
+```
+
+#### Verification
+
+Basic TTS via curl:
+
+```bash
+curl -X POST http://localhost:8095/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "bosonai/higgs-audio-v3-tts-4b",
+        "input": "Hello, how are you?"
+    }' --output hello.wav
+```
+
+Voice clone via Python client:
+
+```bash
+python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
+    --base-url http://localhost:8095 \
+    --model bosonai/higgs-audio-v3-tts-4b \
+    --ref-audio path/to/reference.wav \
+    --ref-text "Transcript of the reference." \
+    --prompts "Text to clone."
+```
+
+#### Notes
+
+- Memory usage: Stage 0 (talker, ~4B) uses ~60% GPU memory; Stage 1 (codec decoder) uses ~25%.
+- Key flags: `--trust-remote-code` and `--omni` are required.
+- Output: 24 kHz mono WAV.
+- Voice cloning: `ref_audio` accepts WAV/FLAC/MP3; `ref_text` is optional but improves fidelity.
+- Deploy config: `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` (auto-discovered from `model_type`).
+  - `max_num_seqs=16` for both stages.
+  - `enforce_eager=true` for both stages (CUDA graph available for stage 0 but no throughput gain at batch>1).
+- Known limitations:
+  - CUDA graph for stage 0 helps batch=1 (-13% RTF) but not batch>1 due to model-owned sampler running outside graph.
+  - Stage 1 (code2wav) must use `enforce_eager=true` (`@torch.inference_mode` incompatible with graph capture).
+  - Async streaming (chunk-based) not yet implemented; sync-only pipeline.

recipes/README.md

@@ -36,8 +36,8 @@ recipes/
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
-| [`nvidia/Cosmos3-Nano.md`](./nvidia/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video generation, text to video with sound  | 1x H200 141GB / B300 |
-| [`nvidia/Cosmos3-Super.md`](./nvidia/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) | 8x H200/H100/A100 / 2x H200 / B300 |
+| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
+| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
 | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
 | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
@@ -50,6 +50,7 @@ recipes/
 | [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
 | [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
 | [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
+| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
 | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
 | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
 |

recipes/StabilityAI/Stable-Diffusion-3.5.md

@@ -0,0 +1,139 @@
+# Stable Diffusion 3.5
+
+> Text-to-image serving with Stability AI's SD 3.5 family
+
+## Summary
+
+- Vendor: Stability AI
+- Model: `stabilityai/stable-diffusion-3.5-medium`, `stabilityai/stable-diffusion-3.5-large`
+- Task: Text-to-image generation
+- Mode: Online serving with the OpenAI-compatible images API
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe when you want a lightweight, fast text-to-image model that
+fits on a single consumer or workstation GPU. SD 3.5-medium requires only
+~17 GiB VRAM and generates 1024×1024 images in ~4 seconds, making it one of
+the most accessible diffusion models in vLLM-Omni. SD 3.5-large produces
+noticeably sharper images at the cost of ~34 GiB VRAM and ~21 seconds per
+image.
+
+## References
+
+- Model (medium): <https://huggingface.co/stabilityai/stable-diffusion-3.5-medium>
+- Model (large): <https://huggingface.co/stabilityai/stable-diffusion-3.5-large>
+- Related example under `examples/`:
+  [`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md)
+- User guide:
+  [`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md)
+
+## Hardware Support
+
+This recipe documents two CUDA GPU configurations on a single RTX A6000 48 GB,
+one for each model variant. Extend it with more hardware sections as community
+validation lands.
+
+## GPU
+
+### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-medium
+
+#### Environment
+
+- OS: Ubuntu 22.04
+- Python: 3.12
+- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
+- vLLM version: 0.22.0
+- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
+
+#### Command
+
+```bash
+vllm serve stabilityai/stable-diffusion-3.5-medium --omni --port 8091
+```
+
+#### Verification
+
+```bash
+curl -X POST http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "stabilityai/stable-diffusion-3.5-medium",
+    "prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus",
+    "negative_prompt": "blurry, out of focus",
+    "guidance_scale": 4.5,
+    "n": 1,
+    "size": "1024x1024"
+  }' | python3 -c "
+import json, base64, sys
+data = json.load(sys.stdin)
+with open('output.png', 'wb') as f:
+    f.write(base64.b64decode(data['data'][0]['b64_json']))
+print('saved output.png')
+"
+```
+
+Or use the example client:
+
+```bash
+python examples/online_serving/text_to_image/openai_chat_client.py \
+  --server http://localhost:8091 \
+  --prompt "a cute cat sitting on a windowsill" \
+  --output output.png
+```
+
+#### Notes
+
+- Memory usage: Model weights ~15.6 GiB, peak VRAM ~17.0 GiB (no CFG), ~20.9 GiB with `guidance_scale > 1` (CFG enabled).
+- Generation time: ~4 seconds without CFG, ~8 seconds with `guidance_scale=4.5`.
+- Model loading time: ~55 seconds (including weight download on first run).
+- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG).
+
+### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-large
+
+#### Environment
+
+- OS: Ubuntu 22.04
+- Python: 3.12
+- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
+- vLLM version: 0.22.0
+- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
+
+#### Command
+
+```bash
+vllm serve stabilityai/stable-diffusion-3.5-large --omni --port 8091
+```
+
+> **Note**: `stable-diffusion-3.5-large` is a gated model. Accept the license
+> on [Hugging Face](https://huggingface.co/stabilityai/stable-diffusion-3.5-large)
+> and set `HF_TOKEN` before running.
+
+#### Verification
+
+```bash
+curl -X POST http://localhost:8091/v1/images/generations \
+  -H "Content-Type: application/json" \
+  -d '{
+    "model": "stabilityai/stable-diffusion-3.5-large",
+    "prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus",
+    "negative_prompt": "blurry, out of focus",
+    "guidance_scale": 4.5,
+    "n": 1,
+    "size": "1024x1024"
+  }' | python3 -c "
+import json, base64, sys
+data = json.load(sys.stdin)
+with open('output.png', 'wb') as f:
+    f.write(base64.b64decode(data['data'][0]['b64_json']))
+print('saved output.png')
+"
+```
+
+#### Notes
+
+- Memory usage: Peak VRAM ~31.6 GiB at 1024×1024, ~45.5 GiB at 2048×2048 with `guidance_scale=4.5`.
+- Generation time: ~19 seconds at 1024×1024, ~102 seconds at 2048×2048.
+- Image quality is noticeably better than medium, with sharper details.
+- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG).
+- 2048×2048 fits within 48 GB but leaves little headroom (~45.5 GiB used out of 49 GiB).

recipes/cosmos3/Cosmos3-Nano.md

@@ -1,28 +1,55 @@
 # Cosmos3-Nano
 
-> Text-to-image, text-to-video, and image-to-video serving
+> Text-to-image, text-to-video, image-to-video, and video-to-video serving
 
 ## Summary
 
 - Vendor: NVIDIA
 - Model: `nvidia/Cosmos3-Nano`
-- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation, with optional synchronized audio (video + sound)
+- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional synchronized audio (video + sound), action policy
 - Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API
 - Maintainer: Community
 
 ## When to use this recipe
 
 Use this recipe to deploy `nvidia/Cosmos3-Nano` for image and video generation.
-A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves all three modes;
-the mode is selected per request:
+A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves these modes; the
+mode is selected per request:
 
 - **T2I** — `POST /v1/images/generations` (or a prompt carrying `modalities=["image"]`).
-- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image.
+- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image/video.
 - **I2V** — `POST /v1/videos/sync` with a reference image (`input_reference` file
   upload, or `image_reference` JSON).
+- **V2V** — `POST /v1/videos/sync` with a reference video (`input_reference`
+  video upload, or `video_reference` JSON). Cosmos3 conditions on selected
+  reference-video latent frames; use `extra_params.condition_frame_indexes_vision`
+  and `extra_params.condition_video_keep` to choose which prefix/tail frames guide
+  generation.
 - **T2VS / I2VS** — add `generate_sound=true` (and optional `sound_duration`) to a
   T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
   the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
+- **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks:
+  - `forward_dynamics` — given a first frame or video **and** an action trajectory,
+    roll out the resulting video. Synchronous: `POST /v1/videos/sync`.
+  - `policy` — given a first frame or video and a language instruction,
+    **predict** the action trajectory (and a rollout video). Use the async
+    `POST /v1/videos` endpoint and read the predicted action from the top-level
+    `action` field.
+  - `inverse_dynamics` — given a video, **recover** the action trajectory. Use
+    the async `POST /v1/videos` endpoint and read the recovered action from
+    the top-level `action` field
+    (`{data, shape, dtype, raw_action_dim, domain_id}`).
+
+  Action requests also take `domain_name` (e.g. `av`, `bridge_orig_lerobot`,
+  `droid_lerobot`, `agibotworld`, …; or a numeric `domain_id`), `raw_action_dim`,
+  and `action_chunk_size` (must equal `num_frames` or `num_frames - 1`). For
+  `forward_dynamics` also pass the `action` array. The dedicated policy checkpoint
+  **`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way
+  (`domain_name=droid_lerobot`).
+
+  Action requests can use `input_reference` or `video_reference` for video input.
+  `policy` and `forward_dynamics` can also use an image reference; `inverse_dynamics`
+  requires a video reference.
 
 ## References
 
@@ -125,6 +152,36 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
   -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
   -o cosmos3_i2v.mp4
 
+# Video-to-video -> /v1/videos/sync with an uploaded reference video.
+# By default Cosmos3 conditions on latent indexes [0, 1]. For the default
+# temporal VAE stride this decodes only the first 5 input frames.
+# The model works best when the prompt describes the actual situation happening in the video.
+# Generic prompts may create sub-standard generations.
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
+  -F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" \
+  -F "num_inference_steps=35" -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true,"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
+  -F "seed=2222" \
+  -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
+  -o cosmos3_v2v.mp4
+
+# V2V can also use a JSON-safe URL/data-URL video reference. Do not combine
+# video_reference with input_reference or image_reference.
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Nano" \
+  -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" \
+  -F "num_inference_steps=35" -F "guidance_scale=6.0" \
+  -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"last"}' \
+  -F 'video_reference={"video_url":"https://example.com/reference.mp4"}' \
+  -o cosmos3_v2v_from_url.mp4
 
 # Text-to-video-with-sound
 curl -sS -X POST http://localhost:8000/v1/videos/sync \
@@ -144,6 +201,50 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
   -F "sound_duration=7.875" \
   -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
   -o cosmos3_t2v_with_sound.mp4
+
+# Action — forward dynamics (first frame + action trajectory -> rollout video).
+# Synchronous; `action` is a JSON array shaped [action_chunk_size, raw_action_dim].
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  --form-string "model=nvidia/Cosmos3-Nano" \
+  --form-string "prompt=You are an autonomous vehicle. This video is captured from a first-person perspective." \
+  -F "input_reference=@first_frame.jpg;type=image/jpeg" \
+  -F "size=640x480" -F "num_frames=61" -F "fps=10" \
+  -F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
+  --form-string "extra_params={\"action_mode\":\"forward_dynamics\",\"domain_name\":\"av\",\"raw_action_dim\":9,\"action_chunk_size\":60,\"action\":$(cat action.json)}" \
+  -F "seed=0" \
+  -o cosmos3_forward_dynamics.mp4
+
+# Action — policy (first frame + instruction -> predicted action trajectory + video).
+# Asynchronous: POST returns a job id; poll, then read the predicted action from
+# the top-level `action` field ({data, shape, dtype, raw_action_dim, domain_id}).
+VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
+  -H "Accept: application/json" \
+  --form-string "model=nvidia/Cosmos3-Nano" \
+  --form-string "prompt=Pick up the banana and place it in the bowl." \
+  -F "input_reference=@first_frame.jpg;type=image/jpeg" \
+  -F "size=640x480" -F "num_frames=17" -F "fps=5" \
+  -F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
+  --form-string 'extra_params={"action_mode":"policy","domain_name":"bridge_orig_lerobot","raw_action_dim":10,"action_chunk_size":16}' \
+  -F "seed=0" | jq -r '.id')
+# poll until status == completed, then:
+curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
+curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_policy.mp4
+
+# Action — inverse dynamics (video -> recovered action trajectory).
+# Asynchronous: use the job metadata to read the recovered action.
+VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
+  -H "Accept: application/json" \
+  --form-string "model=nvidia/Cosmos3-Nano" \
+  --form-string "prompt=Recover the robot action trajectory from this clip." \
+  -F "input_reference=@motion_clip.mp4;type=video/mp4" \
+  -F "size=640x480" -F "num_frames=17" -F "fps=5" \
+  -F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
+  --form-string 'extra_params={"action_mode":"inverse_dynamics","domain_name":"bridge_orig_lerobot","raw_action_dim":10,"action_chunk_size":16}' \
+  -F "seed=0" | jq -r '.id')
+# poll until status == completed, then:
+curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
+curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inverse_dynamics.mp4
 ```
 
 #### Notes
@@ -152,6 +253,7 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
   - T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s**
   - T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s**
   - I2V 1280×720, 189 frames @ 35 steps → ~**99 s**
+  - Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**.
   - Guardrails-on overhead: ~8% on T2I, negligible on video.
 - **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 GPU;
   full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
@@ -159,22 +261,27 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
 - **Determinism:** identical seed reproduces identical output on the same
   hardware; outputs are not bit-identical across different GPU types.
 - **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1,
-  3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V 1280×720,
-  189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
+  3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V/V2V
+  1280×720, 189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
 - **Key flags / params:** `--no-guardrails` (server) or
   `extra_params={"guardrails":false}` (per request) toggles safety. The
   per-request flag only takes effect when the server was launched **with**
   guardrails enabled (it cannot re-enable them on a `--no-guardrails` server).
   `use_resolution_template` / `use_duration_template` are off by default and only
   needed when not using upsampled prompts that already encode resolution/duration.
+  For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent
+  frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the
+  API decodes the first or last needed reference frames (`"first"` by default).
 - **Known limitations:**
   - Guardrails-on requires `cosmos-guardrail` **and** access to the gated
     `nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise
     the server fails at pipeline build with a gated-repo / safety-checker error.
   - A guardrail-blocked prompt currently returns HTTP 500
     (`"Guardrail blocked prompt"`).
-  - Action (policy / forward- / inverse-dynamics) modalities are not part of
-    this integration yet.
+  - Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
+    online. Use async `POST /v1/videos` when you need the predicted/recovered
+    action payload under the top-level `action` field; sync `/v1/videos/sync`
+    returns raw MP4 bytes and does not expose action metadata in the response body.
 
 ### 1x GPU (Offline generation)
 
@@ -238,8 +345,9 @@ python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image
 
 - Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by
   `prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus
-  `num_frames`/`fps`, and `multi_modal_data={"image": ...}` for I2V. For video,
-  frames are returned in `outputs[0].request_output.images` as an
+  `num_frames`/`fps`, `multi_modal_data={"image": ...}` for I2V, and
+  `multi_modal_data={"video": [<PIL frames>]}` or a video tensor/array for V2V.
+  For video, frames are returned in `outputs[0].request_output.images` as an
   `(B, F, H, W, 3)` array.
 - The offline entry must be guarded by `if __name__ == "__main__":` — the engine
   spawns workers with the `spawn` start method.

recipes/cosmos3/Cosmos3-Super.md

@@ -1,12 +1,12 @@
 # Cosmos3-Super
 
-> Frontier 64B world model: text-to-image, text-to-video, image-to-video (+ optional audio)
+> Frontier 64B world model: text-to-image, text-to-video, image-to-video, video-to-video (+ optional audio)
 
 ## Summary
 
 - Vendor: NVIDIA
 - Model: `nvidia/Cosmos3-Super` (64B; also `Cosmos3-Super-Text2Image`, `Cosmos3-Super-Image2Video`)
-- Task: T2I, T2V, I2V generation, with optional synchronized audio (video + sound)
+- Task: T2I, T2V, I2V, V2V generation, with optional synchronized audio (video + sound)
 - Mode: Online serving with the OpenAI-compatible image/video APIs
 - Maintainer: Community
 
@@ -16,8 +16,8 @@ Use this recipe to deploy the 64B `nvidia/Cosmos3-Super` for the highest-quality
 Cosmos3 generation. It shares the same `Cosmos3OmniDiffusersPipeline` and request
 formats as [Cosmos3-Nano](./Cosmos3-Nano.md) — only the checkpoint size and the
 recommended parallelism differ. Mode is selected per request (T2I →
-`/v1/images/generations`; T2V/I2V → `/v1/videos/sync`; add `generate_sound=true`
-for audio).
+`/v1/images/generations`; T2V/I2V/V2V → `/v1/videos/sync`; add
+`generate_sound=true` for audio).
 
 ## References
 
@@ -62,8 +62,9 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs.
 #### Verification
 
 Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full
-T2I/T2V/I2V/T2VS curls); official params: `size=1280x720, num_frames=189, fps=24,
-num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0, max_sequence_length=4096`.
+T2I/T2V/I2V/V2V/T2VS curls); official params: `size=1280x720, num_frames=189,
+fps=24, num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0,
+max_sequence_length=4096`.
 
 ```bash
 curl http://localhost:8000/v1/models
@@ -84,6 +85,15 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
   -F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
   -o cosmos3_super_i2v.mp4
 
+# V2V — add an uploaded reference video. condition_video_keep can be "first" or "last".
+curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
+  -F "model=nvidia/Cosmos3-Super" -F "prompt=Continue the same scene with smooth natural motion." \
+  -F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
+  -F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
+  -F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
+  -F "seed=2222" -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
+  -o cosmos3_super_v2v.mp4
+
 # T2V + sound — add generate_sound/sound_duration (output muxes AAC 48 kHz stereo)
 curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
   -F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
@@ -103,6 +113,12 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
   - T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
   - (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
 - **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
-- Same generation defaults, supported sizes, and `generate_sound`/`sound_duration`
-  semantics as Nano. Action (policy / forward- / inverse-dynamics) modalities are
-  not part of this integration yet.
+- Same generation defaults, supported sizes, V2V reference-video controls
+  (`condition_frame_indexes_vision`, `condition_video_keep`), and
+  `generate_sound`/`sound_duration`
+  semantics as Nano, including the **action** modality: `forward_dynamics`,
+  `policy`, and `inverse_dynamics` — see the Cosmos3-Nano recipe for the request
+  shapes. Use async `/v1/videos` when you need predicted/recovered action metadata
+  under the top-level `action` field. Verified on the 64B Super under
+  `--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`)
+  and the rollout video reliably.

recipes/inclusionAI/Ming-flash-omni-2.0.md

@@ -57,12 +57,6 @@ Adjust `devices` in the YAML to match your hardware.
 
 #### Command
 
-Thinker only (text output):
-
-```bash
-vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091
-```
-
 Thinker + talker (text and/or audio output):
 
 ```bash
@@ -72,6 +66,15 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 \
     --log-stats
 ```
 
+Thinker only (text-only output):
+
+```bash
+vllm serve Jonathan1909/Ming-flash-omni-2.0 \
+    --omni \
+    --deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml \
+    --port 8091
+```
+
 `--log-stats` is optional but recommended while validating the deployment.
 
 #### Verification

tests/benchmarks/test_accuracy_bench_utils.py

@@ -345,6 +345,14 @@ def test_parse_score_payload_handles_raw_json_and_delimited_json():
     assert parse_score_payload(wrapped)["score"] == [6]
 
 
+def test_parse_score_payload_handles_qwen_vl_nested_score_dicts():
+    qwen_style = '{"score": [{"naturalness": 8, "artifact_free": 7}], "reasoning": "good quality"}'
+    assert parse_score_payload(qwen_style)["score"] == [8, 7]
+
+    nested_values = '{"score": [{"score": 9}, {"score": 6}], "reasoning": "ok"}'
+    assert parse_score_payload(nested_values)["score"] == [9, 6]
+
+
 def test_summarize_gedit_generated_records_groups_by_task_and_language():
     records = []
     for group in GEDIT_GROUPS[:2]:

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -61,7 +61,7 @@ def test_stage0_streaming_update_discards_outstanding_async_placeholder_token()
 
     sched._update_request_as_session(session, _make_update([10, 20]))
 
-    assert session.discard_latest_async_tokens is True
+    assert session.async_tokens_to_discard == 1
     assert session.num_output_placeholders == 0
     assert session.spec_token_ids == []
     # The async placeholder makes token 9 unconfirmed, so only 7 and 8 are
@@ -83,7 +83,7 @@ def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder()
 
     sched._update_request_as_session(session, _make_update([10, 20]))
 
-    assert session.discard_latest_async_tokens is False
+    assert getattr(session, "async_tokens_to_discard", 0) == 0
     assert session.num_output_placeholders == 0
     assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20]
     assert list(session._all_token_ids) == [1, 2, 3, 7, 8, 9, 10, 20]

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

@@ -1,7 +1,7 @@
 [
     {
-        "test_name": "test_hunyuan_image3_ti2i_tp4_tp4_inline",
-        "description": "AR (TP=4) + DiT (TP=4) - using inline dataset and deploy-config",
+        "test_name": "test_hunyuan_image3_ti2i_inline",
+        "description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config",
         "server_type": "vllm-omni",
         "benchmark_endpoint": "/v1/images/edits",
         "server_params": {
@@ -21,7 +21,7 @@
                                     "zmq_port": 50051,
                                     "protocol": "rdma",
                                     "device_name": "",
-                                    "memory_pool_size": 4294967296,
+                                    "memory_pool_size": 17179869184,
                                     "memory_pool_device": "cpu"
                                 }
                             }
@@ -32,13 +32,13 @@
                                 "is_comprehension": true,
                                 "final_output": true,
                                 "final_output_type": "text",
-                                "max_num_seqs": 4,
+                                "max_num_seqs": 16,
                                 "gpu_memory_utilization": 0.95,
                                 "enforce_eager": false,
                                 "enable_prefix_caching": false,
                                 "max_num_batched_tokens": 32768,
-                                "devices": "0,1,2,3",
-                                "tensor_parallel_size": 4,
+                                "devices": "0,1",
+                                "tensor_parallel_size": 2,
                                 "hf_overrides": {
                                     "rope_parameters": {
                                         "mrope_section": [0, 32, 32],
@@ -65,13 +65,14 @@
                                 "stage_id": 1,
                                 "max_num_seqs": 1,
                                 "enforce_eager": true,
-                                "devices": "4,5,6,7",
+                                "devices": "2,3",
                                 "distributed_executor_backend": "mp",
                                 "omni_kv_config": {
                                     "need_recv_cache": true
                                 },
                                 "parallel_config": {
-                                    "tensor_parallel_size": 4,
+                                    "tensor_parallel_size": 2,
+                                    "vae_patch_parallel_size": 2,
                                     "enable_expert_parallel": true
                                 },
                                 "input_connectors": {
@@ -114,8 +115,8 @@
                 },
                 "task": "ti2i",
                 "bot-task": "think_recaption",
-                "max-concurrency": 4,
-                "num-prompts": 16,
+                "max-concurrency": 16,
+                "num-prompts": 96,
                 "skip-performance-assertion": true,
                 "baseline": {
                     "throughput_qps": 0.1,

tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json

@@ -23,10 +23,10 @@
                 "ignore_eos": true,
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [1000, 2000, 3000, 4000],
-                    "mean_tpot_ms": [200, 200, 200, 200],
-                    "mean_audio_ttfp_ms": [10000, 15000, 20000, 25000],
-                    "mean_audio_rtf": [1.5, 1.5, 1.5, 1.5]
+                    "mean_ttft_ms": [203, 283, 470, 869],
+                    "mean_tpot_ms": [28, 36, 41, 45],
+                    "mean_audio_ttfp_ms": [861, 1200, 1528, 2100],
+                    "mean_audio_rtf": [0.24, 0.30, 0.35, 0.39]
                 }
             },
             {
@@ -49,9 +49,9 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [2000],
-                    "mean_audio_ttfp_ms": [2000],
-                    "mean_audio_rtf": [0.25]
+                    "mean_ttft_ms": [377],
+                    "mean_audio_ttfp_ms": [684],
+                    "mean_audio_rtf": [0.14]
                 }
             },
             {
@@ -76,9 +76,9 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [6000],
-                    "mean_audio_ttfp_ms": [6000],
-                    "mean_audio_rtf": [0.7]
+                    "mean_ttft_ms": [228],
+                    "mean_audio_ttfp_ms": [686],
+                    "mean_audio_rtf": [0.16]
                 }
             },
             {
@@ -105,9 +105,9 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [12000],
-                    "mean_audio_ttfp_ms": [12000],
-                    "mean_audio_rtf": [1.0]
+                    "mean_ttft_ms": [522],
+                    "mean_audio_ttfp_ms": [1119],
+                    "mean_audio_rtf": [0.21]
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

@@ -31,11 +31,6 @@
                 ],
                 "random_input_len": 2500,
                 "random_output_len": 900,
-                "temperature": 0.4,
-                "top_p": 0.9,
-                "top_k": 1,
-                "seed": 42,
-                "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
                     "modalities": [
@@ -74,11 +69,6 @@
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
-                "temperature": 0.4,
-                "top_p": 0.9,
-                "top_k": 1,
-                "seed": 42,
-                "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
                     "modalities": [
@@ -117,11 +107,6 @@
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
-                "temperature": 0.4,
-                "top_p": 0.9,
-                "top_k": 1,
-                "seed": 42,
-                "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
                     "modalities": [
@@ -162,11 +147,6 @@
                 "random_input_len": 100,
                 "random_output_len": 100,
                 "random_range_ratio": 0.0,
-                "temperature": 0.4,
-                "top_p": 0.9,
-                "top_k": 1,
-                "seed": 42,
-                "repetition_penalty": 1.05,
                 "ignore_eos": true,
                 "extra_body": {
                     "modalities": [

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -1,19 +1,46 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+"""
+Tests for the DiffusersPipelineLoader.
+"""
+
 from types import SimpleNamespace
 
 import pytest
 import torch
 import torch.nn as nn
+from huggingface_hub import snapshot_download
+from vllm.config.load import LoadConfig
 
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
 from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
+from vllm_omni.diffusion.models.helios import HeliosPipeline
 from vllm_omni.diffusion.registry import initialize_model
 
 pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
 
+model_path = "hf-internal-testing/tiny-helios-modular-pipe"
+
+
+@pytest.fixture(scope="module")
+def prefetch_helios_model():
+    """Downloads the tiny helios model prior to running a test."""
+    snapshot_download(model_path)
+
+
+@pytest.fixture(scope="function")
+def mock_tp_group(mocker):
+    """Mocks the tensor parallel group; this is needed to initialize the Helios model."""
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
+    mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
+    mock_group = mocker.MagicMock()
+    mock_group.world_size = 1
+    mock_group.rank_in_group = 0
+    mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
+
 
 class _DummyPipelineModel(nn.Module):
     def __init__(self, *, source_prefix: str):
@@ -42,7 +69,13 @@ class _DummyPipelineModel(nn.Module):
 
 
 def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
-    loader = object.__new__(DiffusersPipelineLoader)
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(use_hsdp=False),
+        quantization_config=None,
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+
     loader.counter_before_loading_weights = 0.0
     loader.counter_after_loading_weights = 0.0
 
@@ -150,16 +183,15 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
         quantization_config=None,
     )
 
-    loader = object.__new__(DiffusersPipelineLoader)
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
     loader.load_weights = lambda model: None  # type: ignore[assignment]
     loader._process_weights_after_loading = lambda model, target_device: None  # type: ignore[assignment]
-    loader._is_gguf_quantization = lambda _od_config: False  # type: ignore[assignment]
+    loader._is_gguf_quantization = lambda: False  # type: ignore[assignment]
 
     monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
     monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
 
     model = loader.load_model(
-        od_config,
         load_device="cpu",
         load_format="custom_pipeline",
         custom_pipeline_name="tests.dummy.ConfigAwarePipeline",
@@ -168,3 +200,33 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
     assert model.captured_config is od_config
     assert model.seen_config_during_init is od_config
     assert get_current_diffusion_config_or_none() is None
+
+
+def test_get_all_weights(prefetch_helios_model, mock_tp_group):
+    """Ensure that get all weights on a tiny model resolves to nonempty weights."""
+    od_config = OmniDiffusionConfig(
+        model_class_name="HeliosPipeline",
+        model=model_path,
+    )
+    loader = DiffusersPipelineLoader(
+        load_config=LoadConfig(),
+        od_config=od_config,
+    )
+    pipeline = HeliosPipeline(od_config=od_config)
+
+    weights = list(loader.get_all_weights(pipeline))
+    assert len(weights) > 0
+
+
+def test_load_model(prefetch_helios_model, mock_tp_group):
+    """Ensure that load model creates an instance of the expected pipeline class."""
+    od_config = OmniDiffusionConfig(
+        model_class_name="HeliosPipeline",
+        model=model_path,
+    )
+    loader = DiffusersPipelineLoader(
+        load_config=LoadConfig(),
+        od_config=od_config,
+    )
+    model = loader.load_model(load_device="cpu")
+    assert isinstance(model, HeliosPipeline)

tests/diffusion/model_loader/test_diffusers_loader_gguf.py

@@ -62,6 +62,7 @@ def _make_loader() -> DiffusersPipelineLoader:
         ignore_patterns=["*.tmp"],
     )
     loader.od_config = None
+    loader.quant_config = None
     loader.counter_before_loading_weights = 0.0
     loader.counter_after_loading_weights = 0.0
     return loader
@@ -69,53 +70,47 @@ def _make_loader() -> DiffusersPipelineLoader:
 
 def test_is_gguf_quantization_accepts_dict_config():
     loader = _make_loader()
-    od_config = SimpleNamespace(
-        quantization_config={
-            "method": "gguf",
-            "gguf_model": "weights.gguf",
-        }
-    )
+    loader.quant_config = {
+        "method": "gguf",
+        "gguf_model": "weights.gguf",
+    }
 
-    assert loader._is_gguf_quantization(od_config) is True
+    assert loader._is_gguf_quantization() is True
 
 
 def test_is_gguf_quantization_rejects_non_gguf_dict_config():
     loader = _make_loader()
-    od_config = SimpleNamespace(
-        quantization_config={
-            "method": "fp8",
-            "gguf_model": "weights.gguf",
-        }
-    )
+    loader.quant_config = {
+        "method": "fp8",
+        "gguf_model": "weights.gguf",
+    }
 
-    assert loader._is_gguf_quantization(od_config) is False
+    assert loader._is_gguf_quantization() is False
 
 
 def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
     loader = _make_loader()
-    od_config = SimpleNamespace(quantization_config={"method": "gguf"})
+    loader.quant_config = {"method": "gguf"}
 
     with pytest.raises(ValueError, match="gguf_model"):
-        loader._is_gguf_quantization(od_config)
+        loader._is_gguf_quantization()
 
 
 def test_is_gguf_quantization_accepts_object_config():
     loader = _make_loader()
-    od_config = SimpleNamespace(
-        quantization_config=SimpleNamespace(
-            get_name=lambda: "gguf",
-            gguf_model="weights.gguf",
-        )
+    loader.quant_config = SimpleNamespace(
+        get_name=lambda: "gguf",
+        gguf_model="weights.gguf",
     )
 
-    assert loader._is_gguf_quantization(od_config) is True
+    assert loader._is_gguf_quantization() is True
 
 
 def test_is_gguf_quantization_uses_fallback_object_without_get_name():
     loader = _make_loader()
-    od_config = SimpleNamespace(quantization_config=SimpleNamespace(gguf_model="weights.gguf"))
+    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
 
-    assert loader._is_gguf_quantization(od_config) is True
+    assert loader._is_gguf_quantization() is True
 
 
 def test_get_model_loadable_names_collects_parameters_and_buffers():
@@ -171,16 +166,14 @@ def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.Mon
 
 def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
     loader = _make_loader()
+    loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
+    loader.od_config = SimpleNamespace(revision="main")
     source = DiffusersPipelineLoader.ComponentSource(
         model_or_path="dummy",
         subfolder="transformer",
         revision=None,
         prefix="transformer.",
     )
-    od_config = SimpleNamespace(
-        revision="main",
-        quantization_config=SimpleNamespace(gguf_model="weights.gguf"),
-    )
 
     class _Adapter:
         def weights_iterator(self):
@@ -190,7 +183,7 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
     monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf")
     monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter())
 
-    weights = list(loader._get_gguf_weights_iterator(source, object(), od_config))
+    weights = list(loader._get_gguf_weights_iterator(source, object()))
 
     assert weights[0][0] == "transformer.weight"
     assert torch.equal(weights[0][1], torch.ones((2, 2)))
@@ -201,13 +194,12 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
 def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch):
     loader = _make_loader()
     model = _DummyModel()
-    od_config = SimpleNamespace()
 
     monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
     monkeypatch.setattr(
         loader,
         "_get_gguf_weights_iterator",
-        lambda source, model, od_config: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
+        lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
     )
 
     def _hf_iter(source):
@@ -222,7 +214,7 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
         lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
     )
 
-    loaded = loader._load_weights_with_gguf(model, od_config)
+    loaded = loader._load_weights_with_gguf(model)
 
     assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
     assert model.calls[0] == ["transformer.weight"]
@@ -233,13 +225,12 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
 def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch):
     loader = _make_loader()
     model = _DummyModel()
-    od_config = SimpleNamespace()
 
     monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
     monkeypatch.setattr(
         loader,
         "_get_gguf_weights_iterator",
-        lambda source, model, od_config: iter(
+        lambda source, model: iter(
             [
                 ("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
                 ("transformer.bias", torch.zeros(2, dtype=torch.float32)),
@@ -262,7 +253,7 @@ def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete
         lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
     )
 
-    loaded = loader._load_weights_with_gguf(model, od_config)
+    loaded = loader._load_weights_with_gguf(model)
 
     assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
     assert hf_calls == ["vae"]

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

@@ -99,12 +99,16 @@ class StubCosmos3Transformer(nn.Module):
         sound_gen: bool = False,
         sound_dim: int = 3,
         sound_latent_fps: float = 25.0,
+        action_gen: bool = False,
+        action_dim: int = 4,
     ) -> None:
         super().__init__()
         self.latent_channel_size = latent_channel_size
         self.sound_gen = sound_gen
         self.sound_dim = sound_dim
         self.sound_latent_fps = sound_latent_fps
+        self.action_gen = action_gen
+        self.action_dim = action_dim
         self.cached_kv: Any | None = None
         self.cached_freqs_gen: Any | None = None
         self.calls: list[dict[str, Any]] = []
@@ -139,7 +143,10 @@ class StubCosmos3Transformer(nn.Module):
             marker = torch.tensor([token], dtype=torch.float32)
             self.cached_kv = [(marker, marker + 100)]
             self.cached_freqs_gen = (marker + 200, marker + 300)
+        action_latents = kwargs.get("action_latents")
         outputs: list[torch.Tensor] = [torch.full_like(hidden_states, float(token))]
+        if action_latents is not None:
+            outputs.append(torch.full_like(action_latents, float(token + 20)))
         if sound_latents is not None:
             outputs.append(torch.full_like(sound_latents, float(token + 10)))
         return outputs[0] if len(outputs) == 1 else tuple(outputs)
@@ -344,7 +351,7 @@ def test_pipeline_init_passes_tokenizer_attrs_into_transformer(
     assert pipeline.transformer.audio_proj_out.out_features == 5
 
 
-def test_preprocess_i2v_image_input() -> None:
+def test_preprocess_i2v_image_and_action_video_inputs() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func
 
     preprocess = get_cosmos3_pre_process_func(SimpleNamespace())
@@ -357,6 +364,29 @@ def test_preprocess_i2v_image_input() -> None:
     assert (result.sampling_params.height, result.sampling_params.width) == (672, 1344)
     assert tuple(result.prompts[0]["additional_information"]["preprocessed_image"].shape[-2:]) == (672, 1344)
 
+    frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue")]
+    action = SimpleNamespace(
+        prompts=[{"prompt": "Move.", "multi_modal_data": {"video": frames}}],
+        sampling_params=SimpleNamespace(height=16, width=32, extra_args={"action_mode": "forward_dynamics"}),
+    )
+
+    additional = preprocess(action).prompts[0]["additional_information"]
+    assert tuple(additional["preprocessed_image"].shape) == (1, 3, 16, 32)
+    assert tuple(additional["preprocessed_video"].shape) == (1, 3, 3, 16, 32)
+
+    frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue", "yellow", "purple", "black")]
+    v2v = SimpleNamespace(
+        prompts=[{"prompt": "Continue.", "multi_modal_data": {"video": frames}}],
+        sampling_params=SimpleNamespace(
+            height=16,
+            width=32,
+            extra_args={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "last"},
+        ),
+    )
+    additional = preprocess(v2v).prompts[0]["additional_information"]
+    assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32)
+    assert additional["condition_frame_indexes_vision"] == [0, 1]
+
 
 def test_postprocess_handles_image_video_audio_and_validation() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func
@@ -434,7 +464,7 @@ def test_prompt_formatting_and_checkpoint_key_remap(make_cosmos3_pipeline) -> No
     assert {key: Cosmos3OmniDiffusersPipeline._remap_ckpt_key(key) for key in remaps} == remaps
 
 
-def test_prepare_latents_for_video_image_and_sound(make_cosmos3_pipeline) -> None:
+def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     latents = pipeline._prepare_latents(16, 24, 5, torch.Generator(device="cpu").manual_seed(0))
     assert latents.shape == (1, 2, 2, 2, 3)
@@ -447,6 +477,19 @@ def test_prepare_latents_for_video_image_and_sound(make_cosmos3_pipeline) -> Non
     assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]]]]]
     assert image_latent.shape == (1, 2, 1, 2, 3)
 
+    pipeline._encode_video_tensor = lambda *args, **kwargs: torch.full((1, 2, 3, 2, 3), 6.0)
+    v2v_latents, v2v_velocity_mask, v2v_condition = pipeline._prepare_latents_v2v(
+        torch.zeros(1, 3, 5, 16, 24),
+        16,
+        24,
+        9,
+        torch.Generator(device="cpu").manual_seed(0),
+        [0, 1],
+    )
+    torch.testing.assert_close(v2v_latents[:, :, 0:2], torch.full((1, 2, 2, 2, 3), 6.0))
+    assert v2v_velocity_mask.tolist() == [[[[[0.0]], [[0.0]], [[1.0]]]]]
+    assert v2v_condition.shape == (1, 2, 3, 2, 3)
+
     pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
     pipeline._sound_tokenizer = SimpleNamespace(
         sample_rate=10,
@@ -463,8 +506,20 @@ def test_prepare_latents_for_video_image_and_sound(make_cosmos3_pipeline) -> Non
     assert (sound_latents.shape, latent_frames) == (torch.Size([1, 3, 6]), 6)
     assert pipeline._decode_sound_latents(torch.zeros(1, 3, 6), target_audio_samples=21).shape == (1, 2, 21)
 
+    pipeline.transformer = pipeline.transformer.__class__(action_gen=True, action_dim=4)
+    action, action_mask, clean, raw_dim = pipeline._prepare_action_latents(
+        mode="forward_dynamics",
+        action_chunk_size=2,
+        raw_action_dim=None,
+        generator=torch.Generator(device="cpu").manual_seed(0),
+        sp=SimpleNamespace(extra_args={"action": [[1.0, 2.0], [3.0, 4.0]]}),
+    )
+    assert raw_dim == 2
+    assert action_mask.tolist() == [[[0.0], [0.0]]]
+    torch.testing.assert_close(action, clean)
 
-def test_diffuse_covers_cfg_i2v_and_sound_steps(make_cosmos3_pipeline) -> None:
+
+def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> None:
     pipeline = make_cosmos3_pipeline()
     latents = torch.zeros(1, 2, 1, 1, 1)
 
@@ -496,20 +551,22 @@ def test_diffuse_covers_cfg_i2v_and_sound_steps(make_cosmos3_pipeline) -> None:
     )
     torch.testing.assert_close(i2v[:, :, 0:1], torch.full((1, 2, 1, 1, 1), 7.0))
 
-    pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
-    video_result, sound_result = pipeline.diffuse(
+    pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
+    video_result, action_result = pipeline.diffuse(
         latents=latents,
-        sound_latents=torch.zeros(1, 3, 4),
+        action_latents=torch.zeros(1, 3, 4),
+        action_velocity_mask=torch.ones(1, 3, 1),
+        action_condition_latents=torch.zeros(1, 3, 4),
         timesteps=torch.tensor([7, 3]),
         cond_ids=_ids(2),
         cond_mask=_mask(),
         uncond_ids=_ids(1),
         uncond_mask=_mask(),
         guidance_scale=1.0,
-        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "action_domain_ids": torch.tensor([0])},
     )
     torch.testing.assert_close(video_result, torch.full_like(latents, 4.0))
-    torch.testing.assert_close(sound_result, torch.full((), 24.0).expand_as(sound_result))
+    torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result))
 
 
 def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -> None:
@@ -568,6 +625,8 @@ class TestForwardRouting:
         def fake_diffuse(**kwargs):
             captured["diffuse_calls"].append(kwargs)
             outputs = [kwargs["latents"] + len(captured["diffuse_calls"])]
+            if kwargs.get("action_latents") is not None:
+                outputs.append(kwargs["action_latents"] + 3.0)
             if kwargs.get("sound_latents") is not None:
                 outputs.append(kwargs["sound_latents"] + 2.0)
             return outputs[0] if len(outputs) == 1 else tuple(outputs)
@@ -612,7 +671,7 @@ class TestForwardRouting:
         assert captured["flow_shifts"] == expected["flow"]
         assert [call[0] for call in pipeline.scheduler.set_timesteps_calls] == expected["steps"]
 
-    def test_forward_i2v_and_sound_routes(self, make_cosmos3_pipeline) -> None:
+    def test_forward_i2v_sound_and_action_routes(self, make_cosmos3_pipeline) -> None:
         pipeline = make_cosmos3_pipeline()
         captured = self._install_forward_stubs(pipeline)
         image_tensor = torch.zeros(1, 3, 16, 16)
@@ -637,6 +696,34 @@ class TestForwardRouting:
         )
         assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
 
+        video_tensor = torch.zeros(1, 3, 5, 16, 16)
+        v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0)
+        v2v_mask = torch.tensor([[[[[0.0]], [[1.0]]]]])
+        pipeline._prepare_latents_v2v = lambda *args, **kwargs: (
+            torch.zeros(1, 2, 2, 1, 1),
+            v2v_mask,
+            v2v_condition,
+        )
+        pipeline.forward(
+            SimpleNamespace(
+                prompts=[
+                    {
+                        "prompt": "continue",
+                        "modalities": ["video"],
+                        "additional_information": {
+                            "preprocessed_video": video_tensor,
+                            "condition_frame_indexes_vision": [0],
+                        },
+                    }
+                ],
+                sampling_params=make_sampling_params(height=16, width=16, num_frames=5),
+            )
+        )
+        assert captured["flow_shifts"][-1] == 10.0
+        assert captured["format"]["negative_prompt"] == ""
+        assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
+        assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
+
         pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
         sound_latents = torch.zeros(1, 3, 4)
         pipeline._resolve_sound_target_samples = lambda *args: (20, 2.0, 10)
@@ -651,6 +738,31 @@ class TestForwardRouting:
         assert captured["diffuse_calls"][-1]["sound_latents"] is sound_latents
         assert output.output["audio_sample_rate"] == 10
 
+        pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, action_gen=True, action_dim=4)
+        output = pipeline.forward(
+            SimpleNamespace(
+                prompts=[
+                    {
+                        "prompt": "Pick the block.",
+                        "modalities": ["video"],
+                        "additional_information": {"preprocessed_image": image_tensor},
+                    }
+                ],
+                sampling_params=make_sampling_params(
+                    height=16,
+                    width=16,
+                    extra_args={
+                        "action_mode": "policy",
+                        "action_chunk_size": 2,
+                        "raw_action_dim": 2,
+                        "domain_name": "bridge_orig_lerobot",
+                    },
+                ),
+            )
+        )
+        assert captured["diffuse_calls"][-1]["shared_kwargs"]["action_domain_ids"].tolist() == [7]
+        assert output.custom_output["action"].shape == (1, 2, 2)
+
     @pytest.mark.parametrize(
         ("prompt", "sampling_params", "message"),
         [

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

@@ -30,8 +30,9 @@ def _tiny_cosmos3_config(**overrides):
     return config
 
 
-def test_mrope_position_ids_cover_text_video_and_sound() -> None:
+def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import (
+        compute_mrope_position_ids_action,
         compute_mrope_position_ids_sound,
         compute_mrope_position_ids_text,
         compute_mrope_position_ids_vision,
@@ -62,6 +63,10 @@ def test_mrope_position_ids_cover_text_video_and_sound() -> None:
     torch.testing.assert_close(sound_ids[0], torch.tensor([10.0, 10.96, 11.92]))
     assert sound_offset == 12
 
+    action_ids, action_offset = compute_mrope_position_ids_action(3, temporal_offset=10, action_fps=None)
+    assert action_ids.tolist() == [[11, 12, 13], [0, 0, 0], [0, 0, 0]]
+    assert action_offset == 14
+
 
 @pytest.mark.parametrize(
     ("key", "value"),
@@ -126,7 +131,7 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
-def test_sound_modules_follow_injected_sound_dim() -> None:
+def test_sound_and_action_modules_follow_config() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
 
     tiny = _tiny_cosmos3_config()
@@ -137,12 +142,22 @@ def test_sound_modules_follow_injected_sound_dim() -> None:
         sound_dim=5,
         sound_latent_fps=40.0,
     )
+    with_action = Cosmos3VFMTransformer(
+        SimpleNamespace(
+            tf_model_config={**tiny, "action_gen": True, "max_action_dim": 6, "num_embodiment_domains": 9},
+            dtype=torch.float32,
+        )
+    )
 
     assert no_modal.sound_gen is False
+    assert no_modal.action_gen is False
     assert not hasattr(no_modal, "audio_proj_in")
+    assert not hasattr(no_modal, "action_proj_in")
     assert with_sound.sound_dim == 5
     assert with_sound.sound_latent_fps == 40.0
     assert with_sound.audio_proj_in.in_features == 5
+    assert with_action.action_dim == 6
+    assert with_action.action_proj_in.num_domains == 9
 
 
 @pytest.mark.parametrize(
@@ -163,33 +178,56 @@ def test_transformer_requires_sound_dim_and_fps_when_sound_gen_true(kwargs: dict
         )
 
 
-def test_sound_pack_unpack_validate_shapes() -> None:
+def test_sound_and_action_pack_unpack_validate_shapes() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
 
     model = object.__new__(Cosmos3VFMTransformer)
     nn.Module.__init__(model)
     model.sound_dim = 3
+    model.action_dim = 3
 
     sound = torch.arange(2 * 3 * 4, dtype=torch.float32).reshape(2, 3, 4)
+    action = torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3)
     torch.testing.assert_close(model.unpack_sound(model.pack_sound(sound)), sound)
+    torch.testing.assert_close(model.unpack_action(model.pack_action(action)), action)
 
     with pytest.raises(ValueError, match="channel mismatch"):
         model.pack_sound(torch.zeros(1, 4, 2))
+    with pytest.raises(ValueError, match="dimension mismatch"):
+        model.pack_action(torch.zeros(1, 2, 4))
 
 
-def test_forward_returns_video_and_sound_predictions(monkeypatch: pytest.MonkeyPatch) -> None:
+@pytest.mark.parametrize(
+    ("config", "transformer_kwargs", "extra_kwargs", "expected_shapes"),
+    [
+        (
+            _tiny_cosmos3_config(),
+            {"sound_gen": True, "sound_dim": 3, "sound_latent_fps": 24.0},
+            {"sound_latents": torch.zeros(1, 3, 4)},
+            [(1, 2, 1, 2, 2), (1, 3, 4)],
+        ),
+        (
+            _tiny_cosmos3_config(action_gen=True, max_action_dim=3, num_embodiment_domains=4),
+            {},
+            {"action_latents": torch.zeros(1, 5, 3), "action_domain_ids": torch.tensor([2])},
+            [(1, 2, 1, 2, 2), (1, 5, 3)],
+        ),
+    ],
+)
+def test_forward_returns_video_plus_optional_modality_predictions(
+    monkeypatch: pytest.MonkeyPatch,
+    config,
+    transformer_kwargs,
+    extra_kwargs,
+    expected_shapes,
+) -> None:
     from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
 
     monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
 
     output = transformer_cosmos3.Cosmos3VFMTransformer(
-        SimpleNamespace(
-            tf_model_config=_tiny_cosmos3_config(),
-            dtype=torch.float32,
-        ),
-        sound_gen=True,
-        sound_dim=3,
-        sound_latent_fps=40.0,
+        SimpleNamespace(tf_model_config=config, dtype=torch.float32),
+        **transformer_kwargs,
     )(
         hidden_states=torch.zeros(1, 2, 1, 2, 2),
         timestep=torch.tensor([1.0]),
@@ -197,11 +235,12 @@ def test_forward_returns_video_and_sound_predictions(monkeypatch: pytest.MonkeyP
         text_mask=torch.ones(1, 2, dtype=torch.long),
         video_shape=(1, 2, 2),
         fps=24.0,
-        sound_latents=torch.zeros(1, 3, 4),
+        action_noisy_mask=torch.ones(1, 5, 1),
+        **extra_kwargs,
     )
 
     assert isinstance(output, tuple)
-    assert [tuple(tensor.shape) for tensor in output] == [(1, 2, 1, 2, 2), (1, 3, 4)]
+    assert [tuple(tensor.shape) for tensor in output] == expected_shapes
 
 
 def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -227,7 +266,7 @@ def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch
         )
 
 
-def test_compute_rope_freqs_places_text_video_and_sound_positions() -> None:
+def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
 
     class FakeRotary:
@@ -275,9 +314,11 @@ def test_compute_rope_freqs_places_text_video_and_sound_positions() -> None:
         fps=24.0,
         device=torch.device("cpu"),
         dtype=torch.float32,
+        t_action=2,
+        action_start_frame_offset=1,
         t_sound=1,
     )
 
     _, gen_pos = rotary.position_ids
-    assert gen_pos.shape == (3, 1, 3)
-    assert gen_pos[0, 0].tolist() == [102, 103, 102]
+    assert gen_pos.shape == (3, 1, 5)
+    assert gen_pos[0, 0].tolist() == [102, 103, 103, 104, 102]

tests/diffusion/models/hunyuan_video/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

tests/diffusion/models/hunyuan_video/test_hunyuan_video_quant_config_propagation.py

@@ -0,0 +1,160 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for HunyuanVideo-1.5 quant_config propagation through transformer creation.
+
+Tests cover:
+- HunyuanVideo15Pipeline passes quant_config to HunyuanVideo15Transformer3DModel
+- HunyuanVideo15I2VPipeline passes quant_config to HunyuanVideo15Transformer3DModel
+"""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 as t2v_module
+import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v as i2v_module
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
+
+
+class TestHunyuanVideoQuantConfigPropagation:
+    """Verify quant_config is propagated to the transformer model in HunyuanVideo-1.5 pipelines."""
+
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_local_device",
+        return_value="cpu",
+    )
+    @patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.prefetch_subfolders")
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.from_pretrained_with_prefetch",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.Qwen2Tokenizer.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.ByT5Tokenizer.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.AutoConfig.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.T5EncoderModel",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.FlowMatchEulerDiscreteScheduler.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_transformer_config_kwargs",
+        return_value={},
+    )
+    def test_t2v_quant_config_passed_through(
+        self,
+        mock_get_kwargs,
+        mock_scheduler,
+        mock_t5_encoder,
+        mock_auto_config,
+        mock_byt5_tok,
+        mock_qwen2_tok,
+        mock_from_pretrained,
+        mock_prefetch,
+        mock_get_device,
+    ):
+        captured_kwargs = {}
+
+        class FakeTransformer:
+            def __init__(self, **kwargs):
+                captured_kwargs.update(kwargs)
+
+        with patch.object(t2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer):
+            fake_qc = MagicMock()
+            od_config = SimpleNamespace(
+                model="fake-t2v-model",
+                dtype=torch.bfloat16,
+                flow_shift=None,
+                tf_model_config=SimpleNamespace(),
+                quantization_config=fake_qc,
+                enable_diffusion_pipeline_profiler=False,
+            )
+
+            t2v_module.HunyuanVideo15Pipeline(od_config=od_config)
+
+            assert captured_kwargs.get("quant_config") is fake_qc
+
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_local_device",
+        return_value="cpu",
+    )
+    @patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.prefetch_subfolders")
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.from_pretrained_with_prefetch",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.Qwen2Tokenizer.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.ByT5Tokenizer.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.AutoConfig.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.T5EncoderModel",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.SiglipImageProcessor.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.FlowMatchEulerDiscreteScheduler.from_pretrained",
+        return_value=MagicMock(),
+    )
+    @patch(
+        "vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_transformer_config_kwargs",
+        return_value={},
+    )
+    def test_i2v_quant_config_passed_through(
+        self,
+        mock_get_kwargs,
+        mock_scheduler,
+        mock_siglip_processor,
+        mock_t5_encoder,
+        mock_auto_config,
+        mock_byt5_tok,
+        mock_qwen2_tok,
+        mock_from_pretrained,
+        mock_prefetch,
+        mock_get_device,
+    ):
+        captured_kwargs = {}
+
+        class FakeTransformer:
+            def __init__(self, **kwargs):
+                captured_kwargs.update(kwargs)
+
+        with patch.object(i2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer):
+            fake_qc = MagicMock()
+            od_config = SimpleNamespace(
+                model="fake-i2v-model",
+                dtype=torch.bfloat16,
+                flow_shift=None,
+                tf_model_config=SimpleNamespace(),
+                quantization_config=fake_qc,
+                enable_diffusion_pipeline_profiler=False,
+            )
+
+            i2v_module.HunyuanVideo15I2VPipeline(od_config=od_config)
+
+            assert captured_kwargs.get("quant_config") is fake_qc

tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py

@@ -72,6 +72,31 @@ class TestPipelineIndependence:
 
         assert issubclass(LTX23Pipeline, CFGParallelMixin)
 
+    def test_ltx23_pipeline_declares_offload_components(self):
+        """LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery."""
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
+        from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
+
+        pipe = object.__new__(LTX23Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.transformer = torch.nn.Linear(1, 1)
+        pipe.text_encoder = torch.nn.Linear(1, 1)
+        pipe.connectors = torch.nn.Linear(1, 1)
+        pipe.vae = torch.nn.Linear(1, 1)
+        pipe.audio_vae = torch.nn.Linear(1, 1)
+        pipe.vocoder = torch.nn.Linear(1, 1)
+
+        modules = ModuleDiscovery.discover(pipe)
+
+        assert LTX23Pipeline._dit_modules == ["transformer"]
+        assert LTX23Pipeline._encoder_modules == ["text_encoder", "connectors"]
+        assert LTX23Pipeline._vae_modules == ["vae", "audio_vae"]
+        assert LTX23Pipeline._resident_modules == ["vocoder"]
+        assert modules.dit_names == ["transformer"]
+        assert modules.encoder_names == ["text_encoder", "connectors"]
+        assert modules.resident_names == ["vocoder"]
+        assert len(modules.vaes) == 2
+
 
 class TestCFGParallelHelpers:
     """Test LTX-2.3 CFG helper math without loading model weights."""

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from torch import nn
+
+from vllm_omni.diffusion.models.ltx2.ltx2_transformer import _make_rms_norm
+
+
+def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer():
+    norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
+
+    assert "weight" not in dict(norm.named_parameters())
+    assert "weight" in dict(norm.named_buffers())
+    assert "weight" not in norm.state_dict()
+
+
+def test_ltx_rms_norm_affine_weight_remains_parameter():
+    norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=True)
+
+    assert isinstance(dict(norm.named_parameters())["weight"], nn.Parameter)
+    assert "weight" not in dict(norm.named_buffers())
+    assert "weight" in norm.state_dict()

tests/diffusion/test_worker_wrapper_base.py

@@ -41,7 +41,7 @@ def mock_od_config(mocker: MockerFixture):
     config.cache_backend = None
     config.cache_config = None
     config.model = "test-model"
-    config.diffusion_load_format = None
+    config.diffusion_load_format = "default"
     config.dtype = "float32"
     config.max_cpu_loras = 0
     config.lora_path = None

tests/distributed/omni_connectors/test_bagel_mooncake_connector.py

@@ -32,8 +32,8 @@ BAGEL_MOONCAKE_CI_DEPLOY = get_deploy_config_path("ci/bagel_mooncake.yaml")
 
 # Reference pixel data extracted from the known-good output image
 # Each entry contains (x, y) position and expected (R, G, B) values
-# "Generated with seed=52, num_inference_steps=15,
-# prompt='A futuristic city skyline at twilight, cyberpunk style'"
+# "Generated with seed=52, num_inference_steps=14,
+# prompt='A cute cat'"
 REFERENCE_PIXELS = [
     {"position": (100, 100), "rgb": (115, 113, 94)},
     {"position": (400, 50), "rgb": (159, 160, 144)},
@@ -63,7 +63,7 @@ def _find_free_port() -> int:
     return port
 
 
-def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
+def _configure_sampling_params(omni: Omni, num_inference_steps: int = 14) -> list:
     """Configure sampling parameters for Bagel text2img generation.
 
     Args:
@@ -269,7 +269,6 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
-@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
 @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."""

tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py

@@ -24,14 +24,13 @@ from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniRunner
 from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.platforms import current_omni_platform
 
 pytestmark = [pytest.mark.usefixtures("clean_gpu_memory_between_tests")]
 
 BAGEL_CI_DEPLOY = get_deploy_config_path("ci/bagel.yaml")
 
 # Reference pixel data extracted from the known-good output image
-# Generated with seed=52, num_inference_steps=15,
+# Generated with seed=52, num_inference_steps=14,
 # prompt='Change the grass color to red',
 # input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
 REFERENCE_PIXELS = [
@@ -47,23 +46,10 @@ REFERENCE_PIXELS = [
     {"position": (256, 256), "rgb": (181, 201, 221)},
 ]
 
-if current_omni_platform.is_rocm():
-    REFERENCE_PIXELS = [
-        {"position": (100, 100), "rgb": (156, 172, 217)},
-        {"position": (400, 50), "rgb": (105, 144, 217)},
-        {"position": (700, 100), "rgb": (118, 159, 232)},
-        {"position": (150, 400), "rgb": (180, 22, 52)},
-        {"position": (512, 336), "rgb": (221, 211, 194)},
-        {"position": (700, 400), "rgb": (192, 10, 46)},
-        {"position": (100, 600), "rgb": (102, 12, 22)},
-        {"position": (400, 600), "rgb": (161, 28, 47)},
-        {"position": (700, 600), "rgb": (100, 87, 94)},
-        {"position": (256, 256), "rgb": (181, 201, 221)},
-    ]
 
 # text2img reference pixels (aligned with offline `bagel/end2end.py` text2img, 15 steps)
-# "Generated with seed=52, num_inference_steps=15,
-# prompt='A futuristic city skyline at twilight, cyberpunk style'"
+# "Generated with seed=52, num_inference_steps=14,
+# prompt='A cute cat'"
 TEXT2IMG_REFERENCE_PIXELS = [
     {"position": (100, 100), "rgb": (115, 113, 94)},
     {"position": (400, 50), "rgb": (159, 160, 144)},
@@ -77,19 +63,6 @@ TEXT2IMG_REFERENCE_PIXELS = [
     {"position": (256, 256), "rgb": (167, 156, 150)},
 ]
 
-if current_omni_platform.is_rocm():
-    TEXT2IMG_REFERENCE_PIXELS = [
-        {"position": (100, 100), "rgb": (115, 113, 94)},
-        {"position": (400, 50), "rgb": (159, 160, 144)},
-        {"position": (700, 100), "rgb": (164, 151, 123)},
-        {"position": (150, 400), "rgb": (120, 121, 107)},
-        {"position": (512, 512), "rgb": (165, 133, 127)},
-        {"position": (700, 400), "rgb": (217, 130, 66)},
-        {"position": (100, 700), "rgb": (191, 168, 152)},
-        {"position": (400, 700), "rgb": (130, 96, 77)},
-        {"position": (700, 700), "rgb": (247, 203, 140)},
-        {"position": (256, 256), "rgb": (167, 156, 150)},
-    ]
 
 PIXEL_TOLERANCE = 10
 TEXT2IMG_PIXEL_TOLERANCE = 5
@@ -115,7 +88,7 @@ def _find_free_port() -> int:
     return port
 
 
-def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
+def _configure_sampling_params(omni: Omni, num_inference_steps: int = 14) -> list:
     """Configure sampling parameters for Bagel img2img generation.
 
     Args:
@@ -258,7 +231,6 @@ def _resolve_deploy_config(config_path: str, run_level: str) -> str:
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
-@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"})
 def test_bagel_img2img_shared_memory_connector(run_level):
     """Test Bagel img2img with shared memory connector."""
@@ -276,7 +248,6 @@ def test_bagel_img2img_shared_memory_connector(run_level):
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.diffusion
-@pytest.mark.skip(reason="Skip failed CI issue 3977: https://github.com/vllm-project/vllm-omni/issues/3977")
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"})
 def test_bagel_text2img_shared_memory_connector(run_level):
     """Test Bagel text2img with shared memory connector."""

tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py

@@ -8,7 +8,7 @@ L4-style smoke runs against an already-running server. The pytest wrappers in
 H100 / MI325) because they launch the live Omni server inside the test.
 
 1. **Daily-Omni** — MCQ accuracy fields in the saved JSON (``daily_omni_accuracy``, …); by default the
-   run **fails** if accuracy is strictly below **0.67** (``--min-daily-omni-accuracy`` / ``ACC_BENCH_MIN_DAILY_OMNI_ACCURACY``).
+   run **fails** if accuracy is strictly below **0.69** (``--min-daily-omni-accuracy`` / ``ACC_BENCH_MIN_DAILY_OMNI_ACCURACY``).
 2. **Seed-TTS** — ``seed-tts-eval``-style metrics when ``--seed-tts-wer-eval`` is used
    (WER / SIM / UTMOS keys from :func:`compute_seed_tts_wer_metrics`).
 
@@ -308,9 +308,9 @@ def build_arg_parser() -> argparse.ArgumentParser:
     p.add_argument(
         "--min-daily-omni-accuracy",
         type=float,
-        default=float((os.environ.get("ACC_BENCH_MIN_DAILY_OMNI_ACCURACY") or "0.67").strip() or "0.67"),
+        default=float((os.environ.get("ACC_BENCH_MIN_DAILY_OMNI_ACCURACY") or "0.69").strip() or "0.69"),
         help="Fail when daily_omni_accuracy is strictly below this threshold (0–1). "
-        "Default baseline 0.67; override with env ACC_BENCH_MIN_DAILY_OMNI_ACCURACY or pass 0 to disable the floor.",
+        "Default baseline 0.69; override with env ACC_BENCH_MIN_DAILY_OMNI_ACCURACY or pass 0 to disable the floor.",
     )
 
     p.add_argument(
@@ -340,7 +340,7 @@ def build_arg_parser() -> argparse.ArgumentParser:
     p.add_argument(
         "--max-seed-tts-mean-wer",
         type=float,
-        default=0.5,
+        default=0.35,
         help="If set, fail when seed_tts_content_error_mean is strictly above this value.",
     )
     p.add_argument(

tests/e2e/accuracy/test_diffusers_backend_similarity.py

@@ -218,7 +218,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> tuple[Image.I
             torch_dtype=torch.bfloat16,
             trust_remote_code=True,
         ).to("cuda")
-
+        pipe.transformer.set_attention_backend("_flash_3_hub")
         _diffusers_dummy_run(pipe)
 
         generator = torch.Generator(device="cuda").manual_seed(SEED)

tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import base64
+import copy
+import io
+import os
+import tempfile
+from pathlib import Path
+
+import pytest
+import requests
+import yaml
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import assert_images_pixel_close, assert_similarity, model_output_dir
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServer
+
+pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
+
+MODEL_NAME = "tencent/HunyuanImage-3.0-Instruct"
+SEED = 42
+NUM_INFERENCE_STEPS = 50
+GUIDANCE_SCALE = 2.5
+HEIGHT = 1024
+WIDTH = 1024
+PROMPT = "A brown and white dog is running on the grass."
+MEAN_THRESHOLD = 3e-2
+P99_THRESHOLD = 3e-1
+SSIM_THRESHOLD = 0.97
+PSNR_THRESHOLD = 30.0
+
+_REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent
+BASELINE_PATH = _REPO_ROOT / "tests" / "assets" / "hunyuan" / "hunyuan_baseline.png"
+_OFFLINE_SCRIPT = _REPO_ROOT / "examples" / "offline_inference" / "hunyuan_image3" / "end2end.py"
+
+# DiT-only deploy config with trust_remote_code (based on hunyuan_image3_dit.yaml).
+_DEPLOY_CONFIG = {
+    "pipeline": "hunyuan_image3_dit",
+    "async_chunk": False,
+    "trust_remote_code": True,
+    "stages": [
+        {
+            "stage_id": 0,
+            "max_num_seqs": 1,
+            "gpu_memory_utilization": 0.9,
+            "enforce_eager": True,
+            "trust_remote_code": True,
+            "devices": "0,1,2,3",  # set dynamically by _write_deploy_config
+            "vae_use_slicing": False,
+            "vae_use_tiling": False,
+            "parallel_config": {
+                "pipeline_parallel_size": 1,
+                "data_parallel_size": 1,
+                "tensor_parallel_size": 4,
+                "enable_expert_parallel": True,
+                "sequence_parallel_size": 1,
+                "ulysses_degree": 1,
+                "ring_degree": 1,
+                "cfg_parallel_size": 1,
+                "vae_patch_parallel_size": 1,
+                "use_hsdp": False,
+                "hsdp_shard_size": -1,
+                "hsdp_replicate_size": 1,
+            },
+            "default_sampling_params": {
+                "seed": SEED,
+            },
+        },
+    ],
+}
+
+
+def _model_name() -> str:
+    return os.environ.get("HUNYUAN_IMAGE3_MODEL", MODEL_NAME)
+
+
+def _devices() -> str:
+    return os.environ.get("HUNYUAN_IMAGE3_DEVICES", "0,1,2,3")
+
+
+def _write_deploy_config(path: Path) -> None:
+    config = copy.deepcopy(_DEPLOY_CONFIG)
+    devices = _devices()
+    config["stages"][0]["devices"] = devices
+    config["stages"][0]["parallel_config"]["tensor_parallel_size"] = len(devices.split(","))
+    path.write_text(yaml.dump(config, default_flow_style=False, sort_keys=False))
+
+
+def _run_vllm_omni_hunyuan_image3_online(*, model: str, deploy_config: str, output_path: Path) -> Image.Image:
+    server_args = [
+        "--deploy-config",
+        deploy_config,
+        "--stage-init-timeout",
+        "300",
+        "--init-timeout",
+        "900",
+        "--enforce-eager",
+        "--trust-remote-code",
+    ]
+    with OmniServer(model, server_args, use_omni=True) as omni_server:
+        response = requests.post(
+            f"http://{omni_server.host}:{omni_server.port}/v1/images/generations",
+            json={
+                "model": omni_server.model,
+                "prompt": PROMPT,
+                "size": f"{WIDTH}x{HEIGHT}",
+                "n": 1,
+                "response_format": "b64_json",
+                "num_inference_steps": NUM_INFERENCE_STEPS,
+                "guidance_scale": GUIDANCE_SCALE,
+                "seed": SEED,
+                "bot_task": "none",
+                "use_system_prompt": "en_unified",
+            },
+            timeout=600,
+        )
+        response.raise_for_status()
+        payload = response.json()
+        assert len(payload["data"]) == 1
+        image_bytes = base64.b64decode(payload["data"][0]["b64_json"])
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        image.load()
+        image.save(output_path)
+        return image
+
+
+def _run_vllm_omni_hunyuan_image3_offline(*, model: str, deploy_config: str, output_path: Path) -> Image.Image:
+    import subprocess
+    import sys
+
+    output_dir = str(output_path.parent)
+    subprocess.run(
+        [
+            sys.executable,
+            str(_OFFLINE_SCRIPT),
+            "--modality",
+            "text2img",
+            "--deploy-config",
+            deploy_config,
+            "--prompts",
+            PROMPT,
+            "--output",
+            output_dir,
+            "--steps",
+            str(NUM_INFERENCE_STEPS),
+            "--guidance-scale",
+            str(GUIDANCE_SCALE),
+            "--seed",
+            str(SEED),
+            "--height",
+            str(HEIGHT),
+            "--width",
+            str(WIDTH),
+            "--bot-task",
+            "none",
+            "--sys-type",
+            "en_unified",
+            "--model",
+            model,
+            "--enforce-eager",
+        ],
+        check=True,
+    )
+    images = sorted(Path(output_dir).glob("output_*.png"))
+    assert images, f"No output image found in {output_dir}"
+    image = Image.open(images[0]).convert("RGB")
+    image.load()
+    image.save(output_path)
+    return image
+
+
+def _assert_against_baseline(image: Image.Image, label: str) -> None:
+    assert BASELINE_PATH.exists(), f"Baseline image not found at {BASELINE_PATH}"
+    baseline_image = Image.open(BASELINE_PATH).convert("RGB")
+
+    assert_images_pixel_close(
+        model_name=f"{MODEL_NAME} ({label} vs baseline)",
+        vllm_image=image,
+        diffusers_image=baseline_image,
+        mean_threshold=MEAN_THRESHOLD,
+        p99_threshold=P99_THRESHOLD,
+    )
+    assert_similarity(
+        model_name=f"{MODEL_NAME} ({label} vs baseline)",
+        vllm_image=image,
+        diffusers_image=baseline_image,
+        ssim_threshold=SSIM_THRESHOLD,
+        psnr_threshold=PSNR_THRESHOLD,
+    )
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+def test_hunyuan_image3_pixel_accuracy_online(accuracy_artifact_root: Path) -> None:
+    model = _model_name()
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        deploy_config_path = Path(tmpdir) / "deploy.yaml"
+        _write_deploy_config(deploy_config_path)
+        image = _run_vllm_omni_hunyuan_image3_online(
+            model=model, deploy_config=str(deploy_config_path), output_path=output_dir / "vllm_omni_online.png"
+        )
+    _assert_against_baseline(image, "online")
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=4)
+def test_hunyuan_image3_pixel_accuracy_offline(accuracy_artifact_root: Path) -> None:
+    model = _model_name()
+    output_dir = model_output_dir(accuracy_artifact_root, MODEL_NAME)
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        deploy_config_path = Path(tmpdir) / "deploy.yaml"
+        _write_deploy_config(deploy_config_path)
+        image = _run_vllm_omni_hunyuan_image3_offline(
+            model=model, deploy_config=str(deploy_config_path), output_path=output_dir / "vllm_omni_offline.png"
+        )
+    _assert_against_baseline(image, "offline")

tests/e2e/accuracy/test_qwen_image.py

@@ -106,6 +106,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image:
             trust_remote_code=True,
             local_files_only=_local_files_only(model),
         ).to("cuda")
+        pipe.transformer.set_attention_backend("_flash_3_hub")
         generator = torch.Generator(device="cuda").manual_seed(SEED)
         result = pipe(  # pyright: ignore[reportCallIssue]
             prompt=PROMPT,
@@ -167,6 +168,7 @@ def _run_diffusers_qwen_image_2512(*, model: str, output_path: Path) -> Image.Im
             trust_remote_code=True,
             local_files_only=_local_files_only(model),
         ).to("cuda")
+        pipe.transformer.set_attention_backend("_flash_3_hub")
         generator = torch.Generator(device="cuda").manual_seed(SEED_2512)
         result = pipe(  # pyright: ignore[reportCallIssue]
             prompt=PROMPT_2512,

tests/e2e/accuracy/test_videogen_modelopt_quant.py

@@ -0,0 +1,273 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end ModelOpt FP8 quality accuracy for video-gen DiTs.
+
+Modeled on ``test_hunyuan_image3.py``'s ``test_quantized_dit_matches_bf16_accuracy``:
+for each calibrated FP8 checkpoint the test generates one video with BF16 and one
+with the FP8 DiT under the same seed, scores prompt-faithfulness with CLIP on the
+middle frame, then asserts both ``CLIP >= absolute_floor`` and
+``CLIP_drop <= drop_threshold``. The test is gated by an opt-in env var (model
+paths are local) so it never runs in generic CI.
+
+To run::
+
+    VIDEOGEN_RUN_QUANT_ACCURACY=1 \\
+    WAN22_A14B_BF16_MODEL=Wan-AI/Wan2.2-T2V-A14B-Diffusers \\
+    WAN22_A14B_FP8_MODEL=/path/to/wan22-a14b-modelopt-fp8 \\
+    HV15_BF16_MODEL=hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
+    HV15_FP8_MODEL=/path/to/hv15-modelopt-fp8 \\
+        pytest -s -v tests/e2e/accuracy/test_videogen_modelopt_quant.py
+"""
+
+from __future__ import annotations
+
+import gc
+import json
+import os
+import tempfile
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import pytest
+import torch
+from PIL import Image
+
+from tests.e2e.accuracy.helpers import CLIPScorer, model_output_dir
+
+pytestmark = [pytest.mark.local_model, pytest.mark.diffusion]
+
+# ---------------------------------------------------------------------------
+# Tunables (override via env to suit local hardware / calibration budget)
+# ---------------------------------------------------------------------------
+SEED = 42
+HEIGHT = int(os.environ.get("VIDEOGEN_QUANT_HEIGHT", "480"))
+WIDTH = int(os.environ.get("VIDEOGEN_QUANT_WIDTH", "832"))
+NUM_FRAMES = int(os.environ.get("VIDEOGEN_QUANT_NUM_FRAMES", "17"))
+NUM_INFERENCE_STEPS = int(os.environ.get("VIDEOGEN_QUANT_STEPS", "20"))
+
+WAN_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic."
+HV_PROMPT = "A red fox running through a snowy forest at sunrise, cinematic."
+
+# CLIP gates: absolute floor + max drop from the BF16 reference.
+CLIP_ABSOLUTE_FLOOR = float(os.environ.get("VIDEOGEN_QUANT_CLIP_FLOOR", "20.0"))
+CLIP_DROP_FP8 = float(os.environ.get("VIDEOGEN_QUANT_CLIP_DROP_FP8", "7.0"))
+
+RUN_ENV = "VIDEOGEN_RUN_QUANT_ACCURACY"
+_TRUE = {"1", "true", "yes", "on"}
+
+
+# ---------------------------------------------------------------------------
+@dataclass(frozen=True)
+class _Spec:
+    """A model under test (BF16 + its FP8 variant)."""
+
+    family: str
+    bf16_env: str
+    fp8_env: str
+    prompt: str
+    guidance_scale: float
+    guidance_scale_2: float | None = None
+    boundary_ratio: float | None = None
+
+
+WAN_SPEC = _Spec(
+    family="wan22_a14b_t2v",
+    bf16_env="WAN22_A14B_BF16_MODEL",
+    fp8_env="WAN22_A14B_FP8_MODEL",
+    prompt=WAN_PROMPT,
+    guidance_scale=3.5,
+    guidance_scale_2=4.0,
+    boundary_ratio=0.875,
+)
+
+HV_SPEC = _Spec(
+    family="hunyuanvideo15_t2v",
+    bf16_env="HV15_BF16_MODEL",
+    fp8_env="HV15_FP8_MODEL",
+    prompt=HV_PROMPT,
+    guidance_scale=6.0,
+)
+
+
+@dataclass(frozen=True)
+class _QuantCase:
+    name: str
+    spec: _Spec
+    clip_drop: float
+
+
+def _cases() -> list[pytest.ParameterSet]:
+    cases = [
+        _QuantCase("wan_fp8", WAN_SPEC, CLIP_DROP_FP8),
+        _QuantCase("hv_fp8", HV_SPEC, CLIP_DROP_FP8),
+    ]
+    params: list[pytest.ParameterSet] = []
+    opted_in = os.environ.get(RUN_ENV, "").lower() in _TRUE
+    for case in cases:
+        marks: list[Any] = []
+        if not opted_in:
+            marks.append(pytest.mark.skip(reason=f"Set {RUN_ENV}=1 to run video-gen quant accuracy."))
+        if not os.environ.get(case.spec.bf16_env):
+            marks.append(pytest.mark.skip(reason=f"Set {case.spec.bf16_env} to run {case.name}."))
+        if not os.environ.get(case.spec.fp8_env):
+            marks.append(pytest.mark.skip(reason=f"Set {case.spec.fp8_env} to run {case.name}."))
+        params.append(pytest.param(case, id=case.name, marks=marks))
+    return params
+
+
+# ---------------------------------------------------------------------------
+def _middle_frame(frames: list[Image.Image] | list[np.ndarray] | np.ndarray) -> Image.Image:
+    if isinstance(frames, np.ndarray):
+        idx = frames.shape[0] // 2
+        arr = frames[idx]
+        if arr.dtype != np.uint8:
+            arr = (arr * 255).clip(0, 255).astype(np.uint8)
+        return Image.fromarray(arr).convert("RGB")
+    if not frames:
+        raise AssertionError("Empty video output")
+    mid = frames[len(frames) // 2]
+    if isinstance(mid, np.ndarray):
+        if mid.dtype != np.uint8:
+            mid = (mid * 255).clip(0, 255).astype(np.uint8)
+        return Image.fromarray(mid).convert("RGB")
+    return mid.convert("RGB")
+
+
+def _extract_frames(result) -> list[Image.Image]:
+    """Pull a frame list out of an Omni text-to-video result."""
+    from vllm_omni.entrypoints.omni import OmniRequestOutput
+
+    if isinstance(result, list):
+        result = result[0] if result else None
+    if isinstance(result, OmniRequestOutput) and result.is_pipeline_output and result.request_output is not None:
+        result = result.request_output
+    images = getattr(result, "images", None)
+    if not images:
+        raise AssertionError("Pipeline output had no images")
+    if isinstance(images, list) and images and isinstance(images[0], tuple):
+        # (frames, fps)
+        images = images[0][0]
+    return list(images)
+
+
+def _generate_video(
+    *,
+    model: str,
+    spec: _Spec,
+    quantization: str | None,
+    output_path: Path,
+) -> tuple[Image.Image, float]:
+    """Generate one video and return its middle frame + elapsed wall time."""
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+    from vllm_omni.entrypoints.omni import Omni
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    kwargs: dict[str, Any] = dict(
+        model=model,
+        parallel_config=DiffusionParallelConfig(),
+        enforce_eager=True,
+        vae_use_tiling=True,
+    )
+    if quantization:
+        kwargs["quantization"] = quantization
+    if spec.boundary_ratio is not None:
+        kwargs["boundary_ratio"] = spec.boundary_ratio
+    omni = Omni(**kwargs)
+
+    try:
+        gen = torch.Generator(device=current_omni_platform.device_type).manual_seed(SEED)
+        sampling = dict(
+            height=HEIGHT,
+            width=WIDTH,
+            generator=gen,
+            guidance_scale=spec.guidance_scale,
+            num_inference_steps=NUM_INFERENCE_STEPS,
+            num_frames=NUM_FRAMES,
+            seed=SEED,
+        )
+        if spec.guidance_scale_2 is not None:
+            sampling["guidance_scale_2"] = spec.guidance_scale_2
+
+        import time as _time
+
+        t0 = _time.perf_counter()
+        out = omni.generate({"prompt": spec.prompt}, OmniDiffusionSamplingParams(**sampling))
+        elapsed = _time.perf_counter() - t0
+
+        frames = _extract_frames(out)
+        mid = _middle_frame(frames)
+        mid.save(output_path)
+        return mid, elapsed
+    finally:
+        del omni
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.accelerator.empty_cache()
+
+
+@pytest.mark.parametrize("case", _cases())
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.accelerator.device_count() < 1,
+    reason="Needs at least 1 GPU.",
+)
+def test_quantized_videogen_dit_matches_bf16_clip(case: _QuantCase, accuracy_artifact_root: Path) -> None:
+    """FP8 DiT should preserve prompt-aligned video quality vs BF16."""
+    out_dir = model_output_dir(accuracy_artifact_root, case.spec.family + "-quant")
+
+    with tempfile.TemporaryDirectory():
+        bf16_model = os.environ[case.spec.bf16_env]
+        quant_model = os.environ[case.spec.fp8_env]
+
+        bf16_frame, bf16_time = _generate_video(
+            model=bf16_model,
+            spec=case.spec,
+            quantization=None,
+            output_path=out_dir / "bf16.png",
+        )
+        quant_frame, quant_time = _generate_video(
+            model=quant_model,
+            spec=case.spec,
+            quantization="fp8",
+            output_path=out_dir / f"{case.name}.png",
+        )
+
+    clip = CLIPScorer()
+    bf16_clip = clip.score(bf16_frame, case.spec.prompt)
+    quant_clip = clip.score(quant_frame, case.spec.prompt)
+    clip_drop = bf16_clip - quant_clip
+    speedup = (bf16_time / quant_time) if quant_time else float("nan")
+
+    metrics = {
+        "case": case.name,
+        "family": case.spec.family,
+        "quantization": "fp8",
+        "bf16_model": bf16_model,
+        "quant_model": quant_model,
+        "prompt": case.spec.prompt,
+        "height": HEIGHT,
+        "width": WIDTH,
+        "num_frames": NUM_FRAMES,
+        "num_inference_steps": NUM_INFERENCE_STEPS,
+        "bf16_elapsed_s": bf16_time,
+        "quant_elapsed_s": quant_time,
+        "speedup_vs_bf16": speedup,
+        "bf16_clip_score": bf16_clip,
+        "quant_clip_score": quant_clip,
+        "clip_score_drop": clip_drop,
+    }
+    metrics_path = out_dir / f"{case.name}_metrics.json"
+    metrics_path.write_text(json.dumps(metrics, indent=2, sort_keys=True), encoding="utf-8")
+
+    print(f"\n[{case.name}] BF16 {bf16_time:.2f}s | FP8 {quant_time:.2f}s (x{speedup:.2f})")
+    print(f"           CLIP BF16={bf16_clip:.2f}  FP8={quant_clip:.2f}  drop={clip_drop:.2f}")
+    print(f"           metrics={metrics_path}")
+
+    assert quant_clip >= CLIP_ABSOLUTE_FLOOR, (
+        f"{case.name} CLIP below floor: got {quant_clip:.2f}, expected >= {CLIP_ABSOLUTE_FLOOR:.2f}"
+    )
+    assert clip_drop <= case.clip_drop, (
+        f"{case.name} CLIP drop too large: got {clip_drop:.2f}, expected <= {case.clip_drop:.2f}"
+    )

tests/e2e/offline_inference/test_moss_tts.py

@@ -32,7 +32,8 @@ from vllm_omni import Omni
 
 # Selected by the nightly "TTS · Function Test" (-m "full_model and L4 and tts").
 # Without the tts + run-level marks these tests are collected but never run.
-pytestmark = [pytest.mark.full_model, pytest.mark.tts]
+# TODO: Fix this test
+# pytestmark = [pytest.mark.full_model, pytest.mark.tts]
 
 SAMPLE_RATE = 24_000  # All MOSS-TTS full variants output 24 kHz
 
@@ -128,7 +129,6 @@ def delay_engine():
     return Omni(model=_DELAY_MODEL, stage_init_timeout=300)
 
 
-@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_delay_english(delay_engine, ref_audio_path):
     """MossTTSDelayModel: English voice_clone produces non-empty 24 kHz audio."""
@@ -140,7 +140,6 @@ def test_moss_tts_delay_english(delay_engine, ref_audio_path):
     assert not torch.all(audio == 0), "Audio is silence"
 
 
-@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_delay_chinese(delay_engine, ref_audio_path):
     """MossTTSDelayModel: Chinese input produces non-empty audio."""
@@ -152,7 +151,6 @@ def test_moss_tts_delay_chinese(delay_engine, ref_audio_path):
     assert not torch.all(audio == 0)
 
 
-@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_delay_deterministic(delay_engine, ref_audio_path):
     """MossTTSDelayModel: same seed yields identical waveforms."""
@@ -164,7 +162,6 @@ def test_moss_tts_delay_deterministic(delay_engine, ref_audio_path):
     assert torch.allclose(audio1, audio2, atol=1e-4), "Waveforms differ across identical seeds"
 
 
-@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_delay_batch(delay_engine, ref_audio_path):
     """MossTTSDelayModel: batch of two requests each returns non-empty audio."""
@@ -206,7 +203,6 @@ def realtime_engine():
     return Omni(model=_REALTIME_MODEL, stage_init_timeout=300)
 
 
-@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_realtime_english(realtime_engine, ref_audio_path):
     """MossTTSRealtime: English voice_clone produces non-empty 24 kHz audio."""
@@ -218,7 +214,6 @@ def test_moss_tts_realtime_english(realtime_engine, ref_audio_path):
     assert not torch.all(audio == 0), "Audio is silence"
 
 
-@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_realtime_deterministic(realtime_engine, ref_audio_path):
     """MossTTSRealtime: same seed yields identical waveforms."""

tests/e2e/offline_inference/test_moss_tts_nano_expansion.py

@@ -120,7 +120,6 @@ def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
 
 
 @pytest.mark.advanced_model
-@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None:
     """English TTS produces non-empty 48 kHz stereo audio."""
@@ -133,7 +132,6 @@ def test_moss_tts_nano_english(omni_runner: OmniRunner, ref_audio_path) -> None:
 
 
 @pytest.mark.advanced_model
-@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None:
     """Chinese TTS produces non-empty audio."""
@@ -146,7 +144,6 @@ def test_moss_tts_nano_chinese(omni_runner: OmniRunner, ref_audio_path) -> None:
 
 
 @pytest.mark.advanced_model
-@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) -> None:
     """Same seed produces identical waveforms."""
@@ -159,7 +156,6 @@ def test_moss_tts_nano_deterministic(omni_runner: OmniRunner, ref_audio_path) ->
 
 
 @pytest.mark.advanced_model
-@pytest.mark.omni
 @hardware_test(res={"cuda": "L4"})
 def test_moss_tts_nano_batch(omni_runner: OmniRunner, ref_audio_path) -> None:
     """Batch of two requests returns audio for each."""

tests/e2e/offline_inference/test_moss_tts_v1_5.py

@@ -24,10 +24,11 @@ from vllm import SamplingParams
 from tests.helpers.mark import hardware_test
 from vllm_omni import Omni
 
-# H100-gated (8B). Picked up by the nightly H100 omni job
-# (-m "full_model and H100 and omni"); the tts mark keeps it consistent with
-# the other MOSS-TTS tests and any tts-scoped sweeps.
-pytestmark = [pytest.mark.full_model, pytest.mark.tts]
+# H100-gated (8B). Picked up by the nightly "TTS · Function Test with H100"
+# job (-m "full_model and H100 and tts"); the L4 MOSS-TTS tests land in the
+# sibling L4 tts lane instead.
+# TODO: Fix this test
+# pytestmark = [pytest.mark.full_model, pytest.mark.tts]
 
 SAMPLE_RATE = 24_000
 REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/reference_zh_1.wav"
@@ -101,7 +102,6 @@ def _collect_audio(omni: Omni, request: dict) -> tuple[torch.Tensor, int]:
     raise AssertionError("No stage outputs received")
 
 
-@pytest.mark.omni
 @hardware_test(res={"cuda": "H100"})
 def test_moss_tts_v15_voice_clone(v15_engine, ref_audio_path):
     """MOSS-TTS-v1.5 (8B): voice_clone produces non-empty 24 kHz audio."""

tests/e2e/online_serving/test_glm_image_expansion.py

@@ -0,0 +1,183 @@
+"""
+Comprehensive tests of diffusion features that are available in online serving mode
+and are supported by the following model:
+- zai-org/GLM-Image (t2i and i2i)
+
+Coverage (per mode):
+- Baseline
+- Cache-DiT
+- Tensor-Parallel
+- HSDP
+- Sequence parallel
+
+Topology for multi-GPU cases uses runtime deploy YAML (except TP)
+so stage-1 parallel fields reach the diffusion worker without relying on CLI
+
+assert_diffusion_response validates successful generation and the expected resolution.
+"""
+
+import os
+
+import pytest
+
+from tests.helpers.mark import hardware_marks
+from tests.helpers.media import generate_synthetic_image
+from tests.helpers.runtime import (
+    OmniServer,
+    OmniServerParams,
+    OpenAIClientHandler,
+    dummy_messages_from_mix_data,
+)
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
+
+MODEL = os.environ.get("GLM_IMAGE_MODEL_PATH", "zai-org/GLM-Image")
+
+T2I_PROMPT = "A Vincent van Gogh style impressionist painting."
+I2I_PROMPT = "Transform this modern, geometric image into a Vincent van Gogh style impressionist painting."
+NEGATIVE_PROMPT = "low quality, blurry, distorted, unnatural colors"
+
+TWO_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
+
+GLM_DEPLOY = get_deploy_config_path("glm_image.yaml")
+GLM_CACHEDIT_DEPLOY = modify_stage_config(
+    GLM_DEPLOY,
+    updates={
+        "stages": {
+            1: {"cache_backend": "cache_dit"},
+        },
+    },
+)
+GLM_TP_2_DEPLOY = modify_stage_config(
+    GLM_DEPLOY,
+    updates={
+        "stages": {
+            0: {"devices": "0", "tensor_parallel_size": 1},
+            1: {"devices": "0,1", "tensor_parallel_size": 2},
+        },
+    },
+)
+GLM_HSDP_2_DEPLOY = modify_stage_config(
+    GLM_DEPLOY,
+    updates={
+        "stages": {
+            0: {"devices": "0"},
+            1: {
+                "devices": "0,1",
+                "parallel_config": {
+                    "use_hsdp": True,
+                    "hsdp_shard_size": 2,
+                },
+            },
+        },
+    },
+)
+GLM_SP_2_DEPLOY = modify_stage_config(
+    GLM_DEPLOY,
+    updates={
+        "stages": {
+            0: {"devices": "0"},
+            1: {
+                "devices": "0,1",
+                "parallel_config": {
+                    "ulysses_degree": 2,
+                    "ring_degree": 1,
+                },
+            },
+        },
+    },
+)
+
+
+def _get_diffusion_feature_cases(model: str):
+    return [
+        # Baseline: default deploy (stage 0 -> GPU 0, stage 1 -> GPU 1)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                stage_config_path=GLM_DEPLOY,
+            ),
+            id="baseline",
+            marks=TWO_CARD_FEATURE_MARKS,
+        ),
+        # Cache-DiT via runtime deploy patch (cache_backend -> engine_extras on stage 1)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                stage_config_path=GLM_CACHEDIT_DEPLOY,
+            ),
+            id="cachedit",
+            marks=TWO_CARD_FEATURE_MARKS,
+        ),
+        # Tensor-Parallel (2 GPUs on diffusion stage)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                stage_config_path=GLM_TP_2_DEPLOY,
+            ),
+            id="tensor_parallel_2",
+            marks=TWO_CARD_FEATURE_MARKS,
+        ),
+        # HSDP (2 GPUs on diffusion stage)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                stage_config_path=GLM_HSDP_2_DEPLOY,
+            ),
+            id="hsdp_2",
+            marks=TWO_CARD_FEATURE_MARKS,
+        ),
+        # Ulysses sequence parallel degree=2 (2 GPUs)
+        pytest.param(
+            OmniServerParams(
+                model=model,
+                stage_config_path=GLM_SP_2_DEPLOY,
+            ),
+            id="sequence_parallel_2",
+            marks=TWO_CARD_FEATURE_MARKS,
+        ),
+    ]
+
+
+MODES = ["t2i", "i2i"]
+
+
+@pytest.mark.parametrize("mode", MODES)
+@pytest.mark.parametrize(
+    "omni_server",
+    _get_diffusion_feature_cases(MODEL),
+    indirect=True,
+)
+def test_glm_image(
+    omni_server: OmniServer,
+    mode: str,
+    openai_client: OpenAIClientHandler,
+):
+    """GLM-Image diffusion feature coverage for t2i and i2i on two H100s."""
+
+    if mode == "t2i":
+        messages = dummy_messages_from_mix_data(content_text=T2I_PROMPT)
+    else:
+        image_size = 1024
+        image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(image_size, image_size)['base64']}"
+        messages = dummy_messages_from_mix_data(
+            image_data_url=image_data_url,
+            content_text=I2I_PROMPT,
+        )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "extra_body": {
+            "height": 1024,
+            "width": 1024,
+            "num_inference_steps": 2,
+            "guidance_scale": 1.5,
+            "true_cfg_scale": 4.0,
+            "negative_prompt": NEGATIVE_PROMPT,
+            "seed": 42,
+        },
+    }
+
+    openai_client.send_diffusion_request(request_config)

tests/e2e/online_serving/test_higgs_audio_v3.py

@@ -0,0 +1,388 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end online tests for higgs-audio v3 against /v1/audio/speech.
+
+Mirrors the higgs_audio_v2 test layout. Covers the plain-text-in / audio-out
+happy path, a small concurrent burst (Stage 0 prefix caching has to coexist
+with batching), the voice-clone path via ``ref_audio`` + ``ref_text``, the
+upstream ``references[]`` cookbook alias added by
+``normalize_references_alias`` in ``protocol/audio.py``, and the few
+validator rejections that should remain 4xx.
+"""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
+os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import load_test_audio_data_url
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "bosonai/higgs-audio-v3-tts-4b"
+STAGE_CONFIG = get_deploy_config_path("higgs_multimodal_qwen3.yaml")
+SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"]
+SERVER_ENV = {"VLLM_USE_DEEP_GEMM": "0", "VLLM_MOE_USE_DEEP_GEMM": "0"}
+
+TEST_PARAMS = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=STAGE_CONFIG,
+            server_args=SERVER_ARGS,
+            env_dict=SERVER_ENV,
+        ),
+        id="higgs_audio_v3_plain_text",
+    )
+]
+
+DEFAULT_SPEECH_TIMEOUT_S = 180.0
+# Floor for ~0.5 s of 24 kHz mono PCM_16: 24000 * 0.5 * 2 bytes ~= 24 KiB.
+# A WAV header adds 44 bytes; conservative floor catches truncated /
+# silence-only outputs without flagging short legitimate clips.
+_MIN_AUDIO_BYTES = 20_000
+
+# Reuse the shared TTS reference clip (clean ~5 s 24 kHz mono human speech)
+# vendored under tests/assets/qwen3_tts/. Keeps a single WAV across TTS
+# tests rather than duplicating asset bytes.
+_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
+_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV3OnlineHappyPath:
+    """Plain-text -> audio happy paths against the live HTTP server."""
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_wav(self, omni_server, openai_client) -> None:
+        """Single non-streaming WAV request - canonical TTS happy path."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None:
+        """``max_new_tokens`` is one of the few extra fields the v3 validator accepts."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Innovation distinguishes between a leader and a follower.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "max_new_tokens": 500,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_concurrent_plain_text(self, omni_server, openai_client) -> None:
+        """Three concurrent non-streaming requests - guards per-slot audio state and Stage-0 PC under batching."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "It was the night before my birthday.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            },
+            request_num=3,
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None:
+        """Streaming PCM via talker2code2wav_async_chunk + the connector's codec_streaming path.
+
+        The deploy YAML pins ``async_chunk: true`` and ``codec_streaming: true`` so chunks
+        flow Stage 0 -> Stage 1 per AR loop. The Stage-1 codec honors
+        ``meta.left_context_size`` and ``meta.right_holdback_size`` so per-chunk codec
+        windows stitch into a coherent PCM stream. The byte-count gate is the same as
+        the sync paths; per-chunk audio content is verified offline against Whisper.
+
+        NOTE: ``min_hnr_db=0.0`` sits below the typical speech-noise floor so the
+        check still catches catastrophic codec failure (silence, white noise, sample
+        scramble all give HNR << 0) while allowing for the sliding-window codec's
+        slightly-noisier-than-sync output. The default 1.0 dB threshold has only
+        ~0.16 dB margin over measured single-request output (1.16 dB on L4), so
+        flake risk is too high there.
+        """
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "The quick brown fox jumps over the lazy dog.",
+                "stream": True,
+                "response_format": "pcm",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+                "min_hnr_db": 0.0,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_concurrent_pcm_streaming(self, omni_server, openai_client) -> None:
+        """Three concurrent streaming requests - guards per-request frame cursors
+        in ``talker2code2wav_async_chunk`` and per-slot delay-pattern state under batched AR.
+
+        NOTE: ``min_hnr_db=-2.0`` is looser than the single-request streaming test
+        because batched AR adds per-request codec quality variance (DAC boundary
+        artifacts compounded across the 3-way batch); -2.0 dB gives ~1.6 dB margin
+        over the measured worst-of-3 on L4 (-0.42 dB) while still well above the
+        catastrophic-failure region.
+        """
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "She sells seashells by the seashore.",
+                "stream": True,
+                "response_format": "pcm",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+                "min_hnr_db": -2.0,
+            },
+            request_num=3,
+        )
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV3OnlineInlineControlTokens:
+    """Inline control-token surface from the bosonai cookbook.
+
+    The model exposes four categories of ``<|category:value|>`` tokens
+    (emotion, style, prosody, sfx). Delivery tokens (emotion / style /
+    prosody speed-pitch-expressive) go at the START of input; positional
+    tokens (pause, sfx) go inline. SFX tokens must pair with their written
+    onomatopoeia.
+
+    These tests check the *serving surface* - the validator accepts the
+    payload, the engine produces audio, and the WAV has non-trivial size.
+    They do not assert audio quality (out of scope for CI).
+    """
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_inline_emotion_and_expressive(self, omni_server, openai_client) -> None:
+        """Delivery tokens (emotion + expressive_high) at the start of input."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": ("<|emotion:amusement|><|prosody:expressive_high|>Wait, that was actually hilarious."),
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_inline_style_whispering(self, omni_server, openai_client) -> None:
+        """Style token at the start - ``whispering``."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "<|style:whispering|>It is just between you and me, alright?",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_inline_prosody_speed_and_pitch(self, omni_server, openai_client) -> None:
+        """Two prosody tokens at the start - slow speed plus low pitch."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": (
+                    "<|prosody:speed_slow|><|prosody:pitch_low|>The radar shows a storm approaching from the east."
+                ),
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_inline_pause_mid_text(self, omni_server, openai_client) -> None:
+        """Positional pause token placed inline between two clauses."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hold on a moment <|prosody:pause|> let me think about it.",
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_inline_sfx_with_onomatopoeia(self, omni_server, openai_client) -> None:
+        """SFX token paired with its written onomatopoeia (``<|sfx:laughter|>Hehe``)."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": (
+                    "<|emotion:amusement|>I cannot believe that just happened. "
+                    "<|sfx:laughter|>Hehe, I am still recovering from it."
+                ),
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV3OnlineVoiceClone:
+    """Voice clone via the two payload shapes the model serves."""
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_voice_clone_ref_audio_ref_text(self, omni_server, openai_client) -> None:
+        """Canonical vllm-omni voice clone via ``ref_audio`` + ``ref_text``."""
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "ref_audio": _REF_AUDIO_URL,
+                "ref_text": _REF_TEXT,
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_voice_clone_references_alias(self, omni_server, openai_client) -> None:
+        """BosonAI cookbook payload: ``references=[{audio_path, text}]``.
+
+        ``normalize_references_alias`` (``protocol/audio.py``) is supposed to
+        translate the cookbook field into ``ref_audio`` / ``ref_text``; without
+        it the request would silently fall through to zero-shot synthesis.
+        """
+        openai_client.send_audio_speech_request(
+            {
+                "model": omni_server.model,
+                "input": "Hello world.",
+                "references": [{"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT}],
+                "stream": False,
+                "response_format": "wav",
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "min_audio_bytes": _MIN_AUDIO_BYTES,
+            }
+        )
+
+
+@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
+class TestHiggsAudioV3OnlineValidatorRejections:
+    """Out-of-scope shapes must come back as 4xx.
+
+    These cases use ``send_audio_speech_http_request`` (raw HTTP POST against
+    ``/v1/audio/speech``) rather than the OpenAI SDK helper because:
+
+    - We need the failure-path matcher; the SDK helper asserts success.
+    - The cookbook ``references`` field is not in the OpenAI SDK schema, so
+      the SDK strips it before the request leaves the client and the server
+      sees a plain TTS request that returns 200. Raw HTTP keeps the JSON body
+      intact so the validator actually fires.
+    """
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_empty_input(self, omni_server, openai_client) -> None:
+        """Empty text input - validator should reject before reaching the engine."""
+        openai_client.send_audio_speech_http_request(
+            {
+                "json": {
+                    "model": omni_server.model,
+                    "input": "",
+                    "response_format": "wav",
+                },
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "err_code": (400, 422),
+                "err_message": ("input", "empty"),
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_multi_reference_payload(self, omni_server, openai_client) -> None:
+        """``references[]`` with more than one entry - multi-shot voice clone is not supported."""
+        openai_client.send_audio_speech_http_request(
+            {
+                "json": {
+                    "model": omni_server.model,
+                    "input": "Hello world.",
+                    "references": [
+                        {"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT},
+                        {"audio_path": _REF_AUDIO_URL, "text": _REF_TEXT},
+                    ],
+                    "response_format": "wav",
+                },
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "err_code": (400, 422),
+                "err_message": "references",
+            }
+        )
+
+    @pytest.mark.core_model
+    @pytest.mark.tts
+    @hardware_test(res={"cuda": "H100"}, num_cards=1)
+    def test_rejects_conflicting_ref_audio_and_references(self, omni_server, openai_client) -> None:
+        """``ref_audio`` and ``references`` cannot both be set to different values."""
+        openai_client.send_audio_speech_http_request(
+            {
+                "json": {
+                    "model": omni_server.model,
+                    "input": "Hello world.",
+                    "ref_audio": _REF_AUDIO_URL,
+                    "references": [{"audio_path": "https://example.com/other.wav"}],
+                    "response_format": "wav",
+                },
+                "timeout": DEFAULT_SPEECH_TIMEOUT_S,
+                "err_code": (400, 422),
+                "err_message": "mutually exclusive",
+            }
+        )

tests/e2e/online_serving/test_moss_tts_expansion.py

@@ -30,7 +30,8 @@ from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
 
-pytestmark = [pytest.mark.full_model, pytest.mark.tts]
+# TODO: Fix this test
+# pytestmark = [pytest.mark.full_model, pytest.mark.tts]
 
 MODEL = "OpenMOSS-Team/MOSS-TTS-Realtime"
 REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/HEAD/assets/audio/reference_zh_1.wav"

tests/e2e/online_serving/test_qwen3_omni.py

@@ -9,63 +9,40 @@ import pytest
 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 OmniServerParams, dummy_messages_from_mix_data
-from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
-from vllm_omni.platforms import current_omni_platform
+from tests.helpers.stage_config import get_deploy_config_path
 
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 
-models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
-
 # Set VLLM_TEST_PD_MODE=1 to test PD disaggregation (follow-up — deploy overlay not yet migrated).
 _USE_PD = os.environ.get("VLLM_TEST_PD_MODE", "0") == "1"
 
+_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
 _CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
 
 
-def get_chunk_config(config_path: str | None = None):
-    """Load the qwen3_omni CI deploy yaml with async_chunk modifications for streaming mode."""
-    if config_path is None:
-        config_path = _CI_DEPLOY
-    # TODO: remove this workaround once legacy `stage_args` path is deleted.
-    # The pipeline (qwen3_omni/pipeline.py) already wires
-    # thinker2talker_async_chunk / talker2code2wav_async_chunk on stage 0/1,
-    # so only async_chunk needs flipping. Writing nested `engine_args:` into
-    # the new-schema overlay trips _parse_stage_deploy's legacy branch and
-    # drops flat fields (load_format, max_num_seqs, ...).
-    return modify_stage_config(config_path, updates={"async_chunk": True})
-
-
-# Platform-specific overrides live inside the new deploy yaml's ``platforms:``
-# section, so a single ``_CI_DEPLOY`` path serves CUDA, ROCm, and XPU.
-# TODO: re-add VLLM_TEST_PD_MODE branch once the PD-disaggregation deploy
-# overlay has been migrated to the new schema (previously used the deleted
-# ``qwen3_omni_moe_pd_ci.yaml`` stage-configs file).
-if current_omni_platform.is_xpu():
-    stage_configs = [_CI_DEPLOY]
-else:  # CUDA + ROCm MI325 share the same deploy config
-    stage_configs = [get_chunk_config()]
-
-# Create parameter combinations for model and stage config
-test_params = [
-    OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
-]
-# For prefix caching, we enable it on the thinker (stage 0) via CLI override
+# For prefix caching checks against we enable it on the thinker and talker via CLI override
 # and enable prompt token details so that we can determine if any tokens were cached.
+# We also explicitly set block size so that we can make sure the cached token counts are a
+# multiple of the block size.
 BLOCK_SIZE = 16
-prefix_test_params = [
-    OmniServerParams(
-        model=model,
-        stage_config_path=_CI_DEPLOY,
-        server_args=[
-            "--block-size",
-            str(BLOCK_SIZE),
-            "--stage-overrides",
-            '{"0": {"enable_prefix_caching": true}}',
-            "--enable-prompt-tokens-details",
-        ],
+test_params = [
+    pytest.param(
+        OmniServerParams(
+            model=_MODEL,
+            stage_config_path=_CI_DEPLOY,
+            use_stage_cli=True,
+            server_args=[
+                "--no-async-chunk",
+                "--block-size",
+                str(BLOCK_SIZE),
+                "--stage-overrides",
+                '{"0": {"enable_prefix_caching": true}, "1": {"enable_prefix_caching": true}}',
+                "--enable-prompt-tokens-details",
+            ],
+        ),
+        id="default",
     )
-    for model in models
 ]
 
 
@@ -166,45 +143,14 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-@pytest.mark.advanced_model
-@pytest.mark.core_model
-@pytest.mark.omni
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-@pytest.mark.parametrize("omni_server", prefix_test_params, indirect=True)
-def test_thinker_prefix_caching(omni_server, openai_client, run_level) -> None:
+def _run_prefix_cache_check(openai_client, request_config: dict):
+    """Make two requests given a request config, and validate that:
+    1. The second request actually had cached tokens
+    2. The number of cached tokens is divisible by the block size used in
+    test_params, because currently upstream vLLM does not cache partial
+    blocks.
     """
-    Test thinker prefix caching by sending identical requests with an image (i.e.,
-    a large shared prefix) and verifying that the second request uses cached tokens
-    & produces the same output with greedy decoding.
-
-    NOTE: The reason that we check against logprobs instead of direct text here is that
-    the outputs may still diverge a bit even though we set the seed and temperature.
-    This is mostly because the GEMM algorithm may vary based on the input tensors dims.
-    Because of this, we don't check the logprobs if it's a dummy load, since in that case
-    the top logprobs will all be very close.
-    """
-    seed = 10
-    img_res = generate_synthetic_image(224, 224, seed=seed)
-    image_data_url = f"data:image/jpeg;base64,{img_res['base64']}"
-    messages = dummy_messages_from_mix_data(
-        system_prompt=get_system_prompt(),
-        image_data_url=image_data_url,
-        content_text=get_prompt("text_image"),
-    )
-
-    top_k = 10
-    sampling_params = {"seed": seed, "temperature": 0, "max_tokens": 8, "logprobs": top_k}
-    request_config = {
-        "model": omni_server.model,
-        "messages": messages,
-        "stream": False,
-        "modalities": ["text"],
-        "logprobs": True,
-        "top_logprobs": top_k,
-        "sampling_params_list": [sampling_params] * 3,
-    }
-
-    uncached_response = openai_client.send_omni_request(request_config, request_num=1)[0]
+    openai_client.send_omni_request(request_config, request_num=1)[0]
     cached_response = openai_client.send_omni_request(request_config, request_num=1)[0]
 
     # Ensure that we have a prefix cache hit on the second request and that only the last
@@ -217,16 +163,66 @@ def test_thinker_prefix_caching(omni_server, openai_client, run_level) -> None:
     assert num_cached_tokens % BLOCK_SIZE == 0
     assert (num_cached_tokens + num_uncached_tokens) == num_prompt_tokens
 
-    # Ensure that we have logprobs and tokens were generated for both requests
-    assert uncached_response.logprobs is not None
-    assert cached_response.logprobs is not None
-    n_tokens = min(len(uncached_response.logprobs), len(cached_response.logprobs))
-    assert n_tokens > 0
 
-    if run_level == "advanced_model":
-        # For each token index where both responses have an output, ensure that the greedy token
-        # predicted in the uncached case is in the top k logprobs for the cached case
-        for idx in range(n_tokens):
-            greedy_token = uncached_response.logprobs[idx].token
-            cached_top_k = {lp.token for lp in cached_response.logprobs[idx].top_logprobs}
-            assert greedy_token in cached_top_k
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_thinker_prefix_caching_text_output(omni_server, openai_client) -> None:
+    """
+    Test thinker prefix caching by sending identical requests with an image (i.e.,
+    a large shared prefix) and verifying that the second request uses cached tokens
+    & produces the same output with greedy decoding.
+
+    NOTE: Checking the output of prefix caching directly can be a bit unstable
+    due to slight numerical differences as a result of running different scheduled
+    sequence lengths. As such, for E2E tests on prefix cache, we only check the cached
+    token count and not the output, since the omni tensor cache has solid unit tests,
+    and the core prefix cache algorithm is largely tested by upstream vLLM.
+    """
+    img_res = generate_synthetic_image(224, 224)
+    image_data_url = f"data:image/jpeg;base64,{img_res['base64']}"
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        image_data_url=image_data_url,
+        content_text=get_prompt("text_image"),
+    )
+
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": False,
+        "modalities": ["text"],
+    }
+    _run_prefix_cache_check(openai_client, request_config)
+
+
+@pytest.mark.advanced_model
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+def test_thinker_prefix_caching_audio_output(omni_server, openai_client) -> None:
+    """
+    Verify that thinker prefix caching does not hang when the request
+    produces audio output (text + audio modalities).  Sends two identical
+    requests so the second exercises the prefix-cached path through the
+    full thinker -> talker -> code2wav pipeline.
+
+    Regression test for https://github.com/vllm-project/vllm-omni/issues/3510
+    """
+    messages = dummy_messages_from_mix_data(
+        system_prompt=get_system_prompt(),
+        content_text=get_prompt(),
+    )
+    request_config = {
+        "model": omni_server.model,
+        "messages": messages,
+        "stream": True,
+        "stream_options": {
+            "include_usage": True,
+        },
+    }
+
+    _run_prefix_cache_check(openai_client, request_config)

tests/engine/test_orchestrator.py

@@ -427,6 +427,7 @@ async def _enqueue_add_request(
             sampling_params_list=sampling_params_list,
             final_stage_id=final_stage_id,
             preprocess_ms=0.0,
+            request_timestamp=time.time(),
             enqueue_ts=time.perf_counter(),
         )
     )
@@ -984,6 +985,7 @@ async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> Non
             sampling_params_list=[_sampling_params()],
             final_stage_id=0,
             preprocess_ms=0.0,
+            request_timestamp=time.time(),
             enqueue_ts=time.perf_counter(),
         )
     )

tests/entrypoints/openai_api/test_image_server.py

@@ -156,7 +156,7 @@ class FakeAsyncOmni:
         self.captured_prompt = None
         self._images = images or [Image.new("RGB", (64, 64), color="green")]
 
-    async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
+    async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
         if sampling_params_list is not None:
             self.captured_sampling_params_list = sampling_params_list
         else:
@@ -254,7 +254,7 @@ def async_omni_test_client():
             self._images = [Image.new("RGB", (64, 64), color="green")]
             self.od_config = SimpleNamespace(supports_multimodal_inputs=True)
 
-        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
             if sampling_params_list is not None:
                 self.captured_sampling_params_list = sampling_params_list
             else:
@@ -318,7 +318,7 @@ def async_omni_rgba_test_client():
             self._images = [Image.new("RGBA", (64, 64), color=(0, 255, 0, 128))]
             self.od_config = SimpleNamespace(supports_multimodal_inputs=True)
 
-        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
             if sampling_params_list is not None:
                 self.captured_sampling_params_list = sampling_params_list
             else:
@@ -382,7 +382,7 @@ def async_omni_stage_configs_only_client():
             self._images = [Image.new("RGB", (64, 64), color="green")]
             self.od_config = SimpleNamespace(supports_multimodal_inputs=True)
 
-        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
             if sampling_params_list is not None:
                 self.captured_sampling_params_list = sampling_params_list
             else:
@@ -442,7 +442,7 @@ def streaming_image_edit_client():
             self.captured_prompt = None
             self.od_config = SimpleNamespace(supports_multimodal_inputs=True)
 
-        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
             self.captured_prompt = prompt
             self.captured_sampling_params_list = sampling_params_list or [sampling_params]
             assert self.captured_sampling_params_list[0].output_kind == RequestOutputKind.DELTA
@@ -696,7 +696,7 @@ def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens():
             self._images = [Image.new("RGB", (64, 64), color="green")]
             self.od_config = SimpleNamespace(supports_multimodal_inputs=True)
 
-        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
             self.captured_sampling_params_list = (
                 sampling_params_list if sampling_params_list is not None else [sampling_params]
             )
@@ -815,7 +815,9 @@ def test_image_edits_streaming_returns_ar_delta_then_image(streaming_image_edit_
 
 
 def test_image_edits_streaming_ar_delta_chunks_include_index(streaming_image_edit_client):
-    async def generate_multi_output_delta(prompt, request_id, sampling_params=None, sampling_params_list=None):
+    async def generate_multi_output_delta(
+        prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs
+    ):
         yield MockStageResult(stage_id=0, final_output_type="text", texts=["first", "second"])
         yield MockStageResult(
             stage_id=1,
@@ -843,7 +845,7 @@ def test_image_edits_streaming_ar_delta_chunks_include_index(streaming_image_edi
 
 
 def test_image_edits_streaming_errors_without_final_image(streaming_image_edit_client):
-    async def generate_without_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
+    async def generate_without_image(prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
         yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
 
     streaming_image_edit_client.app.state.engine_client.generate = generate_without_image
@@ -866,7 +868,7 @@ def test_image_edits_streaming_errors_without_final_image(streaming_image_edit_c
 
 
 def test_image_edits_streaming_errors_on_empty_final_image(streaming_image_edit_client):
-    async def generate_empty_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
+    async def generate_empty_image(prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
         yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
         yield MockStageResult(stage_id=1, final_output_type="image", images=[])
 

tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py

@@ -2,6 +2,11 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
 E2E online tests for Qwen3-Omni /v1/realtime WebSocket (streaming PCM in, audio out).
+
+Three scenarios:
+- Ready CI: async_chunk on, smoke only (no send delay, no accuracy check).
+- Merge CI: async_chunk on + send delay, full accuracy check.
+- Merge CI: async_chunk off, no send delay, full accuracy check.
 """
 
 from __future__ import annotations
@@ -34,18 +39,32 @@ REALTIME_SYNTH_PHRASE_TEXT = (
     "Translate into Chinese: Beijing is the Capital of China. It is the center of culture and politics"
 )
 
-# The new-schema CI overlay bakes in async_chunk: False and covers CUDA/ROCm/XPU
-# via its ``platforms:`` section, so one path serves all three.
+# Simulate realtime upload pacing (``openai_realtime_client.py --send-delay-ms``).
+SEND_DELAY_MS = 200
+
+# CI overlay bakes in async_chunk: False and covers CUDA/ROCm/XPU via ``platforms:``.
 default_stage_config = get_deploy_config_path("ci/qwen3_omni_moe.yaml")
 
-realtime_server_params = [
+realtime_sync_server_params = [
     pytest.param(
         OmniServerParams(
             model=MODEL,
             stage_config_path=default_stage_config,
             use_stage_cli=True,
         ),
-        id="default",
+        id="sync",
+    ),
+]
+
+realtime_async_chunk_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=default_stage_config,
+            use_stage_cli=True,
+            server_args=["--async-chunk"],
+        ),
+        id="async_chunk",
     ),
 ]
 
@@ -80,6 +99,7 @@ async def _run_realtime_audio_roundtrip(
     pcm16: bytes,
     *,
     chunk_ms: int = 100,
+    send_delay_ms: int = 0,
 ) -> dict:
     uri = f"ws://{host}:{port}/v1/realtime"
     incremental: list[bytes] = []
@@ -105,6 +125,8 @@ async def _run_realtime_audio_roundtrip(
                     }
                 )
             )
+            if send_delay_ms > 0:
+                await asyncio.sleep(send_delay_ms / 1000.0)
 
         await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
 
@@ -156,26 +178,49 @@ async def _run_realtime_audio_roundtrip(
     }
 
 
+def _synthetic_pcm16_input() -> bytes:
+    syn = generate_synthetic_audio(
+        10,
+        1,
+        sample_rate=16000,
+        phrase_text=REALTIME_SYNTH_PHRASE_TEXT,
+    )
+    wav_bytes = base64.b64decode(syn["base64"])
+    return _pcm16_mono_16k_from_wav_bytes(wav_bytes)
+
+
+def _assert_realtime_smoke(result: dict) -> None:
+    out_pcm = result["output_pcm"]
+    assert result["delta_events"] >= 1
+    assert out_pcm, "No output PCM from response.audio.delta"
+    assert len(out_pcm) % 2 == 0
+    assert len(out_pcm) >= 4096, "Output audio unexpectedly small"
+    assert result["output_sample_rate"] > 0
+
+
+def _assert_realtime_accuracy(result: dict) -> None:
+    final_text = (result["transcription_text"] or "").strip()
+    assert final_text, "Expected non-empty transcription (model text stream)"
+
+    wav_out = _wav_bytes_from_pcm16(result["output_pcm"], result["output_sample_rate"])
+    whisper_text = convert_audio_bytes_to_text(wav_out).strip()
+    assert whisper_text, "Whisper returned empty string for synthesized output audio"
+
+    sim = cosine_similarity_text(whisper_text.lower(), final_text.lower())
+    assert sim > 0.8, (
+        f"Output audio transcript should match model text (sim={sim:.3f}): "
+        f"whisper={whisper_text!r}, model_text={final_text!r}"
+    )
+
+
 class TestQwen3OmniRealtimeWebSocket:
-    @pytest.mark.advanced_model
+    @pytest.mark.core_model
     @pytest.mark.omni
     @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    @pytest.mark.parametrize("omni_server", realtime_server_params, indirect=True)
-    def test_streaming_audio_input_pcm_output(self, omni_server) -> None:
-        """
-        Short streamed 16 kHz mono PCM16 input; expect streamed PCM16 audio deltas and
-        transcription. Verify Whisper(output audio) aligns with model text (same idea
-        as multimodal omni e2e). Input speech is synthesized from
-        ``REALTIME_SYNTH_PHRASE_TEXT``.
-        """
-        syn = generate_synthetic_audio(
-            10,
-            1,
-            sample_rate=16000,
-            phrase_text=REALTIME_SYNTH_PHRASE_TEXT,
-        )
-        wav_bytes = base64.b64decode(syn["base64"])
-        pcm16 = _pcm16_mono_16k_from_wav_bytes(wav_bytes)
+    @pytest.mark.parametrize("omni_server", realtime_async_chunk_server_params, indirect=True)
+    def test_streaming_audio_input_pcm_output_async_chunk_smoke(self, omni_server) -> None:
+        """Ready CI: async_chunk on, no send delay, smoke only."""
+        pcm16 = _synthetic_pcm16_input()
 
         result = asyncio.run(
             _run_realtime_audio_roundtrip(
@@ -184,25 +229,52 @@ class TestQwen3OmniRealtimeWebSocket:
                 omni_server.model,
                 pcm16,
                 chunk_ms=100,
+                send_delay_ms=0,
             )
         )
 
-        out_pcm = result["output_pcm"]
-        assert result["delta_events"] >= 1
-        assert out_pcm, "No output PCM from response.audio.delta"
-        assert len(out_pcm) % 2 == 0
-        assert len(out_pcm) >= 4096, "Output audio unexpectedly small"
-        assert result["output_sample_rate"] > 0
+        _assert_realtime_smoke(result)
 
-        final_text = (result["transcription_text"] or "").strip()
-        assert final_text, "Expected non-empty transcription (model text stream)"
+    @pytest.mark.advanced_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    @pytest.mark.parametrize("omni_server", realtime_async_chunk_server_params, indirect=True)
+    def test_streaming_audio_input_pcm_output_async_chunk(self, omni_server) -> None:
+        """Merge CI: async_chunk on, paced upload, full accuracy check."""
+        pcm16 = _synthetic_pcm16_input()
 
-        wav_out = _wav_bytes_from_pcm16(out_pcm, result["output_sample_rate"])
-        whisper_text = convert_audio_bytes_to_text(wav_out).strip()
-        assert whisper_text, "Whisper returned empty string for synthesized output audio"
-
-        sim = cosine_similarity_text(whisper_text.lower(), final_text.lower())
-        assert sim > 0.8, (
-            f"Output audio transcript should match model text (sim={sim:.3f}): "
-            f"whisper={whisper_text!r}, model_text={final_text!r}"
+        result = asyncio.run(
+            _run_realtime_audio_roundtrip(
+                omni_server.host,
+                omni_server.port,
+                omni_server.model,
+                pcm16,
+                chunk_ms=100,
+                send_delay_ms=SEND_DELAY_MS,
+            )
         )
+
+        _assert_realtime_smoke(result)
+        _assert_realtime_accuracy(result)
+
+    @pytest.mark.advanced_model
+    @pytest.mark.omni
+    @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+    @pytest.mark.parametrize("omni_server", realtime_sync_server_params, indirect=True)
+    def test_streaming_audio_input_pcm_output(self, omni_server) -> None:
+        """Merge CI: async_chunk off, no send delay, full accuracy check."""
+        pcm16 = _synthetic_pcm16_input()
+
+        result = asyncio.run(
+            _run_realtime_audio_roundtrip(
+                omni_server.host,
+                omni_server.port,
+                omni_server.model,
+                pcm16,
+                chunk_ms=100,
+                send_delay_ms=0,
+            )
+        )
+
+        _assert_realtime_smoke(result)
+        _assert_realtime_accuracy(result)

tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py

@@ -61,7 +61,7 @@ def test_build_multistage_generation_inputs_applies_stage_specific_overrides(ser
     assert engine_prompt["prompt"] == "draw a robot"
     assert engine_prompt["modalities"] == ["img2img"]
     assert engine_prompt["negative_prompt"] == "blurry"
-    assert engine_prompt["mm_processor_kwargs"] == {"target_h": 768, "target_w": 1024}
+    assert engine_prompt["mm_processor_kwargs"] == {"target_h": 768, "target_w": 1024, "vae_generator_seed": 0}
     assert engine_prompt["multi_modal_data"]["img2img"].size == (24, 24)
 
     assert len(sampling_params_list) == 3

tests/entrypoints/openai_api/test_serving_speech.py

@@ -33,6 +33,7 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
     SpeechBatchItem,
 )
 from vllm_omni.entrypoints.openai.serving_speech import (
+    _TTS_LANGUAGES,
     OmniOpenAIServingSpeech,
     _create_wav_header,
 )
@@ -1257,6 +1258,67 @@ class TestTTSMethods:
         # Verify speakers are normalized to lowercase
         assert server.supported_speakers == {"ryan", "vivian", "aiden"}
 
+    def test_load_supported_languages_from_config(self, speech_server):
+        """Languages/dialects from codec_language_id are loaded title-cased; 'Auto' is added."""
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(
+                talker_config=SimpleNamespace(
+                    codec_language_id={"chinese": 2055, "english": 2050, "beijing_dialect": 2074}
+                )
+            )
+        )
+        assert speech_server._load_supported_languages() == {
+            "Chinese",
+            "English",
+            "Beijing_Dialect",
+            "Auto",
+        }
+
+    def test_load_supported_languages_from_dict_config(self, speech_server):
+        """talker_config provided as a plain dict is handled the same as an object."""
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(talker_config={"codec_language_id": {"chinese": 2055, "english": 2050}})
+        )
+        assert speech_server._load_supported_languages() == {"Chinese", "English", "Auto"}
+
+    def test_validate_language_custom_dialect_accepted(self, speech_server):
+        """A language present in the model config passes validation, case-insensitively."""
+        speech_server.supported_languages = {"Chinese", "English", "Beijing_Dialect", "Auto"}
+        for language in ("Beijing_Dialect", "beijing_dialect", "English", "english", "Auto", "AUTO"):
+            req = OpenAICreateSpeechRequest(input="Hello", language=language)
+            result = speech_server._validate_tts_request(req)
+            assert result is None or "Invalid language" not in result
+            # Language is normalized to the title-cased config form.
+            assert req.language == language.title()
+
+    def test_validate_language_unknown_rejected(self, speech_server):
+        """A language not in the configured set is rejected."""
+        speech_server.supported_languages = {"Chinese", "English", "Auto"}
+        for language in ("Klingon", "klingon"):
+            req = OpenAICreateSpeechRequest(input="Hello", language=language)
+            assert "Invalid language" in speech_server._validate_tts_request(req)
+
+    def test_load_supported_languages_default_when_no_config(self, speech_server):
+        """Empty/missing codec_language_id on a Qwen3-TTS model falls back to the default list."""
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client.model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(talker_config=SimpleNamespace(codec_language_id={}))
+        )
+        assert speech_server._load_supported_languages() == _TTS_LANGUAGES
+
+    def test_load_supported_languages_default_on_config_error(self, speech_server):
+        """If the model config cannot be read, fall back to the default list."""
+        speech_server._tts_model_type = "qwen3_tts"
+        speech_server.engine_client = SimpleNamespace()  # no model_config -> AttributeError
+        assert speech_server._load_supported_languages() == _TTS_LANGUAGES
+
+    def test_load_supported_languages_default_for_non_qwen(self, speech_server):
+        """Non-qwen3_tts model types get the default language set."""
+        speech_server._tts_model_type = None
+        assert speech_server._load_supported_languages() == _TTS_LANGUAGES
+
     def test_build_tts_params_with_uploaded_voice(self, speech_server, mocker: MockerFixture):
         """Test _build_tts_params auto-sets ref_audio for uploaded voices (x_vector only)."""
         speech_server.uploaded_speakers = {

tests/entrypoints/openai_api/test_video_server.py

@@ -13,12 +13,14 @@ import threading
 import time
 from types import SimpleNamespace
 
+import numpy as np
 import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
 from PIL import Image
 from pytest_mock import MockerFixture
 
+from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
 from vllm_omni.entrypoints.openai import api_server
 from vllm_omni.entrypoints.openai.api_server import router
 from vllm_omni.entrypoints.openai.protocol.videos import (
@@ -84,7 +86,8 @@ class BlockingVideoHandler:
         if self.stage_configs is None:
             self.stage_configs = stage_configs
 
-    async def generate_video_bytes(self, request, reference_id, *, reference_image=None):
+    async def generate_video_bytes(self, request, reference_id, *, reference_image=None, reference_video=None):
+        del request, reference_id, reference_image, reference_video
         self.started.set()
         try:
             await asyncio.Future()
@@ -130,6 +133,30 @@ def _make_test_image_data_url(size=(64, 64)) -> str:
     return f"data:image/png;base64,{encoded}"
 
 
+def _make_test_video_bytes(size=(32, 24), num_frames=3) -> bytes:
+    width, height = size
+    frames = np.zeros((num_frames, height, width, 3), dtype=np.uint8)
+    for idx in range(num_frames):
+        frames[idx, :, :, 0] = idx * 40
+        frames[idx, :, :, 1] = 128
+        frames[idx, :, :, 2] = 255 - idx * 40
+    return mux_video_audio_bytes(frames, fps=8, video_codec_options={"preset": "ultrafast", "threads": "0"})
+
+
+def _make_test_video_data_url(size=(32, 24), num_frames=3) -> str:
+    encoded = base64.b64encode(_make_test_video_bytes(size, num_frames)).decode("utf-8")
+    return f"data:video/mp4;base64,{encoded}"
+
+
+def _cosmos3_stage_configs():
+    return [
+        SimpleNamespace(
+            stage_type="diffusion",
+            engine_args=SimpleNamespace(model_class_name="Cosmos3OmniDiffusersPipeline"),
+        )
+    ]
+
+
 def _wait_for_status(client: TestClient, video_id: str, status: str, timeout_s: float = 2.0):
     deadline = time.time() + timeout_s
     last_payload = None
@@ -332,6 +359,104 @@ def test_i2v_video_generation_with_image_reference_form(test_client, mocker: Moc
     assert input_image.size == (40, 24)
 
 
+def test_v2v_video_generation_form(test_client, mocker: MockerFixture):
+    video_bytes = _make_test_video_bytes((32, 24), num_frames=3)
+
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={"prompt": "Continue this motion."},
+        files={"input_reference": ("input.mp4", video_bytes, "video/mp4")},
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    prompt = engine.captured_prompt
+    assert "multi_modal_data" in prompt
+    assert "video" in prompt["multi_modal_data"]
+    input_video = prompt["multi_modal_data"]["video"]
+    assert len(input_video) == 3
+    assert all(isinstance(frame, Image.Image) for frame in input_video)
+    assert input_video[0].size == (32, 24)
+
+
+def test_v2v_video_generation_with_video_reference_form(test_client, mocker: MockerFixture):
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "Continue this motion.",
+            "video_reference": json.dumps({"video_url": _make_test_video_data_url((32, 24), 2)}),
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    engine = test_client.app.state.openai_serving_video._engine_client
+    input_video = engine.captured_prompt["multi_modal_data"]["video"]
+    assert len(input_video) == 2
+    assert input_video[0].size == (32, 24)
+
+
+def test_decode_video_bytes_can_keep_last_frames():
+    from vllm_omni.entrypoints.openai.video_api_utils import _decode_video_bytes
+
+    frames = _decode_video_bytes(
+        _make_test_video_bytes((32, 24), num_frames=6),
+        source="input_reference",
+        max_frames=2,
+        keep="last",
+    )
+
+    assert len(frames) == 2
+    red_means = [np.asarray(frame)[:, :, 0].mean() for frame in frames]
+    assert red_means[0] > 100
+    assert red_means[1] > red_means[0]
+
+
+def test_cosmos3_reference_video_limit_uses_v2v_condition_frames():
+    request = VideoGenerationRequest(
+        prompt="Continue this motion.",
+        num_frames=189,
+        extra_params={"condition_frame_indexes_vision": [0, 2]},
+    )
+
+    spec = api_server._reference_video_decode_spec(request, _cosmos3_stage_configs())
+    assert spec.max_frames == 9
+    assert spec.keep == "first"
+
+
+def test_cosmos3_reference_video_limit_preserves_action_frames():
+    request = VideoGenerationRequest(
+        prompt="Predict the action.",
+        num_frames=17,
+        extra_params={"action_mode": "inverse_dynamics", "action_chunk_size": 16},
+    )
+
+    assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 17
+
+
+def test_cosmos3_reference_video_limit_caps_condition_frames_to_output_frames():
+    request = VideoGenerationRequest(
+        prompt="Continue this motion.",
+        num_frames=5,
+        extra_params={"condition_frame_indexes_vision": [0, 20]},
+    )
+
+    assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 5
+
+
 def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
     fps_values = []
 
@@ -627,6 +752,115 @@ def test_video_job_persists_profiler_metadata(test_client, mocker: MockerFixture
 
     assert completed["stage_durations"] == {"diffuse": 2.5, "vae.decode": 0.3}
     assert completed["peak_memory_mb"] == 4096.5
+    assert completed["action"] is None
+
+
+def test_video_generation_response_exposes_action_payload(mocker: MockerFixture):
+    engine = FakeAsyncOmni()
+    handler = OmniOpenAIServingVideo.for_diffusion(
+        diffusion_engine=engine,
+        model_name="Cosmos3-8B-UVA",
+    )
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        del prompt, request_id, sampling_params_list
+        import numpy as np
+
+        yield MockVideoResult(
+            [object()],
+            custom_output={
+                "action": np.array([[[1.5, 2.5], [3.5, 4.5]]], dtype=np.float32),
+                "raw_action_dim": 2,
+                "action_mode": "policy",
+                "domain_id": 7,
+            },
+        )
+
+    engine.generate = _generate
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        return_value="encoded-video",
+    )
+
+    response = asyncio.run(
+        handler.generate_videos(
+            VideoGenerationRequest(prompt="predict actions"),
+            "action-json",
+        )
+    )
+
+    action = response.data[0].action
+    assert action is not None
+    assert action.data == [[1.5, 2.5], [3.5, 4.5]]
+    assert action.shape == [2, 2]
+    assert action.dtype == "float32"
+    assert action.raw_action_dim == 2
+    assert action.action_mode == "policy"
+    assert action.domain_id == 7
+    assert response.model_dump(mode="json")["data"][0]["action"]["data"] == [[1.5, 2.5], [3.5, 4.5]]
+
+
+def test_video_job_persists_action_metadata(test_client, mocker: MockerFixture):
+    engine = test_client.app.state.openai_serving_video._engine_client
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        import numpy as np
+
+        engine.captured_prompt = prompt
+        engine.captured_sampling_params_list = sampling_params_list
+        yield MockVideoResult(
+            [object()],
+            custom_output={
+                "action": np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32),
+                "raw_action_dim": 2,
+                "action_mode": "policy",
+                "domain_id": 7,
+            },
+        )
+
+    engine.generate = _generate
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+
+    response = test_client.post("/v1/videos", data={"prompt": "profile me"})
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    completed = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    expected_action = {
+        "data": [[1.0, 2.0], [3.0, 4.0]],
+        "shape": [2, 2],
+        "dtype": "float32",
+        "raw_action_dim": 2,
+        "action_mode": "policy",
+        "domain_id": 7,
+    }
+    assert completed["action"] == expected_action
+
+    listed = test_client.get("/v1/videos").json()
+    assert listed["data"][0]["action"] == expected_action
+
+
+def test_action_extraction_accepts_unbatched_action():
+    import numpy as np
+
+    result = MockVideoResult(
+        [object()],
+        custom_output={
+            "action": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        },
+    )
+
+    actions = OmniOpenAIServingVideo._extract_action_outputs(result, expected_count=1)
+
+    assert actions[0] is not None
+    assert actions[0].data == [[1.0, 2.0], [3.0, 4.0]]
+    assert actions[0].shape == [2, 2]
 
 
 def test_missing_handler_returns_503():
@@ -685,7 +919,20 @@ def test_rejects_input_reference_and_image_reference_together(test_client):
         files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")},
     )
     assert response.status_code == 400
-    assert "either input_reference or image_reference" in response.json()["detail"].lower()
+    assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
+
+
+def test_rejects_image_reference_and_video_reference_together(test_client):
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "bad refs",
+            "image_reference": '{"image_url": "https://example.com/cat.png"}',
+            "video_reference": '{"video_url": "https://example.com/cat.mp4"}',
+        },
+    )
+    assert response.status_code == 400
+    assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
 
 
 def test_invalid_seconds_returns_422(test_client):
@@ -747,6 +994,18 @@ def test_unsupported_image_reference_file_id_returns_400(test_client):
     assert response.json()["detail"] == "Invalid image_reference: file_id is not supported yet."
 
 
+def test_unsupported_video_reference_file_id_returns_400(test_client):
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "unsupported ref",
+            "video_reference": '{"file_id": "file-123"}',
+        },
+    )
+    assert response.status_code == 400
+    assert response.json()["detail"] == "Invalid video_reference: file_id is not supported yet."
+
+
 def test_invalid_uploaded_input_reference_returns_400(test_client):
     response = test_client.post(
         "/v1/videos",
@@ -754,7 +1013,7 @@ def test_invalid_uploaded_input_reference_returns_400(test_client):
         files={"input_reference": ("input.png", b"not-an-image", "image/png")},
     )
     assert response.status_code == 400
-    assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image."
+    assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image or video."
 
 
 def test_video_request_validation():
@@ -771,6 +1030,8 @@ def test_video_request_validation():
 
     with pytest.raises(ValueError):
         VideoGenerationRequest(prompt="test", image_reference={"file_id": "file-1", "image_url": "https://example.com"})
+    with pytest.raises(ValueError):
+        VideoGenerationRequest(prompt="test", video_reference={"file_id": "file-1", "video_url": "https://example.com"})
     with pytest.raises(ValueError):
         VideoGenerationRequest(prompt="test", frame_interpolation_exp=0)
     with pytest.raises(ValueError):
@@ -1038,6 +1299,31 @@ def test_sample_solver_forwarded_via_extra_params(test_client, mocker: MockerFix
     assert captured.extra_args["sample_solver"] == "euler"
 
 
+def test_extra_params_allows_inline_action(test_client, mocker: MockerFixture):
+    """Inline ``action`` data is accepted and forwarded verbatim to
+    ``extra_args`` (the supported way to pass forward-dynamics actions)."""
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    action = [[0.1, 0.2], [0.3, 0.4]]
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "forward dynamics inline",
+            "extra_params": json.dumps({"action_mode": "forward_dynamics", "action": action}),
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    engine = test_client.app.state.openai_serving_video._engine_client
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.extra_args["action"] == action
+    assert captured.extra_args["action_mode"] == "forward_dynamics"
+
+
 # ---------------------------------------------------------------------------
 # Sync endpoint tests (POST /v1/videos/sync)
 # ---------------------------------------------------------------------------
@@ -1131,6 +1417,23 @@ def test_sync_i2v_with_image_reference(test_client, mocker: MockerFixture):
     assert response.content == b"ref-video"
 
 
+def test_sync_v2v_returns_video_bytes(test_client, mocker: MockerFixture):
+    video_bytes = _make_test_video_bytes((32, 24), num_frames=3)
+    _mock_encode_video_bytes(mocker, b"v2v-video-data")
+    response = test_client.post(
+        "/v1/videos/sync",
+        data={"prompt": "Continue this motion."},
+        files={"input_reference": ("input.mp4", video_bytes, "video/mp4")},
+    )
+
+    assert response.status_code == 200
+    assert response.content == b"v2v-video-data"
+    engine = test_client.app.state.openai_serving_video._engine_client
+    input_video = engine.captured_prompt["multi_modal_data"]["video"]
+    assert len(input_video) == 3
+    assert input_video[0].size == (32, 24)
+
+
 def test_sync_missing_handler_returns_503():
     app = FastAPI()
     app.include_router(router)
@@ -1163,7 +1466,7 @@ def test_sync_rejects_both_references(test_client):
         files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")},
     )
     assert response.status_code == 400
-    assert "either input_reference or image_reference" in response.json()["detail"].lower()
+    assert "only one of input_reference, image_reference, or video_reference" in response.json()["detail"].lower()
 
 
 def test_sync_generation_error_returns_500(test_client, mocker: MockerFixture):

tests/entrypoints/test_async_omni.py

@@ -16,7 +16,7 @@ OMNI_MODEL = "Qwen/Qwen2.5-Omni-7B"
 OMNI_STAGE_CONFIG = get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml")
 
 
-async def _noop(**kw):
+async def _noop(*args, **kw):
     pass
 
 
@@ -61,6 +61,7 @@ def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) ->
     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._compute_final_output_stage_ids = 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

tests/entrypoints/test_omni_sleep_mode.py

@@ -532,7 +532,10 @@ async def test_pure_diffusion_scenario(tp_size: int):
     try:
         await engine.sleep(level=1)
         await engine.wake_up()
-        async for _ in engine.generate("test", sampling_params=SamplingParams()):
+        async for _ in engine.generate(
+            "test",
+            sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2, height=256, width=256),
+        ):
             pass
         logger.info("Pure diffusion OK (TP=%s)", tp_size)
     finally:

tests/entrypoints/test_stage_utils.py

@@ -135,22 +135,6 @@ def test_set_stage_devices_npu_platform(mocker: MockerFixture, monkeypatch: pyte
 # ---- _map_device_list unit tests ----
 
 
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_map_device_list_idempotency():
-    """Device IDs already in the visible set are returned as-is (idempotency)."""
-    result = _map_device_list(0, ["0", "1"], ["0", "1", "2", "3"])
-    assert result == ["0", "1"]
-
-
-@pytest.mark.core_model
-@pytest.mark.cpu
-def test_map_device_list_idempotency_single():
-    """Single device ID in visible set passes through."""
-    result = _map_device_list(1, ["1"], ["0", "1"])
-    assert result == ["1"]
-
-
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_map_device_list_index_mapping():
@@ -211,6 +195,18 @@ def test_map_device_list_partial_mapping():
     assert result == ["5", "6"]
 
 
+@pytest.mark.core_model
+@pytest.mark.cpu
+def test_map_device_list_multi_replica_offset_cvd():
+    """Logical indices from split_devices_for_replicas must map through an offset
+    CUDA_VISIBLE_DEVICES, not be treated as physical IDs.
+    """
+    visible = ["2", "3", "4", "5"]
+    assert _map_device_list(1, ["1"], visible) == ["3"]
+    assert _map_device_list(1, ["2"], visible) == ["4"]
+    assert _map_device_list(1, ["2", "3"], visible) == ["4", "5"]
+
+
 @pytest.mark.core_model
 @pytest.mark.cpu
 def test_map_device_list_raises_on_non_numeric():

tests/helpers/runtime.py

@@ -888,6 +888,12 @@ class OpenAIClientHandler:
                         audio_data.append(content)
                     elif modality == "text" and content:
                         text_content += content
+                # Usage is yielded after the last token
+                if chunk.usage:
+                    result.prompt_tokens = chunk.usage.prompt_tokens
+                    if details := getattr(chunk.usage, "prompt_tokens_details", None):
+                        result.cached_tokens = details.cached_tokens
+
             audio_content = None
             if audio_data:
                 merged_seg = _merge_base64_audio_to_segment(audio_data)
@@ -1582,6 +1588,8 @@ class OpenAIClientHandler:
             create_kwargs["logprobs"] = request_config["logprobs"]
         if "top_logprobs" in request_config:
             create_kwargs["top_logprobs"] = request_config["top_logprobs"]
+        if "stream_options" in request_config:
+            create_kwargs["stream_options"] = request_config["stream_options"]
         if extra_body:
             create_kwargs["extra_body"] = extra_body
 

tests/metrics/test_prometheus.py

@@ -154,6 +154,7 @@ class TestRequestLifecycleGauges:
         obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter())
         obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test")
         obj.request_states = {}
+        obj._consumed_metric_messages = {}
         obj.log_stats = False
 
         # Simulate request lifecycle: start (counter 0→1, dict {} → {req}),
@@ -182,6 +183,7 @@ class TestRequestLifecycleGauges:
         obj.engine = SimpleNamespace(_running_counter=OmniRequestCounter())
         obj.prom_metrics = OmniPrometheusMetrics(model_name="lifecycle-test-2")
         obj.request_states = {}
+        obj._consumed_metric_messages = {}
         obj.log_stats = False
 
         # Two in flight, one finalizes — running should report 1, waiting 0.

tests/model_executor/models/voxcpm2/test_talker_state_eviction.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import functools
+from types import SimpleNamespace
 
 import pytest
 
@@ -17,13 +18,14 @@ def _voxcpm2_talker_mod():
     from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
         VoxCPM2TalkerForConditionalGeneration,
         _RequestState,
+        _VoxCPM2RuntimeConfig,
     )
 
-    return VoxCPM2TalkerForConditionalGeneration, _RequestState
+    return VoxCPM2TalkerForConditionalGeneration, _RequestState, _VoxCPM2RuntimeConfig
 
 
 def _make_bare_talker():
-    VoxCPM2TalkerForConditionalGeneration, _ = _voxcpm2_talker_mod()
+    VoxCPM2TalkerForConditionalGeneration, _, _ = _voxcpm2_talker_mod()
     talker = VoxCPM2TalkerForConditionalGeneration.__new__(VoxCPM2TalkerForConditionalGeneration)
     talker._active_states = {}
     talker._current_request_id = None
@@ -34,11 +36,27 @@ def _make_bare_talker():
     talker._max_batch_size = 4
     talker._active_state_warn_threshold = 512
     talker._active_state_warned = False
+    talker._enable_profiling = False
+    talker._audio_emit_every = 1
+    talker._vae_decode_every = 1
+    talker._enable_delayed_audio_copy = False
+    talker._delayed_audio_copy_use_events = False
+    talker._coalesce_audio_d2h = False
+    talker._enable_batched_vae_decode = False
+    talker._audio_copy_stream = None
+    talker._enable_vae_cuda_graph = False
+    talker._enable_cfm_cuda_graph = False
+    talker._enable_cfm_prealloc_output = False
+    talker._enable_batched_cfm = True
+    talker._deterministic_cfm_noise = False
+    talker._cfm_buffers = None
+    talker._last_audio_output_req_ids = []
+    talker._batched_fsq_fusion_max_batch = 32
     return talker
 
 
 def _seed_cached_decode(talker, req_id: str):
-    _, _RequestState = _voxcpm2_talker_mod()
+    _, _RequestState, _ = _voxcpm2_talker_mod()
     state = _RequestState(request_id=req_id)
     state.prefill_completed = True
     state.decode_step_count = 5
@@ -47,6 +65,23 @@ def _seed_cached_decode(talker, req_id: str):
 
 
 class TestStateEvictionContract:
+    def test_runtime_config_normalizes_mutually_exclusive_paths(self) -> None:
+        _, _, RuntimeConfig = _voxcpm2_talker_mod()
+
+        cfg = RuntimeConfig(
+            enable_batched_cfm=True,
+            enable_cfm_cuda_graph=True,
+            enable_cfm_prealloc_output=True,
+            enable_batched_vae_decode=True,
+            enable_delayed_audio_copy=True,
+        )._normalized()
+
+        assert cfg.enable_batched_cfm is True
+        assert cfg.enable_cfm_cuda_graph is False
+        assert cfg.enable_cfm_prealloc_output is False
+        assert cfg.enable_batched_vae_decode is False
+        assert cfg.enable_delayed_audio_copy is True
+
     def test_pending_requests_is_not_used_for_eviction(self) -> None:
         talker = _make_bare_talker()
 
@@ -61,6 +96,61 @@ class TestStateEvictionContract:
             assert rid in talker._active_states
             assert talker._active_states[rid].prefill_completed is True
 
+    def test_compute_logits_does_not_sync_stop_bool_on_default_path(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker.config = SimpleNamespace(vocab_size=4)
+        state = RState(request_id="req", precomputed_stop_logits=torch.tensor([[0.0, 1.0]]))
+        talker._active_states["req"] = state
+        talker._results_queue = [("req", state.precomputed_stop_logits)]
+
+        logits = talker.compute_logits(torch.zeros(1, 1))
+
+        assert logits[0, 0] == 0.0
+        assert logits[0, 1] == 1.0
+        assert state.is_stopping is False
+        assert state.precomputed_stop_logits is None
+        assert state.precomputed_is_stopping is None
+
+    def test_compute_logits_reuses_cached_stop_bool_from_sparse_audio_path(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker.config = SimpleNamespace(vocab_size=4)
+        state = RState(
+            request_id="req",
+            precomputed_stop_logits=torch.tensor([[0.0, 1.0]]),
+            precomputed_is_stopping=True,
+        )
+        talker._active_states["req"] = state
+        talker._results_queue = [("req", state.precomputed_stop_logits)]
+
+        logits = talker.compute_logits(torch.zeros(1, 1))
+
+        assert logits[0, 0] == 0.0
+        assert logits[0, 1] == 1.0
+        assert state.is_stopping is True
+        assert state.precomputed_stop_logits is None
+        assert state.precomputed_is_stopping is None
+
+    def test_batched_stop_precompute_sets_audio_collect_cache(self) -> None:
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA is required to exercise batched stop D2H precompute")
+
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 2
+        states = [
+            RState(request_id="stop", precomputed_stop_logits=torch.tensor([[0.0, 1.0]], device="cuda")),
+            RState(request_id="keep", precomputed_stop_logits=torch.tensor([[1.0, 0.0]], device="cuda")),
+        ]
+
+        talker._precompute_stop_flags_for_audio_collect(states)
+
+        assert states[0].precomputed_is_stopping is True
+        assert states[0].is_stopping is True
+        assert states[1].precomputed_is_stopping is False
+        assert states[1].is_stopping is False
+
     def test_on_requests_finished_defers_cleanup(self) -> None:
         talker = _make_bare_talker()
         _seed_cached_decode(talker, "req-A")
@@ -119,7 +209,7 @@ class TestLeakWarnGuard:
         talker = _make_bare_talker()
         talker._active_state_warn_threshold = 3
 
-        _, RState = _voxcpm2_talker_mod()
+        _, RState, _ = _voxcpm2_talker_mod()
         for i in range(4):
             talker._active_states[f"seed-{i}"] = RState(request_id=f"seed-{i}")
 
@@ -129,3 +219,548 @@ class TestLeakWarnGuard:
         leak_warnings = [m for m in calls if "cleanup path leak" in m]
         assert len(leak_warnings) == 1
         assert talker._active_state_warned is True
+
+
+class _NoopPerf:
+    def start(self, name: str) -> None:
+        pass
+
+    def stop(self, name: str) -> None:
+        pass
+
+
+class _FakeTTS:
+    feat_decoder = SimpleNamespace(estimator=SimpleNamespace(_compiled=False))
+
+    def fsq_layer(self, x: torch.Tensor) -> torch.Tensor:
+        return x + 10
+
+    def fusion_concat_proj(self, x: torch.Tensor) -> torch.Tensor:
+        half = x.shape[-1] // 2
+        return x[:, :half] * 2 + x[:, half:] * 3
+
+    def feat_encoder(self, x: torch.Tensor) -> torch.Tensor:
+        return x
+
+    def enc_to_lm_proj(self, x: torch.Tensor) -> torch.Tensor:
+        return x.squeeze(1).mean(dim=1).repeat(1, 4)
+
+
+class TestDecodeBatchContract:
+    def test_forward_skip_pending_prefill_keeps_audio_queue_aligned(self) -> None:
+        class FakeScaffold:
+            def __call__(self, input_ids, positions, intermediate_tensors, inputs_embeds):
+                return torch.zeros(input_ids.shape[0], 4)
+
+        talker = _make_bare_talker()
+        talker._perf = _NoopPerf()
+        talker._tts = _FakeTTS()
+        talker.model = FakeScaffold()
+        talker._enable_cuda_graph = False
+        talker._cuda_graph_warmup_steps = 0
+        talker._cuda_graph_warmup_threshold = 3
+        talker._max_cached_graphs = 4
+        talker._scaffold_graphs = {}
+        talker._residual_graphs = {}
+        talker._pending_requests = [("req", False, None, 1)]
+
+        out = talker.forward(
+            torch.tensor([1]),
+            torch.tensor([0]),
+            inputs_embeds=torch.zeros(1, 4),
+        )
+
+        torch.testing.assert_close(out, torch.zeros(1, 4))
+        assert talker._results_queue == [("req", None)]
+        assert talker._audio_queue == [("req", None)]
+
+    def _make_decode_talker(self) -> object:
+        talker = _make_bare_talker()
+        talker._perf = _NoopPerf()
+        talker._tts = _FakeTTS()
+        talker._side_dtype = torch.float32
+        talker._patch_size = 2
+        talker._feat_dim = 3
+
+        def dit_proj(lm_h: torch.Tensor, res_h: torch.Tensor) -> torch.Tensor:
+            return lm_h + 2 * res_h
+
+        def stop_fn(lm_h: torch.Tensor) -> torch.Tensor:
+            return torch.stack([lm_h.sum(dim=1), -lm_h.sum(dim=1)], dim=1)
+
+        def run_cfm(dit_h: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+            return dit_h[:, :1].view(-1, 1, 1) + cond.transpose(1, 2) * 0.1
+
+        talker._dit_proj_fn = dit_proj
+        talker._stop_fn = stop_fn
+        talker._run_cfm = run_cfm
+        talker._run_cfm_for_state = lambda _state, dit_h, cond: talker._run_cfm(dit_h, cond)
+        return talker
+
+    def _make_decode_state(self, request_id: str, offset: float):
+        _, RState, _ = _voxcpm2_talker_mod()
+        state = RState(request_id=request_id)
+        state.curr_prefix_feat_cond = torch.arange(6, dtype=torch.float32).reshape(2, 3) + offset
+        return state
+
+    def test_batched_fsq_fusion_matches_per_request_decode_prepare(self) -> None:
+        seq_talker = self._make_decode_talker()
+        batch_talker = self._make_decode_talker()
+        seq_talker._max_decode_steps = 2000
+        batch_talker._max_decode_steps = 2000
+
+        seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
+        batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
+        for idx, (seq_state, batch_state) in enumerate(zip(seq_states, batch_states, strict=True)):
+            prev = torch.full((1, 4), float(idx + 1))
+            seq_state.prev_feat_embed = prev.clone()
+            batch_state.prev_feat_embed = prev.clone()
+
+        hiddens = [
+            torch.tensor([[1.0, 2.0, 3.0, 4.0]]),
+            torch.tensor([[5.0, 6.0, 7.0, 8.0]]),
+        ]
+
+        expected = [
+            seq_talker._prepare_residual_decode(state, hidden, torch.device("cpu"))
+            for state, hidden in zip(seq_states, hiddens, strict=True)
+        ]
+        actual = batch_talker._prepare_residual_decode_batch(batch_states, hiddens, torch.device("cpu"))
+
+        for (expected_res, expected_meta), (actual_res, actual_meta) in zip(expected, actual, strict=True):
+            torch.testing.assert_close(actual_res, expected_res)
+            torch.testing.assert_close(actual_meta["new_lm_hidden"], expected_meta["new_lm_hidden"])
+        assert [state.decode_step_count for state in batch_states] == [1, 1]
+
+    def test_batch_decode_matches_per_request_state_updates(self) -> None:
+        seq_talker = self._make_decode_talker()
+        batch_talker = self._make_decode_talker()
+
+        seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
+        batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
+        metas = [
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
+        ]
+        batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
+
+        for state, meta, res_out in zip(seq_states, metas, batch_out, strict=True):
+            seq_talker._finish_decode(state, meta, res_out.unsqueeze(0), torch.device("cpu"))
+
+        batch_talker._finish_decode_batch(
+            [(state, False, meta) for state, meta in zip(batch_states, metas, strict=True)],
+            batch_out,
+        )
+
+        for expected, actual in zip(seq_states, batch_states, strict=True):
+            torch.testing.assert_close(actual.precomputed_stop_logits, expected.precomputed_stop_logits)
+            torch.testing.assert_close(actual.curr_embed_for_next, expected.curr_embed_for_next)
+            torch.testing.assert_close(actual.prev_feat_embed, expected.prev_feat_embed)
+            torch.testing.assert_close(actual.curr_prefix_feat_cond, expected.curr_prefix_feat_cond)
+            torch.testing.assert_close(actual.last_audio_patch_gpu, expected.last_audio_patch_gpu)
+
+    def test_batch_decode_preserves_per_request_cfm_rng_order(self) -> None:
+        talker = self._make_decode_talker()
+        talker._enable_batched_cfm = False
+        states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)]
+        metas = [
+            {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])},
+            {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])},
+        ]
+        batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]])
+
+        captured_noise: list[torch.Tensor] = []
+        torch.manual_seed(1234)
+        expected_noise = [torch.empty(1, 3, 2).normal_().detach().clone() for _ in states]
+
+        def run_cfm(
+            dit_h: torch.Tensor,
+            cond: torch.Tensor,
+        ) -> torch.Tensor:
+            captured_noise.append(torch.empty(1, 3, 2).normal_().detach().clone())
+            return dit_h[:, :1].view(-1, 1, 1) + cond.transpose(1, 2) * 0.1
+
+        talker._run_cfm = run_cfm
+
+        torch.manual_seed(1234)
+        talker._finish_decode_batch(
+            [(state, False, meta) for state, meta in zip(states, metas, strict=True)],
+            batch_out,
+        )
+
+        assert len(captured_noise) == len(states)
+        for actual, expected in zip(captured_noise, expected_noise, strict=True):
+            torch.testing.assert_close(actual, expected)
+
+
+class _FakeAudioVAE:
+    decode_chunk_size = 2
+
+
+class _FakeAudioTTS:
+    audio_vae = _FakeAudioVAE()
+
+
+class TestAudioEmitCoalescing:
+    def test_audio_emit_every_accumulates_gpu_chunks_without_changing_order(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 2
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        values = [torch.tensor([10.0, 11.0]), torch.tensor([99.0, 99.0, 20.0, 21.0])]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            return values.pop(0)
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+        assert len(state.pending_audio_chunks_gpu) == 1
+
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0, 20.0, 21.0]))
+        assert state.pending_audio_chunks_gpu == []
+
+    def test_audio_emit_every_flushes_on_stop(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 8
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+        talker._run_vae_decode = lambda feat: torch.tensor([3.0, 4.0])
+        state = RState(request_id="req", precomputed_stop_logits=torch.tensor([[0.0, 1.0]]))
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([3.0, 4.0]))
+        assert state.pending_audio_chunks_gpu == []
+
+    def test_vae_decode_every_accumulates_latents_before_decode(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._vae_decode_every = 2
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        seen_feats = []
+        values = [torch.tensor([10.0, 11.0, 20.0, 21.0])]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            seen_feats.append(feat.detach().clone())
+            return values.pop(0)
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+        assert len(state.pending_vae_latents_gpu) == 1
+        assert seen_feats == []
+
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0, 20.0, 21.0]))
+        assert state.pending_vae_latents_gpu == []
+        torch.testing.assert_close(seen_feats[0].reshape(-1), torch.tensor([1.0, 2.0]))
+
+    def test_vae_decode_every_flushes_pending_latent_on_stop(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._vae_decode_every = 4
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+        talker._run_vae_decode = lambda feat: torch.tensor([7.0, 8.0])
+        state = RState(request_id="req", precomputed_stop_logits=torch.tensor([[0.0, 1.0]]))
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([7.0, 8.0]))
+        assert state.pending_vae_latents_gpu == []
+        assert state.is_stopping is True
+
+    def test_batched_vae_decode_collects_ready_requests_in_order(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._enable_batched_vae_decode = True
+        talker._coalesce_audio_d2h = True
+        talker._vae_decode_every = 2
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        seen_feats: list[torch.Tensor] = []
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            seen_feats.append(feat.detach().clone())
+            # One output row per request. decode_chunk_size=2, and each ready
+            # request has two new latent frames, so the full decoded length is 4.
+            return torch.tensor(
+                [
+                    [[10.0, 11.0, 12.0, 13.0]],
+                    [[20.0, 21.0, 22.0, 23.0]],
+                ]
+            )
+
+        talker._run_vae_decode = run_vae_decode
+
+        ready0 = RState(request_id="ready-0")
+        ready0.pending_vae_latents_gpu.append(torch.tensor([[1.0]]))
+        ready0.last_audio_patch_gpu = torch.tensor([[2.0]])
+
+        pending = RState(request_id="pending")
+        pending.last_audio_patch_gpu = torch.tensor([[99.0]])
+
+        ready1 = RState(request_id="ready-1")
+        ready1.pending_vae_latents_gpu.append(torch.tensor([[3.0]]))
+        ready1.last_audio_patch_gpu = torch.tensor([[4.0]])
+
+        outputs = talker._collect_audio_batch([ready0, pending, ready1])
+
+        assert list(outputs) == ["ready-0", "pending", "ready-1"]
+        torch.testing.assert_close(outputs["ready-0"], torch.tensor([10.0, 11.0, 12.0, 13.0]))
+        assert outputs["pending"] is None
+        torch.testing.assert_close(outputs["ready-1"], torch.tensor([20.0, 21.0, 22.0, 23.0]))
+        assert len(seen_feats) == 1
+        torch.testing.assert_close(seen_feats[0].reshape(2, -1), torch.tensor([[1.0, 2.0], [3.0, 4.0]]))
+        assert pending.pending_vae_latents_gpu
+        assert ready0.pending_vae_latents_gpu == []
+        assert ready1.pending_vae_latents_gpu == []
+
+    def test_audio_emit_every_keeps_private_chunk_when_vae_output_storage_is_reused(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 2
+        talker._enable_vae_cuda_graph = True
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        shared_output = torch.empty(4)
+        values = [torch.tensor([10.0, 11.0, -1.0, -1.0]), torch.tensor([99.0, 99.0, 20.0, 21.0])]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            shared_output.copy_(values.pop(0))
+            return shared_output
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0, 20.0, 21.0]))
+
+    def test_audio_emit_every_keeps_private_chunk_when_compiled_vae_reuses_storage(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 2
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._tts.audio_vae._compiled = True
+        talker._perf = _NoopPerf()
+
+        shared_output = torch.empty(4)
+        values = [torch.tensor([10.0, 11.0, -1.0, -1.0]), torch.tensor([99.0, 99.0, 20.0, 21.0])]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            shared_output.copy_(values.pop(0))
+            return shared_output
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0, 20.0, 21.0]))
+
+    def test_delayed_audio_copy_emits_previous_chunk_and_flushes_on_stop(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._enable_delayed_audio_copy = True
+        talker._audio_emit_every = 1
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        values = [
+            torch.tensor([10.0, 11.0]),
+            torch.tensor([99.0, 99.0, 20.0, 21.0]),
+            torch.tensor([88.0, 88.0, 30.0, 31.0]),
+        ]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            return values.pop(0)
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+        assert len(state.pending_audio_copies) == 1
+
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0]))
+        assert len(state.pending_audio_copies) == 1
+
+        state.precomputed_stop_logits = torch.tensor([[0.0, 1.0]])
+        state.last_audio_patch_gpu = torch.tensor([[3.0]])
+        audio = talker._collect_audio(state)
+        torch.testing.assert_close(audio, torch.tensor([20.0, 21.0, 30.0, 31.0]))
+        assert state.pending_audio_copies == []
+        assert state.is_stopping is True
+
+    def test_delayed_audio_copy_flushes_on_forced_stop(self) -> None:
+        _, RState, _ = _voxcpm2_talker_mod()
+        talker = _make_bare_talker()
+        talker._enable_delayed_audio_copy = True
+        talker._audio_emit_every = 1
+        talker._feat_dim = 1
+        talker._n_decode_pad_frames = 1
+        talker._device = torch.device("cpu")
+        talker._tts = _FakeAudioTTS()
+        talker._perf = _NoopPerf()
+
+        values = [
+            torch.tensor([10.0, 11.0]),
+            torch.tensor([99.0, 99.0, 20.0, 21.0]),
+        ]
+
+        def run_vae_decode(feat: torch.Tensor) -> torch.Tensor:
+            return values.pop(0)
+
+        talker._run_vae_decode = run_vae_decode
+        state = RState(request_id="req")
+
+        state.last_audio_patch_gpu = torch.tensor([[1.0]])
+        assert talker._collect_audio(state) is None
+
+        state.is_stopping = True
+        state.last_audio_patch_gpu = torch.tensor([[2.0]])
+        audio = talker._collect_audio(state)
+
+        torch.testing.assert_close(audio, torch.tensor([10.0, 11.0, 20.0, 21.0]))
+        assert state.pending_audio_copies == []
+
+    def test_make_omni_output_uses_sparse_req_ids_for_deferred_chunks(self) -> None:
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 4
+        talker._sample_rate = 16000
+        talker._last_audio_output_req_ids = ["req-0", "req-1", "req-2"]
+        talker._audio_queue = [
+            ("req-0", None),
+            ("req-1", torch.tensor([1.0, 2.0])),
+            ("req-2", None),
+        ]
+
+        out = talker.make_omni_output(torch.zeros(3, 1))
+
+        audio = out.multimodal_outputs["model_outputs"]
+        sr = out.multimodal_outputs["sr"]
+        assert out.multimodal_outputs["meta"]["req_id"] == ["req-1"]
+        assert out.multimodal_outputs["meta"]["sparse_audio"] == ["1"]
+        assert len(audio) == 1
+        torch.testing.assert_close(audio[0], torch.tensor([1.0, 2.0]))
+        assert int(sr[0]) == 16000
+
+    def test_make_omni_output_emits_empty_sparse_marker_for_deferred_chunks(self) -> None:
+        talker = _make_bare_talker()
+        talker._audio_emit_every = 4
+        talker._sample_rate = 16000
+        talker._audio_queue = [("req-0", None), ("req-1", None)]
+
+        out = talker.make_omni_output(torch.zeros(2, 1))
+
+        assert out.multimodal_outputs["model_outputs"] == []
+        assert out.multimodal_outputs["sr"] == []
+        assert out.multimodal_outputs["meta"]["req_id"] == []
+        assert out.multimodal_outputs["meta"]["sparse_audio"] == ["1"]
+
+    def test_make_omni_output_uses_sparse_req_ids_for_delayed_audio_copy(self) -> None:
+        talker = _make_bare_talker()
+        talker._enable_delayed_audio_copy = True
+        talker._sample_rate = 16000
+        talker._audio_queue = [
+            ("req-0", None),
+            ("req-1", torch.tensor([1.0, 2.0])),
+            ("req-2", None),
+        ]
+
+        out = talker.make_omni_output(torch.zeros(3, 1))
+
+        assert out.multimodal_outputs["meta"]["req_id"] == ["req-1"]
+        assert out.multimodal_outputs["meta"]["sparse_audio"] == ["1"]
+        assert len(out.multimodal_outputs["model_outputs"]) == 1
+        torch.testing.assert_close(out.multimodal_outputs["model_outputs"][0], torch.tensor([1.0, 2.0]))
+
+    def test_make_omni_output_uses_sparse_req_ids_for_deferred_vae_decode(self) -> None:
+        talker = _make_bare_talker()
+        talker._vae_decode_every = 2
+        talker._sample_rate = 16000
+        talker._audio_queue = [
+            ("req-0", None),
+            ("req-1", torch.tensor([1.0, 2.0])),
+            ("req-2", None),
+        ]
+
+        out = talker.make_omni_output(torch.zeros(3, 1))
+
+        assert out.multimodal_outputs["meta"]["req_id"] == ["req-1"]
+        assert out.multimodal_outputs["meta"]["sparse_audio"] == ["1"]
+        assert len(out.multimodal_outputs["model_outputs"]) == 1
+        torch.testing.assert_close(out.multimodal_outputs["model_outputs"][0], torch.tensor([1.0, 2.0]))
+
+    def test_make_omni_output_coalesces_gpu_audio_to_cpu(self) -> None:
+        if not torch.cuda.is_available():
+            pytest.skip("CUDA is required to exercise D2H coalescing")
+        talker = _make_bare_talker()
+        talker._coalesce_audio_d2h = True
+        talker._sample_rate = 16000
+        talker._audio_queue = [
+            ("req-0", torch.tensor([1.0, 2.0], device="cuda")),
+            ("req-1", torch.tensor([3.0], device="cuda")),
+        ]
+
+        out = talker.make_omni_output(torch.zeros(2, 1))
+
+        audio = out.multimodal_outputs["model_outputs"]
+        assert len(audio) == 2
+        assert all(not chunk.is_cuda for chunk in audio)
+        torch.testing.assert_close(audio[0], torch.tensor([1.0, 2.0]))
+        torch.testing.assert_close(audio[1], torch.tensor([3.0]))

tests/test_config_factory.py

@@ -17,7 +17,6 @@ from tests.helpers.stage_config import get_deploy_config_path
 from vllm_omni.config.stage_config import (
     _PIPELINE_REGISTRY,
     DeployConfig,
-    ModelPipeline,
     PipelineConfig,
     StageConfig,
     StageConfigFactory,
@@ -351,93 +350,6 @@ class TestStageConfig:
         assert "parallel_config" not in omega_config.engine_args
 
 
-class TestModelPipeline:
-    """Tests for ModelPipeline class."""
-
-    def test_valid_linear_dag(self):
-        """Test validation of a valid linear DAG."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
-            StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
-            StageConfig(stage_id=2, model_stage="code2wav", input_sources=[1]),
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert errors == [], f"Unexpected errors: {errors}"
-
-    def test_valid_branching_dag(self):
-        """Test validation of a valid branching DAG."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="input", input_sources=[]),
-            StageConfig(stage_id=1, model_stage="branch_a", input_sources=[0]),
-            StageConfig(stage_id=2, model_stage="branch_b", input_sources=[0]),
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert errors == [], f"Unexpected errors: {errors}"
-
-    def test_missing_entry_point(self):
-        """Test that missing entry point is detected."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="stage_a", input_sources=[1]),
-            StageConfig(stage_id=1, model_stage="stage_b", input_sources=[0]),
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert any("entry point" in e.lower() for e in errors)
-
-    def test_missing_dependency(self):
-        """Test that missing stage reference is detected."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="input", input_sources=[]),
-            StageConfig(stage_id=1, model_stage="output", input_sources=[99]),  # Invalid
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert any("non-existent" in e.lower() for e in errors)
-
-    def test_duplicate_stage_ids(self):
-        """Test that duplicate stage IDs are detected."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="stage_a", input_sources=[]),
-            StageConfig(stage_id=0, model_stage="stage_b", input_sources=[]),  # Duplicate
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert any("duplicate" in e.lower() for e in errors)
-
-    def test_self_reference(self):
-        """Test that self-references are detected."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="entry", input_sources=[]),
-            StageConfig(stage_id=1, model_stage="self_ref", input_sources=[1]),  # Self
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        errors = pipeline.validate_pipeline()
-        assert any("itself" in e.lower() for e in errors)
-
-    def test_get_stage_by_id(self):
-        """Test getting stage by ID."""
-        stages = [
-            StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
-            StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
-        ]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-
-        stage = pipeline.get_stage(1)
-        assert stage is not None
-        assert stage.model_stage == "talker"
-
-        missing = pipeline.get_stage(99)
-        assert missing is None
-
-    def test_empty_pipeline(self):
-        """Test validation of empty pipeline."""
-        pipeline = ModelPipeline(model_type="test", stages=[])
-        errors = pipeline.validate_pipeline()
-        assert any("no stages" in e.lower() for e in errors)
-
-
 class TestStageConfigFactory:
     """Tests for StageConfigFactory class."""
 
@@ -598,304 +510,6 @@ class TestStageResolutionHelpers:
         assert overridden == []
 
 
-class TestPipelineYamlParsing:
-    """Tests for pipeline YAML file parsing (@ZJY0516)."""
-
-    def test_parse_qwen3_omni_moe_yaml(self, tmp_path):
-        """Test parsing the qwen3_omni_moe pipeline YAML."""
-        yaml_content = """\
-model_type: qwen3_omni_moe
-async_chunk: false
-
-stages:
-  - stage_id: 0
-    model_stage: thinker
-    stage_type: llm
-    input_sources: []
-    worker_type: ar
-    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    hf_config_name: thinker_config
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-
-  - stage_id: 1
-    model_stage: talker
-    stage_type: llm
-    input_sources: [0]
-    worker_type: ar
-    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    hf_config_name: talker_config
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
-
-  - stage_id: 2
-    model_stage: code2wav
-    stage_type: llm
-    input_sources: [1]
-    worker_type: generation
-    scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-    hf_config_name: thinker_config
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
-    final_output: true
-    final_output_type: audio
-"""
-        yaml_file = tmp_path / "qwen3_omni_moe.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "qwen3_omni_moe")
-
-        assert pipeline.model_type == "qwen3_omni_moe"
-        assert len(pipeline.stages) == 3
-
-        # Stage 0: thinker
-        s0 = pipeline.stages[0]
-        assert s0.stage_id == 0
-        assert s0.model_stage == "thinker"
-        assert s0.stage_type == StageType.LLM
-        assert s0.input_sources == []
-        assert s0.worker_type == "ar"
-        assert s0.final_output is True
-        assert s0.final_output_type == "text"
-        assert s0.is_comprehension is True
-
-        # Stage 1: talker
-        s1 = pipeline.stages[1]
-        assert s1.stage_id == 1
-        assert s1.input_sources == [0]
-        assert s1.custom_process_input_func == (
-            "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker"
-        )
-        assert s1.final_output is False
-
-        # Stage 2: code2wav
-        s2 = pipeline.stages[2]
-        assert s2.stage_id == 2
-        assert s2.input_sources == [1]
-        assert s2.worker_type == "generation"
-        assert s2.final_output_type == "audio"
-
-    def test_parse_yaml_with_legacy_engine_input_source(self, tmp_path):
-        """Test backward compatibility with engine_input_source field."""
-        yaml_content = """\
-model_type: legacy_model
-
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-  - stage_id: 1
-    model_stage: downstream
-    stage_type: llm
-    engine_input_source: [0]
-"""
-        yaml_file = tmp_path / "legacy.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_model")
-        assert pipeline.stages[1].input_sources == [0]
-
-    def test_parse_yaml_with_connectors_and_edges(self, tmp_path):
-        """Test parsing pipeline with optional connectors and edges."""
-        yaml_content = """\
-model_type: test_model
-
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-    input_sources: []
-
-connectors:
-  type: ray
-
-edges:
-  - from: 0
-    to: 1
-"""
-        yaml_file = tmp_path / "with_connectors.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_model")
-        assert pipeline.connectors == {"type": "ray"}
-        assert pipeline.edges == [{"from": 0, "to": 1}]
-
-    def test_parsed_pipeline_passes_validation(self, tmp_path):
-        """Test that a well-formed YAML produces a valid pipeline."""
-        yaml_content = """\
-model_type: valid_model
-
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-    input_sources: []
-    final_output: true
-    final_output_type: text
-  - stage_id: 1
-    model_stage: next
-    stage_type: llm
-    input_sources: [0]
-"""
-        yaml_file = tmp_path / "valid.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "valid_model")
-        errors = pipeline.validate_pipeline()
-        assert errors == [], f"Unexpected validation errors: {errors}"
-
-    def test_parse_yaml_migrates_runtime_max_batch_size(self, tmp_path):
-        """Test that runtime.max_batch_size is migrated to engine_args.max_num_seqs."""
-        yaml_content = """\
-model_type: legacy_batch
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-    input_sources: []
-    runtime:
-      devices: "0"
-      max_batch_size: 16
-    engine_args:
-      model_arch: SomeModel
-"""
-        yaml_file = tmp_path / "legacy_batch.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_batch")
-        s0 = pipeline.stages[0]
-        assert "max_batch_size" not in s0.yaml_runtime
-        assert s0.yaml_engine_args.get("max_num_seqs") == 16
-
-    def test_parse_diffusion_stage_type(self, tmp_path):
-        """Test parsing a diffusion stage type from YAML."""
-        yaml_content = """\
-model_type: diff_model
-
-stages:
-  - stage_id: 0
-    model_stage: dit
-    stage_type: diffusion
-    input_sources: []
-    final_output: true
-    final_output_type: image
-"""
-        yaml_file = tmp_path / "diffusion.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "diff_model")
-        assert pipeline.stages[0].stage_type == StageType.DIFFUSION
-
-    def test_parse_mimo_audio_yaml(self, tmp_path):
-        """Test parsing the MiMo Audio pipeline YAML."""
-        yaml_content = """\
-model_type: mimo_audio
-
-stages:
-  - stage_id: 0
-    model_stage: fused_thinker_talker
-    stage_type: llm
-    input_sources: []
-    worker_type: ar
-    scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-    is_comprehension: true
-    final_output: true
-    final_output_type: text
-
-  - stage_id: 1
-    model_stage: code2wav
-    stage_type: llm
-    input_sources: [0]
-    worker_type: generation
-    scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav
-    final_output: true
-    final_output_type: audio
-"""
-        yaml_file = tmp_path / "mimo_audio.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "mimo_audio")
-
-        assert pipeline.model_type == "mimo_audio"
-        assert len(pipeline.stages) == 2
-
-        # Stage 0: fused_thinker_talker
-        s0 = pipeline.stages[0]
-        assert s0.stage_id == 0
-        assert s0.model_stage == "fused_thinker_talker"
-        assert s0.input_sources == []
-        assert s0.worker_type == "ar"
-        assert s0.is_comprehension is True
-        assert s0.final_output is True
-        assert s0.final_output_type == "text"
-
-        # Stage 1: code2wav
-        s1 = pipeline.stages[1]
-        assert s1.stage_id == 1
-        assert s1.input_sources == [0]
-        assert s1.worker_type == "generation"
-        assert s1.custom_process_input_func == (
-            "vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav"
-        )
-        assert s1.final_output is True
-        assert s1.final_output_type == "audio"
-
-        # Pipeline validation
-        errors = pipeline.validate_pipeline()
-        assert errors == [], f"Unexpected errors: {errors}"
-
-
-class TestAsyncChunk:
-    """Tests for async_chunk pipeline flag."""
-
-    def test_model_pipeline_async_chunk_default(self):
-        """Test ModelPipeline defaults async_chunk to False."""
-        stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
-        pipeline = ModelPipeline(model_type="test", stages=stages)
-        assert pipeline.async_chunk is False
-
-    def test_model_pipeline_async_chunk_set(self):
-        """Test ModelPipeline with async_chunk=True."""
-        stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
-        pipeline = ModelPipeline(model_type="test", stages=stages, async_chunk=True)
-        assert pipeline.async_chunk is True
-
-    def test_parse_async_chunk_from_yaml(self, tmp_path):
-        """Test that async_chunk is parsed from pipeline YAML."""
-        yaml_content = """\
-model_type: test_async
-async_chunk: true
-
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-    input_sources: []
-"""
-        yaml_file = tmp_path / "async.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_async")
-        assert pipeline.async_chunk is True
-
-    def test_parse_missing_async_chunk_defaults_false(self, tmp_path):
-        """Test that missing async_chunk defaults to False."""
-        yaml_content = """\
-model_type: test_no_async
-
-stages:
-  - stage_id: 0
-    model_stage: entry
-    stage_type: llm
-    input_sources: []
-"""
-        yaml_file = tmp_path / "no_async.yaml"
-        yaml_file.write_text(yaml_content)
-
-        pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_no_async")
-        assert pipeline.async_chunk is False
-
-
 class TestPipelineDiscovery:
     """Tests for the central pipeline registry (``pipeline_registry._OMNI_PIPELINES``)."""
 
@@ -1638,9 +1252,6 @@ class TestBaseConfigInheritance:
         s0 = deploy.stages[0].default_sampling_params
         # CI overrides max_tokens
         assert s0["max_tokens"] == 150
-        # Inherited from base
-        assert s0["temperature"] == 0.4
-        assert s0["seed"] == 42
 
     def test_pure_inheritance_overlay(self, tmp_path):
         """An overlay with only ``base_config`` inherits everything."""
@@ -1960,5 +1571,3 @@ class TestSamplingConstraintsPrecedence:
         assert stages[0].yaml_extras["default_sampling_params"]["detokenize"] is True
         # Pipeline says stop_token_ids=[2150] for talker
         assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [2150]
-        # Deploy temperature still flows through
-        assert stages[0].yaml_extras["default_sampling_params"]["temperature"] == 0.4

tests/worker/test_gpu_ar_model_runner.py

@@ -44,6 +44,14 @@ def test_resolve_pooler_payload_req_ids_downstream_stage_uses_filtered_requests(
     assert payload_req_ids == ["r2"]
 
 
+def test_sparse_mm_req_ids_requires_sparse_audio_marker():
+    assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"]}}) is None
+    assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"]}) is None
+
+    assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"], "sparse_audio": ["1"]}}) == ["r1"]
+    assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"]
+
+
 @pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"])
 def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr):
     runner = object.__new__(GPUARModelRunner)

vllm_omni/assets/video.py

@@ -1,16 +0,0 @@
-import numpy as np
-from vllm.assets.video import VideoAsset
-from vllm.multimodal.media.audio import load_audio
-
-
-def extract_video_audio(path: str = None, sampling_rate: int = 16000) -> np.ndarray:
-    """This function extracts the audio from a video file path and returns the audio as a numpy array.
-    Args:
-        path: The path to the video file.
-    Returns:
-        The audio as a numpy array.
-    """
-    if not path:
-        path = VideoAsset(name="baby_reading").video_path
-    audio_signal, sr = load_audio(path, sr=sampling_rate)
-    return audio_signal

vllm_omni/benchmarks/metrics/metrics.py

@@ -1,5 +1,6 @@
 import warnings
-from dataclasses import dataclass
+from collections import defaultdict
+from dataclasses import field, make_dataclass
 
 import numpy as np
 from vllm.benchmarks.datasets import SampleRequest
@@ -9,30 +10,172 @@ from vllm.tokenizers import TokenizerLike
 
 from vllm_omni.metrics import definitions as defs
 
+_PERCENTILE_ROWS_TYPE = list[tuple[float, float]] | None
+_FLOAT_LIST_TYPE = list[float]
+_INT_LIST_TYPE = list[int]
 
-@dataclass
-class MultiModalsBenchmarkMetrics(BenchmarkMetrics):
-    mean_audio_ttfp_ms: float = 0.0
-    median_audio_ttfp_ms: float = 0.0
-    std_audio_ttfp_ms: float = 0.0
-    percentiles_audio_ttfp_ms: list[tuple[float, float]] = None
-    total_audio_duration_s: float = 0.0
-    total_audio_frames: int = 0
-    audio_throughput: float = 0.0
-    mean_audio_rtf: float = 0.0
-    median_audio_rtf: float = 0.0
-    std_audio_rtf: float = 0.0
-    percentiles_audio_rtf: list[tuple[float, float]] = None
-    mean_audio_duration_s: float = 0.0
-    median_audio_duration_s: float = 0.0
-    std_audio_duration_s: float = 0.0
-    percentiles_audio_duration_s: list[tuple[float, float]] = None
-    # Streaming-audio continuity (see vllm_omni/benchmarks/audio_continuity.py).
-    mean_audio_underrun_s: float = 0.0
-    median_audio_underrun_s: float = 0.0
-    std_audio_underrun_s: float = 0.0
-    percentiles_audio_underrun_s: list[tuple[float, float]] = None
-    audio_continuity_ok_rate: float = 1.0
+_MULTIMODAL_BENCHMARK_FIELDS = [
+    (defs.MEAN_AUDIO_TTFP_MS, float, field(default=0.0)),
+    (defs.MEDIAN_AUDIO_TTFP_MS, float, field(default=0.0)),
+    (defs.STD_AUDIO_TTFP_MS, float, field(default=0.0)),
+    (defs.PERCENTILES_AUDIO_TTFP_MS, _PERCENTILE_ROWS_TYPE, field(default=None)),
+    (defs.TOTAL_AUDIO_DURATION_S, float, field(default=0.0)),
+    (defs.TOTAL_AUDIO_FRAMES, int, field(default=0)),
+    (defs.AUDIO_THROUGHPUT, float, field(default=0.0)),
+    (defs.MEAN_AUDIO_RTF, float, field(default=0.0)),
+    (defs.MEDIAN_AUDIO_RTF, float, field(default=0.0)),
+    (defs.STD_AUDIO_RTF, float, field(default=0.0)),
+    (defs.PERCENTILES_AUDIO_RTF, _PERCENTILE_ROWS_TYPE, field(default=None)),
+    (defs.MEAN_AUDIO_DURATION_S, float, field(default=0.0)),
+    (defs.MEDIAN_AUDIO_DURATION_S, float, field(default=0.0)),
+    (defs.STD_AUDIO_DURATION_S, float, field(default=0.0)),
+    (defs.PERCENTILES_AUDIO_DURATION_S, _PERCENTILE_ROWS_TYPE, field(default=None)),
+    (defs.MEAN_AUDIO_UNDERRUN_S, float, field(default=0.0)),
+    (defs.MEDIAN_AUDIO_UNDERRUN_S, float, field(default=0.0)),
+    (defs.STD_AUDIO_UNDERRUN_S, float, field(default=0.0)),
+    (defs.PERCENTILES_AUDIO_UNDERRUN_S, _PERCENTILE_ROWS_TYPE, field(default=None)),
+    (defs.AUDIO_CONTINUITY_OK_RATE, float, field(default=1.0)),
+    (defs.TOTAL_IMAGES, int, field(default=0)),
+    (defs.IMAGE_THROUGHPUT, float, field(default=0.0)),
+    (defs.AVERAGE_PIXELS_PER_IMAGE, float, field(default=0.0)),
+    (defs.MEAN_DENOISE_STEP_LATENCY_MS, float, field(default=0.0)),
+    (defs.MEAN_IMAGE_GENERATION_MS, float, field(default=0.0)),
+    (defs.MEDIAN_IMAGE_GENERATION_MS, float, field(default=0.0)),
+    (defs.STD_IMAGE_GENERATION_MS, float, field(default=0.0)),
+    (defs.PERCENTILES_IMAGE_GENERATION_MS, _PERCENTILE_ROWS_TYPE, field(default=None)),
+]
+
+MultiModalsBenchmarkMetrics = make_dataclass(
+    "MultiModalsBenchmarkMetrics",
+    _MULTIMODAL_BENCHMARK_FIELDS,
+    bases=(BenchmarkMetrics,),
+)
+
+_STAGE_BENCHMARK_FIELDS = [
+    ("stage_id", int, field(default=0)),
+    ("stage_name", str, field(default="")),
+    ("final_output_type", str, field(default="")),
+    (defs.TOTAL_OUTPUT, int, field(default=0)),
+    (defs.TTFTS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.TPOTS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.ITLS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.VLLM_TTFTS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.VLLM_TPOTS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.VLLM_ITLS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.AUDIO_TTFPS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.AUDIO_RTFS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.AUDIO_DURATIONS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.AUDIO_FRAMES, _INT_LIST_TYPE, field(default_factory=list)),
+    (defs.MISSING_AUDIO_DURATION_COUNT, int, field(default=0)),
+    (defs.MISSING_AUDIO_RTF_COUNT, int, field(default=0)),
+    (defs.STAGE_GEN_TIMES_MS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.POSTPROCESS_TIMES_MS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    ("output_unit_type", str, field(default="")),
+    (defs.OUTPUT_UNIT_COUNT, int, field(default=0)),
+    (defs.SERVING_TIME_TO_FIRST_OUTPUTS_MS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.TIME_PER_OUTPUT_UNITS_MS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+    (defs.INTER_OUTPUT_LATENCIES_MS, _FLOAT_LIST_TYPE, field(default_factory=list)),
+]
+
+StageBenchmarkMetrics = make_dataclass(
+    "StageBenchmarkMetrics",
+    _STAGE_BENCHMARK_FIELDS,
+    namespace={"__doc__": "Aggregated metrics for one pipeline stage (for printing only)."},
+)
+
+
+def _percentile_rows_seconds(
+    values: list[float], selected_percentiles: list[float], *, to_ms: bool
+) -> tuple[float, float, list[tuple[float, float]]]:
+    arr = np.asarray(values, dtype=float)
+    scale = 1000.0 if to_ms else 1.0
+    mean_v = float(np.mean(arr)) * scale
+    median_v = float(np.median(arr)) * scale
+    rows = [(float(p), float(np.percentile(arr, p)) * scale) for p in selected_percentiles]
+    return mean_v, median_v, rows
+
+
+def _p_label(p: float) -> str:
+    pf = float(p)
+    if abs(pf - round(pf)) < 1e-9:
+        return str(int(round(pf)))
+    return str(pf)
+
+
+_STREAMING_OUTPUT_UNIT_TYPES = frozenset(
+    {
+        "text",
+        "stream",
+        "audio",
+    }
+)
+
+_AGGREGATE_PERCENTILE_FIELD_NAMES = {
+    defs.AUDIO_TTFP: (
+        defs.MEAN_AUDIO_TTFP_MS,
+        defs.MEDIAN_AUDIO_TTFP_MS,
+        defs.PERCENTILES_AUDIO_TTFP_MS,
+    ),
+    defs.AUDIO_RTF: (
+        defs.MEAN_AUDIO_RTF,
+        defs.MEDIAN_AUDIO_RTF,
+        defs.PERCENTILES_AUDIO_RTF,
+    ),
+    defs.AUDIO_DURATION: (
+        defs.MEAN_AUDIO_DURATION_S,
+        defs.MEDIAN_AUDIO_DURATION_S,
+        defs.PERCENTILES_AUDIO_DURATION_S,
+    ),
+    defs.AUDIO_UNDERRUN: (
+        defs.MEAN_AUDIO_UNDERRUN_S,
+        defs.MEDIAN_AUDIO_UNDERRUN_S,
+        defs.PERCENTILES_AUDIO_UNDERRUN_S,
+    ),
+}
+
+
+def _wants_text_ttft(metrics: list[str]) -> bool:
+    return "ttft" in metrics
+
+
+def _wants_text_tpot(metrics: list[str]) -> bool:
+    return "tpot" in metrics or "tpop" in metrics
+
+
+def _wants_text_itl(metrics: list[str]) -> bool:
+    return "itl" in metrics
+
+
+def _wants_stream_ttfc(metrics: list[str]) -> bool:
+    return "ttfc" in metrics
+
+
+def _wants_stream_tpoc(metrics: list[str]) -> bool:
+    return "tpoc" in metrics or "tpop" in metrics
+
+
+def _wants_stream_icl(metrics: list[str]) -> bool:
+    return "icl" in metrics
+
+
+def _stage_modality_flags(
+    final_output_type: str,
+    output_unit_type: str,
+) -> tuple[bool, bool, bool, bool, bool]:
+    is_text_stage = final_output_type == "text" or output_unit_type == "text"
+    is_audio_stage = final_output_type == "audio" or output_unit_type == "audio"
+    is_image_stage = final_output_type in {"image", "images"} or output_unit_type == "image"
+    is_video_stage = final_output_type in {"video", "videos"} or output_unit_type == "video"
+    is_internal_stream_stage = (
+        output_unit_type in _STREAMING_OUTPUT_UNIT_TYPES and not is_text_stage and not is_audio_stage
+    )
+    return (
+        is_text_stage,
+        is_audio_stage,
+        is_image_stage,
+        is_video_stage,
+        is_internal_stream_stage,
+    )
 
 
 def print_metrics(
@@ -43,6 +186,10 @@ def print_metrics(
     benchmark_duration,
     goodput_config_dict,
     metrics: MultiModalsBenchmarkMetrics,
+    *,
+    outputs: list[RequestFuncOutput] | None = None,
+    selected_percentiles: list[float] | None = None,
+    print_stage: bool = False,
 ):
     print("{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="="))
     print("{:<40} {:<10}".format("Successful requests:", metrics.completed))
@@ -61,7 +208,19 @@ def print_metrics(
         process_one_metric("e2el", metrics)
     print_text_metrics(task_type, selected_percentile_metrics, metrics)
     if task_type == TaskType.GENERATION:
-        print_audio_metrics(selected_percentile_metrics, metrics)
+        if _has_audio_output(metrics):
+            print_audio_metrics(selected_percentile_metrics, metrics)
+        if _has_image_output(metrics):
+            print_image_metrics(selected_percentiles or [], metrics)
+        if print_stage and outputs and selected_percentiles is not None:
+            print("\n{s:{c}^{n}}".format(s=" Stage Benchmark Result ", n=50, c="="))
+            for sm in _build_stage_metrics_from_outputs(outputs):
+                print_stage_metrics(
+                    task_type,
+                    selected_percentile_metrics,
+                    selected_percentiles,
+                    sm,
+                )
     print("=" * 50)
 
 
@@ -79,22 +238,36 @@ def print_text_metrics(task_type, selected_percentile_metrics, metrics: MultiMod
         # No text tokens generated (e.g. pure TTS speech endpoint): per-token
         # latency metrics (ttft/tpot/itl) are undefined, so skip them.
         has_text_output = metrics.total_output > 0
-        for metric in selected_percentile_metrics:
-            if metric == "e2el":
-                continue
-            if not metric.startswith("audio") and has_text_output:
-                process_one_metric(metric, metrics)
+        if has_text_output:
+            if _wants_text_ttft(selected_percentile_metrics):
+                process_one_metric("ttft", metrics)
+            if _wants_text_tpot(selected_percentile_metrics):
+                process_one_metric("tpot", metrics)
+            if _wants_text_itl(selected_percentile_metrics):
+                process_one_metric("itl", metrics)
+
+
+def _has_audio_output(metrics: MultiModalsBenchmarkMetrics) -> bool:
+    return bool(
+        getattr(metrics, defs.TOTAL_AUDIO_DURATION_S, 0.0) > 0 or getattr(metrics, defs.TOTAL_AUDIO_FRAMES, 0) > 0
+    )
+
+
+def _has_image_output(metrics: MultiModalsBenchmarkMetrics) -> bool:
+    return int(getattr(metrics, defs.TOTAL_IMAGES, 0) or 0) > 0
 
 
 def print_audio_metrics(selected_percentile_metrics, metrics: MultiModalsBenchmarkMetrics):
     print("{s:{c}^{n}}".format(s=" Audio Result ", n=50, c="="))
-    print("{:<40} {:<10.2f}".format("Total audio duration generated(s):", metrics.total_audio_duration_s))
-    print("{:<40} {:<10}".format("Total audio frames generated:", metrics.total_audio_frames))
-    print("{:<40} {:<10.2f}".format("Audio throughput(audio duration/s):", metrics.audio_throughput))
+    print(
+        "{:<40} {:<10.2f}".format("Total audio duration generated(s):", getattr(metrics, defs.TOTAL_AUDIO_DURATION_S))
+    )
+    print("{:<40} {:<10}".format("Total audio frames generated:", getattr(metrics, defs.TOTAL_AUDIO_FRAMES)))
+    print("{:<40} {:<10.2f}".format("Audio throughput(audio duration/s):", getattr(metrics, defs.AUDIO_THROUGHPUT)))
     print(
         "{:<40} {:<10.2%}".format(
             "Streaming continuity OK rate:",
-            metrics.audio_continuity_ok_rate,
+            getattr(metrics, defs.AUDIO_CONTINUITY_OK_RATE),
         )
     )
     for metric in selected_percentile_metrics:
@@ -102,6 +275,37 @@ def print_audio_metrics(selected_percentile_metrics, metrics: MultiModalsBenchma
             process_one_metric(metric, metrics)
 
 
+def print_image_metrics(selected_percentiles: list[float], metrics: MultiModalsBenchmarkMetrics):
+    print("{s:{c}^{n}}".format(s=" Image Result ", n=50, c="="))
+    print("{:<40} {:<10}".format("Total images generated:", getattr(metrics, defs.TOTAL_IMAGES)))
+    print("{:<40} {:<10.2f}".format("Image throughput (img/s):", getattr(metrics, defs.IMAGE_THROUGHPUT)))
+    print(
+        "{:<40} {:<10.2f}".format(
+            "Average pixels per image:",
+            getattr(metrics, defs.AVERAGE_PIXELS_PER_IMAGE),
+        )
+    )
+    if getattr(metrics, defs.MEAN_DENOISE_STEP_LATENCY_MS) > 0:
+        print(
+            "{:<40} {:<10.2f}".format(
+                "Mean denoise step latency (ms):",
+                getattr(metrics, defs.MEAN_DENOISE_STEP_LATENCY_MS),
+            )
+        )
+    if getattr(metrics, defs.MEAN_IMAGE_GENERATION_MS) > 0:
+        print("-----------------Image Generation-----------------")
+        print("{:<40} {:<10.2f}".format("Mean IMAGE_GENERATION (ms):", getattr(metrics, defs.MEAN_IMAGE_GENERATION_MS)))
+        print(
+            "{:<40} {:<10.2f}".format(
+                "Median IMAGE_GENERATION (ms):",
+                getattr(metrics, defs.MEDIAN_IMAGE_GENERATION_MS),
+            )
+        )
+        for p, value in getattr(metrics, defs.PERCENTILES_IMAGE_GENERATION_MS) or []:
+            if p in selected_percentiles:
+                print("{:<40} {:<10.2f}".format(f"P{_p_label(p)} IMAGE_GENERATION (ms):", value))
+
+
 def process_one_metric(
     metric_attribute_name: str,
     metrics: MultiModalsBenchmarkMetrics,
@@ -131,15 +335,16 @@ def process_one_metric(
     elif is_audio_rtf:
         suffix = ""
         unit_suffix = ""
-    mean_attr_name = f"mean_{metric_attribute_name}{suffix}"
+    field_names = _AGGREGATE_PERCENTILE_FIELD_NAMES.get(metric_attribute_name)
+    mean_attr_name = field_names[0] if field_names else f"mean_{metric_attribute_name}{suffix}"
     mean_value = getattr(metrics, mean_attr_name, 0.0)
     print(f"{f'Mean {metric_attribute_name.upper()}{unit_suffix}:':<40} {mean_value:<10.2f}")
 
-    median_attr_name = f"median_{metric_attribute_name}{suffix}"
+    median_attr_name = field_names[1] if field_names else f"median_{metric_attribute_name}{suffix}"
     median_value = getattr(metrics, median_attr_name, 0.0)
     print(f"{f'Median {metric_attribute_name.upper()}{unit_suffix}:':<40} {median_value:<10.2f}")
 
-    percentiles_attr_name = f"percentiles_{metric_attribute_name}{suffix}"
+    percentiles_attr_name = field_names[2] if field_names else f"percentiles_{metric_attribute_name}{suffix}"
     percentiles = getattr(metrics, percentiles_attr_name, [])
 
     for percentile, value in percentiles:
@@ -148,6 +353,356 @@ def process_one_metric(
         print(f"{label:<40} {value:<10.2f}")
 
 
+def _print_percentile_metric(
+    header: str,
+    label: str,
+    values: list[float],
+    selected_percentiles: list[float],
+    *,
+    values_are_ms: bool = False,
+    to_ms: bool = True,
+) -> None:
+    if not values:
+        return
+    values_s = [v / 1000.0 for v in values] if values_are_ms else values
+    mean_v, med_v, pcts = _percentile_rows_seconds(values_s, selected_percentiles, to_ms=to_ms)
+    unit_suffix = " (ms)" if to_ms else ""
+    print("{s:{c}^{n}}".format(s=header, n=50, c="-"))
+    print(f"{f'Mean {label}{unit_suffix}:':<40} {mean_v:<10.2f}")
+    print(f"{f'Median {label}{unit_suffix}:':<40} {med_v:<10.2f}")
+    for p, v in pcts:
+        p_str = _p_label(p)
+        print(f"{f'P{p_str} {label}{unit_suffix}:':<40} {v:<10.2f}")
+
+
+def _print_stage_timing(sm: StageBenchmarkMetrics, selected_percentiles: list[float]) -> None:
+    _print_percentile_metric(
+        "Stage Timing",
+        "stage_gen_time",
+        getattr(sm, defs.STAGE_GEN_TIMES_MS),
+        selected_percentiles,
+        values_are_ms=True,
+        to_ms=True,
+    )
+    postprocess_times_ms = getattr(sm, defs.POSTPROCESS_TIMES_MS)
+    if postprocess_times_ms and any(v > 0 for v in postprocess_times_ms):
+        _print_percentile_metric(
+            "Stage Postprocess",
+            "postprocess",
+            postprocess_times_ms,
+            selected_percentiles,
+            values_are_ms=True,
+            to_ms=True,
+        )
+
+
+def _print_text_stage_metrics(
+    task_type: TaskType,
+    selected_percentile_metrics: list[str],
+    selected_percentiles: list[float],
+    sm: StageBenchmarkMetrics,
+) -> None:
+    print("{s:{c}^{n}}".format(s=" Text Result ", n=50, c="="))
+    print("{:<40} {:<10}".format("Stage generated tokens:", getattr(sm, defs.TOTAL_OUTPUT)))
+
+    if task_type != TaskType.GENERATION:
+        return
+    if _wants_text_ttft(selected_percentile_metrics):
+        _print_percentile_metric(
+            "Serving Time to First Token",
+            "Serving TTFT",
+            getattr(sm, defs.TTFTS),
+            selected_percentiles,
+        )
+        _print_percentile_metric("Time to First Token", "TTFT", getattr(sm, defs.VLLM_TTFTS), selected_percentiles)
+    if _wants_text_tpot(selected_percentile_metrics):
+        _print_percentile_metric(
+            "Time per Output Token (excl. 1st token)",
+            "TPOT",
+            getattr(sm, defs.VLLM_TPOTS),
+            selected_percentiles,
+        )
+    if _wants_text_itl(selected_percentile_metrics):
+        _print_percentile_metric("Inter-token Latency", "ITL", getattr(sm, defs.VLLM_ITLS), selected_percentiles)
+
+
+def _print_audio_stage_metrics(
+    selected_percentile_metrics: list[str],
+    selected_percentiles: list[float],
+    sm: StageBenchmarkMetrics,
+) -> None:
+    total_duration = float(sum(getattr(sm, defs.AUDIO_DURATIONS)))
+    total_frames = int(sum(getattr(sm, defs.AUDIO_FRAMES)))
+    print("{s:{c}^{n}}".format(s=" Audio Result ", n=50, c="="))
+    print("{:<40} {:<10.2f}".format("Stage audio duration generated(s):", total_duration))
+    print("{:<40} {:<10}".format("Stage audio frames generated:", total_frames))
+    if defs.AUDIO_TTFP in selected_percentile_metrics:
+        _print_percentile_metric(
+            "Serving Time to First Packet",
+            "Serving AUDIO_TTFP",
+            getattr(sm, defs.AUDIO_TTFPS),
+            selected_percentiles,
+        )
+    if defs.AUDIO_RTF in selected_percentile_metrics:
+        audio_rtfs = getattr(sm, defs.AUDIO_RTFS)
+        if audio_rtfs:
+            _print_percentile_metric("Real Time Factor", "AUDIO_RTF", audio_rtfs, selected_percentiles, to_ms=False)
+        elif getattr(sm, defs.MISSING_AUDIO_RTF_COUNT) > 0:
+            print("{:<40} {:<10}".format("AUDIO_RTF skipped:", "missing stage-local audio duration/sample rate"))
+    if defs.AUDIO_DURATION in selected_percentile_metrics:
+        audio_durations = getattr(sm, defs.AUDIO_DURATIONS)
+        if audio_durations:
+            _print_percentile_metric(
+                "Audio Duration",
+                "AUDIO_DURATION",
+                audio_durations,
+                selected_percentiles,
+                to_ms=False,
+            )
+        elif getattr(sm, defs.MISSING_AUDIO_DURATION_COUNT) > 0:
+            print("{:<40} {:<10}".format("AUDIO_DURATION skipped:", "missing stage-local audio duration/sample rate"))
+
+
+def _print_image_stage_metrics(
+    selected_percentiles: list[float],
+    sm: StageBenchmarkMetrics,
+) -> None:
+    image_num = getattr(sm, defs.OUTPUT_UNIT_COUNT) if getattr(sm, "output_unit_type") == "image" else 0
+    print("{s:{c}^{n}}".format(s=" Image Result ", n=50, c="="))
+    print("{:<40} {:<10}".format("Total images generated:", image_num))
+    _print_percentile_metric(
+        "Image Generation",
+        "IMAGE_GENERATION",
+        getattr(sm, defs.STAGE_GEN_TIMES_MS),
+        selected_percentiles,
+        values_are_ms=True,
+        to_ms=True,
+    )
+
+
+def _print_internal_stream_stage_metrics(
+    selected_percentile_metrics: list[str],
+    selected_percentiles: list[float],
+    sm: StageBenchmarkMetrics,
+) -> None:
+    print("{s:{c}^{n}}".format(s=" Internal Stream Result ", n=50, c="="))
+    serving_time_to_first_outputs_ms = getattr(sm, defs.SERVING_TIME_TO_FIRST_OUTPUTS_MS)
+    if _wants_stream_ttfc(selected_percentile_metrics) and serving_time_to_first_outputs_ms:
+        _print_percentile_metric(
+            "Serving Time to First Chunk",
+            "Serving TTFC",
+            serving_time_to_first_outputs_ms,
+            selected_percentiles,
+            values_are_ms=True,
+            to_ms=True,
+        )
+    time_per_output_units_ms = getattr(sm, defs.TIME_PER_OUTPUT_UNITS_MS)
+    if _wants_stream_tpoc(selected_percentile_metrics) and time_per_output_units_ms:
+        _print_percentile_metric(
+            "Time per Output Chunk (excl. 1st chunk)",
+            "TPOP",
+            time_per_output_units_ms,
+            selected_percentiles,
+            values_are_ms=True,
+            to_ms=True,
+        )
+    inter_output_latencies_ms = getattr(sm, defs.INTER_OUTPUT_LATENCIES_MS)
+    if _wants_stream_icl(selected_percentile_metrics) and inter_output_latencies_ms:
+        _print_percentile_metric(
+            "Inter-chunk Latency",
+            "ICL",
+            inter_output_latencies_ms,
+            selected_percentiles,
+            values_are_ms=True,
+            to_ms=True,
+        )
+
+
+def print_stage_metrics(
+    task_type: TaskType,
+    selected_percentile_metrics: list[str],
+    selected_percentiles: list[float],
+    sm: StageBenchmarkMetrics,
+):
+    title = f" Stage {getattr(sm, 'stage_id')} ({getattr(sm, 'stage_name')}) "
+    (
+        is_text_stage,
+        is_audio_stage,
+        is_image_stage,
+        is_video_stage,
+        is_internal_stream_stage,
+    ) = _stage_modality_flags(getattr(sm, "final_output_type"), getattr(sm, "output_unit_type"))
+
+    print("{s:{c}^{n}}".format(s=title, n=50, c="="))
+    if is_image_stage:
+        _print_image_stage_metrics(selected_percentiles, sm)
+        return
+
+    if is_video_stage:
+        return
+
+    _print_stage_timing(sm, selected_percentiles)
+    if is_text_stage:
+        _print_text_stage_metrics(task_type, selected_percentile_metrics, selected_percentiles, sm)
+    elif is_audio_stage:
+        _print_audio_stage_metrics(
+            selected_percentile_metrics,
+            selected_percentiles,
+            sm,
+        )
+    elif is_internal_stream_stage:
+        _print_internal_stream_stage_metrics(selected_percentile_metrics, selected_percentiles, sm)
+
+
+def _build_stage_metrics_from_outputs(
+    outputs: list[RequestFuncOutput],
+) -> list[StageBenchmarkMetrics]:
+    """Aggregate per ``stage_id`` using ``stage_metrics`` snapshots from the client."""
+    buckets: dict[str, list[tuple[RequestFuncOutput, dict]]] = defaultdict(list)
+    for out in outputs:
+        if not getattr(out, "success", False):
+            continue
+        smap = getattr(out, "stage_metrics", None) or {}
+        if not isinstance(smap, dict) or not smap:
+            continue
+        for sid, info in smap.items():
+            if isinstance(info, dict):
+                buckets[str(sid)].append((out, info))
+
+    result: list[StageBenchmarkMetrics] = []
+    for sid in sorted(buckets.keys(), key=lambda x: int(x)):
+        rows = buckets[sid]
+        stage_id_int = int(sid)
+        stage_name = str((rows[0][1] or {}).get("stage_name") or f"stage_{stage_id_int}")
+        final_output_type = str((rows[0][1] or {}).get("final_output_type") or "unknown")
+        output_unit_types = [
+            str((info or {}).get("output_unit_type") or "") for _, info in rows if (info or {}).get("output_unit_type")
+        ]
+        output_unit_type = output_unit_types[0] if output_unit_types else "other"
+
+        is_text_stage, is_audio_stage, _, _, is_internal_stream_stage = _stage_modality_flags(
+            final_output_type, output_unit_type
+        )
+        total_output = 0
+
+        ttfts: list[float] = []
+        tpots: list[float] = []
+        itls: list[float] = []
+        vllm_ttfts: list[float] = []
+        vllm_tpots: list[float] = []
+        vllm_itls: list[float] = []
+        stage_gen_times_ms = [float((info or {}).get(defs.STAGE_GEN_TIME_MS) or 0.0) for _, info in rows]
+        postprocess_times_ms = [float((info or {}).get(defs.POSTPROCESS_TIME_MS) or 0.0) for _, info in rows]
+        output_unit_count = sum(int((info or {}).get(defs.OUTPUT_UNIT_COUNT) or 0) for _, info in rows)
+        inter_output_latencies_ms: list[float] = []
+        serving_time_to_first_outputs_ms: list[float] = []
+        time_per_output_units_ms: list[float] = []
+        if is_text_stage or is_internal_stream_stage:
+            serving_time_to_first_outputs_ms = [
+                float((info or {}).get(defs.SERVING_TIME_TO_FIRST_OUTPUT_MS) or 0.0)
+                for _, info in rows
+                if (info or {}).get(defs.SERVING_TIME_TO_FIRST_OUTPUT_MS) is not None
+            ]
+            time_per_output_units_ms = [
+                float((info or {}).get(defs.TIME_PER_OUTPUT_UNIT_MS) or 0.0)
+                for _, info in rows
+                if int((info or {}).get(defs.OUTPUT_UNIT_COUNT) or 0) > 1
+                and (info or {}).get(defs.TIME_PER_OUTPUT_UNIT_MS) is not None
+            ]
+            for _, info in rows:
+                values = (info or {}).get(defs.INTER_OUTPUT_LATENCIES_MS)
+                if isinstance(values, list):
+                    inter_output_latencies_ms.extend(float(v or 0.0) for v in values)
+                elif (info or {}).get(defs.INTER_OUTPUT_LATENCY_MS) is not None:
+                    inter_output_latencies_ms.append(float((info or {}).get(defs.INTER_OUTPUT_LATENCY_MS) or 0.0))
+
+        if is_text_stage:
+            for _, info in rows:
+                ttft_ms = float((info or {}).get(defs.SERVING_TIME_TO_FIRST_OUTPUT_MS) or 0.0)
+                if ttft_ms > 0:
+                    ttfts.append(ttft_ms / 1000.0)
+                tpot_ms = float((info or {}).get(defs.TIME_PER_OUTPUT_UNIT_MS) or 0.0)
+                if int((info or {}).get(defs.OUTPUT_UNIT_COUNT) or 0) > 1 and tpot_ms > 0:
+                    tpots.append(tpot_ms / 1000.0)
+                values = (info or {}).get(defs.INTER_OUTPUT_LATENCIES_MS)
+                if isinstance(values, list):
+                    itls.extend(float(v or 0.0) / 1000.0 for v in values)
+                elif (info or {}).get(defs.INTER_OUTPUT_LATENCY_MS) is not None:
+                    itls.append(float((info or {}).get(defs.INTER_OUTPUT_LATENCY_MS) or 0.0) / 1000.0)
+                vllm_ttft_ms = float((info or {}).get(defs.VLLM_TTFT_MS) or 0.0)
+                if vllm_ttft_ms > 0:
+                    vllm_ttfts.append(vllm_ttft_ms / 1000.0)
+                vllm_tpot_ms = float((info or {}).get(defs.VLLM_TPOT_MS) or 0.0)
+                if vllm_tpot_ms > 0:
+                    vllm_tpots.append(vllm_tpot_ms / 1000.0)
+                vllm_values = (info or {}).get(defs.VLLM_ITLS_MS)
+                if isinstance(vllm_values, list):
+                    vllm_itls.extend(float(v or 0.0) / 1000.0 for v in vllm_values)
+                elif (info or {}).get(defs.VLLM_ITL_MS) is not None:
+                    vllm_itls.append(float((info or {}).get(defs.VLLM_ITL_MS) or 0.0) / 1000.0)
+            total_output = sum(int((info or {}).get(defs.NUM_TOKENS_OUT) or 0) for _, info in rows)
+
+        audio_ttfps: list[float] = []
+        audio_rtfs: list[float] = []
+        audio_durations: list[float] = []
+        audio_frames: list[int] = []
+        missing_audio_duration_count = 0
+        missing_audio_rtf_count = 0
+        if is_audio_stage:
+            for _, info in rows:
+                audio_ttfp_ms = float((info or {}).get(defs.SERVING_TIME_TO_FIRST_OUTPUT_MS) or 0.0)
+                if audio_ttfp_ms > 0:
+                    audio_ttfps.append(audio_ttfp_ms / 1000.0)
+
+                frame_count = int((info or {}).get(defs.AUDIO_FRAMES) or 0)
+                if frame_count <= 0 and (info or {}).get("output_unit_type") == "audio":
+                    frame_count = int((info or {}).get(defs.OUTPUT_UNIT_COUNT) or 0)
+                audio_frames.append(frame_count)
+
+                duration_s = float((info or {}).get(f"{defs.AUDIO_DURATION}_s") or 0.0)
+                if duration_s > 0:
+                    audio_durations.append(duration_s)
+                else:
+                    missing_audio_duration_count += 1
+
+                rtf = float((info or {}).get(defs.AUDIO_RTF) or 0.0)
+                if rtf > 0:
+                    audio_rtfs.append(rtf)
+                else:
+                    missing_audio_rtf_count += 1
+
+        result.append(
+            StageBenchmarkMetrics(
+                **{
+                    "stage_id": stage_id_int,
+                    "stage_name": stage_name,
+                    "final_output_type": final_output_type,
+                    defs.TOTAL_OUTPUT: total_output,
+                    defs.TTFTS: ttfts,
+                    defs.TPOTS: tpots,
+                    defs.ITLS: itls,
+                    defs.VLLM_TTFTS: vllm_ttfts,
+                    defs.VLLM_TPOTS: vllm_tpots,
+                    defs.VLLM_ITLS: vllm_itls,
+                    defs.AUDIO_TTFPS: audio_ttfps,
+                    defs.AUDIO_RTFS: audio_rtfs,
+                    defs.AUDIO_DURATIONS: audio_durations,
+                    defs.AUDIO_FRAMES: audio_frames,
+                    defs.MISSING_AUDIO_DURATION_COUNT: missing_audio_duration_count,
+                    defs.MISSING_AUDIO_RTF_COUNT: missing_audio_rtf_count,
+                    defs.STAGE_GEN_TIMES_MS: stage_gen_times_ms,
+                    defs.POSTPROCESS_TIMES_MS: postprocess_times_ms,
+                    "output_unit_type": output_unit_type,
+                    defs.OUTPUT_UNIT_COUNT: output_unit_count,
+                    defs.SERVING_TIME_TO_FIRST_OUTPUTS_MS: serving_time_to_first_outputs_ms,
+                    defs.TIME_PER_OUTPUT_UNITS_MS: time_per_output_units_ms,
+                    defs.INTER_OUTPUT_LATENCIES_MS: inter_output_latencies_ms,
+                }
+            )
+        )
+    return result
+
+
 def calculate_metrics(
     input_requests: list[SampleRequest],
     outputs: list[RequestFuncOutput],
@@ -160,6 +715,7 @@ def calculate_metrics(
     max_concurrency,
     request_rate,
     benchmark_duration,
+    print_stage: bool = False,
 ) -> tuple[BenchmarkMetrics, list[int]]:
     """Calculate the metrics for the benchmark.
 
@@ -187,6 +743,10 @@ def calculate_metrics(
     audio_rtfs: list[float] = []
     audio_duration: list[float] = []
     audio_frames: list[int] = []
+    image_generation_times_ms: list[float] = []
+    denoise_step_latencies_ms: list[float] = []
+    total_images = 0
+    total_image_pixels = 0
     audio_underruns: list[float] = []
     audio_continuity_ok: list[bool] = []
     input_audio_duration = 0.0
@@ -222,6 +782,15 @@ def calculate_metrics(
             audio_rtfs.append(getattr(outputs[i], defs.AUDIO_RTF, 0.0))
             audio_duration.append(getattr(outputs[i], defs.AUDIO_DURATION, 0.0))
             audio_frames.append(getattr(outputs[i], defs.AUDIO_FRAMES, 0.0))
+            image_count = int(getattr(outputs[i], defs.IMAGE_COUNT, 0) or 0)
+            total_images += image_count
+            image_generation_time_ms = float(getattr(outputs[i], defs.IMAGE_GENERATION_TIME_MS, 0.0) or 0.0)
+            if image_generation_time_ms > 0:
+                image_generation_times_ms.append(image_generation_time_ms)
+            total_image_pixels += int(getattr(outputs[i], defs.IMAGE_PIXELS, 0) or 0)
+            denoise_step_latency_ms = float(getattr(outputs[i], defs.DENOISE_STEP_LATENCY_MS, 0.0) or 0.0)
+            if denoise_step_latency_ms > 0:
+                denoise_step_latencies_ms.append(denoise_step_latency_ms)
             audio_underruns.append(getattr(outputs[i], f"{defs.AUDIO_UNDERRUN}_s", 0.0))
             audio_continuity_ok.append(bool(getattr(outputs[i], defs.AUDIO_CONTINUITY_OK, True)))
             e2els.append(outputs[i].latency)
@@ -336,26 +905,6 @@ def calculate_metrics(
         std_ttft_ms=np.std(ttfts or 0) * 1000,
         median_ttft_ms=np.median(ttfts or 0) * 1000,
         percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles],
-        mean_audio_ttfp_ms=np.mean(audio_ttfps or 0) * 1000,
-        std_audio_ttfp_ms=np.std(audio_ttfps or 0) * 1000,
-        median_audio_ttfp_ms=np.median(audio_ttfps or 0) * 1000,
-        percentiles_audio_ttfp_ms=[(p, np.percentile(audio_ttfps or 0, p) * 1000) for p in selected_percentiles],
-        mean_audio_duration_s=np.mean(audio_duration or 0),
-        std_audio_duration_s=np.std(audio_duration or 0),
-        median_audio_duration_s=np.median(audio_duration or 0),
-        percentiles_audio_duration_s=[(p, np.percentile(audio_duration or 0, p)) for p in selected_percentiles],
-        total_audio_duration_s=sum(audio_duration),
-        total_audio_frames=sum(audio_frames),
-        audio_throughput=sum(audio_duration) / dur_s,
-        mean_audio_rtf=np.mean(audio_rtfs or 0),
-        std_audio_rtf=np.std(audio_rtfs or 0),
-        median_audio_rtf=np.median(audio_rtfs or 0),
-        percentiles_audio_rtf=[(p, np.percentile(audio_rtfs or 0, p)) for p in selected_percentiles],
-        mean_audio_underrun_s=np.mean(audio_underruns or 0),
-        std_audio_underrun_s=np.std(audio_underruns or 0),
-        median_audio_underrun_s=np.median(audio_underruns or 0),
-        percentiles_audio_underrun_s=[(p, np.percentile(audio_underruns or 0, p)) for p in selected_percentiles],
-        audio_continuity_ok_rate=(sum(audio_continuity_ok) / len(audio_continuity_ok)) if audio_continuity_ok else 1.0,
         mean_tpot_ms=np.mean(tpots or 0) * 1000,
         std_tpot_ms=np.std(tpots or 0) * 1000,
         median_tpot_ms=np.median(tpots or 0) * 1000,
@@ -371,6 +920,46 @@ def calculate_metrics(
         max_output_tokens_per_s=max_output_tokens_per_s,
         max_concurrent_requests=max_concurrent_requests,
         rtfx=input_audio_duration / dur_s,
+        **{
+            defs.MEAN_AUDIO_TTFP_MS: np.mean(audio_ttfps or 0) * 1000,
+            defs.STD_AUDIO_TTFP_MS: np.std(audio_ttfps or 0) * 1000,
+            defs.MEDIAN_AUDIO_TTFP_MS: np.median(audio_ttfps or 0) * 1000,
+            defs.PERCENTILES_AUDIO_TTFP_MS: [
+                (p, np.percentile(audio_ttfps or 0, p) * 1000) for p in selected_percentiles
+            ],
+            defs.MEAN_AUDIO_DURATION_S: np.mean(audio_duration or 0),
+            defs.STD_AUDIO_DURATION_S: np.std(audio_duration or 0),
+            defs.MEDIAN_AUDIO_DURATION_S: np.median(audio_duration or 0),
+            defs.PERCENTILES_AUDIO_DURATION_S: [
+                (p, np.percentile(audio_duration or 0, p)) for p in selected_percentiles
+            ],
+            defs.TOTAL_AUDIO_DURATION_S: sum(audio_duration),
+            defs.TOTAL_AUDIO_FRAMES: sum(audio_frames),
+            defs.AUDIO_THROUGHPUT: sum(audio_duration) / dur_s,
+            defs.MEAN_AUDIO_RTF: np.mean(audio_rtfs or 0),
+            defs.STD_AUDIO_RTF: np.std(audio_rtfs or 0),
+            defs.MEDIAN_AUDIO_RTF: np.median(audio_rtfs or 0),
+            defs.PERCENTILES_AUDIO_RTF: [(p, np.percentile(audio_rtfs or 0, p)) for p in selected_percentiles],
+            defs.TOTAL_IMAGES: total_images,
+            defs.IMAGE_THROUGHPUT: total_images / dur_s,
+            defs.AVERAGE_PIXELS_PER_IMAGE: (total_image_pixels / total_images) if total_images > 0 else 0.0,
+            defs.MEAN_DENOISE_STEP_LATENCY_MS: np.mean(denoise_step_latencies_ms or 0),
+            defs.MEAN_IMAGE_GENERATION_MS: np.mean(image_generation_times_ms or 0),
+            defs.STD_IMAGE_GENERATION_MS: np.std(image_generation_times_ms or 0),
+            defs.MEDIAN_IMAGE_GENERATION_MS: np.median(image_generation_times_ms or 0),
+            defs.PERCENTILES_IMAGE_GENERATION_MS: [
+                (p, np.percentile(image_generation_times_ms or 0, p)) for p in selected_percentiles
+            ],
+            defs.MEAN_AUDIO_UNDERRUN_S: np.mean(audio_underruns or 0),
+            defs.STD_AUDIO_UNDERRUN_S: np.std(audio_underruns or 0),
+            defs.MEDIAN_AUDIO_UNDERRUN_S: np.median(audio_underruns or 0),
+            defs.PERCENTILES_AUDIO_UNDERRUN_S: [
+                (p, np.percentile(audio_underruns or 0, p)) for p in selected_percentiles
+            ],
+            defs.AUDIO_CONTINUITY_OK_RATE: (
+                (sum(audio_continuity_ok) / len(audio_continuity_ok)) if audio_continuity_ok else 1.0
+            ),
+        },
     )
     print_metrics(
         task_type,
@@ -380,5 +969,8 @@ def calculate_metrics(
         benchmark_duration,
         goodput_config_dict,
         metrics,
+        outputs=outputs,
+        selected_percentiles=selected_percentiles,
+        print_stage=print_stage,
     )
     return metrics, actual_output_lens

vllm_omni/benchmarks/patch/patch.py

@@ -2,6 +2,7 @@ import asyncio
 import contextlib
 import io
 import json
+import mimetypes
 import os
 import random
 import ssl
@@ -12,7 +13,7 @@ import wave
 from collections.abc import Iterable
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Literal
+from typing import Any, Literal
 
 import aiohttp
 import pybase64 as base64
@@ -48,9 +49,41 @@ from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
 )
 from vllm_omni.benchmarks.data_modules.sound_effect_dataset import SoundEffectDataset
 from vllm_omni.benchmarks.data_modules.ttsd_dataset import TTSDDataset
+from vllm_omni.metrics import definitions as defs
 
 _AUDIO_CONTINUITY_THRESHOLD_ENV = "VLLM_OMNI_BENCH_AUDIO_CONTINUITY_THRESHOLD_S"
-_DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S = 0.1
+RETURN_STAGE_METRICS_FIELD = "return_stage_metrics"
+_IMAGE_STAGE_METRICS_BACKENDS = frozenset({"openai-image-edits-omni"})
+_PRINT_STAGE = False
+
+
+def maybe_enable_stage_metrics(extra_body: dict[str, Any] | None, *, enabled: bool) -> dict[str, Any] | None:
+    """Return extra_body with stage-metric opt-in when benchmark metrics need it."""
+    if not enabled:
+        return extra_body
+    body = dict(extra_body or {})
+    body.setdefault(RETURN_STAGE_METRICS_FIELD, True)
+    return body
+
+
+def should_request_stage_metrics(args: Any) -> bool:
+    """Whether this benchmark run needs server-side stage metrics in responses."""
+    if getattr(args, "print_stage", False):
+        return True
+
+    backend = getattr(args, "backend", None)
+    if backend in _IMAGE_STAGE_METRICS_BACKENDS:
+        return True
+
+    extra_body = getattr(args, "extra_body", None) or {}
+    modalities = extra_body.get("modalities") if isinstance(extra_body, dict) else None
+    return backend == "openai-chat-omni" and "image" in (modalities or [])
+
+
+def set_print_stage(enabled: bool) -> None:
+    """Set whether this benchmark run prints the stage benchmark section."""
+    global _PRINT_STAGE
+    _PRINT_STAGE = bool(enabled)
 
 
 def _audio_continuity_threshold_s() -> float:
@@ -62,7 +95,7 @@ def _audio_continuity_threshold_s() -> float:
     """
     raw = os.environ.get(_AUDIO_CONTINUITY_THRESHOLD_ENV)
     if not raw:
-        return _DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S
+        return defs.AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S
     try:
         value = float(raw)
     except ValueError:
@@ -70,9 +103,9 @@ def _audio_continuity_threshold_s() -> float:
             "Invalid %s=%r; using default %.3fs",
             _AUDIO_CONTINUITY_THRESHOLD_ENV,
             raw,
-            _DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S,
+            defs.AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S,
         )
-        return _DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S
+        return defs.AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S
     return max(value, 0.0)
 
 
@@ -352,6 +385,10 @@ class MixRequestFuncOutput(RequestFuncOutput):
     audio_duration: float = 0.0
     audio_frames: int = 0
     audio_rtf: float = 0.0
+    image_count: int = 0
+    image_generation_time_ms: float = 0.0
+    image_pixels: int = 0
+    denoise_step_latency_ms: float = 0.0
     text_latency: float = 0.0
     #: Worst-case streaming-audio underrun (wall-clock seconds the player
     #: would have been starved). Populated by the audio-speech backend; ``0.0``
@@ -366,6 +403,200 @@ class MixRequestFuncOutput(RequestFuncOutput):
     #: Raw PCM s16le mono at 24 kHz for Seed-TTS WER: from ``/v1/audio/speech`` stream or
     #: resampled export after ``openai-chat-omni`` audio deltas.
     tts_output_pcm_bytes: bytes | None = None
+    #: Per-stage snapshot from orchestrator ``metrics["stage_metrics"]`` (merged across SSE chunks).
+    stage_metrics: dict[str, dict] | None = None
+    stage_id: int | None = None
+    final_output_type: str | None = None
+
+
+_IMAGE_EDITS_EXTRA_BODY_FORM_FIELDS = (
+    "negative_prompt",
+    "num_inference_steps",
+    "guidance_scale",
+    "strength",
+    "true_cfg_scale",
+    "seed",
+    "generator_device",
+    "lora",
+    "layers",
+    "resolution",
+    "bot_task",
+    "sys_type",
+    "system_prompt",
+    RETURN_STAGE_METRICS_FIELD,
+)
+
+
+def _guess_mime_type(path: str) -> str:
+    mime, _ = mimetypes.guess_type(path)
+    return mime or "application/octet-stream"
+
+
+def _iter_image_edit_inputs(value: Any) -> Iterable[Any]:
+    """Yield image references from benchmark multimodal content."""
+    if value is None:
+        return
+    if isinstance(value, list):
+        for item in value:
+            yield from _iter_image_edit_inputs(item)
+        return
+    if not isinstance(value, dict):
+        yield value
+        return
+
+    content_type = value.get("type")
+    if content_type == "image_url":
+        image_url = value.get("image_url")
+        if isinstance(image_url, dict):
+            url = image_url.get("url")
+            if url:
+                yield url
+        elif image_url:
+            yield image_url
+        return
+
+    for key in ("image", "images"):
+        if key in value:
+            yield from _iter_image_edit_inputs(value[key])
+
+
+def _add_image_edit_input_to_form(form: aiohttp.FormData, image_input: Any) -> None:
+    if isinstance(image_input, dict) and "bytes" in image_input:
+        form.add_field(
+            "image",
+            image_input["bytes"],
+            filename="benchmark.png",
+            content_type="image/png",
+        )
+        return
+
+    if isinstance(image_input, str):
+        if image_input.startswith(("data:image", "http://", "https://")):
+            form.add_field("url", image_input)
+            return
+        local_path = image_input.removeprefix("file://")
+        if os.path.exists(local_path):
+            with open(local_path, "rb") as f:
+                image_bytes = f.read()
+            form.add_field(
+                "image",
+                image_bytes,
+                filename=os.path.basename(local_path),
+                content_type=_guess_mime_type(local_path),
+            )
+            return
+
+    raise ValueError(f"Unsupported image edit input: {type(image_input).__name__}")
+
+
+def _add_image_edit_extra_body_to_form(form: aiohttp.FormData, extra_body: dict[str, Any]) -> None:
+    for key in _IMAGE_EDITS_EXTRA_BODY_FORM_FIELDS:
+        value = extra_body.get(key)
+        if value is None:
+            continue
+        if isinstance(value, (dict, list)):
+            form.add_field(key, json.dumps(value))
+        else:
+            form.add_field(key, str(value))
+
+
+def _coerce_positive_int(value: Any) -> int | None:
+    try:
+        coerced = int(value)
+    except (TypeError, ValueError):
+        return None
+    return coerced if coerced > 0 else None
+
+
+def _extract_output_tokens_from_metrics(metrics: dict[str, Any]) -> int | None:
+    top_level_tokens = _coerce_positive_int(metrics.get(defs.NUM_TOKENS_OUT))
+    if top_level_tokens is not None:
+        return top_level_tokens
+
+    stage_snapshot = metrics.get("stage_metrics")
+    if not isinstance(stage_snapshot, dict):
+        return None
+
+    fallback_tokens: list[int] = []
+    for info in stage_snapshot.values():
+        if not isinstance(info, dict):
+            continue
+        num_tokens_out = _coerce_positive_int(info.get(defs.NUM_TOKENS_OUT))
+        if num_tokens_out is None:
+            continue
+        if info.get("final_output_type") == "text" or info.get("output_unit_type") == "token":
+            return num_tokens_out
+        fallback_tokens.append(num_tokens_out)
+    return max(fallback_tokens, default=None)
+
+
+def _update_output_stage_metrics_from_payload(
+    output: MixRequestFuncOutput,
+    data: dict[str, Any],
+    *,
+    update_output_tokens: bool = True,
+) -> None:
+    metrics = data.get("metrics")
+    if not isinstance(metrics, dict):
+        return
+    if update_output_tokens:
+        if (num_tokens_out := _extract_output_tokens_from_metrics(metrics)) is not None:
+            output.output_tokens = num_tokens_out
+    if isinstance(sid := metrics.get("stage_id"), int):
+        output.stage_id = sid
+    if isinstance(final_output_type := metrics.get("final_output_type"), str):
+        output.final_output_type = final_output_type
+    stage_snapshot = metrics.get("stage_metrics")
+    if isinstance(stage_snapshot, dict):
+        if output.stage_metrics is None:
+            output.stage_metrics = {}
+        output.stage_metrics.update(stage_snapshot)
+
+
+def _image_metrics_from_stage_metrics(metrics: dict[str, Any] | None) -> tuple[int, float, int, float]:
+    if not isinstance(metrics, dict):
+        return 0, 0.0, 0, 0.0
+    stage_snapshot = metrics.get("stage_metrics")
+    if not isinstance(stage_snapshot, dict):
+        return 0, 0.0, 0, 0.0
+    image_count = 0
+    image_generation_ms = 0.0
+    image_pixels = 0
+    denoise_step_latency_ms = 0.0
+    for info in stage_snapshot.values():
+        if not isinstance(info, dict):
+            continue
+        final_output_type = info.get("final_output_type")
+        output_unit_type = info.get("output_unit_type")
+        if final_output_type not in {"image", "images"} and output_unit_type != "image":
+            continue
+        image_count += int(info.get(defs.OUTPUT_UNIT_COUNT) or 0)
+        image_generation_ms += float(info.get(defs.STAGE_GEN_TIME_MS) or 0.0)
+        image_pixels += int(info.get(defs.IMAGE_PIXELS) or 0)
+        denoise_step_latency_ms = max(
+            denoise_step_latency_ms,
+            float(info.get(defs.DENOISE_STEP_LATENCY_MS) or 0.0),
+        )
+    return image_count, image_generation_ms, image_pixels, denoise_step_latency_ms
+
+
+def _image_generation_ms_from_content(content: Any) -> float:
+    if not isinstance(content, list):
+        return 0.0
+    for item in content:
+        if not isinstance(item, dict):
+            continue
+        stage_durations = item.get("stage_durations")
+        if not isinstance(stage_durations, dict):
+            continue
+        gen_values = [
+            float(value)
+            for key, value in stage_durations.items()
+            if str(key).endswith("_gen_ms") and isinstance(value, (int, float))
+        ]
+        if gen_values:
+            return max(gen_values)
+    return 0.0
 
 
 async def async_request_openai_chat_omni_completions(
@@ -448,6 +679,13 @@ async def async_request_openai_chat_omni_completions(
         output.output_tokens = 0
         output.error = ""
         output.success = False
+        output.stage_metrics = {}
+        output.stage_id = None
+        output.final_output_type = None
+        output.image_count = 0
+        output.image_generation_time_ms = 0.0
+        output.image_pixels = 0
+        output.denoise_step_latency_ms = 0.0
         try:
             async with session.post(url=api_url, json=payload, headers=headers) as response:
                 if response.status == 200:
@@ -519,9 +757,27 @@ async def async_request_openai_chat_omni_completions(
                                                     logger.warning("Failed to parse wav audio chunk: %s", ex)
                                             else:
                                                 audio_bytes_buffer.extend(audio_bytes)
+                                    elif modality == "image":
+                                        output.image_count += 1
+                                        content_image_ms = _image_generation_ms_from_content(content)
+                                        if content_image_ms > 0:
+                                            output.image_generation_time_ms += content_image_ms
 
-                                if metrics := data.get("metrics"):
-                                    output.output_tokens = metrics.get("num_tokens_out", 0)
+                                _update_output_stage_metrics_from_payload(output, data)
+                                (
+                                    metrics_image_count,
+                                    metrics_image_ms,
+                                    metrics_image_pixels,
+                                    metrics_denoise_step_ms,
+                                ) = _image_metrics_from_stage_metrics(data.get("metrics"))
+                                if metrics_image_count > output.image_count:
+                                    output.image_count = metrics_image_count
+                                if metrics_image_ms > output.image_generation_time_ms:
+                                    output.image_generation_time_ms = metrics_image_ms
+                                if metrics_image_pixels > output.image_pixels:
+                                    output.image_pixels = metrics_image_pixels
+                                if metrics_denoise_step_ms > output.denoise_step_latency_ms:
+                                    output.denoise_step_latency_ms = metrics_denoise_step_ms
 
                                 if usage := data.get("usage"):
                                     if (pt := usage.get("prompt_tokens")) is not None:
@@ -610,6 +866,138 @@ async def async_request_openai_chat_omni_completions(
     return output
 
 
+async def async_request_openai_image_edits_omni(
+    request_func_input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> MixRequestFuncOutput:
+    """Streaming request to /v1/images/edits for multi-stage image-edit benchmarks."""
+    api_url = request_func_input.api_url
+    _validate_api_url(api_url, "OpenAI Image Edits API", "images/edits")
+
+    extra_body = dict(request_func_input.extra_body or {})
+    model = request_func_input.model_name if request_func_input.model_name else request_func_input.model
+    output = MixRequestFuncOutput()
+    output.prompt_len = request_func_input.prompt_len
+    output.itl = []
+    output.stage_metrics = {}
+    output.output_tokens = 0
+    output.image_count = 0
+    output.image_generation_time_ms = 0.0
+    output.image_pixels = 0
+    output.denoise_step_latency_ms = 0.0
+
+    form = aiohttp.FormData()
+    form.add_field("model", model)
+    form.add_field("prompt", request_func_input.prompt)
+    form.add_field("response_format", "b64_json")
+    form.add_field("output_format", str(extra_body.get("output_format", "png")))
+    form.add_field("stream", "true")
+
+    size = extra_body.get("size")
+    if size is None:
+        width, height = extra_body.get("width"), extra_body.get("height")
+        size = f"{width}x{height}" if width is not None and height is not None else "auto"
+    form.add_field("size", str(size))
+
+    _add_image_edit_extra_body_to_form(form, extra_body)
+
+    try:
+        image_inputs = list(_iter_image_edit_inputs(request_func_input.multi_modal_content))
+        if not image_inputs:
+            raise ValueError(
+                "openai-image-edits-omni requires image multimodal content. "
+                "For synthetic inputs, use --dataset-name random-mm with an image bucket."
+            )
+        for image_input in image_inputs:
+            _add_image_edit_input_to_form(form, image_input)
+    except Exception:
+        output.success = False
+        output.error = traceback.format_exc()
+        if pbar:
+            pbar.update(1)
+        return output
+
+    headers = {
+        "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
+    }
+    _update_headers_common(headers, request_func_input)
+
+    st = time.perf_counter()
+    output.start_time = st
+    timestamp = st
+    most_recent_text_timestamp = st
+    generated_text = ""
+    try:
+        async with session.post(url=api_url, data=form, headers=headers) as response:
+            if response.status == 200:
+                handler = StreamedResponseHandler()
+                async for chunk_bytes in response.content.iter_any():
+                    if not chunk_bytes:
+                        continue
+                    for message in handler.add_chunk(chunk_bytes):
+                        if type(message) is bytes:
+                            message = message.decode("utf-8")
+                        if message.startswith(":"):
+                            continue
+                        chunk = message.removeprefix("data: ")
+                        if chunk == "[DONE]":
+                            continue
+
+                        timestamp = time.perf_counter()
+                        data = json.loads(chunk)
+                        _update_output_stage_metrics_from_payload(
+                            output,
+                            data,
+                            update_output_tokens=(data.get("type") == "ar_delta"),
+                        )
+
+                        chunk_type = data.get("type")
+                        if chunk_type == "ar_delta":
+                            if output.ttft == 0.0:
+                                output.ttft = timestamp - st
+                            else:
+                                output.itl.append(timestamp - most_recent_text_timestamp)
+                            delta = data.get("delta") or ""
+                            generated_text += delta
+                            most_recent_text_timestamp = timestamp
+                            output.text_latency = timestamp - st
+                        elif chunk_type == "image":
+                            image_data = data.get("data")
+                            output.image_count += len(image_data) if isinstance(image_data, list) else 1
+                            content_image_ms = _image_generation_ms_from_content(data.get("data"))
+                            if content_image_ms > 0:
+                                output.image_generation_time_ms += content_image_ms
+                        (
+                            metrics_image_count,
+                            metrics_image_ms,
+                            metrics_image_pixels,
+                            metrics_denoise_step_ms,
+                        ) = _image_metrics_from_stage_metrics(data.get("metrics"))
+                        if metrics_image_count > output.image_count:
+                            output.image_count = metrics_image_count
+                        if metrics_image_ms > output.image_generation_time_ms:
+                            output.image_generation_time_ms = metrics_image_ms
+                        if metrics_image_pixels > output.image_pixels:
+                            output.image_pixels = metrics_image_pixels
+                        if metrics_denoise_step_ms > output.denoise_step_latency_ms:
+                            output.denoise_step_latency_ms = metrics_denoise_step_ms
+                output.latency = timestamp - st
+                output.generated_text = generated_text
+                output.success = True
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception:
+        output.success = False
+        output.error = traceback.format_exc()
+        logger.error(f"ERROR: send image edit request failed, reason is: {output.error}")
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
 async def async_request_openai_audio_speech(
     request_func_input: RequestFuncInput, session: aiohttp.ClientSession, pbar: tqdm | None = None
 ) -> MixRequestFuncOutput:
@@ -727,6 +1115,10 @@ ASYNC_REQUEST_FUNCS["openai-audio-speech"] = async_request_openai_audio_speech
 if "openai-audio-speech" not in OPENAI_COMPATIBLE_BACKENDS:
     OPENAI_COMPATIBLE_BACKENDS.append("openai-audio-speech")
 
+ASYNC_REQUEST_FUNCS["openai-image-edits-omni"] = async_request_openai_image_edits_omni
+if "openai-image-edits-omni" not in OPENAI_COMPATIBLE_BACKENDS:
+    OPENAI_COMPATIBLE_BACKENDS.append("openai-image-edits-omni")
+
 # Daily-Omni backend for audio-visual reasoning benchmark
 # Reuses openai-chat-omni completions for video+text understanding
 ASYNC_REQUEST_FUNCS["daily-omni"] = async_request_openai_chat_omni_completions
@@ -739,7 +1131,10 @@ from vllm.benchmarks import serve
 from vllm.benchmarks.lib.ready_checker import wait_for_endpoint
 from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, get_request
 
-from vllm_omni.benchmarks.metrics.metrics import MultiModalsBenchmarkMetrics, calculate_metrics
+from vllm_omni.benchmarks.metrics.metrics import (
+    MultiModalsBenchmarkMetrics,
+    calculate_metrics,
+)
 
 # ruff: noqa: E402
 
@@ -996,6 +1391,7 @@ async def benchmark(
             max_concurrency=max_concurrency,
             request_rate=request_rate,
             benchmark_duration=benchmark_duration,
+            print_stage=_PRINT_STAGE,
         )
     else:
         metrics = calculate_metrics_for_embeddings(
@@ -1016,9 +1412,13 @@ async def benchmark(
             "request_goodput": metrics.request_goodput if goodput_config_dict else None,
             "output_throughput": metrics.output_throughput,
             "total_token_throughput": metrics.total_token_throughput,
-            "total_audio_duration_s": metrics.total_audio_duration_s,
-            "total_audio_frames": metrics.total_audio_frames,
-            "audio_throughput": metrics.audio_throughput,
+            defs.TOTAL_AUDIO_DURATION_S: getattr(metrics, defs.TOTAL_AUDIO_DURATION_S),
+            defs.TOTAL_AUDIO_FRAMES: getattr(metrics, defs.TOTAL_AUDIO_FRAMES),
+            defs.AUDIO_THROUGHPUT: getattr(metrics, defs.AUDIO_THROUGHPUT),
+            defs.TOTAL_IMAGES: getattr(metrics, defs.TOTAL_IMAGES),
+            defs.IMAGE_THROUGHPUT: getattr(metrics, defs.IMAGE_THROUGHPUT),
+            defs.AVERAGE_PIXELS_PER_IMAGE: getattr(metrics, defs.AVERAGE_PIXELS_PER_IMAGE),
+            defs.MEAN_DENOISE_STEP_LATENCY_MS: getattr(metrics, defs.MEAN_DENOISE_STEP_LATENCY_MS),
             "input_lens": [output.prompt_len for output in outputs],
             "output_lens": actual_output_lens,
             "ttfts": [output.ttft for output in outputs],
@@ -1074,21 +1474,34 @@ async def benchmark(
     if rps_change_events:
         result["rps_change_events"] = rps_change_events
 
+    result_percentile_metrics: list[str] = []
+    if "ttft" in selected_percentile_metrics:
+        result_percentile_metrics.append("ttft")
+    if "tpot" in selected_percentile_metrics or "tpop" in selected_percentile_metrics:
+        result_percentile_metrics.append("tpot")
+    if "itl" in selected_percentile_metrics:
+        result_percentile_metrics.append("itl")
+    if "e2el" in selected_percentile_metrics:
+        result_percentile_metrics.append("e2el")
+    for metric in selected_percentile_metrics:
+        if metric.startswith("audio") and metric not in result_percentile_metrics:
+            result_percentile_metrics.append(metric)
+
     def process_one_metric(
         # E.g., "ttft"
         metric_attribute_name: str,
     ):
         # This function prints and adds statistics of the specified
         # metric.
-        if metric_attribute_name not in selected_percentile_metrics:
+        if metric_attribute_name not in result_percentile_metrics:
             return
         # No text tokens generated (e.g. pure TTS speech endpoint): per-token
         # latency metrics (ttft/tpot/itl) are undefined, so skip them.
         is_text_token_metric = not (metric_attribute_name == "e2el" or metric_attribute_name.startswith("audio"))
         if is_text_token_metric and getattr(metrics, "total_output", 0) == 0:
             return
-        is_audio_rtf = metric_attribute_name == "audio_rtf"
-        is_audio_duration_or_underrun = metric_attribute_name in ("audio_duration", "audio_underrun")
+        is_audio_rtf = metric_attribute_name == defs.AUDIO_RTF
+        is_audio_duration_or_underrun = metric_attribute_name in (defs.AUDIO_DURATION, defs.AUDIO_UNDERRUN)
 
         suffix = "_ms"
         if is_audio_duration_or_underrun:
@@ -1107,9 +1520,10 @@ async def benchmark(
             result[f"p{p_word}_{metric_attribute_name}{suffix}"] = value
 
     if task_type == TaskType.GENERATION:
-        for metric in selected_percentile_metrics:
+        for metric in result_percentile_metrics:
             process_one_metric(metric)
     else:
+        result_percentile_metrics.append("e2el")
         process_one_metric("e2el")
 
     if profile:

vllm_omni/benchmarks/serve.py

@@ -9,6 +9,11 @@ from vllm.benchmarks.serve import main_async
 # This monkey-patches vllm.benchmarks.datasets.get_samples before it's used
 # Must be imported before any vllm.benchmarks module usage
 import vllm_omni.benchmarks.patch.patch  # noqa: F401
+from vllm_omni.benchmarks.patch.patch import (
+    maybe_enable_stage_metrics,
+    set_print_stage,
+    should_request_stage_metrics,
+)
 
 
 def main(args: argparse.Namespace) -> dict[str, Any]:
@@ -18,4 +23,9 @@ def main(args: argparse.Namespace) -> dict[str, Any]:
         os.environ["SEED_TTS_WER_SAVE_ITEMS"] = "1"
     if getattr(args, "daily_omni_save_eval_items", False):
         os.environ["DAILY_OMNI_SAVE_EVAL_ITEMS"] = "1"
+    set_print_stage(getattr(args, "print_stage", False))
+    args.extra_body = maybe_enable_stage_metrics(
+        getattr(args, "extra_body", None),
+        enabled=should_request_stage_metrics(args),
+    )
     return asyncio.run(main_async(args))

vllm_omni/config/__init__.py

@@ -6,7 +6,6 @@ from vllm_omni.config.lora import LoRAConfig
 from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.config.stage_config import (
     DeployConfig,
-    ModelPipeline,
     PipelineConfig,
     StageConfig,
     StageConfigFactory,
@@ -30,7 +29,6 @@ __all__ = [
     "LoRAConfig",
     "StageConfig",
     "StageConfigFactory",
-    "ModelPipeline",
     "StageType",
     "StageExecutionType",
     "StagePipelineConfig",

vllm_omni/config/pipeline_registry.py

@@ -149,6 +149,10 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
         "vllm_omni.model_executor.models.higgs_audio_v2.pipeline",
         "HIGGS_AUDIO_V2_PIPELINE",
     ),
+    "higgs_multimodal_qwen3": (
+        "vllm_omni.model_executor.models.higgs_audio_v3.pipeline",
+        "HIGGS_AUDIO_V3_PIPELINE",
+    ),
     "dynin_omni": (
         "vllm_omni.model_executor.models.dynin_omni.pipeline",
         "DYNIN_OMNI_PIPELINE",

vllm_omni/config/stage_config.py

@@ -20,13 +20,6 @@ from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 
-_MODELS_DIR = Path(__file__).resolve().parent.parent / "model_executor" / "models"
-
-
-def get_pipeline_path(model_dir: str, filename: str) -> Path:
-    return _MODELS_DIR / model_dir / filename
-
-
 logger = init_logger(__name__)
 
 
@@ -1061,77 +1054,6 @@ class StageConfig:
         return create_config(config_dict)
 
 
-@dataclass
-class ModelPipeline:
-    """Complete pipeline definition for a multi-stage model (legacy).
-
-    TODO(@lishunyang12): remove once all models migrate to PipelineConfig.
-    """
-
-    model_type: str
-    stages: list[StageConfig]
-
-    # Pipeline-wide behavior flags
-    async_chunk: bool = False
-
-    # Optional distributed configuration
-    connectors: dict[str, Any] | None = None
-    edges: list[dict[str, Any]] | None = None
-
-    def get_stage(self, stage_id: int) -> StageConfig | None:
-        """Look up a stage by its ID.
-
-        Args:
-            stage_id: The stage ID to search for.
-
-        Returns:
-            The matching StageConfig, or None if not found.
-        """
-        for stage in self.stages:
-            if stage.stage_id == stage_id:
-                return stage
-        return None
-
-    def validate_pipeline(self) -> list[str]:
-        """Validate pipeline topology at model integration time (not runtime).
-
-        Checks:
-        - All stage IDs are unique
-        - All input_sources reference valid stage IDs
-        - At least one entry point (stage with empty input_sources)
-
-        Returns:
-            List of validation error messages. Empty list if valid.
-        """
-        errors: list[str] = []
-
-        if not self.stages:
-            errors.append("Topology has no stages defined")
-            return errors
-
-        # Check for unique stage IDs
-        stage_ids = [s.stage_id for s in self.stages]
-        if len(stage_ids) != len(set(stage_ids)):
-            errors.append("Duplicate stage IDs found")
-
-        stage_id_set = set(stage_ids)
-
-        # Check input_sources reference valid stages
-        for stage in self.stages:
-            for source_id in stage.input_sources:
-                if source_id not in stage_id_set:
-                    errors.append(f"Stage {stage.stage_id} references non-existent input source {source_id}")
-                if source_id == stage.stage_id:
-                    errors.append(f"Stage {stage.stage_id} references itself as input source")
-
-        # Check for at least one entry point
-        entry_points = [s for s in self.stages if not s.input_sources]
-        if not entry_points:
-            errors.append("No entry point found (stage with empty input_sources)")
-
-        return errors
-
-
 class StageConfigFactory:
     """Factory that loads pipeline YAML and merges CLI overrides.
 
@@ -1206,33 +1128,9 @@ class StageConfigFactory:
                             continue
                     return cls._create_from_registry(registered.model_type, cli_overrides, deploy_config_path)
 
-        # --- Legacy path: load from pipeline YAML ---
-        pipeline = cls._load_pipeline(model, trust_remote_code=trust_remote_code)
-
-        if pipeline is None:
-            return None
-
-        errors = pipeline.validate_pipeline()
-        if errors:
-            logger.warning(f"Pipeline validation warnings for {model}: {errors}")
-
-        # Materialize the resolved pipeline-wide async_chunk value into every
-        # stage so build_engine_args_dict() can inject the stage connector
-        # spec and explicit False overrides are preserved.
-        resolved_async_chunk = cli_overrides.get("async_chunk")
-        if resolved_async_chunk is None:
-            resolved_async_chunk = bool(pipeline.async_chunk)
-        for stage in pipeline.stages:
-            stage.yaml_engine_args["async_chunk"] = bool(resolved_async_chunk)
-
-        # Apply CLI overrides
-        result: list[StageConfig] = []
-        for stage in pipeline.stages:
-            # Merge global CLI overrides
-            stage.runtime_overrides = cls._merge_cli_overrides(stage, cli_overrides)
-            result.append(stage)
-
-        return result
+        # Not in the pipeline registry — let the caller fall back to the
+        # legacy ``stage_configs/*.yaml`` path (resolve_model_config_path).
+        return None
 
     @classmethod
     def _create_from_registry(
@@ -1355,193 +1253,6 @@ class StageConfigFactory:
 
         return [config_dict]
 
-    @classmethod
-    def _load_pipeline(cls, model: str, trust_remote_code: bool = True) -> ModelPipeline | None:
-        """Load a legacy ``pipeline.yaml`` for the model.
-
-        Searches ``model_executor/models/<dir>/pipeline.yaml`` by trying
-        (a) the raw ``model_type`` as the directory name, then
-        (b) ``model_type`` with hyphens replaced by underscores,
-        and finally (c) scanning every ``pipeline.yaml`` for one that
-        declares a matching ``model_type`` or ``hf_architectures``.
-
-        Returns None if no pipeline.yaml is found — caller handles the
-        ``resolve_model_config_path`` fallback via stage_configs/ YAMLs.
-        """
-        model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
-        if model_type is None:
-            return None
-
-        # Direct lookups by convention
-        candidates = [model_type, model_type.replace("-", "_")]
-        for dir_name in candidates:
-            pipeline_path = get_pipeline_path(dir_name, "pipeline.yaml")
-            if pipeline_path.exists():
-                return cls._parse_pipeline_yaml(pipeline_path, model_type)
-
-        # Scan fallback: read every pipeline.yaml and match on declared fields
-        hf_archs = set(getattr(hf_config, "architectures", []) or []) if hf_config else set()
-        if _MODELS_DIR.exists():
-            for subdir in sorted(_MODELS_DIR.iterdir()):
-                if not subdir.is_dir():
-                    continue
-                pipeline_path = subdir / "pipeline.yaml"
-                if not pipeline_path.exists():
-                    continue
-                try:
-                    cfg = load_yaml_config(pipeline_path)
-                except Exception as exc:
-                    logger.debug("Skip %s: %s", pipeline_path, exc)
-                    continue
-                declared_type = getattr(cfg, "model_type", None)
-                declared_archs = set(getattr(cfg, "hf_architectures", None) or [])
-                if declared_type == model_type or (hf_archs and hf_archs.intersection(declared_archs)):
-                    return cls._parse_pipeline_yaml(pipeline_path, declared_type or model_type)
-
-        logger.debug("No pipeline.yaml found for model_type %s (archs=%s)", model_type, sorted(hf_archs))
-        return None
-
-    # Keys consumed as explicit StageConfig fields — everything else is
-    # passed through via yaml_extras.
-    _KNOWN_STAGE_KEYS: set[str] = {
-        "stage_id",
-        "model_stage",
-        "stage_type",
-        "input_sources",
-        "engine_input_source",
-        "custom_process_input_func",
-        "final_output",
-        "final_output_type",
-        "worker_type",
-        "scheduler_cls",
-        "hf_config_name",
-        "is_comprehension",
-        "engine_args",
-        "runtime",
-    }
-
-    @classmethod
-    def _parse_pipeline_yaml(cls, path: Path, model_type: str) -> ModelPipeline:
-        """Parse a pipeline YAML file.
-
-        Args:
-            path: Path to the YAML file.
-            model_type: Model type identifier.
-
-        Returns:
-            ModelPipeline object.
-        """
-        config_data = load_yaml_config(path)
-
-        stages: list[StageConfig] = []
-        for stage_data in config_data.stages:
-            # Use .get() for optional fields — idiomatic for OmegaConf DictConfig
-            stage_type_str = stage_data.get("stage_type", "llm")
-            stage_type = StageType(stage_type_str) if stage_type_str else StageType.LLM
-
-            # Handle both 'input_sources' (new) and 'engine_input_source' (legacy)
-            input_sources = stage_data.get("input_sources", None)
-            if input_sources is None:
-                input_sources = stage_data.get("engine_input_source", [])
-            if input_sources is None:
-                input_sources = []
-            input_sources = list(input_sources)
-
-            # Extract per-stage engine_args and runtime dicts
-            raw_ea = stage_data.get("engine_args", None)
-            yaml_engine_args = to_dict(raw_ea) if raw_ea is not None else {}
-            raw_rt = stage_data.get("runtime", None)
-            yaml_runtime = to_dict(raw_rt) if raw_rt is not None else {}
-
-            # Migrate legacy runtime.max_batch_size → engine_args.max_num_seqs
-            if "max_batch_size" in yaml_runtime:
-                mbs = yaml_runtime.pop("max_batch_size")
-                yaml_engine_args.setdefault("max_num_seqs", int(mbs))
-                logger.debug(
-                    "Stage %s: migrated runtime.max_batch_size=%s to engine_args.max_num_seqs",
-                    stage_data.get("stage_id", "?"),
-                    mbs,
-                )
-
-            # Topology-level fields that also live inside engine_args in legacy
-            # YAMLs (worker_type, scheduler_cls, etc.) — read from both places.
-            worker_type = stage_data.get("worker_type", None) or yaml_engine_args.pop("worker_type", None)
-            scheduler_cls = stage_data.get("scheduler_cls", None) or yaml_engine_args.pop("scheduler_cls", None)
-            if scheduler_cls:
-                async_sched = yaml_engine_args.get("async_scheduling")
-                if async_sched is not None:
-                    logger.warning(
-                        "Stage %s: async_scheduling=%r and scheduler_cls=%r "
-                        "should not be set together. scheduler_cls will take "
-                        "precedence for which scheduler is used.",
-                        stage_data.stage_id,
-                        async_sched,
-                        scheduler_cls,
-                    )
-                else:
-                    logger.warning(
-                        "Stage %s: scheduler_cls=%r is deprecated. Use async_scheduling instead.",
-                        stage_data.stage_id,
-                        scheduler_cls,
-                    )
-            hf_config_name = stage_data.get("hf_config_name", None) or yaml_engine_args.pop("hf_config_name", None)
-            model_stage = getattr(stage_data, "model_stage", None) or yaml_engine_args.pop("model_stage", None)
-
-            # Collect pass-through fields (default_sampling_params,
-            # output_connectors, input_connectors, tts_args, etc.)
-            yaml_extras: dict[str, Any] = {}
-            for key in stage_data:
-                if key not in cls._KNOWN_STAGE_KEYS:
-                    val = stage_data[key]
-                    try:
-                        yaml_extras[key] = to_dict(val)
-                    except ValueError:
-                        yaml_extras[key] = val
-
-            stage = StageConfig(
-                stage_id=stage_data.stage_id,
-                model_stage=model_stage or "",
-                stage_type=stage_type,
-                input_sources=input_sources,
-                custom_process_input_func=stage_data.get("custom_process_input_func", None),
-                final_output=stage_data.get("final_output", False),
-                final_output_type=stage_data.get("final_output_type", None),
-                worker_type=worker_type,
-                scheduler_cls=scheduler_cls,
-                hf_config_name=hf_config_name,
-                is_comprehension=stage_data.get("is_comprehension", False),
-                yaml_engine_args=yaml_engine_args,
-                yaml_runtime=yaml_runtime,
-                yaml_extras=yaml_extras,
-            )
-            stages.append(stage)
-
-        # Get pipeline-wide flags
-        async_chunk = config_data.get("async_chunk", False)
-
-        # Get optional connector config — check both top-level and nested
-        # under ``runtime`` (legacy stage_configs format).
-        connectors = None
-        edges = None
-        if hasattr(config_data, "connectors"):
-            connectors = to_dict(config_data.connectors)
-        if hasattr(config_data, "edges"):
-            edges = to_dict(config_data.edges)
-        if hasattr(config_data, "runtime") and config_data.runtime is not None:
-            top_runtime = config_data.runtime
-            if connectors is None and hasattr(top_runtime, "connectors"):
-                connectors = to_dict(top_runtime.connectors)
-            if edges is None and hasattr(top_runtime, "edges"):
-                edges = to_dict(top_runtime.edges)
-
-        return ModelPipeline(
-            model_type=getattr(config_data, "model_type", model_type),
-            stages=stages,
-            async_chunk=async_chunk,
-            connectors=connectors,
-            edges=edges,
-        )
-
     @classmethod
     def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]:
         """Auto-detect model_type from model directory.

vllm_omni/core/prefix_cache.py

@@ -391,27 +391,33 @@ class OmniTensorPrefixCache:
         self,
         query_start_loc: torch.Tensor,
         input_batch: InputBatch,
-        multimodal_outputs: dict,
+        multimodal_outputs: dict | None,
         num_scheduled_tokens: dict[str, int],
     ):
         """Get the merged multimodal states if hidden state prefix caching is enabled."""
         combined_multimodal_outputs = {}
+        # Talkers that produce multimodal outputs only at decode (e.g. Higgs
+        # Audio v3's audio codes) have ``multimodal_outputs is None`` at the
+        # initial prefill step. Treat it as an empty mapping so the merger
+        # short-circuits cleanly and the cache merge still runs once codes
+        # start arriving.
+        mm_outputs = multimodal_outputs if multimodal_outputs is not None else {}
         # First get the prefix cached tensors that are present in the mm data
         for mm_key in self.mm_cache_keys:
-            if mm_key in multimodal_outputs:
+            if mm_key in mm_outputs:
                 combined_multimodal_outputs[mm_key] = self._get_merged_tensors(
                     query_start_loc=query_start_loc,
                     input_batch=input_batch,
                     cache=self.mm_outputs_cache[mm_key],
-                    hidden_states=multimodal_outputs[mm_key],
+                    hidden_states=mm_outputs[mm_key],
                     num_scheduled_tokens=num_scheduled_tokens,
                 )
 
         # Then, get everything else (passthrough data); first, convert to CPU
         # tensors similarly to the non prefix cached path, and then populate
         # the subdicts mapping request IDs -> payload objects
-        passthrough_keys = set(multimodal_outputs.keys()) - self.mm_cache_keys
-        passthrough_mm_data = {k: v for k, v in multimodal_outputs.items() if k in passthrough_keys}
+        passthrough_keys = set(mm_outputs.keys()) - self.mm_cache_keys
+        passthrough_mm_data = {k: v for k, v in mm_outputs.items() if k in passthrough_keys}
         mm_cpu = build_mm_cpu(multimodal_outputs=passthrough_mm_data)
 
         for mm_key, mm_val in mm_cpu.items():

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -13,7 +13,7 @@ from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMSchedu
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs, FinishReason
 from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
@@ -69,6 +69,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # alive.  Stopped on the first scheduler step after extraction ack.
         self.pending_stop_after_extraction: set[str] = set()
 
+        self.finished_req_ids_dict = defaultdict(set)
+
         # [Omni] Pre-parse KV transfer criteria
         self.kv_transfer_criteria = self._get_kv_transfer_criteria()
 
@@ -434,7 +436,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                             # Downstream async-chunk stages receive real payloads from the
                             # connector. This update only resumes polling for the next segment.
                             self.chunk_transfer_adapter.segment_finished_requests.discard(request.request_id)
-                    request.discard_latest_async_tokens = True
+                    request.async_tokens_to_discard = 1
                     request.num_output_placeholders = 0
                     request.spec_token_ids = []
                     request._output_token_ids.clear()
@@ -556,16 +558,24 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # outputs in this step.
         engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
 
+        # FIXME: finished_req_ids_dict is unconditionally initialized as
+        # defaultdict(set) in __init__ (not gated by include_finished_set).
+        # This branch is therefore always eligible once any client_index is
+        # populated; revisit when wiring streaming-only / upstream semantics.
         finished_req_ids = self.finished_req_ids_dict
         if finished_req_ids:
             # Include ids of requests that finished since last outputs
             # were sent.
             for client_index, finished_set in finished_req_ids.items():
-                # Set finished request set in EngineCoreOutputs for this client.
-                if (eco := engine_core_outputs.get(client_index)) is not None:
-                    eco.finished_requests = finished_set
-                else:
-                    engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
+                eco = engine_core_outputs.get(client_index)
+                if eco is None:
+                    eco = EngineCoreOutputs()
+                    engine_core_outputs[client_index] = eco
+                emitted = {o.request_id for o in eco.outputs}
+                for req_id in finished_set:
+                    if req_id not in emitted:
+                        eco.outputs.append(EngineCoreOutput(req_id, [], finish_reason=FinishReason.ABORT))
+                eco.finished_requests = finished_set
             finished_req_ids.clear()
 
         if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats)) is not None:
@@ -632,13 +642,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         """
         req_id = session.request_id
         self._new_prompt_len_snapshot[req_id] = len(update.prompt_token_ids)
-        session.discard_latest_async_tokens = False
         outstanding_async_tokens = getattr(session, "num_output_placeholders", 0)
         if outstanding_async_tokens > 0:
             # Async scheduling may already have sampled the previous
             # segment's next token. Drop that late token instead of
             # appending it to the new streaming segment.
-            session.discard_latest_async_tokens = True
+            session.async_tokens_to_discard = 1
             session.num_computed_tokens -= session.num_output_placeholders
             session.num_output_placeholders = 0
             session.spec_token_ids = []

vllm_omni/data_entry_keys.py

@@ -63,6 +63,7 @@ class OmniPayloadMeta(TypedDict, total=False):
     stream_finished: torch.Tensor
     req_id: list[str]
     left_context_size: int
+    right_holdback_size: int
     override_keys: list[tuple[str, str]]
     num_processed_tokens: int
     next_stage_prompt_len: int
@@ -150,6 +151,7 @@ class MetaStruct(_StructBase):
     stream_finished: torch.Tensor | None = None
     req_id: list[str] | None = None
     left_context_size: int | None = None
+    right_holdback_size: int | None = None
     override_keys: list[tuple[str, str]] | None = None
     num_processed_tokens: int | None = None
     next_stage_prompt_len: int | None = None

vllm_omni/deploy/higgs_multimodal_qwen3.yaml

@@ -0,0 +1,74 @@
+# Higgs-audio v3 deploy: Stage 0 (talker) -> Stage 1 (code2wav) via shared memory.
+#
+# Plain TTS only (no voice cloning). async_chunk streams codec chunks from
+# Stage 0 to Stage 1 as they are produced (TTFA approaches one AR loop's
+# worth of work). Stage 1 honors meta.left_context_size + meta.right_holdback_size
+# so the per-chunk codec windows stitch into a coherent PCM stream.
+async_chunk: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      # Stream PCM bytes back to the client as Stage 1 emits each chunk.
+      codec_streaming: true
+      # Streaming chunk shape for talker2code2wav_async_chunk:
+      #   codec_chunk_frames        steady-state new frames per emit
+      #   codec_left_context_frames sliding past context to avoid leading-edge artifacts
+      #   codec_right_holdback_frames trailing frames held back so the codec
+      #                              has future context for emit-region tail samples
+      #   initial_codec_chunk_frames smaller first chunk to minimize TTFA
+      codec_chunk_frames: 25
+      codec_left_context_frames: 25
+      codec_right_holdback_frames: 4
+      initial_codec_chunk_frames: 1
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 5000
+      connector_get_max_wait: 500
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 16
+    gpu_memory_utilization: 0.6
+    enforce_eager: true
+    trust_remote_code: true
+    # Stage-0 prefix caching is on by default. The talker's HS opt-out and
+    # deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py)
+    # keep per-step bookkeeping negligible; any shared text prefix across
+    # concurrent requests is a net throughput win. Stage 1 (codec decoder)
+    # has no KV cache and stays disabled.
+    enable_prefix_caching: true
+    async_scheduling: false
+    max_num_batched_tokens: 4096
+    max_model_len: 8192
+    devices: "0"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 1.0
+      top_p: 0.95
+      top_k: 50
+      max_tokens: 2048
+      seed: 42
+      repetition_penalty: 1.0
+
+  - stage_id: 1
+    max_num_seqs: 16
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 65536
+    max_model_len: 65536
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      top_p: 1.0
+      top_k: -1
+      max_tokens: 65536
+      seed: 42
+      repetition_penalty: 1.0

vllm_omni/deploy/qwen3_omni_moe.yaml

@@ -18,6 +18,7 @@ connectors:
   connector_of_shared_memory:
     name: SharedMemoryConnector
     extra:
+      initial_codec_chunk_frames: 4
       codec_chunk_frames: 25
       codec_left_context_frames: 25
 
@@ -30,12 +31,8 @@ stages:
     enable_prefix_caching: false
     devices: "0"
     default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
+      temperature: 0.0
       max_tokens: 2048
-      seed: 42
-      repetition_penalty: 1.05
 
   - stage_id: 1
     max_num_batched_tokens: 32768

vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml

@@ -18,6 +18,7 @@ connectors:
   connector_of_shared_memory:
     name: SharedMemoryConnector
     extra:
+      initial_codec_chunk_frames: 4
       codec_chunk_frames: 25
       codec_left_context_frames: 25
 

vllm_omni/deploy/voxcpm2.yaml

@@ -1,19 +1,13 @@
 # VoxCPM2 deploy: single-stage AR pipeline with per-request state batching.
 # Verified on 1x H20 141GB.
 #
-# Fields omitted from a stage fall back to StageDeployConfig defaults (see
-# vllm_omni/config/stage_config.py). Notable choices:
-#   * enforce_eager=true because the talker's KV-cache save/restore loop is
-#     not cudagraph-compatible (captured separately via voxcpm2_talker's own
-#     _CapturedGraph path).
-#   * max_num_seqs=4 matches the legacy max_batch_size=4 contract: concurrent
-#     requests are supported via per-request StaticKVCache save/restore.
+# VoxCPM2-local runtime knobs are passed through hf_overrides.
 async_chunk: false
 dtype: bfloat16
 
 stages:
   - stage_id: 0
-    max_num_seqs: 4
+    max_num_seqs: 32
     gpu_memory_utilization: 0.9
     enforce_eager: true
     trust_remote_code: true
@@ -22,6 +16,18 @@ stages:
     max_num_batched_tokens: 4096
     max_model_len: 4096
     devices: "0"
+    engine_extras:
+      hf_overrides:
+        voxcpm2_runtime_config:
+          audio_emit_every: 1
+          vae_decode_every: 3
+          enable_batched_cfm: true
+          enable_batched_vae_decode: true
+          coalesce_audio_d2h: true
+          enable_loc_dit_fused_qkv: true
+          enable_loc_dit_fused_mlp: true
+          enable_loc_dit_zero_dt_cache: true
+          enable_loc_dit_skip_qkv_contig: true
     default_sampling_params:
       temperature: 0.0
       top_p: 1.0

vllm_omni/diffusion/cache/cache_dit_backend.py

@@ -242,6 +242,80 @@ def enable_cache_for_wan22(pipeline: Any, cache_config: Any) -> Callable[[int],
     return refresh_cache_context
 
 
+def enable_cache_for_wan22_s2v(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
+    """Enable cache-dit for Wan2.2 S2V.
+
+    S2V uses a single transformer, but unlike the other Wan2.2 variants its
+    block loop calls each block as ``block(hidden_states, **kwargs)`` and keeps
+    the timestep modulation state in ``e`` rather than a second positional
+    tensor. CacheDiT Pattern_3 matches that contract: cache hidden states only
+    and pass the remaining conditioning through kwargs unchanged.
+
+    The S2V transformer has an ``after_transformer_block`` method that injects
+    audio embeddings after specific layers. The cached blocks wrapper
+    (Wan22S2VCachedBlocks._run_block) calls the original internally, so we
+    permanently replace it with a no-op on the transformer to prevent double
+    injection from the main forward loop.
+    """
+    db_cache_config = _build_db_cache_config(cache_config)
+    calibrator_config = None
+    if cache_config.enable_taylorseer:
+        taylorseer_order = cache_config.taylorseer_order
+        calibrator_config = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
+        logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
+
+    # Save the original after_transformer_block before cache-dit wrapping
+    transformer = pipeline.transformer
+    if hasattr(transformer, "after_transformer_block"):
+        transformer._cache_dit_original_after_transformer_block = transformer.after_transformer_block
+
+    Wan22S2VCachedAdapter.apply(
+        BlockAdapter(
+            transformer=transformer,
+            blocks=[transformer.blocks],
+            forward_pattern=[ForwardPattern.Pattern_3],
+            params_modifiers=[
+                ParamsModifier(cache_config=db_cache_config, calibrator_config=calibrator_config),
+            ],
+            has_separate_cfg=True,
+        ),
+        cache_config=db_cache_config,
+        calibrator_config=calibrator_config,
+    )
+
+    # Permanently replace after_transformer_block with a no-op.
+    # The cached blocks wrapper (Wan22S2VCachedBlocks._run_block) already calls
+    # the original via _cache_dit_original_after_transformer_block.
+    def _noop_after_transformer_block(block_idx: int, hidden_states: torch.Tensor) -> torch.Tensor:
+        return hidden_states
+
+    transformer.after_transformer_block = _noop_after_transformer_block
+
+    def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
+        """Refresh cache context for the S2V transformer."""
+        if cache_config.scm_steps_mask_policy is None:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                num_inference_steps=num_inference_steps,
+                verbose=verbose,
+            )
+        else:
+            cache_dit.refresh_context(
+                pipeline.transformer,
+                cache_config=DBCacheConfig().reset(
+                    num_inference_steps=num_inference_steps,
+                    steps_computation_mask=cache_dit.steps_mask(
+                        mask_policy=cache_config.scm_steps_mask_policy,
+                        total_steps=num_inference_steps,
+                    ),
+                    steps_computation_policy=cache_config.scm_steps_policy,
+                ),
+                verbose=verbose,
+            )
+
+    return refresh_cache_context
+
+
 def enable_cache_for_longcat_image(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
     """Enable cache-dit for LongCatImage pipeline.
 
@@ -986,6 +1060,151 @@ class BagelCachedBlocks(CachedBlocks_Pattern_0_1_2):
         return hidden_states, encoder_hidden_states
 
 
+class Wan22S2VCachedBlocks(CachedBlocks_Pattern_3_4_5):
+    """CacheDiT blocks wrapper that preserves S2V per-layer audio injection."""
+
+    def _run_block(self, block_id: int, block: torch.nn.Module, hidden_states: torch.Tensor, *args, **kwargs):
+        hidden_states = block(hidden_states, *args, **kwargs)
+        hidden_states, new_encoder_hidden_states = self._process_block_outputs(hidden_states)
+        original_after_transformer_block = getattr(
+            self.transformer,
+            "_cache_dit_original_after_transformer_block",
+            getattr(self.transformer, "after_transformer_block", None),
+        )
+        if original_after_transformer_block is not None:
+            hidden_states = original_after_transformer_block(block_id, hidden_states)
+        return hidden_states, new_encoder_hidden_states
+
+    def call_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        new_encoder_hidden_states = None
+        for block_id, block in enumerate(self.transformer_blocks):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+        return hidden_states, new_encoder_hidden_states
+
+    def call_Fn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        new_encoder_hidden_states = None
+        for block_id, block in enumerate(self._Fn_blocks()):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+        return hidden_states, new_encoder_hidden_states
+
+    def call_Mn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        original_hidden_states = hidden_states
+        new_encoder_hidden_states = None
+        start_idx = self.context_manager.Fn_compute_blocks()
+        for block_id, block in enumerate(self._Mn_blocks(), start=start_idx):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+
+        hidden_states = hidden_states.contiguous()
+        hidden_states_residual = hidden_states - original_hidden_states.to(hidden_states.device)
+
+        return (
+            hidden_states,
+            new_encoder_hidden_states,
+            hidden_states_residual,
+        )
+
+    def call_Bn_blocks(
+        self,
+        hidden_states: torch.Tensor,
+        *args,
+        **kwargs,
+    ):
+        new_encoder_hidden_states = None
+        if self.context_manager.Bn_compute_blocks() == 0:
+            return hidden_states, new_encoder_hidden_states
+
+        start_idx = len(self.transformer_blocks) - self.context_manager.Bn_compute_blocks()
+        for block_id, block in enumerate(self._Bn_blocks(), start=start_idx):
+            hidden_states, new_encoder_hidden_states = self._run_block(
+                block_id,
+                block,
+                hidden_states,
+                *args,
+                **kwargs,
+            )
+
+        return hidden_states, new_encoder_hidden_states
+
+
+class Wan22S2VCachedAdapter(CachedAdapter):
+    """CacheDiT adapter that uses Wan22S2VCachedBlocks for S2V audio injection.
+
+    Only overrides collect_unified_blocks to use Wan22S2VCachedBlocks (which
+    calls after_transformer_block per-layer internally). The base class
+    mock_transformer handles the forward wrapping — after_transformer_block is
+    permanently replaced with a no-op in enable_cache_for_wan22_s2v() to prevent
+    double injection.
+    """
+
+    @classmethod
+    def collect_unified_blocks(
+        cls,
+        block_adapter: BlockAdapter,
+        contexts_kwargs: list[dict],
+    ) -> list[dict[str, torch.nn.ModuleList]]:
+        BlockAdapter.assert_normalized(block_adapter)
+
+        total_cached_blocks: list[dict[str, torch.nn.ModuleList]] = []
+        assert hasattr(block_adapter.pipe, "_context_manager")
+
+        for i in range(len(block_adapter.transformer)):
+            unified_blocks_bind_context = {}
+            for j in range(len(block_adapter.blocks[i])):
+                cache_config: BasicCacheConfig = contexts_kwargs[i * len(block_adapter.blocks[i]) + j]["cache_config"]
+                unified_blocks_bind_context[block_adapter.unique_blocks_name[i][j]] = torch.nn.ModuleList(
+                    [
+                        Wan22S2VCachedBlocks(
+                            block_adapter.blocks[i][j],
+                            transformer=block_adapter.transformer[i],
+                            forward_pattern=block_adapter.forward_pattern[i][j],
+                            check_forward_pattern=block_adapter.check_forward_pattern,
+                            check_num_outputs=block_adapter.check_num_outputs,
+                            cache_prefix=block_adapter.blocks_name[i][j],
+                            cache_context=block_adapter.unique_blocks_name[i][j],
+                            context_manager=block_adapter.pipe._context_manager,
+                            cache_type=cache_config.cache_type,
+                        )
+                    ]
+                )
+
+            total_cached_blocks.append(unified_blocks_bind_context)
+
+        return total_cached_blocks
+
+
 class BagelCachedAdapter(CachedAdapter):
     """
     Custom CachedAdapter for Bagel that uses BagelCachedContextManager and BagelCachedBlocks.
@@ -1808,6 +2027,7 @@ CUSTOM_DIT_ENABLERS.update(
     {
         "Wan22Pipeline": enable_cache_for_wan22,
         "Wan22I2VPipeline": enable_cache_for_wan22,
+        "Wan22S2VPipeline": enable_cache_for_wan22_s2v,
         "HunyuanImage3Pipeline": enable_cache_for_hunyuan_image3,
         "FluxPipeline": enable_cache_for_flux,
         "Flux2KleinPipeline": enable_cache_for_flux2_klein,

vllm_omni/diffusion/cache/teacache/coefficient_estimator.py

@@ -81,7 +81,7 @@ class DefaultAdapter:
 
         loader = DiffusersPipelineLoader(LoadConfig(), od_config=od_config)
         # load_model will handle dtypes / device placement, put in .eval() mode
-        return loader.load_model(od_config=od_config, load_device=device)
+        return loader.load_model(load_device=device)
 
     @staticmethod
     def get_transformer(pipeline: Any) -> tuple[Any, str]: