[Profiler] Add SGLANG_PROFILE_BY_STAGE_DECODE_MIN_BS to defer the decode-stage capture (#38067)

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"
     )

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: