[Bugfix] Report stepwise Qwen-Image profiler metrics (#4517)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

tests/diffusion/test_diffusion_step_pipeline.py

@@ -30,6 +30,9 @@ from vllm_omni.diffusion.ipc import (
     pack_diffusion_output_shm,
     unpack_diffusion_output_shm,
 )
+from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import (
+    DiffusionPipelineProfilerMixin,
+)
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched import StepScheduler
 from vllm_omni.diffusion.sched.interface import (
@@ -97,6 +100,51 @@ class _StepPipeline:
         return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
 
 
+class _ProfilingStepPipeline(_StepPipeline):
+    enable_diffusion_pipeline_profiler = True
+
+    def __init__(self):
+        super().__init__()
+        self._stage_durations: dict[str, float] = {}
+
+    @property
+    def stage_durations(self) -> dict[str, float]:
+        return dict(self._stage_durations)
+
+    def clear_profiler_records(self) -> None:
+        self._stage_durations.clear()
+
+    def prepare_encode(self, state, **kwargs):
+        result = super().prepare_encode(state, **kwargs)
+        self._stage_durations["QwenImagePipeline.text_encoder.forward"] = 1.0
+        return result
+
+    def denoise_step(self, input_batch, **kwargs):
+        result = super().denoise_step(input_batch, **kwargs)
+        self._stage_durations["QwenImagePipeline.diffuse"] = 2.0
+        return result
+
+    def post_decode(self, state, **kwargs):
+        result = super().post_decode(state, **kwargs)
+        self._stage_durations["QwenImagePipeline.vae.decode"] = 3.0
+        return result
+
+
+class _AutoDenoiseProfilerPipeline(DiffusionPipelineProfilerMixin):
+    _PROFILER_TARGETS: list[str] = []
+
+    def __init__(self):
+        self.setup_diffusion_pipeline_profiler(
+            enable_diffusion_pipeline_profiler=True,
+        )
+
+    def forward(self):
+        return None
+
+    def denoise_step(self):
+        return "ok"
+
+
 class _InterruptingStepPipeline(_StepPipeline):
     interrupt = True
 
@@ -130,6 +178,9 @@ class _FakePeakMemoryPlatform:
         index = min(self.reset_calls - 1, len(self._reserved_mb) - 1)
         return int((self._reserved_mb[index] - 100) * 1024**2)
 
+    def is_available(self) -> bool:
+        return True
+
 
 class _IdentityNoiseTransformer(torch.nn.Module):
     def forward(self, x: torch.Tensor, **kwargs):
@@ -437,6 +488,16 @@ def test_input_batch_cached_repack_keeps_static_prompt_fields_for_same_compositi
     torch.testing.assert_close(repacked.prompt_embeds, torch.ones(1, 2, 3))
 
 
+@pytest.mark.cpu
+def test_step_profiler_reports_denoise_step_as_diffuse():
+    pipeline = _AutoDenoiseProfilerPipeline()
+
+    assert pipeline.denoise_step() == "ok"
+
+    assert any(key.endswith(".diffuse") for key in pipeline.stage_durations)
+    assert not any(key.endswith(".denoise_step") for key in pipeline.stage_durations)
+
+
 @pytest.mark.cpu
 class TestRunner:
     """DiffusionModelRunner.execute_stepwise"""
@@ -469,6 +530,53 @@ class TestRunner:
         assert runner.pipeline.scheduler_calls == 2
         assert runner.pipeline.decode_calls == 1
 
+    def test_stepwise_output_includes_stage_and_peak_metrics(self, monkeypatch):
+        runner = _make_runner()
+        runner.pipeline = _ProfilingStepPipeline()
+        req = _make_step_request()
+        reset_calls = []
+        monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "is_available",
+            lambda: True,
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "reset_peak_memory_stats",
+            lambda: reset_calls.append(True),
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "max_memory_reserved",
+            lambda: 2 * 1024**2,
+        )
+        monkeypatch.setattr(
+            model_runner_module.current_omni_platform,
+            "max_memory_allocated",
+            lambda: 1024**2,
+        )
+
+        DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_scheduler_output(req, step_id=0),
+        )
+        result = DiffusionModelRunner.execute_stepwise(
+            runner,
+            _make_cached_scheduler_output(step_id=1),
+        )
+
+        output = result.get_request_output("req-1")
+        assert output.finished is True
+        assert output.result is not None
+        assert output.result.peak_memory_mb == 2
+        assert output.result.stage_durations == {
+            "QwenImagePipeline.text_encoder.forward": 1.0,
+            "QwenImagePipeline.diffuse": 4.0,
+            "QwenImagePipeline.vae.decode": 3.0,
+        }
+        assert reset_calls == [True]
+
     def test_carries_peak_memory_across_stepwise_request_lifecycle(self, monkeypatch):
         runner = _make_runner()
         req = _make_step_request()

vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py

@@ -14,6 +14,7 @@ logger = init_logger(__name__)
 
 def profiler(name: str, func: Callable, instance: Any) -> Callable:
     """Timing a function execution."""
+    metric_name = f"{name.rsplit('.', 1)[0]}.diffuse" if name.endswith(".denoise_step") else name
 
     @functools.wraps(func)
     def wrapper(*args, **kwargs) -> Any:
@@ -28,10 +29,10 @@ def profiler(name: str, func: Callable, instance: Any) -> Callable:
             if current_omni_platform.is_available():
                 current_omni_platform.synchronize()
             duration = time.perf_counter() - start_time
-            logger.info(f"[DiffusionPipelineProfiler] {name} took {duration:.6f}s")
+            logger.info(f"[DiffusionPipelineProfiler] {metric_name} took {duration:.6f}s")
             # record the profiling data: duration of stages
             with instance._profiler_lock:
-                instance._stage_durations[name] = instance._stage_durations.get(name, 0.0) + duration
+                instance._stage_durations[metric_name] = instance._stage_durations.get(metric_name, 0.0) + duration
 
     return wrapper
 
@@ -93,7 +94,11 @@ class DiffusionPipelineProfilerMixin:
         self._stage_durations: dict[str, float] = {}
         targets = profiler_targets if profiler_targets is not None else self._PROFILER_TARGETS
         if not targets:
