fix: harden omni weight snapshot downloads

Serialize full-repo Hugging Face snapshot materialization with a repo-wide
node lock and retry incomplete snapshots so CI does not race on shared caches.

Co-authored-by: Cursor <cursoragent@cursor.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 with L4"
+      - label: ":full_moon: TTS · Function Test"
         timeout_in_minutes: 90
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
@@ -413,45 +413,6 @@ 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:
@@ -801,45 +762,6 @@ 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:
@@ -937,46 +859,6 @@ 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,42 +273,6 @@ 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:
@@ -399,29 +363,6 @@ 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,7 +16,6 @@ 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,9 +151,6 @@ 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]
@@ -173,12 +170,6 @@ 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,47 +64,6 @@ 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)
@@ -114,7 +73,10 @@ def parse_score_payload(raw_text: str) -> dict[str, Any]:
             parsed = {"score": json.loads(stripped), "reasoning": ""}
         else:
             raise
-    parsed["score"] = _normalize_score_list(parsed.get("score", []))
+    score = parsed.get("score", [])
+    if not isinstance(score, list):
+        score = [score]
+    parsed["score"] = [int(value) for value in score]
     return parsed
 
 
@@ -429,9 +391,6 @@ 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
@@ -441,9 +400,6 @@ 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(
@@ -522,9 +478,6 @@ 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 {
@@ -680,25 +633,6 @@ 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)
@@ -741,9 +675,6 @@ 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,7 +2,7 @@
 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" and "openai-image-edits-omni" backend.
+We currently only support using the "openai-chat-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:
@@ -52,11 +52,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
 - `--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", "ttfc", "tpoc", "icl", '
+                    'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
-                    '"tpop", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
-
-- `-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
@@ -280,9 +276,7 @@ Generate synthetic image、video、audio inputs alongside random text prompts to
 
 Notes:
 
-- Works only with online benchmark via:
+- Works only with online benchmark via the OpenAI backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
-  - 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):
 
@@ -363,105 +357,3 @@ 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 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.
+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.
 
 #### 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 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix)
+#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_28_x86_64.whl
+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
 ```
 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+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
+uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721
 
 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, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
+| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, T2V with sound | `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,8 +50,6 @@ 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` | ✅︎ | ✅︎ | | |
@@ -64,10 +62,6 @@ 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` | ✅︎ | ✅︎ | | |
@@ -75,13 +69,11 @@ 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,7 +145,6 @@ 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,10 +16,9 @@ Usage:
 """
 
 import argparse
-import json
+import os
 import time
 from pathlib import Path
-from typing import Any
 
 import torch
 
@@ -30,16 +29,6 @@ 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)."
@@ -98,12 +87,6 @@ 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,
@@ -195,10 +178,30 @@ def parse_args() -> argparse.Namespace:
         help="Enable cache-dit summary logging after diffusion forward passes.",
     )
     parser.add_argument(
-        "--profiler-config",
+        "--profile-dir",
-        type=parse_profiler_config,
+        type=str,
         default=None,
-        help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
+        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.",
     )
     return parser.parse_args()
 
@@ -236,9 +239,23 @@ def main():
         ulysses_degree=args.ulysses_degree,
     )
 
-    profiler_enabled = args.profiler_config is not None
+    # 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
 
-    omni_kwargs = dict(
+        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(
         model=args.model,
         model_class_name="WanS2VPipeline",
         flow_shift=args.flow_shift,
@@ -252,12 +269,8 @@ def main():
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_cache_dit_summary=args.enable_cache_dit_summary,
-        profiler_config=args.profiler_config,
+        profiler_config=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}")
@@ -312,16 +325,13 @@ 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, dict):
+        if profile_results and isinstance(profile_results, list):
-            traces = profile_results.get("traces", [])
             print("\n" + "=" * 60)
             print("PROFILING RESULTS:")
-            for rank, trace in enumerate(traces):
+            for i, result in enumerate(profile_results):
-                print(f"\nRank {rank}:")
+                print(f"\nStage {i}:")
-                if trace:
+                if result:
-                    print(f"  • Trace: {trace}")
+                    print(f"  {result}")
-            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

@@ -1,204 +0,0 @@
-# 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

@@ -1,40 +0,0 @@
-# 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

@@ -1,115 +0,0 @@
-# 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

@@ -1,32 +0,0 @@
-#!/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,9 +64,7 @@ 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

@@ -1,110 +0,0 @@
-# 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 |
-| [`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 |
+| [`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 |
-| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
+| [`nvidia/Cosmos3-Super.md`](./nvidia/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) | 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/inclusionAI/Ming-flash-omni-2.0.md

@@ -57,6 +57,12 @@ 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
@@ -66,15 +72,6 @@ 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

recipes/nvidia/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), action policy
+- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation, with optional synchronized audio (video + sound)
 - Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API
 - Maintainer: Community
 
@@ -23,23 +23,6 @@ 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
 
@@ -161,35 +144,6 @@ 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
@@ -198,7 +152,6 @@ curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_polic
   - 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)
@@ -220,10 +173,8 @@ curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_polic
     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 `forward_dynamics` (sync `/v1/videos/sync`) and `policy` (async
+  - Action (policy / forward- / inverse-dynamics) modalities are not part of
-    `/v1/videos`, returns the predicted action under the top-level `action`
+    this integration yet.
-    field) are supported online. **Online inference of inverse dynamics will be
-    added in a follow-up MR.**
 
 ### 1x GPU (Offline generation)
 

recipes/nvidia/Cosmos3-Super.md

@@ -104,9 +104,5 @@ 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, including the **action** modality: `forward_dynamics`
+  semantics as Nano. Action (policy / forward- / inverse-dynamics) modalities are
-  (sync `/v1/videos/sync`) and `policy` (async `/v1/videos`, predicted action under
+  not part of this integration yet.
-  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.

tests/benchmarks/test_accuracy_bench_utils.py

@@ -345,14 +345,6 @@ 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.async_tokens_to_discard == 1
+    assert session.discard_latest_async_tokens is True
     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 getattr(session, "async_tokens_to_discard", 0) == 0
+    assert session.discard_latest_async_tokens is False
     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_inline",
+        "test_name": "test_hunyuan_image3_ti2i_tp4_tp4_inline",
-        "description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config",
+        "description": "AR (TP=4) + DiT (TP=4) - 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": 17179869184,
+                                    "memory_pool_size": 4294967296,
                                     "memory_pool_device": "cpu"
                                 }
                             }
@@ -32,13 +32,13 @@
                                 "is_comprehension": true,
                                 "final_output": true,
                                 "final_output_type": "text",
-                                "max_num_seqs": 16,
+                                "max_num_seqs": 4,
                                 "gpu_memory_utilization": 0.95,
                                 "enforce_eager": false,
                                 "enable_prefix_caching": false,
                                 "max_num_batched_tokens": 32768,
-                                "devices": "0,1",
+                                "devices": "0,1,2,3",
-                                "tensor_parallel_size": 2,
+                                "tensor_parallel_size": 4,
                                 "hf_overrides": {
                                     "rope_parameters": {
                                         "mrope_section": [0, 32, 32],
@@ -65,14 +65,13 @@
                                 "stage_id": 1,
                                 "max_num_seqs": 1,
                                 "enforce_eager": true,
-                                "devices": "2,3",
+                                "devices": "4,5,6,7",
                                 "distributed_executor_backend": "mp",
                                 "omni_kv_config": {
                                     "need_recv_cache": true
                                 },
                                 "parallel_config": {
-                                    "tensor_parallel_size": 2,
+                                    "tensor_parallel_size": 4,
-                                    "vae_patch_parallel_size": 2,
                                     "enable_expert_parallel": true
                                 },
                                 "input_connectors": {
@@ -115,8 +114,8 @@
                 },
                 "task": "ti2i",
                 "bot-task": "think_recaption",
-                "max-concurrency": 16,
+                "max-concurrency": 4,
-                "num-prompts": 96,
+                "num-prompts": 16,
                 "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": [203, 283, 470, 869],
+                    "mean_ttft_ms": [1000, 2000, 3000, 4000],
-                    "mean_tpot_ms": [28, 36, 41, 45],
+                    "mean_tpot_ms": [200, 200, 200, 200],
-                    "mean_audio_ttfp_ms": [861, 1200, 1528, 2100],
+                    "mean_audio_ttfp_ms": [10000, 15000, 20000, 25000],
-                    "mean_audio_rtf": [0.24, 0.30, 0.35, 0.39]
+                    "mean_audio_rtf": [1.5, 1.5, 1.5, 1.5]
                 }
             },
             {
@@ -49,9 +49,9 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [377],
+                    "mean_ttft_ms": [2000],
-                    "mean_audio_ttfp_ms": [684],
+                    "mean_audio_ttfp_ms": [2000],
-                    "mean_audio_rtf": [0.14]
+                    "mean_audio_rtf": [0.25]
                 }
             },
             {
@@ -76,9 +76,9 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [228],
+                    "mean_ttft_ms": [6000],
-                    "mean_audio_ttfp_ms": [686],
+                    "mean_audio_ttfp_ms": [6000],
-                    "mean_audio_rtf": [0.16]
+                    "mean_audio_rtf": [0.7]
                 }
             },
             {
@@ -105,9 +105,9 @@
                 },
                 "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
                 "baseline": {
-                    "mean_ttft_ms": [522],
+                    "mean_ttft_ms": [12000],
-                    "mean_audio_ttfp_ms": [1119],
+                    "mean_audio_ttfp_ms": [12000],
-                    "mean_audio_rtf": [0.21]
+                    "mean_audio_rtf": [1.0]
                 }
             }
         ]

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

@@ -31,6 +31,11 @@
                 ],
                 "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": [
@@ -69,6 +74,11 @@
                 "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": [
@@ -107,6 +117,11 @@
                 "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": [
@@ -147,6 +162,11 @@
                 "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,16 +99,12 @@ 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]] = []
@@ -143,10 +139,7 @@ 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)
@@ -351,7 +344,7 @@ def test_pipeline_init_passes_tokenizer_attrs_into_transformer(
     assert pipeline.transformer.audio_proj_out.out_features == 5
 
 
-def test_preprocess_i2v_image_and_action_video_inputs() -> None:
+def test_preprocess_i2v_image_input() -> None:
     from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_pre_process_func
 
     preprocess = get_cosmos3_pre_process_func(SimpleNamespace())
@@ -364,16 +357,6 @@ def test_preprocess_i2v_image_and_action_video_inputs() -> 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
@@ -451,7 +434,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_sound_and_action(make_cosmos3_pipeline) -> None:
+def test_prepare_latents_for_video_image_and_sound(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)
@@ -480,20 +463,8 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline)
     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)
 
@@ -525,22 +496,20 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N
     )
     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, action_gen=True, action_dim=4)
+    pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
-    video_result, action_result = pipeline.diffuse(
+    video_result, sound_result = pipeline.diffuse(
         latents=latents,
-        action_latents=torch.zeros(1, 3, 4),
+        sound_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, "action_domain_ids": torch.tensor([0])},
+        shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
     )
     torch.testing.assert_close(video_result, torch.full_like(latents, 4.0))
-    torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result))
+    torch.testing.assert_close(sound_result, torch.full((), 24.0).expand_as(sound_result))
 
 
 def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -> None:
@@ -599,8 +568,6 @@ 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)
@@ -645,7 +612,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_sound_and_action_routes(self, make_cosmos3_pipeline) -> None:
+    def test_forward_i2v_and_sound_routes(self, make_cosmos3_pipeline) -> None:
         pipeline = make_cosmos3_pipeline()
         captured = self._install_forward_stubs(pipeline)
         image_tensor = torch.zeros(1, 3, 16, 16)
@@ -684,31 +651,6 @@ 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,9 +30,8 @@ def _tiny_cosmos3_config(**overrides):
     return config
 
 
-def test_mrope_position_ids_cover_text_video_sound_and_action() -> None:
+def test_mrope_position_ids_cover_text_video_and_sound() -> 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,
@@ -63,10 +62,6 @@ def test_mrope_position_ids_cover_text_video_sound_and_action() -> 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"),
@@ -131,7 +126,7 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
     assert tuple(output.shape) == (1, 2, 1, 2, 2)
 
 
-def test_sound_and_action_modules_follow_config() -> None:
+def test_sound_modules_follow_injected_sound_dim() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
 
     tiny = _tiny_cosmos3_config()
@@ -142,22 +137,12 @@ def test_sound_and_action_modules_follow_config() -> 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(
@@ -178,56 +163,33 @@ def test_transformer_requires_sound_dim_and_fps_when_sound_gen_true(kwargs: dict
         )
 
 
-def test_sound_and_action_pack_unpack_validate_shapes() -> None:
+def test_sound_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))
 
 
-@pytest.mark.parametrize(
+def test_forward_returns_video_and_sound_predictions(monkeypatch: pytest.MonkeyPatch) -> None:
-    ("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=config, dtype=torch.float32),
+        SimpleNamespace(
-        **transformer_kwargs,
+            tf_model_config=_tiny_cosmos3_config(),
+            dtype=torch.float32,
+        ),
+        sound_gen=True,
+        sound_dim=3,
+        sound_latent_fps=40.0,
     )(
         hidden_states=torch.zeros(1, 2, 1, 2, 2),
         timestep=torch.tensor([1.0]),
@@ -235,12 +197,11 @@ def test_forward_returns_video_plus_optional_modality_predictions(
         text_mask=torch.ones(1, 2, dtype=torch.long),
         video_shape=(1, 2, 2),
         fps=24.0,
-        action_noisy_mask=torch.ones(1, 5, 1),
+        sound_latents=torch.zeros(1, 3, 4),
-        **extra_kwargs,
     )
 
     assert isinstance(output, tuple)
-    assert [tuple(tensor.shape) for tensor in output] == expected_shapes
+    assert [tuple(tensor.shape) for tensor in output] == [(1, 2, 1, 2, 2), (1, 3, 4)]
 
 
 def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch: pytest.MonkeyPatch) -> None:
@@ -266,7 +227,7 @@ def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch
         )
 
 
-def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> None:
+def test_compute_rope_freqs_places_text_video_and_sound_positions() -> None:
     from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
 
     class FakeRotary:
@@ -314,11 +275,9 @@ def test_compute_rope_freqs_places_text_video_action_and_sound_positions() -> No
         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, 5)
+    assert gen_pos.shape == (3, 1, 3)
-    assert gen_pos[0, 0].tolist() == [102, 103, 103, 104, 102]
+    assert gen_pos[0, 0].tolist() == [102, 103, 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=14,
+# "Generated with seed=52, num_inference_steps=15,
-# prompt='A cute cat'"
+# prompt='A futuristic city skyline at twilight, cyberpunk style'"
 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 = 14) -> list:
+def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
     """Configure sampling parameters for Bagel text2img generation.
 
     Args:
@@ -269,6 +269,7 @@ 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,13 +24,14 @@ 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=14,
+# Generated with seed=52, num_inference_steps=15,
 # prompt='Change the grass color to red',
 # input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
 REFERENCE_PIXELS = [
@@ -46,10 +47,23 @@ 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=14,
+# "Generated with seed=52, num_inference_steps=15,
-# prompt='A cute cat'"
+# prompt='A futuristic city skyline at twilight, cyberpunk style'"
 TEXT2IMG_REFERENCE_PIXELS = [
     {"position": (100, 100), "rgb": (115, 113, 94)},
     {"position": (400, 50), "rgb": (159, 160, 144)},
@@ -63,6 +77,19 @@ 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
@@ -88,7 +115,7 @@ def _find_free_port() -> int:
     return port
 
 
-def _configure_sampling_params(omni: Omni, num_inference_steps: int = 14) -> list:
+def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
     """Configure sampling parameters for Bagel img2img generation.
 
     Args:
@@ -231,6 +258,7 @@ 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."""
@@ -248,6 +276,7 @@ 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.69** (``--min-daily-omni-accuracy`` / ``ACC_BENCH_MIN_DAILY_OMNI_ACCURACY``).
+   run **fails** if accuracy is strictly below **0.67** (``--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.69").strip() or "0.69"),
+        default=float((os.environ.get("ACC_BENCH_MIN_DAILY_OMNI_ACCURACY") or "0.67").strip() or "0.67"),
         help="Fail when daily_omni_accuracy is strictly below this threshold (0–1). "
-        "Default baseline 0.69; override with env ACC_BENCH_MIN_DAILY_OMNI_ACCURACY or pass 0 to disable the floor.",
+        "Default baseline 0.67; 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.35,
+        default=0.5,
         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

@@ -1,218 +0,0 @@
-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,7 +106,6 @@ 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,
@@ -168,7 +167,6 @@ 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,8 +32,7 @@ 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.
-# TODO: Fix this test
+pytestmark = [pytest.mark.full_model, pytest.mark.tts]
-# pytestmark = [pytest.mark.full_model, pytest.mark.tts]
 
 SAMPLE_RATE = 24_000  # All MOSS-TTS full variants output 24 kHz
 
@@ -129,6 +128,7 @@ 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,6 +140,7 @@ 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."""
@@ -151,6 +152,7 @@ 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."""
@@ -162,6 +164,7 @@ 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."""
@@ -203,6 +206,7 @@ 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."""
@@ -214,6 +218,7 @@ 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,6 +120,7 @@ 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."""
@@ -132,6 +133,7 @@ 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."""
@@ -144,6 +146,7 @@ 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."""
@@ -156,6 +159,7 @@ 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,11 +24,10 @@ from vllm import SamplingParams
 from tests.helpers.mark import hardware_test
 from vllm_omni import Omni
 
-# H100-gated (8B). Picked up by the nightly "TTS · Function Test with H100"
+# H100-gated (8B). Picked up by the nightly H100 omni job
-# job (-m "full_model and H100 and tts"); the L4 MOSS-TTS tests land in the
+# (-m "full_model and H100 and omni"); the tts mark keeps it consistent with
-# sibling L4 tts lane instead.
+# the other MOSS-TTS tests and any tts-scoped sweeps.
-# TODO: Fix this test
+pytestmark = [pytest.mark.full_model, pytest.mark.tts]
-# 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"
@@ -102,6 +101,7 @@ 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

@@ -1,183 +0,0 @@
-"""
-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

@@ -1,388 +0,0 @@
-# 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,8 +30,7 @@ from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServerParams
 from tests.helpers.stage_config import get_deploy_config_path
 
-# TODO: Fix this test
+pytestmark = [pytest.mark.full_model, pytest.mark.tts]
-# 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,40 +9,63 @@ 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
+from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
+from vllm_omni.platforms import current_omni_platform
 
 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")
 
 
-# For prefix caching checks against we enable it on the thinker and talker via CLI override
+def get_chunk_config(config_path: str | None = None):
-# and enable prompt token details so that we can determine if any tokens were cached.
+    """Load the qwen3_omni CI deploy yaml with async_chunk modifications for streaming mode."""
-# We also explicitly set block size so that we can make sure the cached token counts are a
+    if config_path is None:
-# multiple of the block size.
+        config_path = _CI_DEPLOY
-BLOCK_SIZE = 16
+    # 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 = [
-    pytest.param(
+    OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
-        OmniServerParams(
+]
-            model=_MODEL,
+# For prefix caching, we enable it on the thinker (stage 0) via CLI override
-            stage_config_path=_CI_DEPLOY,
+# and enable prompt token details so that we can determine if any tokens were cached.
-            use_stage_cli=True,
+BLOCK_SIZE = 16
-            server_args=[
+prefix_test_params = [
-                "--no-async-chunk",
+    OmniServerParams(
-                "--block-size",
+        model=model,
-                str(BLOCK_SIZE),
+        stage_config_path=_CI_DEPLOY,
-                "--stage-overrides",
+        server_args=[
-                '{"0": {"enable_prefix_caching": true}, "1": {"enable_prefix_caching": true}}',
+            "--block-size",
-                "--enable-prompt-tokens-details",
+            str(BLOCK_SIZE),
-            ],
+            "--stage-overrides",
-        ),
+            '{"0": {"enable_prefix_caching": true}}',
-        id="default",
+            "--enable-prompt-tokens-details",
+        ],
     )
+    for model in models
 ]
 
 
@@ -143,14 +166,45 @@ def test_text_to_text_001(omni_server, openai_client) -> None:
     openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
 
 
-def _run_prefix_cache_check(openai_client, request_config: dict):
+@pytest.mark.advanced_model
-    """Make two requests given a request config, and validate that:
+@pytest.mark.core_model
-    1. The second request actually had cached tokens
+@pytest.mark.omni
-    2. The number of cached tokens is divisible by the block size used in
+@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
-    test_params, because currently upstream vLLM does not cache partial
+@pytest.mark.parametrize("omni_server", prefix_test_params, indirect=True)
-    blocks.
+def test_thinker_prefix_caching(omni_server, openai_client, run_level) -> None:
     """
-    openai_client.send_omni_request(request_config, request_num=1)[0]
+    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]
     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
@@ -163,66 +217,16 @@ def _run_prefix_cache_check(openai_client, request_config: dict):
     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
 
-@pytest.mark.advanced_model
+    if run_level == "advanced_model":
-@pytest.mark.core_model
+        # For each token index where both responses have an output, ensure that the greedy token
-@pytest.mark.omni
+        # predicted in the uncached case is in the top k logprobs for the cached case
-@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
+        for idx in range(n_tokens):
-@pytest.mark.parametrize("omni_server", test_params, indirect=True)
+            greedy_token = uncached_response.logprobs[idx].token
-def test_thinker_prefix_caching_text_output(omni_server, openai_client) -> None:
+            cached_top_k = {lp.token for lp in cached_response.logprobs[idx].top_logprobs}
-    """
+            assert greedy_token in cached_top_k
-    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,7 +427,6 @@ 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(),
         )
     )
@@ -985,7 +984,6 @@ 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, **kwargs):
+    async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
         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, **kwargs):
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
             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, **kwargs):
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
             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, **kwargs):
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
             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, **kwargs):
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
             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, **kwargs):
+        async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
             self.captured_sampling_params_list = (
                 sampling_params_list if sampling_params_list is not None else [sampling_params]
             )
@@ -815,9 +815,7 @@ 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(
+    async def generate_multi_output_delta(prompt, request_id, sampling_params=None, sampling_params_list=None):
-        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,
@@ -845,7 +843,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, **kwargs):
+    async def generate_without_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
         yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
 
     streaming_image_edit_client.app.state.engine_client.generate = generate_without_image
@@ -868,7 +866,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, **kwargs):
+    async def generate_empty_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
         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,11 +2,6 @@
 # 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
@@ -39,32 +34,18 @@ REALTIME_SYNTH_PHRASE_TEXT = (
     "Translate into Chinese: Beijing is the Capital of China. It is the center of culture and politics"
 )
 
-# Simulate realtime upload pacing (``openai_realtime_client.py --send-delay-ms``).
+# The new-schema CI overlay bakes in async_chunk: False and covers CUDA/ROCm/XPU
-SEND_DELAY_MS = 200
+# via its ``platforms:`` section, so one path serves all three.
-
-# 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_sync_server_params = [
+realtime_server_params = [
     pytest.param(
         OmniServerParams(
             model=MODEL,
             stage_config_path=default_stage_config,
             use_stage_cli=True,
         ),
-        id="sync",
+        id="default",
-    ),
-]
-
-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",
     ),
 ]
 
@@ -99,7 +80,6 @@ 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] = []
@@ -125,8 +105,6 @@ 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}))
 
@@ -178,92 +156,26 @@ 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.core_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_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(
-                omni_server.host,
-                omni_server.port,
-                omni_server.model,
-                pcm16,
-                chunk_ms=100,
-                send_delay_ms=0,
-            )
-        )
-
-        _assert_realtime_smoke(result)
-
     @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)
+    @pytest.mark.parametrize("omni_server", realtime_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()
-
-        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()
+        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)
 
         result = asyncio.run(
             _run_realtime_audio_roundtrip(
@@ -272,9 +184,25 @@ class TestQwen3OmniRealtimeWebSocket:
                 omni_server.model,
                 pcm16,
                 chunk_ms=100,
-                send_delay_ms=0,
             )
         )
 
-        _assert_realtime_smoke(result)
+        out_pcm = result["output_pcm"]
-        _assert_realtime_accuracy(result)
+        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
+
+        final_text = (result["transcription_text"] or "").strip()
+        assert final_text, "Expected non-empty transcription (model text stream)"
+
+        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}"
+        )

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, "vae_generator_seed": 0}
+    assert engine_prompt["mm_processor_kwargs"] == {"target_h": 768, "target_w": 1024}
     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,115 +627,6 @@ 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():
@@ -1147,31 +1038,6 @@ 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(*args, **kw):
+async def _noop(**kw):
     pass
 
 
@@ -61,7 +61,6 @@ 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,6 +135,22 @@ 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():
@@ -195,18 +211,6 @@ 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,12 +888,6 @@ 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)
@@ -1588,8 +1582,6 @@ 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,7 +154,6 @@ 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}),
@@ -183,7 +182,6 @@ 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,7 +5,6 @@
 from __future__ import annotations
 
 import functools
-from types import SimpleNamespace
 
 import pytest
 
@@ -18,14 +17,13 @@ def _voxcpm2_talker_mod():
     from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import (
         VoxCPM2TalkerForConditionalGeneration,
         _RequestState,
-        _VoxCPM2RuntimeConfig,
     )
 
-    return VoxCPM2TalkerForConditionalGeneration, _RequestState, _VoxCPM2RuntimeConfig
+    return VoxCPM2TalkerForConditionalGeneration, _RequestState
 
 
 def _make_bare_talker():
-    VoxCPM2TalkerForConditionalGeneration, _, _ = _voxcpm2_talker_mod()
+    VoxCPM2TalkerForConditionalGeneration, _ = _voxcpm2_talker_mod()
     talker = VoxCPM2TalkerForConditionalGeneration.__new__(VoxCPM2TalkerForConditionalGeneration)
     talker._active_states = {}
     talker._current_request_id = None
@@ -36,27 +34,11 @@ 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
@@ -65,23 +47,6 @@ 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()
 
@@ -96,61 +61,6 @@ 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")
@@ -209,7 +119,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}")
 
@@ -219,548 +129,3 @@ 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,6 +1638,9 @@ 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."""
@@ -1957,3 +1960,5 @@ 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

@@ -1,910 +0,0 @@
-# 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,14 +44,6 @@ 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

@@ -0,0 +1,16 @@
+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,6 +1,5 @@
 import warnings
-from collections import defaultdict
+from dataclasses import dataclass
-from dataclasses import field, make_dataclass
 
 import numpy as np
 from vllm.benchmarks.datasets import SampleRequest
@@ -10,172 +9,30 @@ 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]
 
-_MULTIMODAL_BENCHMARK_FIELDS = [
+@dataclass
-    (defs.MEAN_AUDIO_TTFP_MS, float, field(default=0.0)),
+class MultiModalsBenchmarkMetrics(BenchmarkMetrics):
-    (defs.MEDIAN_AUDIO_TTFP_MS, float, field(default=0.0)),
+    mean_audio_ttfp_ms: float = 0.0
-    (defs.STD_AUDIO_TTFP_MS, float, field(default=0.0)),
+    median_audio_ttfp_ms: float = 0.0
-    (defs.PERCENTILES_AUDIO_TTFP_MS, _PERCENTILE_ROWS_TYPE, field(default=None)),
+    std_audio_ttfp_ms: float = 0.0
-    (defs.TOTAL_AUDIO_DURATION_S, float, field(default=0.0)),
+    percentiles_audio_ttfp_ms: list[tuple[float, float]] = None
-    (defs.TOTAL_AUDIO_FRAMES, int, field(default=0)),
+    total_audio_duration_s: float = 0.0
-    (defs.AUDIO_THROUGHPUT, float, field(default=0.0)),
+    total_audio_frames: int = 0
-    (defs.MEAN_AUDIO_RTF, float, field(default=0.0)),
+    audio_throughput: float = 0.0
-    (defs.MEDIAN_AUDIO_RTF, float, field(default=0.0)),
+    mean_audio_rtf: float = 0.0
-    (defs.STD_AUDIO_RTF, float, field(default=0.0)),
+    median_audio_rtf: float = 0.0
-    (defs.PERCENTILES_AUDIO_RTF, _PERCENTILE_ROWS_TYPE, field(default=None)),
+    std_audio_rtf: float = 0.0
-    (defs.MEAN_AUDIO_DURATION_S, float, field(default=0.0)),
+    percentiles_audio_rtf: list[tuple[float, float]] = None
-    (defs.MEDIAN_AUDIO_DURATION_S, float, field(default=0.0)),
+    mean_audio_duration_s: float = 0.0
-    (defs.STD_AUDIO_DURATION_S, float, field(default=0.0)),
+    median_audio_duration_s: float = 0.0
-    (defs.PERCENTILES_AUDIO_DURATION_S, _PERCENTILE_ROWS_TYPE, field(default=None)),
+    std_audio_duration_s: float = 0.0
-    (defs.MEAN_AUDIO_UNDERRUN_S, float, field(default=0.0)),
+    percentiles_audio_duration_s: list[tuple[float, float]] = None
-    (defs.MEDIAN_AUDIO_UNDERRUN_S, float, field(default=0.0)),
+    # Streaming-audio continuity (see vllm_omni/benchmarks/audio_continuity.py).
-    (defs.STD_AUDIO_UNDERRUN_S, float, field(default=0.0)),
+    mean_audio_underrun_s: float = 0.0
-    (defs.PERCENTILES_AUDIO_UNDERRUN_S, _PERCENTILE_ROWS_TYPE, field(default=None)),
+    median_audio_underrun_s: float = 0.0
-    (defs.AUDIO_CONTINUITY_OK_RATE, float, field(default=1.0)),
+    std_audio_underrun_s: float = 0.0
-    (defs.TOTAL_IMAGES, int, field(default=0)),
+    percentiles_audio_underrun_s: list[tuple[float, float]] = None
-    (defs.IMAGE_THROUGHPUT, float, field(default=0.0)),
+    audio_continuity_ok_rate: float = 1.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(
@@ -186,10 +43,6 @@ 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))
@@ -208,19 +61,7 @@ def print_metrics(
         process_one_metric("e2el", metrics)
     print_text_metrics(task_type, selected_percentile_metrics, metrics)
     if task_type == TaskType.GENERATION:
-        if _has_audio_output(metrics):
+        print_audio_metrics(selected_percentile_metrics, 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)
 
 
@@ -238,36 +79,22 @@ 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
-        if has_text_output:
+        for metric in selected_percentile_metrics:
-            if _wants_text_ttft(selected_percentile_metrics):
+            if metric == "e2el":
-                process_one_metric("ttft", metrics)
+                continue
-            if _wants_text_tpot(selected_percentile_metrics):
+            if not metric.startswith("audio") and has_text_output:
-                process_one_metric("tpot", metrics)
+                process_one_metric(metric, 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(
+    print("{:<40} {:<10.2f}".format("Total audio duration generated(s):", metrics.total_audio_duration_s))
-        "{:<40} {:<10.2f}".format("Total audio duration generated(s):", getattr(metrics, defs.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}".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:",
-            getattr(metrics, defs.AUDIO_CONTINUITY_OK_RATE),
+            metrics.audio_continuity_ok_rate,
         )
     )
     for metric in selected_percentile_metrics:
@@ -275,37 +102,6 @@ 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,
@@ -335,16 +131,15 @@ def process_one_metric(
     elif is_audio_rtf:
         suffix = ""
         unit_suffix = ""
-    field_names = _AGGREGATE_PERCENTILE_FIELD_NAMES.get(metric_attribute_name)
+    mean_attr_name = f"mean_{metric_attribute_name}{suffix}"
-    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 = field_names[1] if field_names else f"median_{metric_attribute_name}{suffix}"
+    median_attr_name = 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 = field_names[2] if field_names else f"percentiles_{metric_attribute_name}{suffix}"
+    percentiles_attr_name = f"percentiles_{metric_attribute_name}{suffix}"
     percentiles = getattr(metrics, percentiles_attr_name, [])
 
     for percentile, value in percentiles:
@@ -353,356 +148,6 @@ 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],
@@ -715,7 +160,6 @@ def calculate_metrics(
     max_concurrency,
     request_rate,
     benchmark_duration,
-    print_stage: bool = False,
 ) -> tuple[BenchmarkMetrics, list[int]]:
     """Calculate the metrics for the benchmark.
 
@@ -743,10 +187,6 @@ 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
@@ -782,15 +222,6 @@ 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)
@@ -905,6 +336,26 @@ 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,
@@ -920,46 +371,6 @@ 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,
@@ -969,8 +380,5 @@ 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,7 +2,6 @@ import asyncio
 import contextlib
 import io
 import json
-import mimetypes
 import os
 import random
 import ssl
@@ -13,7 +12,7 @@ import wave
 from collections.abc import Iterable
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any, Literal
+from typing import Literal
 
 import aiohttp
 import pybase64 as base64
@@ -49,41 +48,9 @@ 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"
-RETURN_STAGE_METRICS_FIELD = "return_stage_metrics"
+_DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S = 0.1
-_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:
@@ -95,7 +62,7 @@ def _audio_continuity_threshold_s() -> float:
     """
     raw = os.environ.get(_AUDIO_CONTINUITY_THRESHOLD_ENV)
     if not raw:
-        return defs.AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S
+        return _DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S
     try:
         value = float(raw)
     except ValueError:
@@ -103,9 +70,9 @@ def _audio_continuity_threshold_s() -> float:
             "Invalid %s=%r; using default %.3fs",
             _AUDIO_CONTINUITY_THRESHOLD_ENV,
             raw,
-            defs.AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S,
+            _DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S,
         )
-        return defs.AUDIO_CONTINUITY_DEFAULT_THRESHOLD_S
+        return _DEFAULT_AUDIO_CONTINUITY_THRESHOLD_S
     return max(value, 0.0)
 
 
@@ -385,10 +352,6 @@ 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``
@@ -403,200 +366,6 @@ 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(
@@ -679,13 +448,6 @@ 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:
@@ -757,27 +519,9 @@ 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
 
-                                _update_output_stage_metrics_from_payload(output, data)
+                                if metrics := data.get("metrics"):
-                                (
+                                    output.output_tokens = metrics.get("num_tokens_out", 0)
-                                    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:
@@ -866,138 +610,6 @@ 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:
@@ -1115,10 +727,6 @@ 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
@@ -1131,10 +739,7 @@ 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 (
+from vllm_omni.benchmarks.metrics.metrics import MultiModalsBenchmarkMetrics, calculate_metrics
-    MultiModalsBenchmarkMetrics,
-    calculate_metrics,
-)
 
 # ruff: noqa: E402
 
@@ -1391,7 +996,6 @@ 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(
@@ -1412,13 +1016,9 @@ 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,
-            defs.TOTAL_AUDIO_DURATION_S: getattr(metrics, defs.TOTAL_AUDIO_DURATION_S),
+            "total_audio_duration_s": metrics.total_audio_duration_s,
-            defs.TOTAL_AUDIO_FRAMES: getattr(metrics, defs.TOTAL_AUDIO_FRAMES),
+            "total_audio_frames": metrics.total_audio_frames,
-            defs.AUDIO_THROUGHPUT: getattr(metrics, defs.AUDIO_THROUGHPUT),
+            "audio_throughput": metrics.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],
@@ -1474,34 +1074,21 @@ 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 result_percentile_metrics:
+        if metric_attribute_name not in selected_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 == defs.AUDIO_RTF
+        is_audio_rtf = metric_attribute_name == "audio_rtf"
-        is_audio_duration_or_underrun = metric_attribute_name in (defs.AUDIO_DURATION, defs.AUDIO_UNDERRUN)
+        is_audio_duration_or_underrun = metric_attribute_name in ("audio_duration", "audio_underrun")
 
         suffix = "_ms"
         if is_audio_duration_or_underrun:
@@ -1520,10 +1107,9 @@ async def benchmark(
             result[f"p{p_word}_{metric_attribute_name}{suffix}"] = value
 
     if task_type == TaskType.GENERATION:
-        for metric in result_percentile_metrics:
+        for metric in selected_percentile_metrics:
             process_one_metric(metric)
     else:
-        result_percentile_metrics.append("e2el")
         process_one_metric("e2el")
 
     if profile:

vllm_omni/benchmarks/serve.py

@@ -9,11 +9,6 @@ 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]:
@@ -23,9 +18,4 @@ 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,10 +149,6 @@ _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,33 +391,27 @@ class OmniTensorPrefixCache:
         self,
         query_start_loc: torch.Tensor,
         input_batch: InputBatch,
-        multimodal_outputs: dict | None,
+        multimodal_outputs: dict,
         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 mm_outputs:
+            if mm_key in multimodal_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=mm_outputs[mm_key],
+                    hidden_states=multimodal_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(mm_outputs.keys()) - self.mm_cache_keys
+        passthrough_keys = set(multimodal_outputs.keys()) - self.mm_cache_keys
-        passthrough_mm_data = {k: v for k, v in mm_outputs.items() if k in passthrough_keys}
+        passthrough_mm_data = {k: v for k, v in multimodal_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, FinishReason
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
 from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
@@ -69,8 +69,6 @@ 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()
 
@@ -436,7 +434,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.async_tokens_to_discard = 1
+                    request.discard_latest_async_tokens = True
                     request.num_output_placeholders = 0
                     request.spec_token_ids = []
                     request._output_token_ids.clear()
@@ -558,24 +556,16 @@ 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():
-                eco = engine_core_outputs.get(client_index)
+                # Set finished request set in EngineCoreOutputs for this client.
-                if eco is None:
+                if (eco := engine_core_outputs.get(client_index)) is not None:
-                    eco = EngineCoreOutputs()
+                    eco.finished_requests = finished_set
-                    engine_core_outputs[client_index] = eco
+                else:
-                emitted = {o.request_id for o in eco.outputs}
+                    engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
-                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:
@@ -642,12 +632,13 @@ 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.async_tokens_to_discard = 1
+            session.discard_latest_async_tokens = True
             session.num_computed_tokens -= session.num_output_placeholders
             session.num_output_placeholders = 0
             session.spec_token_ids = []

vllm_omni/data_entry_keys.py

@@ -63,7 +63,6 @@ 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
@@ -151,7 +150,6 @@ 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

@@ -1,74 +0,0 @@
-# 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,7 +18,6 @@ connectors:
   connector_of_shared_memory:
     name: SharedMemoryConnector
     extra:
-      initial_codec_chunk_frames: 4
       codec_chunk_frames: 25
       codec_left_context_frames: 25
 
@@ -31,8 +30,12 @@ stages:
     enable_prefix_caching: false
     devices: "0"
     default_sampling_params:
-      temperature: 0.0
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
       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,7 +18,6 @@ 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,13 +1,19 @@
 # VoxCPM2 deploy: single-stage AR pipeline with per-request state batching.
 # Verified on 1x H20 141GB.
 #
-# VoxCPM2-local runtime knobs are passed through hf_overrides.
+# 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.
 async_chunk: false
 dtype: bfloat16
 
 stages:
   - stage_id: 0
-    max_num_seqs: 32
+    max_num_seqs: 4
     gpu_memory_utilization: 0.9
     enforce_eager: true
     trust_remote_code: true
@@ -16,18 +22,6 @@ 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,80 +242,6 @@ 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.
 
@@ -1060,151 +986,6 @@ 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.
@@ -2027,7 +1808,6 @@ 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,8 +96,6 @@ 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"))
@@ -724,9 +722,6 @@ 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,7 +52,6 @@ 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,7 +5,6 @@ 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__)
 
@@ -301,234 +300,3 @@ 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/models/bagel/pipeline_bagel.py

@@ -777,9 +777,7 @@ class BagelPipeline(nn.Module, SupportsComponentDiscovery, DiffusionPipelineProf
             if torch.is_tensor(v):
                 generation_input[k] = v.to(self.device)
 
-        # NOTE: For now we disable device specific noise regeneration so that e2e tests can run
+        self._regen_init_noise_on_device(generation_input, req.sampling_params.seed)
-        # 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

@@ -1,214 +0,0 @@
-# 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,19 +46,6 @@ 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__)
@@ -111,67 +98,14 @@ 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 action preprocessing expected PIL image or image path, got {type(value)!r}.")
+        raise TypeError(f"Cosmos3 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", "")
@@ -182,63 +116,39 @@ 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")
-            raw_video = multi_modal_data.get("video")
+            if raw_image is None:
-            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"] = {}
 
-            if raw_image is None:
+            image = _pil_to_rgb(raw_image)
-                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:
-                if action_mode is not None:
+                max_area = 720 * 1280
-                    _set_action_size_from_image(request, image)
+                aspect_ratio = image.height / image.width
-                else:
+                mod_value = 16
-                    max_area = 720 * 1280
+                height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
-                    aspect_ratio = image.height / image.width
+                width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
-                    mod_value = 16
+                if request.sampling_params.height is None:
-                    height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
+                    request.sampling_params.height = height
-                    width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
+                if request.sampling_params.width is None:
-                    if request.sampling_params.height is None:
+                    request.sampling_params.width = width
-                        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
-            if action_mode is not None:
+            scale = max(target_w / image.width, target_h / image.height)
-                prompt["additional_information"]["preprocessed_image"] = _preprocess_action_image(
+            resize_w = int(np.ceil(scale * image.width))
-                    image,
+            resize_h = int(np.ceil(scale * image.height))
-                    int(target_h),
+            image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS)
-                    int(target_w),
+            left = (resize_w - target_w) // 2
-                )
+            top = (resize_h - target_h) // 2
-            else:
+            image = image.crop((left, top, left + target_w, top + target_h))
-                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(
+            prompt["additional_information"]["preprocessed_image"] = video_processor.preprocess(
-                    image, height=target_h, width=target_w
+                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
@@ -511,17 +421,11 @@ 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."):
@@ -624,22 +528,13 @@ 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, action, sound order for multimodal generation.
+        or a tuple in video, sound order for sound generation.
         """
         return self.transformer(**kwargs)
 
@@ -716,12 +611,6 @@ 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
@@ -1128,30 +1017,6 @@ 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,
@@ -1188,95 +1053,6 @@ 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(
@@ -1290,15 +1066,11 @@ 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).
 
@@ -1343,13 +1115,10 @@ 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]
@@ -1369,86 +1138,57 @@ class Cosmos3OmniDiffusersPipeline(
 
         def _split_noise_pred(
             noise_pred: torch.Tensor | tuple[torch.Tensor, ...],
-        ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        ) -> tuple[torch.Tensor, torch.Tensor | None]:
-            has_action = action_latents is not None
             has_sound = sound_latents is not None
-            if not has_action and not has_sound:
+            if not has_sound:
                 if isinstance(noise_pred, tuple):
                     raise ValueError("Cosmos3 video-only diffusion received tuple predictions.")
-                return noise_pred, None, None
+                return noise_pred, None
             if not isinstance(noise_pred, tuple):
                 raise ValueError("Cosmos3 multimodal diffusion expects transformer predictions as a tuple.")
-            expected = 1 + int(has_action) + int(has_sound)
+            if len(noise_pred) != 2:
-            if len(noise_pred) != expected:
+                raise ValueError(f"Cosmos3 sound diffusion expected 2 predictions, got {len(noise_pred)}.")
-                raise ValueError(
+            return noise_pred[0], noise_pred[1]
-                    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, action_pred, sound_pred = _split_noise_pred(noise_pred)
+            video_pred, sound_pred = _split_noise_pred(noise_pred)
             if velocity_mask is not None:
                 video_pred = video_pred * velocity_mask
-            if action_pred is not None and action_velocity_mask is not None:
+            if sound_latents is 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, action_pred, sound_pred)
+                packed_noise, shapes, numels = _pack_joint(video_pred, sound_pred)
-                packed_latents, _, _ = _pack_joint(latents, action_latents, sound_latents)
+                packed_latents, _, _ = _pack_joint(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[idx]
+                    sound_latents = unpacked[1]
             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, action_latents, sound_latents
+            nonlocal latents, sound_latents
-            if action_latents is None and sound_latents is None:
+            if 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[idx]
+                sound_latents = step_out[1]
 
         if cfg_parallel:
             for t in self.progress_bar(timesteps):
@@ -1466,7 +1206,6 @@ class Cosmos3OmniDiffusersPipeline(
                         timestep=timestep,
                         text_ids=cond_ids,
                         text_mask=cond_mask,
-                        action_latents=action_latents,
                         sound_latents=sound_latents,
                         **shared_kwargs,
                     ),
@@ -1475,13 +1214,12 @@ 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, action_latents, sound_latents))
+                _assign_step_out(_step(noise_pred, t, latents, sound_latents))
 
         elif do_cfg:
             cond_cache: tuple = (None, None)
@@ -1499,7 +1237,6 @@ class Cosmos3OmniDiffusersPipeline(
                     timestep=timestep,
                     text_ids=cond_ids,
                     text_mask=cond_mask,
-                    action_latents=action_latents,
                     sound_latents=sound_latents,
                     **shared_kwargs,
                 )
@@ -1513,7 +1250,6 @@ class Cosmos3OmniDiffusersPipeline(
                         timestep=timestep,
                         text_ids=uncond_ids,
                         text_mask=uncond_mask,
-                        action_latents=action_latents,
                         sound_latents=sound_latents,
                         **shared_kwargs,
                     )
@@ -1527,7 +1263,7 @@ class Cosmos3OmniDiffusersPipeline(
                 else:
                     noise_pred = noise_cond
 
-                _assign_step_out(_step(noise_pred, t, latents, action_latents, sound_latents))
+                _assign_step_out(_step(noise_pred, t, latents, sound_latents))
 
         else:
             for t in self.progress_bar(timesteps):
@@ -1537,15 +1273,12 @@ 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, action_latents, sound_latents))
+                _assign_step_out(_step(noise_pred, t, 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)
@@ -1567,29 +1300,15 @@ 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 "
@@ -1630,36 +1349,6 @@ 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))
@@ -1672,23 +1361,6 @@ 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
 
@@ -1724,58 +1396,7 @@ 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.
-        action_latents = None
+        if image_tensor is not None and not is_t2i:
-        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,
@@ -1806,13 +1427,6 @@ 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)
@@ -1825,15 +1439,11 @@ 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:
@@ -1851,11 +1461,7 @@ class Cosmos3OmniDiffusersPipeline(
             latents = torch.cat(samples, dim=0)
         else:
             diffusion_output = _run_diffusion(latents)
-            if action_enabled and sound_enabled:
+            if 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
@@ -1877,18 +1483,4 @@ 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,47 +136,6 @@ 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)
 # ---------------------------------------------------------------------------
@@ -204,7 +163,6 @@ 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.
 
@@ -222,17 +180,10 @@ def compute_mrope_position_ids_vision(
         )
         base_tps = base_fps / effective_base_tcf
         frame_indices = torch.arange(grid_t, dtype=torch.float32)
-        t_index = (
+        t_index = (frame_indices / tps * base_tps + temporal_offset).view(-1, 1).expand(-1, grid_h * grid_w).flatten()
-            ((frame_indices + start_frame_offset) / tps * base_tps + temporal_offset)
-            .view(-1, 1)
-            .expand(-1, grid_h * grid_w)
-            .flatten()
-        )
     else:
-        t_index = (
+        t_index = torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten() + int(
-            torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand(-1, grid_h * grid_w).flatten()
+            temporal_offset
-            + 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()
@@ -269,30 +220,6 @@ 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."""
 
@@ -1020,28 +947,19 @@ 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(
@@ -1074,20 +992,6 @@ 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)
@@ -1175,21 +1079,6 @@ 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(
@@ -1201,9 +1090,6 @@ 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."""
@@ -1211,7 +1097,6 @@ 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 = []
@@ -1231,17 +1116,6 @@ 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,
@@ -1287,9 +1161,7 @@ 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:
@@ -1297,16 +1169,14 @@ 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, action chunk size, "
+            "Adjust the spatial resolution, frame count, sound duration, "
-            "sound duration, or sound latent FPS so the combined media sequence is a "
+            "or sound latent FPS so the combined media sequence is a "
             "multiple of ulysses_degree."
-            if has_action or has_sound
+            if has_sound
             else (
                 "Adjust the spatial resolution so that "
                 "t * ceil(h/patch) * ceil(w/patch) is a multiple "
@@ -1328,11 +1198,6 @@ 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,
@@ -1345,10 +1210,6 @@ 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
@@ -1357,7 +1218,7 @@ class Cosmos3VFMTransformer(nn.Module):
 
         Returns:
             [B, C, t, h, w] velocity prediction, or
-            tuple outputs in video, action, sound order when extra modalities are provided.
+            tuple outputs in video, sound order when sound latents are provided.
         """
         t, h, w = video_shape
         hp, wp, _, _ = self._pad_to_patch_size(h, w)
@@ -1369,14 +1230,7 @@ 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 "
@@ -1390,21 +1244,8 @@ 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(
@@ -1436,22 +1277,9 @@ 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)
@@ -1466,9 +1294,6 @@ 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)
@@ -1486,9 +1311,7 @@ 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,
         )
@@ -1523,12 +1346,10 @@ class Cosmos3VFMTransformer(nn.Module):
 
         # Final norm and project back to latent space
         hidden_gen = self.norm_moe_gen(hidden_gen)
-        if not has_action and not has_sound:
+        if 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)
@@ -1536,11 +1357,6 @@ 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/hunyuan_image3/hunyuan_image3_transformer.py

@@ -32,6 +32,7 @@ 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,
@@ -72,7 +73,6 @@ 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,9 +573,7 @@ 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()}
-            # Siglip2VisionTransformer now returns a plain tensor (B, max_patches,
+            image_embed = self.vision_model(image, **cur_kwargs).last_hidden_state
-            # 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)
@@ -626,7 +624,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, generator=None):
+    def vae_encode(self, image, cfg_factor=1):
         config = self.vae.config
 
         if image.ndim == 3:
@@ -641,7 +639,11 @@ class HunyuanImage3Pipeline(
             if isinstance(vae_encode_result, torch.Tensor):
                 latents = vae_encode_result
             else:
-                latents = vae_encode_result.latent_dist.sample(generator)
+                # 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)
             if hasattr(config, "shift_factor") and config.shift_factor:
                 latents.sub_(config.shift_factor)
             if hasattr(config, "scaling_factor") and config.scaling_factor:
@@ -664,17 +666,14 @@ 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 batch_idx, cond_image_info_list in enumerate(batch_cond_image_info_list):
+        for cond_image_info_list in 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))
@@ -864,11 +863,9 @@ class HunyuanImage3Pipeline(
         output, sections = out["output"], out["sections"]
 
         # 4. Encode conditional images
-        # Skip encoding if AR KV reuse is enabled
+        if batch_cond_image_info is not None and len(batch_cond_image_info[0]) > 0:
-        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], generator=generator
+                batch_cond_image_info, cfg_factor[mode]
             )
             vit_kwargs = {"spatial_shapes": [], "attention_mask": []}
             for cond_image_info in batch_cond_image_info:
