[diffusion] fix: fix stages not logged when perf_dump_path is provided (#16016)

python/sglang/multimodal_gen/configs/sample/sampling_params.py

@@ -16,7 +16,6 @@ from dataclasses import dataclass
 from enum import Enum, auto
 from typing import Any
 
-from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.multimodal_gen.utils import StoreBoolean
 
@@ -208,10 +207,6 @@ class SamplingParams:
         if env_steps is not None and self.num_inference_steps is not None:
             self.num_inference_steps = int(env_steps)
 
-        # Auto-enable stage logging if dump path is provided
-        if self.perf_dump_path:
-            envs.SGLANG_DIFFUSION_STAGE_LOGGING = True
-
     def _validate(self):
         """
         check if the sampling params is correct by itself

python/sglang/multimodal_gen/runtime/managers/gpu_worker.py

@@ -9,6 +9,7 @@ from typing import List
 import torch
 from setproctitle import setproctitle
 
+from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.runtime.distributed import (
     get_sp_group,
     maybe_init_distributed_environment_and_model_parallel,
@@ -31,10 +32,7 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import (
     globally_suppress_loggers,
     init_logger,
 )
-from sglang.multimodal_gen.runtime.utils.perf_logger import (
-    PerformanceLogger,
-    StageProfiler,
-)
+from sglang.multimodal_gen.runtime.utils.perf_logger import PerformanceLogger
 
 logger = init_logger(__name__)
 
@@ -129,7 +127,8 @@ class GPUWorker:
             duration_ms = (time.monotonic() - start_time) * 1000
             output_batch.timings.total_duration_ms = duration_ms
 
-            if StageProfiler.metrics_enabled():
+            # TODO: extract to avoid duplication
+            if req.perf_dump_path is not None or envs.SGLANG_DIFFUSION_STAGE_LOGGING:
                 PerformanceLogger.log_request_summary(timings=output_batch.timings)
         except Exception as e:
             logger.error(

python/sglang/multimodal_gen/runtime/pipelines_core/stages/base.py

@@ -201,7 +201,12 @@ class PipelineStage(ABC):
             raise
 
         # Execute the actual stage logic with unified profiling
-        with StageProfiler(stage_name, logger=logger, timings=batch.timings):
+        with StageProfiler(
+            stage_name,
+            logger=logger,
+            timings=batch.timings,
+            perf_dump_path_provided=batch.perf_dump_path is not None,
+        ):
             result = self.forward(batch, server_args)
 
         # Post-execution output verification

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -959,7 +959,10 @@ class DenoisingStage(PipelineStage):
             with self.progress_bar(total=num_inference_steps) as progress_bar:
                 for i, t_host in enumerate(timesteps_cpu):
                     with StageProfiler(
-                        f"denoising_step_{i}", logger=logger, timings=batch.timings
+                        f"denoising_step_{i}",
+                        logger=logger,
+                        timings=batch.timings,
+                        perf_dump_path_provided=batch.perf_dump_path is not None,
                     ):
                         t_int = int(t_host.item())
                         t_device = timesteps[i]

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_dmd.py

@@ -98,7 +98,10 @@ class DmdDenoisingStage(DenoisingStage):
                     continue
 
                 with StageProfiler(
-                    f"denoising_step_{i}", logger=logger, timings=batch.timings
+                    f"denoising_step_{i}",
+                    logger=logger,
+                    timings=batch.timings,
+                    perf_dump_path_provided=batch.perf_dump_path is not None,
                 ):
                     # Expand latents for I2V
                     noise_latents = latents.clone()

python/sglang/multimodal_gen/runtime/utils/perf_logger.py

@@ -131,31 +131,26 @@ class StageProfiler:
         logger: _SGLDiffusionLogger,
         timings: Optional["RequestTimings"],
         simple_log: bool = False,
+        perf_dump_path_provided: bool = False,
     ):
         self.stage_name = stage_name
         self.timings = timings
         self.logger = logger
         self.simple_log = simple_log
         self.start_time = 0.0
-
-        self._metrics_enabled = StageProfiler.metrics_enabled()
-
-    @staticmethod
-    def metrics_enabled():
-        # Check env var at runtime to ensure we pick up changes (e.g. from CLI args)
-        return envs.SGLANG_DIFFUSION_STAGE_LOGGING
+        self.enabled = perf_dump_path_provided or envs.SGLANG_DIFFUSION_STAGE_LOGGING
 
     def __enter__(self):
         if self.simple_log:
             self.logger.info(f"[{self.stage_name}] started...")
 
-        if (self._metrics_enabled and self.timings) or self.simple_log:
+        if (self.enabled and self.timings) or self.simple_log:
             self.start_time = time.perf_counter()
 
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if not ((self._metrics_enabled and self.timings) or self.simple_log):
+        if not ((self.enabled and self.timings) or self.simple_log):
             return False
 
         execution_time_s = time.perf_counter() - self.start_time
@@ -175,7 +170,7 @@ class StageProfiler:
                 f"[{self.stage_name}] finished in {execution_time_s:.4f} seconds",
             )
 
-        if self._metrics_enabled and self.timings:
+        if self.enabled and self.timings:
             if "denoising_step_" in self.stage_name:
                 index = int(self.stage_name[len("denoising_step_") :])
                 self.timings.record_steps(index, execution_time_s)