[CI] Revert "[Feature] gen-only FP8 Quantization Support for SenseNova-U1" (#4196)

Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>

.buildkite/test-nightly.yml

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

.buildkite/test-ready.yml

@@ -273,6 +273,42 @@ steps:
                       path: /mnt/hf-cache
                       type: DirectoryOrCreate
 
+      - label: "Omni · Qwen3-Omni Realtime WebSocket"
+        commands:
+          - |
+            timeout 30m bash -c "
+              export VLLM_WORKER_MULTIPROC_METHOD=spawn
+              pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_streaming_audio_input_pcm_output_async_chunk_smoke -m 'core_model' --run-level 'core_model'
+            "
+        agents:
+          queue: "mithril-h100-pool"
+        plugins:
+          - kubernetes:
+              podSpec:
+                containers:
+                  - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                    resources:
+                      limits:
+                        nvidia.com/gpu: 2
+                    volumeMounts:
+                      - name: devshm
+                        mountPath: /dev/shm
+                      - name: hf-cache
+                        mountPath: /root/.cache/huggingface
+                    env:
+                      - name: HF_HOME
+                        value: /root/.cache/huggingface
+                nodeSelector:
+                  node.kubernetes.io/instance-type: gpu-h100-sxm
+                volumes:
+                  - name: devshm
+                    emptyDir:
+                      medium: Memory
+                  - name: hf-cache
+                    hostPath:
+                      path: /mnt/hf-cache
+                      type: DirectoryOrCreate
+
       - label: "Omni · VoxCPM2 Test"
         timeout_in_minutes: 20
         commands:

benchmarks/accuracy/common.py

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

benchmarks/accuracy/image_to_image/gedit_bench.py

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

docs/cli/bench/serve.md

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

docs/models/supported_models.md

@@ -33,7 +33,7 @@ th {
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, T2V with sound | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
+| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
 | `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |

docs/user_guide/diffusion_features.md

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

examples/offline_inference/speech_to_video/speech_to_video.py

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

pyproject.toml

@@ -64,7 +64,9 @@ dev = [
     "lpips>=0.1.4", # for quantization fp8
     "orjson>=3.11.7", # to make v1/embeddings API fast
     "voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests
-    "torchdiffeq>=0.2.5" # for covo_audio Test
+    "torchdiffeq>=0.2.5", # for covo_audio Test
+    "FlagEmbedding", # for hunyuan_image3 tests
+    "kernels<0.15.0" # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
 ]
 
 demo = [

recipes/README.md

@@ -36,8 +36,8 @@ recipes/
 | [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
 | [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
 | [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
-| [`nvidia/Cosmos3-Nano.md`](./nvidia/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video generation, text to video with sound  | 1x H200 141GB / B300 |
-| [`nvidia/Cosmos3-Super.md`](./nvidia/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) | 8x H200/H100/A100 / 2x H200 / B300 |
+| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video generation, text to video with sound, action policy  | 1x H200 141GB / B300 |
+| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
 | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
 | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |

recipes/cosmos3/Cosmos3-Nano.md

@@ -6,7 +6,7 @@
 
 - Vendor: NVIDIA
 - Model: `nvidia/Cosmos3-Nano`
-- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation, with optional synchronized audio (video + sound)
+- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation, with optional synchronized audio (video + sound), action policy
 - Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API
 - Maintainer: Community
 
@@ -23,6 +23,23 @@ the mode is selected per request:
 - **T2VS / I2VS** — add `generate_sound=true` (and optional `sound_duration`) to a
   T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
   the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
+- **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks:
+  - `forward_dynamics` — given a first frame **and** an action trajectory, roll out
+    the resulting video. Synchronous: `POST /v1/videos/sync`.
+  - `policy` — given a first frame and a language instruction, **predict** the action
+    trajectory (and a rollout video). Use the async `POST /v1/videos` endpoint and
+    read the predicted action from the top-level `action` field
+    (`{data, shape, dtype, raw_action_dim, domain_id}`).
+
+  Action requests also take `domain_name` (e.g. `av`, `bridge_orig_lerobot`,
+  `droid_lerobot`, `agibotworld`, …; or a numeric `domain_id`), `raw_action_dim`,
+  and `action_chunk_size` (must equal `num_frames` or `num_frames - 1`). For
+  `forward_dynamics` also pass the `action` array. The dedicated policy checkpoint
+  **`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way
+  (`domain_name=droid_lerobot`).
+
+  `inverse_dynamics` (recover the action from a given video) is supported by the
+  pipeline; **online inference of inverse dynamics will be added in a follow-up MR.**
 
 ## References
 
@@ -144,6 +161,35 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
   -F "sound_duration=7.875" \
   -F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
   -o cosmos3_t2v_with_sound.mp4
+
+# Action — forward dynamics (first frame + action trajectory -> rollout video).
+# Synchronous; `action` is a JSON array shaped [action_chunk_size, raw_action_dim].
+curl -sS -X POST http://localhost:8000/v1/videos/sync \
+  -H "Accept: video/mp4" \
+  --form-string "model=nvidia/Cosmos3-Nano" \
+  --form-string "prompt=You are an autonomous vehicle. This video is captured from a first-person perspective." \
+  -F "input_reference=@first_frame.jpg;type=image/jpeg" \
+  -F "size=640x480" -F "num_frames=61" -F "fps=10" \
+  -F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
+  --form-string "extra_params={\"action_mode\":\"forward_dynamics\",\"domain_name\":\"av\",\"raw_action_dim\":9,\"action_chunk_size\":60,\"action\":$(cat action.json)}" \
+  -F "seed=0" \
+  -o cosmos3_forward_dynamics.mp4
+
+# Action — policy (first frame + instruction -> predicted action trajectory + video).
+# Asynchronous: POST returns a job id; poll, then read the predicted action from
+# the top-level `action` field ({data, shape, dtype, raw_action_dim, domain_id}).
+VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
+  -H "Accept: application/json" \
+  --form-string "model=nvidia/Cosmos3-Nano" \
+  --form-string "prompt=Pick up the banana and place it in the bowl." \
+  -F "input_reference=@first_frame.jpg;type=image/jpeg" \
+  -F "size=640x480" -F "num_frames=17" -F "fps=5" \
+  -F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
+  --form-string 'extra_params={"action_mode":"policy","domain_name":"bridge_orig_lerobot","raw_action_dim":10,"action_chunk_size":16}' \
+  -F "seed=0" | jq -r '.id')
+# poll until status == completed, then:
+curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
+curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_policy.mp4
 ```
 
 #### Notes
@@ -152,6 +198,7 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
   - T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s**
   - T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s**
   - I2V 1280×720, 189 frames @ 35 steps → ~**99 s**
+  - Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**.
   - Guardrails-on overhead: ~8% on T2I, negligible on video.
 - **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 GPU;
   full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
@@ -173,8 +220,10 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
     the server fails at pipeline build with a gated-repo / safety-checker error.
   - A guardrail-blocked prompt currently returns HTTP 500
     (`"Guardrail blocked prompt"`).
-  - Action (policy / forward- / inverse-dynamics) modalities are not part of
-    this integration yet.
+  - Action `forward_dynamics` (sync `/v1/videos/sync`) and `policy` (async
+    `/v1/videos`, returns the predicted action under the top-level `action`
+    field) are supported online. **Online inference of inverse dynamics will be
+    added in a follow-up MR.**
 
 ### 1x GPU (Offline generation)
 

recipes/cosmos3/Cosmos3-Super.md

@@ -104,5 +104,9 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
   - (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
 - **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
 - Same generation defaults, supported sizes, and `generate_sound`/`sound_duration`
-  semantics as Nano. Action (policy / forward- / inverse-dynamics) modalities are
-  not part of this integration yet.
+  semantics as Nano, including the **action** modality: `forward_dynamics`
+  (sync `/v1/videos/sync`) and `policy` (async `/v1/videos`, predicted action under
+  the top-level `action` field) — see the Cosmos3-Nano recipe for the request shape.
+  Online inference of inverse dynamics will be added in a follow-up MR. Verified on
+  the 64B Super under `--cfg-parallel-size 2`: async `policy` returns the predicted
+  action (`[16, 10]`) and the rollout video reliably.

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

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

tests/benchmarks/test_accuracy_bench_utils.py

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

tests/core/sched/test_omni_ar_scheduler_streaming.py

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

tests/dfx/perf/tests/test_hunyuan_image3_it2i.json

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

tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json

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

tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json

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

tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py

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

tests/diffusion/models/cosmos3/test_cosmos3_transformer.py

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

tests/distributed/omni_connectors/test_bagel_mooncake_connector.py

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

tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py

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

tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py

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

tests/e2e/accuracy/test_diffusers_backend_similarity.py

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

tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

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

tests/e2e/accuracy/test_qwen_image.py

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

tests/e2e/offline_inference/test_moss_tts.py

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

tests/e2e/offline_inference/test_moss_tts_nano_expansion.py

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

tests/e2e/offline_inference/test_moss_tts_v1_5.py

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

tests/e2e/online_serving/test_glm_image_expansion.py

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

tests/e2e/online_serving/test_moss_tts_expansion.py

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

tests/e2e/online_serving/test_qwen3_omni.py

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

tests/engine/test_orchestrator.py

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

tests/entrypoints/openai_api/test_image_server.py

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

tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py

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

tests/entrypoints/openai_api/test_serving_chat_multistage_generation.py

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

tests/entrypoints/openai_api/test_video_server.py

@@ -627,6 +627,115 @@ def test_video_job_persists_profiler_metadata(test_client, mocker: MockerFixture
 
     assert completed["stage_durations"] == {"diffuse": 2.5, "vae.decode": 0.3}
     assert completed["peak_memory_mb"] == 4096.5
+    assert completed["action"] is None
+
+
+def test_video_generation_response_exposes_action_payload(mocker: MockerFixture):
+    engine = FakeAsyncOmni()
+    handler = OmniOpenAIServingVideo.for_diffusion(
+        diffusion_engine=engine,
+        model_name="Cosmos3-8B-UVA",
+    )
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        del prompt, request_id, sampling_params_list
+        import numpy as np
+
+        yield MockVideoResult(
+            [object()],
+            custom_output={
+                "action": np.array([[[1.5, 2.5], [3.5, 4.5]]], dtype=np.float32),
+                "raw_action_dim": 2,
+                "action_mode": "policy",
+                "domain_id": 7,
+            },
+        )
+
+    engine.generate = _generate
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
+        return_value="encoded-video",
+    )
+
+    response = asyncio.run(
+        handler.generate_videos(
+            VideoGenerationRequest(prompt="predict actions"),
+            "action-json",
+        )
+    )
+
+    action = response.data[0].action
+    assert action is not None
+    assert action.data == [[1.5, 2.5], [3.5, 4.5]]
+    assert action.shape == [2, 2]
+    assert action.dtype == "float32"
+    assert action.raw_action_dim == 2
+    assert action.action_mode == "policy"
+    assert action.domain_id == 7
+    assert response.model_dump(mode="json")["data"][0]["action"]["data"] == [[1.5, 2.5], [3.5, 4.5]]
+
+
+def test_video_job_persists_action_metadata(test_client, mocker: MockerFixture):
+    engine = test_client.app.state.openai_serving_video._engine_client
+
+    async def _generate(prompt, request_id, sampling_params_list):
+        import numpy as np
+
+        engine.captured_prompt = prompt
+        engine.captured_sampling_params_list = sampling_params_list
+        yield MockVideoResult(
+            [object()],
+            custom_output={
+                "action": np.array([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32),
+                "raw_action_dim": 2,
+                "action_mode": "policy",
+                "domain_id": 7,
+            },
+        )
+
+    engine.generate = _generate
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+
+    response = test_client.post("/v1/videos", data={"prompt": "profile me"})
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    completed = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+
+    expected_action = {
+        "data": [[1.0, 2.0], [3.0, 4.0]],
+        "shape": [2, 2],
+        "dtype": "float32",
+        "raw_action_dim": 2,
+        "action_mode": "policy",
+        "domain_id": 7,
+    }
+    assert completed["action"] == expected_action
+
+    listed = test_client.get("/v1/videos").json()
+    assert listed["data"][0]["action"] == expected_action
+
+
+def test_action_extraction_accepts_unbatched_action():
+    import numpy as np
+
+    result = MockVideoResult(
+        [object()],
+        custom_output={
+            "action": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32),
+            "raw_action_dim": 2,
+            "action_mode": "policy",
+            "domain_id": 7,
+        },
+    )
+
+    actions = OmniOpenAIServingVideo._extract_action_outputs(result, expected_count=1)
+
+    assert actions[0] is not None
+    assert actions[0].data == [[1.0, 2.0], [3.0, 4.0]]
+    assert actions[0].shape == [2, 2]
 
 
 def test_missing_handler_returns_503():
@@ -1038,6 +1147,31 @@ def test_sample_solver_forwarded_via_extra_params(test_client, mocker: MockerFix
     assert captured.extra_args["sample_solver"] == "euler"
 
 
+def test_extra_params_allows_inline_action(test_client, mocker: MockerFixture):
+    """Inline ``action`` data is accepted and forwarded verbatim to
+    ``extra_args`` (the supported way to pass forward-dynamics actions)."""
+    mocker.patch(
+        "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
+        return_value=b"fake-video",
+    )
+    action = [[0.1, 0.2], [0.3, 0.4]]
+    response = test_client.post(
+        "/v1/videos",
+        data={
+            "prompt": "forward dynamics inline",
+            "extra_params": json.dumps({"action_mode": "forward_dynamics", "action": action}),
+        },
+    )
+
+    assert response.status_code == 200
+    video_id = response.json()["id"]
+    _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
+    engine = test_client.app.state.openai_serving_video._engine_client
+    captured = engine.captured_sampling_params_list[0]
+    assert captured.extra_args["action"] == action
+    assert captured.extra_args["action_mode"] == "forward_dynamics"
+
+
 # ---------------------------------------------------------------------------
 # Sync endpoint tests (POST /v1/videos/sync)
 # ---------------------------------------------------------------------------

tests/entrypoints/test_async_omni.py

@@ -16,7 +16,7 @@ OMNI_MODEL = "Qwen/Qwen2.5-Omni-7B"
 OMNI_STAGE_CONFIG = get_deploy_config_path("ci/qwen2_5_omni_thinker_only.yaml")
 
 
-async def _noop(**kw):
+async def _noop(*args, **kw):
     pass
 
 
@@ -61,6 +61,7 @@ def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) ->
     omni._final_output_handler = lambda: None
     omni.resolve_sampling_params_list = lambda params, allow_delta_coercion: params
     omni._compute_final_stage_id = lambda output_modalities: 0
+    omni._compute_final_output_stage_ids = lambda output_modalities: [0]
     omni._process_orchestrator_results = fake_process_results
     omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None)
     return omni

tests/entrypoints/test_stage_utils.py

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

tests/helpers/runtime.py

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

tests/metrics/test_prometheus.py

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

tests/model_executor/models/voxcpm2/test_talker_state_eviction.py

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

tests/test_config_factory.py

@@ -1638,9 +1638,6 @@ class TestBaseConfigInheritance:
         s0 = deploy.stages[0].default_sampling_params
         # CI overrides max_tokens
         assert s0["max_tokens"] == 150
-        # Inherited from base
-        assert s0["temperature"] == 0.4
-        assert s0["seed"] == 42
 
     def test_pure_inheritance_overlay(self, tmp_path):
         """An overlay with only ``base_config`` inherits everything."""
@@ -1960,5 +1957,3 @@ class TestSamplingConstraintsPrecedence:
         assert stages[0].yaml_extras["default_sampling_params"]["detokenize"] is True
         # Pipeline says stop_token_ids=[2150] for talker
         assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [2150]
-        # Deploy temperature still flows through
-        assert stages[0].yaml_extras["default_sampling_params"]["temperature"] == 0.4

tests/worker/test_gpu_ar_model_runner.py

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

vllm_omni/assets/video.py

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

vllm_omni/benchmarks/metrics/metrics.py

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

vllm_omni/benchmarks/patch/patch.py

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

vllm_omni/benchmarks/serve.py

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

vllm_omni/core/sched/omni_ar_scheduler.py

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

vllm_omni/deploy/qwen3_omni_moe.yaml

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

vllm_omni/deploy/qwen3_omni_moe_multi_replicas.yaml

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

vllm_omni/deploy/voxcpm2.yaml

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

vllm_omni/diffusion/cache/cache_dit_backend.py

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

vllm_omni/diffusion/diffusion_engine.py

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

vllm_omni/diffusion/inline_stage_diffusion_client.py

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

vllm_omni/diffusion/layers/rope.py

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

vllm_omni/diffusion/models/bagel/pipeline_bagel.py

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

vllm_omni/diffusion/models/cosmos3/action.py

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

vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py

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

vllm_omni/diffusion/models/cosmos3/transformer_cosmos3.py

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

vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py

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

vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py

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

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py

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

vllm_omni/diffusion/stage_diffusion_client.py

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

vllm_omni/engine/arg_utils.py

@@ -77,13 +77,10 @@ def register_omni_models_to_vllm():
 
     _register_omni_hf_configs()
 
-    # 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.
+    supported_archs = ModelRegistry.get_supported_archs()
     for arch, (mod_folder, mod_relname, cls_name) in _OMNI_MODELS.items():
-        ModelRegistry.register_model(arch, f"vllm_omni.model_executor.models.{mod_folder}.{mod_relname}:{cls_name}")
+        if arch not in supported_archs:
+            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,6 +123,7 @@ class StageRuntimeInfo:
     final_output: bool
     final_output_type: FinalOutputModalityType | None
     stage_type: str
+    model_stage: str | None = None
 
 
 # ============================================================================
@@ -1247,6 +1248,7 @@ 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),
                 )
             )
 
@@ -1616,6 +1618,7 @@ 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,
@@ -1628,6 +1631,7 @@ 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
@@ -1713,7 +1717,9 @@ 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(),
         )
 
@@ -2186,6 +2192,7 @@ 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,
@@ -2209,6 +2216,7 @@ 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,
@@ -2236,6 +2244,7 @@ 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,
@@ -2253,6 +2262,7 @@ 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,
@@ -2270,6 +2280,7 @@ 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,
@@ -2281,6 +2292,7 @@ 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",
@@ -2294,6 +2306,7 @@ 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,
@@ -2305,6 +2318,7 @@ 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,7 +24,9 @@ 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,6 +126,11 @@ 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)
@@ -451,6 +456,7 @@ 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 "
@@ -469,6 +475,8 @@ 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
@@ -512,7 +520,9 @@ 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)
@@ -558,6 +568,8 @@ 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()
@@ -665,6 +677,7 @@ 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:
@@ -684,6 +697,8 @@ 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.
@@ -775,7 +790,9 @@ 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)
 
@@ -814,6 +831,34 @@ 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)
@@ -852,6 +897,11 @@ 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(
@@ -860,7 +910,7 @@ class Orchestrator:
                     replica_id=replica_id,
                     engine_outputs=output,
                     metrics=stage_metrics,
-                    finished=finished and stage_id == req_state.final_stage_id,
+                    finished=request_finished,
                     stage_submit_ts=submit_ts,
                 )
             )
@@ -915,7 +965,7 @@ class Orchestrator:
                         is_final_update=True,
                     )
 
-        if finished and stage_id == req_state.final_stage_id:
+        if request_finished:
             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:
@@ -1341,7 +1391,12 @@ class Orchestrator:
                     resumable=downstream_resumable,
                 )
                 request.external_req_id = request.request_id
-                await next_pool.submit_initial(request_id, req_state, request, prompt_text=None)
+                await next_pool.submit_initial(
+                    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
+from vllm.v1.metrics.stats import IterationStats, RequestStateStats
 
 from vllm_omni.data_entry_keys import unflatten_payload
 from vllm_omni.engine.output_modality import DRAINABLE_MODALITIES
@@ -35,13 +35,21 @@ 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."""
@@ -318,6 +326,32 @@ 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,
@@ -371,6 +405,7 @@ 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]
@@ -402,3 +437,73 @@ 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,6 +428,8 @@ 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,
@@ -440,12 +442,11 @@ 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=None,
+            model_stage=model_stage,
             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,6 +24,7 @@ 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
@@ -96,6 +97,10 @@ 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] = {}
@@ -428,6 +433,10 @@ 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."""
@@ -485,13 +494,66 @@ 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:
@@ -504,6 +566,7 @@ 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,
@@ -519,8 +582,321 @@ 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(
@@ -658,11 +1034,16 @@ 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,6 +261,7 @@ 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.
 
@@ -334,11 +335,12 @@ class AsyncOmni(EngineClient, OmniBase):
             )
 
             # Track per-request metrics
-            wall_start_ts = time.time()
+            wall_start_ts = float(arrival_time) if arrival_time is not None else 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,
@@ -372,6 +374,8 @@ 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(
@@ -379,6 +383,8 @@ 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
@@ -421,6 +427,8 @@ 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:
@@ -456,6 +464,8 @@ 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
@@ -466,6 +476,8 @@ 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):
@@ -492,6 +504,8 @@ 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:
@@ -501,6 +515,8 @@ 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,6 +131,31 @@ 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",
@@ -173,16 +198,39 @@ 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."
-                    "This argument specifies the metrics to report percentiles."
-                    'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf". '
+                    "Comma-separated list of selected metrics to report percentiles. "
+                    'For text metrics, "ttft", "tpot", and "itl" affect the global benchmark and text '
+                    'stage metrics. "tpop" also requests text TPOT/TPOP globally and per stage, and internal '
+                    '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 = (
@@ -216,4 +264,9 @@ 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,6 +100,9 @@ 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,6 +121,7 @@ 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(
@@ -145,6 +146,7 @@ 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 as OrchestratorMetrics
+from vllm_omni.metrics.stats import OrchestratorAggregator
 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,6 +189,7 @@ 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)
 
@@ -308,6 +309,7 @@ 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.
@@ -326,6 +328,17 @@ 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)
@@ -333,7 +346,8 @@ class OmniBase(PDDisaggregationMixin):
             return
         _m = msg.metrics
         stage_id = msg.stage_id
-        req_state.metrics.on_stage_metrics(stage_id, req_id, _m)
+        stage_meta = self.engine.get_stage_metadata(stage_id)
+        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:
@@ -378,6 +392,20 @@ 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(
@@ -415,7 +443,7 @@ class OmniBase(PDDisaggregationMixin):
         self,
         result: OutputMessage,
         stage_id: int,
-        metrics: OrchestratorMetrics,
+        metrics: OrchestratorAggregator,
         req_start_ts: dict[str, float],
         wall_start_ts: float,
         final_stage_id_for_e2e: int,
@@ -457,10 +485,11 @@ class OmniBase(PDDisaggregationMixin):
         metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, now)
 
         _m = result.metrics
-        if finished and _m is not None:
-            metrics.on_stage_metrics(stage_id, req_id, _m)
-
         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)
+
         if not stage_meta.final_output:
             return None
 
@@ -510,6 +539,31 @@ 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,
@@ -521,6 +575,7 @@ 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,6 +29,7 @@ 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
@@ -445,6 +446,7 @@ 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
@@ -477,8 +479,29 @@ 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(
-            app,
+            _TimestampMiddleware(app),
             sock=sock,
             enable_ssl_refresh=args.enable_ssl_refresh,
             host=args.host,
@@ -1592,6 +1615,7 @@ 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)
 
@@ -1654,6 +1678,7 @@ 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(
@@ -1736,6 +1761,7 @@ 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:
@@ -1826,12 +1852,14 @@ 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(
@@ -1987,6 +2015,16 @@ 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)
@@ -2050,6 +2088,8 @@ 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(
@@ -2057,6 +2097,7 @@ 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,
@@ -2282,6 +2323,28 @@ 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)
@@ -2545,7 +2608,7 @@ async def _run_video_generation_job(
     started_at = time.perf_counter()
     output_path = None
     try:
-        video_bytes, stage_durations, peak_memory_mb = await handler.generate_video_bytes(
+        video_bytes, stage_durations, peak_memory_mb, action = await handler.generate_video_bytes(
             request, video_id, reference_image=reference_image
         )
 
@@ -2563,6 +2626,7 @@ 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:
@@ -2774,7 +2838,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 = await asyncio.wait_for(
+        video_bytes, stage_durations, peak_memory_mb, _action = 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,6 +13,7 @@ from vllm_omni.entrypoints.openai.protocol.images import (
     ResponseFormat,
 )
 from vllm_omni.entrypoints.openai.protocol.videos import (
+    VideoAction,
     VideoData,
     VideoGenerationRequest,
     VideoGenerationResponse,
@@ -27,6 +28,7 @@ __all__ = [
     "ImageGenerationRequest",
     "ImageGenerationResponse",
     "ResponseFormat",
+    "VideoAction",
     "VideoData",
     "VideoGenerationRequest",
     "VideoGenerationResponse",

vllm_omni/entrypoints/openai/protocol/images.py

@@ -193,6 +193,10 @@ 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):
@@ -205,6 +209,10 @@ 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,12 +220,24 @@ 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):
@@ -298,6 +310,7 @@ 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,9 +86,14 @@ 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 [], int(sr)
+            return [], sample_rate_hz
 
         raw_audio = mm.get(key)
         chunks: list[np.ndarray] = []
@@ -101,7 +106,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, int(sr)
+        return chunks, sample_rate_hz
 
     @staticmethod
     def _pcm16_b64(audio_f32: np.ndarray) -> str:

vllm_omni/entrypoints/openai/serving_chat.py

@@ -216,6 +216,27 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         self._supported_speakers = set()
         return self._supported_speakers
 
+    @staticmethod
+    def _truthy_extra_body_flag(request: Any, key: str) -> bool:
+        if isinstance(request, dict):
+            extra_body = request
+            model_extra = {}
+        else:
+            extra_body = getattr(request, "extra_body", None) or {}
+            model_extra = getattr(request, "model_extra", None) or {}
+        value = extra_body.get(key, model_extra.get(key))
+        if isinstance(value, str):
+            return value.lower() in {"1", "true", "yes", "on"}
+        return bool(value)
+
+    @classmethod
+    def _filter_stage_metrics_detail(cls, metrics: dict[str, Any] | None, request: Any) -> dict[str, Any] | None:
+        if not metrics:
+            return metrics
+        if cls._truthy_extra_body_flag(request, "return_stage_metrics"):
+            return metrics
+        return None
+
     async def create_chat_completion(
         self,
         request: ChatCompletionRequest,
@@ -244,6 +265,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         if self._diffusion_mode:
             return await self._create_diffusion_chat_completion(request, raw_request)
 
+        request_timestamp = time.time()
+        if raw_request is not None:
+            request_timestamp = float(getattr(raw_request.state, "request_timestamp", request_timestamp))
+
         error_check_ret = await self._check_model(request)
         if error_check_ret is not None:
             logger.error("Error with model %s", error_check_ret)
@@ -555,6 +580,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     request_id=request_id,
                     sampling_params_list=sampling_params_list,
                     output_modalities=output_modalities,
+                    arrival_time=request_timestamp,
                 )
 
                 generators.append(generator)
@@ -1065,11 +1091,6 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 if omni_res.metrics:
                     last_metrics = omni_res.metrics
 
-                if res.prompt_token_ids is not None:
-                    num_prompt_tokens = len(res.prompt_token_ids)
-                    if res.encoder_prompt_token_ids is not None:
-                        num_prompt_tokens += len(res.encoder_prompt_token_ids)
-
                 # Initialize role before conditional blocks to avoid UnboundLocalError
                 # when handling audio/image responses
                 role = self.get_chat_request_role(request)
@@ -1081,6 +1102,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 # the result_generator, it needs to be sent as the FIRST
                 # response (by the try...catch).
                 if first_iteration_dict[final_output_type] and final_output_type == "text":
+                    # NOTE: prompt token IDs / cached tokens are based on the first iteration
+                    # of the stage that producing text output for consistency for current
+                    # non-stream behaviors.
+                    if res.prompt_token_ids is not None:
+                        num_prompt_tokens = len(res.prompt_token_ids)
+                        if res.encoder_prompt_token_ids is not None:
+                            num_prompt_tokens += len(res.encoder_prompt_token_ids)
+
                     num_cached_tokens = res.num_cached_tokens
                     # Send first response for each choice with role
                     # NOTE: num_choices defaults to 1 so this usually executes once per request
@@ -1621,7 +1650,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                             choices=[choice_data],
                             model=model_name,
                             modality=final_output_type,
-                            metrics=omni_res.metrics,
+                            metrics=self._filter_stage_metrics_detail(omni_res.metrics, request),
                         )
 
                         # handle usage stats if requested & if continuous
@@ -1702,6 +1731,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         choices=choices_data,
                         model=model_name,
                         modality=final_output_type,
+                        metrics=self._filter_stage_metrics_detail(omni_res.metrics, request),
                     )
                     chunk.usage = UsageInfo(
                         prompt_tokens=num_prompt_tokens,
@@ -1711,6 +1741,45 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     data = chunk.model_dump_json(exclude_unset=True)
                     yield f"data: {data}\n\n"
 
+                elif final_output_type == "image":
+                    role = self.get_chat_request_role(request)
+                    choices_data = []
+                    for choice in self._create_image_choice(omni_res, role, request, stream=True):
+                        delta = DeltaMessage.model_construct(role=role)
+                        object.__setattr__(delta, "content", choice.message.content)
+                        if hasattr(delta, "__pydantic_fields_set__"):
+                            delta.__pydantic_fields_set__.add("content")
+                        stream_choice = ChatCompletionResponseStreamChoice(
+                            index=choice.index,
+                            delta=delta,
+                            logprobs=None,
+                            finish_reason=choice.finish_reason,
+                            stop_reason=choice.stop_reason,
+                        )
+                        if stream_choice.finish_reason is not None:
+                            modality_finished[stream_choice.index].add("image")
+                        if modality_seen[stream_choice.index] < set(request.modalities) or not all(
+                            m in modality_finished[stream_choice.index] for m in modality_seen[stream_choice.index]
+                        ):
+                            stream_choice.finish_reason = None
+                        else:
+                            stop_reason_emitted[stream_choice.index] = True
+                        choices_data.append(stream_choice)
+                    chunk = OmniChatCompletionStreamResponse(
+                        id=request_id,
+                        object=chunk_object_type,
+                        created=created_time,
+                        choices=choices_data,
+                        model=model_name,
+                        modality=final_output_type,
+                        metrics=self._filter_stage_metrics_detail(omni_res.metrics, request),
+                    )
+                    # NOTE: Currently usage is only set the text stages to align with the behavior
+                    # of the full generator. TODO (Alex): Add support for usage on all stages for
+                    # both streaming and non-streaming.
+                    data = chunk.model_dump_json(exclude_unset=True)
+                    yield f"data: {data}\n\n"
+
                 else:
                     logger.warning(f"Unsupported streaming final output type: {final_output_type}")
                     continue
@@ -1771,7 +1840,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     choices=[],
                     model=model_name,
                     usage=final_usage,
-                    metrics=last_metrics,
+                    metrics=self._filter_stage_metrics_detail(last_metrics, request),
                 )
                 final_usage_data = final_usage_chunk.model_dump_json(exclude_unset=True, exclude_none=True)
                 yield f"data: {final_usage_data}\n\n"
@@ -1925,6 +1994,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     response_metrics.update(extra)
             choices.extend(choices_data)
 
+        response_metrics = self._filter_stage_metrics_detail(response_metrics, request)
+
         # Compute prompt_text for non-streaming response (upstream #42052)
         prompt_text = (
             getattr(final_res, "prompt", None)
@@ -2247,6 +2318,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         # omits "sr" and runs at 24kHz; Ming-flash-omni surfaces a 44.1kHz
         # AudioVAE rate via multimodal_output["sr"].
         sr_raw = mm_output.get("sr")
+        if isinstance(sr_raw, (list, tuple)):
+            sr_raw = next((item for item in sr_raw if item is not None), None)
         if sr_raw is None:
             sample_rate = 24000
         elif hasattr(sr_raw, "item"):
@@ -2526,6 +2599,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             mm_processor_kwargs["target_h"] = height
         if width is not None:
             mm_processor_kwargs["target_w"] = width
+        if seed is not None and engine_prompt_data is not None:
+            mm_processor_kwargs["vae_generator_seed"] = int(seed)
         if mm_processor_kwargs:
             engine_prompt["mm_processor_kwargs"] = mm_processor_kwargs
         if engine_prompt_data is not None:
@@ -2695,6 +2770,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         extra_body: dict[str, Any] | None = None,
         reference_images: list[str] | None = None,
         request_id: str | None = None,
+        arrival_time: float | None = None,
         stream: bool = False,
         model: str | None = None,
         output_format: str = "png",
@@ -2716,6 +2792,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         engine, gen_prompt, gen_params, pil_images = prepared
         if extra_body is None:
             extra_body = {}
+        return_stage_metrics = self._truthy_extra_body_flag(extra_body, "return_stage_metrics")
 
         if isinstance(engine, AsyncOmni):
             diffusion_engine = cast(AsyncOmni, engine)
@@ -2754,6 +2831,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                 prompt=engine_prompt,
                 sampling_params_list=sampling_params_list,
                 request_id=request_id,
+                arrival_time=arrival_time,
             )
             if stream:
                 return self._stream_diffusion_image_chunks(
@@ -2762,6 +2840,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                     output_format=output_format,
                     output_compression=output_compression,
                     size=size,
+                    return_stage_metrics=return_stage_metrics,
                     raw_request=raw_request,
                 )
 
@@ -2817,6 +2896,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
         output_format: str,
         output_compression: int,
         size: str,
+        return_stage_metrics: bool,
         raw_request: Request | None = None,
     ) -> AsyncIterator[str]:
         """Yield image edit SSE chunks from multi-stage diffusion outputs."""
@@ -2826,6 +2906,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
             async for output in result_generator:
                 final_output_type = getattr(output, "final_output_type", None)
                 stage_id = getattr(output, "stage_id", None)
+                metrics = getattr(output, "metrics", None) if return_stage_metrics else None
                 if final_output_type == "text" and stage_id == 0:
                     request_output = output.request_output
                     for completion in request_output.outputs:
@@ -2837,6 +2918,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                             index=completion.index,
                             created=created,
                             model=model,
+                            metrics=metrics,
                         )
                         yield f"data: {chunk.model_dump_json()}\n\n"
                 elif final_output_type == "image":
@@ -2860,6 +2942,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin):
                         size=size,
                         created=created,
                         model=model,
+                        metrics=metrics,
                     )
                     yield f"data: {chunk.model_dump_json()}\n\n"
                     emitted_image = True

vllm_omni/entrypoints/openai/serving_video.py

@@ -16,6 +16,7 @@ from vllm.logger import init_logger
 
 from vllm_omni.entrypoints.async_omni import AsyncOmni
 from vllm_omni.entrypoints.openai.protocol.videos import (
+    VideoAction,
     VideoData,
     VideoGenerationRequest,
     VideoGenerationResponse,
@@ -44,6 +45,7 @@ class VideoGenerationArtifacts:
 
     videos: list[Any]
     audios: list[Any | None]
+    actions: list[VideoAction | None]
     audio_sample_rate: int
     output_fps: int
     stage_durations: dict[str, float]
@@ -162,7 +164,20 @@ class OmniOpenAIServingVideo:
                 )
             # Merge extra_params into extra_args
             gen_params.extra_args.update(request.extra_params)
-            logger.info("Applied extra_params: %s", request.extra_params)
+
+            # Redact the inline ``action`` array (hundreds of floats) when
+            # logging so it doesn't flood the logs; everything else is logged
+            # verbatim.
+            loggable = request.extra_params
+            action_val = loggable.get("action")
+            if action_val is not None:
+                summary = (
+                    f"<{type(action_val).__name__} len={len(action_val)}>"
+                    if hasattr(action_val, "__len__")
+                    else f"<{type(action_val).__name__}>"
+                )
+                loggable = {**loggable, "action": summary}
+            logger.info("Applied extra_params: %s", loggable)
 
         self._apply_lora(request.lora, gen_params)
 
@@ -177,11 +192,13 @@ class OmniOpenAIServingVideo:
         result = await self._run_generation(prompt, gen_params, reference_id)
         videos = self._extract_video_outputs(result)
         audios = self._extract_audio_outputs(result, expected_count=len(videos))
+        actions = self._extract_action_outputs(result, expected_count=len(videos))
         audio_sample_rate = self._resolve_audio_sample_rate(result)
         output_fps = (vp.fps or self._resolve_fps(result) or 24) * self._resolve_video_fps_multiplier(result)
         return VideoGenerationArtifacts(
             videos=videos,
             audios=audios,
+            actions=actions,
             audio_sample_rate=audio_sample_rate,
             output_fps=output_fps,
             stage_durations=self._extract_stage_durations(result),
@@ -215,7 +232,8 @@ class OmniOpenAIServingVideo:
                         audio_sample_rate=artifacts.audio_sample_rate,
                         video_codec_options=video_codec_options,
                     )
-                )
+                ),
+                action=artifacts.actions[idx],
             )
             for idx, video in enumerate(artifacts.videos)
         ]
@@ -234,7 +252,7 @@ class OmniOpenAIServingVideo:
         reference_id: str,
         *,
         reference_image: ReferenceImage | None = None,
-    ) -> tuple[bytes, dict[str, float], float]:
+    ) -> tuple[bytes, dict[str, float], float, VideoAction | None]:
         """Generate a video and return raw MP4 bytes, bypassing base64 encoding."""
         artifacts = await self._run_and_extract(request, reference_id, reference_image=reference_image)
         if len(artifacts.videos) > 1:
@@ -259,7 +277,7 @@ class OmniOpenAIServingVideo:
         )
         _t_encode_ms = (time.perf_counter() - _t_encode_start) * 1000
         logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms)
-        return video_bytes, artifacts.stage_durations, artifacts.peak_memory_mb
+        return video_bytes, artifacts.stage_durations, artifacts.peak_memory_mb, artifacts.actions[0]
 
     @staticmethod
     def _resolve_video_fps_multiplier(result: Any) -> int:
@@ -440,6 +458,18 @@ class OmniOpenAIServingVideo:
 
         return 24000
 
+    @classmethod
+    def _extract_action_outputs(cls, result: Any, expected_count: int) -> list[VideoAction | None]:
+        custom_output = cls._extract_custom_output(result)
+        if not custom_output or "action" not in custom_output:
+            return [None] * expected_count
+
+        action_items = cls._split_action_payload(custom_output["action"], expected_count)
+        return [
+            cls._make_video_action(action_item, custom_output) if action_item is not None else None
+            for action_item in action_items
+        ]
+
     @staticmethod
     def _extract_custom_output(result: Any) -> dict[str, Any]:
         custom_output = getattr(result, "custom_output", None)
@@ -458,6 +488,89 @@ class OmniOpenAIServingVideo:
 
         return custom_output if isinstance(custom_output, dict) else {}
 
+    @classmethod
+    def _split_action_payload(cls, action: Any, expected_count: int) -> list[Any | None]:
+        if expected_count <= 0:
+            return []
+
+        shape = cls._shape_of(action)
+        if len(shape) >= 3:
+            count = min(shape[0], expected_count)
+            actions = [cls._index_action(action, i) for i in range(count)]
+            actions.extend([None] * (expected_count - count))
+            return actions
+
+        return [action] + [None] * (expected_count - 1)
+
+    @classmethod
+    def _make_video_action(cls, action: Any, custom_output: dict[str, Any]) -> VideoAction:
+        data = cls._to_jsonable(action)
+        if not isinstance(data, list):
+            data = [data]
+
+        action_mode = custom_output.get("action_mode")
+        return VideoAction(
+            data=data,
+            shape=cls._shape_of(action),
+            dtype=cls._dtype_of(action),
+            raw_action_dim=cls._coerce_optional_int(custom_output.get("raw_action_dim")),
+            action_mode=str(action_mode) if action_mode is not None else None,
+            domain_id=cls._coerce_optional_int(custom_output.get("domain_id")),
+        )
+
+    @staticmethod
+    def _index_action(action: Any, index: int) -> Any:
+        try:
+            return action[index]
+        except (IndexError, KeyError, TypeError):
+            return None
+
+    @classmethod
+    def _to_jsonable(cls, value: Any) -> Any:
+        if hasattr(value, "detach"):
+            value = value.detach()
+        if hasattr(value, "cpu"):
+            value = value.cpu()
+        if hasattr(value, "tolist"):
+            return cls._to_jsonable(value.tolist())
+        if isinstance(value, (list, tuple)):
+            return [cls._to_jsonable(item) for item in value]
+        if hasattr(value, "item"):
+            try:
+                return value.item()
+            except (TypeError, ValueError):
+                pass
+        return value
+
+    @classmethod
+    def _shape_of(cls, value: Any) -> list[int]:
+        shape = getattr(value, "shape", None)
+        if shape is not None:
+            try:
+                return [int(dim) for dim in shape]
+            except (TypeError, ValueError):
+                pass
+        if isinstance(value, (list, tuple)):
+            if not value:
+                return [0]
+            return [len(value)] + cls._shape_of(value[0])
+        return []
+
+    @staticmethod
+    def _dtype_of(value: Any) -> str | None:
+        dtype = getattr(value, "dtype", None)
+        return str(dtype) if dtype is not None else None
+
+    @staticmethod
+    def _coerce_optional_int(value: Any) -> int | None:
+        if value is None:
+            return None
+        try:
+            value = value.item() if hasattr(value, "item") else value
+            return int(value)
+        except (TypeError, ValueError):
+            return None
+
     @staticmethod
     def _resolve_fps(result: Any) -> int | None:
         """Extract fps from multimodal_output if the model reported it."""

vllm_omni/entrypoints/openai/video_api_utils.py

@@ -94,7 +94,9 @@ def _normalize_video_tensor(video_tensor: torch.Tensor) -> np.ndarray:
         video_tensor = video_tensor.permute(1, 2, 3, 0)
 
     if video_tensor.is_floating_point():
-        video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
+        # Cast to float32 first: bf16 (e.g. SANA-WM's refiner output) has no
+        # numpy dtype, so ``.numpy()`` below raises on it.
+        video_tensor = video_tensor.float().clamp(-1, 1) * 0.5 + 0.5
     else:
         video_tensor = video_tensor.to(torch.float32) / 255.0
     video_array = video_tensor.numpy()

vllm_omni/entrypoints/stage_utils.py

@@ -95,24 +95,6 @@ def _map_device_list(stage_id: int, device_list: list[str], visible_device_list:
     if not all(device.isdigit() for device in device_list):
         raise ValueError("Logical devices must be non-negative integers")
 
-    # Idempotency: if every requested id already names an entry in the visible
-    # pool (by value, not by index), treat the YAML field as physical and skip
-    # the remap. This avoids a double-mapping crash when a parent harness has
-    # already narrowed CUDA_VISIBLE_DEVICES to specific physical ids before
-    # spawning the subprocess (the test runtime's OmniServerStageCli does this
-    # — it sets each stage subprocess's env via ``_set_stage_device_env``).
-    # Without this guard, a YAML ``devices: "1"`` against a narrowed visible
-    # list ``["1"]`` would try ``visible[1]`` (out of range) and raise.
-    visible_set = set(visible_device_list)
-    if all(device in visible_set for device in device_list):
-        logger.info(
-            "Stage %s requested devices %s are already in the visible set %s; using as-is (no remap).",
-            stage_id,
-            device_list,
-            visible_device_list,
-        )
-        return list(device_list)
-
     logical_ids = [int(device) for device in device_list]
 
     # If ALL logical device IDs are >= num_visible, none can be 0-based indices

vllm_omni/metrics/definitions.py

@@ -11,6 +11,8 @@ time-bearing metrics use the ``_s`` suffix (values in seconds), counters use
 ``_total`` (auto-suffixed by the prometheus client), sizes use ``_bytes``.
 """
 
+from typing import Any
+
 # vllm:omni_ avoids upstream's unregister_vllm_metrics() stripping, which
 # removes every collector whose ``_name`` does not start with ``vllm``.
 METRIC_PREFIX = "vllm:omni_"
@@ -27,6 +29,78 @@ AUDIO_UNDERRUN = "audio_underrun"
 AUDIO_CONTINUITY_OK = "audio_continuity_ok"
 AUDIO_SKIPPED_REQUESTS = "audio_skipped_requests"
 
+# Bench-side aggregate field names. Keep these centralized so new benchmark
+# metrics reuse the same vocabulary instead of inventing parallel spellings.
+MEAN_AUDIO_TTFP_MS = f"mean_{AUDIO_TTFP}_ms"
+MEDIAN_AUDIO_TTFP_MS = f"median_{AUDIO_TTFP}_ms"
+STD_AUDIO_TTFP_MS = f"std_{AUDIO_TTFP}_ms"
+PERCENTILES_AUDIO_TTFP_MS = f"percentiles_{AUDIO_TTFP}_ms"
+TOTAL_AUDIO_DURATION_S = f"total_{AUDIO_DURATION}_s"
+TOTAL_AUDIO_FRAMES = f"total_{AUDIO_FRAMES}"
+AUDIO_THROUGHPUT = "audio_throughput"
+MEAN_AUDIO_RTF = f"mean_{AUDIO_RTF}"
+MEDIAN_AUDIO_RTF = f"median_{AUDIO_RTF}"
+STD_AUDIO_RTF = f"std_{AUDIO_RTF}"
+PERCENTILES_AUDIO_RTF = f"percentiles_{AUDIO_RTF}"
+MEAN_AUDIO_DURATION_S = f"mean_{AUDIO_DURATION}_s"
+MEDIAN_AUDIO_DURATION_S = f"median_{AUDIO_DURATION}_s"
+STD_AUDIO_DURATION_S = f"std_{AUDIO_DURATION}_s"
+PERCENTILES_AUDIO_DURATION_S = f"percentiles_{AUDIO_DURATION}_s"
+MEAN_AUDIO_UNDERRUN_S = f"mean_{AUDIO_UNDERRUN}_s"
+MEDIAN_AUDIO_UNDERRUN_S = f"median_{AUDIO_UNDERRUN}_s"
+STD_AUDIO_UNDERRUN_S = f"std_{AUDIO_UNDERRUN}_s"
+PERCENTILES_AUDIO_UNDERRUN_S = f"percentiles_{AUDIO_UNDERRUN}_s"
+AUDIO_CONTINUITY_OK_RATE = f"{AUDIO_CONTINUITY_OK}_rate"
+
+IMAGE_COUNT = "image_count"
+IMAGE_GENERATION = "image_generation"
+IMAGE_GENERATION_TIME_MS = f"{IMAGE_GENERATION}_time_ms"
+IMAGE_PIXELS = "image_pixels"
+TOTAL_IMAGES = "total_images"
+IMAGE_THROUGHPUT = "image_throughput"
+AVERAGE_PIXELS_PER_IMAGE = "average_pixels_per_image"
+DENOISE_STEP_LATENCY = "denoise_step_latency"
+DENOISE_STEP_LATENCY_MS = f"{DENOISE_STEP_LATENCY}_ms"
+MEAN_DENOISE_STEP_LATENCY_MS = f"mean_{DENOISE_STEP_LATENCY}_ms"
+MEAN_IMAGE_GENERATION_MS = f"mean_{IMAGE_GENERATION}_ms"
+MEDIAN_IMAGE_GENERATION_MS = f"median_{IMAGE_GENERATION}_ms"
+STD_IMAGE_GENERATION_MS = f"std_{IMAGE_GENERATION}_ms"
+PERCENTILES_IMAGE_GENERATION_MS = f"percentiles_{IMAGE_GENERATION}_ms"
+
+# Stage snapshot / StageBenchmarkMetrics field names.
+TOTAL_OUTPUT = "total_output"
+TTFTS = "ttfts"
+TPOTS = "tpots"
+ITLS = "itls"
+VLLM_TTFTS = "vllm_ttfts"
+VLLM_TPOTS = "vllm_tpots"
+VLLM_ITLS = "vllm_itls"
+AUDIO_TTFPS = "audio_ttfps"
+AUDIO_RTFS = "audio_rtfs"
+AUDIO_DURATIONS = "audio_durations"
+MISSING_AUDIO_DURATION_COUNT = "missing_audio_duration_count"
+MISSING_AUDIO_RTF_COUNT = "missing_audio_rtf_count"
+STAGE_GEN_TIME = "stage_gen_time"
+STAGE_GEN_TIME_MS = f"{STAGE_GEN_TIME}_ms"
+STAGE_GEN_TIMES_MS = f"{STAGE_GEN_TIME}s_ms"
+POSTPROCESS_TIME = "postprocess_time"
+POSTPROCESS_TIME_MS = f"{POSTPROCESS_TIME}_ms"
+POSTPROCESS_TIMES_MS = f"{POSTPROCESS_TIME}s_ms"
+OUTPUT_UNIT_COUNT = "output_unit_count"
+SERVING_TIME_TO_FIRST_OUTPUT_MS = "serving_time_to_first_output_ms"
+SERVING_TIME_TO_FIRST_OUTPUTS_MS = "serving_time_to_first_outputs_ms"
+TIME_PER_OUTPUT_UNIT_MS = "time_per_output_unit_ms"
+TIME_PER_OUTPUT_UNITS_MS = "time_per_output_units_ms"
+INTER_OUTPUT_LATENCY_MS = "inter_output_latency_ms"
+INTER_OUTPUT_LATENCIES_MS = "inter_output_latencies_ms"
+VLLM_TTFT_MS = "vllm_ttft_ms"
+VLLM_TPOT_MS = "vllm_tpot_ms"
+VLLM_ITL_MS = "vllm_itl_ms"
+VLLM_ITLS_MS = "vllm_itls_ms"
+NUM_TOKENS_IN = "num_tokens_in"
+NUM_TOKENS_OUT = "num_tokens_out"
+AUDIO_SAMPLE_RATE = "audio_sample_rate"
+
 
 # ============================================================================
 # Pipeline-level metric families (request counts + e2e latency)
@@ -186,6 +260,16 @@ def compute_audio_rtf(stage_gen_time_s: float, audio_duration_s: float) -> float
     return stage_gen_time_s / audio_duration_s
 
 
+def compute_denoise_step_latency(stage_gen_time: float, num_inference_steps: int) -> float:
+    """Mean denoise step latency = image stage generation time / step count.
+
+    The returned value uses the same time unit as ``stage_gen_time``.
+    """
+    if num_inference_steps <= 0:
+        return 0.0
+    return stage_gen_time / float(num_inference_steps)
+
+
 # ============================================================================
 # Audio sample-rate resolution
 # ============================================================================
@@ -195,21 +279,22 @@ def compute_audio_rtf(stage_gen_time_s: float, audio_duration_s: float) -> float
 # at runtime so this default only kicks in when the field is missing.
 DEFAULT_AUDIO_SAMPLE_RATE = 24000
 
-_SAMPLE_RATE_KEYS = ("audio_sample_rate", "sample_rate", "sampling_rate", "sr")
+_SAMPLE_RATE_KEYS = ("output_sample_rate", "audio_sample_rate", "sample_rate", "sampling_rate", "sr")
 
 
-def resolve_audio_sample_rate(multimodal_output: dict | None) -> int:
-    """Extract audio sample_rate from a multimodal_output dict, with fallbacks.
+def resolve_audio_sample_rate(source: dict[str, Any] | Any | None) -> int:
+    """Extract audio sample_rate from a dict or config object, with fallbacks.
 
     Tries the same key chain as serving_chat.py's audio response path so
     /metrics audio_duration_s = audio_frames / sample_rate stays consistent
-    with what the OpenAI streaming endpoint reports back to clients.
+    with what the OpenAI streaming endpoint reports back to clients. Also
+    accepts config objects that expose the same values as attributes.
     Returns DEFAULT_AUDIO_SAMPLE_RATE when no usable value is present.
     """
-    if not multimodal_output:
+    if not source:
         return DEFAULT_AUDIO_SAMPLE_RATE
     for key in _SAMPLE_RATE_KEYS:
-        raw = multimodal_output.get(key)
+        raw = source.get(key) if isinstance(source, dict) else getattr(source, key, None)
         if raw is None:
             continue
         try:

vllm_omni/metrics/stats.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any
 
 from vllm.logger import init_logger
 
+from vllm_omni.metrics import definitions as defs
 from vllm_omni.metrics.utils import _build_field_defs, _build_row, _format_table
 
 if TYPE_CHECKING:
@@ -45,7 +46,22 @@ class StageRequestStats:
     postprocess_time_ms: float = 0.0
     diffusion_metrics: dict[str, int] = None
     audio_generated_frames: int = 0
+    audio_sample_rate: int = 0
+    audio_duration_s: float = 0.0
+    audio_rtf: float = 0.0
+    image_pixels: int = 0
+    denoise_step_latency_ms: float = 0.0
     pipeline_timings: dict[str, float] | None = None
+    output_unit_type: str | None = None
+    output_unit_count: int = 0
+    serving_time_to_first_output_ms: float = 0.0
+    time_per_output_unit_ms: float = 0.0
+    inter_output_latency_ms: float = 0.0
+    inter_output_latencies_ms: list[float] | None = None
+    vllm_ttft_ms: float = 0.0
+    vllm_tpot_ms: float = 0.0
+    vllm_itl_ms: float = 0.0
+    vllm_itls_ms: list[float] | None = None
 
     @property
     def rx_mbps(self) -> float:
@@ -332,14 +348,15 @@ class OrchestratorAggregator:
                 and len(multimodal_output) > 0
             ):
                 nframes = sum(
-                    int(t.shape[0]) if t.ndim > 0 else 1
+                    int(t.shape[-1]) if t.ndim > 0 else 1
                     for t in (multimodal_output if isinstance(multimodal_output, list) else [multimodal_output])
                 )
                 stage_events_for_req = self.stage_events.get(request_id, [])
                 if stage_events_for_req:
                     for stage_event in stage_events_for_req:
                         if stage_event.stage_id == stage_id:
-                            stage_event.audio_generated_frames += nframes
+                            if stage_event.audio_generated_frames <= 0:
+                                stage_event.audio_generated_frames = nframes
                             break
                 else:
                     logger.warning(
@@ -392,24 +409,40 @@ class OrchestratorAggregator:
             if output_to_yield is None:
                 return
 
-            # 3. Not finished yet — empty metrics, skip audio recording
+            rid_key = str(req_id)
+            stage_snapshot = self._build_stage_metrics_snapshot(rid_key)
+
+            # 3. Not finished yet — expose incremental per-stage snapshot for streaming clients.
             if not finished:
-                output_to_yield.metrics = {}
+                if stage_snapshot:
+                    output_to_yield.metrics = {"stage_metrics": stage_snapshot}
+                else:
+                    output_to_yield.metrics = {}
                 return
 
-            # 4. Finished with output: assign text metrics if available
-            output_to_yield.metrics = {}
+            # 4. Finished with output: always attach per-stage snapshot; keep legacy top-level
+            # token fields only for text stages (OpenAI-style completion token accounting).
+            output_to_yield.metrics = {"stage_metrics": stage_snapshot}
             stage_event = next(
-                (evt for evt in reversed(self.stage_events.get(req_id, [])) if evt.stage_id == stage_id),
+                (evt for evt in reversed(self.stage_events.get(rid_key, [])) if evt.stage_id == stage_id),
                 None,
             )
             if stage_event is not None and stage_event.final_output_type == "text":
-                output_to_yield.metrics = {
-                    "num_tokens_in": stage_event.num_tokens_in,
-                    "num_tokens_out": stage_event.num_tokens_out,
-                    "stage_id": stage_event.stage_id,
-                    "final_output_type": stage_event.final_output_type,
-                }
+                output_to_yield.metrics.update(
+                    {
+                        defs.NUM_TOKENS_IN: stage_event.num_tokens_in,
+                        defs.NUM_TOKENS_OUT: stage_event.num_tokens_out,
+                        "stage_id": stage_event.stage_id,
+                        "final_output_type": stage_event.final_output_type,
+                    }
+                )
+            elif stage_event is not None:
+                output_to_yield.metrics.update(
+                    {
+                        "stage_id": stage_event.stage_id,
+                        "final_output_type": stage_event.final_output_type,
+                    }
+                )
 
             # 5. Finished: record audio generated frames
             self.record_audio_generated_frames(output_to_yield, stage_id, req_id)
@@ -421,6 +454,110 @@ class OrchestratorAggregator:
                 req_id,
             )
 
+    @staticmethod
+    def _merge_stage_metric_event(
+        current: dict[str, Any] | None,
+        evt: StageRequestStats,
+    ) -> dict[str, Any]:
+        sid = int(evt.stage_id) if evt.stage_id is not None else -1
+        if current is None:
+            current = {
+                "stage_id": sid,
+                "final_output_type": evt.final_output_type,
+                defs.NUM_TOKENS_IN: int(evt.num_tokens_in),
+                defs.NUM_TOKENS_OUT: int(evt.num_tokens_out),
+                defs.STAGE_GEN_TIME_MS: float(evt.stage_gen_time_ms),
+                defs.POSTPROCESS_TIME_MS: float(evt.postprocess_time_ms),
+                defs.AUDIO_FRAMES: int(evt.audio_generated_frames),
+                defs.AUDIO_SAMPLE_RATE: int(evt.audio_sample_rate),
+                f"{defs.AUDIO_DURATION}_s": float(evt.audio_duration_s),
+                defs.AUDIO_RTF: float(evt.audio_rtf),
+                defs.IMAGE_PIXELS: int(evt.image_pixels),
+                defs.DENOISE_STEP_LATENCY_MS: float(evt.denoise_step_latency_ms),
+                "output_unit_type": evt.output_unit_type,
+                defs.OUTPUT_UNIT_COUNT: int(evt.output_unit_count),
+                defs.SERVING_TIME_TO_FIRST_OUTPUT_MS: float(evt.serving_time_to_first_output_ms),
+                defs.TIME_PER_OUTPUT_UNIT_MS: float(evt.time_per_output_unit_ms),
+                defs.INTER_OUTPUT_LATENCY_MS: float(evt.inter_output_latency_ms),
+                defs.INTER_OUTPUT_LATENCIES_MS: list(evt.inter_output_latencies_ms or []),
+                defs.VLLM_TTFT_MS: float(evt.vllm_ttft_ms),
+                defs.VLLM_TPOT_MS: float(evt.vllm_tpot_ms),
+                defs.VLLM_ITL_MS: float(evt.vllm_itl_ms),
+                defs.VLLM_ITLS_MS: list(evt.vllm_itls_ms or []),
+            }
+            return current
+
+        current[defs.NUM_TOKENS_IN] = int(current.get(defs.NUM_TOKENS_IN, 0)) + int(evt.num_tokens_in)
+        current[defs.NUM_TOKENS_OUT] = int(current.get(defs.NUM_TOKENS_OUT, 0)) + int(evt.num_tokens_out)
+        current[defs.STAGE_GEN_TIME_MS] = float(current.get(defs.STAGE_GEN_TIME_MS, 0.0)) + float(evt.stage_gen_time_ms)
+        current[defs.POSTPROCESS_TIME_MS] = float(current.get(defs.POSTPROCESS_TIME_MS, 0.0)) + float(
+            evt.postprocess_time_ms
+        )
+        current[defs.AUDIO_FRAMES] = int(current.get(defs.AUDIO_FRAMES, 0)) + int(evt.audio_generated_frames)
+        if int(current.get(defs.AUDIO_SAMPLE_RATE, 0)) <= 0 and int(evt.audio_sample_rate) > 0:
+            current[defs.AUDIO_SAMPLE_RATE] = int(evt.audio_sample_rate)
+        current[f"{defs.AUDIO_DURATION}_s"] = float(current.get(f"{defs.AUDIO_DURATION}_s", 0.0)) + float(
+            evt.audio_duration_s
+        )
+        current[defs.IMAGE_PIXELS] = int(current.get(defs.IMAGE_PIXELS, 0)) + int(evt.image_pixels)
+        denoise_step_latency_ms = float(evt.denoise_step_latency_ms)
+        if denoise_step_latency_ms > 0:
+            current[defs.DENOISE_STEP_LATENCY_MS] = denoise_step_latency_ms
+        current[defs.OUTPUT_UNIT_COUNT] = int(current.get(defs.OUTPUT_UNIT_COUNT, 0)) + int(evt.output_unit_count)
+
+        first_output_ms = float(evt.serving_time_to_first_output_ms)
+        current_first_output_ms = float(current.get(defs.SERVING_TIME_TO_FIRST_OUTPUT_MS, 0.0))
+        if current_first_output_ms <= 0 < first_output_ms:
+            current[defs.SERVING_TIME_TO_FIRST_OUTPUT_MS] = first_output_ms
+
+        if evt.output_unit_type:
+            current["output_unit_type"] = evt.output_unit_type
+        if evt.final_output_type:
+            current["final_output_type"] = evt.final_output_type
+
+        inter_latencies = list(current.get(defs.INTER_OUTPUT_LATENCIES_MS) or [])
+        inter_latencies.extend(list(evt.inter_output_latencies_ms or []))
+        current[defs.INTER_OUTPUT_LATENCIES_MS] = inter_latencies
+        current[defs.INTER_OUTPUT_LATENCY_MS] = (
+            sum(inter_latencies) / float(len(inter_latencies)) if inter_latencies else 0.0
+        )
+
+        vllm_ttft_ms = float(evt.vllm_ttft_ms)
+        current_vllm_ttft_ms = float(current.get(defs.VLLM_TTFT_MS, 0.0))
+        if current_vllm_ttft_ms <= 0 < vllm_ttft_ms:
+            current[defs.VLLM_TTFT_MS] = vllm_ttft_ms
+        vllm_tpot_ms = float(evt.vllm_tpot_ms)
+        if vllm_tpot_ms > 0:
+            current[defs.VLLM_TPOT_MS] = vllm_tpot_ms
+        vllm_itls = list(current.get(defs.VLLM_ITLS_MS) or [])
+        vllm_itls.extend(list(evt.vllm_itls_ms or []))
+        current[defs.VLLM_ITLS_MS] = vllm_itls
+        current[defs.VLLM_ITL_MS] = sum(vllm_itls) / float(len(vllm_itls)) if vllm_itls else 0.0
+
+        output_count = int(current.get(defs.OUTPUT_UNIT_COUNT, 0))
+        remaining_ms = max(
+            float(current.get(defs.STAGE_GEN_TIME_MS, 0.0))
+            - float(current.get(defs.SERVING_TIME_TO_FIRST_OUTPUT_MS, 0.0)),
+            0.0,
+        )
+        current[defs.TIME_PER_OUTPUT_UNIT_MS] = remaining_ms / float(output_count - 1) if output_count > 1 else 0.0
+        duration_s = float(current.get(f"{defs.AUDIO_DURATION}_s", 0.0))
+        current[defs.AUDIO_RTF] = defs.compute_audio_rtf(
+            float(current.get(defs.STAGE_GEN_TIME_MS, 0.0)) / 1000.0, duration_s
+        )
+        return current
+
+    def _build_stage_metrics_snapshot(self, req_id: str) -> dict[str, dict[str, Any]]:
+        """Aggregate per-stage metrics for ``req_id`` (string key), for streaming/benchmark clients."""
+        snapshot: dict[str, dict[str, Any]] = {}
+        for evt in self.stage_events.get(req_id, []):
+            sid = int(evt.stage_id) if evt.stage_id is not None else -1
+            if sid < 0:
+                continue
+            sid_key = str(sid)
+            snapshot[sid_key] = self._merge_stage_metric_event(snapshot.get(sid_key), evt)
+        return snapshot
+
     def _as_stage_request_stats(
         self,
         stage_id: int,
@@ -432,7 +569,8 @@ class OrchestratorAggregator:
         stats = metrics
         stats.stage_id = stage_id
         stats.request_id = req_id
-        stats.final_output_type = final_output_type
+        if final_output_type is not None:
+            stats.final_output_type = final_output_type
         stats.diffusion_metrics = (
             {k: int(v) for k, v in self.diffusion_metrics.pop(req_id, {}).items()}
             if req_id in self.diffusion_metrics

vllm_omni/metrics/utils.py

@@ -94,7 +94,13 @@ def _format_table(
         if isinstance(value, float):
             return f"{value:,.3f}"
         if isinstance(value, list):
-            return ", ".join(str(f"{v:,.3f}") for v in value)
+            if not value:
+                return ""
+            try:
+                avg = sum(float(v) for v in value) / len(value)
+                return f"{avg:,.3f} (n={len(value)})"
+            except (TypeError, ValueError):
+                return ", ".join(str(v) for v in value)
         return str(value)
 
     table = PrettyTable()
@@ -117,6 +123,17 @@ def _format_table(
         if column_key is None:
             raise ValueError("column_key is required for multi-column mode")
         col_headers = [f"{column_prefix}{row.get(column_key, '?')}" for row in data]
+        # PrettyTable requires unique field names.  When the same column key
+        # appears more than once (e.g. two events with the same stage_id),
+        # deduplicate by appending _2, _3, … to the later occurrences.
+        if len(col_headers) != len(set(col_headers)):
+            seen: dict[str, int] = {}
+            deduped: list[str] = []
+            for h in col_headers:
+                n = seen.get(h, 0)
+                seen[h] = n + 1
+                deduped.append(h if n == 0 else f"{h}_{n + 1}")
+            col_headers = deduped
         table.field_names = ["Field"] + col_headers
         table.align["Field"] = "l"
         for col in col_headers:
@@ -134,9 +151,15 @@ def count_tokens_from_outputs(engine_outputs: list[Any]) -> int:
         try:
             outs = getattr(_ro, "outputs", None)
             if outs and len(outs) > 0:
-                tokens = getattr(outs[0], "token_ids", None)
-                if tokens is not None:
-                    total += len(tokens)
+                for output in outs:
+                    # In DELTA mode token_ids only contains the latest chunk.
+                    # Omni's output processor attaches the cumulative sequence
+                    # for inter-stage routing and accurate metrics.
+                    tokens = getattr(output, "cumulative_token_ids", None)
+                    if tokens is None:
+                        tokens = getattr(output, "token_ids", None)
+                    if tokens is not None:
+                        total += len(tokens)
         except Exception:
             # Ignore any issues with individual outputs to keep token counting best-effort.
             pass

vllm_omni/model_executor/model_loader/weight_utils.py

@@ -1,8 +1,4 @@
-import contextlib
-import errno
-import os
 import time
-from collections.abc import Iterator
 from pathlib import Path
 
 import huggingface_hub
@@ -17,159 +13,6 @@ else:
 
 logger = init_logger(__name__)
 
-_DOWNLOAD_MAX_ATTEMPTS = 3
-_DOWNLOAD_BACKOFF_BASE_S = 1.0
-_FULL_SNAPSHOT_METADATA_MARKERS = (
-    "config.json",
-    "preprocessor_config.json",
-    "processor_config.json",
-    "tokenizer_config.json",
-    "model_index.json",
-)
-
-
-def _node_lock_dir() -> str:
-    candidates: list[str] = []
-    hf_home = os.environ.get("HF_HOME")
-    if hf_home:
-        candidates.append(os.path.join(hf_home, "locks", "vllm-omni-weight-prefetch"))
-    xdg_cache = os.environ.get("XDG_CACHE_HOME")
-    if xdg_cache:
-        candidates.append(os.path.join(xdg_cache, "huggingface", "locks", "vllm-omni-weight-prefetch"))
-    candidates.append(
-        os.path.join(os.path.expanduser("~"), ".cache", "huggingface", "locks", "vllm-omni-weight-prefetch")
-    )
-    candidates.append(os.path.join("/tmp", "vllm-omni-weight-prefetch-locks"))
-
-    for cand in candidates:
-        try:
-            os.makedirs(cand, exist_ok=True)
-            probe = os.path.join(cand, ".write_check")
-            with open(probe, "a"):
-                pass
-            with contextlib.suppress(OSError):
-                os.remove(probe)
-            return cand
-        except OSError:
-            continue
-    fallback = os.path.join("/tmp", "vllm-omni-weight-prefetch-locks")
-    os.makedirs(fallback, exist_ok=True)
-    return fallback
-
-
-def _safe_repo_filename(model: str) -> str:
-    return model.replace("/", "__").replace(os.sep, "__") + ".lock"
-
-
-def _dotfile_lock_acquire(lock_dir: str, model: str, timeout: float = 300.0, poll_interval: float = 0.5) -> str | None:
-    lock_path = os.path.join(lock_dir, _safe_repo_filename(model) + ".dir")
-    deadline = time.monotonic() + timeout
-    while True:
-        try:
-            os.makedirs(lock_path, exist_ok=False)
-            logger.info("Acquired dotfile weight prefetch lock for %s at %s", model, lock_path)
-            return lock_path
-        except FileExistsError:
-            if time.monotonic() >= deadline:
-                logger.warning(
-                    "Timed out waiting for dotfile weight prefetch lock %s after %.0fs; proceeding unlocked",
-                    lock_path,
-                    timeout,
-                )
-                return None
-            time.sleep(poll_interval)
-
-
-@contextlib.contextmanager
-def _repo_download_lock(model: str) -> Iterator[None]:
-    """Serialize full-repo snapshot materialization across Omni processes."""
-    lock_dir = None
-    dotfile_held = None
-    fd = None
-    flock_held = False
-
-    try:
-        import fcntl  # type: ignore[import-not-found]
-    except ImportError:  # pragma: no cover - non-POSIX
-        fcntl = None
-
-    if fcntl is not None:
-        try:
-            lock_dir = _node_lock_dir()
-            lock_path = os.path.join(lock_dir, _safe_repo_filename(model))
-            fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644)
-            fcntl.flock(fd, fcntl.LOCK_EX)
-            flock_held = True
-            logger.info("Acquired flock weight prefetch lock for %s at %s", model, lock_path)
-        except OSError as exc:
-            if fd is not None:
-                with contextlib.suppress(OSError):
-                    os.close(fd)
-                fd = None
-            if exc.errno not in (errno.ENOLCK, errno.EOPNOTSUPP, errno.EACCES, errno.EINVAL):
-                raise
-            logger.warning("fcntl.flock unavailable for weight prefetch of %s (%s); using dotfile lock", model, exc)
-
-    if not flock_held:
-        try:
-            lock_dir = lock_dir or _node_lock_dir()
-            dotfile_held = _dotfile_lock_acquire(lock_dir, model)
-        except OSError as exc:
-            logger.warning("Could not allocate weight prefetch lock dir for %s (%s); proceeding unlocked", model, exc)
-
-    try:
-        yield
-    finally:
-        if flock_held and fd is not None:
-            with contextlib.suppress(OSError):
-                fcntl.flock(fd, fcntl.LOCK_UN)
-        if fd is not None:
-            with contextlib.suppress(OSError):
-                os.close(fd)
-        if dotfile_held is not None:
-            with contextlib.suppress(OSError):
-                os.rmdir(dotfile_held)
-
-
-def _looks_like_auth_error(exc: BaseException) -> bool:
-    try:
-        from huggingface_hub.errors import GatedRepoError, RepositoryNotFoundError  # type: ignore[import-not-found]
-
-        if isinstance(exc, GatedRepoError | RepositoryNotFoundError):
-            return True
-    except ImportError:  # pragma: no cover - very old huggingface_hub
-        pass
-
-    status = getattr(getattr(exc, "response", None), "status_code", None)
-    if status in (401, 403):
-        return True
-    msg = str(exc).lower()
-    return "401 client error" in msg or "403 client error" in msg or "gatedrepo" in msg
-
-
-def _is_full_snapshot_request(allow_patterns: list[str]) -> bool:
-    return any(pattern.strip() in {"*", "**", "./*", "**/*"} for pattern in allow_patterns)
-
-
-def _verify_snapshot_materialized(hf_folder: str, allow_patterns: list[str]) -> None:
-    root = Path(hf_folder)
-    if not root.exists():
-        raise OSError(f"Downloaded snapshot folder does not exist: {hf_folder}")
-
-    if not any(root.glob(pattern) for pattern in allow_patterns):
-        raise OSError(f"Downloaded snapshot {hf_folder} has no files matching {allow_patterns}")
-
-    # Full-repo downloads are used before transformers loads tokenizers,
-    # processors, and feature extractors. A half-materialized shared cache can
-    # contain weights while missing these small metadata files, which later
-    # surfaces as "Can't load feature extractor ... preprocessor_config.json".
-    if _is_full_snapshot_request(allow_patterns) and not any(
-        (root / marker).exists() for marker in _FULL_SNAPSHOT_METADATA_MARKERS
-    ):
-        raise OSError(
-            f"Downloaded full snapshot {hf_folder} is missing expected metadata files {_FULL_SNAPSHOT_METADATA_MARKERS}"
-        )
-
 
 def download_weights_from_hf_specific(
     model_name_or_path: str,
@@ -201,77 +44,44 @@ def download_weights_from_hf_specific(
         str: The path to the downloaded model weights.
     """
     assert len(allow_patterns) > 0
-    allow_patterns = list(allow_patterns)
     local_only = huggingface_hub.constants.HF_HUB_OFFLINE
     download_kwargs = {"tqdm_class": DisabledTqdm} if not envs.VLLM_USE_MODELSCOPE else {}
 
     logger.info("Using model weights format %s", allow_patterns)
-    start_time = time.perf_counter()
-    last_exc: BaseException | None = None
-
-    for attempt in range(1, _DOWNLOAD_MAX_ATTEMPTS + 1):
-        try:
-            # Use both vLLM's cache lock and a repo-wide node lock. The latter
-            # protects the transformers v5 eager metadata resolver from seeing
-            # a peer process' half-written shared HF cache.
-            with get_lock(model_name_or_path, cache_dir), _repo_download_lock(model_name_or_path):
-                if require_all:
-                    hf_folder = snapshot_download(
-                        model_name_or_path,
-                        allow_patterns=allow_patterns,
-                        ignore_patterns=ignore_patterns,
-                        cache_dir=cache_dir,
-                        revision=revision,
-                        local_files_only=local_only,
-                        **download_kwargs,
-                    )
-                    _verify_snapshot_materialized(hf_folder, allow_patterns)
-                    break
-
-                hf_folder = None
-                for allow_pattern in allow_patterns:
-                    hf_folder = snapshot_download(
-                        model_name_or_path,
-                        allow_patterns=allow_pattern,
-                        ignore_patterns=ignore_patterns,
-                        cache_dir=cache_dir,
-                        revision=revision,
-                        local_files_only=local_only,
-                        **download_kwargs,
-                    )
-                    _verify_snapshot_materialized(hf_folder, [allow_pattern])
-                    # If we have downloaded weights for this allow_pattern,
-                    # we don't need to check the rest, unless require_all is set.
-                    if any(Path(hf_folder).glob(allow_pattern)):
-                        break
-                if hf_folder is None:
-                    raise OSError(f"No snapshot downloaded for patterns {allow_patterns}")
-                break
-        except Exception as exc:
-            last_exc = exc
-            if _looks_like_auth_error(exc) or attempt == _DOWNLOAD_MAX_ATTEMPTS or local_only:
-                raise
-            backoff = _DOWNLOAD_BACKOFF_BASE_S * attempt
-            logger.warning(
-                "Downloading weights for %s with patterns %s failed on attempt %d/%d (%s: %s); retrying in %.1fs",
+    # Use file lock to prevent multiple processes from
+    # downloading the same model weights at the same time.
+    with get_lock(model_name_or_path, cache_dir):
+        start_time = time.perf_counter()
+        if require_all:
+            hf_folder = snapshot_download(
                 model_name_or_path,
-                allow_patterns,
-                attempt,
-                _DOWNLOAD_MAX_ATTEMPTS,
-                type(exc).__name__,
-                exc,
-                backoff,
+                allow_patterns=allow_patterns,
+                ignore_patterns=ignore_patterns,
+                cache_dir=cache_dir,
+                revision=revision,
+                local_files_only=local_only,
+                **download_kwargs,
+            )
+        else:
+            for allow_pattern in allow_patterns:
+                hf_folder = snapshot_download(
+                    model_name_or_path,
+                    allow_patterns=allow_pattern,
+                    ignore_patterns=ignore_patterns,
+                    cache_dir=cache_dir,
+                    revision=revision,
+                    local_files_only=local_only,
+                    **download_kwargs,
+                )
+                # If we have downloaded weights for this allow_pattern,
+                # we don't need to check the rest, unless require_all is set.
+                if any(Path(hf_folder).glob(allow_pattern)):
+                    break
+        time_taken = time.perf_counter() - start_time
+        if time_taken > 0.5:
+            logger.info(
+                "Time spent downloading weights for %s: %.6f seconds",
+                model_name_or_path,
+                time_taken,
             )
-            time.sleep(backoff)
-
-    if last_exc is not None and "hf_folder" not in locals():
-        raise last_exc
-
-    time_taken = time.perf_counter() - start_time
-    if time_taken > 0.5:
-        logger.info(
-            "Time spent downloading weights for %s: %.6f seconds",
-            model_name_or_path,
-            time_taken,
-        )
     return hf_folder

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -643,14 +643,70 @@ class CosyVoice3Model(
     ) -> torch.Tensor:
         if self.model_stage == "cosyvoice3_talker":
             if is_multimodal is not None and any(is_multimodal):
+                # Per-request rearrange for the talker prompt layout
+                # [SOS_ph, TASK_ID_ph, AUDIO_ph * pstoken_len, ...text_tokens...]
+                # → [SOS_emb, ...text_embs..., TASK_ID_emb, AUDIO_embs].
+                #
+                # In batched prefill ``input_ids`` is the flat concatenation of
+                # N sequences and ``multimodal_embeddings`` is a list with one
+                # tensor per multimodal request. Each request's audio
+                # placeholders form one contiguous ``True`` group in
+                # ``is_multimodal``; the 2 positions immediately before each
+                # group are the SOS / TASK_ID placeholders. Walking the groups
+                # lets us locate every request's segment without needing a
+                # ``seq_token_counts`` kwarg (which only the generation runner
+                # injects, not the AR runner used here).
                 embed_tokens = self.model.llm.model.embed_tokens(input_ids)
                 sos = self.model.speech_embedding.weight[self.model.sos].reshape(1, -1)
                 task_id = self.model.speech_embedding.weight[self.model.task_id].reshape(1, -1)
-                prompt_speech_token_emb = multimodal_embeddings[0]
-                pstoken_len = prompt_speech_token_emb.shape[0]  # Get length from tensor shape
-                embed_tokens = torch.cat(
-                    [sos, embed_tokens[2 + pstoken_len :], task_id, prompt_speech_token_emb], dim=0
-                )
+
+                is_mm = is_multimodal.to(torch.bool).reshape(-1)
+                # vLLM v1's multimodal placeholder range covers the FULL
+                # ``[SOS_ph, TASK_ID_ph, AUDIO_ph * pstoken_len]`` block
+                # (length ``2 + pstoken_len``, see ``_get_prompt_updates``'s
+                # ``insertion_end`` which inserts ``[1] * (1 + 1 + token_len)``
+                # and ``mm_position.length`` matches that). So each contiguous
+                # ``True`` group in ``is_mm`` starts at a request's SOS
+                # placeholder, NOT at its first audio placeholder.
+                prev_false = torch.cat([torch.ones(1, dtype=torch.bool, device=is_mm.device), ~is_mm[:-1]])
+                group_starts = (is_mm & prev_false).nonzero(as_tuple=True)[0].tolist()
+                if len(group_starts) != len(multimodal_embeddings):
+                    raise RuntimeError(
+                        f"cosyvoice3 talker: found {len(group_starts)} placeholder "
+                        f"blocks in is_multimodal but {len(multimodal_embeddings)} "
+                        f"multimodal_embeddings tensors — these must match 1:1."
+                    )
+
+                total_tokens = int(embed_tokens.shape[0])
+                # vLLM v1 packs cached (decode) requests first, then new
+                # (prefill) requests. The decode tokens at positions
+                # ``[0:group_starts[0]]`` get a speech_embedding lookup (the
+                # talker is generating codec tokens here, so input ids are
+                # codec ids, not text ids). Positions starting at
+                # ``group_starts[0]`` are the first new prefill request's
+                # SOS placeholder.
+                segments: list[torch.Tensor] = []
+                first_prefill = group_starts[0]
+                if first_prefill > 0:
+                    decode_ids = input_ids[:first_prefill].to(dtype=torch.long)
+                    segments.append(self.model.speech_embedding.weight[decode_ids])
+
+                for i, req_start in enumerate(group_starts):
+                    pstoken_len_i = int(multimodal_embeddings[i].shape[0])
+                    # Real text starts after the 2 + pstoken_len_i leading
+                    # placeholders of this request's prompt.
+                    text_start = req_start + 2 + pstoken_len_i
+                    # Next request's segment starts at its own SOS placeholder,
+                    # which is exactly the next True group's first index.
+                    if i + 1 < len(group_starts):
+                        req_end = group_starts[i + 1]
+                    else:
+                        req_end = total_tokens
+                    segments.append(sos)
+                    segments.append(embed_tokens[text_start:req_end])
+                    segments.append(task_id)
+                    segments.append(multimodal_embeddings[i])
+                embed_tokens = torch.cat(segments, dim=0)
             else:
                 embed_tokens = self.model.speech_embedding.weight[input_ids]
             return embed_tokens

vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py

@@ -23,7 +23,9 @@ from vllm.config.multimodal import BaseDummyOptions
 from vllm.distributed import (
     get_ep_group,
     get_pp_group,
+    get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_gather,
 )
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
@@ -1041,6 +1043,9 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc
         images = mm_data.get("images", [])
         logger.debug(f"process image count: {len(images)}")
         batch_feature = image_processor(prompt, images, **tok_kwargs)
+        vae_generator_seed = mm_kwargs.get("vae_generator_seed")
+        if vae_generator_seed is not None and images:
+            batch_feature["vae_generator_seed"] = torch.full((len(images),), int(vae_generator_seed), dtype=torch.long)
         return batch_feature
 
     def _hf_processor_applies_updates(
@@ -1078,6 +1083,8 @@ class HunyuanImage3MultiModalProcessor(BaseMultiModalProcessor[HunyuanImage3Proc
             config["base_size"] = MultiModalFieldConfig.batched("image")
         if "ratio_index" in hf_inputs:
             config["ratio_index"] = MultiModalFieldConfig.batched("image")
+        if "vae_generator_seed" in hf_inputs:
+            config["vae_generator_seed"] = MultiModalFieldConfig.batched("image")
         return config
 
     def _get_prompt_updates(
@@ -1451,6 +1458,12 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
 
     prefer_model_sampler = True
 
+    # Siglip2 ViT supports data-parallel encoding (mm_encoder_tp_mode="data"):
+    # weights are replicated and the image batch is sharded across TP ranks.
+    # Without this flag, vLLM silently falls back to "weights" (TP). See
+    # _vit_encode_dp.
+    supports_encoder_tp_data = True
+
     packed_modules_mapping = {
         "qkv_proj": [
             "q_proj",
@@ -1519,7 +1532,14 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         self.time_embed = TimestepEmbedder(hidden_size=config.hidden_size)
 
         # vision
-        self.vision_model = Siglip2VisionTransformer(config.vit)
+        multimodal_config = vllm_config.model_config.multimodal_config
+        self.multimodal_config = multimodal_config
+        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
+        self.vision_model = Siglip2VisionTransformer(
+            config.vit,
+            quant_config=quant_config,
+            prefix="vision_model",
+        )
         self.vision_aligner = LightProjector(config.vit_aligner)
 
         # Used to embed timestep information into the input sequence.
@@ -1717,6 +1737,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         # we reconstruct per-image shapes from vae_token_grid_hw below.
         kwargs.pop("vae_pixel_size", None)
         vae_token_grid_hw = kwargs.pop("vae_token_grid_hw", None)
+        vae_generator_seed = kwargs.pop("vae_generator_seed", None)
 
         if vit_pixel_values is None or vae_pixel_values is None:
             return None
@@ -1756,6 +1777,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
                 "vit_spatial_shapes": vit_spatial_shapes,
                 "vae_pixel_values": vae_image_list,
                 "vae_token_grid_hw": vae_token_grid_hw,
+                "vae_generator_seed": vae_generator_seed,
             },
         )
 
@@ -1763,6 +1785,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         self,
         images: torch.Tensor,
         cfg_factor: int = 1,
+        generator: torch.Generator | None = None,
     ) -> tuple[tuple[torch.Tensor, torch.Tensor], ...]:
         """
         Encode images through VAE encoder.
@@ -1782,11 +1805,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
             images = images.to(dtype=self.vae.dtype)
 
         vae_encode_result = self.vae.encode(images)
-        # Match HunyuanImage-3's cond encode path: sample the posterior, but
-        # use a fixed generator so online requests do not consume the global
-        # RNG and drift across a long-running server.
-        _cond_vae_gen = torch.Generator(device=images.device).manual_seed(0)
-        latents = vae_encode_result.latent_dist.sample(_cond_vae_gen)
+        latents = vae_encode_result.latent_dist.sample(generator)
 
         # Apply shift and scaling factors if present
         if hasattr(config, "shift_factor") and config.shift_factor:
@@ -1825,17 +1844,69 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         """
         Encode pixel_values through ViT encoder (vision_model and vision_aligner).
         """
-        # Handle empty batch
+        # Handle empty batch. Multimodal inputs are identical across TP ranks,
+        # so this returns on all ranks together without collective-op deadlock.
         if pixel_values.shape[0] == 0:
             return None
 
-        vision_output = self.vision_model(
-            pixel_values, attention_mask=vit_attention_mask, spatial_shapes=vit_spatial_shapes
-        )
-        image_embed = vision_output.last_hidden_state
+        if self.use_data_parallel:
+            # ViT weights are replicated (disable_tp on every linear); shard the
+            # image batch across ranks instead of doing redundant full-batch
+            # work on each one.
+            image_embed = self._vit_encode_dp(pixel_values, vit_attention_mask, vit_spatial_shapes)
+        else:
+            image_embed = self.vision_model(
+                pixel_values,
+                attention_mask=vit_attention_mask,
+                spatial_shapes=vit_spatial_shapes,
+            )
         image_embed = self.vision_aligner(image_embed)
         return image_embed
 
+    def _vit_encode_dp(
+        self,
+        pixel_values: torch.Tensor,
+        vit_attention_mask: torch.Tensor,
+        vit_spatial_shapes: torch.Tensor,
+    ) -> torch.Tensor:
+        """Data-parallel ViT: each rank runs a slice of the image batch, then
+        outputs are all-gathered. Mirrors ``run_dp_sharded_vision_model`` but
+        shards the three batched ViT inputs together (pixel_values /
+        attention_mask / spatial_shapes) instead of a single tensor.
+        """
+        num_images = pixel_values.shape[0]
+        world_size = get_tensor_model_parallel_world_size()
+        rank = get_tensor_model_parallel_rank()
+        num_per_rank = (num_images + world_size - 1) // world_size
+        num_padded = num_per_rank * world_size - num_images
+
+        if num_padded > 0:
+            # Pad the batch so every rank gets an equal slice (required for
+            # all_gather). Dummy images use spatial_shapes (1, 1) with a single
+            # valid patch: Siglip2 interpolates position embeddings per image
+            # and rejects size 0, so (0, 0) padding would crash. Their outputs
+            # are dropped after the gather.
+            pad_pixels = pixel_values.new_zeros((num_padded, *pixel_values.shape[1:]))
+            pixel_values = torch.cat([pixel_values, pad_pixels], dim=0)
+
+            pad_mask = vit_attention_mask.new_zeros((num_padded, *vit_attention_mask.shape[1:]))
+            pad_mask[:, 0] = 1
+            vit_attention_mask = torch.cat([vit_attention_mask, pad_mask], dim=0)
+
+            pad_shapes = vit_spatial_shapes.new_ones((num_padded, *vit_spatial_shapes.shape[1:]))
+            vit_spatial_shapes = torch.cat([vit_spatial_shapes, pad_shapes], dim=0)
+
+        start = rank * num_per_rank
+        end = start + num_per_rank
+        image_embed = self.vision_model(
+            pixel_values[start:end],
+            attention_mask=vit_attention_mask[start:end],
+            spatial_shapes=vit_spatial_shapes[start:end],
+        )
+        # max_patches (dim 1) is identical across ranks, so dim-0 gather is safe.
+        image_embed = tensor_model_parallel_all_gather(image_embed.contiguous(), dim=0)
+        return image_embed[:num_images]
+
     def _timestep_encode(
         self,
         timestep: torch.Tensor,
@@ -1866,6 +1937,12 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         vit_pixel_attention_mask = pixel_values["vit_pixel_attention_mask"]
         vit_spatial_shapes = pixel_values["vit_spatial_shapes"]
         vae_pixel_values = pixel_values["vae_pixel_values"]
+        vae_generator_seed = pixel_values.get("vae_generator_seed")
+        vae_generator = None
+        if vae_generator_seed is not None and vae_generator_seed.numel() > 0:
+            vae_generator = torch.Generator(device=vae_pixel_values[0].device).manual_seed(
+                int(vae_generator_seed.reshape(-1)[0].item())
+            )
 
         # Perform ViT encoding
         vit_embeddings = self._vit_encode(vit_pixel_values, vit_pixel_attention_mask, vit_spatial_shapes)
@@ -1874,7 +1951,7 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo
         # `reso_group` bucket so shapes are ragged across the image-batch dim.
         vae_token_embeddings = []
         for vae_image_i in vae_pixel_values:
-            t_i, latents_i = self._vae_encode(vae_image_i.unsqueeze(0), vae_cfg_factor)
+            t_i, latents_i = self._vae_encode(vae_image_i.unsqueeze(0), vae_cfg_factor, generator=vae_generator)
             t_emb = self.time_embed(t_i[0])
             vae_tokens, _, _ = self.patch_embed(latents_i, t_emb)
             vae_token_embeddings.append(vae_tokens)

vllm_omni/model_executor/models/hunyuan_image3/pipeline.py

@@ -26,7 +26,8 @@ HUNYUAN_IMAGE3_PIPELINE = PipelineConfig(
             model_stage="AR",
             execution_type=StageExecutionType.LLM_AR,
             input_sources=(),
-            final_output=False,
+            final_output=True,
+            final_output_type="text",
             owns_tokenizer=False,
             requires_multimodal_data=True,
             model_arch=_HUNYUAN_IMAGE3_MODEL_ARCH,

vllm_omni/model_executor/models/hunyuan_image3/siglip2.py

@@ -25,18 +25,37 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+"""SigLIP2 Vision Transformer for HunyuanImage3, rewritten in vLLM style.
 
-import warnings
+Key optimizations over the original HuggingFace-style implementation:
+- QKVParallelLinear: fused QKV projection with tensor parallelism support
+- MMEncoderAttention: FlashAttention / xFormers backend
+- ColumnParallelLinear / RowParallelLinear: TP-aware MLP layers
+- Packed sequence processing: eliminates padding waste via cu_seqlens
+- Data parallel support for multi-GPU ViT inference
+"""
+
+from collections.abc import Iterable
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from transformers.activations import ACT2FN
-from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask
-from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from vllm.distributed import divide, get_tensor_model_parallel_world_size
+from vllm.model_executor.layers.activation import get_act_fn
+from vllm.model_executor.layers.attention import MMEncoderAttention
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    QKVParallelLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.vision import is_vit_use_data_parallel
 
 
 class Config:
+    """Convert dict config to object with attribute access."""
+
     def __init__(self, config):
         if config is not None:
             for key, value in config.items():
@@ -65,470 +84,336 @@ class Siglip2VisionEmbeddings(nn.Module):
         self.position_embedding_size = int(self.num_patches**0.5)
         self.position_embedding = nn.Embedding(self.num_patches, self.embed_dim)
 
-    @staticmethod
-    def resize_positional_embeddings(
-        positional_embeddings: torch.Tensor,
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor,
         spatial_shapes: torch.LongTensor,
-        max_length: int,
     ) -> torch.Tensor:
         """
-        Resize positional embeddings to image-specific size and pad to a fixed size.
+        Process packed pixel values with per-image position embeddings.
+
         Args:
-            positional_embeddings (`torch.Tensor`):
-                Position embeddings of shape (height, width, embed_dim)
-            spatial_shapes (`torch.LongTensor`):
-                Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to
-            max_length (`int`):
-                Maximum length of the positional embeddings to pad resized positional embeddings to
+            pixel_values: Packed pixel values
+                (total_real_patches, num_channels * patch_size * patch_size)
+            spatial_shapes: Per-image spatial shapes (B, 2) as [(h, w), ...]
+
         Returns:
-            `torch.Tensor`: Embeddings of shape (batch_size, max_length, embed_dim)
+            Packed embeddings (total_real_patches, embed_dim)
         """
-        batch_size = spatial_shapes.shape[0]
-        embed_dim = positional_embeddings.shape[-1]
-        source_dtype = positional_embeddings.dtype
+        target_dtype = self.patch_embedding.weight.dtype
+        patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
 
-        resulted_positional_embeddings = torch.empty(
-            (batch_size, max_length, embed_dim),
-            device=positional_embeddings.device,
-            dtype=source_dtype,
+        # Resize position embeddings per image and concatenate (packed)
+        positional_embeddings = self.position_embedding.weight.reshape(
+            self.position_embedding_size, self.position_embedding_size, -1
         )
+        # (H_pe, W_pe, embed_dim) → (1, embed_dim, H_pe, W_pe) for interpolation
+        pe_for_resize = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
+        # Upcast on CPU since antialias is not supported for bfloat16/float16
+        if pe_for_resize.device.type == "cpu":
+            pe_for_resize = pe_for_resize.to(torch.float32)
 
-        # (height, width, embed_dim) -> (1, embed_dim, height, width) for interpolation
-        positional_embeddings = positional_embeddings.permute(2, 0, 1).unsqueeze(0)
-
-        # Upcast to float32 on CPU because antialias is not supported for bfloat16/float16 on CPU
-        if positional_embeddings.device.type == "cpu":
-            positional_embeddings = positional_embeddings.to(torch.float32)
-
-        for i in range(batch_size):
-            # (1, dim, height, width) -> (1, dim, target_height, target_width)
-            height, width = spatial_shapes[i]
-            resized_embeddings = F.interpolate(
-                positional_embeddings,
+        position_embs: list[torch.Tensor] = []
+        for i in range(spatial_shapes.shape[0]):
+            height, width = int(spatial_shapes[i, 0]), int(spatial_shapes[i, 1])
+            resized = F.interpolate(
+                pe_for_resize,
                 size=(height, width),
                 mode="bilinear",
                 align_corners=False,
                 antialias=True,
             )
+            # (1, embed_dim, h, w) → (h*w, embed_dim)
+            resized = resized.reshape(self.embed_dim, height * width).transpose(0, 1)
+            position_embs.append(resized.to(target_dtype))
 
-            # (1, dim, target_height, target_width) -> (target_height * target_width, dim)
-            resized_embeddings = resized_embeddings.reshape(embed_dim, height * width).transpose(0, 1)
-
-            # Cast to original dtype
-            resized_embeddings = resized_embeddings.to(source_dtype)
-
-            resulted_positional_embeddings[i, : height * width] = resized_embeddings
-            resulted_positional_embeddings[i, height * width :] = resized_embeddings[0]
-
-        return resulted_positional_embeddings
-
-    def forward(self, pixel_values: torch.FloatTensor, spatial_shapes: torch.LongTensor) -> torch.Tensor:
-        """
-        Args:
-            pixel_values (`torch.FloatTensor`):
-                Pixel values of shape (batch_size, max_num_patches, num_channels * patch_size * patch_size)
-            spatial_shapes (`List[Tuple[int, int]]`):
-                Spatial shapes of shape (batch_size, 2) to resize the positional embeddings to
-        """
-
-        # Apply patch embeddings to already patchified pixel values
-        target_dtype = self.patch_embedding.weight.dtype
-        patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))
-
-        # Get positional resized and padded positional embeddings
-        positional_embeddings = self.position_embedding.weight.reshape(
-            self.position_embedding_size, self.position_embedding_size, -1
-        )
-        resized_positional_embeddings = self.resize_positional_embeddings(
-            positional_embeddings, spatial_shapes, max_length=pixel_values.shape[1]
-        )
-
-        # Add positional embeddings to patch embeddings
-        embeddings = patch_embeds + resized_positional_embeddings
-        return embeddings
+        packed_position_embs = torch.cat(position_embs, dim=0)
+        return patch_embeds + packed_position_embs
 
 
 class Siglip2Attention(nn.Module):
-    """Multi-headed attention from 'Attention Is All You Need' paper"""
+    """Multi-headed attention using QKVParallelLinear and MMEncoderAttention."""
 
-    def __init__(self, config):
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
-        self.config = config
         self.embed_dim = config.hidden_size
         self.num_heads = config.num_attention_heads
         self.head_dim = self.embed_dim // self.num_heads
-        if self.head_dim * self.num_heads != self.embed_dim:
-            raise ValueError(
-                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
-                f" {self.num_heads})."
-            )
         self.scale = self.head_dim**-0.5
-        self.dropout = config.attention_dropout
 
-        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
-        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
-        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
-        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        output_attentions: bool | None = False,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """Input shape: Batch x Time x Channel"""
-
-        batch_size, q_len, _ = hidden_states.size()
-
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
-
-        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-
-        k_v_seq_len = key_states.shape[-2]
-        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) * self.scale
-
-        if attn_weights.size() != (batch_size, self.num_heads, q_len, k_v_seq_len):
-            raise ValueError(
-                f"Attention weights should be of size {(batch_size, self.num_heads, q_len, k_v_seq_len)}, but is"
-                f" {attn_weights.size()}"
-            )
-
-        if attention_mask is not None:
-            if attention_mask.size() != (batch_size, 1, q_len, k_v_seq_len):
-                raise ValueError(
-                    f"Attention mask should be of size {(batch_size, 1, q_len, k_v_seq_len)}, "
-                    f"but is {attention_mask.size()}"
-                )
-            attn_weights = attn_weights + attention_mask
-
-        # upcast attention to fp32
-        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
-        attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
-        attn_output = torch.matmul(attn_weights, value_states)
-
-        if attn_output.size() != (batch_size, self.num_heads, q_len, self.head_dim):
-            raise ValueError(
-                f"`attn_output` should be of size {(batch_size, self.num_heads, q_len, self.head_dim)}, but is"
-                f" {attn_output.size()}"
-            )
-
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.reshape(batch_size, q_len, self.embed_dim)
-
-        attn_output = self.out_proj(attn_output)
-
-        return attn_output, attn_weights
-
-
-class Siglip2SdpaAttention(Siglip2Attention):
-    """
-    Siglip2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
-    `Siglip2Attention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt
-    to SDPA API.
-    """
-
-    is_causal = False
-
-    # Adapted from Siglip2Attention.forward and transformers.models.llama.modeling_llama.LlamaSdpaAttention.forward
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
-        output_attentions: bool | None = False,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        if output_attentions:
-            # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"`
-            #  once this is implemented.
-            warnings.warn(
-                "Siglip2Model is using Siglip2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` "
-                "does not support `output_attentions=True`. Falling back to the manual attention implementation, "
-                "but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. "
-                'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
-            )
-            return super().forward(
-                hidden_states=hidden_states,
-                attention_mask=attention_mask,
-                output_attentions=output_attentions,
-            )
-
-        batch_size, q_len, _ = hidden_states.size()
-
-        query_states = self.q_proj(hidden_states)
-        key_states = self.k_proj(hidden_states)
-        value_states = self.v_proj(hidden_states)
-
-        query_states = query_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        key_states = key_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-        value_states = value_states.view(batch_size, q_len, self.num_heads, self.head_dim).transpose(1, 2)
-
-        # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with
-        # custom attn_mask,
-        # Reference: https://github.com/pytorch/pytorch/issues/112577.
-        if query_states.device.type == "cuda" and attention_mask is not None:
-            query_states = query_states.contiguous()
-            key_states = key_states.contiguous()
-            value_states = value_states.contiguous()
-
-        # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of an
-        # inline conditional assignment in SDPA to support both torch.compile's dynamic shapes and full graph options.
-        # An inline conditional prevents dynamic shapes from compiling.
-        is_causal = True if self.is_causal and q_len > 1 else False
-
-        attn_output = torch.nn.functional.scaled_dot_product_attention(
-            query_states,
-            key_states,
-            value_states,
-            attn_mask=attention_mask,
-            dropout_p=self.dropout if self.training else 0.0,
-            is_causal=is_causal,
+        use_data_parallel = is_vit_use_data_parallel()
+        self.qkv_proj = QKVParallelLinear(
+            hidden_size=self.embed_dim,
+            head_size=self.head_dim,
+            total_num_heads=self.num_heads,
+            quant_config=quant_config,
+            prefix=f"{prefix}.qkv_proj",
+            disable_tp=use_data_parallel,
+        )
+        self.out_proj = RowParallelLinear(
+            input_size=self.embed_dim,
+            output_size=self.embed_dim,
+            quant_config=quant_config,
+            prefix=f"{prefix}.out_proj",
+            disable_tp=use_data_parallel,
         )
 
-        attn_output = attn_output.transpose(1, 2).contiguous()
-        attn_output = attn_output.view(batch_size, q_len, self.embed_dim)
+        self.tp_size = 1 if use_data_parallel else get_tensor_model_parallel_world_size()
+        self.num_heads_per_partition = divide(self.num_heads, self.tp_size)
 
-        attn_output = self.out_proj(attn_output)
+        self.attn = MMEncoderAttention(
+            num_heads=self.num_heads_per_partition,
+            head_size=self.head_dim,
+            scale=self.scale,
+            prefix=f"{prefix}.attn",
+        )
 
-        return attn_output, None
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Args:
+            hidden_states: Packed input (total_tokens, embed_dim)
+            cu_seqlens: Cumulative sequence lengths (B+1,)
+        """
+        seq_length = hidden_states.shape[0]
+
+        qkv, _ = self.qkv_proj(hidden_states)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        q = q.view(seq_length, self.num_heads_per_partition, self.head_dim)
+        k = k.view(seq_length, self.num_heads_per_partition, self.head_dim)
+        v = v.view(seq_length, self.num_heads_per_partition, self.head_dim)
+
+        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+        attn_output = self.attn(
+            query=q.unsqueeze(0),
+            key=k.unsqueeze(0),
+            value=v.unsqueeze(0),
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+        attn_output = attn_output.reshape(seq_length, self.num_heads_per_partition * self.head_dim)
+
+        attn_output, _ = self.out_proj(attn_output)
+        return attn_output
 
 
 class Siglip2MLP(nn.Module):
-    def __init__(self, config):
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
-        self.config = config
-        self.activation_fn = ACT2FN[config.hidden_act]
-        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
-        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+        use_data_parallel = is_vit_use_data_parallel()
+        self.activation_fn = get_act_fn(config.hidden_act)
+        self.fc1 = ColumnParallelLinear(
+            config.hidden_size,
+            config.intermediate_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc1",
+            disable_tp=use_data_parallel,
+        )
+        self.fc2 = RowParallelLinear(
+            config.intermediate_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.fc2",
+            disable_tp=use_data_parallel,
+        )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.fc1(hidden_states)
+        hidden_states, _ = self.fc1(hidden_states)
         hidden_states = self.activation_fn(hidden_states)
-        hidden_states = self.fc2(hidden_states)
+        hidden_states, _ = self.fc2(hidden_states)
         return hidden_states
 
 
 class Siglip2EncoderLayer(nn.Module):
-    def __init__(self, config):
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
         self.embed_dim = config.hidden_size
-        self.self_attn = Siglip2Attention(config=config)
+        self.self_attn = Siglip2Attention(
+            config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn",
+        )
         self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
-        self.mlp = Siglip2MLP(config)
+        self.mlp = Siglip2MLP(
+            config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mlp",
+        )
         self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
 
-    # Ignore copy
     def forward(
         self,
         hidden_states: torch.Tensor,
-        attention_mask: torch.Tensor,
-        output_attentions: bool | None = False,
-    ) -> tuple[torch.FloatTensor]:
+        cu_seqlens: torch.Tensor,
+    ) -> torch.Tensor:
         """
         Args:
-            hidden_states (`torch.FloatTensor`):
-                Input to the layer of shape `(batch, seq_len, embed_dim)`.
-            attention_mask (`torch.FloatTensor`):
-                Attention mask of shape `(batch, 1, q_len, k_v_seq_len)` where padding elements are indicated by very
-                large negative values.
-            output_attentions (`bool`, *optional*, defaults to `False`):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
+            hidden_states: Packed input (total_tokens, embed_dim)
+            cu_seqlens: Cumulative sequence lengths (B+1,)
         """
         residual = hidden_states
-
         hidden_states = self.layer_norm1(hidden_states)
-        hidden_states, attn_weights = self.self_attn(
-            hidden_states=hidden_states,
-            attention_mask=attention_mask,
-            output_attentions=output_attentions,
-        )
+        hidden_states = self.self_attn(hidden_states, cu_seqlens=cu_seqlens)
         hidden_states = residual + hidden_states
 
         residual = hidden_states
         hidden_states = self.layer_norm2(hidden_states)
         hidden_states = self.mlp(hidden_states)
         hidden_states = residual + hidden_states
-
-        outputs = (hidden_states,)
-
-        if output_attentions:
-            outputs += (attn_weights,)
-
-        return outputs
+        return hidden_states
 
 
 class Siglip2Encoder(nn.Module):
-    """
-    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
-    [`Siglip2EncoderLayer`].
-    Args:
-        config: Siglip2Config
-    """
+    """Transformer encoder with packed sequence processing."""
 
-    def __init__(self, config):
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
-        self.config = config
-        self.layers = nn.ModuleList([Siglip2EncoderLayer(config) for _ in range(config.num_hidden_layers)])
-        self.gradient_checkpointing = True
+        self.layers = nn.ModuleList(
+            [
+                Siglip2EncoderLayer(
+                    config,
+                    quant_config=quant_config,
+                    prefix=f"{prefix}.layers.{idx}",
+                )
+                for idx in range(config.num_hidden_layers)
+            ]
+        )
 
-    # Ignore copy
     def forward(
         self,
-        inputs_embeds,
-        attention_mask: torch.Tensor | None = None,
-        output_attentions: bool | None = None,
-        output_hidden_states: bool | None = None,
-        return_dict: bool | None = None,
-    ) -> tuple | BaseModelOutput:
-        r"""
-        Args:
-            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
-                Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.
-                This is useful if you want more control over how to convert `input_ids` indices into associated vectors
-                than the model's internal embedding lookup matrix.
-            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
-                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
-                - 1 for tokens that are **not masked**,
-                - 0 for tokens that are **masked**.
-                What are attention masks? See HuggingFace documentation for details.
-            output_attentions (`bool`, *optional*):
-                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
-                returned tensors for more detail.
-            output_hidden_states (`bool`, *optional*):
-                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
-                for more detail.
-            return_dict (`bool`, *optional*):
-                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        hidden_states: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+    ) -> torch.Tensor:
         """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        encoder_states = () if output_hidden_states else None
-        all_attentions = () if output_attentions else None
-
-        hidden_states = inputs_embeds
-        for layer_index, encoder_layer in enumerate(self.layers):  # len(self.layers): 27
-            if output_hidden_states:
-                encoder_states = encoder_states + (hidden_states,)
-
-            layer_outputs = encoder_layer(
-                hidden_states,
-                attention_mask,
-                output_attentions=output_attentions,
-            )
-
-            hidden_states = layer_outputs[0]
-
-            if output_attentions:
-                all_attentions = all_attentions + (layer_outputs[1],)
-
-        if output_hidden_states:
-            encoder_states = encoder_states + (hidden_states,)
-
-        if not return_dict:
-            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
-        return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions)
-
-
-class Siglip2MultiheadAttentionPoolingHead(nn.Module):
-    """Multihead Attention Pooling."""
-
-    def __init__(self, config):
-        super().__init__()
-
-        self.probe = nn.Parameter(torch.randn(1, 1, config.hidden_size))
-        self.attention = torch.nn.MultiheadAttention(config.hidden_size, config.num_attention_heads, batch_first=True)
-        self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.mlp = Siglip2MLP(config)
-        self.num_heads = config.num_attention_heads
-
-    def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor | None = None):
-        batch_size = hidden_state.shape[0]
-        probe = self.probe.repeat(batch_size, 1, 1)
-
-        if attention_mask is not None:
-            target_len, source_len = probe.shape[1], hidden_state.shape[1]
-            attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_state.dtype, target_len)
-            attention_mask = attention_mask.repeat(1, self.num_heads, target_len, 1)
-            attention_mask = attention_mask.reshape(-1, target_len, source_len)
-
-        hidden_state = self.attention(probe, hidden_state, hidden_state, attn_mask=attention_mask)[0]
-
-        residual = hidden_state
-        hidden_state = self.layernorm(hidden_state)
-        hidden_state = residual + self.mlp(hidden_state)
-
-        return hidden_state[:, 0]
+        Args:
+            hidden_states: Packed input (total_tokens, embed_dim)
+            cu_seqlens: Cumulative sequence lengths (B+1,)
+        """
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, cu_seqlens)
+        return hidden_states
 
 
 class Siglip2VisionTransformer(nn.Module):
-    def __init__(self, config):
+    def __init__(
+        self,
+        config,
+        quant_config: QuantizationConfig | None = None,
+        prefix: str = "",
+    ):
         super().__init__()
         config = Config(config)
         self.config = config
-        embed_dim = config.hidden_size
+        self.embed_dim = config.hidden_size
 
         self.embeddings = Siglip2VisionEmbeddings(config)
-        self.encoder = Siglip2Encoder(config)
-        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
-        self.use_head = True if not hasattr(config, "vision_use_head") else config.vision_use_head
-        if self.use_head:
-            self.head = Siglip2MultiheadAttentionPoolingHead(config)
-        self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
+        self.encoder = Siglip2Encoder(
+            config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.encoder" if prefix else "encoder",
+        )
+        self.post_layernorm = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
 
     def forward(
         self,
         pixel_values: torch.FloatTensor,
         attention_mask: torch.Tensor,
         spatial_shapes: torch.LongTensor,
-        output_attentions: bool | None = None,
-        output_hidden_states: bool | None = None,
-        return_dict: bool | None = None,
-    ) -> tuple | BaseModelOutputWithPooling:
-        r"""
-        Returns:
+    ) -> torch.Tensor:
         """
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        Args:
+            pixel_values: Batched pixel values
+                (B, max_num_patches, num_channels * patch_size * patch_size)
+            attention_mask: (B, max_num_patches) with 1 for real, 0 for padding
+            spatial_shapes: (B, 2) with (height, width) per image
+
+        Returns:
+            (B, max_num_patches, hidden_size) with zeros at padding positions
+        """
+        batch_size, max_patches, _ = pixel_values.shape
+
+        # Pack: extract real tokens using attention_mask
+        mask_bool = attention_mask.bool()
+        packed_pixels = pixel_values[mask_bool]  # (total_real_patches, C*P*P)
+
+        # Compute cu_seqlens from spatial_shapes (h * w per image)
+        seq_lens = (spatial_shapes[:, 0] * spatial_shapes[:, 1]).to(torch.int32)
+        cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=pixel_values.device)
+        cu_seqlens[1:] = seq_lens.cumsum(0)
+
+        # Embeddings (packed)
+        hidden_states = self.embeddings(packed_pixels, spatial_shapes)
+
+        # Encoder (packed)
+        hidden_states = self.encoder(hidden_states, cu_seqlens)
+
+        # Post layernorm
+        hidden_states = self.post_layernorm(hidden_states)
+
+        # Unpack: scatter back to (B, max_patches, hidden_size)
+        output = torch.zeros(
+            batch_size,
+            max_patches,
+            self.embed_dim,
+            dtype=hidden_states.dtype,
+            device=hidden_states.device,
         )
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output[mask_bool] = hidden_states
 