@@ -1392,23 +1389,18 @@ 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 prompt kwargs forwarded by ar2diffusion
+        # Fall back to per-prompt use_system_prompt 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, bot_task or "image", system_prompt)
+            system_prompt = get_system_prompt(use_system_prompt, "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
 
@@ -1455,18 +1447,6 @@ 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,
@@ -1477,8 +1457,6 @@ 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/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, ClassVar
+from typing import Any
 
 import numpy as np
 import PIL.Image
@@ -482,7 +482,6 @@ class Wan22S2VPipeline(
     _DEFAULT_MOTION_FRAMES = 73
     _DEFAULT_INFER_FRAMES = 80
     _DEFAULT_FPS = 16
-    dummy_run_num_frames: ClassVar[int] = 0
 
     def __init__(
         self,
@@ -930,19 +929,21 @@ 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.
 
-        With return_dict=False, the model returns (output,) where output is [B, C, T, H, W].
+        The S2V model returns a list of tensors (one per batch element).
-        We squeeze batch dim since S2V processes one sample at a time.
+        We take the first element 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]
-        if isinstance(result, tuple):
+        return result[0] if isinstance(result, tuple) else result
-            return result[0].squeeze(0)
-        return result.sample.squeeze(0)
 
     def diffuse(
         self,
@@ -995,29 +996,32 @@ 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 = {
-                    "hidden_states": latent_model_input,
+                    "x": latent_model_input,
-                    "timestep": timestep,
+                    "t": timestep,
-                    "encoder_hidden_states": prompt_embeds[0:1],
+                    "context": prompt_embeds[0:1],
+                    "seq_len": max_seq_len,
                     "cond_states": cond_latents,
                     "motion_latents": input_motion_latents,
                     "ref_latents": ref_latents,
+                    "motion_frames": motion_frames,
                     "drop_motion_frames": drop_first_motion,
-                    "encoder_hidden_states_audio": positive_audio_emb,
+                    "audio_emb": positive_audio_emb,
-                    "return_dict": False,
                 }
 
                 if do_true_cfg:
                     negative_kwargs = {
-                        "hidden_states": latent_model_input,
+                        "x": latent_model_input,
-                        "timestep": timestep,
+                        "t": timestep,
-                        "encoder_hidden_states": negative_prompt_embeds[0:1],
+                        "context": negative_prompt_embeds[0:1],
+                        "seq_len": max_seq_len,
                         "cond_states": cond_latents,
                         "motion_latents": input_motion_latents,
                         "ref_latents": ref_latents,
+                        "motion_frames": motion_frames,
                         "drop_motion_frames": drop_first_motion,
-                        "encoder_hidden_states_audio": negative_audio_emb,
+                        "audio_emb": negative_audio_emb,
-                        "return_dict": False,
                     }
                 else:
                     negative_kwargs = None
@@ -1170,7 +1174,7 @@ class Wan22S2VPipeline(
         ref_latents = self.encode_ref_image(image, height, width, device=device).to(dtype=dtype)
 
         # ---- 4. Initial motion latents (zeros) ----
-        motion_frames = min(self.motion_frames, infer_frames - 1)
+        motion_frames = self.motion_frames
         motion_pixels = torch.zeros(
             [1, 3, motion_frames, height, width],
             dtype=dtype,
@@ -1284,10 +1288,6 @@ 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)

vllm_omni/diffusion/stage_diffusion_client.py

@@ -119,7 +119,6 @@ class StageDiffusionClient(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 = getattr(metadata, "requires_multimodal_data", False)
         self.custom_process_input_func = getattr(metadata, "custom_process_input_func", None)

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -542,8 +542,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                     # Requests that have loaded chunk from last round
                     # of schedule, but have not scheduled
                     continue
-                if self.is_done_receiving_chunks(request.request_id):
+                if request.request_id in self.finished_requests:
-                    request.additional_information = None
                     continue
                 # Requests that waiting for chunk
                 self.load_async(request)

vllm_omni/engine/arg_utils.py

@@ -77,10 +77,13 @@ def register_omni_models_to_vllm():
 
     _register_omni_hf_configs()
 
-    supported_archs = ModelRegistry.get_supported_archs()
+    # Always register every arch declared in _OMNI_MODELS — these are the
+    # archs vllm-omni explicitly owns, so we want our implementation to
+    # win even if a future vllm release ships a same-named built-in.
+    # vllm's ModelRegistry.register_model already emits a warning and
+    # overrides on collision; we let it handle that.
     for arch, (mod_folder, mod_relname, cls_name) in _OMNI_MODELS.items():
-        if arch not in supported_archs:
+        ModelRegistry.register_model(arch, f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}:{cls_name}")
-            ModelRegistry.register_model(arch, f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}:{cls_name}")
 
 
 @dataclass

vllm_omni/engine/async_omni_engine.py

@@ -123,7 +123,6 @@ class StageRuntimeInfo:
     final_output: bool
     final_output_type: FinalOutputModalityType | None
     stage_type: str
-    model_stage: str | None = None
 
 
 # ============================================================================
@@ -1248,7 +1247,6 @@ class AsyncOmniEngine:
                     final_output=first_client.final_output,
                     final_output_type=first_client.final_output_type,
                     stage_type=first_client.stage_type,
-                    model_stage=getattr(first_client, "model_stage", None),
                 )
             )
 
@@ -1618,7 +1616,6 @@ class AsyncOmniEngine:
         prompt_text: str | None = None,
         sampling_params_list: Sequence[Any] | None = None,
         final_stage_id: int = 0,
-        final_output_stage_ids: Sequence[int] | None = None,
         arrival_time: float | None = None,
         lora_request: Any = None,
         tokenization_kwargs: dict[str, Any] | None = None,
@@ -1631,7 +1628,6 @@ class AsyncOmniEngine:
         message_type: Literal["add_request", "streaming_update"] = "add_request",
     ) -> StageSubmissionMessage:
         """Build an add_request message after stage-0 preprocessing."""
-        request_timestamp = float(arrival_time) if arrival_time is not None else time.time()
         effective_sampling_params_list: list[OmniSamplingParams] = (
             list(cast(Sequence[OmniSamplingParams], sampling_params_list))
             if sampling_params_list is not None
@@ -1717,9 +1713,7 @@ class AsyncOmniEngine:
             output_prompt_text=output_prompt_text,
             sampling_params_list=effective_sampling_params_list,
             final_stage_id=final_stage_id,
-            final_output_stage_ids=list(final_output_stage_ids) if final_output_stage_ids is not None else None,
             preprocess_ms=_preprocess_ms,
-            request_timestamp=request_timestamp,
             enqueue_ts=time.perf_counter(),
         )
 
@@ -2192,7 +2186,6 @@ class AsyncOmniEngine:
         prompt_text: str | None = None,
         sampling_params_list: Sequence[Any] | None = None,
         final_stage_id: int = 0,
-        final_output_stage_ids: Sequence[int] | None = None,
         arrival_time: float | None = None,
         lora_request: Any = None,
         tokenization_kwargs: dict[str, Any] | None = None,
@@ -2216,7 +2209,6 @@ class AsyncOmniEngine:
             prompt_text=prompt_text,
             sampling_params_list=sampling_params_list,
             final_stage_id=final_stage_id,
-            final_output_stage_ids=final_output_stage_ids,
             arrival_time=arrival_time,
             lora_request=lora_request,
             tokenization_kwargs=tokenization_kwargs,
@@ -2244,7 +2236,6 @@ class AsyncOmniEngine:
         prompt_text: str | None = None,
         sampling_params_list: Sequence[Any] | None = None,
         final_stage_id: int = 0,
-        final_output_stage_ids: Sequence[int] | None = None,
         arrival_time: float | None = None,
         lora_request: Any = None,
         tokenization_kwargs: dict[str, Any] | None = None,
@@ -2262,7 +2253,6 @@ class AsyncOmniEngine:
             prompt_text=prompt_text,
             sampling_params_list=sampling_params_list,
             final_stage_id=final_stage_id,
-            final_output_stage_ids=final_output_stage_ids,
             arrival_time=arrival_time,
             lora_request=lora_request,
             tokenization_kwargs=tokenization_kwargs,
@@ -2280,7 +2270,6 @@ class AsyncOmniEngine:
         prompt_text: str | None = None,
         sampling_params_list: Sequence[Any] | None = None,
         final_stage_id: int = 0,
-        final_output_stage_ids: Sequence[int] | None = None,
         arrival_time: float | None = None,
         *,
         resumable: bool = True,
@@ -2292,7 +2281,6 @@ class AsyncOmniEngine:
             prompt_text=prompt_text,
             sampling_params_list=sampling_params_list,
             final_stage_id=final_stage_id,
-            final_output_stage_ids=final_output_stage_ids,
             arrival_time=arrival_time,
             resumable=resumable,
             message_type="streaming_update",
@@ -2306,7 +2294,6 @@ class AsyncOmniEngine:
         prompt_text: str | None = None,
         sampling_params_list: Sequence[Any] | None = None,
         final_stage_id: int = 0,
-        final_output_stage_ids: Sequence[int] | None = None,
         arrival_time: float | None = None,
         *,
         resumable: bool = True,
@@ -2318,7 +2305,6 @@ class AsyncOmniEngine:
             prompt_text=prompt_text,
             sampling_params_list=sampling_params_list,
             final_stage_id=final_stage_id,
-            final_output_stage_ids=final_output_stage_ids,
             arrival_time=arrival_time,
             resumable=resumable,
         )

vllm_omni/engine/messages.py

@@ -24,9 +24,7 @@ class StageSubmissionMessage(EngineQueueMessage, kw_only=True):
     sampling_params_list: list[OmniSamplingParams]
     final_stage_id: int
     preprocess_ms: float
-    request_timestamp: float
     enqueue_ts: float
-    final_output_stage_ids: list[int] | None = None
 
 
 class AddCompanionRequestMessage(EngineQueueMessage, kw_only=True):

vllm_omni/engine/orchestrator.py

@@ -126,11 +126,6 @@ class OrchestratorRequestState:
     prompt: Any = None
     sampling_params_list: list[Any] = field(default_factory=list)
     final_stage_id: int = -1
-    final_output_stage_ids: set[int] = field(default_factory=set)
-    finished_final_output_stage_ids: set[int] = field(default_factory=set)
-
-    # Wall-clock timestamp when the client-facing engine request was accepted.
-    request_timestamp: float = 0.0
 
     # Metrics: timestamp when request was submitted to each stage.
     stage_submit_ts: dict[int, float] = field(default_factory=dict)
@@ -456,7 +451,6 @@ class Orchestrator:
         if not sampling_params_list:
             raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.")
         final_stage_id = msg.final_stage_id
-        final_output_stage_ids = set(msg.final_output_stage_ids or [final_stage_id])
 
         logger.debug(
             "[Orchestrator] _handle_add_request: stage=%s req=%s "
@@ -475,8 +469,6 @@ class Orchestrator:
             prompt=original_prompt,
             sampling_params_list=sampling_params_list,
             final_stage_id=final_stage_id,
-            final_output_stage_ids=final_output_stage_ids,
-            request_timestamp=float(msg.request_timestamp or _time.time()),
             mm_features=getattr(prompt, "mm_features", None),
         )
         self.request_states[request_id] = req_state
@@ -520,9 +512,7 @@ class Orchestrator:
                 output_prompt_text=msg.output_prompt_text,
                 sampling_params_list=msg.sampling_params_list,
                 final_stage_id=msg.final_stage_id,
-                final_output_stage_ids=msg.final_output_stage_ids,
                 preprocess_ms=msg.preprocess_ms,
-                request_timestamp=msg.request_timestamp,
                 enqueue_ts=msg.enqueue_ts,
             )
             await self._handle_add_request(fallback_msg)
