A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 06:30 UTC
Add PyTorch profiler schedule support with warmup/active iterations (#35240)
fenypatel99 committed
Mar 4, 2026 at 12:53 UTC
7eca85911072b9732293c3d4181e20a4c9394b21
2 files changed
+86
-3
vllm/config/profiler.py
+25
-2
@@ -45,8 +45,10 @@ class ProfilerConfig:
45
worker's traces (CPU & GPU) will be saved under this directory. Note that
46
it must be an absolute path."""
47
48
- torch_profiler_with_stack: bool = True
49
- """If `True`, enables stack tracing in the torch profiler. Enabled by default."""
48
+ torch_profiler_with_stack: bool = False
49
+ """If `True`, enables stack tracing in the torch profiler. Disabled by default
50
+ to reduce overhead. Can be enabled via VLLM_TORCH_PROFILER_WITH_STACK=1 env var
51
+ or --profiler-config.torch_profiler_with_stack=true CLI flag."""
52
53
torch_profiler_with_flops: bool = False
54
"""If `True`, enables FLOPS counting in the torch profiler. Disabled by default."""
@@ -81,6 +83,27 @@ class ProfilerConfig:
83
Defaults to 0, meaning no limit.
84
"""
85
86
+ warmup_iterations: int = Field(default=0, ge=0)
87
+ """Number of warmup iterations for PyTorch profiler schedule.
88
+ During warmup, the profiler runs but data is discarded. This helps reduce
89
+ noise from JIT compilation and other one-time costs in the profiled trace.
90
+ Defaults to 0 (schedule-based profiling disabled, recording all iterations).
91
+ Set to a positive value (e.g., 2) to enable schedule-based profiling.
92
+ """
93
+
94
+ active_iterations: int = Field(default=5, ge=1)
95
+ """Number of active iterations for PyTorch profiler schedule.
96
+ This is the number of iterations where profiling data is actually collected.
97
+ Defaults to 5 active iterations.
98
+ """
99
+
100
+ wait_iterations: int = Field(default=0, ge=0)
101
+ """Number of wait iterations for PyTorch profiler schedule.
102
+ During wait, the profiler is completely off with zero overhead.
103
+ This allows skipping initial iterations before warmup begins.
104
+ Defaults to 0 (no wait period).
105
+ """
106
+
107
def compute_hash(self) -> str:
108
"""
109
WARNING: Whenever a new field is added to this config,
vllm/profiler/wrapper.py
+61
-1
@@ -96,7 +96,9 @@ class WorkerProfiler(ABC):
96
logger.info_once("Starting profiler after delay...", scope="local")
97
self._call_start()
98
99
- if self._running:
99
+ # Call profiler step for schedule-based profiling
100
+ # Only count iterations where data is actually recorded (not warmup)
101
+ if self._running and self._profiler_step():
102
self._profiling_for_iters += 1
103
104
if (
@@ -113,6 +115,16 @@ class WorkerProfiler(ABC):
115
self._call_stop()
116
return
117
118
+ def _profiler_step(self) -> bool:
119
+ """Called each step when profiler is running.
120
+ Override in subclasses to handle schedule-based profiling.
121
+
122
+ Returns:
123
+ True if the step was an active profiling step (data recorded),
124
+ False if the step was a warmup step (data discarded).
125
+ """
126
+ return True
127
+
128
def stop(self) -> None:
129
"""Attempt to stop the profiler, accounting for overlapped calls."""
130
if not self._active:
@@ -187,8 +199,29 @@ class TorchProfilerWrapper(WorkerProfiler):
199
)
200
201
self.dump_cpu_time_total = "CPU" in activities and len(activities) == 1
202
+
203
+ # Create profiler schedule if warmup or wait iterations are configured
204
+ profiler_schedule = None
205
+ if profiler_config.warmup_iterations > 0 or profiler_config.wait_iterations > 0:
206
+ profiler_schedule = torch.profiler.schedule(
207
+ skip_first=0,
208
+ wait=profiler_config.wait_iterations,
209
+ warmup=profiler_config.warmup_iterations,
210
+ active=profiler_config.active_iterations,
211
+ repeat=1,
212
+ )
213
+ if local_rank in (None, 0):
214
+ logger.info_once(
215
+ "Profiler schedule configured: wait=%d, warmup=%d, active=%d",
216
+ profiler_config.wait_iterations,
217
+ profiler_config.warmup_iterations,
218
+ profiler_config.active_iterations,
219
+ scope="local",
220
+ )
221
+
222
self.profiler = torch.profiler.profile(
223
activities=[TorchProfilerActivityMap[activity] for activity in activities],
224
+ schedule=profiler_schedule,
225
record_shapes=profiler_config.torch_profiler_record_shapes,
226
profile_memory=profiler_config.torch_profiler_with_memory,
227
with_stack=profiler_config.torch_profiler_with_stack,
@@ -196,6 +229,17 @@ class TorchProfilerWrapper(WorkerProfiler):
229
on_trace_ready=trace_handler,
230
)
231
232
+ # Track if we're using a schedule (need to call step())
233
+ self._uses_schedule = profiler_schedule is not None
234
+ self._warmup_iterations = profiler_config.warmup_iterations
235
+ # Subtract 1 because profiler.start() already consumes step 0
236
+ # (WAIT or WARMUP), so only wait + warmup - 1 non-active steps
237
+ # remain to be advanced through via profiler.step() calls.
238
+ self._warmup_steps_remaining = max(
239
+ profiler_config.wait_iterations + profiler_config.warmup_iterations - 1,
240
+ 0,
241
+ )
242
+
243
@override
244
def _start(self) -> None:
245
self.profiler.start()
@@ -228,6 +272,22 @@ class TorchProfilerWrapper(WorkerProfiler):
272
)
273
)
274
275
+ @override
276
+ def _profiler_step(self) -> bool:
277
+ """Call profiler.step() when using schedule-based profiling.
278
+
279
+ Returns:
280
+ True if the step was an active profiling step (data recorded),
281
+ False if the step was a warmup step (data discarded).
282
+ """
283
+ if self._uses_schedule:
284
+ self.profiler.step()
285
+ # Track warmup steps - only count active steps toward max_iterations
286
+ if self._warmup_steps_remaining > 0:
287
+ self._warmup_steps_remaining -= 1
288
+ return False
289
+ return True
290
+
291
@override
292
def annotate_context_manager(self, name: str):
293
return torch.profiler.record_function(name)