Merge branch 'main' into fix/omni-weight-snapshot-prefetch

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>

pyproject.toml

@@ -64,7 +64,8 @@ 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
 ]
 
 demo = [

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_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/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.25,
         help="If set, fail when seed_tts_content_error_mean is strictly above this value.",
     )
     p.add_argument(

tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py

@@ -26,7 +26,7 @@ HEIGHT = 1024
 WIDTH = 1024
 PROMPT = "A brown and white dog is running on the grass."
 MEAN_THRESHOLD = 3e-2
-P99_THRESHOLD = 4e-1
+P99_THRESHOLD = 3e-1
 SSIM_THRESHOLD = 0.97
 PSNR_THRESHOLD = 30.0
 

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

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/deploy/qwen3_omni_moe.yaml

@@ -31,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/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/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/hunyuan_image3/pipeline_hunyuan_image3.py

@@ -865,7 +865,9 @@ 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]
             )
@@ -1459,6 +1461,7 @@ class HunyuanImage3Pipeline(
             num_inference_steps=num_inference_steps,
             guidance_scale=guidance_scale,
             batch_cond_image_info=batch_cond_image_info,
+            **ar_kv_kwargs,
         )
 
         model_inputs.update(ar_kv_kwargs)

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/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:
@@ -775,7 +788,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)
 
@@ -852,6 +867,12 @@ 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 +881,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 +936,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 +1362,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/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(
@@ -2050,6 +2078,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 +2087,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,

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/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)
@@ -1624,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
@@ -1705,6 +1731,48 @@ 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,
+                        completion_tokens=0,
+                        total_tokens=num_prompt_tokens,
+                    )
+                    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
@@ -1772,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"
@@ -1926,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)
@@ -2248,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"):
@@ -2696,6 +2768,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",
@@ -2717,6 +2790,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)
@@ -2755,6 +2829,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(
@@ -2763,6 +2838,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,
                 )
 
@@ -2818,6 +2894,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."""
@@ -2827,6 +2904,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:
@@ -2838,6 +2916,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":
@@ -2861,6 +2940,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/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/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

@@ -134,9 +134,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/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/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/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