[diffusion] profile: early exit when enough steps are captured to reduce the size of the trace file (#14803)

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

@@ -219,6 +219,17 @@ class InputValidationStage(PipelineStage):
                 batch, server_args, condition_image_width, condition_image_height
             )
 
+        # if height or width is not specified at this point, set default to 720p
+        default_height = 720
+        default_width = 1080
+        if batch.height is None and batch.width is None:
+            batch.height = default_height
+            batch.width = default_width
+        elif batch.height is None:
+            batch.height = batch.width * default_height // default_width
+        elif batch.width is None:
+            batch.width = batch.height * default_width // default_height
+
         return batch
 
     def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
@@ -251,6 +262,7 @@ class InputValidationStage(PipelineStage):
         result.add_check("height", batch.height, V.positive_int)
         result.add_check("width", batch.width, V.positive_int)
         # Validate height and width
+
         if batch.height % 8 != 0 or batch.width % 8 != 0:
             raise ValueError(
                 f"Height and width must be divisible by 8 but are {batch.height} and {batch.width}."

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

@@ -56,14 +56,14 @@ class SGLDiffusionProfiler:
             # profile denoising stage only
             warmup = 1
             num_actual_steps = num_inference_steps if num_steps == -1 else num_steps
-            num_active_steps = num_actual_steps + warmup
+            self.num_active_steps = num_actual_steps + warmup
             self.profiler = torch.profiler.profile(
                 **common_torch_profiler_args,
                 schedule=torch.profiler.schedule(
                     skip_first=0,
                     wait=0,
                     warmup=warmup,
-                    active=num_active_steps,
+                    active=self.num_active_steps,
                     repeat=1,
                 ),
             )
@@ -89,7 +89,12 @@ class SGLDiffusionProfiler:
 
     def step_denoising_step(self):
         if not self.full_profile:
-            self._step()
+            if self.num_active_steps >= 0:
+                self._step()
+                self.num_active_steps -= 1
+            else:
+                # early exit when enough steps are captured, to reduce the trace file size
+                self.stop(dump_rank=0)
 
     @classmethod
     def get_instance(cls) -> "SGLDiffusionProfiler":