build: bump docker base images and install docs to v0.22.0

Update Dockerfiles, quickstart, and installation guides to reference
vLLM v0.22.0 (CUDA 12.9, ROCm 7.2.2) instead of v0.20.0/v0.21.0.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.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:
           - |

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.
 - [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/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/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, 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/user_guide/diffusion_features.md

@@ -145,6 +145,7 @@ 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**                  |     ❌     |     ✅      |           ✅           |       ✅        |         ✅         |         ❌         |   ❌    |             ❌             |          ❌           |       ❌        |        ❌         |

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

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 = [

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 generation, text to video with sound, action policy  | 1x H200 141GB / B300 |
+| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V 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 |

recipes/cosmos3/Cosmos3-Nano.md

@@ -6,7 +6,7 @@
 
 - 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), and image-to-video (I2V) 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
 
@@ -23,6 +23,23 @@ the mode is selected per request:
 - **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 **and** an action trajectory, roll out
+    the resulting video. Synchronous: `POST /v1/videos/sync`.
+  - `policy` — given a first frame 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
+    (`{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`).
+
+  `inverse_dynamics` (recover the action from a given video) is supported by the
+  pipeline; **online inference of inverse dynamics will be added in a follow-up MR.**
 
 ## References
 
@@ -144,6 +161,35 @@ 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
 ```
 
 #### Notes
@@ -152,6 +198,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)
@@ -173,8 +220,10 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
     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` (sync `/v1/videos/sync`) and `policy` (async
+    `/v1/videos`, returns the predicted action under the top-level `action`
+    field) are supported online. **Online inference of inverse dynamics will be
+    added in a follow-up MR.**
 
 ### 1x GPU (Offline generation)
 

recipes/cosmos3/Cosmos3-Super.md

@@ -104,5 +104,9 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
   - (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.
+  semantics as Nano, including the **action** modality: `forward_dynamics`
+  (sync `/v1/videos/sync`) and `policy` (async `/v1/videos`, predicted action under
+  the top-level `action` field) — see the Cosmos3-Nano recipe for the request shape.
+  Online inference of inverse dynamics will be added in a follow-up MR. 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/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,16 @@ 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)
+
 
 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 +451,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)
@@ -463,8 +480,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 +525,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 +599,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 +645,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)
@@ -651,6 +684,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/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/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_video_server.py

@@ -627,6 +627,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():
@@ -1038,6 +1147,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)
 # ---------------------------------------------------------------------------

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

@@ -1638,9 +1638,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 +1957,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/unit/higgs_audio_v3/test_higgs_audio_v3.py

@@ -0,0 +1,910 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Deterministic unit tests for higgs-audio v3.
+
+These tests verify AC-1 (config), AC-3 (prompt), AC-4 (fused modules),
+AC-5 (delay pattern), AC-7 (stage processor), and AC-10 (registry)
+without requiring the actual checkpoint or GPU.
+"""
+
+import pytest
+import torch
+
+# ---- AC-1: Configuration ----
+
+
+class TestHiggsAudioV3Config:
+    def test_default_config_loads(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
+            HiggsAudioV3Config,
+        )
+
+        config = HiggsAudioV3Config()
+        assert config.num_codebooks == 8
+        assert config.codebook_size == 1026
+        assert config.audio_stream_bos_id == 1024
+        assert config.audio_stream_eos_id == 1025
+        assert config.sample_rate == 24000
+        assert config.frame_rate == 25
+        assert config.num_real_codes == 1024
+        assert config.tie_modality_embeddings is True
+
+    def test_custom_audio_encoder_config(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
+            HiggsAudioV3Config,
+        )
+
+        config = HiggsAudioV3Config(
+            audio_encoder_config={
+                "encoder_type": "discrete",
+                "num_codebooks": 4,
+                "vocab_size": 512,
+                "tie_word_embeddings": False,
+            }
+        )
+        assert config.num_codebooks == 4
+        assert config.codebook_size == 512
+        assert config.tie_modality_embeddings is False
+
+    def test_invalid_num_codebooks_rejected(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
+            HiggsAudioV3Config,
+        )
+
+        with pytest.raises(ValueError, match="num_codebooks must be > 0"):
+            HiggsAudioV3Config(audio_encoder_config={"num_codebooks": 0, "vocab_size": 1026})
+
+    def test_negative_num_codebooks_rejected(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
+            HiggsAudioV3Config,
+        )
+
+        with pytest.raises(ValueError, match="num_codebooks must be > 0"):
+            HiggsAudioV3Config(audio_encoder_config={"num_codebooks": -1, "vocab_size": 1026})
+
+    def test_special_token_ids_initially_none(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
+            HiggsAudioV3Config,
+        )
+
+        config = HiggsAudioV3Config()
+        assert config.tts_token_id is None
+        assert config.text_token_id is None
+        assert config.audio_continuation_id is None
+
+    def test_hidden_size_from_text_config(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
+            HiggsAudioV3Config,
+        )
+
+        config = HiggsAudioV3Config(text_config={"model_type": "qwen3", "hidden_size": 2560})
+        assert config.hidden_size == 2560
+        assert config.audio_hidden_size == 2560
+
+
+# ---- AC-4: Fused Multi-Codebook Modules ----
+
+
+class TestFusedModules:
+    def test_embedding_shape(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            HiggsFusedMultiTextEmbedding,
+        )
+
+        embed = HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=256)
+        assert embed.weight.shape == (8 * 1026, 256)
+
+        codes = torch.randint(0, 1024, (5, 8))
+        out = embed(codes)
+        assert out.shape == (5, 256)
+
+    def test_embedding_offset_indexing(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            HiggsFusedMultiTextEmbedding,
+        )
+
+        embed = HiggsFusedMultiTextEmbedding(num_codebooks=2, vocab_size=4, hidden_size=3)
+        # Set weights to identity-like pattern so we can verify offsets
+        with torch.no_grad():
+            embed.weight.zero_()
+            # Codebook 0, vocab [0..3] -> row 0..3
+            # Codebook 1, vocab [0..3] -> row 4..7
+            for i in range(8):
+                embed.weight[i, i % 3] = float(i)
+
+        # codes[0] = [1, 2] -> embed[1] + embed[4+2=6]
+        codes = torch.tensor([[1, 2]])
+        out = embed(codes)
+        expected = embed.weight[1] + embed.weight[6]
+        assert torch.allclose(out[0], expected)
+
+    def test_head_shape(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            HiggsFusedMultiTextHead,
+        )
+
+        head = HiggsFusedMultiTextHead(num_codebooks=8, vocab_size=1026, hidden_size=256)
+        assert head.weight.shape == (8 * 1026, 256)
+
+        hidden = torch.randn(3, 256)
+        out = head.generate(hidden)
+        assert out.shape == (3, 8, 1026)
+
+    def test_tying(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            HiggsFusedMultiTextEmbedding,
+            HiggsFusedMultiTextHead,
+        )
+
+        embed = HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=64)
+        head = HiggsFusedMultiTextHead(num_codebooks=8, vocab_size=1026, hidden_size=64)
+        head.weight = embed.weight
+        assert head.weight is embed.weight
+        # Modification to one should reflect in the other
+        with torch.no_grad():
+            embed.weight[0, 0] = 42.0
+        assert head.weight[0, 0] == 42.0
+
+    def test_codes_out_of_range(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            HiggsFusedMultiTextEmbedding,
+        )
+
+        embed = HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16)
+        # Code 1026 + offset for codebook 7 = 1026 + 7*1026 = 8208
+        # But weight only has 8*1026 = 8208 rows, so index 8208 is out of bounds
+        codes = torch.tensor([[0, 0, 0, 0, 0, 0, 0, 1026]])  # Last cb: 1026 is OOB
+        with pytest.raises(IndexError):
+            embed(codes)
+
+
+# ---- AC-5: Delay Pattern Behavior ----
+
+
+class TestDelayPatternBehavior:
+    """Test the delay pattern masking logic extracted from the talker."""
+
+    def test_boc_eoc_ids(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            BOC_ID,
+            EOC_ID,
+        )
+
+        assert BOC_ID == 1024
+        assert EOC_ID == 1025
+
+    def test_delay_phase_boc_masking(self):
+        """During delay phase, codebooks beyond delay_count are forced to BOC."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
+
+        num_codebooks = 8
+        # Simulate delay phase at step 3 (delay_count=2, so CBs 3-7 should be BOC)
+        codes = torch.randint(0, 1024, (num_codebooks,))
+        delay_count = 2
+        next_cb = delay_count + 1
+        if next_cb < num_codebooks:
+            codes[next_cb:] = BOC_ID
+        # CBs 0-2 should have original codes, CBs 3-7 should be BOC
+        assert all(codes[i] != BOC_ID or i >= 3 for i in range(num_codebooks))
+        assert all(codes[i] == BOC_ID for i in range(3, num_codebooks))
+
+    def test_cb0_eoc_triggers_rampdown(self):
+        """EOC on codebook 0 starts ramp-down; EOC on other codebooks does not."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            EOC_ID,
+        )
+
+        num_codebooks = 8
+        codes = torch.randint(0, 1024, (num_codebooks,))
+
+        # Simulate: cb0 emits EOC
+        codes[0] = EOC_ID
+        assert int(codes[0].item()) == EOC_ID
+        # Ramp-down should start with countdown = N-2 = 6
+        eoc_countdown = num_codebooks - 2
+        assert eoc_countdown == 6
+
+        # If cb3 emits EOC but cb0 doesn't, no ramp-down
+        codes2 = torch.randint(0, 1024, (num_codebooks,))
+        codes2[3] = EOC_ID
+        assert int(codes2[0].item()) != EOC_ID  # cb0 is not EOC
+
+    def test_rampdown_termination(self):
+        """After N-2 steps of ramp-down, generation_done becomes True."""
+        num_codebooks = 8
+        eoc_countdown = num_codebooks - 2  # = 6
+        generation_done = False
+        for _ in range(6):
+            eoc_countdown -= 1
+            if eoc_countdown <= 0:
+                generation_done = True
+        assert generation_done is True
+
+    def test_all_boc_seed(self):
+        """First audio step should seed last_codes with all-BOC."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
+
+        seeded = torch.full((8,), BOC_ID, dtype=torch.long)
+        assert seeded.shape == (8,)
+        assert all(seeded[i] == BOC_ID for i in range(8))
+
+
+# ---- AC-5: Real Sampler Method Tests ----
+
+
+class TestSamplerMethods:
+    """Test the actual _sample_audio_codes and _apply_delay_pattern_masking methods."""
+
+    def _make_minimal_talker(self):
+        """Create a minimal talker-like object with sampler/masking methods."""
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+        class FakeTalker:
+            num_codebooks = 8
+            codebook_size = 1026
+            _audio_state = {}
+
+        t = FakeTalker()
+        t._sample_audio_codes = mod.HiggsAudioV3TalkerForConditionalGeneration._sample_audio_codes.__get__(t)
+        t._apply_delay_pattern_masking = (
+            mod.HiggsAudioV3TalkerForConditionalGeneration._apply_delay_pattern_masking.__get__(t)
+        )
+        return t
+
+    def test_sample_respects_mask(self):
+        """Tokens masked to -inf must never be sampled."""
+        t = self._make_minimal_talker()
+        logits = torch.full((10, 1026), float("-inf"))
+        # Only token 500 is allowed for each row
+        logits[:, 500] = 0.0
+        result = t._sample_audio_codes(logits)
+        assert result.shape == (10,)
+        assert (result == 500).all(), f"Expected all 500, got {result.tolist()}"
+
+    def test_sample_all_masked_falls_back_to_argmax(self):
+        """All-masked row should fall back to argmax (least-negative logit)."""
+        t = self._make_minimal_talker()
+        logits = torch.full((2, 1026), float("-inf"))
+        # Row 0: all masked
+        # Row 1: only token 42 allowed
+        logits[1, 42] = 0.0
+        result = t._sample_audio_codes(logits)
+        assert result.shape == (2,)
+        assert result[1].item() == 42
+
+    def test_delay_masking_forces_boc_during_delay(self):
+        """During delay phase, codebooks beyond delay_count must have only BOC allowed."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
+
+        t = self._make_minimal_talker()
+        # Simulate delay_count=2 for batch row 0
+        t._audio_state[0] = {"num_delay": 2, "num_remaining_delays": None}
+        cb_logits = torch.zeros(1, 8, 1026)  # [1 audio row, 8 codebooks, 1026 vocab]
+        t._apply_delay_pattern_masking(cb_logits, [0])
+        # CBs 3-7 should have only BOC allowed (everything else -inf)
+        for q in range(3, 8):
+            row = cb_logits[0, q]
+            assert row[BOC_ID].item() == 0.0  # BOC kept
+            # All non-BOC should be -inf
+            non_boc = torch.cat([row[:BOC_ID], row[BOC_ID + 1 :]])
+            assert (non_boc == float("-inf")).all(), f"CB{q} has non-inf non-BOC values"
+
+    def test_delay_masking_disallows_boc_for_active_codebooks(self):
+        """Active codebooks during delay should have BOC disallowed."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
+
+        t = self._make_minimal_talker()
+        t._audio_state[0] = {"num_delay": 3, "num_remaining_delays": None}
+        cb_logits = torch.zeros(1, 8, 1026)
+        t._apply_delay_pattern_masking(cb_logits, [0])
+        # CBs 0-3 are active: BOC should be -inf
+        for q in range(4):
+            assert cb_logits[0, q, BOC_ID].item() == float("-inf")
+
+    def test_delay_masking_only_cb0_allows_eoc(self):
+        """Only codebook 0 should allow EOC during normal generation."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
+
+        t = self._make_minimal_talker()
+        t._audio_state[0] = {"num_delay": 8, "num_remaining_delays": None}
+        cb_logits = torch.zeros(1, 8, 1026)
+        t._apply_delay_pattern_masking(cb_logits, [0])
+        # CB0 should keep EOC
+        assert cb_logits[0, 0, EOC_ID].item() != float("-inf")
+        # CB1-7 should have EOC masked
+        for q in range(1, 8):
+            assert cb_logits[0, q, EOC_ID].item() == float("-inf"), f"CB{q} EOC not masked"
+
+    def test_rampdown_masking_locks_to_eoc(self):
+        """During ramp-down, locked codebooks should only allow EOC."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            BOC_ID,
+            EOC_ID,
+        )
+
+        t = self._make_minimal_talker()
+        # Ramp-down with 4 remaining delays: lock CBs 0-3 to EOC
+        t._audio_state[0] = {"num_delay": 8, "num_remaining_delays": 4}
+        cb_logits = torch.zeros(1, 8, 1026)
+        t._apply_delay_pattern_masking(cb_logits, [0])
+        # CBs 0-3 locked: only EOC allowed
+        for q in range(4):
+            row = cb_logits[0, q]
+            assert row[EOC_ID].item() == 0.0
+            non_eoc = torch.cat([row[:EOC_ID], row[EOC_ID + 1 :]])
+            assert (non_eoc == float("-inf")).all(), f"Locked CB{q} has non-inf non-EOC"
+        # CBs 4-7 active: BOC and EOC disallowed
+        for q in range(4, 8):
+            assert cb_logits[0, q, BOC_ID].item() == float("-inf")
+            assert cb_logits[0, q, EOC_ID].item() == float("-inf")
+
+
+# ---- AC-6: Feedback Method Tests ----
+
+
+class TestFeedbackMethods:
+    def test_postprocess_emits_audio_codes(self):
+        """postprocess() should return codes from _last_audio_codes."""
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+        class FakeTalker:
+            _last_audio_codes = torch.tensor([[100, 200, 300, 400, 500, 600, 700, 800]])
+            _postprocess_cursor = 0
+
+        t = FakeTalker()
+        t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
+        result = t.postprocess(torch.zeros(1, 64))
+        assert "codes" in result
+        assert "audio" in result["codes"]
+        assert result["codes"]["audio"].shape == (1, 8)
+
+    def test_postprocess_skips_negative_rows(self):
+        """postprocess() should skip rows with -1 (no audio)."""
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+        class FakeTalker:
+            _last_audio_codes = torch.tensor([[-1, -1, -1, -1, -1, -1, -1, -1]])
+            _postprocess_cursor = 0
+
+        t = FakeTalker()
+        t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
+        result = t.postprocess(torch.zeros(1, 64))
+        assert result == {}
+
+    def test_postprocess_advances_cursor(self):
+        """postprocess() should advance cursor by 1 per call."""
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+        class FakeTalker:
+            _last_audio_codes = torch.tensor(
+                [[100, 200, 300, 400, 500, 600, 700, 800], [-1, -1, -1, -1, -1, -1, -1, -1]]
+            )
+            _postprocess_cursor = 0
+
+        t = FakeTalker()
+        t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
+        r1 = t.postprocess(torch.zeros(1, 64))
+        assert "codes" in r1
+        r2 = t.postprocess(torch.zeros(1, 64))
+        assert r2 == {}  # Second row is all -1
+
+
+# ---- AC-6: Audio Feedback Method Tests ----
+
+
+class TestAudioFeedback:
+    """Test _apply_audio_feedback with minimal fake talker."""
+
+    def _make_feedback_talker(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+        embed = mod.HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16)
+        torch.nn.init.ones_(embed.weight)
+
+        class FakeTalker:
+            num_codebooks = 8
+            codebook_size = 1026
+            _audio_continuation_id = 99999  # fake audio token
+            _last_step_query_start_loc = None
+            _audio_state = {}
+            multimodal_embedding = embed
+            model = type("M", (), {"embed_tokens": lambda self, ids: torch.zeros(ids.shape[0], 16)})()
+
+        t = FakeTalker()
+        t._apply_audio_feedback = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_audio_feedback.__get__(t)
+        return t
+
+    def test_text_positions_unchanged(self):
+        """Non-audio positions should not be modified."""
+        t = self._make_feedback_talker()
+        # No audio state → all positions unchanged
+        input_ids = torch.tensor([1, 2, 3, 4])
+        hidden = torch.randn(4, 16)
+        result = t._apply_audio_feedback(hidden, input_ids)
+        assert torch.equal(result, hidden)
+
+    def test_audio_position_replaced_with_state(self):
+        """Audio position with state should have its embedding replaced."""
+        t = self._make_feedback_talker()
+        audio_id = t._audio_continuation_id
+        t._audio_state[1] = {"last_codes": torch.zeros(8, dtype=torch.long)}
+        input_ids = torch.tensor([1, audio_id, 3])
+        hidden = torch.zeros(3, 16)
+        result = t._apply_audio_feedback(hidden, input_ids)
+        # Position 1 (audio) should be non-zero (embedding of all-0 codes)
+        assert result[1].abs().sum() > 0
+        # Positions 0 and 2 should remain zero
+        assert result[0].abs().sum() == 0
+        assert result[2].abs().sum() == 0
+
+    def test_no_state_no_replacement(self):
+        """Audio position without state should not be replaced."""
+        t = self._make_feedback_talker()
+        audio_id = t._audio_continuation_id
+        # No state for row 0
+        input_ids = torch.tensor([audio_id])
+        hidden = torch.zeros(1, 16)
+        result = t._apply_audio_feedback(hidden, input_ids)
+        assert torch.equal(result, hidden)
+
+
+# ---- AC-8: Codec Strictness ----
+
+
+class TestCodecStrictness:
+    def test_bundled_missing_quantizer_key_raises(self):
+        """Bundled codec load must fail when a quantizer codebook key is missing."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
+            HiggsAudioV3Config,
+        )
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
+            HiggsAudioV3Code2Wav,
+        )
+
+        config = HiggsAudioV3Config()
+        c2w = HiggsAudioV3Code2Wav(config=config)
+
+        # State with only 7 quantizers (missing quantizer.quantizers.7.*)
+        codec_state = {}
+        for i in range(7):
+            codec_state[f"quantizer.quantizers.{i}.codebook.embed"] = torch.randn(1024, 64)
+            codec_state[f"quantizer.quantizers.{i}.project_out.weight"] = torch.randn(1024, 64)
+            codec_state[f"quantizer.quantizers.{i}.project_out.bias"] = torch.randn(1024)
+        codec_state["fc2.weight"] = torch.randn(256, 1024)
+        codec_state["fc2.bias"] = torch.randn(256)
+
+        with pytest.raises(KeyError, match="quantizer 7"):
+            c2w._load_from_bundled_state(codec_state, device=torch.device("cpu"))
+
+    def test_bundled_missing_fc2_raises(self):
+        """Bundled codec load must fail when fc2 keys are missing."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
+            HiggsAudioV3Config,
+        )
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
+            HiggsAudioV3Code2Wav,
+        )
+
+        config = HiggsAudioV3Config()
+        c2w = HiggsAudioV3Code2Wav(config=config)
+
+        codec_state = {}
+        for i in range(8):
+            codec_state[f"quantizer.quantizers.{i}.codebook.embed"] = torch.randn(1024, 64)
+            codec_state[f"quantizer.quantizers.{i}.project_out.weight"] = torch.randn(1024, 64)
+            codec_state[f"quantizer.quantizers.{i}.project_out.bias"] = torch.randn(1024)
+        # No fc2 keys
+
+        with pytest.raises(KeyError, match="fc2"):
+            c2w._load_from_bundled_state(codec_state, device=torch.device("cpu"))
+
+    def test_bundled_unknown_decoder_key_rejected(self, monkeypatch):
+        """Unknown decoder-side key should be rejected as unconsumed."""
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod
+        from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
+            HiggsAudioV3Config,
+        )
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
+            HiggsAudioV3Code2Wav,
+        )
+
+        class FakeDAC(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.dummy = torch.nn.Parameter(torch.zeros(1))
+
+        monkeypatch.setattr(c2w_mod, "build_boson_dac_decoder", lambda device: FakeDAC().to(device))
+        monkeypatch.setattr(c2w_mod, "build_higgs_audio_acoustic_decoder", lambda cfg, device: FakeDAC().to(device))
+
+        config = HiggsAudioV3Config()
+        c2w = HiggsAudioV3Code2Wav(config=config)
+
+        codec_state = _build_full_codec_state()
+        codec_state["unknown_decoder_module.weight"] = torch.randn(10)
+
+        with pytest.raises(RuntimeError, match="unexpected decoder-side keys"):
+            c2w._load_from_bundled_state(codec_state, device=torch.device("cpu"))
+
+    def test_bundled_encoder_side_keys_accepted(self, monkeypatch):
+        """Known encoder-side keys should not trigger unconsumed-key rejection."""
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_code2wav as c2w_mod
+        from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import (
+            HiggsAudioV3Config,
+        )
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import (
+            HiggsAudioV3Code2Wav,
+        )
+
+        class FakeDAC(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.dummy = torch.nn.Parameter(torch.zeros(1))
+
+        monkeypatch.setattr(c2w_mod, "build_boson_dac_decoder", lambda device: FakeDAC().to(device))
+        monkeypatch.setattr(c2w_mod, "build_higgs_audio_acoustic_decoder", lambda cfg, device: FakeDAC().to(device))
+
+        config = HiggsAudioV3Config()
+        c2w = HiggsAudioV3Code2Wav(config=config)
+
+        codec_state = _build_full_codec_state()
+        codec_state["acoustic_encoder.block.0.conv_t1.weight"] = torch.randn(10)
+        codec_state["semantic_model.encoder.layers.0.weight"] = torch.randn(10)
+        codec_state["fc.weight"] = torch.randn(10)
+
+        # Should NOT raise — encoder-side keys are allowed
+        c2w._load_from_bundled_state(codec_state, device=torch.device("cpu"))
+        assert c2w._loaded
+
+
+def _build_minimal_codec_state() -> dict[str, torch.Tensor]:
+    """Build a minimal codec state dict (quantizer + fc2, no DAC)."""
+    state: dict[str, torch.Tensor] = {}
+    for i in range(8):
+        state[f"quantizer.quantizers.{i}.codebook.embed"] = torch.randn(1024, 64)
+        state[f"quantizer.quantizers.{i}.project_out.weight"] = torch.randn(1024, 64)
+        state[f"quantizer.quantizers.{i}.project_out.bias"] = torch.randn(1024)
+    state["fc2.weight"] = torch.randn(256, 1024)
+    state["fc2.bias"] = torch.randn(256)
+    return state
+
+
+def _build_full_codec_state() -> dict[str, torch.Tensor]:
+    """Build a full codec state dict including acoustic_decoder keys.
+
+    The acoustic_decoder keys here match what a FakeDAC module expects
+    (just 'dummy'). Real DAC has many more keys.
+    """
+    state = _build_minimal_codec_state()
+    state["acoustic_decoder.dummy"] = torch.zeros(1)
+    return state
+
+
+# ---- AC-2: Loader Required-Key Validation ----
+
+
+class TestLoaderRequiredKeys:
+    """Test exact required-key validation in load_weights."""
+
+    def test_build_required_keys_count(self):
+        """Required key set should have 1 text_emb + 36*11 layers + 1 norm + 1 modality = 399."""
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            HiggsAudioV3TalkerForConditionalGeneration,
+        )
+
+        keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(36)
+        assert len(keys) == 1 + 36 * 11 + 1 + 1  # text + layers + norm + modality = 399
+
+    def test_build_required_keys_contains_norm(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            HiggsAudioV3TalkerForConditionalGeneration,
+        )
+
+        keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(2)
+        assert "body.norm.weight" in keys
+
+    def test_build_required_keys_contains_layer_subkeys(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            HiggsAudioV3TalkerForConditionalGeneration,
+        )
+
+        keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(2)
+        assert "body.layers.0.self_attn.k_proj.weight" in keys
+        assert "body.layers.1.mlp.gate_proj.weight" in keys
+
+    def test_build_required_keys_contains_modality_embedding(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            HiggsAudioV3TalkerForConditionalGeneration,
+        )
+
+        keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(1)
+        assert "tied.embedding.modality_embeddings.0.embedding.weight" in keys
+
+    def test_build_required_keys_contains_text_embedding(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
+            HiggsAudioV3TalkerForConditionalGeneration,
+        )
+
+        keys = HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(1)
+        assert "tied.embedding.text_embedding.weight" in keys
+
+
+class TestLoadWeightsEnforcement:
+    """Test load_weights() itself with synthetic key sets via a minimal fake talker."""
+
+    def _make_fake_talker_and_keys(self):
+        """Build a minimal talker-like object and the full 1-layer key set."""
+        from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
+
+        # Minimal fake talker with just enough to run load_weights
+        talker = object.__new__(mod.HiggsAudioV3TalkerForConditionalGeneration)
+        torch.nn.Module.__init__(talker)
+
+        # Attach minimal attributes load_weights needs
+        talker.multimodal_embedding = mod.HiggsFusedMultiTextEmbedding(8, 1026, 16)
+        talker.modality_head = mod.HiggsFusedMultiTextHead(8, 1026, 16)
+        talker.tie_modality = True
+        talker.modality_head.weight = talker.multimodal_embedding.weight
+
+        # Fake backbone config with 1 layer
+        class FakeConfig:
+            num_hidden_layers = 1
+            tie_word_embeddings = True
+
+        talker._backbone_config = FakeConfig()
+
+        # Minimal model and lm_head so _BackboneWrapper can be constructed
+        talker.model = torch.nn.Module()
+        talker.lm_head = torch.nn.Module()
+
+        # Fake vllm_config for resolve_special_tokens (no-op)
+        class FakeModelConfig:
+            model = None
+
+        class FakeVllmConfig:
+            model_config = FakeModelConfig()
+
+        talker.vllm_config = FakeVllmConfig()
+        talker.config = type("C", (), {"resolve_special_tokens": lambda self, p: None})()
+        talker._resolved_tokens = True
+        talker._audio_continuation_id = None
+        talker._eos_token_id = None
+
+        # Build the full required key set for 1 layer
+        required = mod.HiggsAudioV3TalkerForConditionalGeneration._build_required_keys(1)
+
+        # Build synthetic weights: tiny tensors for every required key
+        weights = {}
+        for key in required:
+            if "modality_embeddings.0.embedding" in key:
+                weights[key] = torch.randn(8 * 1026, 16)
+            else:
+                weights[key] = torch.randn(4)  # tiny placeholder
+
+        # Monkeypatch _BackboneWrapper.load_weights to accept and return keys
+        def fake_backbone_load(self_wrapper, ws):
+            consumed = set()
+            for name, _ in ws:
+                consumed.add(name)
+            return consumed
+
+        talker._fake_backbone_load = fake_backbone_load
+        return talker, weights, mod, fake_backbone_load
+
+    def test_load_weights_succeeds_with_all_keys(self, monkeypatch):
+        """load_weights succeeds when all required keys are present."""
+        talker, weights, mod, fake_load = self._make_fake_talker_and_keys()
+        monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load)
+        result = talker.load_weights(iter(weights.items()))
+        assert isinstance(result, set)
+        assert "multimodal_embedding.weight" in result
+
+    def test_load_weights_fails_missing_norm(self, monkeypatch):
+        """load_weights raises when body.norm.weight is missing."""
+        talker, weights, mod, fake_load = self._make_fake_talker_and_keys()
+        monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load)
+        del weights["body.norm.weight"]
+        with pytest.raises(RuntimeError, match="body.norm.weight"):
+            talker.load_weights(iter(weights.items()))
+
+    def test_load_weights_fails_missing_layer_subkey(self, monkeypatch):
+        """load_weights raises when a specific layer subkey is missing."""
+        talker, weights, mod, fake_load = self._make_fake_talker_and_keys()
+        monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load)
+        del weights["body.layers.0.self_attn.k_proj.weight"]
+        with pytest.raises(RuntimeError, match="body.layers.0.self_attn.k_proj.weight"):
+            talker.load_weights(iter(weights.items()))
+
+    def test_load_weights_fails_missing_modality_embedding(self, monkeypatch):
+        """load_weights raises when modality embedding is missing."""
+        talker, weights, mod, fake_load = self._make_fake_talker_and_keys()
+        monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load)
+        del weights["tied.embedding.modality_embeddings.0.embedding.weight"]
+        with pytest.raises(RuntimeError, match="modality_embeddings"):
+            talker.load_weights(iter(weights.items()))
+
+    def test_load_weights_fails_unknown_prefix(self, monkeypatch):
+        """load_weights raises on unknown non-codec checkpoint prefix."""
+        talker, weights, mod, fake_load = self._make_fake_talker_and_keys()
+        monkeypatch.setattr(mod._BackboneWrapper, "load_weights", fake_load)
+        weights["unknown.prefix.weight"] = torch.randn(4)
+        with pytest.raises(ValueError, match="Unexpected checkpoint key"):
+            talker.load_weights(iter(weights.items()))
+
+
+# ---- AC-7: Stage Input Processor ----
+
+
+class TestStageInputProcessor:
+    def test_revert_delay_pattern(self):
+        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+            _revert_delay_pattern,
+        )
+
+        # 8 codebooks, 3 real frames -> delayed shape [8, 3+8-1=10]
+        Q, T = 8, 3
+        delayed = torch.full((Q, T + Q - 1), 1024)  # Fill with BOC
+        for i in range(Q):
+            for t in range(T):
+                delayed[i, i + t] = i * 100 + t  # Real codes at shifted positions
+        result = _revert_delay_pattern(delayed)
+        assert result.shape == (Q, T)
+        for i in range(Q):
+            for t in range(T):
+                assert result[i, t].item() == i * 100 + t
+
+    def test_revert_delay_pattern_rejects_wrong_codebooks(self):
+        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+            _revert_delay_pattern,
+        )
+
+        # 7 codebooks should be rejected
+        codes = torch.zeros(7, 20)
+        with pytest.raises(ValueError, match="Expected exactly 8 codebook rows"):
+            _revert_delay_pattern(codes)
+
+    def test_revert_delay_pattern_rejects_too_few_frames(self):
+        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+            _revert_delay_pattern,
+        )
+
+        # 8 codebooks but only 5 frames (need at least 8)
+        codes = torch.zeros(8, 5)
+        with pytest.raises(ValueError, match="Not enough frames"):
+            _revert_delay_pattern(codes)
+
+    def test_filter_real_code_frames(self):
+        from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
+            _filter_real_code_frames,
+        )
+
+        # 8 codebooks, 4 frames
+        codes = torch.tensor(
+            [
+                [100, 200, 1024, 300],  # cb0: frame 2 has BOC
+                [101, 201, 1024, 301],
+                [102, 202, 1024, 302],
+                [103, 203, 1024, 303],
+                [104, 204, 1024, 304],
+                [105, 205, 1024, 305],
+                [106, 206, 1024, 306],
+                [107, 207, 1024, 307],
+            ]
+        )
+        result = _filter_real_code_frames(codes)
+        # Frame 2 (column 2) has BOC in all codebooks -> filtered out
+        assert result.shape == (8, 3)
+        assert result[0, 0].item() == 100
+        assert result[0, 1].item() == 200
+        assert result[0, 2].item() == 300
+
+
+# ---- AC-10: Registry ----
+
+
+class TestRegistry:
+    def test_talker_registered(self):
+        from vllm_omni.model_executor.models.registry import _OMNI_MODELS
+
+        assert "HiggsMultimodalQwen3ForConditionalGeneration" in _OMNI_MODELS
+        assert "HiggsAudioV3TalkerForConditionalGeneration" in _OMNI_MODELS
+
+    def test_code2wav_registered(self):
+        from vllm_omni.model_executor.models.registry import _OMNI_MODELS
+
+        assert "HiggsAudioV3Code2WavForConditionalGeneration" in _OMNI_MODELS
+
+    def test_pipeline_registered(self):
+        from vllm_omni.config.pipeline_registry import _OMNI_PIPELINES
+
+        assert "higgs_multimodal_qwen3" in _OMNI_PIPELINES
+
+    def test_deploy_yaml_exists(self):
+        import os
+
+        # __file__ = tests/unit/higgs_audio_v3/test_*.py → 4 dirnames to repo root
+        repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+        yaml_path = os.path.join(repo_root, "vllm_omni", "deploy", "higgs_multimodal_qwen3.yaml")
+        assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}"
+
+
+# ---- AC-3: Prompt Builder ----
+
+
+class TestPromptBuilder:
+    def _make_mock_tokenizer(self):
+        """Create a mock tokenizer with the required special tokens."""
+
+        class MockTokenizer:
+            def __init__(self):
+                self._added_vocab = {
+                    "<|tts|>": 151700,
+                    "<|text|>": 151701,
+                    "<|audio|>": 151702,
+                    "<|ref_audio|>": 151703,
+                    "<|ref_text|>": 151704,
+                }
+
+            def get_added_vocab(self):
+                return self._added_vocab
+
+            def encode(self, text, add_special_tokens=True):
+                # Simple word-level tokenization for testing
+                return list(range(100, 100 + len(text.split())))
+
+        return MockTokenizer()
+
+    def test_plain_tts_prompt(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
+            HiggsAudioV3TokenizerAdapter,
+        )
+
+        tok = self._make_mock_tokenizer()
+        adapter = HiggsAudioV3TokenizerAdapter(tok)
+        ids = adapter.build_prompt("Hello world")
+        assert ids[0] == 151700  # <|tts|>
+        assert ids[1] == 151701  # <|text|>
+        assert ids[-1] == 151702  # <|audio|>
+        assert len(ids) == 2 + 2 + 1  # tts + text + 2 word tokens + audio
+
+    def test_empty_text_rejected(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
+            HiggsAudioV3TokenizerAdapter,
+        )
+
+        tok = self._make_mock_tokenizer()
+        adapter = HiggsAudioV3TokenizerAdapter(tok)
+        with pytest.raises(ValueError, match="non-empty"):
+            adapter.build_prompt("")
+
+    def test_whitespace_only_rejected(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
+            HiggsAudioV3TokenizerAdapter,
+        )
+
+        tok = self._make_mock_tokenizer()
+        adapter = HiggsAudioV3TokenizerAdapter(tok)
+        with pytest.raises(ValueError, match="non-empty"):
+            adapter.build_prompt("   ")
+
+    def test_missing_specials_rejected(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
+            HiggsAudioV3TokenizerAdapter,
+        )
+
+        class BadTokenizer:
+            def get_added_vocab(self):
+                return {"<|tts|>": 1}  # Missing <|text|> and <|audio|>
+
+        with pytest.raises(ValueError, match="missing"):
+            HiggsAudioV3TokenizerAdapter(BadTokenizer())
+
+    def test_no_voice_clone_tokens_in_plain_tts(self):
+        from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
+            HiggsAudioV3TokenizerAdapter,
+        )
+
+        tok = self._make_mock_tokenizer()
+        adapter = HiggsAudioV3TokenizerAdapter(tok)
+        ids = adapter.build_prompt("Hello")
+        # Should not contain ref_audio or ref_text token IDs
+        assert 151703 not in ids  # <|ref_audio|>
+        assert 151704 not in ids  # <|ref_text|>

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/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/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/diffusion_engine.py