@@ -568,8 +558,6 @@ class Orchestrator:
             prompt=companion_prompt,
             sampling_params_list=sampling_params_list,
             final_stage_id=0,
-            final_output_stage_ids={0},
-            request_timestamp=parent_state.request_timestamp,
         )
         self.request_states[companion_id] = companion_state
         companion_state.stage_submit_ts[0] = _time.time()
@@ -677,7 +665,6 @@ class Orchestrator:
                         if output is None:
                             continue
 
-                        pool.record_output_timestamps([output])
                         await self._handle_processed_outputs(stage_id, replica_id, [output])
                         idle = False
                     else:
@@ -697,8 +684,6 @@ class Orchestrator:
                                     "new_prompt_len_snapshot",
                                     None,
                                 )
-                                if req_state.streaming.enabled:
-                                    await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state)
                             # OmniSchedulerMixin.make_stats() already throttles
                             # per-scheduler at 1 Hz, so raw_outputs.scheduler_stats
                             # being non-None means this replica passed its own gate.
@@ -790,9 +775,7 @@ class Orchestrator:
                 stage_metrics = pool.build_stage_metrics(
                     [output],
                     submit_ts=req_state.stage_submit_ts.get(stage_id, _time.time()),
-                    request_timestamp=req_state.request_timestamp,
                     replica_id=replica_id,
-                    sampling_params=req_state.sampling_params_list[stage_id],
                 )
                 stage_metrics.pipeline_timings = dict(req_state.pipeline_timings)
 
@@ -831,34 +814,6 @@ class Orchestrator:
             if self.request_states.pop(request_id, None) is not None and self._running_counter is not None:
                 self._running_counter.decrement()
 
-    async def _apply_raw_terminal_stage_finish(
-        self,
-        stage_id: int,
-        eco: Any,
-        req_state: OrchestratorRequestState,
-    ) -> None:
-        """Record session-level finish markers dropped by the streaming output processor.
-
-        Streaming segment stops set ``is_segment_finished=True`` and are handled
-        via processed outputs. Session termination (e.g. ``finish_requests`` after
-        ``resumable=False``) emits a terminal ``finish_reason`` with
-        ``is_segment_finished=False``, but vLLM's output processor may remove the
-        request state before that EngineCoreOutput is processed.
-
-        Only update ``finished_final_output_stage_ids`` here. Request cleanup stays
-        in ``_route_output`` so downstream async-chunk stages can still deliver
-        outputs after stage-0 session end.
-        """
-        if getattr(eco, "finish_reason", None) is None:
-            return
-        if getattr(eco, "is_segment_finished", False):
-            return
-
-        final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id}
-        if stage_id not in final_output_stage_ids:
-            return
-        req_state.finished_final_output_stage_ids.add(stage_id)
-
     def _maybe_clone_diffusion_params_for_cfg(self, request_id: str, params: Any) -> Any:
         """Attach CFG companion ids to diffusion sampling params when needed."""
         companion_request_ids = self._cfg_tracker.get_companion_request_ids(request_id)
@@ -897,11 +852,6 @@ class Orchestrator:
             await self._cleanup_request_ids([req_id])
             return
 
-        request_finished = False
-        if finished and self.stage_pools[stage_id].final_output:
-            req_state.finished_final_output_stage_ids.add(stage_id)
-            final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id}
-            request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids)
         if self.stage_pools[stage_id].final_output:
             await self.output_async_queue.put(
                 OutputMessage(
@@ -910,7 +860,7 @@ class Orchestrator:
                     replica_id=replica_id,
                     engine_outputs=output,
                     metrics=stage_metrics,
-                    finished=request_finished,
+                    finished=finished and stage_id == req_state.final_stage_id,
                     stage_submit_ts=submit_ts,
                 )
             )
@@ -965,7 +915,7 @@ class Orchestrator:
                         is_final_update=True,
                     )
 
-        if request_finished:
+        if finished and stage_id == req_state.final_stage_id:
             await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)])
 
     def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool:
@@ -1391,12 +1341,7 @@ class Orchestrator:
                     resumable=downstream_resumable,
                 )
                 request.external_req_id = request.request_id
-                await next_pool.submit_initial(
+                await next_pool.submit_initial(request_id, req_state, request, prompt_text=None)
-                    request_id,
-                    req_state,
-                    request,
-                    prompt_text=None,
-                )
 
             # async_chunk pre-submit fires per stage edge (N-1 -> N). Source
             # replica is stage 0's bound replica (single-replica thinker in

vllm_omni/engine/output_processor.py

@@ -13,7 +13,7 @@ from vllm.v1.engine.output_processor import (
     RequestState,
 )
 from vllm.v1.engine.parallel_sampling import ParentRequest
-from vllm.v1.metrics.stats import IterationStats, RequestStateStats
+from vllm.v1.metrics.stats import IterationStats
 
 from vllm_omni.data_entry_keys import unflatten_payload
 from vllm_omni.engine.output_modality import DRAINABLE_MODALITIES
@@ -35,21 +35,13 @@ class OmniRequestState(RequestState):
         *args,
         **kwargs,
     ):
-        arrival_time = kwargs.get("arrival_time")
-        if arrival_time is None and len(args) > 12:
-            arrival_time = args[12]
         super().__init__(*args, **kwargs)
-        self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0))
         # Omni-specific: multimodal output accumulation
         # NOTE: Keys in mm_accumulated matter, because they dictate which
         # outputs are drained (i.e., only drain modality keys, don't drain
         # hidden states).
         self.mm_accumulated: dict[str, Any] = {}
 
-    def apply_streaming_update(self, update) -> None:
-        super().apply_streaming_update(update)
-        self.native_text_stats.arrival_time = update.arrival_time
-
     @staticmethod
     def _to_cpu(x):
         """Try to convert to CPU tensor if needed."""
@@ -326,32 +318,6 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             tracing_enabled=tracing_enabled,
         )
         self.engine_core_output_type = engine_core_output_type
-        self._native_text_metrics_by_request: dict[str, dict[str, Any]] = {}
-
-    def _native_text_metric_record(self, request_id: str) -> dict[str, Any]:
-        return self._native_text_metrics_by_request.setdefault(
-            request_id,
-            {
-                "vllm_ttft_ms": 0.0,
-                "vllm_tpot_ms": 0.0,
-                "vllm_itl_ms": 0.0,
-                "vllm_itls_ms": [],
-            },
-        )
-
-    def pop_native_text_metrics(self, request_id: str) -> dict[str, Any]:
-        return self._native_text_metrics_by_request.pop(request_id, {})
-
-    def abort_requests(self, request_ids, internal: bool) -> list[str]:
-        request_ids = list(request_ids)
-        for request_id in request_ids:
-            if internal:
-                req_state = self.request_states.get(request_id)
-                if req_state is not None:
-                    self._native_text_metrics_by_request.pop(req_state.external_req_id, None)
-            else:
-                self._native_text_metrics_by_request.pop(request_id, None)
-        return super().abort_requests(request_ids, internal)
 
     def add_request(
         self,
@@ -405,7 +371,6 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
 
         external_req_id = getattr(req_state, "external_req_id", None)
         if external_req_id is not None:
-            self._native_text_metrics_by_request.pop(external_req_id, None)
             request_ids = self.external_req_ids.get(external_req_id)
             if request_ids is not None:
                 self.external_req_ids[external_req_id] = [rid for rid in request_ids if rid != request_id]
@@ -437,73 +402,3 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             engine_core_timestamp=engine_core_timestamp,
             iteration_stats=iteration_stats,
         )
-
-    def _update_stats_from_output(
-        self,
-        req_state: RequestState,
-        engine_core_output: EngineCoreOutput,
-        engine_core_timestamp: float | None,
-        iteration_stats: IterationStats | None,
-    ):
-        was_prefilling = req_state.is_prefilling
-        native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None
-        previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0
-
-        super()._update_stats_from_output(
-            req_state,
-            engine_core_output,
-            engine_core_timestamp,
-            None,
-        )
-
-        if iteration_stats is None or engine_core_timestamp is None or native_stats is None:
-            return
-
-        iteration_stats.update_from_output(
-            engine_core_output,
-            engine_core_timestamp,
-            was_prefilling,
-            native_stats,
-            self.lora_states,
-            req_state.lora_name,
-        )
-        record = self._native_text_metric_record(req_state.external_req_id)
-        if was_prefilling:
-            record["vllm_ttft_ms"] = max(float(native_stats.first_token_latency) * 1000.0, 0.0)
-            return
-
-        if previous_last_token_ts > 0:
-            itl_ms = max((engine_core_timestamp - previous_last_token_ts) * 1000.0, 0.0)
-            itls_ms = record.setdefault("vllm_itls_ms", [])
-            itls_ms.append(itl_ms)
-            record["vllm_itl_ms"] = sum(itls_ms) / float(len(itls_ms))
-
-    def _update_stats_from_finished(
-        self,
-        req_state: RequestState,
-        finish_reason: FinishReason | None,
-        iteration_stats: IterationStats | None,
-    ):
-        super()._update_stats_from_finished(req_state, finish_reason, None)
-
-        native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None
-        if iteration_stats is None or native_stats is None or finish_reason is None:
-            return
-
-        iteration_stats.update_from_finished_request(
-            finish_reason=finish_reason,
-            request_id=req_state.external_req_id,
-            num_prompt_tokens=req_state.prompt_len,
-            max_tokens_param=req_state.max_tokens_param,
-            req_stats=native_stats,
-            num_cached_tokens=req_state.num_cached_tokens,
-        )
-        if not iteration_stats.finished_requests:
-            return
-
-        finished_request = iteration_stats.finished_requests[-1]
-        if finished_request.request_id != req_state.external_req_id:
-            return
-        self._native_text_metric_record(req_state.external_req_id)["vllm_tpot_ms"] = (
-            float(finished_request.mean_time_per_output_token) * 1000.0
-        )

vllm_omni/engine/stage_init_utils.py

@@ -428,8 +428,6 @@ def extract_stage_metadata(stage_config: Any) -> StageMetadata:
         _mod, _fn = _ckf_path.rsplit(".", 1)
         cfg_kv_collect_func = getattr(importlib.import_module(_mod), _fn)
 
-    model_stage = getattr(engine_args, "model_stage", None)
-
     if stage_type == "diffusion":
         return StageMetadata(
             stage_id=stage_id,
@@ -442,11 +440,12 @@ def extract_stage_metadata(stage_config: Any) -> StageMetadata:
             final_output_type=final_output_type,
             default_sampling_params=default_sampling_params,
             custom_process_input_func=custom_process_input_func,
-            model_stage=model_stage,
+            model_stage=None,
             runtime_cfg=runtime_cfg,
             cfg_kv_collect_func=cfg_kv_collect_func,
         )
 
+    model_stage = getattr(engine_args, "model_stage", None)
     engine_output_type = getattr(engine_args, "engine_output_type", None)
     is_comprehension = getattr(stage_config, "is_comprehension", False)
     requires_multimodal_data = getattr(runtime_cfg, "requires_multimodal_data", False)

vllm_omni/engine/stage_pool.py

@@ -24,7 +24,6 @@ from vllm_omni.engine.stage_client import (
     StagePoolLLMClient,
 )
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.metrics import definitions as defs
 from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics
 from vllm_omni.metrics.stats import StageStats
 from vllm_omni.metrics.utils import count_tokens_from_outputs
@@ -97,10 +96,6 @@ class StagePool:
         self._next_replica_id = 0
         self._request_bindings: dict[str, int] = {}
         self._replica_metrics: list[_ReplicaMetrics] = [_ReplicaMetrics() for _ in self.clients]
-        self._output_timestamps_by_request: dict[str, list[float]] = {}
-        self._non_empty_first_output_timestamps_by_request: dict[str, float] = {}
-        self._audio_frames_by_request: dict[str, int] = {}
-        self._audio_sample_rate_by_request: dict[str, int] = {}
 
         # Distributed-mode state. Populated by add_client / remove_client.
         self._addr_to_replica_id: dict[str, int] = {}
@@ -433,10 +428,6 @@ class StagePool:
         """Drop the route binding for *request_id* in this stage."""
         self._request_bindings.pop(request_id, None)
         self._affinity.pop(request_id, None)
-        self._output_timestamps_by_request.pop(str(request_id), None)
-        self._non_empty_first_output_timestamps_by_request.pop(str(request_id), None)
-        self._audio_frames_by_request.pop(str(request_id), None)
-        self._audio_sample_rate_by_request.pop(str(request_id), None)
 
     def release_bindings(self, request_ids: list[str]) -> None:
         """Drop route bindings for the given request ids in this stage."""
@@ -494,66 +485,13 @@ class StagePool:
         request_outputs: list[Any],
         *,
         submit_ts: float,
-        request_timestamp: float,
         replica_id: int,
-        sampling_params: Any | None = None,
     ) -> StageRequestMetrics:
         """Build stage metrics for outputs produced on one replica."""
         now = _time.time()
         stage_gen_time_ms = (now - submit_ts) * 1000.0
 
-        request_id = str(getattr(request_outputs[0], "request_id", "")) if request_outputs else ""
-        output_timestamps = self._output_timestamps_by_request.pop(request_id, []) if request_id else []
-        non_empty_first_output_ts = (
-            self._non_empty_first_output_timestamps_by_request.pop(request_id, None) if request_id else None
-        )
         num_tokens_out = count_tokens_from_outputs(request_outputs)