-        hidden_states = self.embeddings(pixel_values, spatial_shapes)
+        return output
 
-        if attention_mask is not None and not self._use_flash_attention_2:
-            # [batch_size, seq_len] -> [batch_size, 1, tgt_seq_len, src_seq_len]
-            encoder_attention_mask = _prepare_4d_attention_mask(attention_mask, hidden_states.dtype)
-        else:
-            encoder_attention_mask = attention_mask
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("qkv_proj", "q_proj", "q"),
+            ("qkv_proj", "k_proj", "k"),
+            ("qkv_proj", "v_proj", "v"),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
 
-        encoder_outputs = self.encoder(
-            inputs_embeds=hidden_states,
-            attention_mask=encoder_attention_mask,
-            output_attentions=output_attentions,
-            output_hidden_states=output_hidden_states,
-            return_dict=return_dict,
-        )
-
-        last_hidden_state = encoder_outputs[0]
-        last_hidden_state = self.post_layernorm(last_hidden_state)
-
-        pooler_output = self.head(last_hidden_state, attention_mask) if self.use_head else None
-        if not return_dict:
-            return (last_hidden_state, pooler_output) + encoder_outputs[1:]
-
-        return BaseModelOutputWithPooling(
-            last_hidden_state=last_hidden_state,
-            pooler_output=pooler_output,
-            hidden_states=encoder_outputs.hidden_states,
-            attentions=encoder_outputs.attentions,
-        )
+        for name, loaded_weight in weights:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                if name not in params_dict:
+                    # Skip weights for removed modules (e.g., pooling head)
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
 
 
 class LightProjector(nn.Module):

vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py

@@ -72,6 +72,7 @@ from vllm_omni.transformers_utils.processors.ming import (
     PLACEHOLDER_VIDEO_TOKEN_IN_TEXT,
     MingFlashOmniProcessor,
     MingWhisperFeatureExtractor,
+    raise_missing_video_processor,
 )
 
 from .audio_encoder import WhisperAudioEncoder
@@ -540,20 +541,20 @@ class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashO
         if images is not None:
             image_outputs = hf_processor.image_processor(
                 images=images,
-                videos=None,
                 return_tensors="pt",
             )
             data.update(image_outputs)
 
         videos = mm_data.get("videos", None)
         if videos is not None:
-            # TODO: ``videos=`` on image_processor is deprecated since
-            # transformers v4.57 (removed in v5); migrate to Qwen2VLVideoProcessor.
-            video_outputs = hf_processor.image_processor(
-                images=None,
-                videos=videos,
-                return_tensors="pt",
-            )
+            video_processor = getattr(hf_processor, "video_processor", None)
+            if video_processor is not None:
+                video_outputs = video_processor(
+                    videos=videos,
+                    return_tensors="pt",
+                )
+            else:
+                raise_missing_video_processor()
             # Rename keys to distinguish from images
             if "pixel_values" in video_outputs:
                 video_outputs["pixel_values_videos"] = video_outputs.pop("pixel_values")

vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py

@@ -714,6 +714,9 @@ class BailingMoeV2Model(nn.Module):
     def get_input_embeddings(self):
         return self.word_embeddings
 
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.word_embeddings(input_ids)
+
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -794,6 +797,9 @@ class BailingMoeV2ForCausalLM(nn.Module, CustomProcessMixin):
         self.sampler = Sampler()
         self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors
 
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
     def forward(
         self,
         input_ids: torch.Tensor,
@@ -818,7 +824,7 @@ class BailingMoeV2ForCausalLM(nn.Module, CustomProcessMixin):
         hidden_states: torch.Tensor,
         sampling_metadata,
     ) -> torch.Tensor | None:
-        logits = self.logits_processor(self.lm_head, hidden_states, sampling_metadata)
+        logits = self.logits_processor(self.lm_head, hidden_states)
         return logits
 
     def sample(

vllm_omni/model_executor/models/ming_flash_omni/talker_module.py

@@ -902,6 +902,7 @@ class MingAudioGenerator:
         if self._audio_vae is None:
             return requested_max_steps
 
+        # Transformers >=5.x may expose these config values as 0-d tensors.
         sample_rate = float(self._audio_vae.config.sample_rate)
         vae_patch_size = float(getattr(self._audio_vae.config, "patch_size", 4))
         hop_size = float(getattr(self._audio_vae.decoder, "hop_length", 320))
@@ -1041,7 +1042,7 @@ class MingAudioGenerator:
                 use_cache=True,
             )
         else:
-            past_seen_tokens = past_key_values.get_seq_length()
+            past_seen_tokens = int(past_key_values.get_seq_length())
             cache_position = torch.arange(
                 past_seen_tokens,
                 past_seen_tokens + inputs_embeds.shape[1],

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -36,6 +36,7 @@ from vllm.v1.sample.metadata import SamplingMetadata
 from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.data_entry_keys import Embeddings, HiddenStates, Ids, OmniPayload, OmniPayloadMeta
+from vllm_omni.metrics import definitions as defs
 from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import (
@@ -527,9 +528,16 @@ class Qwen3OmniMoeForConditionalGeneration(
             return OmniOutput(text_hidden_states=talker_hidden, multimodal_outputs=multimodal_outputs)
         elif self.model_stage == "code2wav":
             audio_tensors = model_outputs
+            sample_rate = defs.resolve_audio_sample_rate(self.code2wav_config)
+            # `sr` is the audio sample rate metadata consumed by downstream
+            # audio serving and stage-local audio metrics.
+            sr_tensors = [torch.tensor(sample_rate, dtype=torch.int32) for _ in audio_tensors]
             return OmniOutput(
                 text_hidden_states=None,
-                multimodal_outputs={"model_outputs": [audio_tensor.reshape(1, -1) for audio_tensor in audio_tensors]},
+                multimodal_outputs={
+                    "model_outputs": [audio_tensor.reshape(1, -1) for audio_tensor in audio_tensors],
+                    "sr": sr_tensors,
+                },
             )
 
         return model_outputs