diff --git a/tests/diffusion/test_diffusion_step_pipeline.py b/tests/diffusion/test_diffusion_step_pipeline.py index bd2289283..7e18996ed 100644 --- a/tests/diffusion/test_diffusion_step_pipeline.py +++ b/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() diff --git a/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py b/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py index 8d4bf41de..5a293b905 100644 --- a/vllm_omni/diffusion/profiler/diffusion_pipeline_profiler.py +++ b/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" diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 84eb32044..cb16bdeac 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/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( diff --git a/vllm_omni/diffusion/worker/utils.py b/vllm_omni/diffusion/worker/utils.py index 4fd374d34..ba6165a9d 100644 --- a/vllm_omni/diffusion/worker/utils.py +++ b/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