diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 7598f2773..9e38409a0 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -433,6 +433,10 @@ class Envs: SGLANG_PROFILE_WITH_STACK = EnvBool(True) SGLANG_PROFILE_RECORD_SHAPES = EnvBool(True) SGLANG_PROFILE_V2 = EnvBool(False) + # profile_by_stage: do not start the decode-stage capture until a decode batch + # reaches this many requests (0 = first decode batch). Lets a batch-size bench + # capture steady-state full-admission decode steps instead of the ramp-up. + SGLANG_PROFILE_BY_STAGE_DECODE_MIN_BS = EnvInt(0) SGLANG_ENABLE_NVTX_SCHEDULER = EnvBoolWithAlias( False, deprecated_name="SGLANG_ENABLE_NVTX" ) diff --git a/python/sglang/srt/managers/scheduler_components/profiler_manager.py b/python/sglang/srt/managers/scheduler_components/profiler_manager.py index 63aa5aac3..59462aa64 100644 --- a/python/sglang/srt/managers/scheduler_components/profiler_manager.py +++ b/python/sglang/srt/managers/scheduler_components/profiler_manager.py @@ -421,8 +421,12 @@ class SchedulerProfilerManager: elif batch.forward_mode.is_decode(): if self.profiler_decode_ct == 0: if self.profile_in_progress: - # force trace flush + # force trace flush (a prefill capture must not absorb decode steps) self._stop_profile(stage=ForwardMode.EXTEND) + min_bs = envs.SGLANG_PROFILE_BY_STAGE_DECODE_MIN_BS.get() + if min_bs > 0 and batch.batch_size() < min_bs: + # Wait for full admission before capturing the decode stage + return self._start_profile(batch.forward_mode) self.profiler_decode_ct += 1 if self.profiler_decode_ct > self.profiler_target_decode_ct: