diff --git a/docs/cli/bench/serve.md b/docs/cli/bench/serve.md index f6e43be09..d6b3ef393 100644 --- a/docs/cli/bench/serve.md +++ b/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`. + +### Multi-Stage Benchmark + +
+Show more + +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. + +
diff --git a/pyproject.toml b/pyproject.toml index 203475ad4..0950b2682 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/tests/assets/hunyuan/hunyuan_baseline.png b/tests/assets/hunyuan/hunyuan_baseline.png index fe49f87f6..b2bbb60f2 100644 Binary files a/tests/assets/hunyuan/hunyuan_baseline.png and b/tests/assets/hunyuan/hunyuan_baseline.png differ diff --git a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json b/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json index 661f6d263..b95aacd05 100644 --- a/tests/dfx/perf/tests/test_hunyuan_image3_it2i.json +++ b/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, diff --git a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json b/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json index 76ac1d858..a973bb0e6 100644 --- a/tests/dfx/perf/tests/test_qwen3_omni_vllm_text.json +++ b/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": [ diff --git a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py b/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py index fe705734b..106507114 100644 --- a/tests/distributed/omni_connectors/test_bagel_mooncake_connector.py +++ b/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.""" diff --git a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py b/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py index 3906611b1..eb5a79d6a 100644 --- a/tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py +++ b/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.""" diff --git a/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py b/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py index b0dc2af2d..5e872e73a 100644 --- a/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py +++ b/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( diff --git a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py b/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py index b5f1bc4c3..6ac930f46 100644 --- a/tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py +++ b/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 diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 0f6163b73..b86cc0af9 100644 --- a/tests/engine/test_orchestrator.py +++ b/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(), ) ) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index d6a7dbaf7..38cb86c6e 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/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=[]) diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 5d7176490..b8637c1e9 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/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 diff --git a/tests/metrics/test_prometheus.py b/tests/metrics/test_prometheus.py index 63becf036..1f79b186b 100644 --- a/tests/metrics/test_prometheus.py +++ b/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. diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 9119d78cd..e541d3628 100644 --- a/tests/test_config_factory.py +++ b/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 diff --git a/vllm_omni/benchmarks/metrics/metrics.py b/vllm_omni/benchmarks/metrics/metrics.py index 55f4b8830..100ab19ef 100644 --- a/vllm_omni/benchmarks/metrics/metrics.py +++ b/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 diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 7238cdf88..3357fa8bc 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/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: diff --git a/vllm_omni/benchmarks/serve.py b/vllm_omni/benchmarks/serve.py index d3f3510c5..8e3b19ea9 100644 --- a/vllm_omni/benchmarks/serve.py +++ b/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)) diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index 3bfa6042a..06869a657 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/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 diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py index d8eb2ac53..a2717bcc8 100644 --- a/vllm_omni/diffusion/inline_stage_diffusion_client.py +++ b/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 diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 02caed9e3..56a5b0dcc 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/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( diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index a0f80e8a1..9cc2aafe5 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/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) diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 6ab44b4b0..c8da748eb 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/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) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index d5abe7809..da1497c0c 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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, ) diff --git a/vllm_omni/engine/messages.py b/vllm_omni/engine/messages.py index 0ced58391..536f8d525 100644 --- a/vllm_omni/engine/messages.py +++ b/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): diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 99ee7f93b..017362812 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/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 diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 5fb46ed0d..369fe3c7b 100644 --- a/vllm_omni/engine/output_processor.py +++ b/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 + ) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 81ac09048..53267e23b 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/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) diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index c6a40f55e..4e1f3d5a5 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/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) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index ae770346d..93f2dadff 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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, ) diff --git a/vllm_omni/entrypoints/cli/benchmark/serve.py b/vllm_omni/entrypoints/cli/benchmark/serve.py index 3ba03e297..eed7d4f4e 100644 --- a/vllm_omni/entrypoints/cli/benchmark/serve.py +++ b/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) diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 7b1b15e48..020b483f0 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/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 diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index b10dd777d..d415c07ea 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/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, ) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index b19360624..1be7aceba 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/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, diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index 7f82151bf..83338624a 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/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): diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index afe623f34..c7b24e11f 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 3fb991225..303b3b4a1 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/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() diff --git a/vllm_omni/metrics/definitions.py b/vllm_omni/metrics/definitions.py index b45217abc..6c1288edf 100644 --- a/vllm_omni/metrics/definitions.py +++ b/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: diff --git a/vllm_omni/metrics/stats.py b/vllm_omni/metrics/stats.py index 88e7e770a..3db2fda14 100644 --- a/vllm_omni/metrics/stats.py +++ b/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 diff --git a/vllm_omni/metrics/utils.py b/vllm_omni/metrics/utils.py index ea83d959e..1ed987a19 100644 --- a/vllm_omni/metrics/utils.py +++ b/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 diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 02326f456..d20312a97 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/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 diff --git a/vllm_omni/model_executor/models/hunyuan_image3/pipeline.py b/vllm_omni/model_executor/models/hunyuan_image3/pipeline.py index 94b4cd837..06ebff9c7 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/pipeline.py +++ b/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, diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 3966f2140..4a144bff0 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/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