chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2025-12-10 16:11:22 +08:00
committed by GitHub
parent 56e5c07424
commit 87dbdddc93
2 changed files with 20 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
python/sglang/multimodal_gen/runtime/utils/profiler.py
View File

@@ -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":