@@ -96,6 +96,8 @@ def _move_tensor_tree_to_cpu(value: object) -> object:
 
 
 def get_dummy_run_num_frames(model_class_name: str, supports_audio_input: bool) -> int:
+    """Get num_frames for the dummy warmup run. Returns 0 to skip warmup."""
+
     model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
     if model_cls is not None and hasattr(model_cls, "dummy_run_num_frames"):
         return int(getattr(model_cls, "dummy_run_num_frames"))
@@ -722,6 +724,9 @@ class DiffusionEngine:
             prompt.setdefault("multi_modal_data", {})["audio"] = dummy_audio
 
         num_frames = get_dummy_run_num_frames(self.od_config.model_class_name, supports_audio_input)
+        if num_frames <= 0:
+            logger.info("Skipping dummy warmup run (num_frames=0)")
+            return
         req = OmniDiffusionRequest(
             prompts=[prompt],
             request_id=DUMMY_DIFFUSION_REQUEST_ID,

vllm_omni/diffusion/inline_stage_diffusion_client.py

@@ -52,6 +52,7 @@ class InlineStageDiffusionClient(StageClientBase):
         self.replica_id = metadata.replica_id
         self.final_output = metadata.final_output
         self.final_output_type = metadata.final_output_type
+        self.model_stage = getattr(metadata, "model_stage", None)
         self.default_sampling_params = metadata.default_sampling_params
         self.requires_multimodal_data = metadata.requires_multimodal_data
         self.custom_process_input_func = metadata.custom_process_input_func

vllm_omni/diffusion/layers/rope.py

@@ -5,6 +5,7 @@ from einops import rearrange, repeat
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.layers.custom_op import CustomOp
+from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
@@ -300,3 +301,234 @@ def apply_rope_to_qk(
         query = rope(query, cos, sin)
         key = rope(key, cos, sin)
     return query, key
+
+
+class WanS2VRotaryPosEmbed(torch.nn.Module):
+    """Precompute complex-valued RoPE embeddings for S2V multi-grid positions.
+
+    Owns the base frequency buffer and provides forward() to compute position
+    embeddings given hidden_states (for shape) and grid_sizes.
+    """
+
+    def __init__(self, num_heads: int, head_dim: int, max_seq_len: int = 1024):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        d = head_dim
+        freqs = torch.cat(
+            [
+                self._rope_params(max_seq_len, d - 4 * (d // 6)),
+                self._rope_params(max_seq_len, 2 * (d // 6)),
+                self._rope_params(max_seq_len, 2 * (d // 6)),
+            ],
+            dim=1,
+        )
+        self.register_buffer("freqs", freqs.to(torch.complex64), persistent=False)
+
+    @staticmethod
+    @torch.amp.autocast(current_omni_platform.device_type, enabled=False)
+    def _rope_params(max_seq_len, dim, theta=10000):
+        if dim % 2 != 0:
+            raise ValueError(f"dim ({dim}) must be even")
+        freqs = torch.outer(
+            torch.arange(max_seq_len), 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim))
+        )
+        return torch.polar(torch.ones_like(freqs), freqs)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        grid_sizes: list,
+        trainable_freqs: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Precompute RoPE embeddings for the given grid layout.
+
+        Args:
+            hidden_states: Tensor [B, S, ...] (used for batch/seq shape and device)
+            grid_sizes: Grid specification (list of [offsets, sizes, totals])
+            trainable_freqs: Optional trainable frequency overrides for t_f < 0
+
+        Returns:
+            Complex tensor [B, S, 1, head_dim//2] of precomputed position embeddings
+        """
+        b, s = hidden_states.shape[0], hidden_states.shape[1]
+        c = self.head_dim // 2
+        device = hidden_states.device
+
+        freqs = self.freqs.to(device)
+        if trainable_freqs is not None:
+            freqs_input = [freqs, trainable_freqs]
+        else:
+            freqs_input = freqs
+
+        if isinstance(freqs_input, list):
+            trainable_f = freqs_input[1]
+            freqs_split = freqs_input[0]
+        else:
+            trainable_f = None
+            freqs_split = freqs_input
+        freqs_split = freqs_split.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+        output = torch.empty((b, s, 1, c), device=device, dtype=torch.complex64)
+        seq_bucket = [0]
+        if not isinstance(grid_sizes, list):
+            grid_sizes = [grid_sizes]
+        for g in grid_sizes:
+            if not isinstance(g, list):
+                g = [torch.zeros_like(g), g]
+            batch_size = g[0].shape[0]
+            for i in range(batch_size):
+                f_o, h_o, w_o = g[0][i]
+                f, h, w = g[1][i]
+                t_f, t_h, t_w = g[2][i]
+                seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o
+                seq_len = int(seq_f * seq_h * seq_w)
+                if seq_len > 0:
+                    if t_f > 0:
+                        assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0
+                        seq_f_int = int(seq_f)
+                        seq_h_int = int(seq_h)
+                        seq_w_int = int(seq_w)
+
+                        if f_o >= 0:
+                            f_sam = torch.linspace(int(f_o), int(t_f + f_o) - 1, seq_f_int, device=device).long()
+                        else:
+                            f_sam = torch.linspace(int(-f_o), int(-t_f - f_o) + 1, seq_f_int, device=device).long()
+                        h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long()
+                        w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long()
+
+                        freqs_0 = freqs_split[0][f_sam] if f_o >= 0 else freqs_split[0][f_sam].conj()
+                        freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1)
+
+                        freqs_i = torch.cat(
+                            [
+                                freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                freqs_split[1][h_sam]
+                                .view(1, seq_h_int, 1, -1)
+                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                freqs_split[2][w_sam]
+                                .view(1, 1, seq_w_int, -1)
+                                .expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                            ],
+                            dim=-1,
+                        ).reshape(seq_len, 1, -1)
+                    elif t_f < 0:
+                        freqs_i = trainable_f.unsqueeze(1)
+                    output[i, seq_bucket[-1] : seq_bucket[-1] + seq_len] = freqs_i
+            seq_bucket.append(seq_bucket[-1] + seq_len)
+        return output
+
+
+class RotaryEmbeddingWanS2V(RotaryEmbeddingWan):
+    """Apply RoPE using precomputed complex freqs for Wan S2V main transformer.
+
+    Converts complex freqs (from WanS2VRotaryPosEmbed) to cos/sin and delegates
+    to RotaryEmbeddingWan for platform-optimized application (float32 kernel).
+    Under TP, freqs has 1 head — broadcasts automatically via cos/sin.
+    """
+
+    def __init__(self) -> None:
+        super().__init__(is_neox_style=False, half_head_dim=True)
+
+    def forward(self, x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor:
+        freqs_sliced = freqs[:, : x.size(1)]
+        cos = freqs_sliced.real.to(x.dtype)
+        sin = freqs_sliced.imag.to(x.dtype)
+        return super().forward(x, cos, sin)
+
+
+class RotaryEmbeddingS2VGrid(torch.nn.Module):
+    """Grid-based RoPE for S2V motioner/init attention.
+
+    Applies complex-valued rotary embeddings using 3D grid sampling
+    (frame, height, width). Used by SimpleSelfAttention, SwinSelfAttention,
+    CausalSelfAttention in motioner blocks.
+    """
+
+    @staticmethod
+    @torch.amp.autocast(current_omni_platform.device_type, enabled=False)
+    def precompute(
+        seq_len: int, num_heads: int, head_dim: int, grid_sizes, freqs: torch.Tensor, device: torch.device, start=None
+    ) -> torch.Tensor:
+        """Precompute position frequency tensor from grid specification.
+
+        Returns a complex tensor that can be reused across layers via apply_precomputed().
+        """
+        c = head_dim // 2
+
+        trainable_freqs = None
+        if isinstance(freqs, list):
+            trainable_freqs = freqs[1]
+            freqs = freqs[0]
+        freqs = freqs.to(device)
+        freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+        if not isinstance(grid_sizes, list):
+            grid_sizes = [grid_sizes]
+
+        batch_size = grid_sizes[0][0].shape[0] if isinstance(grid_sizes[0], list) else grid_sizes[0].shape[0]
+        precomputed = torch.empty((batch_size, seq_len, 1, c), device=device, dtype=torch.complex64)
+        seq_bucket = [0]
+
+        for g in grid_sizes:
+            if not isinstance(g, list):
+                g = [torch.zeros_like(g), g]
+            g_batch_size = g[0].shape[0]
+            for i in range(g_batch_size):
+                if start is None:
+                    f_o, h_o, w_o = g[0][i]
+                else:
+                    f_o, h_o, w_o = start[i]
+
+                f, h, w = g[1][i]
+                t_f, t_h, t_w = g[2][i]
+                seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o
+                seg_len = int(seq_f * seq_h * seq_w)
+                if seg_len > 0:
+                    if t_f > 0:
+                        seq_f_int = int(seq_f)
+                        seq_h_int = int(seq_h)
+                        seq_w_int = int(seq_w)
+
+                        if f_o >= 0:
+                            f_sam = torch.linspace(int(f_o), int(t_f + f_o) - 1, seq_f_int, device=device).long()
+                        else:
+                            f_sam = torch.linspace(int(-f_o), int(-t_f - f_o) + 1, seq_f_int, device=device).long()
+                        h_sam = torch.linspace(int(h_o), int(t_h + h_o) - 1, seq_h_int, device=device).long()
+                        w_sam = torch.linspace(int(w_o), int(t_w + w_o) - 1, seq_w_int, device=device).long()
+
+                        freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][f_sam].conj()
+                        freqs_0 = freqs_0.view(seq_f_int, 1, 1, -1)
+
+                        freqs_i = torch.cat(
+                            [
+                                freqs_0.expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                freqs[1][h_sam].view(1, seq_h_int, 1, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                                freqs[2][w_sam].view(1, 1, seq_w_int, -1).expand(seq_f_int, seq_h_int, seq_w_int, -1),
+                            ],
+                            dim=-1,
+                        ).reshape(seg_len, 1, -1)
+                    elif t_f < 0:
+                        freqs_i = trainable_freqs.unsqueeze(1)
+                    precomputed[i, seq_bucket[-1] : seq_bucket[-1] + seg_len] = freqs_i
+            seq_bucket.append(seq_bucket[-1] + seg_len)
+        return precomputed
+
+    @staticmethod
+    @torch.amp.autocast(current_omni_platform.device_type, enabled=False)
+    def apply_precomputed(x: torch.Tensor, precomputed_freqs: torch.Tensor) -> torch.Tensor:
+        """Apply precomputed position frequencies to input tensor."""
+        n = x.size(2)
+        input_dtype = x.dtype
+        seq_len = x.size(1)
+        precomputed_freqs = precomputed_freqs[:, :seq_len]
+        x_c = torch.view_as_complex(x.to(torch.float64).reshape(x.size(0), seq_len, n, -1, 2))
+        x_c = torch.view_as_real(x_c * precomputed_freqs).flatten(3)
+        return x_c.reshape(x.size(0), seq_len, n, -1).to(input_dtype)
+
+    @staticmethod
+    @torch.amp.autocast(current_omni_platform.device_type, enabled=False)
+    def forward(x: torch.Tensor, grid_sizes, freqs: torch.Tensor, start=None) -> torch.Tensor:
+        n, d = x.size(2), x.size(3)
+        precomputed = RotaryEmbeddingS2VGrid.precompute(x.size(1), n, d, grid_sizes, freqs, x.device, start=start)
+        return RotaryEmbeddingS2VGrid.apply_precomputed(x, precomputed)

vllm_omni/diffusion/model_loader/hub_prefetch.py

@@ -65,10 +65,22 @@ import errno
 import logging
 import os
 import time
-from collections.abc import Iterable, Iterator
+from collections.abc import Callable, Iterable, Iterator
+from typing import Any
 
 logger = logging.getLogger(__name__)
 
+# A racy / partially-evicted HF cache (the exact failure this module defends
+# against) is transient: re-running ``snapshot_download`` blocks on the peer
+# writer's per-blob ``.lock`` and then returns a complete tree. So a bounded
+# retry with linear backoff is what actually closes the window that a single
+# best-effort attempt left open (Buildkite vllm-omni-rebase #1858: both the
+# ``cuda_ti2v_hsdp`` missing-shard ``OSError`` and the ``wan_2_1_vace`` default
+# ``UMT5Config`` size-mismatch were a swallowed prefetch followed by a
+# ``from_pretrained`` against the half-written cache).
+_PREFETCH_MAX_ATTEMPTS = 3
+_PREFETCH_BACKOFF_BASE_S = 1.0
+
 
 def _node_lock_dir() -> str:
     """Return a node-local directory suitable for the prefetch lock files.
@@ -276,7 +288,11 @@ def prefetch_subfolders(
     if local_files_only or not model or os.path.isdir(model):
         return
 
-    logger.info("Prefetching %s subfolders: %s", model, list(subfolders))
+    # Materialise ``subfolders`` up-front: it may be a one-shot generator and
+    # we reference it again in the retry / logging paths below.
+    subfolders = list(subfolders)
+
+    logger.info("Prefetching %s subfolders: %s", model, subfolders)
 
     try:
         from huggingface_hub import snapshot_download
@@ -306,39 +322,65 @@ def prefetch_subfolders(
     # snapshot is now warm. This is what makes the prefetch race-free
     # even when many ``DiffusionWorker`` subprocesses (or multiple
     # OmniServer instances on the same node) hit this code in parallel.
-    try:
-        with _repo_prefetch_lock(model):
-            snapshot_download(
-                repo_id=model,
-                allow_patterns=allow_patterns,
-            )
-        logger.info("Prefetch complete for %s", model)
-    except Exception as exc:
-        # Best-effort: propagate only via logging. The subsequent
-        # ``from_pretrained`` call will raise a clearer, call-site-specific
-        # error (auth, 404, disk full, ...) that we'd rather surface - EXCEPT
-        # for auth/gating, which we escalate here with an explicit hint so
-        # readers of CI logs don't have to correlate the generic "OSError:
-        # <repo> does not appear to have a file named ..." that
-        # ``from_pretrained`` would otherwise emit much later with an
-        # unrelated-looking message.
-        if _looks_like_auth_error(exc):
-            logger.error(
-                "Hub prefetch for '%s' failed with an authentication / gated "
-                "repository error (%s: %s). The CI HF_TOKEN must (1) be set "
-                "in the step env, (2) be valid, and (3) belong to an account "
-                "that has accepted the model license on huggingface.co. See "
-                "docs/contributing/ci/hf_credentials.md.",
-                model,
-                type(exc).__name__,
-                exc,
-            )
-        else:
+    # A single best-effort attempt is not enough: when several diffusion
+    # workers race a cold cache (and the node-wide lock fails to serialise
+    # them, as observed for HSDP / ring launches), ``snapshot_download`` can
+    # raise on a half-written tree. Swallowing that and proceeding straight
+    # to ``from_pretrained`` is exactly what turned a recoverable prefetch
+    # hiccup into a hard server crash. Retry with backoff so the snapshot
+    # actually completes before any loader reads the cache.
+    for attempt in range(1, _PREFETCH_MAX_ATTEMPTS + 1):
+        try:
+            with _repo_prefetch_lock(model):
+                snapshot_download(
+                    repo_id=model,
+                    allow_patterns=allow_patterns,
+                )
+            logger.info("Prefetch complete for %s", model)
+            return
+        except Exception as exc:
+            # Auth / gating never heals on retry - escalate immediately with
+            # an explicit hint so readers of CI logs don't have to correlate
+            # the generic "OSError: <repo> does not appear to have a file
+            # named ..." that ``from_pretrained`` would otherwise emit later.
+            if _looks_like_auth_error(exc):
+                logger.error(
+                    "Hub prefetch for '%s' failed with an authentication / gated "
+                    "repository error (%s: %s). The CI HF_TOKEN must (1) be set "
+                    "in the step env, (2) be valid, and (3) belong to an account "
+                    "that has accepted the model license on huggingface.co. See "
+                    "docs/contributing/ci/hf_credentials.md.",
+                    model,
+                    type(exc).__name__,
+                    exc,
+                )
+                return
+
+            if attempt < _PREFETCH_MAX_ATTEMPTS:
+                backoff = _PREFETCH_BACKOFF_BASE_S * attempt
+                logger.warning(
+                    "Hub prefetch for repo '%s' subfolders %s failed on attempt %d/%d (%s: %s); retrying in %.1fs",
+                    model,
+                    subfolders,
+                    attempt,
+                    _PREFETCH_MAX_ATTEMPTS,
+                    type(exc).__name__,
+                    exc,
+                    backoff,
+                )
+                time.sleep(backoff)
+                continue
+
+            # Exhausted retries. Stay best-effort: propagate only via logging
+            # so the subsequent ``from_pretrained`` call surfaces the real,
+            # call-site-specific error (and ``from_pretrained_with_prefetch``
+            # gets a final chance to heal the cache).
             logger.warning(
-                "Hub prefetch for repo '%s' subfolders %s failed (%s: %s); "
-                "falling back to on-demand download in from_pretrained",
+                "Hub prefetch for repo '%s' subfolders %s failed after %d attempts "
+                "(%s: %s); falling back to on-demand download in from_pretrained",
                 model,
-                list(subfolders),
+                subfolders,
+                _PREFETCH_MAX_ATTEMPTS,
                 type(exc).__name__,
                 exc,
             )
@@ -376,6 +418,78 @@ def _looks_like_auth_error(exc: BaseException) -> bool:
     return "401 client error" in msg or "403 client error" in msg or "gatedrepo" in msg
 
 
+def from_pretrained_with_prefetch(
+    factory: Callable[..., Any],
+    model: str,
+    *,
+    subfolder: str,
+    prefetch_list: Iterable[str],
+    local_files_only: bool = False,
+    max_attempts: int = _PREFETCH_MAX_ATTEMPTS,
+    **from_pretrained_kwargs: Any,
+) -> Any:
+    """Call ``factory.from_pretrained`` healing a racy / partial HF cache.
+
+    ``factory`` is a bound ``SomeModel.from_pretrained`` (or any callable with
+    the same ``(model, *, subfolder, local_files_only, **kwargs)`` signature).
+
+    This is a stronger sibling of :func:`retry_on_missing_shard`: that helper
+    only retries the missing-shard ``OSError`` and never re-prefetches, so it
+    cannot recover the second face of the same race. Two shapes of partial
+    -cache failure crash the diffusion server outright:
+
+    * ``OSError: <repo> does not appear to have a file named
+      text_encoder/model-0000X-of-0000Y.safetensors`` - a shard is still under
+      its ``.incomplete`` name.
+    * ``RuntimeError: You set 'ignore_mismatched_sizes' to 'False' ...`` -
+      ``text_encoder/config.json`` was not present yet, so ``transformers`` v5
+      silently fell back to the default (tiny) config and then could not load
+      the real checkpoint into it.
+
+    Both heal once the cache is complete. So on those errors we re-run a
+    *verified* prefetch (which blocks on the peer writer and retries the
+    download) and reload, instead of letting the worker die. Local paths and
+    ``local_files_only`` loads cannot be healed by re-fetching, so they raise
+    on the first failure exactly as before.
+    """
+    prefetch_list = list(prefetch_list)
+    can_heal = not local_files_only and bool(model) and not os.path.isdir(model)
+    last_exc: BaseException | None = None
+
+    for attempt in range(1, max_attempts + 1):
+        try:
+            return factory(
+                model,
+                subfolder=subfolder,
+                local_files_only=local_files_only,
+                **from_pretrained_kwargs,
+            )
+        except (OSError, RuntimeError, ValueError) as exc:
+            last_exc = exc
+            if not can_heal or attempt >= max_attempts:
+                break
+            backoff = _PREFETCH_BACKOFF_BASE_S * attempt
+            logger.warning(
+                "from_pretrained(%s, subfolder=%s) failed on attempt %d/%d "
+                "(%s: %s); re-prefetching repo and retrying in %.1fs",
+                model,
+                subfolder,
+                attempt,
+                max_attempts,
+                type(exc).__name__,
+                exc,
+                backoff,
+            )
+            time.sleep(backoff)
+            # Force a fresh, verified snapshot of every component this pipeline
+            # needs - not just ``subfolder`` - so a sibling component that was
+            # also half-written gets repaired in the same pass.
+            prefetch_subfolders(model, prefetch_list, local_files_only=False)
+
+    assert last_exc is not None  # loop only exits via return or a caught exc
+    raise last_exc
+
+
 def retry_on_missing_shard(load_fn, *, max_retries: int = 3, base_delay: float = 5.0):
     """Call *load_fn* with retry on the transformers v5 shard-resolution race.
 

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

@@ -777,7 +777,9 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
             if torch.is_tensor(v):
                 generation_input[k] = v.to(self.device)
 
-        self._regen_init_noise_on_device(generation_input, req.sampling_params.seed)
+        # NOTE: For now we disable device specific noise regeneration so that e2e tests can run
+        # on both CUDA and ROCm. Context: https://github.com/vllm-project/vllm-omni/pull/4081
+        # self._regen_init_noise_on_device(generation_input, req.sampling_params.seed)
 
         # text cfg
         generation_input_cfg_text = self.bagel.prepare_vae_latent_cfg(

vllm_omni/diffusion/models/cosmos3/action.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Action-token helpers for Cosmos3 UVA/action generation."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+
+ACTION_MODE_POLICY = "policy"
+ACTION_MODE_FORWARD_DYNAMICS = "forward_dynamics"
+ACTION_MODE_INVERSE_DYNAMICS = "inverse_dynamics"
+ACTION_MODES = {
+    ACTION_MODE_POLICY,
+    ACTION_MODE_FORWARD_DYNAMICS,
+    ACTION_MODE_INVERSE_DYNAMICS,
+}
+
+
+EMBODIMENT_TO_DOMAIN_ID: dict[str, int] = {
+    "no_action": 0,
+    "av": 1,
+    "camera_pose": 2,
+    "hand_pose": 3,
+    "pusht": 4,
+    "libero": 5,
+    "umi": 6,
+    "bridge_orig_lerobot": 7,
+    "droid_lerobot": 8,
+    "robomind-franka": 8,
+    "galbot": 9,
+    "robomind-franka-dual": 12,
+    "robomind-ur": 13,
+    "agibotworld": 15,
+    "agibot_gear_gripper": 15,
+    "agibot_gear_gripper_ext": 15,
+    "fractal": 20,
+}
+
+
+VIDEO_RES_SIZE_INFO: dict[str, dict[str, tuple[int, int]]] = {
+    "256": {
+        "1,1": (256, 256),
+        "4,3": (320, 256),
+        "3,4": (256, 320),
+        "16,9": (320, 192),
+        "9,16": (192, 320),
+    },
+    "480": {
+        "1,1": (640, 640),
+        "4,3": (736, 544),
+        "3,4": (544, 736),
+        "16,9": (832, 480),
+        "9,16": (480, 832),
+    },
+    "704": {
+        "1,1": (960, 960),
+        "4,3": (1088, 832),
+        "3,4": (832, 1088),
+        "16,9": (1280, 704),
+        "9,16": (704, 1280),
+    },
+    "720": {
+        "1,1": (960, 960),
+        "4,3": (1104, 832),
+        "3,4": (832, 1104),
+        "16,9": (1280, 720),
+        "9,16": (720, 1280),
+    },
+}
+
+
+def normalize_action_mode(mode: Any) -> str | None:
+    if mode is None:
+        return None
+    normalized = str(mode).strip().lower()
+    if not normalized:
+        return None
+    if normalized not in ACTION_MODES:
+        raise ValueError(f"Unsupported Cosmos3 action_mode={mode!r}; expected one of {sorted(ACTION_MODES)}.")
+    return normalized
+
+
+def resolve_domain_id(
+    *,
+    domain_id: Any = None,
+    domain_name: Any = None,
+    require_explicit: bool = False,
+) -> int:
+    if domain_id is not None:
+        resolved = int(domain_id)
+        if resolved < 0:
+            raise ValueError(f"Cosmos3 domain_id must be non-negative, got {resolved}.")
+        return resolved
+
+    if domain_name is None or str(domain_name).strip() == "":
+        if require_explicit:
+            raise ValueError(
+                "Cosmos3 action generation requires extra_args['domain_id'] or non-empty extra_args['domain_name']."
+            )
+        return 0
+
+    key = str(domain_name).strip().lower()
+    if key not in EMBODIMENT_TO_DOMAIN_ID:
+        raise ValueError(
+            f"Unknown Cosmos3 action domain_name={domain_name!r}; "
+            f"expected one of {sorted(EMBODIMENT_TO_DOMAIN_ID)} or pass domain_id directly."
+        )
+    return EMBODIMENT_TO_DOMAIN_ID[key]
+
+
+def action_condition_indexes(mode: str, action_length: int) -> list[int]:
+    mode = normalize_action_mode(mode)
+    if mode == ACTION_MODE_FORWARD_DYNAMICS:
+        return list(range(action_length))
+    if mode in {ACTION_MODE_POLICY, ACTION_MODE_INVERSE_DYNAMICS}:
+        return []
+    raise AssertionError(f"Unexpected action mode: {mode!r}")
+
+
+def vision_condition_indexes(mode: str, video_length: int, temporal_compression_factor: int) -> list[int]:
+    mode = normalize_action_mode(mode)
+    latent_frames = (video_length - 1) // temporal_compression_factor + 1
+    if mode in {ACTION_MODE_POLICY, ACTION_MODE_FORWARD_DYNAMICS}:
+        return [0]
+    if mode == ACTION_MODE_INVERSE_DYNAMICS:
+        return list(range(latent_frames))
+    raise AssertionError(f"Unexpected action mode: {mode!r}")
+
+
+def action_start_frame_offset(mode: str, action_length: int, video_length: int) -> int:
+    del mode
+    if action_length == video_length - 1:
+        return 1
+    if action_length == video_length:
+        return 0
+    raise ValueError(
+        "Cosmos3 action_chunk_size must equal num_frames - 1 or num_frames; "
+        f"got action_chunk_size={action_length}, num_frames={video_length}."
+    )
+
+
+def build_action_condition_mask(
+    mode: str,
+    action_length: int,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    mask = torch.zeros(1, action_length, 1, device=device, dtype=dtype)
+    for idx in action_condition_indexes(mode, action_length):
+        mask[:, idx, :] = 1.0
+    return mask
+
+
+def build_vision_condition_mask(
+    mode: str,
+    video_length: int,
+    temporal_compression_factor: int,
+    *,
+    device: torch.device,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    latent_frames = (video_length - 1) // temporal_compression_factor + 1
+    mask = torch.zeros(1, 1, latent_frames, 1, 1, device=device, dtype=dtype)
+    for idx in vision_condition_indexes(mode, video_length, temporal_compression_factor):
+        mask[:, :, idx, :, :] = 1.0
+    return mask
+
+
+def pad_action_to_dim(action: torch.Tensor, action_dim: int) -> torch.Tensor:
+    if action.shape[-1] > action_dim:
+        raise ValueError(f"Cosmos3 action dimension {action.shape[-1]} exceeds model action_dim={action_dim}.")
+    if action.shape[-1] == action_dim:
+        return action
+    padding = torch.zeros(*action.shape[:-1], action_dim - action.shape[-1], dtype=action.dtype, device=action.device)
+    return torch.cat([action, padding], dim=-1)
+
+
+def load_action_tensor(action: Any = None) -> torch.Tensor:
+    if action is None:
+        raise ValueError(
+            "Cosmos3 forward_dynamics action mode requires extra_args['action'] "
+            "(a JSON array / nested list / tensor of shape [T, D])."
+        )
+    if isinstance(action, torch.Tensor):
+        tensor = action.detach().to(dtype=torch.float32)
+    else:
+        tensor = torch.as_tensor(np.asarray(action), dtype=torch.float32)
+    if tensor.ndim == 3 and tensor.shape[0] == 1:
+        tensor = tensor.squeeze(0)
+    if tensor.ndim != 2:
+        raise ValueError(f"Cosmos3 action must have shape [T, D], got {tuple(tensor.shape)}.")
+    return tensor
+
+
+def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]:
+    key = str(resolution)
+    if key not in VIDEO_RES_SIZE_INFO:
+        raise ValueError(
+            f"Unknown Cosmos3 action resolution={resolution!r}; expected one of {sorted(VIDEO_RES_SIZE_INFO)}."
+        )
+    input_ratio = h / w
+    best_size = None
+    best_diff = float("inf")
+    for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values():
+        diff = abs(input_ratio - cand_h / cand_w)
+        if diff < best_diff:
+            best_diff = diff
+            best_size = (cand_w, cand_h)
+    assert best_size is not None
+    return best_size

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

@@ -46,6 +46,19 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
+from .action import (
+    ACTION_MODE_FORWARD_DYNAMICS,
+    ACTION_MODE_INVERSE_DYNAMICS,
+    ACTION_MODE_POLICY,
+    action_start_frame_offset,
+    build_action_condition_mask,
+    build_vision_condition_mask,
+    find_closest_target_size,
+    load_action_tensor,
+    normalize_action_mode,
+    pad_action_to_dim,
+    resolve_domain_id,
+)
 from .transformer_cosmos3 import Cosmos3VFMTransformer, resolve_sound_gen
 
 logger = init_logger(__name__)
@@ -98,14 +111,67 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
     if is_guardrails_enabled(od_config):
         ensure_initialized(od_config)
 
+    def _extra_args(request: OmniDiffusionRequest) -> dict[str, Any]:
+        extra = getattr(getattr(request, "sampling_params", None), "extra_args", None)
+        return extra if isinstance(extra, dict) else {}
+
+    def _request_action_mode(request: OmniDiffusionRequest) -> str | None:
+        return normalize_action_mode(_extra_args(request).get("action_mode"))
+
+    def _set_action_size_from_image(request: OmniDiffusionRequest, image: PIL.Image.Image) -> tuple[int, int]:
+        sp = request.sampling_params
+        if sp.height is not None and sp.width is not None:
+            return int(sp.height), int(sp.width)
+
+        extra = _extra_args(request)
+        resolution = extra.get("resolution", extra.get("image_size", 480))
+        target_w, target_h = find_closest_target_size(image.height, image.width, resolution)
+        if sp.height is None:
+            sp.height = target_h
+        if sp.width is None:
+            sp.width = target_w
+        return int(sp.height), int(sp.width)
+
     def _pil_to_rgb(value: Any) -> PIL.Image.Image:
         if isinstance(value, str):
             return PIL.Image.open(value).convert("RGB")
         if isinstance(value, PIL.Image.Image):
             return value.convert("RGB")
-        raise TypeError(f"Cosmos3 preprocessing expected PIL image or image path, got {type(value)!r}.")
+        raise TypeError(f"Cosmos3 action preprocessing expected PIL image or image path, got {type(value)!r}.")
+
+    def _resize_and_pad_action_image(image: PIL.Image.Image, target_h: int, target_w: int) -> PIL.Image.Image:
+        scale = min(target_w / image.width, target_h / image.height, 1.0)
+        resize_w = max(1, int(scale * image.width + 0.5))
+        resize_h = max(1, int(scale * image.height + 0.5))
+        if (resize_w, resize_h) != image.size:
+            image = image.resize((resize_w, resize_h), PIL.Image.Resampling.BICUBIC)
+
+        array = np.asarray(image)
+        pad_h = target_h - resize_h
+        pad_w = target_w - resize_w
+        if pad_h < 0 or pad_w < 0:
+            raise ValueError(
+                f"Cosmos3 action image resize exceeded target size: resized={(resize_h, resize_w)}, "
+                f"target={(target_h, target_w)}."
+            )
+        if pad_h == 0 and pad_w == 0:
+            return image
+        pad_mode = "reflect" if pad_h < resize_h and pad_w < resize_w else "edge"
+        padded = np.pad(array, ((0, pad_h), (0, pad_w), (0, 0)), mode=pad_mode)
+        return PIL.Image.fromarray(padded)
+
+    def _preprocess_action_image(image: PIL.Image.Image, target_h: int, target_w: int) -> torch.Tensor:
+        image = _resize_and_pad_action_image(image, target_h, target_w)
+        return video_processor.preprocess(image, height=target_h, width=target_w)
+
+    def _preprocess_action_video(frames: list[Any], target_h: int, target_w: int) -> torch.Tensor:
+        if not frames:
+            raise ValueError("Cosmos3 action video input must contain at least one frame.")
+        processed = [_preprocess_action_image(_pil_to_rgb(frame), target_h, target_w).squeeze(0) for frame in frames]
+        return torch.stack(processed, dim=1).unsqueeze(0).contiguous()
 
     def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
+        action_mode = _request_action_mode(request)
         if is_guardrails_enabled(od_config, request.sampling_params):
             for prompt in request.prompts:
                 text = prompt if isinstance(prompt, str) else prompt.get("prompt", "")
@@ -116,39 +182,63 @@ def get_cosmos3_pre_process_func(od_config: OmniDiffusionConfig):
                 continue
             multi_modal_data = prompt.get("multi_modal_data", {}) or {}
             raw_image = multi_modal_data.get("image")
-            if raw_image is None:
+            raw_video = multi_modal_data.get("video")
+            if raw_image is None and not (action_mode is not None and raw_video is not None):
                 continue
 
             if "additional_information" not in prompt:
                 prompt["additional_information"] = {}
 
-            image = _pil_to_rgb(raw_image)
+            if raw_image is None:
+                if not isinstance(raw_video, list) or not raw_video:
+                    raise TypeError("Cosmos3 action video input must be a non-empty list of PIL images or image paths.")
+                image = _pil_to_rgb(raw_video[0])
+            else:
+                image = _pil_to_rgb(raw_image)
 
             # Auto-calculate H/W from aspect ratio (720p max area)
             if request.sampling_params.height is None or request.sampling_params.width is None:
-                max_area = 720 * 1280
-                aspect_ratio = image.height / image.width
-                mod_value = 16
-                height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
-                width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
-                if request.sampling_params.height is None:
-                    request.sampling_params.height = height
-                if request.sampling_params.width is None:
-                    request.sampling_params.width = width
+                if action_mode is not None:
+                    _set_action_size_from_image(request, image)
+                else:
+                    max_area = 720 * 1280
+                    aspect_ratio = image.height / image.width
+                    mod_value = 16
+                    height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+                    width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+                    if request.sampling_params.height is None:
+                        request.sampling_params.height = height
+                    if request.sampling_params.width is None:
+                        request.sampling_params.width = width
 
             target_w = request.sampling_params.width
             target_h = request.sampling_params.height
-            scale = max(target_w / image.width, target_h / image.height)
-            resize_w = int(np.ceil(scale * image.width))
-            resize_h = int(np.ceil(scale * image.height))
-            image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS)
-            left = (resize_w - target_w) // 2
-            top = (resize_h - target_h) // 2
-            image = image.crop((left, top, left + target_w, top + target_h))
+            if action_mode is not None:
+                prompt["additional_information"]["preprocessed_image"] = _preprocess_action_image(
+                    image,
+                    int(target_h),
+                    int(target_w),
+                )
+            else:
+                scale = max(target_w / image.width, target_h / image.height)
+                resize_w = int(np.ceil(scale * image.width))
+                resize_h = int(np.ceil(scale * image.height))
+                image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS)
+                left = (resize_w - target_w) // 2
+                top = (resize_h - target_h) // 2
+                image = image.crop((left, top, left + target_w, top + target_h))
 
-            prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
-                image, height=target_h, width=target_w
-            )
+                prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
+                    image, height=target_h, width=target_w
+                )
+            if action_mode is not None and raw_video is not None:
+                if not isinstance(raw_video, list):
+                    raise TypeError("Cosmos3 action video input must be a list of PIL images or image paths.")
+                prompt["additional_information"]["preprocessed_video"] = _preprocess_action_video(
+                    raw_video,
+                    int(target_h),
+                    int(target_w),
+                )
             request.prompts[i] = prompt
 
         return request
@@ -421,11 +511,17 @@ class Cosmos3OmniDiffusersPipeline(
                 "time_embedder.",
                 "audio_proj_in.",
                 "audio_proj_out.",
+                "action_proj_in.",
+                "action_proj_out.",
             )
         ):
             return f"transformer.{k}"
         if k in ("audio_modality_embed", "audio_modality_embed.weight"):
             return "transformer.audio_modality_embed"
+        if k in ("action_modality_embed", "action_modality_embed.weight"):
+            return "transformer.action_modality_embed"
+        if k.startswith("action_pos_embed."):
+            return None
 
         # Skip lm_head
         if k.startswith("lm_head."):
@@ -528,13 +624,22 @@ class Cosmos3OmniDiffusersPipeline(
                     f"the checkpoint is missing sound weights for {missing}. "
                     "Use a sound-capable transformer checkpoint."
                 )
+        if getattr(self.transformer, "action_gen", False):
+            action_markers = ("action_proj_in.", "action_proj_out.", "action_modality_embed")
+            missing = [marker.rstrip(".") for marker in action_markers if not any(marker in name for name in loaded)]
+            if missing:
+                raise ValueError(
+                    "Cosmos3 transformer config enables action generation, but "
+                    f"the checkpoint is missing action weights for {missing}. "
+                    "Use an action-capable transformer checkpoint."
+                )
         return loaded
 
     def predict_noise(self, **kwargs) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Override CFGParallelMixin.predict_noise for Cosmos3.
 
         The transformer returns the raw prediction: video-only as a tensor,
-        or a tuple in video, sound order for sound generation.
+        or a tuple in video, action, sound order for multimodal generation.
         """
         return self.transformer(**kwargs)
 
@@ -611,6 +716,12 @@ class Cosmos3OmniDiffusersPipeline(
                 return True
         return False
 
+    @classmethod
+    def _get_action_mode(cls, prompt_data, sp) -> str | None:
+        return normalize_action_mode(
+            cls._get_sp_param(sp, "action_mode", cls._get_prompt_param(prompt_data, "action_mode", None))
+        )
+
     def _get_sound_tokenizer(self):
         if self._sound_tokenizer is None:
             from .sound_tokenizer import Cosmos3SoundTokenizer
@@ -1017,6 +1128,30 @@ class Cosmos3OmniDiffusersPipeline(
 
         return latent.to(self.dtype)
 
+    def _encode_video_tensor(self, video_tensor: torch.Tensor) -> torch.Tensor:
+        """VAE-encode a preprocessed pixel video [1, 3, T, H, W]."""
+        if video_tensor.ndim == 4:
+            video_tensor = video_tensor.unsqueeze(0)
+        if video_tensor.ndim != 5:
+            raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
+        if video_tensor.shape[0] != 1 or video_tensor.shape[1] != 3:
+            raise ValueError(f"Cosmos3 video tensor must have shape [1, 3, T, H, W], got {tuple(video_tensor.shape)}.")
+
+        video = video_tensor.to(device=self.device, dtype=self.vae.dtype)
+        latent = self.vae.encode(video).latent_dist.mode()
+
+        if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"):
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype)
+            )
+            latents_std = torch.tensor(self.vae.config.latents_std).view(1, -1, 1, 1, 1).to(latent.device, latent.dtype)
+            latent = (latent - latents_mean) / latents_std
+        else:
+            scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0)
+            latent = latent * scaling_factor
+
+        return latent.to(self.dtype)
+
     def _prepare_latents_i2v(
         self,
         image_tensor: torch.Tensor,
@@ -1053,6 +1188,95 @@ class Cosmos3OmniDiffusersPipeline(
         velocity_mask = 1.0 - condition_mask
         return latents, velocity_mask, image_latent
 
+    def _prepare_latents_action_video(
+        self,
+        video_tensor: torch.Tensor,
+        mode: str,
+        height: int,
+        width: int,
+        num_frames: int,
+        generator: torch.Generator,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Prepare video latents for action modes with mode-specific conditioning."""
+        del height, width
+        C = self.transformer.latent_channel_size
+        T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1
+        H_lat = video_tensor.shape[-2] // self.vae_scale_factor_spatial
+        W_lat = video_tensor.shape[-1] // self.vae_scale_factor_spatial
+
+        noise = randn_tensor(
+            (1, C, T_lat, H_lat, W_lat),
+            generator=generator,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        cond_latent = self._encode_video_tensor(video_tensor)
+        if cond_latent.shape[2:] != noise.shape[2:]:
+            raise ValueError(
+                "Cosmos3 action video latent shape mismatch: "
+                f"encoded={tuple(cond_latent.shape)}, expected={tuple(noise.shape)}."
+            )
+        condition_mask = build_vision_condition_mask(
+            mode,
+            num_frames,
+            self.vae_scale_factor_temporal,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise
+        velocity_mask = 1.0 - condition_mask
+        return latents, velocity_mask, cond_latent
+
+    def _prepare_action_latents(
+        self,
+        *,
+        mode: str,
+        action_chunk_size: int,
+        raw_action_dim: int | None,
+        generator: torch.Generator,
+        sp,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]:
+        action_dim = int(getattr(self.transformer, "action_dim", 64))
+        if mode == ACTION_MODE_FORWARD_DYNAMICS:
+            action = load_action_tensor(self._get_sp_param(sp, "action", None))
+            if action.shape[0] < action_chunk_size:
+                pad = action[-1:].repeat(action_chunk_size - action.shape[0], 1)
+                action = torch.cat([action, pad], dim=0)
+            elif action.shape[0] > action_chunk_size:
+                action = action[:action_chunk_size]
+            if raw_action_dim is None:
+                raw_action_dim = int(action.shape[-1])
+            clean_action = pad_action_to_dim(action, action_dim)
+        else:
+            if raw_action_dim is None:
+                raise ValueError(
+                    "Cosmos3 action_mode='policy' and 'inverse_dynamics' require extra_args['raw_action_dim']."
+                )
+            clean_action = torch.zeros(action_chunk_size, action_dim, dtype=torch.float32)
+
+        raw_action_dim = int(raw_action_dim)
+        if raw_action_dim <= 0 or raw_action_dim > action_dim:
+            raise ValueError(f"Cosmos3 raw_action_dim must be in [1, {action_dim}], got {raw_action_dim}.")
+
+        clean_action = clean_action.to(device=self.device, dtype=self.dtype).unsqueeze(0)
+        condition_mask = build_action_condition_mask(
+            mode,
+            action_chunk_size,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        noise = randn_tensor(
+            (1, action_chunk_size, action_dim),
+            generator=generator,
+            device=self.device,
+            dtype=self.dtype,
+        )
+        noise[:, :, raw_action_dim:] = 0
+        clean_action[:, :, raw_action_dim:] = 0
+        action_latents = condition_mask * clean_action + (1.0 - condition_mask) * noise
+        action_velocity_mask = 1.0 - condition_mask
+        return action_latents, action_velocity_mask, clean_action, raw_action_dim
+
     # -- Denoising loop (shared by T2V and I2V) -----------------------------
 
     def diffuse(
@@ -1066,11 +1290,15 @@ class Cosmos3OmniDiffusersPipeline(
         guidance_scale: float,
         shared_kwargs: dict,
         *,
+        action_latents: torch.Tensor | None = None,
+        action_velocity_mask: torch.Tensor | None = None,
+        action_condition_latents: torch.Tensor | None = None,
         sound_latents: torch.Tensor | None = None,
         velocity_mask: torch.Tensor | None = None,
         image_latent: torch.Tensor | None = None,
         condition_latents: torch.Tensor | None = None,
         guidance_interval: tuple[float, float] | None = None,
+        raw_action_dim: int | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Denoising loop with 3-mode CFG support (parallel, sequential, none).
 
@@ -1115,10 +1343,13 @@ class Cosmos3OmniDiffusersPipeline(
         # scheduler with cross-element dependencies (e.g. per-modality timestep).
         def _pack_joint(
             video_tensor: torch.Tensor,
+            action_tensor: torch.Tensor | None = None,
             sound_tensor: torch.Tensor | None = None,
         ):
             batch = video_tensor.shape[0]
             tensors = [video_tensor]
+            if action_tensor is not None:
+                tensors.append(action_tensor)
             if sound_tensor is not None:
                 tensors.append(sound_tensor)
             flats = [tensor.reshape(batch, -1) for tensor in tensors]
@@ -1138,57 +1369,86 @@ class Cosmos3OmniDiffusersPipeline(
 
         def _split_noise_pred(
             noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
-        ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+            has_action = action_latents is not None
             has_sound = sound_latents is not None
-            if not has_sound:
+            if not has_action and not has_sound:
                 if isinstance(noise_pred, tuple):
                     raise ValueError("Cosmos3 video-only diffusion received tuple predictions.")
-                return noise_pred, None
+                return noise_pred, None, None
             if not isinstance(noise_pred, tuple):
                 raise ValueError("Cosmos3 multimodal diffusion expects transformer predictions as a tuple.")
-            if len(noise_pred) != 2:
-                raise ValueError(f"Cosmos3 sound diffusion expected 2 predictions, got {len(noise_pred)}.")
-            return noise_pred[0], noise_pred[1]
+            expected = 1 + int(has_action) + int(has_sound)
+            if len(noise_pred) != expected:
+                raise ValueError(
+                    f"Cosmos3 multimodal diffusion expected {expected} predictions, got {len(noise_pred)}."
+                )
+            video_pred = noise_pred[0]
+            idx = 1
+            action_pred = noise_pred[idx] if has_action else None
+            if has_action:
+                idx += 1
+            sound_pred = noise_pred[idx] if has_sound else None
+            return video_pred, action_pred, sound_pred
 
         def _step(
             noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
             t: torch.Tensor,
             latents: torch.Tensor,
+            action_latents: torch.Tensor | None,
             sound_latents: torch.Tensor | None,
         ) -> torch.Tensor | tuple[torch.Tensor, ...]:
-            video_pred, sound_pred = _split_noise_pred(noise_pred)
+            video_pred, action_pred, sound_pred = _split_noise_pred(noise_pred)
             if velocity_mask is not None:
                 video_pred = video_pred * velocity_mask
-            if sound_latents is None:
+            if action_pred is not None and action_velocity_mask is not None:
+                action_pred = action_pred * action_velocity_mask
+                if raw_action_dim is not None and 0 < raw_action_dim < action_pred.shape[-1]:
+                    action_pred[..., raw_action_dim:] = 0
+            if action_latents is None and sound_latents is None:
                 latents = self.scheduler.step(video_pred, t, latents, return_dict=False)[0]
             else:
-                packed_noise, shapes, numels = _pack_joint(video_pred, sound_pred)
-                packed_latents, _, _ = _pack_joint(latents, sound_latents)
+                packed_noise, shapes, numels = _pack_joint(video_pred, action_pred, sound_pred)
+                packed_latents, _, _ = _pack_joint(latents, action_latents, sound_latents)
                 packed_next = self.scheduler.step(packed_noise, t, packed_latents, return_dict=False)[0]
                 unpacked = _unpack_joint(packed_next, shapes, numels)
                 latents = unpacked[0]
+                idx = 1
+                if action_latents is not None:
+                    action_latents = unpacked[idx]
+                    idx += 1
                 if sound_latents is not None:
-                    sound_latents = unpacked[1]
+                    sound_latents = unpacked[idx]
             if condition_latents is not None and velocity_mask is not None:
                 latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents
             elif image_latent is not None:
                 latents[:, :, 0:1, :, :] = image_latent
+            if action_latents is not None and action_condition_latents is not None and action_velocity_mask is not None:
+                action_latents = (
+                    action_velocity_mask * action_latents + (1.0 - action_velocity_mask) * action_condition_latents
+                )
             outputs = [latents]
+            if action_latents is not None:
+                outputs.append(action_latents)
             if sound_latents is not None:
                 outputs.append(sound_latents)
             return outputs[0] if len(outputs) == 1 else tuple(outputs)
 
         def _assign_step_out(step_out: torch.Tensor | tuple[torch.Tensor, ...]) -> None:
-            nonlocal latents, sound_latents
-            if sound_latents is None:
+            nonlocal latents, action_latents, sound_latents
+            if action_latents is None and sound_latents is None:
                 assert isinstance(step_out, torch.Tensor)
                 latents = step_out
                 return
             if not isinstance(step_out, tuple):
                 raise ValueError("Cosmos3 multimodal diffusion step returned a non-tuple result.")
             latents = step_out[0]
+            idx = 1
+            if action_latents is not None:
+                action_latents = step_out[idx]
+                idx += 1
             if sound_latents is not None:
-                sound_latents = step_out[1]
+                sound_latents = step_out[idx]
 
         if cfg_parallel:
             for t in self.progress_bar(timesteps):
@@ -1206,6 +1466,7 @@ class Cosmos3OmniDiffusersPipeline(
                         timestep=timestep,
                         text_ids=cond_ids,
                         text_mask=cond_mask,
+                        action_latents=action_latents,
                         sound_latents=sound_latents,
                         **shared_kwargs,
                     ),
@@ -1214,12 +1475,13 @@ class Cosmos3OmniDiffusersPipeline(
                         timestep=timestep,
                         text_ids=uncond_ids,
                         text_mask=uncond_mask,
+                        action_latents=action_latents,
                         sound_latents=sound_latents,
                         **shared_kwargs,
                     ),
                     cfg_normalize=False,
                 )
-                _assign_step_out(_step(noise_pred, t, latents, sound_latents))
+                _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents))
 
         elif do_cfg:
             cond_cache: tuple = (None, None)
@@ -1237,6 +1499,7 @@ class Cosmos3OmniDiffusersPipeline(
                     timestep=timestep,
                     text_ids=cond_ids,
                     text_mask=cond_mask,
+                    action_latents=action_latents,
                     sound_latents=sound_latents,
                     **shared_kwargs,
                 )
@@ -1250,6 +1513,7 @@ class Cosmos3OmniDiffusersPipeline(
                         timestep=timestep,
                         text_ids=uncond_ids,
                         text_mask=uncond_mask,
+                        action_latents=action_latents,
                         sound_latents=sound_latents,
                         **shared_kwargs,
                     )
@@ -1263,7 +1527,7 @@ class Cosmos3OmniDiffusersPipeline(
                 else:
                     noise_pred = noise_cond
 
-                _assign_step_out(_step(noise_pred, t, latents, sound_latents))
+                _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents))
 
         else:
             for t in self.progress_bar(timesteps):
@@ -1273,12 +1537,15 @@ class Cosmos3OmniDiffusersPipeline(
                     timestep=timestep,
                     text_ids=cond_ids,
                     text_mask=cond_mask,
+                    action_latents=action_latents,
                     sound_latents=sound_latents,
                     **shared_kwargs,
                 )
-                _assign_step_out(_step(noise_pred, t, latents, sound_latents))
+                _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents))
 
         outputs = [latents]
+        if action_latents is not None:
+            outputs.append(action_latents)
         if sound_latents is not None:
             outputs.append(sound_latents)
         return outputs[0] if len(outputs) == 1 else tuple(outputs)
@@ -1300,15 +1567,29 @@ class Cosmos3OmniDiffusersPipeline(
             prompt = prompt_data
             negative_prompt = None
             image_tensor = None
+            action_video_tensor = None
         else:
             prompt = prompt_data.get("prompt", "")
             negative_prompt = prompt_data.get("negative_prompt")
             additional_info = prompt_data.get("additional_information", {}) or {}
             image_tensor = additional_info.get("preprocessed_image")
+            action_video_tensor = additional_info.get("preprocessed_video")
 
         sp = req.sampling_params
         is_t2i = self._is_t2i_request(req)
         sound_enabled = self._is_sound_request(prompt_data, sp)
+        action_mode = self._get_action_mode(prompt_data, sp)
+        action_enabled = action_mode is not None
+        if action_enabled and is_t2i:
+            raise ValueError("Cosmos3 action generation is supported only for video outputs.")
+        if action_enabled and sound_enabled:
+            raise ValueError("Cosmos3 action+sound joint generation is not supported in this phase.")
+        if action_enabled and not getattr(self.transformer, "action_gen", False):
+            raise ValueError(
+                "Cosmos3 action generation was requested, but the transformer was "
+                "initialized without action modules. Check that the checkpoint config "
+                "enables action_gen and includes action weights."
+            )
         if sound_enabled and is_t2i:
             raise ValueError(
                 "Cosmos3 sound generation is supported only for video outputs in "
@@ -1349,6 +1630,36 @@ class Cosmos3OmniDiffusersPipeline(
             default_guidance_interval = None
             batch_size = 1  # Existing video pipeline assumes B=1.
 
+        if action_enabled:
+            action_chunk_param = self._get_sp_param(sp, "action_chunk_size", None)
+            if action_chunk_param is not None:
+                action_chunk_size = int(action_chunk_param)
+                if sp.num_frames is None:
+                    num_frames = action_chunk_size + 1
+            elif sp.num_frames is None:
+                action_chunk_size = 16
+                num_frames = action_chunk_size + 1
+            else:
+                action_chunk_size = int(num_frames) - 1
+            if action_chunk_size <= 0:
+                raise ValueError(f"Cosmos3 action_chunk_size must be positive, got {action_chunk_size}.")
+            if num_frames not in (action_chunk_size, action_chunk_size + 1):
+                raise ValueError(
+                    "Cosmos3 action requests require num_frames to equal action_chunk_size "
+                    f"or action_chunk_size + 1; got num_frames={num_frames}, action_chunk_size={action_chunk_size}."
+                )
+            num_inference_steps = sp.num_inference_steps or 30
+            guidance_scale = sp.guidance_scale if sp.guidance_scale is not None else 1.0
+            default_flow_shift = 5.0
+
+        domain_id = None
+        if action_enabled:
+            domain_id = resolve_domain_id(
+                domain_id=self._get_sp_param(sp, "domain_id", None),
+                domain_name=self._get_sp_param(sp, "domain_name", None),
+                require_explicit=True,
+            )
+
         # Runtime controls: prefer ``extra_args`` (OpenAI endpoints write
         # there) over direct attrs.
         flow_shift_target = float(self._get_sp_param(sp, "flow_shift", default_flow_shift))
@@ -1361,6 +1672,23 @@ class Cosmos3OmniDiffusersPipeline(
         )
         use_system_prompt = bool(self._get_sp_param(sp, "use_system_prompt", False))
 
+        if action_enabled and action_video_tensor is None:
+            extra_action_video = self._get_sp_param(sp, "action_video", None)
+            if isinstance(extra_action_video, torch.Tensor):
+                action_video_tensor = extra_action_video
+        if action_enabled and isinstance(action_video_tensor, torch.Tensor):
+            if action_video_tensor.ndim == 4:
+                action_video_tensor = action_video_tensor.unsqueeze(0)
+            if action_video_tensor.ndim != 5:
+                raise ValueError(
+                    "Cosmos3 extra_args['action_video'] must have shape [1, 3, T, H, W] "
+                    f"or [3, T, H, W], got {tuple(action_video_tensor.shape)}."
+                )
+            if sp.height is None:
+                height = int(action_video_tensor.shape[-2])
+            if sp.width is None:
+                width = int(action_video_tensor.shape[-1])
+
         self._guidance_scale = guidance_scale
         self._num_timesteps = num_inference_steps
 
@@ -1396,7 +1724,58 @@ class Cosmos3OmniDiffusersPipeline(
         # batching B=N together would require expanding text K/V (UND
         # pathway is text-only and cached) and is left as a future
         # optimization.
-        if image_tensor is not None and not is_t2i:
+        action_latents = None
+        action_velocity_mask = None
+        action_condition_latents = None
+        raw_action_dim = None
+        action_offset = 1
+        if action_enabled:
+            if action_video_tensor is not None and action_video_tensor.ndim == 4:
+                action_video_tensor = action_video_tensor.unsqueeze(0)
+            if action_video_tensor is not None and action_video_tensor.ndim != 5:
+                raise ValueError(
+                    "Cosmos3 action video tensor must have shape [1, 3, T, H, W] "
+                    f"or [3, T, H, W], got {tuple(action_video_tensor.shape)}."
+                )
+            if action_video_tensor is not None and action_video_tensor.shape[2] < num_frames:
+                pad = action_video_tensor[:, :, -1:].repeat(1, 1, num_frames - action_video_tensor.shape[2], 1, 1)
+                action_video_tensor = torch.cat([action_video_tensor, pad], dim=2)
+            elif action_video_tensor is not None and action_video_tensor.shape[2] > num_frames:
+                action_video_tensor = action_video_tensor[:, :, :num_frames]
+
+            if action_mode == ACTION_MODE_INVERSE_DYNAMICS and action_video_tensor is None:
+                raise ValueError("Cosmos3 inverse_dynamics action mode requires multi_modal_data['video'].")
+            if action_mode in {ACTION_MODE_POLICY, ACTION_MODE_FORWARD_DYNAMICS} and image_tensor is None:
+                if action_video_tensor is None:
+                    raise ValueError(
+                        f"Cosmos3 action_mode={action_mode!r} requires multi_modal_data['image'] "
+                        "or multi_modal_data['video']."
+                    )
+                image_tensor = action_video_tensor[:, :, 0]
+
+            raw_action_dim_param = self._get_sp_param(sp, "raw_action_dim", None)
+            raw_action_dim = int(raw_action_dim_param) if raw_action_dim_param is not None else None
+            action_prepared = self._prepare_action_latents(
+                mode=action_mode,
+                action_chunk_size=action_chunk_size,
+                raw_action_dim=raw_action_dim,
+                generator=generator,
+                sp=sp,
+            )
+            action_latents, action_velocity_mask, action_condition_latents, raw_action_dim = action_prepared
+            action_offset = action_start_frame_offset(action_mode, action_chunk_size, num_frames)
+
+        if action_enabled and action_video_tensor is not None:
+            latents, velocity_mask, condition_latents = self._prepare_latents_action_video(
+                action_video_tensor,
+                action_mode,
+                height,
+                width,
+                num_frames,
+                generator,
+            )
+            image_latent = condition_latents[:, :, 0:1]
+        elif image_tensor is not None and not is_t2i:
             latents, velocity_mask, image_latent = self._prepare_latents_i2v(
                 image_tensor,
                 height,
@@ -1427,6 +1806,13 @@ class Cosmos3OmniDiffusersPipeline(
         shared_kwargs = dict(video_shape=video_shape, fps=frame_rate)
         if velocity_mask is not None:
             shared_kwargs["noisy_frame_mask"] = velocity_mask
+        if action_enabled:
+            shared_kwargs.update(
+                action_domain_ids=torch.tensor([domain_id], dtype=torch.long, device=self.device),
+                action_noisy_mask=action_velocity_mask,
+                action_start_frame_offset=action_offset,
+                action_fps=float(self._get_sp_param(sp, "action_fps", frame_rate) or frame_rate),
+            )
 
         def _run_diffusion(start_latents):
             self.scheduler.set_timesteps(num_inference_steps, device=self.device)
@@ -1439,11 +1825,15 @@ class Cosmos3OmniDiffusersPipeline(
                 uncond_mask=uncond_mask,
                 guidance_scale=guidance_scale,
                 shared_kwargs=shared_kwargs,
+                action_latents=action_latents,
+                action_velocity_mask=action_velocity_mask,
+                action_condition_latents=action_condition_latents,
                 sound_latents=sound_latents,
                 velocity_mask=velocity_mask,
                 image_latent=image_latent,
                 condition_latents=condition_latents,
                 guidance_interval=guidance_interval,
+                raw_action_dim=raw_action_dim,
             )
 
         if is_t2i and batch_size > 1:
@@ -1461,7 +1851,11 @@ class Cosmos3OmniDiffusersPipeline(
             latents = torch.cat(samples, dim=0)
         else:
             diffusion_output = _run_diffusion(latents)
-            if sound_enabled:
+            if action_enabled and sound_enabled:
+                latents, action_latents, sound_latents = diffusion_output
+            elif action_enabled:
+                latents, action_latents = diffusion_output
+            elif sound_enabled:
                 latents, sound_latents = diffusion_output
             else:
                 latents = diffusion_output
@@ -1483,4 +1877,18 @@ class Cosmos3OmniDiffusersPipeline(
             audio = self._decode_sound_latents(sound_latents, target_audio_samples)
             return DiffusionOutput(output={"video": video, "audio": audio, "audio_sample_rate": sound_sample_rate})
 
+        if action_enabled:
+            if action_latents is None or raw_action_dim is None or domain_id is None:
+                raise ValueError("Cosmos3 action generation finished without action latents.")
+            action = action_latents[:, :, :raw_action_dim].detach().cpu()
+            return DiffusionOutput(
+                output={"video": video},
+                custom_output={
+                    "action": action,
+                    "raw_action_dim": raw_action_dim,
+                    "action_mode": action_mode,
+                    "domain_id": domain_id,
+                },
+            )
+
         return DiffusionOutput(output={"image": video} if is_t2i else {"video": video})

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

@@ -136,6 +136,47 @@ def resolve_sound_gen(od_config: Any) -> bool:
     return False
 
 
+class DomainAwareLinear(nn.Module):
+    """Linear projection with one weight/bias pair per action embodiment domain."""
+
+    def __init__(
+        self,
+        input_size: int,
+        output_size: int,
+        num_domains: int,
+        *,
+        dtype: torch.dtype = torch.bfloat16,
+    ) -> None:
+        super().__init__()
+        self.input_size = int(input_size)
+        self.output_size = int(output_size)
+        self.num_domains = int(num_domains)
+        self.fc = nn.Embedding(self.num_domains, self.output_size * self.input_size, dtype=dtype)
+        self.bias = nn.Embedding(self.num_domains, self.output_size, dtype=dtype)
+        nn.init.xavier_uniform_(self.fc.weight)
+        nn.init.zeros_(self.bias.weight)
+
+    def forward(self, x: torch.Tensor, domain_id: torch.Tensor) -> torch.Tensor:
+        if domain_id.ndim == 0:
+            domain_id = domain_id.unsqueeze(0)
+        domain_id = domain_id.to(device=x.device, dtype=torch.long).reshape(-1)
+        if x.shape[0] != domain_id.shape[0]:
+            raise ValueError(
+                "Cosmos3 action domain_id batch size must match action tokens: "
+                f"tokens={x.shape[0]}, domain_id={domain_id.shape[0]}."
+            )
+        if torch.any((domain_id < 0) | (domain_id >= self.num_domains)):
+            raise ValueError(f"Cosmos3 action domain_id must be in [0, {self.num_domains}), got {domain_id.tolist()}.")
+
+        weight = self.fc(domain_id).view(domain_id.shape[0], self.input_size, self.output_size)
+        bias = self.bias(domain_id).view(domain_id.shape[0], self.output_size)
+        if x.ndim == 2:
+            return torch.bmm(x.unsqueeze(1), weight).squeeze(1) + bias
+        if x.ndim == 3:
+            return torch.bmm(x, weight) + bias.unsqueeze(1)
+        raise ValueError(f"Cosmos3 DomainAwareLinear expected rank-2 or rank-3 input, got {tuple(x.shape)}.")
+
+
 # ---------------------------------------------------------------------------
 # Rotary Position Embeddings (mRoPE)
 # ---------------------------------------------------------------------------
@@ -163,6 +204,7 @@ def compute_mrope_position_ids_vision(
     temporal_compression_factor: int = 4,
     base_temporal_compression_factor: int | None = None,
     enable_fps_modulation: bool = True,
+    start_frame_offset: int = 0,
 ) -> tuple[torch.Tensor, int | float]:
     """Generate 3D mRoPE position IDs for vision tokens.
 
@@ -180,10 +222,17 @@ def compute_mrope_position_ids_vision(
         )
         base_tps = base_fps / effective_base_tcf
         frame_indices = torch.arange(grid_t, dtype=torch.float32)
-        t_index = (frame_indices / tps * base_tps + temporal_offset).view(-1, 1).expand(-1, grid_h * grid_w).flatten()
+        t_index = (
+            ((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset)
+            .view(-1, 1)
+            .expand(-1, grid_h * grid_w)
+            .flatten()
+        )
     else:
-        t_index = torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten() + int(
-            temporal_offset
+        t_index = (
+            torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten()
+            + int(temporal_offset)
+            + start_frame_offset
         )
 
     h_index = torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten()
@@ -220,6 +269,30 @@ def compute_mrope_position_ids_sound(
     )
 
 
+def compute_mrope_position_ids_action(
+    grid_t: int,
+    temporal_offset: int | float,
+    action_fps: float | None,
+    base_fps: float = 24.0,
+    base_temporal_compression_factor: int = 4,
+    enable_fps_modulation: bool = True,
+    start_frame_offset: int = 1,
+) -> tuple[torch.Tensor, int | float]:
+    """Generate mRoPE IDs for action tokens as a frame-rate (T, 1, 1) grid."""
+    return compute_mrope_position_ids_vision(
+        grid_t=grid_t,
+        grid_h=1,
+        grid_w=1,
+        temporal_offset=temporal_offset,
+        fps=action_fps,
+        base_fps=base_fps,
+        temporal_compression_factor=1,
+        base_temporal_compression_factor=base_temporal_compression_factor,
+        enable_fps_modulation=enable_fps_modulation,
+        start_frame_offset=start_frame_offset,
+    )
+
+
 class Qwen3VLTextRotaryEmbedding(nn.Module):
     """Multi-dimensional rotary position embedding for Qwen3-VL."""
 
@@ -947,19 +1020,28 @@ class Cosmos3VFMTransformer(nn.Module):
         self.sound_gen = sound_gen
         self.sound_dim = sound_dim
         self.sound_latent_fps = sound_latent_fps
-
         if self.sound_gen and (sound_dim is None or sound_latent_fps is None):
             raise ValueError(
                 "Cosmos3VFMTransformer requires an explicit sound_dim and sound_latent_fps when sound_gen is True; "
                 "the pipeline must pass Cosmos3SoundTokenizer.latent_ch so the audio projection "
                 "layers are sized from the authoritative AVAE latent width."
             )
+        action_gen_value = _od_config_get(od_config, "action_gen", None)
+        action_dim_value = _od_config_get(od_config, "action_dim", None)
+        if action_dim_value is None:
+            action_dim_value = _od_config_get(od_config, "max_action_dim", None)
+        self.action_gen = _as_bool(action_gen_value) if action_gen_value is not None else False
+        self.action_dim = int(action_dim_value if action_dim_value is not None else 64)
+        self.num_embodiment_domains = int(_od_config_get(od_config, "num_embodiment_domains", 32))
         if temporal_compression_factor is None:
             temporal_compression_factor = _tf_config_get(model_config, "temporal_compression_factor", 4)
         self.temporal_compression_factor = int(temporal_compression_factor)
         self.temporal_compression_factor_sound = int(
             _tf_config_get(model_config, "temporal_compression_factor_sound", 1)
         )
+        self.temporal_compression_factor_sound = int(
+            _tf_config_get(model_config, "temporal_compression_factor_sound", 1)
+        )
         self.enable_fps_modulation = bool(_tf_config_get(model_config, "enable_fps_modulation", True))
         self.temporal_modality_margin = int(
             _tf_config_get(
@@ -992,6 +1074,20 @@ class Cosmos3VFMTransformer(nn.Module):
         self.proj_in = nn.Linear(self.patch_latent_dim, self.hidden_size)
         self.proj_out = nn.Linear(self.hidden_size, self.patch_latent_dim)
         self.time_embedder = TimestepEmbedder(self.hidden_size, target_dtype=dtype)
+        if self.action_gen:
+            self.action_proj_in = DomainAwareLinear(
+                self.action_dim,
+                self.hidden_size,
+                self.num_embodiment_domains,
+                dtype=dtype,
+            )
+            self.action_proj_out = DomainAwareLinear(
+                self.hidden_size,
+                self.action_dim,
+                self.num_embodiment_domains,
+                dtype=dtype,
+            )
+            self.action_modality_embed = nn.Parameter(torch.zeros(self.hidden_size, dtype=dtype))
         if self.sound_gen:
             self.audio_proj_in = nn.Linear(self.sound_dim, self.hidden_size)
             self.audio_proj_out = nn.Linear(self.hidden_size, self.sound_dim)
@@ -1079,6 +1175,21 @@ class Cosmos3VFMTransformer(nn.Module):
         """[B, T_sound, C_sound] -> [B, C_sound, T_sound]."""
         return tokens.permute(0, 2, 1).contiguous()
 
+    def pack_action(self, action_latents: torch.Tensor) -> torch.Tensor:
+        """Validate and return action latents as [B, T_action, D_action] tokens."""
+        if action_latents.ndim != 3:
+            raise ValueError(f"Cosmos3 action latents must have shape [B, T, D], got {tuple(action_latents.shape)}.")
+        if action_latents.shape[-1] != self.action_dim:
+            raise ValueError(
+                f"Cosmos3 action latent dimension mismatch: expected {self.action_dim}, got {action_latents.shape[-1]}."
+            )
+        return action_latents.contiguous()
+
+    @staticmethod
+    def unpack_action(tokens: torch.Tensor) -> torch.Tensor:
+        """Return [B, T_action, D_action] action predictions."""
+        return tokens.contiguous()
+
     # -- RoPE computation ----------------------------------------------------
 
     def _compute_rope_freqs(
@@ -1090,6 +1201,9 @@ class Cosmos3VFMTransformer(nn.Module):
         fps: float | None,
         device: torch.device,
         dtype: torch.dtype,
+        t_action: int | None = None,
+        action_start_frame_offset: int = 1,
+        action_fps: float | None = None,
         t_sound: int | None = None,
     ) -> tuple[tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor]]:
         """Compute mRoPE cos/sin for UND text and GEN media pathways."""
@@ -1097,6 +1211,7 @@ class Cosmos3VFMTransformer(nn.Module):
         S_text = text_mask.shape[1]
         text_lengths = text_mask.sum(dim=1).long()
         effective_fps = fps if fps is not None and t > 1 else None
+        action_frames = int(t_action or 0)
         sound_frames = int(t_sound or 0)
 
         text_pos_list = []
@@ -1116,6 +1231,17 @@ class Cosmos3VFMTransformer(nn.Module):
                 enable_fps_modulation=self.enable_fps_modulation,
             )
             gen_positions = [v_pos]
+            if action_frames > 0:
+                a_pos, _ = compute_mrope_position_ids_action(
+                    action_frames,
+                    temporal_offset=media_temporal_offset,
+                    action_fps=action_fps if action_fps is not None else fps,
+                    base_fps=self.base_fps,
+                    base_temporal_compression_factor=self.temporal_compression_factor,
+                    enable_fps_modulation=self.enable_fps_modulation,
+                    start_frame_offset=action_start_frame_offset,
+                )
+                gen_positions.append(a_pos)
             if sound_frames > 0:
                 s_pos, _ = compute_mrope_position_ids_sound(
                     sound_frames,
@@ -1161,7 +1287,9 @@ class Cosmos3VFMTransformer(nn.Module):
         *,
         s_gen: int,
         s_video: int,
+        s_action: int,
         s_sound: int,
+        has_action: bool,
         has_sound: bool,
         ulysses_size: int,
     ) -> None:
@@ -1169,14 +1297,16 @@ class Cosmos3VFMTransformer(nn.Module):
             return
 
         detail_parts = [f"video tokens {s_video}"]
+        if has_action:
+            detail_parts.append(f"action tokens {s_action}")
         if has_sound:
             detail_parts.append(f"sound tokens {s_sound}")
         detail = " = " + " + ".join(detail_parts) if len(detail_parts) > 1 else ""
         adjust_detail = (
-            "Adjust the spatial resolution, frame count, sound duration, "
-            "or sound latent FPS so the combined media sequence is a "
+            "Adjust the spatial resolution, frame count, action chunk size, "
+            "sound duration, or sound latent FPS so the combined media sequence is a "
             "multiple of ulysses_degree."
-            if has_sound
+            if has_action or has_sound
             else (
                 "Adjust the spatial resolution so that "
                 "t * ceil(h/patch) * ceil(w/patch) is a multiple "
@@ -1198,6 +1328,11 @@ class Cosmos3VFMTransformer(nn.Module):
         text_mask: torch.Tensor,
         video_shape: tuple[int, int, int],
         fps: float | None = None,
+        action_latents: torch.Tensor | None = None,
+        action_domain_ids: torch.Tensor | None = None,
+        action_noisy_mask: torch.Tensor | None = None,
+        action_start_frame_offset: int = 1,
+        action_fps: float | None = None,
         sound_latents: torch.Tensor | None = None,
         noisy_frame_mask: torch.Tensor | None = None,
         **kwargs,
@@ -1210,6 +1345,10 @@ class Cosmos3VFMTransformer(nn.Module):
             text_mask: [B, S_text] attention mask (1=real, 0=pad)
             video_shape: (t, h, w) in latent space
             fps: video frame rate for temporal mRoPE modulation
+            action_latents: Optional [B, T_action, D_action] noisy action latents.
+            action_domain_ids: Optional [B] embodiment domain IDs for action projections.
+            action_noisy_mask: Optional [B, T_action, 1] mask where 1=noisy
+                action token and 0=clean conditioned token.
             sound_latents: Optional [B, C_sound, T_sound] noisy sound latents.
             noisy_frame_mask: Optional [B, 1, t, 1, 1] mask where 1=noisy (add
                 timestep embedding, predict velocity) and 0=conditioned (clean
@@ -1218,7 +1357,7 @@ class Cosmos3VFMTransformer(nn.Module):
 
         Returns:
             [B, C, t, h, w] velocity prediction, or
-            tuple outputs in video, sound order when sound latents are provided.
+            tuple outputs in video, action, sound order when extra modalities are provided.
         """
         t, h, w = video_shape
         hp, wp, _, _ = self._pad_to_patch_size(h, w)
@@ -1230,7 +1369,14 @@ class Cosmos3VFMTransformer(nn.Module):
                 f"Cosmos3 requires identical real text lengths within a batch "
                 f"(got min={min_real_len}, max={max_real_len})."
             )
+        has_action = action_latents is not None
         has_sound = sound_latents is not None
+        if has_action and not self.action_gen:
+            raise ValueError(
+                "Cosmos3 action generation was requested, but this transformer "
+                "was initialized without action modules. Check that the "
+                "transformer config enables action_gen."
+            )
         if has_sound and not self.sound_gen:
             raise ValueError(
                 "Cosmos3 sound generation was requested, but this transformer "
@@ -1244,8 +1390,21 @@ class Cosmos3VFMTransformer(nn.Module):
         # Patchify latents and project to hidden space
         hidden_video = self.proj_in(self.patchify(hidden_states, t, h, w))
         s_video = hidden_video.shape[1]
+        s_action = 0
+        hidden_action = None
         s_sound = 0
         hidden_sound = None
+        if action_latents is not None:
+            if action_latents.shape[0] != hidden_states.shape[0]:
+                raise ValueError(
+                    "Cosmos3 action and video batch sizes must match: "
+                    f"video={hidden_states.shape[0]}, action={action_latents.shape[0]}."
+                )
+            if action_domain_ids is None:
+                action_domain_ids = torch.zeros(action_latents.shape[0], dtype=torch.long, device=action_latents.device)
+            hidden_action = self.action_proj_in(self.pack_action(action_latents), action_domain_ids)
+            hidden_action = hidden_action + self.action_modality_embed.to(hidden_action.dtype)
+            s_action = hidden_action.shape[1]
         if sound_latents is not None:
             if sound_latents.shape[0] != hidden_states.shape[0]:
                 raise ValueError(
@@ -1277,9 +1436,22 @@ class Cosmos3VFMTransformer(nn.Module):
         else:
             hidden_video = hidden_video + time_embed.unsqueeze(1)
 
+        if hidden_action is not None:
+            if action_noisy_mask is None:
+                hidden_action = hidden_action + time_embed.unsqueeze(1)
+            else:
+                if action_noisy_mask.shape != (hidden_action.shape[0], hidden_action.shape[1], 1):
+                    raise ValueError(
+                        "Cosmos3 action_noisy_mask must have shape [B, T_action, 1], "
+                        f"got {tuple(action_noisy_mask.shape)}."
+                    )
+                hidden_action = hidden_action + time_embed.unsqueeze(1) * action_noisy_mask.to(hidden_action.dtype)
+
         if hidden_sound is not None:
             hidden_sound = hidden_sound + time_embed.unsqueeze(1)
         hidden_parts = [hidden_video]
+        if hidden_action is not None:
+            hidden_parts.append(hidden_action)
         if hidden_sound is not None:
             hidden_parts.append(hidden_sound)
         hidden_gen = torch.cat(hidden_parts, dim=1)
@@ -1294,6 +1466,9 @@ class Cosmos3VFMTransformer(nn.Module):
                 fps,
                 hidden_states.device,
                 hidden_states.dtype,
+                t_action=s_action,
+                action_start_frame_offset=action_start_frame_offset,
+                action_fps=action_fps,
                 t_sound=s_sound,
             )
             cached_kv_full = self.language_model(text_ids, freqs_und)
@@ -1311,7 +1486,9 @@ class Cosmos3VFMTransformer(nn.Module):
         self._validate_gen_sequence_parallel(
             s_gen=hidden_gen.shape[1],
             s_video=s_video,
+            s_action=s_action,
             s_sound=s_sound,
+            has_action=has_action,
             has_sound=has_sound,
             ulysses_size=ulysses_size,
         )
@@ -1346,10 +1523,12 @@ class Cosmos3VFMTransformer(nn.Module):
 
         # Final norm and project back to latent space
         hidden_gen = self.norm_moe_gen(hidden_gen)
-        if not has_sound:
+        if not has_action and not has_sound:
             return self.unpatchify(self.proj_out(hidden_gen), t, h, w)
 
         split_sizes = [s_video]
+        if has_action:
+            split_sizes.append(s_action)
         if has_sound:
             split_sizes.append(s_sound)
         split_hidden = hidden_gen.split(split_sizes, dim=1)
@@ -1357,6 +1536,11 @@ class Cosmos3VFMTransformer(nn.Module):
         video_pred = self.unpatchify(self.proj_out(hidden_video), t, h, w)
         outputs: list[torch.Tensor] = [video_pred]
         split_idx = 1
+        if has_action:
+            hidden_action = split_hidden[split_idx]
+            split_idx += 1
+            assert action_domain_ids is not None
+            outputs.append(self.unpack_action(self.action_proj_out(hidden_action, action_domain_ids)))
         if has_sound:
             hidden_sound = split_hidden[split_idx]
             outputs.append(self.unpack_sound(self.audio_proj_out(hidden_sound)))

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py

@@ -39,7 +39,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
@@ -213,9 +213,10 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
         # Avoid the transformers v5 multi-worker subfolder race (see
         # ``vllm_omni/diffusion/model_loader/hub_prefetch.py`` for the full
         # analysis; L4 build #1043 hit this on FLUX.2-klein-4B's text_encoder).
+        flux2_subfolders = ["scheduler", "text_encoder", "tokenizer", "vae"]
         prefetch_subfolders(
             model,
-            ["scheduler", "text_encoder", "tokenizer", "vae"],
+            flux2_subfolders,
             local_files_only=local_files_only,
         )
 
@@ -224,9 +225,15 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
             subfolder="scheduler",
             local_files_only=local_files_only,
         )
-        self.text_encoder = Qwen3ForCausalLM.from_pretrained(
+        # ``from_pretrained_with_prefetch`` re-prefetches and retries if a peer
+        # worker left the cache half-written (missing-shard ``OSError`` or the
+        # default-config size-mismatch ``RuntimeError``) instead of crashing
+        # the worker - FLUX.2-klein-4B's sharded text_encoder hit this on #1043.
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen3ForCausalLM.from_pretrained,
             model,
             subfolder="text_encoder",
+            prefetch_list=flux2_subfolders,
             local_files_only=local_files_only,
         ).to(self._execution_device)
         self.tokenizer = Qwen2TokenizerFast.from_pretrained(
@@ -234,9 +241,11 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
             subfolder="tokenizer",
             local_files_only=local_files_only,
         )
-        self.vae = AutoencoderKLFlux2.from_pretrained(
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKLFlux2.from_pretrained,
             model,
             subfolder="vae",
+            prefetch_list=flux2_subfolders,
             local_files_only=local_files_only,
         ).to(self._execution_device)
 

vllm_omni/diffusion/models/hidream_image/pipeline_hidream_i1_image.py

@@ -31,6 +31,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -177,24 +178,55 @@ class HiDreamI1ImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfi
         # Check if model is a local path
         local_files_only = os.path.exists(model)
 
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the in-repo component set before any from_pretrained
+        # (``text_encoder_4`` lives in a separate Llama repo and is unaffected).
+        hidream_subfolders = [
+            "scheduler",
+            "vae",
+            "text_encoder",
+            "tokenizer",
+            "text_encoder_2",
+            "tokenizer_2",
+            "text_encoder_3",
+            "tokenizer_3",
+        ]
+        prefetch_subfolders(model, hidream_subfolders, local_files_only=local_files_only)
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
-            self.device
-        )
-        self.text_encoder = CLIPTextModelWithProjection.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKL.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
+        self.text_encoder = from_pretrained_with_prefetch(
+            CLIPTextModelWithProjection.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
         )
         self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
-            model, subfolder="text_encoder_2", local_files_only=local_files_only
+        self.text_encoder_2 = from_pretrained_with_prefetch(
+            CLIPTextModelWithProjection.from_pretrained,
+            model,
+            subfolder="text_encoder_2",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
         )
         self.tokenizer_2 = CLIPTokenizer.from_pretrained(
             model, subfolder="tokenizer_2", local_files_only=local_files_only
         )
-        self.text_encoder_3 = T5EncoderModel.from_pretrained(
-            model, subfolder="text_encoder_3", local_files_only=local_files_only
+        self.text_encoder_3 = from_pretrained_with_prefetch(
+            T5EncoderModel.from_pretrained,
+            model,
+            subfolder="text_encoder_3",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
         )
         self.tokenizer_3 = T5Tokenizer.from_pretrained(
             model, subfolder="tokenizer_3", local_files_only=local_files_only

vllm_omni/diffusion/models/hidream_image/pipeline_hidream_image.py

@@ -31,6 +31,7 @@ from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.hidream_image import HiDreamImageTransformer2DModel
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
@@ -177,24 +178,55 @@ class HiDreamImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile
         # Check if model is a local path
         local_files_only = os.path.exists(model)
 
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the in-repo component set before any from_pretrained
+        # (``text_encoder_4`` lives in a separate Llama repo and is unaffected).
+        hidream_subfolders = [
+            "scheduler",
+            "vae",
+            "text_encoder",
+            "tokenizer",
+            "text_encoder_2",
+            "tokenizer_2",
+            "text_encoder_3",
+            "tokenizer_3",
+        ]
+        prefetch_subfolders(model, hidream_subfolders, local_files_only=local_files_only)
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
-            self.device
-        )
-        self.text_encoder = CLIPTextModelWithProjection.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKL.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
+        self.text_encoder = from_pretrained_with_prefetch(
+            CLIPTextModelWithProjection.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
         )
         self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
-            model, subfolder="text_encoder_2", local_files_only=local_files_only
+        self.text_encoder_2 = from_pretrained_with_prefetch(
+            CLIPTextModelWithProjection.from_pretrained,
+            model,
+            subfolder="text_encoder_2",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
         )
         self.tokenizer_2 = CLIPTokenizer.from_pretrained(
             model, subfolder="tokenizer_2", local_files_only=local_files_only
         )
-        self.text_encoder_3 = T5EncoderModel.from_pretrained(
-            model, subfolder="text_encoder_3", local_files_only=local_files_only
+        self.text_encoder_3 = from_pretrained_with_prefetch(
+            T5EncoderModel.from_pretrained,
+            model,
+            subfolder="text_encoder_3",
+            prefetch_list=hidream_subfolders,
+            local_files_only=local_files_only,
         )
         self.tokenizer_3 = T5Tokenizer.from_pretrained(
             model, subfolder="tokenizer_3", local_files_only=local_files_only

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

@@ -32,7 +32,6 @@ from vllm.distributed import (
     get_tensor_model_parallel_world_size,
 )
 from vllm.model_executor.layers.activation import SiluAndMul
-from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     QKVParallelLinear,
@@ -73,6 +72,7 @@ from vllm_omni.diffusion.distributed.sp_plan import (
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
+from vllm_omni.diffusion.layers.norm import RMSNorm
 from vllm_omni.diffusion.layers.rope import RotaryEmbedding
 from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE
 from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -573,7 +573,9 @@ class HunyuanImage3Pipeline(
         cond_vit_image_embeds = []
         for batch_idx, image in enumerate(cond_vit_images):
             cur_kwargs = {k: v[batch_idx] for k, v in vit_kwargs.items()}
-            image_embed = self.vision_model(image, **cur_kwargs).last_hidden_state
+            # Siglip2VisionTransformer now returns a plain tensor (B, max_patches,
+            # hidden_size) after the vLLM-layers refactor, not an HF-style output.
+            image_embed = self.vision_model(image, **cur_kwargs)
             image_embed = self.vision_aligner(image_embed)
             n, seq_len, dim = image_embed.shape
             image_embed = image_embed.reshape(n * seq_len, dim)
@@ -624,7 +626,7 @@ class HunyuanImage3Pipeline(
             rope_image_info.append(list(zip(image_slices, image_shapes)))
         return rope_image_info
 
-    def vae_encode(self, image, cfg_factor=1):
+    def vae_encode(self, image, cfg_factor=1, generator=None):
         config = self.vae.config
 
         if image.ndim == 3:
@@ -639,11 +641,7 @@ class HunyuanImage3Pipeline(
             if isinstance(vae_encode_result, torch.Tensor):
                 latents = vae_encode_result
             else:
-                # Match HunyuanImage-3's cond encode path: sample the
-                # posterior, but use a fixed generator so repeated online
-                # requests are deterministic.
-                _cond_vae_gen = torch.Generator(device=image.device).manual_seed(0)
-                latents = vae_encode_result.latent_dist.sample(_cond_vae_gen)
+                latents = vae_encode_result.latent_dist.sample(generator)
             if hasattr(config, "shift_factor") and config.shift_factor:
                 latents.sub_(config.shift_factor)
             if hasattr(config, "scaling_factor") and config.scaling_factor:
@@ -666,14 +664,17 @@ class HunyuanImage3Pipeline(
         self,
         batch_cond_image_info_list: list[list[JointImageInfo]],
         cfg_factor: int = 1,
+        generator=None,
     ):
         # VAE encode one by one, as we assume cond images have different sizes
         batch_cond_vae_images, batch_cond_t, batch_cond_vit_images = [], [], []
-        for cond_image_info_list in batch_cond_image_info_list:
+        for batch_idx, cond_image_info_list in enumerate(batch_cond_image_info_list):
             cond_vae_image_list, cond_t_list, cond_vit_image_list = [], [], []
+            cond_generator = generator[batch_idx] if isinstance(generator, list) else generator
             for image_info in cond_image_info_list:
                 cond_t_, cond_vae_image_ = self.vae_encode(
                     image_info.vae_image_info.image_tensor.to(self.device),
+                    generator=cond_generator,
                 )
                 cond_vit_image_list.append(image_info.vision_image_info.image_tensor)
                 cond_vae_image_list.append(cond_vae_image_.squeeze(0))
@@ -863,9 +864,11 @@ class HunyuanImage3Pipeline(
         output, sections = out["output"], out["sections"]
 
         # 4. Encode conditional images
-        if batch_cond_image_info is not None and len(batch_cond_image_info[0]) > 0:
+        # Skip encoding if AR KV reuse is enabled
+        has_ar_kv = kwargs.get("ar_kv_data")
+        if batch_cond_image_info is not None and len(batch_cond_image_info[0]) > 0 and not has_ar_kv:
             cond_vae_images, cond_timestep, cond_vit_images = self._encode_cond_image(
-                batch_cond_image_info, cfg_factor[mode]
+                batch_cond_image_info, cfg_factor[mode], generator=generator
             )
             vit_kwargs = {"spatial_shapes": [], "attention_mask": []}
             for cond_image_info in batch_cond_image_info:
@@ -1389,18 +1392,23 @@ class HunyuanImage3Pipeline(
         **kwargs,
     ) -> DiffusionOutput:
         extra_args = getattr(getattr(req, "sampling_params", None), "extra_args", {}) or {}
+        bot_task = extra_args.get("bot_task")
         use_system_prompt = extra_args.get("use_system_prompt")
         system_prompt = extra_args.get("system_prompt")
-        # Fall back to per-prompt use_system_prompt forwarded by ar2diffusion
+        # Fall back to per-prompt prompt kwargs forwarded by ar2diffusion
         if req.prompts:
             first_prompt = req.prompts[0]
             if isinstance(first_prompt, dict):
+                if bot_task is None:
+                    bot_task = first_prompt.get("bot_task")
                 if use_system_prompt is None:
                     use_system_prompt = first_prompt.get("use_system_prompt")
                 if system_prompt is None:
                     system_prompt = first_prompt.get("system_prompt")
+        if isinstance(bot_task, str) and bot_task.lower() == "none":
+            bot_task = None
         if use_system_prompt is not None:
-            system_prompt = get_system_prompt(use_system_prompt, "image", system_prompt)
+            system_prompt = get_system_prompt(use_system_prompt, bot_task or "image", system_prompt)
             system_prompt = system_prompt.strip() if system_prompt is not None else ""
         prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt
 
@@ -1447,6 +1455,18 @@ class HunyuanImage3Pipeline(
         # ---- AR KV Reuse: extract injected KV from request ----
         ar_kv_kwargs = self._extract_ar_kv_from_request(req)
 
+        tokenizer_bot_task = bot_task
+        if tokenizer_bot_task == "think_recaption":
+            tokenizer_bot_task = "think"
+        elif tokenizer_bot_task == "vanilla":
+            tokenizer_bot_task = "image"
+        supported_bot_tasks = {"auto", "image", "think", "recaption", "img_ratio"}
+        if tokenizer_bot_task is not None and tokenizer_bot_task not in supported_bot_tasks:
+            raise ValueError(
+                f"Unsupported HunyuanImage3 single-stage bot_task: {tokenizer_bot_task!r}. "
+                f"Supported values are: {sorted(supported_bot_tasks)}."
+            )
+
         model_inputs = self.prepare_model_inputs(
             prompt=prompt,
             cot_text=cot_text,
@@ -1457,6 +1477,8 @@ class HunyuanImage3Pipeline(
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
             batch_cond_image_info=batch_cond_image_info,
+            bot_task=tokenizer_bot_task or "auto",
+            **ar_kv_kwargs,
         )
 
         model_inputs.update(ar_kv_kwargs)

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5.py

@@ -25,6 +25,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_1
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel
@@ -101,9 +102,20 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, Diff
         model = od_config.model
         local_files_only = os.path.exists(model)
 
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the whole component set before any from_pretrained so
+        # peer workers load from a warm, complete cache.
+        hv_subfolders = ["tokenizer", "text_encoder", "tokenizer_2", "text_encoder_2", "vae", "scheduler"]
+        prefetch_subfolders(model, hv_subfolders, local_files_only=local_files_only)
+
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.text_encoder = Qwen2_5_VLTextModel.from_pretrained(
-            model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen2_5_VLTextModel.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=hv_subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
 
         self.tokenizer_2 = ByT5Tokenizer.from_pretrained(
@@ -112,8 +124,13 @@ class HunyuanVideo15Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, Diff
         t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only)
         self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device)
 
-        self.vae = DistributedAutoencoderKLHunyuanVideo15.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+        self.vae = from_pretrained_with_prefetch(
+            DistributedAutoencoderKLHunyuanVideo15.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=hv_subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=torch.float32,
         ).to(self.device)
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(

vllm_omni/diffusion/models/hunyuan_video/pipeline_hunyuan_video_1_5_i2v.py

@@ -32,6 +32,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan_video_1
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.hunyuan_video.hunyuan_video_15_transformer import HunyuanVideo15Transformer3DModel
 from vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 import (
     extract_glyph_texts,
@@ -123,9 +124,29 @@ class HunyuanVideo15I2VPipeline(
         model = od_config.model
         local_files_only = os.path.exists(model)
 
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the whole component set before any from_pretrained so
+        # peer workers load from a warm, complete cache.
+        hv_subfolders = [
+            "tokenizer",
+            "text_encoder",
+            "tokenizer_2",
+            "text_encoder_2",
+            "image_encoder",
+            "feature_extractor",
+            "vae",
+            "scheduler",
+        ]
+        prefetch_subfolders(model, hv_subfolders, local_files_only=local_files_only)
+
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.text_encoder = Qwen2_5_VLTextModel.from_pretrained(
-            model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen2_5_VLTextModel.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=hv_subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
 
         self.tokenizer_2 = ByT5Tokenizer.from_pretrained(
@@ -134,15 +155,25 @@ class HunyuanVideo15I2VPipeline(
         t5_config = AutoConfig.from_pretrained(model, subfolder="text_encoder_2", local_files_only=local_files_only)
         self.text_encoder_2 = T5EncoderModel(t5_config, prefix="text_encoder_2").to(dtype=dtype, device=self.device)
 
-        self.image_encoder = SiglipVisionModel.from_pretrained(
-            model, subfolder="image_encoder", torch_dtype=dtype, local_files_only=local_files_only
+        self.image_encoder = from_pretrained_with_prefetch(
+            SiglipVisionModel.from_pretrained,
+            model,
+            subfolder="image_encoder",
+            prefetch_list=hv_subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
         self.feature_extractor = SiglipImageProcessor.from_pretrained(
             model, subfolder="feature_extractor", local_files_only=local_files_only
         )
 
-        self.vae = DistributedAutoencoderKLHunyuanVideo15.from_pretrained(
-            model, subfolder="vae", torch_dtype=torch.float32, local_files_only=local_files_only
+        self.vae = from_pretrained_with_prefetch(
+            DistributedAutoencoderKLHunyuanVideo15.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=hv_subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=torch.float32,
         ).to(self.device)
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(

vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image.py

@@ -27,6 +27,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import LongCatImageTransformer2DModel
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -223,19 +224,32 @@ class LongCatImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfile
         model = od_config.model
         local_files_only = os.path.exists(model)
 
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the whole component set before any from_pretrained.
+        longcat_subfolders = ["scheduler", "text_encoder", "tokenizer", "vae"]
+        prefetch_subfolders(model, longcat_subfolders, local_files_only=local_files_only)
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
 
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen2_5_VLForConditionalGeneration.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=longcat_subfolders,
+            local_files_only=local_files_only,
         )
         self.text_processor = Qwen2VLProcessor.from_pretrained(
             model, subfolder="tokenizer", local_files_only=local_files_only
         )
-        self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
-            self.device
-        )
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKL.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=longcat_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
         self.transformer = LongCatImageTransformer2DModel(od_config=od_config)
         self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 

vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py

@@ -29,6 +29,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.longcat_image.longcat_image_transformer import (
     LongCatImageTransformer2DModel,
@@ -244,19 +245,32 @@ class LongCatImageEditPipeline(nn.Module, CFGParallelMixin, SupportImageInput, D
         model = od_config.model
         local_files_only = os.path.exists(model)
 
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the whole component set before any from_pretrained.
+        longcat_subfolders = ["scheduler", "text_encoder", "text_processor", "tokenizer", "vae"]
+        prefetch_subfolders(model, longcat_subfolders, local_files_only=local_files_only)
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model, subfolder="text_encoder", local_files_only=local_files_only
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen2_5_VLForConditionalGeneration.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=longcat_subfolders,
+            local_files_only=local_files_only,
         )
         self.text_processor = Qwen2VLProcessor.from_pretrained(
             model, subfolder="text_processor", local_files_only=local_files_only
         )
 
-        self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
-            self.device
-        )
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKL.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=longcat_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
         self.transformer = LongCatImageTransformer2DModel(od_config=od_config)
         self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -36,6 +36,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
@@ -185,6 +186,19 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
             ),
         ]
 
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the whole component set before any from_pretrained.
+        ltx2_subfolders = [
+            "tokenizer",
+            "text_encoder",
+            "connectors",
+            "vae",
+            "audio_vae",
+            "vocoder",
+            "scheduler",
+        ]
+        prefetch_subfolders(model, ltx2_subfolders, local_files_only=local_files_only)
+
         self.tokenizer = AutoTokenizer.from_pretrained(
             model,
             subfolder="tokenizer",
@@ -193,36 +207,46 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
         # prefer mmap loading as default device is cuda, and the output of text encoder
         # could be deterministic.
         with torch.device("cpu"):
-            self.text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
+            self.text_encoder = from_pretrained_with_prefetch(
+                Gemma3ForConditionalGeneration.from_pretrained,
                 model,
                 subfolder="text_encoder",
-                torch_dtype=dtype,
+                prefetch_list=ltx2_subfolders,
                 local_files_only=local_files_only,
+                torch_dtype=dtype,
             ).to(self.device)
-        self.connectors = LTX2TextConnectors.from_pretrained(
+        self.connectors = from_pretrained_with_prefetch(
+            LTX2TextConnectors.from_pretrained,
             model,
             subfolder="connectors",
-            torch_dtype=dtype,
+            prefetch_list=ltx2_subfolders,
             local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
 
-        self.vae = AutoencoderKLLTX2Video.from_pretrained(
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKLLTX2Video.from_pretrained,
             model,
             subfolder="vae",
-            torch_dtype=dtype,
+            prefetch_list=ltx2_subfolders,
             local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
-        self.audio_vae = AutoencoderKLLTX2Audio.from_pretrained(
+        self.audio_vae = from_pretrained_with_prefetch(
+            AutoencoderKLLTX2Audio.from_pretrained,
             model,
             subfolder="audio_vae",
-            torch_dtype=dtype,
+            prefetch_list=ltx2_subfolders,
             local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
-        self.vocoder = LTX2Vocoder.from_pretrained(
+        self.vocoder = from_pretrained_with_prefetch(
+            LTX2Vocoder.from_pretrained,
             model,
             subfolder="vocoder",
-            torch_dtype=dtype,
+            prefetch_list=ltx2_subfolders,
             local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
 
         transformer_config = load_transformer_config(model, "transformer", local_files_only)

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py

@@ -45,6 +45,7 @@ from vllm_omni.diffusion.distributed.parallel_state import (
 )
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
@@ -149,26 +150,59 @@ class LTX23Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
             ),
         ]
 
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the whole component set before any from_pretrained.
+        ltx2_subfolders = [
+            "tokenizer",
+            "text_encoder",
+            "connectors",
+            "vae",
+            "audio_vae",
+            "vocoder",
+            "scheduler",
+        ]
+        prefetch_subfolders(model, ltx2_subfolders, local_files_only=local_files_only)
+
         # --- Tokenizer (lightweight, stays wherever) ---
         self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 
         # --- Text encoder: load on CPU, move to GPU only during encoding ---
         with torch.device("cpu"):
-            self.text_encoder = Gemma3ForConditionalGeneration.from_pretrained(
-                model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
+            self.text_encoder = from_pretrained_with_prefetch(
+                Gemma3ForConditionalGeneration.from_pretrained,
+                model,
+                subfolder="text_encoder",
+                prefetch_list=ltx2_subfolders,
+                local_files_only=local_files_only,
+                torch_dtype=dtype,
             )
 
         # --- Connectors: CPU (LTX-2.3 connectors include caption projection) ---
-        self.connectors = LTX2TextConnectors.from_pretrained(
-            model, subfolder="connectors", torch_dtype=dtype, local_files_only=local_files_only
+        self.connectors = from_pretrained_with_prefetch(
+            LTX2TextConnectors.from_pretrained,
+            model,
+            subfolder="connectors",
+            prefetch_list=ltx2_subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=dtype,
         )
 
         # --- VAE, Audio VAE: CPU ---
-        self.vae = AutoencoderKLLTX2Video.from_pretrained(
-            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKLLTX2Video.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=ltx2_subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=dtype,
         )
-        self.audio_vae = AutoencoderKLLTX2Audio.from_pretrained(
-            model, subfolder="audio_vae", torch_dtype=dtype, local_files_only=local_files_only
+        self.audio_vae = from_pretrained_with_prefetch(
+            AutoencoderKLLTX2Audio.from_pretrained,
+            model,
+            subfolder="audio_vae",
+            prefetch_list=ltx2_subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=dtype,
         )
 
         # --- Vocoder: prefer BWE vocoder (48kHz) for LTX-2.3 ---

vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py

@@ -32,6 +32,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.omnigen2.omnigen2_transformer import (
     OmniGen2RotaryPosEmbed,
@@ -673,23 +674,40 @@ class OmniGen2Pipeline(CFGParallelMixin, nn.Module, SupportsComponentDiscovery):
             )
         ]
 
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the whole component set before any from_pretrained.
+        omnigen2_subfolders = ["scheduler", "vae", "mllm", "processor"]
+        prefetch_subfolders(model, omnigen2_subfolders, local_files_only=local_files_only)
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
-            self.device
-        )
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKL.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=omnigen2_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
 
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, OmniGen2Transformer2DModel)
         self.transformer = OmniGen2Transformer2DModel(
             **transformer_kwargs,
             quant_config=od_config.quantization_config,
         )
-        self.mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained(
-            model, subfolder="mllm", local_files_only=local_files_only
+        self.mllm = from_pretrained_with_prefetch(
+            Qwen2_5_VLForConditionalGeneration.from_pretrained,
+            model,
+            subfolder="mllm",
+            prefetch_list=omnigen2_subfolders,
+            local_files_only=local_files_only,
         ).to(self.device)
-        self.processor = Qwen2_5_VLProcessor.from_pretrained(
-            model, subfolder="processor", local_files_only=local_files_only
+        self.processor = from_pretrained_with_prefetch(
+            Qwen2_5_VLProcessor.from_pretrained,
+            model,
+            subfolder="processor",
+            prefetch_list=omnigen2_subfolders,
+            local_files_only=local_files_only,
         )
         self.vae_scale_factor = (
             2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self, "vae") and self.vae is not None else 8

vllm_omni/diffusion/models/ovis_image/pipeline_ovis_image.py

@@ -38,6 +38,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.ovis_image.ovis_image_transformer import OvisImageTransformer2DModel
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.request import OmniDiffusionRequest
@@ -163,20 +164,32 @@ class OvisImagePipeline(nn.Module, CFGParallelMixin, DiffusionPipelineProfilerMi
         self._execution_device = get_local_device()
         model = od_config.model
         local_files_only = os.path.exists(model)
+
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the whole component set before any from_pretrained.
+        ovis_subfolders = ["scheduler", "text_encoder", "vae", "tokenizer"]
+        prefetch_subfolders(model, ovis_subfolders, local_files_only=local_files_only)
+
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
 
-        self.text_encoder = Qwen3Model.from_pretrained(
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen3Model.from_pretrained,
             model,
             subfolder="text_encoder",
+            prefetch_list=ovis_subfolders,
             local_files_only=local_files_only,
             torch_dtype=od_config.dtype,
         )
 
-        self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
-            self._execution_device
-        )
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKL.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=ovis_subfolders,
+            local_files_only=local_files_only,
+        ).to(self._execution_device)
 
         self.tokenizer = Qwen2TokenizerFast.from_pretrained(
             model, subfolder="tokenizer", local_files_only=local_files_only

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py

@@ -26,7 +26,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_qwenimage import DistributedAutoencoderKLQwenImage
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
@@ -279,18 +279,25 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
 
         # See pipeline_qwen_image_edit_plus: guard against transformers v5
         # multi-worker race on partial subfolder shard sets (Buildkite #1043).
+        qwen_subfolders = ["scheduler", "text_encoder", "vae", "tokenizer"]
         prefetch_subfolders(
             model,
-            ["scheduler", "text_encoder", "vae", "tokenizer"],
+            qwen_subfolders,
         )
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = retry_on_missing_shard(
-            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
-                model, subfolder="text_encoder", local_files_only=local_files_only
-            )
+        # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
+        # half-written cache (missing-shard ``OSError`` *and* the default
+        # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
+        # could not recover) instead of crashing the worker.
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen2_5_VLForConditionalGeneration.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=qwen_subfolders,
+            local_files_only=local_files_only,
         )
         # Qwen2.5-VL ships a vision tower that text-to-image does not use.
         # Drop it while the model is still on CPU, before moving to GPU, so
@@ -307,8 +314,12 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
         else:
             logger.warning("Qwen-Image: vision tower not found on text encoder; skipping drop")
         self.text_encoder = self.text_encoder.to(self.device)
-        self.vae = DistributedAutoencoderKLQwenImage.from_pretrained(
-            model, subfolder="vae", local_files_only=local_files_only
+        self.vae = from_pretrained_with_prefetch(
+            DistributedAutoencoderKLQwenImage.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=qwen_subfolders,
+            local_files_only=local_files_only,
         ).to(self.device)
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
         self.transformer = QwenImageTransformer2DModel(

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -27,7 +27,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
     QwenImageCFGParallelMixin,
@@ -247,28 +247,43 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi
 
         # See pipeline_qwen_image_edit_plus: guard against transformers v5
         # multi-worker race on partial subfolder shard sets (Buildkite #1043).
+        qwen_subfolders = ["scheduler", "text_encoder", "vae", "tokenizer", "processor"]
         prefetch_subfolders(
             model,
-            ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
+            qwen_subfolders,
         )
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = retry_on_missing_shard(
-            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
-                model, subfolder="text_encoder", local_files_only=local_files_only
-            ).to(self.device)
-        )
+        # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
+        # half-written cache (missing-shard ``OSError`` *and* the default
+        # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
+        # could not recover) instead of crashing the worker.
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen2_5_VLForConditionalGeneration.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=qwen_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
 
-        self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
-            self.device
-        )
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKLQwenImage.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=qwen_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
         self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.processor = retry_on_missing_shard(
-            lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only)
+        self.processor = from_pretrained_with_prefetch(
+            Qwen2VLProcessor.from_pretrained,
+            model,
+            subfolder="processor",
+            prefetch_list=qwen_subfolders,
+            local_files_only=local_files_only,
         )
 
         self.stage = None

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py

@@ -25,7 +25,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.model_metadata import QWEN_IMAGE_EDIT_PLUS_MAX_INPUT_IMAGES
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.cfg_parallel import (
@@ -216,29 +216,44 @@ class QwenImageEditPlusPipeline(
         # #1043 for Qwen/Qwen-Image-Edit-2509). snapshot_download takes a
         # per-repo file lock so the first worker downloads and the rest
         # wait on a warm cache before loading.
+        qwen_subfolders = ["scheduler", "text_encoder", "vae", "tokenizer", "processor"]
         prefetch_subfolders(
             model,
-            ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
+            qwen_subfolders,
         )
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = retry_on_missing_shard(
-            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
-                model, subfolder="text_encoder", local_files_only=local_files_only
-            ).to(self.device)
-        )
+        # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
+        # half-written cache (missing-shard ``OSError`` *and* the default
+        # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
+        # could not recover) instead of crashing the worker.
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen2_5_VLForConditionalGeneration.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=qwen_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
 
-        self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
-            self.device
-        )
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKLQwenImage.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=qwen_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
 
         transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, QwenImageTransformer2DModel)
         self.transformer = QwenImageTransformer2DModel(od_config=od_config, **transformer_kwargs)
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.processor = retry_on_missing_shard(
-            lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only)
+        self.processor = from_pretrained_with_prefetch(
+            Qwen2VLProcessor.from_pretrained,
+            model,
+            subfolder="processor",
+            prefetch_list=qwen_subfolders,
+            local_files_only=local_files_only,
         )
 
         self.stage = None

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -24,7 +24,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders, retry_on_missing_shard
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.qwen_image.autoencoder_kl_qwenimage import (
     AutoencoderKLQwenImage,
@@ -220,26 +220,41 @@ class QwenImageLayeredPipeline(nn.Module, SupportImageInput, QwenImageCFGParalle
 
         # See pipeline_qwen_image_edit_plus: guard against transformers v5
         # multi-worker race on partial subfolder shard sets (Buildkite #1043).
+        qwen_subfolders = ["scheduler", "text_encoder", "vae", "tokenizer", "processor"]
         prefetch_subfolders(
             model,
-            ["scheduler", "text_encoder", "vae", "tokenizer", "processor"],
+            qwen_subfolders,
         )
 
         # modules keep same as transformers & diffusers
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only
         )
-        self.text_encoder = retry_on_missing_shard(
-            lambda: Qwen2_5_VLForConditionalGeneration.from_pretrained(
-                model, subfolder="text_encoder", local_files_only=local_files_only
-            ).to(self.device)
-        )
-        self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to(
-            self.device
-        )
+        # ``from_pretrained_with_prefetch`` re-prefetches and retries on a
+        # half-written cache (missing-shard ``OSError`` *and* the default
+        # -config size-mismatch ``RuntimeError`` that ``retry_on_missing_shard``
+        # could not recover) instead of crashing the worker.
+        self.text_encoder = from_pretrained_with_prefetch(
+            Qwen2_5_VLForConditionalGeneration.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=qwen_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderKLQwenImage.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=qwen_subfolders,
+            local_files_only=local_files_only,
+        ).to(self.device)
         self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.processor = retry_on_missing_shard(
-            lambda: Qwen2VLProcessor.from_pretrained(model, subfolder="processor", local_files_only=local_files_only)
+        self.processor = from_pretrained_with_prefetch(
+            Qwen2VLProcessor.from_pretrained,
+            model,
+            subfolder="processor",
+            prefetch_list=qwen_subfolders,
+            local_files_only=local_files_only,
         )
 
         # modules re-implemented for vLLM-Omni

vllm_omni/diffusion/models/sd3/pipeline_sd3.py

@@ -19,7 +19,7 @@ from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl import Distribu
 from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.sd3.sd3_transformer import (
     SD3Transformer2DModel,
 )
@@ -156,18 +156,19 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
         # SD3.5 loads six subfolders in a row, each with multi-shard
         # safetensors - it is the worst-case fan-out for the race window.
+        sd3_subfolders = [
+            "scheduler",
+            "tokenizer",
+            "tokenizer_2",
+            "tokenizer_3",
+            "text_encoder",
+            "text_encoder_2",
+            "text_encoder_3",
+            "vae",
+        ]
         prefetch_subfolders(
             model,
-            [
-                "scheduler",
-                "tokenizer",
-                "tokenizer_2",
-                "tokenizer_3",
-                "text_encoder",
-                "text_encoder_2",
-                "text_encoder_3",
-                "vae",
-            ],
+            sd3_subfolders,
             local_files_only=local_files_only,
         )
 
@@ -198,31 +199,43 @@ class StableDiffusion3Pipeline(nn.Module, CFGParallelMixin, DiffusionPipelinePro
         self.tokenizer_3 = T5Tokenizer.from_pretrained(
             model, subfolder="tokenizer_3", local_files_only=local_files_only
         )
-        self.text_encoder = CLIPTextModelWithProjection.from_pretrained(
+        # ``from_pretrained_with_prefetch`` re-prefetches and retries if a
+        # peer worker left the cache half-written (missing-shard ``OSError`` or
+        # the default-config size-mismatch ``RuntimeError``) instead of
+        # crashing the worker - critical here given the six-subfolder fan-out.
+        self.text_encoder = from_pretrained_with_prefetch(
+            CLIPTextModelWithProjection.from_pretrained,
             model,
             subfolder="text_encoder",
-            torch_dtype=dtype,
+            prefetch_list=sd3_subfolders,
             local_files_only=local_files_only,
+            torch_dtype=dtype,
         )
-        self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(
+        self.text_encoder_2 = from_pretrained_with_prefetch(
+            CLIPTextModelWithProjection.from_pretrained,
             model,
             subfolder="text_encoder_2",
-            torch_dtype=dtype,
+            prefetch_list=sd3_subfolders,
             local_files_only=local_files_only,
+            torch_dtype=dtype,
         )
-        self.text_encoder_3 = T5EncoderModel.from_pretrained(
+        self.text_encoder_3 = from_pretrained_with_prefetch(
+            T5EncoderModel.from_pretrained,
             model,
             subfolder="text_encoder_3",
-            torch_dtype=dtype,
+            prefetch_list=sd3_subfolders,
             local_files_only=local_files_only,
+            torch_dtype=dtype,
         )
         self.transformer = SD3Transformer2DModel(od_config=od_config)
 
-        self.vae = DistributedAutoencoderKL.from_pretrained(
+        self.vae = from_pretrained_with_prefetch(
+            DistributedAutoencoderKL.from_pretrained,
             model,
             subfolder="vae",
-            torch_dtype=dtype,
+            prefetch_list=sd3_subfolders,
             local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
 
         self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8

vllm_omni/diffusion/models/stable_audio/pipeline_stable_audio.py

@@ -28,6 +28,7 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.interface import SupportAudioOutput
 from vllm_omni.diffusion.models.stable_audio.stable_audio_transformer import (
     StableAudioDiTModel,
@@ -109,6 +110,11 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil
             ),
         ]
 
+        # See ``hub_prefetch.py`` for the transformers v5 multi-worker subfolder
+        # race; prefetch the whole component set before any from_pretrained.
+        sa_subfolders = ["tokenizer", "text_encoder", "vae", "projection_model", "scheduler"]
+        prefetch_subfolders(model, sa_subfolders, local_files_only=local_files_only)
+
         # Load tokenizer
         self.tokenizer = T5TokenizerFast.from_pretrained(
             model,
@@ -117,27 +123,33 @@ class StableAudioPipeline(nn.Module, SupportAudioOutput, DiffusionPipelineProfil
         )
 
         # Load text encoder
-        self.text_encoder = T5EncoderModel.from_pretrained(
+        self.text_encoder = from_pretrained_with_prefetch(
+            T5EncoderModel.from_pretrained,
             model,
             subfolder="text_encoder",
-            torch_dtype=dtype,
+            prefetch_list=sa_subfolders,
             local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
 
         # Load VAE (AutoencoderOobleck for audio)
-        self.vae = AutoencoderOobleck.from_pretrained(
+        self.vae = from_pretrained_with_prefetch(
+            AutoencoderOobleck.from_pretrained,
             model,
             subfolder="vae",
-            torch_dtype=torch.float32,
+            prefetch_list=sa_subfolders,
             local_files_only=local_files_only,
+            torch_dtype=torch.float32,
         ).to(self.device)
 
         # Load projection model (using diffusers implementation)
-        self.projection_model = StableAudioProjectionModel.from_pretrained(
+        self.projection_model = from_pretrained_with_prefetch(
+            StableAudioProjectionModel.from_pretrained,
             model,
             subfolder="projection_model",
-            torch_dtype=dtype,
+            prefetch_list=sa_subfolders,
             local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
 
         # Initialize transformer from HF config to keep architecture aligned with checkpoint.

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -26,7 +26,7 @@ from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, Pipe
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
 from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
@@ -336,18 +336,39 @@ class Wan22Pipeline(
             )
 
         # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
+        component_subfolders = ["tokenizer", "text_encoder", "vae"]
         prefetch_subfolders(
             model,
-            ["tokenizer", "text_encoder", "vae"],
+            component_subfolders,
             local_files_only=local_files_only,
         )
 
-        self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.text_encoder = UMT5EncoderModel.from_pretrained(
-            model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
+        # ``from_pretrained_with_prefetch`` re-prefetches and retries if the
+        # cache is still half-written (the missing-shard ``OSError`` and the
+        # default-``UMT5Config`` size-mismatch ``RuntimeError`` seen on multi
+        # -worker HSDP / ring launches), instead of crashing the worker.
+        self.tokenizer = from_pretrained_with_prefetch(
+            AutoTokenizer.from_pretrained,
+            model,
+            subfolder="tokenizer",
+            prefetch_list=component_subfolders,
+            local_files_only=local_files_only,
+        )
+        self.text_encoder = from_pretrained_with_prefetch(
+            UMT5EncoderModel.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=component_subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
-        self.vae = DistributedAutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
+        self.vae = from_pretrained_with_prefetch(
+            DistributedAutoencoderKLWan.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=component_subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
 
         # Initialize transformers with correct config (weights loaded via load_weights)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py

@@ -26,7 +26,7 @@ from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, Pipe
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
-from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
+from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
 from vllm_omni.diffusion.models.interface import SupportImageInput
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero
@@ -197,35 +197,64 @@ class Wan22I2VPipeline(
                 )
             )
 
-        # Text encoder
-        self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
-        self.text_encoder = UMT5EncoderModel.from_pretrained(
-            model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
-        ).to(self.device)
-
         # Image encoder (CLIP) - optional, for Wan2.1-style I2V
         self.has_image_encoder = "image_encoder" in model_index and model_index["image_encoder"][0] is not None
 
-        # See ``hub_prefetch.py`` for the transformers v5 subfolder race.
+        # See ``hub_prefetch.py`` for the transformers v5 subfolder race. The
+        # prefetch MUST run before any ``from_pretrained`` below - previously
+        # the tokenizer / text_encoder were loaded ahead of this call, so they
+        # never benefited from the prefetch and could hit the half-written
+        # cache directly.
         subfolders = ["tokenizer", "text_encoder", "vae"]
         if self.has_image_encoder:
             subfolders.extend(["image_processor", "image_encoder"])
         prefetch_subfolders(model, subfolders, local_files_only=local_files_only)
 
+        # Text encoder
+        self.tokenizer = from_pretrained_with_prefetch(
+            AutoTokenizer.from_pretrained,
+            model,
+            subfolder="tokenizer",
+            prefetch_list=subfolders,
+            local_files_only=local_files_only,
+        )
+        self.text_encoder = from_pretrained_with_prefetch(
+            UMT5EncoderModel.from_pretrained,
+            model,
+            subfolder="text_encoder",
+            prefetch_list=subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=dtype,
+        ).to(self.device)
+
         if self.has_image_encoder:
-            self.image_processor = CLIPImageProcessor.from_pretrained(
-                model, subfolder="image_processor", local_files_only=local_files_only
+            self.image_processor = from_pretrained_with_prefetch(
+                CLIPImageProcessor.from_pretrained,
+                model,
+                subfolder="image_processor",
+                prefetch_list=subfolders,
+                local_files_only=local_files_only,
             )
-            self.image_encoder = CLIPVisionModel.from_pretrained(
-                model, subfolder="image_encoder", torch_dtype=dtype, local_files_only=local_files_only
+            self.image_encoder = from_pretrained_with_prefetch(
+                CLIPVisionModel.from_pretrained,
+                model,
+                subfolder="image_encoder",
+                prefetch_list=subfolders,
+                local_files_only=local_files_only,
+                torch_dtype=dtype,
             ).to(self.device)
         else:
             self.image_processor = None
             self.image_encoder = None
 
         # VAE
-        self.vae = DistributedAutoencoderKLWan.from_pretrained(
-            model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
+        self.vae = from_pretrained_with_prefetch(
+            DistributedAutoencoderKLWan.from_pretrained,
+            model,
+            subfolder="vae",
+            prefetch_list=subfolders,
+            local_files_only=local_files_only,
+            torch_dtype=dtype,
         ).to(self.device)
 
         # Transformers (weights loaded via load_weights)

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py

@@ -13,7 +13,7 @@ import logging
 import math
 import os
 from collections.abc import Iterable
-from typing import Any
+from typing import Any, ClassVar
 
 import numpy as np
 import PIL.Image
@@ -482,6 +482,7 @@ class Wan22S2VPipeline(
     _DEFAULT_MOTION_FRAMES = 73
     _DEFAULT_INFER_FRAMES = 80
     _DEFAULT_FPS = 16
+    dummy_run_num_frames: ClassVar[int] = 0
 
     def __init__(
         self,
@@ -929,21 +930,19 @@ class Wan22S2VPipeline(
     def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
         """Forward pass through the S2V transformer to predict noise.
 
-        The S2V model returns a list of tensors (one per batch element).
-        We take the first element since S2V processes one sample at a time.
+        With return_dict=False, the model returns (output,) where output is [B, C, T, H, W].
+        We squeeze batch dim since S2V processes one sample at a time.
         """
         if current_model is None:
             current_model = self.transformer
-        # WanModel_S2V's norm layers compute in float32; autocast ensures
-        # the float32→bfloat16 casts happen automatically (matching the
-        # original Wan2.2 inference path).
         param_dtype = next(current_model.parameters()).dtype
         with torch.amp.autocast(self.device.type, dtype=param_dtype):
             result = current_model(**kwargs)
-        # WanModel_S2V.forward returns a list of tensors
         if isinstance(result, list):
             return result[0]
-        return result[0] if isinstance(result, tuple) else result
+        if isinstance(result, tuple):
+            return result[0].squeeze(0)
+        return result.sample.squeeze(0)
 
     def diffuse(
         self,
@@ -996,32 +995,29 @@ class Wan22S2VPipeline(
                 latent_model_input = [latents.to(device)]
                 timestep = t.unsqueeze(0).to(device) if t.dim() == 0 else t.to(device)
 
-                # -- Positive (conditional) prediction kwargs --
                 positive_kwargs = {
-                    "x": latent_model_input,
-                    "t": timestep,
-                    "context": prompt_embeds[0:1],
-                    "seq_len": max_seq_len,
+                    "hidden_states": latent_model_input,
+                    "timestep": timestep,
+                    "encoder_hidden_states": prompt_embeds[0:1],
                     "cond_states": cond_latents,
                     "motion_latents": input_motion_latents,
                     "ref_latents": ref_latents,
-                    "motion_frames": motion_frames,
                     "drop_motion_frames": drop_first_motion,
-                    "audio_emb": positive_audio_emb,
+                    "encoder_hidden_states_audio": positive_audio_emb,
+                    "return_dict": False,
                 }
 
                 if do_true_cfg:
                     negative_kwargs = {
-                        "x": latent_model_input,
-                        "t": timestep,
-                        "context": negative_prompt_embeds[0:1],
-                        "seq_len": max_seq_len,
+                        "hidden_states": latent_model_input,
+                        "timestep": timestep,
+                        "encoder_hidden_states": negative_prompt_embeds[0:1],
                         "cond_states": cond_latents,
                         "motion_latents": input_motion_latents,
                         "ref_latents": ref_latents,
-                        "motion_frames": motion_frames,
                         "drop_motion_frames": drop_first_motion,
-                        "audio_emb": negative_audio_emb,
+                        "encoder_hidden_states_audio": negative_audio_emb,
+                        "return_dict": False,
                     }
                 else:
                     negative_kwargs = None
@@ -1174,7 +1170,7 @@ class Wan22S2VPipeline(
         ref_latents = self.encode_ref_image(image, height, width, device=device).to(dtype=dtype)
 
         # ---- 4. Initial motion latents (zeros) ----
-        motion_frames = self.motion_frames
+        motion_frames = min(self.motion_frames, infer_frames - 1)
         motion_pixels = torch.zeros(
             [1, 3, motion_frames, height, width],
             dtype=dtype,
@@ -1288,6 +1284,10 @@ class Wan22S2VPipeline(
             )
 
             # ---- Decode this clip ----
+            if self.od_config.enable_cpu_offload:
+                self.transformer.to("cpu")
+                current_omni_platform.empty_cache()
+
             latents_for_decode = latents.unsqueeze(0)  # [1, C, T, H, W]
             if not (drop_first_motion and r == 0):
                 decode_latents = torch.cat([motion_latents, latents_for_decode], dim=2)