-        output_unit_type = self._infer_output_unit_type(request_outputs, token_count=num_tokens_out)
-        output_unit_count = self._count_output_units(
-            request_outputs,
-            unit_type=output_unit_type,
-            fallback_token_count=num_tokens_out,
-        )
-        native_text_metrics = {}
-        if request_id:
-            pop_native_text_metrics = getattr(self.output_processor, "pop_native_text_metrics", None)
-            if callable(pop_native_text_metrics):
-                native_text_metrics = pop_native_text_metrics(request_id)
-        current_audio_frames, current_audio_sample_rate, _ = self._collect_audio_metrics(request_outputs)
-        accumulated_audio_frames = self._audio_frames_by_request.pop(request_id, 0) if request_id else 0
-        accumulated_audio_sample_rate = self._audio_sample_rate_by_request.pop(request_id, 0) if request_id else 0
-        audio_generated_frames = max(accumulated_audio_frames, current_audio_frames)
-        audio_sample_rate = accumulated_audio_sample_rate or current_audio_sample_rate
-        audio_duration_s = (
-            float(audio_generated_frames) / float(audio_sample_rate)
-            if audio_generated_frames > 0 and audio_sample_rate > 0
-            else 0.0
-        )
-        image_pixels = self._count_image_pixels(request_outputs) if output_unit_type == "image" else 0
-        num_inference_steps = self._coerce_int_scalar(getattr(sampling_params, "num_inference_steps", None))
-        denoise_step_latency_ms = (
-            defs.compute_denoise_step_latency(stage_gen_time_ms, num_inference_steps)
-            if output_unit_type == "image"
-            else 0.0
-        )
-        has_output_timestamps = bool(output_timestamps)
-        first_ts = output_timestamps[0] if has_output_timestamps else now
-        serving_time_to_first_output_ms = (
-            max((non_empty_first_output_ts - request_timestamp) * 1000.0, 0.0)
-            if non_empty_first_output_ts is not None
-            else 0.0
-        )
-        remaining_ms = max((now - first_ts) * 1000.0, 0.0)
-        if output_unit_count > 1 and self._is_streaming_output_unit_type(output_unit_type):
-            time_per_output_unit_ms = remaining_ms / float(output_unit_count - 1)
-        else:
-            time_per_output_unit_ms = 0.0
-        inter_output_latencies_ms = [
-            max((cur_ts - prev_ts) * 1000.0, 0.0) for prev_ts, cur_ts in zip(output_timestamps, output_timestamps[1:])
-        ]
-        inter_output_latency_ms = (
-            sum(inter_output_latencies_ms) / float(len(inter_output_latencies_ms)) if inter_output_latencies_ms else 0.0
-        )
         num_tokens_in = 0
         if self.stage_id == 0:
             for ro in request_outputs:
@@ -566,7 +504,6 @@ class StagePool:
         batch_id = metrics.batch_seq
         metrics.agg_total_tokens += num_tokens_out
         metrics.agg_total_gen_time_ms += stage_gen_time_ms
-        audio_rtf = defs.compute_audio_rtf(stage_gen_time_ms / 1000.0, audio_duration_s)
 
         return StageRequestMetrics(
             num_tokens_in=num_tokens_in,
@@ -582,321 +519,8 @@ class StagePool:
                 total_token=metrics.agg_total_tokens,
                 total_gen_time_ms=metrics.agg_total_gen_time_ms,
             ),
-            audio_generated_frames=audio_generated_frames,
-            audio_sample_rate=audio_sample_rate,
-            audio_duration_s=audio_duration_s,
-            audio_rtf=audio_rtf,
-            image_pixels=image_pixels,
-            denoise_step_latency_ms=denoise_step_latency_ms,
-            output_unit_type=output_unit_type,
-            output_unit_count=output_unit_count,
-            serving_time_to_first_output_ms=serving_time_to_first_output_ms,
-            time_per_output_unit_ms=time_per_output_unit_ms,
-            inter_output_latency_ms=inter_output_latency_ms,
-            inter_output_latencies_ms=inter_output_latencies_ms,
-            vllm_ttft_ms=float(native_text_metrics.get("vllm_ttft_ms") or 0.0),
-            vllm_tpot_ms=float(native_text_metrics.get("vllm_tpot_ms") or 0.0),
-            vllm_itl_ms=float(native_text_metrics.get("vllm_itl_ms") or 0.0),
-            vllm_itls_ms=list(native_text_metrics.get("vllm_itls_ms") or []),
         )
 
