[Feature][Diffusion] Expose detailed pipeline timings (#6822)

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>

tests/diffusion/models/minimax_h3/test_minimax_h3_offload.py

@@ -13,12 +13,14 @@ from PIL import Image
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
 
 
-def test_h3_profiler_targets_cover_aggregate_reference_encoders():
+def test_h3_profiler_targets_cover_reference_encoders_and_vae_decoders():
     from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
 
     targets = MiniMaxH3Pipeline._PROFILER_TARGETS
     assert "_encode_visual_conditions" in targets
     assert "_encode_reference_audio_conditions" in targets
+    assert "video_vae.decode_latent" in targets
+    assert "audio_vae.decode_latent" in targets
     assert "_encode_video_conditions" not in targets
     assert "_encode_video_audio_conditions" not in targets
     assert "_encode_audio_conditions" not in targets

tests/diffusion/test_diffusion_engine_metrics.py

@@ -140,6 +140,7 @@ class TestMetricKeys:
                 raise AssertionError("diffusion_engine_total_time_ms should be attached in step_streaming()")
 
         assert '"diffusion_engine_exec_time_ms": exec_total_time * 1000' in step_streaming_source
+        assert '"output_ready_wait_time_ms": output_ready_wait_time * 1000' in step_streaming_source
         # step_total_ms is no longer emitted as a metric (no family consumes
         # it); it lives only in the debug log breakdown below.
         assert '"diffusion_engine_total_time_ms"' not in step_streaming_source

tests/metrics/test_modality.py

@@ -532,6 +532,7 @@ class TestDiffusionUnitConversion:
             metrics={
                 "preprocess_time_ms": 50.0,
                 "diffusion_engine_exec_time_ms": 1500.0,
+                "output_ready_wait_time_ms": 250.0,
                 "postprocess_time_ms": 20.0,
                 "vae_decode_time_ms": 300.0,
                 "forward_time_ms": 800.0,
@@ -545,6 +546,7 @@ class TestDiffusionUnitConversion:
         assert dict(aggregator.diffusion_metrics["req-units"]) == {
             "preprocess_time_s": pytest.approx(0.05),
             "diffusion_engine_exec_time_s": pytest.approx(1.5),
+            "output_ready_wait_time_s": pytest.approx(0.25),
             "postprocess_time_s": pytest.approx(0.02),
             "vae_decode_time_s": pytest.approx(0.3),
             "forward_time_s": pytest.approx(0.8),

vllm_omni/diffusion/diffusion_engine.py

@@ -439,8 +439,10 @@ class DiffusionEngine:
         generator = self.get_output_stream(request_id)
         async for output in generator:
             exec_total_time = time.perf_counter() - exec_start_time
+            output_ready_wait_time = 0.0
             # Async mode: wait for background D2H/SHM to complete.
             if output.async_output_id:
+                output_ready_wait_start_time = time.perf_counter()
                 fut = self.executor.wait_output_ready(output.async_output_id)
                 timeout = _async_output_timeout()
                 try:
@@ -455,6 +457,7 @@ class DiffusionEngine:
                         describe(output.async_output_id) if describe else "unavailable",
                     )
                     raise
+                output_ready_wait_time = time.perf_counter() - output_ready_wait_start_time
             postprocess_start_time = time.perf_counter()
             scheduler_metrics = diffusion_scheduler_waiting_metrics(getattr(self, "_scheduler_num_waiting_reqs", 0))
             try:
@@ -468,9 +471,11 @@ class DiffusionEngine:
             step_total_ms = (time.perf_counter() - diffusion_engine_start_time) * 1000
             logger.debug(
                 "DiffusionEngine.step_streaming breakdown: preprocess=%.2f ms, "
-                "add_req_and_wait=%.2f ms, postprocess=%.2f ms, total=%.2f ms",
+                "add_req_and_wait=%.2f ms, output_ready_wait=%.2f ms, "
+                "postprocess=%.2f ms, total=%.2f ms",
                 preprocess_time * 1000,
                 exec_total_time * 1000,
+                output_ready_wait_time * 1000,
                 postprocess_time * 1000,
                 step_total_ms,
             )
@@ -478,6 +483,7 @@ class DiffusionEngine:
                 metrics_update = {
                     "preprocess_time_ms": preprocess_time * 1000,
                     "diffusion_engine_exec_time_ms": exec_total_time * 1000,
+                    "output_ready_wait_time_ms": output_ready_wait_time * 1000,
                     "postprocess_time_ms": postprocess_time * 1000,
                     **scheduler_metrics,
                 }

vllm_omni/diffusion/models/minimax_h3/pipeline_minimax_h3.py

@@ -683,6 +683,8 @@ class MiniMaxH3Pipeline(
         "_encode_reference_audio_conditions",
         "diffuse",
         "decode",
+        "video_vae.decode_latent",
+        "audio_vae.decode_latent",
         "prepare_encode",
         "denoise_step",
         "post_decode",

vllm_omni/metrics/stats.py

@@ -656,6 +656,7 @@ class OrchestratorAggregator:
     _MS_TO_S: dict[str, str] = {
         "preprocess_time_ms": "preprocess_time_s",
         "diffusion_engine_exec_time_ms": "diffusion_engine_exec_time_s",
+        "output_ready_wait_time_ms": "output_ready_wait_time_s",
         "postprocess_time_ms": "postprocess_time_s",
         "vae_decode_time_ms": "vae_decode_time_s",
         "forward_time_ms": "forward_time_s",