From 2278a321ca9d04b44c04cff945d14775db18b913 Mon Sep 17 00:00:00 2001 From: Mick Date: Fri, 3 Apr 2026 01:16:38 +0800 Subject: [PATCH] [diffusion] chore: fix stage profiler for multi-stage denoising (#21955) --- .../pipelines_core/stages/denoising.py | 1 + .../pipelines_core/stages/denoising_av.py | 1 + .../pipelines_core/stages/denoising_dmd.py | 1 + .../stages/model_specific_stages/mova.py | 1 + .../runtime/utils/perf_logger.py | 19 +++++++++++-------- 5 files changed, 15 insertions(+), 8 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 2611c2733..2f348d68c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py @@ -1025,6 +1025,7 @@ class DenoisingStage(PipelineStage): logger=logger, metrics=batch.metrics, perf_dump_path_provided=batch.perf_dump_path is not None, + record_as_step=True, ): t_int = int(t_host.item()) t_device = timesteps[i] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py index 55deb10f4..7bd5a7e95 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py @@ -407,6 +407,7 @@ class LTX2AVDenoisingStage(DenoisingStage): logger=logger, metrics=batch.metrics, perf_dump_path_provided=batch.perf_dump_path is not None, + record_as_step=True, ): t_int = int(t_host.item()) t_device = timesteps[i] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py index 504fc429e..1b0223d51 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py @@ -104,6 +104,7 @@ class DmdDenoisingStage(DenoisingStage): logger=logger, metrics=batch.metrics, perf_dump_path_provided=batch.perf_dump_path is not None, + record_as_step=True, ): t_int = int(t.item()) if self.transformer_2 is not None: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py index afd43238e..e05048e41 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/mova.py @@ -430,6 +430,7 @@ class MOVADenoisingStage(PipelineStage): logger=logger, metrics=metrics, perf_dump_path_provided=perf_dump_path_provided, + record_as_step=True, ): pair_t = paired_timesteps[idx_step] if getattr(pair_t, "shape", None) == (2,): diff --git a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py index d44a50d37..9b1eb68cf 100644 --- a/python/sglang/multimodal_gen/runtime/utils/perf_logger.py +++ b/python/sglang/multimodal_gen/runtime/utils/perf_logger.py @@ -64,9 +64,8 @@ class RequestMetrics: """Records the duration of a pipeline stage""" self.stages[stage_name] = duration_s * 1000 # Store as milliseconds - def record_steps(self, index: int, duration_s: float): - """Records the duration of a denoising step""" - assert index == len(self.steps) + def record_step(self, duration_s: float): + """Records the duration of a denoising step in execution order.""" self.steps.append(duration_s * 1000) def record_memory_snapshot(self, checkpoint_name: str, snapshot: MemorySnapshot): @@ -192,6 +191,7 @@ class StageProfiler: log_stage_start_end: bool = False, perf_dump_path_provided: bool = False, capture_memory: bool = False, + record_as_step: bool = False, ): self.stage_name = stage_name self.metrics = metrics @@ -200,6 +200,10 @@ class StageProfiler: self.log_timing = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING self.log_stage_start_end = log_stage_start_end self.capture_memory = capture_memory + self.record_as_step = record_as_step + + def _should_record_as_step(self) -> bool: + return self.record_as_step or self.stage_name.startswith("denoising_step_") def __enter__(self): if self.log_stage_start_end: @@ -211,7 +215,7 @@ class StageProfiler: if (self.log_timing and self.metrics) or self.log_stage_start_end: if ( os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1" - and self.stage_name.startswith("denoising_step_") + and self._should_record_as_step() and torch.get_device_module().is_available() ): torch.get_device_module().synchronize() @@ -225,7 +229,7 @@ class StageProfiler: if ( os.environ.get("SGLANG_DIFFUSION_SYNC_STAGE_PROFILING", "0") == "1" - and self.stage_name.startswith("denoising_step_") + and self._should_record_as_step() and torch.get_device_module().is_available() ): torch.get_device_module().synchronize() @@ -247,9 +251,8 @@ class StageProfiler: ) if self.log_timing and self.metrics: - if "denoising_step_" in self.stage_name: - index = int(self.stage_name[len("denoising_step_") :]) - self.metrics.record_steps(index, execution_time_s) + if self._should_record_as_step(): + self.metrics.record_step(execution_time_s) else: self.metrics.record_stage(self.stage_name, execution_time_s)