-            return
+            targets = []
+        else:
+            targets = list(targets)
+        if not profiler_targets and hasattr(self, "denoise_step"):
+            targets.append("denoise_step")
 
         targets = ["forward"] + [
             t for t in targets if t != "forward"

vllm_omni/diffusion/worker/diffusion_model_runner.py

@@ -39,7 +39,15 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.input_batch import InputBatch, scatter_latents
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, DiffusionRequestState, RunnerOutput
+from vllm_omni.diffusion.worker.utils import (
+    BatchRunnerOutput,
+    DiffusionRequestState,
+    RunnerOutput,
+    attach_stage_durations,
+    clear_pipeline_stage_durations,
+    consume_pipeline_stage_durations,
+    merge_stage_durations,
+)
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.platforms import current_omni_platform
 from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
@@ -487,6 +495,19 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
 
         return self._runner_output_from_outputs(reqs, outputs)
 
+    def _attach_stepwise_metrics(
+        self,
+        state: DiffusionRequestState,
+        output: DiffusionOutput,
+        *,
+        is_primary: bool,
+    ) -> None:
+        merge_stage_durations(
+            state,
+            consume_pipeline_stage_durations(self.pipeline),
+        )
+        attach_stage_durations(state, output)
+
     def execute_model(self, req: OmniDiffusionRequest, kv_prefetch_jobs: dict | None = None) -> DiffusionOutput:
         """
         Execute a forward pass for the given requests.
@@ -601,8 +622,13 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                     else:
                         gen_device = self.device
                     state.sampling.generator = torch.Generator(device=gen_device).manual_seed(state.sampling.seed)
+                clear_pipeline_stage_durations(self.pipeline)
                 # encode
                 self.pipeline.prepare_encode(state)
+                merge_stage_durations(
+                    state,
+                    consume_pipeline_stage_durations(self.pipeline),
+                )
 
         input_batch = InputBatch.make_batch(
             states,
@@ -658,19 +684,27 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
         use_hsdp = self.od_config.parallel_config.use_hsdp
         grad_context = torch.no_grad() if use_hsdp else torch.inference_mode()
         with grad_context:
+            had_active_states = bool(self.state_cache)
             states, new_request_ids = self._update_states(scheduler_output)
+            is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
+            if new_request_ids and not had_active_states and is_primary and current_omni_platform.is_available():
+                current_omni_platform.reset_peak_memory_stats()
             input_batch = self._prepare_batch_inputs(states, new_request_ids)
             attn_metadata = self._prepare_attn_metadata(input_batch)
-            is_primary = not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0
-            if is_primary:
-                current_omni_platform.reset_peak_memory_stats()
 
             with set_forward_context(
                 vllm_config=self.vllm_config,
                 omni_diffusion_config=self.od_config,
                 attn_metadata=attn_metadata,
             ):
+                clear_pipeline_stage_durations(self.pipeline)
                 noise_pred = self.pipeline.denoise_step(input_batch, states=states)
+                denoise_stage_durations = consume_pipeline_stage_durations(self.pipeline)
+                for state in states:
+                    merge_stage_durations(
+                        state,
+                        denoise_stage_durations,
+                    )
 
                 runner_output_list = []
                 pipeline_interrupted = getattr(self.pipeline, "interrupt", False)
@@ -693,14 +727,20 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                             self.pipeline.step_scheduler(
                                 req, noise_pred[offset : offset + row_num] if noise_pred is not None else None
                             )
-                            offset = offset + row_num
                             if self.od_config.streaming_output:
                                 should_decode = req.chunk_denoise_completed
                             else:
                                 should_decode = req.denoise_completed
 
                             if should_decode:
+                                clear_pipeline_stage_durations(self.pipeline)
                                 result = self.pipeline.post_decode(req)
+                                if result is not None:
+                                    self._attach_stepwise_metrics(
+                                        req,
+                                        result,
+                                        is_primary=is_primary,
+                                    )
                             else:
                                 result = None
                             # finished should be computed after post_decode() advanced chunk_index
@@ -717,6 +757,7 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
                                     result=result,
                                 )
                             )
+                            offset = offset + row_num
                         except Exception as per_req_exc:
                             offset = offset + row_num
                             logger.error(

vllm_omni/diffusion/worker/utils.py

@@ -15,6 +15,41 @@ if TYPE_CHECKING:
     from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
 
 
+def clear_pipeline_stage_durations(pipeline: Any) -> None:
+    clear_records = getattr(pipeline, "clear_profiler_records", None)
+    if getattr(pipeline, "enable_diffusion_pipeline_profiler", False) and callable(clear_records):
+        clear_records()
+
+
+def consume_pipeline_stage_durations(pipeline: Any) -> dict[str, float]:
+    if not getattr(pipeline, "enable_diffusion_pipeline_profiler", False):
+        return {}
+    stage_durations = getattr(pipeline, "stage_durations", None)
+    if not isinstance(stage_durations, dict):
+        return {}
+    result = {stage: float(duration) for stage, duration in stage_durations.items()}
+    clear_pipeline_stage_durations(pipeline)
+    return result
+
+
+def merge_stage_durations(
+    state: DiffusionRequestState,
+    stage_durations: dict[str, float],
+) -> None:
+    if not stage_durations:
+        return
+    for stage, duration in stage_durations.items():
+        state.stage_durations[stage] = float(state.stage_durations.get(stage, 0.0)) + float(duration)
+
+
+def attach_stage_durations(
+    state: DiffusionRequestState,
+    output: DiffusionOutput,
+) -> None:
+    if state.stage_durations:
+        output.stage_durations = dict(state.stage_durations)
+
+
 @dataclass
 class DiffusionRequestState:
     """Per-request mutable state across all pipeline stages.
@@ -79,6 +114,9 @@ class DiffusionRequestState:
     # For example: Wan condition tensors / masks, or Bagel KV contexts.
     extra: dict[str, Any] = field(default_factory=dict)
 
+    # ── Runner-owned profiling metadata ──
+    stage_durations: dict[str, float] = field(default_factory=dict)
+
     # Peak device memory observed while this request is active in step mode.
     peak_memory_mb: float = 0.0