-    def _infer_output_unit_type(self, request_outputs: list[Any], *, token_count: int) -> str:
-        final_output_type = getattr(self.stage_client, "final_output_type", None)
-
-        if self._has_image_output(request_outputs) or final_output_type in {"image", "images"}:
-            return "image"
-        if self._has_video_output(request_outputs) or final_output_type in {"video", "videos"}:
-            return "video"
-        if self._has_audio_output(request_outputs) or final_output_type == "audio":
-            return "audio"
-        if self._has_trajectory_latent_output(request_outputs) or self._has_latent_output(request_outputs):
-            return "latent"
-        if final_output_type in {"latent", "latents"}:
-            return "latent"
-        if final_output_type == "text":
-            return "text"
-        if token_count > 0:
-            return "stream"
-        return "other"
-
-    @staticmethod
-    def _is_streaming_output_unit_type(unit_type: str) -> bool:
-        return unit_type in {"text", "stream"}
-
-    def _count_output_units(
-        self,
-        request_outputs: list[Any],
-        *,
-        unit_type: str,
-        fallback_token_count: int,
-    ) -> int:
-        if unit_type == "audio":
-            total_frames, _, _ = self._collect_audio_metrics(request_outputs)
-            if total_frames > 0:
-                return total_frames
-        if unit_type == "image":
-            total_images = sum(self._count_images(ro) for ro in request_outputs)
-            if total_images > 0:
-                return total_images
-        if unit_type == "video":
-            total_videos = sum(self._count_videos(ro) for ro in request_outputs)
-            if total_videos > 0:
-                return total_videos
-        if unit_type == "latent":
-            total_latents = sum(
-                self._count_value_units(getattr(ro, "trajectory_latents", None))
-                + self._count_value_units(getattr(ro, "latents", None))
-                for ro in request_outputs
-            )
-            if total_latents > 0:
-                return total_latents
-        return int(fallback_token_count)
-
-    def _has_audio_output(self, request_outputs: list[Any]) -> bool:
-        for ro in request_outputs:
-            for mm_output in self._iter_multimodal_outputs(ro):
-                if isinstance(mm_output, dict) and mm_output.get("audio") is not None:
-                    return True
-        return False
-
-    def _collect_audio_metrics(self, request_outputs: list[Any]) -> tuple[int, int, float]:
-        total_frames = 0
-        sample_rate = 0
-        for ro in request_outputs:
-            for mm_output in self._iter_multimodal_outputs(ro):
-                if not isinstance(mm_output, dict):
-                    continue
-                if sample_rate <= 0:
-                    sample_rate = self._infer_audio_sample_rate(mm_output)
-                audio_output = mm_output.get("audio")
-                if audio_output is None:
-                    continue
-                items = audio_output if isinstance(audio_output, list) else [audio_output]
-                for item in items:
-                    total_frames += self._count_audio_frames(item)
-        duration_s = (float(total_frames) / float(sample_rate)) if total_frames > 0 and sample_rate > 0 else 0.0
-        return int(total_frames), int(sample_rate), duration_s
-
-    @staticmethod
-    def _count_audio_frames(audio_output: Any) -> int:
-        shape = getattr(audio_output, "shape", None)
-        if shape is not None and len(shape) > 0:
-            # Audio chunks are concatenated on dim=-1 in the output processor,
-            # so the frame/sample axis is the last dim (e.g. [channels, frames]).
-            # Keep this aligned with serving_chat.py: audio tensors are consumed
-            # as (T,), (C, T), or (B, C, T). Flattening would corrupt
-            # multi-channel audio.
-            return int(shape[-1])
-        try:
-            return len(audio_output)
-        except TypeError:
-            return 1
-
-    def _infer_audio_sample_rate(self, mm_output: dict[str, Any]) -> int:
-        for key in ("audio_sample_rate", "sample_rate", "sampling_rate", "sr"):
-            rate = self._coerce_int_scalar(mm_output.get(key))
-            if rate > 0:
-                return rate
-        for attr in ("audio_sample_rate", "sample_rate", "sampling_rate", "output_sample_rate"):
-            rate = self._coerce_int_scalar(getattr(self.stage_client, attr, None))
-            if rate > 0:
-                return rate
-            rate = self._coerce_int_scalar(getattr(self._stage_vllm_config, attr, None))
-            if rate > 0:
-                return rate
-        return 0
-
-    @classmethod
-    def _coerce_int_scalar(cls, value: Any) -> int:
-        if value is None:
-            return 0
-        if isinstance(value, (list, tuple)):
-            for item in value:
-                coerced = cls._coerce_int_scalar(item)
-                if coerced > 0:
-                    return coerced
-            return 0
-        item = getattr(value, "item", None)
-        if callable(item):
-            try:
-                return int(item())
-            except Exception:
-                return 0
-        try:
-            return int(value)
-        except (TypeError, ValueError):
-            return 0
-
-    def _has_image_output(self, request_outputs: list[Any]) -> bool:
-        for ro in request_outputs:
-            if self._count_images(ro) > 0:
-                return True
-        return False
-
-    def _has_video_output(self, request_outputs: list[Any]) -> bool:
-        for ro in request_outputs:
-            if self._count_videos(ro) > 0:
-                return True
-        return False
-
-    def _has_trajectory_latent_output(self, request_outputs: list[Any]) -> bool:
-        return any(self._is_non_empty_value(getattr(ro, "trajectory_latents", None)) for ro in request_outputs)
-
-    def _has_latent_output(self, request_outputs: list[Any]) -> bool:
-        return any(self._is_non_empty_value(getattr(ro, "latents", None)) for ro in request_outputs)
-
-    def _count_images(self, request_output: Any) -> int:
-        total_images = self._count_value_units(getattr(request_output, "images", None))
-        for mm_output in self._iter_multimodal_outputs(request_output):
-            if isinstance(mm_output, dict):
-                total_images += self._count_value_units(mm_output.get("image"))
-                total_images += self._count_value_units(mm_output.get("images"))
-        return total_images
-
-    def _count_image_pixels(self, request_outputs: list[Any]) -> int:
-        total_pixels = 0
-        for ro in request_outputs:
-            total_pixels += self._count_image_value_pixels(getattr(ro, "images", None))
-            for mm_output in self._iter_multimodal_outputs(ro):
-                if isinstance(mm_output, dict):
-                    total_pixels += self._count_image_value_pixels(mm_output.get("image"))
-                    total_pixels += self._count_image_value_pixels(mm_output.get("images"))
-        return total_pixels
-
-    @classmethod
-    def _count_image_value_pixels(cls, value: Any) -> int:
-        if value is None:
-            return 0
-        if isinstance(value, (list, tuple)):
-            return sum(cls._count_image_value_pixels(item) for item in value)
-
-        size = getattr(value, "size", None)
-        if isinstance(size, tuple) and len(size) >= 2:
-            try:
-                return int(size[0]) * int(size[1])
-            except (TypeError, ValueError):
-                return 0
-
-        shape = getattr(value, "shape", None)
-        if shape is None or len(shape) < 2:
-            return 0
-        dims = [int(dim) for dim in shape]
-        if len(dims) >= 4:
-            return dims[0] * dims[-2] * dims[-1]
-        if len(dims) == 3 and dims[0] in (1, 3, 4):
-            return dims[1] * dims[2]
-        if len(dims) == 3 and dims[-1] in (1, 3, 4):
-            return dims[0] * dims[1]
-        return dims[-2] * dims[-1]
-
-    def _count_videos(self, request_output: Any) -> int:
-        total_videos = self._count_video_units(getattr(request_output, "video", None))
-        total_videos += self._count_video_units(getattr(request_output, "videos", None))
-        for mm_output in self._iter_multimodal_outputs(request_output):
-            if isinstance(mm_output, dict):
-                total_videos += self._count_video_units(mm_output.get("video"))
-                total_videos += self._count_video_units(mm_output.get("videos"))
-        return total_videos
-
-    def has_non_empty_output(self, request_output: Any) -> bool:
-        return self._has_non_empty_output(request_output)
-
-    def _has_non_empty_output(self, request_output: Any) -> bool:
-        final_output_type = getattr(request_output, "final_output_type", None)
-        if final_output_type is None:
-            final_output_type = getattr(self.stage_client, "final_output_type", None)
-        if final_output_type == "text":
-            return any(bool(getattr(output, "text", "")) for output in getattr(request_output, "outputs", None) or [])
-
-        if getattr(request_output, "images", None):
-            return True
-        if getattr(request_output, "video", None) or getattr(request_output, "videos", None):
-            return True
-        if getattr(request_output, "trajectory_latents", None) is not None:
-            return True
-        custom_output = getattr(request_output, "_custom_output", None)
-        if isinstance(custom_output, dict) and custom_output:
-            return True
-        for mm_output in self._iter_multimodal_outputs(request_output):
-            if isinstance(mm_output, dict) and any(self._is_non_empty_value(value) for value in mm_output.values()):
-                return True
-        for output in getattr(request_output, "outputs", None) or []:
-            if getattr(output, "text", None):
-                return True
-            token_ids = getattr(output, "token_ids", None)
-            if self._is_non_empty_value(token_ids):
-                return True
-            cumulative_token_ids = getattr(output, "cumulative_token_ids", None)
-            if self._is_non_empty_value(cumulative_token_ids):
-                return True
-        return False
-
-    @staticmethod
-    def _is_non_empty_value(value: Any) -> bool:
-        if value is None:
-            return False
-        shape = getattr(value, "shape", None)
-        if shape is not None:
-            return all(int(dim) > 0 for dim in shape)
-        try:
-            return len(value) > 0
-        except TypeError:
-            return True
-
-    @staticmethod
-    def _count_value_units(value: Any) -> int:
-        if value is None:
-            return 0
-        shape = getattr(value, "shape", None)
-        if shape is not None:
-            return int(shape[0]) if len(shape) > 0 and int(shape[0]) > 0 else 0
-        try:
-            return len(value)
-        except TypeError:
-            return 1
-
-    @staticmethod
-    def _count_video_units(value: Any) -> int:
-        if value is None:
-            return 0
-        if isinstance(value, (list, tuple)):
-            return len(value)
-        shape = getattr(value, "shape", None)
-        if shape is not None:
-            return 1 if all(int(dim) > 0 for dim in shape) else 0
-        try:
-            return len(value)
-        except TypeError:
-            return 1
-
-    def record_output_timestamps(self, request_outputs: list[Any], *, output_ts: float | None = None) -> None:
-        """Record all output timestamps and the first non-empty output timestamp."""
-        output_ts = _time.time() if output_ts is None else output_ts
-        for request_output in request_outputs:
-            request_id = getattr(request_output, "request_id", None)
-            if request_id is None:
-                continue
-            rid = str(request_id)
-            self._output_timestamps_by_request.setdefault(rid, []).append(output_ts)
-            if self._has_non_empty_output(request_output):
-                self._non_empty_first_output_timestamps_by_request.setdefault(rid, output_ts)
-            audio_frames, audio_sample_rate, _ = self._collect_audio_metrics([request_output])
-            if audio_frames > 0:
-                self._audio_frames_by_request[rid] = self._audio_frames_by_request.get(rid, 0) + audio_frames
-            if self._audio_sample_rate_by_request.get(rid, 0) <= 0 and audio_sample_rate > 0:
-                self._audio_sample_rate_by_request[rid] = audio_sample_rate
-
-    def _iter_multimodal_outputs(self, request_output: Any) -> list[dict[str, Any]]:
-        multimodal_outputs: list[dict[str, Any]] = []
-        outer_mm = getattr(request_output, "multimodal_output", None)
-        if isinstance(outer_mm, dict) and outer_mm:
-            multimodal_outputs.append(outer_mm)
-        for output in getattr(request_output, "outputs", None) or []:
-            inner_mm = getattr(output, "multimodal_output", None)
-            if isinstance(inner_mm, dict) and inner_mm:
-                multimodal_outputs.append(inner_mm)
-        return multimodal_outputs
-
     # ---- Stage-local admission ----
 
     async def submit_initial(
@@ -1034,16 +658,11 @@ class StagePool:
             return []
         client = cast(StagePoolLLMClient, raw_client)
         processor = self.output_processor
-        iteration_stats = IterationStats()
         processed = processor.process_outputs(
             raw_outputs.outputs,
             raw_outputs.timestamp,
             iteration_stats,
         )
-        # Use the same wall-clock source as OrchestratorRequestState.stage_submit_ts.
-        # EngineCoreOutputs.timestamp may use a different clock base, which would
-        # make TTFO negative and get clamped to 0.
-        self.record_output_timestamps(processed.request_outputs, output_ts=_time.time())
 
         if processed.reqs_to_abort:
             await client.abort_requests_async(processed.reqs_to_abort)

vllm_omni/entrypoints/async_omni.py

@@ -261,7 +261,6 @@ class AsyncOmni(EngineClient, OmniBase):
         data_parallel_rank: int | None = None,
         reasoning_ended: bool | None = None,
         reasoning_parser_kwargs: dict[str, Any] | None = None,
-        arrival_time: float | None = None,
     ) -> AsyncGenerator[OmniRequestOutput, None]:
         """Generate outputs for the given prompt(s) asynchronously.
 
@@ -335,12 +334,11 @@ class AsyncOmni(EngineClient, OmniBase):
             )
 
             # Track per-request metrics
-            wall_start_ts = float(arrival_time) if arrival_time is not None else time.time()
+            wall_start_ts = time.time()
             req_start_ts: dict[str, float] = {}
 
             # Determine the final stage for E2E stats
             final_stage_id_for_e2e = self._compute_final_stage_id(output_modalities)
-            final_output_stage_ids = self._compute_final_output_stage_ids(output_modalities) or [final_stage_id_for_e2e]
 
             metrics = OrchestratorMetrics(
                 self.num_stages,
@@ -374,8 +372,6 @@ class AsyncOmni(EngineClient, OmniBase):
                     input_stream=prompt,
                     sampling_params_list=req_sp_list,
                     final_stage_id=final_stage_id_for_e2e,
-                    final_output_stage_ids=final_output_stage_ids,
-                    arrival_time=wall_start_ts,
                 )
             else:
                 await self.engine.add_request_async(
@@ -383,8 +379,6 @@ class AsyncOmni(EngineClient, OmniBase):
                     prompt=prompt,
                     sampling_params_list=req_sp_list,
                     final_stage_id=final_stage_id_for_e2e,
-                    final_output_stage_ids=final_output_stage_ids,
-                    arrival_time=wall_start_ts,
                 )
             submit_ts = time.time()
             req_state.metrics.stage_first_ts[0] = submit_ts
@@ -427,8 +421,6 @@ class AsyncOmni(EngineClient, OmniBase):
         input_stream: AsyncGenerator[StreamingInput, None],
         sampling_params_list: Sequence[OmniSamplingParams],
         final_stage_id: int,
-        final_output_stage_ids: Sequence[int],
-        arrival_time: float,
     ) -> asyncio.Task:
         """Submit a streaming input generator as incremental stage-0 updates."""
         if not sampling_params_list:
@@ -464,8 +456,6 @@ class AsyncOmni(EngineClient, OmniBase):
                             prompt_text=prompt_text,
                             sampling_params_list=chunk_sampling_params_list,
                             final_stage_id=final_stage_id,
-                            final_output_stage_ids=final_output_stage_ids,
-                            arrival_time=arrival_time,
                             resumable=True,
                         )
                         has_submitted_first_chunk = True
@@ -476,8 +466,6 @@ class AsyncOmni(EngineClient, OmniBase):
                             prompt_text=prompt_text,
                             sampling_params_list=chunk_sampling_params_list,
                             final_stage_id=final_stage_id,
-                            final_output_stage_ids=final_output_stage_ids,
-                            arrival_time=arrival_time,
                             resumable=True,
                         )
             except (asyncio.CancelledError, GeneratorExit):
@@ -504,8 +492,6 @@ class AsyncOmni(EngineClient, OmniBase):
                             prompt_text=None,
                             sampling_params_list=final_sampling_params_list,
                             final_stage_id=final_stage_id,
-                            final_output_stage_ids=final_output_stage_ids,
-                            arrival_time=arrival_time,
                             resumable=False,
                         )
                     else:
@@ -515,8 +501,6 @@ class AsyncOmni(EngineClient, OmniBase):
                             prompt_text=None,
                             sampling_params_list=final_sampling_params_list,
                             final_stage_id=final_stage_id,
-                            final_output_stage_ids=final_output_stage_ids,
-                            arrival_time=arrival_time,
                             resumable=False,
                         )
 

vllm_omni/entrypoints/cli/benchmark/serve.py

@@ -131,31 +131,6 @@ def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None:
     )
 
 
-def add_omni_benchmark_cli_args(parser: argparse.ArgumentParser) -> None:
-    """Add vLLM-Omni specific serving benchmark options."""
-    group = parser.add_argument_group("vLLM-Omni Multi-stage Benchmark Options")
-    group.add_argument(
-        "--print-stage",
-        action="store_true",
-        default=False,
-        help=(
-            "Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. "
-            "Disabled by default. The latency sections follow --percentile-metrics by modality: "
-            "ttft/tpot/itl control text stages, ttfc/tpoc/icl control internal stream stages, "
-            "and tpop controls both text TPOP and internal stream TPOP."
-        ),
-    )
-    group.add_argument(
-        "--image-edits-bot-task",
-        type=str,
-        default="think",
-        help=(
-            "Default bot_task form field for --backend openai-image-edits-omni "
-            '(/v1/images/edits). Use --extra-body \'{"bot_task":"..."}\' to override per run.'
-        ),
-    )
-
-
 _OMNI_BENCH_DATASET_CHOICES = (
     "daily-omni",
     "seed-tts",
@@ -198,39 +173,16 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
         # Add Daily-Omni specific arguments
         add_daily_omni_cli_args(parser)
         add_seed_tts_cli_args(parser)
-        add_omni_benchmark_cli_args(parser)
 
-        for action in parser._actions:
-            if action.dest == "dataset_name" and action.choices is not None:
-                extra = [
-                    c
-                    for c in (
-                        "daily-omni",
-                        "seed-tts",
-                        "seed-tts-text",
-                        "seed-tts-design",
-                        "ttsd",
-                        "sound-effect",
-                    )
-                    if c not in action.choices
-                ]
-                if extra:
-                    action.choices = list(action.choices) + extra
-            if action.dest == "backend" and action.choices is not None:
-                extra = [c for c in ("openai-image-edits-omni",) if c not in action.choices]
-                if extra:
-                    action.choices = list(action.choices) + extra
         _extend_omni_dataset_name_choices(parser)
 
         # Update help messages for omni-specific features
         for action in parser._actions:
             if action.dest == "percentile_metrics":
                 action.help = (
-                    "Comma-separated list of selected metrics to report percentiles. "
+                    "Comma-separated list of selected metrics to report percentiles."
-                    'For text metrics, "ttft", "tpot", and "itl" affect the global benchmark and text '
+                    "This argument specifies the metrics to report percentiles."
-                    'stage metrics. "tpop" also requests text TPOT/TPOP globally and per stage, and internal '
+                    'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf". '
-                    'stream TPOP. "ttfc", "tpoc", and "icl" only affect internal stream stage metrics. '
-                    'Audio metrics include "audio_ttfp", "audio_rtf", "audio_duration", and "audio_underrun".'
                 )
             if action.dest == "random_mm_limit_mm_per_prompt":
                 action.help = (
@@ -264,9 +216,4 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase):
             os.environ["SEED_TTS_WER_EVAL"] = "1"
         if getattr(args, "seed_tts_wer_save_items", False):
             os.environ["SEED_TTS_WER_SAVE_ITEMS"] = "1"
-        image_edits_bot_task = getattr(args, "image_edits_bot_task", None)
-        if image_edits_bot_task is not None:
-            extra_body = dict(getattr(args, "extra_body", None) or {})
-            extra_body.setdefault("bot_task", image_edits_bot_task)
-            args.extra_body = extra_body
         main(args)

vllm_omni/entrypoints/cli/serve.py

@@ -100,9 +100,6 @@ class OmniServeCommand(CLISubcommand):
             model_config = dict(existing) if isinstance(existing, dict) else {}
             model_config["guardrails"] = False
             args.model_config = model_config
-            explicit_keys = getattr(args, "explicit_keys", None)
-            if explicit_keys is not None:
-                args.explicit_keys = explicit_keys | {"model_config"}
 
         if args.headless:
             run_headless(args)

vllm_omni/entrypoints/omni.py

@@ -121,7 +121,6 @@ class Omni(OmniBase):
             for req_id, prompt in zip(request_ids, request_prompts):
                 prompt_modalities = prompt.get("modalities", None) if isinstance(prompt, dict) else None
                 final_stage_id = self._compute_final_stage_id(prompt_modalities)
-                final_output_stage_ids = self._compute_final_output_stage_ids(prompt_modalities) or [final_stage_id]
                 req_final_stage_ids[req_id] = final_stage_id
 
                 metrics = OrchestratorMetrics(
@@ -146,7 +145,6 @@ class Omni(OmniBase):
                     prompt=prompt,
                     sampling_params_list=req_sp_list,
                     final_stage_id=final_stage_id,
-                    final_output_stage_ids=final_output_stage_ids,
                 )
                 submit_ts = time.time()
                 req_state.metrics.stage_first_ts[0] = submit_ts

vllm_omni/entrypoints/omni_base.py

@@ -22,7 +22,7 @@ from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin
 from vllm_omni.entrypoints.utils import coerce_param_message_types, get_final_stage_id_for_e2e
 from vllm_omni.metrics.modality import OmniModalityMetrics, observe_modality_at_finalize
 from vllm_omni.metrics.prometheus import OmniPrometheusMetrics
-from vllm_omni.metrics.stats import OrchestratorAggregator
+from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics
 from vllm_omni.metrics.transfer import OmniTransferMetrics
 from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
 from vllm_omni.outputs import OmniRequestOutput
@@ -189,7 +189,6 @@ class OmniBase(PDDisaggregationMixin):
         self.async_chunk = bool(getattr(self.engine, "async_chunk", False))
 
         self.request_states: dict[str, ClientRequestState] = {}
-        self._consumed_metric_messages: dict[str, set[int]] = {}
         self.prom_metrics = OmniPrometheusMetrics(model_name=model, log_stats=log_stats)
         self.mod_metrics = OmniModalityMetrics(model_name=model, log_stats=log_stats)
 
@@ -309,7 +308,6 @@ class OmniBase(PDDisaggregationMixin):
             )
         finally:
             self.request_states.pop(request_id, None)
-            self._consumed_metric_messages.pop(request_id, None)
             # Republish gauges so any stale value left by the per-stage
             # publish in _process_single_result (which runs while the request
             # is still in self.request_states) is corrected after the pop.
@@ -328,17 +326,6 @@ class OmniBase(PDDisaggregationMixin):
             self._stage_meta_list,
         )
 
-    def _compute_final_output_stage_ids(self, output_modalities: list[str] | None) -> list[int]:
-        requested_modalities = output_modalities or self.output_modalities
-        requested_modalities = [m for m in requested_modalities if m in self.output_modalities]
-        if not requested_modalities:
-            requested_modalities = self.output_modalities
-        return [
-            sid
-            for sid, stage in enumerate(self._stage_meta_list)
-            if getattr(stage, "final_output", False) and stage.final_output_type in requested_modalities
-        ]
-
     def _process_stage_metrics_message(self, msg: StageMetricsMessage) -> None:
         req_id = msg.request_id
         req_state = self.request_states.get(req_id)
@@ -346,8 +333,7 @@ class OmniBase(PDDisaggregationMixin):
             return
         _m = msg.metrics
         stage_id = msg.stage_id
-        stage_meta = self.engine.get_stage_metadata(stage_id)
+        req_state.metrics.on_stage_metrics(stage_id, req_id, _m)
-        req_state.metrics.on_stage_metrics(stage_id, req_id, _m, stage_meta.final_output_type)
         submit_ts = msg.stage_submit_ts
         now = time.time()
         if req_state.metrics.stage_first_ts[stage_id] is None:
@@ -392,20 +378,6 @@ class OmniBase(PDDisaggregationMixin):
 
         req_state.stage_id = stage_id
 
-        if msg.metrics is not None and not msg.finished and req_state.metrics is not None:
-            stage_meta = self.engine.get_stage_metadata(stage_id)
-            output_type = getattr(msg.engine_outputs, "final_output_type", stage_meta.final_output_type)
-            msg_id = id(msg)
-            consumed = self._consumed_metric_messages.setdefault(req_id, set())
-            if msg_id not in consumed:
-                req_state.metrics.on_stage_metrics(stage_id, req_id, msg.metrics, output_type)
-                submit_ts = msg.stage_submit_ts
-                now = time.time()
-                if req_state.metrics.stage_first_ts[stage_id] is None:
-                    req_state.metrics.stage_first_ts[stage_id] = submit_ts if submit_ts is not None else now
-                req_state.metrics.stage_last_ts[stage_id] = max(req_state.metrics.stage_last_ts[stage_id] or 0.0, now)
-                consumed.add(msg_id)
-
         return False, req_id, stage_id, req_state
 
     def _check_engine_output_error(
@@ -443,7 +415,7 @@ class OmniBase(PDDisaggregationMixin):
         self,
         result: OutputMessage,
         stage_id: int,
-        metrics: OrchestratorAggregator,
+        metrics: OrchestratorMetrics,
         req_start_ts: dict[str, float],
         wall_start_ts: float,
         final_stage_id_for_e2e: int,
@@ -485,11 +457,10 @@ class OmniBase(PDDisaggregationMixin):
         metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, now)
 
         _m = result.metrics
-        stage_meta = self.engine.get_stage_metadata(stage_id)
-        output_type = getattr(engine_outputs, "final_output_type", stage_meta.final_output_type)
         if finished and _m is not None:
-            metrics.on_stage_metrics(stage_id, req_id, _m, output_type)
+            metrics.on_stage_metrics(stage_id, req_id, _m)
 
+        stage_meta = self.engine.get_stage_metadata(stage_id)
         if not stage_meta.final_output:
             return None
 
@@ -539,31 +510,6 @@ class OmniBase(PDDisaggregationMixin):
         self.prom_metrics.set_waiting(max(0, total - running))
 
         images = getattr(engine_outputs, "images", []) if output_type == "image" else []
-        response_metrics: dict[str, Any] = {}
-        stage_metrics: dict[str, dict[str, Any]] = {}
-        rid_key = str(req_id)
-        for evt in metrics.stage_events.get(rid_key, []):
-            if evt.stage_id is None:
-                continue
-            sid = int(evt.stage_id)
-            evt_stage_meta = self.engine.get_stage_metadata(sid)
-            evt_output_type = evt.final_output_type
-            if not evt_output_type:
-                evt_output_type = evt_stage_meta.final_output_type
-            stage_name = evt_stage_meta.model_stage
-            sid_key = str(sid)
-            stage_metrics[sid_key] = OrchestratorAggregator._merge_stage_metric_event(stage_metrics.get(sid_key), evt)
-            stage_metrics[sid_key]["stage_name"] = stage_name or f"stage_{sid}"
-            stage_metrics[sid_key]["final_output_type"] = evt_output_type
-        if stage_metrics:
-            response_metrics["stage_metrics"] = stage_metrics
-            current_stage_metrics = stage_metrics.get(str(stage_id))
-            if current_stage_metrics is not None:
-                response_metrics["stage_id"] = current_stage_metrics["stage_id"]
-                response_metrics["final_output_type"] = current_stage_metrics["final_output_type"]
-                if current_stage_metrics["final_output_type"] == "text":
-                    response_metrics["num_tokens_in"] = current_stage_metrics["num_tokens_in"]
-                    response_metrics["num_tokens_out"] = current_stage_metrics["num_tokens_out"]
         return OmniRequestOutput(
             request_id=req_id or "",
             stage_id=stage_id,
@@ -575,7 +521,6 @@ class OmniBase(PDDisaggregationMixin):
             trajectory_log_probs=getattr(engine_outputs, "trajectory_log_probs", None),
             trajectory_decoded=getattr(engine_outputs, "trajectory_decoded", None),
             _custom_output=getattr(engine_outputs, "_custom_output", {}),
-            metrics=response_metrics,
             stage_durations=stage_durations,
             peak_memory_mb=peak_memory_mb,
         )

vllm_omni/entrypoints/openai/api_server.py

@@ -29,7 +29,6 @@ from PIL import Image
 from pydantic import BaseModel, Field
 from starlette.datastructures import State
 from starlette.routing import Route
-from starlette.types import ASGIApp, Receive, Scope, Send
 from vllm.engine.protocol import EngineClient
 from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
 from vllm.entrypoints.chat_utils import load_chat_template
@@ -446,7 +445,6 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
 
         # OMNI: Pass supported_tasks to build_app (required by upstream vLLM)
         app = build_openai_app(args, supported_tasks)
-
         # OMNI: Remove upstream routes that we override with omni-specific handlers
         _remove_route_from_app(app, "/v1/chat/completions", {"POST"})
         _remove_route_from_app(app, "/v1/models", {"GET"})  # Remove upstream /v1/models to use omni's handler
@@ -479,29 +477,8 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None,
                 vllm_config.parallel_config._api_process_rank,
                 listen_address,
             )
-
-        class _TimestampMiddleware:
-            """Pure-ASGI outermost wrapper that stamps HTTP request arrival time.
-
-            Wraps the fully-built Starlette app as an outer ASGI layer so no
-            Starlette internals (user_middleware, middleware_stack, etc.) are
-            touched. Websocket and lifespan scopes pass through unchanged.
-            """
-
-            def __init__(self, inner: ASGIApp) -> None:
-                self._inner = inner
-
-            def __getattr__(self, name: str) -> Any:
-                return getattr(self._inner, name)
-
-            async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
-                if scope["type"] == "http":
-                    scope.setdefault("state", {})
-                    scope["state"]["request_timestamp"] = time.time()
-                await self._inner(scope, receive, send)
-
         shutdown_task = await serve_http(
-            _TimestampMiddleware(app),
+            app,
             sock=sock,
             enable_ssl_refresh=args.enable_ssl_refresh,
             host=args.host,
@@ -1615,7 +1592,6 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
     Raises:
         HTTPException: For validation errors, missing engine, or generation failures
     """
-    request_timestamp = float(getattr(raw_request.state, "request_timestamp", time.time()))
     # Get engine client (AsyncOmni) from app state
     engine_client, model_name, stage_configs = _get_engine_and_model(raw_request)
 
@@ -1678,7 +1654,6 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
                 extra_body=extra_body,
                 request_id=f"img_gen-{random_uuid()}",
                 raw_request=raw_request,
-                arrival_time=request_timestamp,
             )
             if isinstance(generation_result, ErrorResponse):
                 return JSONResponse(
@@ -1761,7 +1736,6 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
             stage_configs=stage_configs,
             prompt=prompt,
             request_id=request_id,
-            arrival_time=request_timestamp,
         )
 
         if result is None:
@@ -1852,14 +1826,12 @@ async def edit_images(
     bot_task: str | None = Form(None),
     sys_type: str | None = Form(None),
     system_prompt: str | None = Form(None),
-    return_stage_metrics: bool | None = Form(None),
 ) -> ImageGenerationResponse:
     """
     OpenAI-compatible image edit endpoint.
     """
 
     # 1. get engine and model
-    request_timestamp = float(getattr(raw_request.state, "request_timestamp", time.time()))
     engine_client, model_name, stage_configs = _get_engine_and_model(raw_request)
     if model is not None and model != model_name:
         raise HTTPException(
@@ -2015,16 +1987,6 @@ async def edit_images(
         _update_if_not_none(gen_params, "layers", layers)
         _update_if_not_none(gen_params, "resolution", resolution)
 
-        extra_args = dict(getattr(gen_params, "extra_args", {}) or {})
-        edit_extra_args = _build_hunyuan_edit_extra_args(
-            bot_task=bot_task,
-            sys_type=sys_type,
-            system_prompt=system_prompt,
-        )
-        extra_args.update(edit_extra_args)
-        if extra_args:
-            gen_params.extra_args = extra_args
-
         # 4. Generate images
         request_id = f"img_edit-{random_uuid()}"
         raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id)
@@ -2088,8 +2050,6 @@ async def edit_images(
                 extra_body["sys_type"] = sys_type
             if system_prompt is not None:
                 extra_body["system_prompt"] = system_prompt
-            if return_stage_metrics is not None:
-                extra_body["return_stage_metrics"] = return_stage_metrics
 
             prompt_text = prompt.get("prompt", "")
             generation_result = await chat_handler.generate_diffusion_images(
@@ -2097,7 +2057,6 @@ async def edit_images(
                 extra_body=extra_body,
                 reference_images=ref_b64_list,
                 request_id=request_id,
-                arrival_time=request_timestamp,
                 stream=stream,
                 model=model_name,
                 output_format=output_format,
@@ -2323,28 +2282,6 @@ def _check_max_generated_image_size(
             )
 
 
-def _build_hunyuan_edit_extra_args(
-    *,
-    bot_task: str | None,
-    sys_type: str | None,
-    system_prompt: str | None,
-) -> dict[str, Any]:
-    """Map Hunyuan /v1/images/edits form fields to DiT ``extra_args``."""
-    extra_args: dict[str, Any] = {}
-    effective_use_system_prompt = sys_type
-    if effective_use_system_prompt is None and bot_task is not None:
-        from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import resolve_sys_type
-
-        effective_use_system_prompt = resolve_sys_type(bot_task)
-    if effective_use_system_prompt is not None:
-        extra_args["use_system_prompt"] = effective_use_system_prompt
-    if system_prompt is not None:
-        extra_args["system_prompt"] = system_prompt
-    if bot_task is not None:
-        extra_args["bot_task"] = bot_task
-    return extra_args
-
-
 def _update_if_not_none(object: Any, key: str, val: Any) -> None:
     if val is not None:
         setattr(object, key, val)
@@ -2608,7 +2545,7 @@ async def _run_video_generation_job(
     started_at = time.perf_counter()
     output_path = None
     try:
-        video_bytes, stage_durations, peak_memory_mb, action = await handler.generate_video_bytes(
+        video_bytes, stage_durations, peak_memory_mb = await handler.generate_video_bytes(
             request, video_id, reference_image=reference_image
         )
 
@@ -2626,7 +2563,6 @@ async def _run_video_generation_job(
                 "inference_time_s": time.perf_counter() - started_at,
                 "stage_durations": stage_durations,
                 "peak_memory_mb": peak_memory_mb,
-                "action": action,
             },
         )
     except (EngineGenerateError, EngineDeadError) as exc:
@@ -2838,7 +2774,7 @@ async def create_video_sync(
     raw_request.state.request_metadata = RequestResponseMetadata(request_id=request_id)
     started_at = time.perf_counter()
     try:
-        video_bytes, stage_durations, peak_memory_mb, _action = await asyncio.wait_for(
+        video_bytes, stage_durations, peak_memory_mb = await asyncio.wait_for(
             handler.generate_video_bytes(request, request_id, reference_image=reference_image),
             timeout=VIDEO_SYNC_TIMEOUT_S,
         )

vllm_omni/entrypoints/openai/protocol/__init__.py

@@ -13,7 +13,6 @@ from vllm_omni.entrypoints.openai.protocol.images import (
     ResponseFormat,
 )
 from vllm_omni.entrypoints.openai.protocol.videos import (
-    VideoAction,
     VideoData,
     VideoGenerationRequest,
     VideoGenerationResponse,
@@ -28,7 +27,6 @@ __all__ = [
     "ImageGenerationRequest",
     "ImageGenerationResponse",
     "ResponseFormat",
-    "VideoAction",
     "VideoData",
     "VideoGenerationRequest",
     "VideoGenerationResponse",

vllm_omni/entrypoints/openai/protocol/audio.py

@@ -115,54 +115,6 @@ class OpenAICreateSpeechRequest(BaseModel):
             raise ValueError("'speaker_embedding' values must be finite (no NaN or Inf)")
         return v
 
-    @model_validator(mode="before")
-    @classmethod
-    def normalize_references_alias(cls, data):
-        """Map the BosonAI Higgs Audio v3 cookbook ``references`` array onto ``ref_audio`` / ``ref_text``.
-
-        Upstream Higgs Audio v3 examples and the boson.ai cookbook send voice
-        clones as::
-
-            {"references": [{"audio_path": str, "text": str}]}
-
-        vllm-omni's existing convention is ``ref_audio`` + ``ref_text``.
-        Without this alias pydantic silently drops the unknown ``references``
-        field and the request falls through to zero-shot synthesis with no
-        error, which is hard to debug from the client side. The normalizer
-        translates a single reference and rejects multi-reference payloads
-        (vllm-omni does not yet support multi-shot voice clone) or conflicts
-        with the explicit ``ref_audio`` / ``ref_text`` fields.
-        """
-        if not isinstance(data, dict):
-            return data
-        refs = data.get("references")
-        if refs is None:
-            return data
-        if not isinstance(refs, list) or len(refs) == 0:
-            raise ValueError("'references' must be a non-empty array of {audio_path, text} objects")
-        if len(refs) > 1:
-            raise ValueError("'references' only supports a single reference; multi-shot voice clone is not supported")
-        first = refs[0]
-        if not isinstance(first, dict):
-            raise ValueError("'references[0]' must be an object with 'audio_path' and optional 'text'")
-        audio_path = first.get("audio_path")
-        if not isinstance(audio_path, str) or not audio_path:
-            raise ValueError(
-                "'references[0].audio_path' is required and must be a string (URL, data: URI, or file path)"
-            )
-        existing_ref_audio = data.get("ref_audio")
-        if existing_ref_audio and existing_ref_audio != audio_path:
-            raise ValueError("'references' and 'ref_audio' are mutually exclusive")
-        data["ref_audio"] = audio_path
-        text = first.get("text")
-        if isinstance(text, str) and text.strip():
-            existing_ref_text = data.get("ref_text")
-            if existing_ref_text and existing_ref_text != text:
-                raise ValueError("'references[0].text' and 'ref_text' conflict; supply only one")
-            data["ref_text"] = text
-        data.pop("references", None)
-        return data
-
     @model_validator(mode="before")
     @classmethod
     def reject_higgs_audio_v2_unsupported_aliases(cls, data):

vllm_omni/entrypoints/openai/protocol/images.py

@@ -193,10 +193,6 @@ class ImageEditARDeltaChunk(BaseModel):
     index: int = Field(default=0, description="Completion index for the AR stage output")
     created: int = Field(..., description="Unix timestamp of when the stream was created")
     model: str = Field(..., description="Model used for the image edit request")
-    metrics: dict[str, Any] | None = Field(
-        default=None,
-        description="Optional vLLM-Omni per-stage metrics snapshot for benchmark clients.",
-    )
 
 
 class ImageEditImageChunk(BaseModel):
@@ -209,10 +205,6 @@ class ImageEditImageChunk(BaseModel):
     size: str = Field(..., description="The generated image size")
     created: int = Field(..., description="Unix timestamp of when the stream was created")
     model: str = Field(..., description="Model used for the image edit request")
-    metrics: dict[str, Any] | None = Field(
-        default=None,
-        description="Optional vLLM-Omni per-stage metrics snapshot for benchmark clients.",
-    )
 
 
 class ImageEditStreamError(BaseModel):

vllm_omni/entrypoints/openai/protocol/videos.py

@@ -220,24 +220,12 @@ class VideoGenerationRequest(BaseModel):
         return vp
 
 
-class VideoAction(BaseModel):
-    """Generated action sequence returned by action-capable video models."""
-
-    data: list[Any] = Field(..., description="JSON-serializable nested action values")
-    shape: list[int] = Field(..., description="Shape of the returned action data")
-    dtype: str | None = Field(default=None, description="Source action dtype, if available")
-    raw_action_dim: int | None = Field(default=None, description="Raw action dimension requested by the model")
-    action_mode: str | None = Field(default=None, description="Action generation mode")
-    domain_id: int | None = Field(default=None, description="Action embodiment domain id")
-
-
 class VideoData(BaseModel):
     """Single generated video data."""
 
     b64_json: str | None = Field(default=None, description="Base64-encoded MP4 video")
     url: str | None = Field(default=None, description="Video URL (not implemented)")
     revised_prompt: str | None = Field(default=None, description="Revised prompt (OpenAI compatibility, always null)")
-    action: VideoAction | None = Field(default=None, description="Generated action sequence metadata, if any")
 
 
 class VideoGenerationResponse(BaseModel):
@@ -310,7 +298,6 @@ class VideoResponse(BaseModel):
         default=0.0,
         description="Peak device memory usage in MB reported by the diffusion pipeline.",
     )
-    action: VideoAction | None = Field(default=None, description="Generated action sequence metadata, if any")
 
     @property
     def file_extension(self) -> str:

vllm_omni/entrypoints/openai/realtime_connection.py

@@ -86,14 +86,9 @@ class RealtimeConnection(VllmRealtimeConnection):
             return [], 24000
 
         sr = mm.get("sr") or mm.get("sample_rate") or mm.get("audio_sample_rate") or 24000
-        if isinstance(sr, (list, tuple)) and sr:
-            sr = sr[-1]
-        if hasattr(sr, "item"):
-            sr = sr.item()
-        sample_rate_hz = int(sr)
         key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None)
         if key is None:
-            return [], sample_rate_hz
+            return [], int(sr)
 
         raw_audio = mm.get(key)
         chunks: list[np.ndarray] = []
@@ -106,7 +101,7 @@ class RealtimeConnection(VllmRealtimeConnection):
             arr = self._tensor_to_numpy(raw_audio)
             if arr is not None and arr.size > 0:
                 chunks.extend(self._raw_waveform_to_deltas(arr))
-        return chunks, sample_rate_hz
+        return chunks, int(sr)
 
     @staticmethod
     def _pcm16_b64(audio_f32: np.ndarray) -> str: