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 12, 22:15 UTC
[Feat] Iteration-level profiling for Torch and CUDA profiler (#28987)
Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Benjamin Chislett committed
Nov 19, 2025 at 22:17 UTC
fcbcba6c70a3308705aa21adebb443bf9015b486
5 files changed
+435
-61
tests/v1/worker/test_gpu_profiler.py
new
+203
@@ -0,0 +1,203 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+
5
+import vllm.envs as envs
6
+from vllm.profiler.gpu_profiler import WorkerProfiler
7
+
8
+
9
+class ConcreteWorkerProfiler(WorkerProfiler):
10
+ """
11
+ A basic implementation of a worker profiler for testing purposes.
12
+ """
13
+
14
+ def __init__(self):
15
+ self.start_call_count = 0
16
+ self.stop_call_count = 0
17
+ self.should_fail_start = False
18
+ super().__init__()
19
+
20
+ def _start(self) -> None:
21
+ if self.should_fail_start:
22
+ raise RuntimeError("Simulated start failure")
23
+ self.start_call_count += 1
24
+
25
+ def _stop(self) -> None:
26
+ self.stop_call_count += 1
27
+
28
+
29
+@pytest.fixture(autouse=True)
30
+def reset_mocks():
31
+ """Fixture to reset mocks and env variables before each test."""
32
+ envs.VLLM_PROFILER_DELAY_ITERS = 0
33
+ envs.VLLM_PROFILER_MAX_ITERS = 0
34
+
35
+
36
+def test_immediate_start_stop():
37
+ """Test standard start without delay."""
38
+ profiler = ConcreteWorkerProfiler()
39
+
40
+ profiler.start()
41
+ assert profiler._running is True
42
+ assert profiler._active is True
43
+ assert profiler.start_call_count == 1
44
+
45
+ profiler.stop()
46
+ assert profiler._running is False
47
+ assert profiler._active is False
48
+ assert profiler.stop_call_count == 1
49
+
50
+
51
+def test_delayed_start():
52
+ """Test that profiler waits for N steps before actually starting."""
53
+ envs.VLLM_PROFILER_DELAY_ITERS = 2
54
+ profiler = ConcreteWorkerProfiler()
55
+
56
+ # User requests start
57
+ profiler.start()
58
+
59
+ # Should be active (request accepted) but not running (waiting for delay)
60
+ assert profiler._active is True
61
+ assert profiler._running is False
62
+ assert profiler.start_call_count == 0
63
+
64
+ # Step 1
65
+ profiler.step()
66
+ assert profiler._running is False
67
+
68
+ # Step 2 (Threshold reached)
69
+ profiler.step()
70
+ assert profiler._running is True
71
+ assert profiler.start_call_count == 1
72
+
73
+
74
+def test_max_iterations():
75
+ """Test that profiler stops automatically after max iterations."""
76
+ envs.VLLM_PROFILER_MAX_ITERS = 2
77
+ profiler = ConcreteWorkerProfiler()
78
+
79
+ profiler.start()
80
+ assert profiler._running is True
81
+
82
+ # Iteration 1
83
+ profiler.step() # profiling_count becomes 1
84
+ assert profiler._running is True
85
+
86
+ # Iteration 2
87
+ profiler.step() # profiling_count becomes 2
88
+ assert profiler._running is True
89
+
90
+ # Iteration 3 (Exceeds max)
91
+ profiler.step() # profiling_count becomes 3
92
+
93
+ # Should have stopped now
94
+ assert profiler._running is False
95
+ assert profiler.stop_call_count == 1
96
+
97
+
98
+def test_delayed_start_and_max_iters():
99
+ """Test combined delayed start and max iterations."""
100
+ envs.VLLM_PROFILER_DELAY_ITERS = 2
101
+ envs.VLLM_PROFILER_MAX_ITERS = 2
102
+ profiler = ConcreteWorkerProfiler()
103
+
104
+ profiler.start()
105
+
106
+ # Step 1
107
+ profiler.step()
108
+ assert profiler._running is False
109
+ assert profiler._active is True
110
+
111
+ # Step 2 (Starts now)
112
+ profiler.step()
113
+ assert profiler._profiling_for_iters == 1
114
+ assert profiler._running is True
115
+ assert profiler._active is True
116
+
117
+ # Next iteration
118
+ profiler.step()
119
+ assert profiler._profiling_for_iters == 2
120
+ assert profiler._running is True
121
+
122
+ # Iteration 2 (exceeds max)
123
+ profiler.step()
124
+
125
+ # Should have stopped now
126
+ assert profiler._running is False
127
+ assert profiler.stop_call_count == 1
128
+
129
+
130
+def test_idempotency():
131
+ """Test that calling start/stop multiple times doesn't break logic."""
132
+ profiler = ConcreteWorkerProfiler()
133
+
134
+ # Double Start
135
+ profiler.start()
136
+ profiler.start()
137
+ assert profiler.start_call_count == 1 # Should only start once
138
+
139
+ # Double Stop
140
+ profiler.stop()
141
+ profiler.stop()
142
+ assert profiler.stop_call_count == 1 # Should only stop once
143
+
144
+
145
+def test_step_inactive():
146
+ """Test that stepping while inactive does nothing."""
147
+ envs.VLLM_PROFILER_DELAY_ITERS = 2
148
+ profiler = ConcreteWorkerProfiler()
149
+
150
+ # Not started yet
151
+ profiler.step()
152
+ profiler.step()
153
+
154
+ # Even though we stepped 2 times, start shouldn't happen because active=False
155
+ assert profiler.start_call_count == 0
156
+
157
+
158
+def test_start_failure():
159
+ """Test behavior when the underlying _start method raises exception."""
160
+ profiler = ConcreteWorkerProfiler()
161
+ profiler.should_fail_start = True
162
+
163
+ profiler.start()
164
+
165
+ # Exception caught in _call_start
166
+ assert profiler._running is False # Should not mark as running
167
+ assert profiler._active is True # Request is still considered active
168
+ assert profiler.start_call_count == 0 # Logic failed inside start
169
+
170
+
171
+def test_shutdown():
172
+ """Test that shutdown calls stop only if running."""
173
+ profiler = ConcreteWorkerProfiler()
174
+
175
+ # Case 1: Not running
176
+ profiler.shutdown()
177
+ assert profiler.stop_call_count == 0
178
+
179
+ # Case 2: Running
180
+ profiler.start()
181
+ profiler.shutdown()
182
+ assert profiler.stop_call_count == 1
183
+
184
+
185
+def test_mixed_delay_and_stop():
186
+ """Test manual stop during the delay period."""
187
+ envs.VLLM_PROFILER_DELAY_ITERS = 5
188
+ profiler = ConcreteWorkerProfiler()
189
+
190
+ profiler.start()
191
+ profiler.step()
192
+ profiler.step()
193
+
194
+ # User cancels before delay finishes
195
+ profiler.stop()
196
+ assert profiler._active is False
197
+
198
+ # Further steps should not trigger start
199
+ profiler.step()
200
+ profiler.step()
201
+ profiler.step()
202
+
203
+ assert profiler.start_call_count == 0
vllm/envs.py
+16
@@ -92,11 +92,14 @@ if TYPE_CHECKING:
92
VLLM_TORCH_PROFILER_DIR: str | None = None
93
VLLM_TORCH_PROFILER_RECORD_SHAPES: bool = False
94
VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY: bool = False
95
+ VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM: bool = False
96
VLLM_USE_AOT_COMPILE: bool = False
97
VLLM_USE_BYTECODE_HOOK: bool = False
98
VLLM_FORCE_AOT_LOAD: bool = False
99
VLLM_TORCH_PROFILER_WITH_STACK: bool = True
100
VLLM_TORCH_PROFILER_WITH_FLOPS: bool = False
101
+ VLLM_PROFILER_DELAY_ITERS: int = 0
102
+ VLLM_PROFILER_MAX_ITERS: int = 0
103
VLLM_USE_TRITON_AWQ: bool = False
104
VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
105
VLLM_SKIP_P2P_CHECK: bool = False
@@ -872,6 +875,19 @@ environment_variables: dict[str, Callable[[], Any]] = {
875
"VLLM_TORCH_PROFILER_WITH_FLOPS": lambda: bool(
876
os.getenv("VLLM_TORCH_PROFILER_WITH_FLOPS", "0") != "0"
877
),
878
+ # Disable torch profiling of the AsyncLLMEngine process.
879
+ # If set to 1, will not profile the engine process.
880
+ "VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM": lambda: bool(
881
+ os.getenv("VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM", "0") != "0"
882
+ ),
883
+ # Delay number of iterations before starting profiling when using
884
+ # the torch/torch CUDA profiler. If set to 0, will start profiling immediately.
885
+ "VLLM_PROFILER_DELAY_ITERS": lambda: int(
886
+ os.getenv("VLLM_PROFILER_DELAY_ITERS", "0")
887
+ ),
888
+ # Maximum number of iterations to profile when using the torch/torch CUDA profiler.
889
+ # If set to 0, will not limit the number of iterations.
890
+ "VLLM_PROFILER_MAX_ITERS": lambda: int(os.getenv("VLLM_PROFILER_MAX_ITERS", "0")),
891
# If set, vLLM will use Triton implementations of AWQ.
892
"VLLM_USE_TRITON_AWQ": lambda: bool(int(os.getenv("VLLM_USE_TRITON_AWQ", "0"))),
893
# If set, allow loading or unloading lora adapters in runtime,
vllm/profiler/gpu_profiler.py
+194
-19
@@ -1,37 +1,212 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from abc import ABC, abstractmethod
5
+from contextlib import nullcontext
6
+
7
+import torch
8
+from typing_extensions import override
9
+
10
+import vllm.envs as envs
11
from vllm.logger import init_logger
12
13
logger = init_logger(__name__)
14
15
9
-class CudaProfilerWrapper:
16
+class WorkerProfiler(ABC):
17
def __init__(self) -> None:
11
- self._profiler_running = False
12
- # Note: lazy import to avoid dependency issues if CUDA is not available.
13
- import torch.cuda.profiler as cuda_profiler
18
+ self._delay_iters = envs.VLLM_PROFILER_DELAY_ITERS
19
+ if self._delay_iters > 0:
20
+ logger.info_once(
21
+ "GPU profiling will start "
22
+ f"{self._delay_iters} steps after start_profile."
23
+ )
24
15
- self._cuda_profiler = cuda_profiler
25
+ self._max_iters = envs.VLLM_PROFILER_MAX_ITERS
26
+ if self._max_iters > 0:
27
+ logger.info_once(
28
+ "GPU profiling will stop "
29
+ f"after {self._max_iters} worker steps, "
30
+ "or when stop_profile is received."
31
+ )
32
17
- def start(self) -> None:
33
+ # Track when the profiler gets triggered by start_profile
34
+ self._active_iteration_count = 0
35
+ self._active = False
36
+
37
+ # Track when the profiler is actually running
38
+ self._profiling_for_iters = 0
39
+ self._running = False
40
+
41
+ @abstractmethod
42
+ def _start(self) -> None:
43
+ """Start the profiler."""
44
+ pass
45
+
46
+ @abstractmethod
47
+ def _stop(self) -> None:
48
+ """Stop the profiler."""
49
+ pass
50
+
51
+ def _call_start(self) -> None:
52
+ """Call _start with error handling but no safeguards."""
53
try:
19
- self._cuda_profiler.start()
20
- self._profiler_running = True
21
- logger.info_once("Started CUDA profiler")
54
+ self._start()
55
+ self._running = True # Only mark as running if start succeeds
56
except Exception as e:
23
- logger.warning_once("Failed to start CUDA profiler: %s", e)
57
+ logger.warning("Failed to start profiler: %s", e)
58
+
59
+ def _call_stop(self) -> None:
60
+ """Call _stop with error handling but no safeguards."""
61
+ try:
62
+ self._stop()
63
+ logger.info("Profiler stopped successfully.")
64
+ except Exception as e:
65
+ logger.warning("Failed to stop profiler: %s", e)
66
+ self._running = False # Always mark as not running, assume stop worked
67
+
68
+ def start(self) -> None:
69
+ """Attempt to start the profiler, accounting for delayed starts."""
70
+ if self._active:
71
+ logger.debug(
72
+ "start_profile received when profiler is already active. "
73
+ "Ignoring request."
74
+ )
75
+ return
76
+ self._active = True
77
+ if self._delay_iters == 0:
78
+ self._call_start()
79
+
80
+ def step(self) -> None:
81
+ """Update the profiler state at each worker step,
82
+ to handle delayed starts and max iteration limits."""
83
+ if not self._active:
84
+ return
85
+
86
+ self._active_iteration_count += 1
87
+
88
+ if (
89
+ not self._running
90
+ and self._delay_iters > 0
91
+ and self._active_iteration_count == self._delay_iters
92
+ ):
93
+ logger.info("Starting profiler after delay...")
94
+ self._call_start()
95
+
96
+ if self._running:
97
+ self._profiling_for_iters += 1
98
+
99
+ if (
100
+ self._max_iters > 0
101
+ and self._running
102
+ and self._profiling_for_iters > self._max_iters
103
+ ):
104
+ # Automatically stop the profiler after max iters
105
+ # will be marked as not running, but leave as active so that stop
106
+ # can clean up properly
107
+ logger.info("Max profiling iterations reached. Stopping profiler...")
108
+ self._call_stop()
109
+ return
110
111
def stop(self) -> None:
26
- if self._profiler_running:
27
- try:
28
- self._cuda_profiler.stop()
29
- logger.info_once("Stopped CUDA profiler")
30
- except Exception as e:
31
- logger.warning_once("Failed to stop CUDA profiler: %s", e)
32
- finally:
33
- self._profiler_running = False
112
+ """Attempt to stop the profiler, accounting for overlapped calls."""
113
+ if not self._active:
114
+ logger.debug(
115
+ "stop_profile received when profiler is not active. Ignoring request."
116
+ )
117
+ return
118
+ self._active = False
119
+ self._active_iteration_count = 0
120
+ self._profiling_for_iters = 0
121
+
122
+ if self._running:
123
+ self._call_stop()
124
125
def shutdown(self) -> None:
126
"""Ensure profiler is stopped when shutting down."""
37
- self.stop()
127
+ logger.info_once("Shutting down profiler")
128
+ if self._running:
129
+ self.stop()
130
+
131
+ def annotate_context_manager(self, name: str):
132
+ """Return a context manager to annotate profiler traces."""
133
+ return nullcontext()
134
+
135
+
136
+class TorchProfilerWrapper(WorkerProfiler):
137
+ def __init__(self, worker_name: str, local_rank: int) -> None:
138
+ super().__init__()
139
+
140
+ self.local_rank = local_rank
141
+ torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
142
+ logger.info(
143
+ "Torch profiling enabled. Traces will be saved to: %s",
144
+ torch_profiler_trace_dir,
145
+ )
146
+ logger.debug(
147
+ "Profiler config: record_shapes=%s,"
148
+ "profile_memory=%s,with_stack=%s,with_flops=%s",
149
+ envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
150
+ envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
151
+ envs.VLLM_TORCH_PROFILER_WITH_STACK,
152
+ envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
153
+ )
154
+ self.profiler = torch.profiler.profile(
155
+ activities=[
156
+ torch.profiler.ProfilerActivity.CPU,
157
+ torch.profiler.ProfilerActivity.CUDA,
158
+ ],
159
+ record_shapes=envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
160
+ profile_memory=envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
161
+ with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
162
+ with_flops=envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
163
+ on_trace_ready=torch.profiler.tensorboard_trace_handler(
164
+ torch_profiler_trace_dir, worker_name=worker_name, use_gzip=True
165
+ ),
166
+ )
167
+
168
+ @override
169
+ def _start(self) -> None:
170
+ self.profiler.start()
171
+
172
+ @override
173
+ def _stop(self) -> None:
174
+ self.profiler.stop()
175
+
176
+ rank = self.local_rank
177
+ profiler_dir = envs.VLLM_TORCH_PROFILER_DIR
178
+ profiler_out_file = f"{profiler_dir}/profiler_out_{rank}.txt"
179
+ sort_key = "self_cuda_time_total"
180
+ table = self.profiler.key_averages().table(sort_by=sort_key)
181
+
182
+ with open(profiler_out_file, "w") as f:
183
+ print(table, file=f)
184
+
185
+ # only print profiler results on rank 0
186
+ if rank == 0:
187
+ print(table)
188
+
189
+ @override
190
+ def annotate_context_manager(self, name: str):
191
+ return torch.profiler.record_function(name)
192
+
193
+
194
+class CudaProfilerWrapper(WorkerProfiler):
195
+ def __init__(self) -> None:
196
+ super().__init__()
197
+ # Note: lazy import to avoid dependency issues if CUDA is not available.
198
+ import torch.cuda.profiler as cuda_profiler
199
+
200
+ self._cuda_profiler = cuda_profiler
201
+
202
+ @override
203
+ def _start(self) -> None:
204
+ self._cuda_profiler.start()
205
+
206
+ @override
207
+ def _stop(self) -> None:
208
+ self._cuda_profiler.stop()
209
+
210
+ @override
211
+ def annotate_context_manager(self, name: str):
212
+ return torch.cuda.nvtx.range(name)
vllm/v1/engine/async_llm.py
+13
-1
@@ -160,11 +160,23 @@ class AsyncLLM(EngineClient):
160
except RuntimeError:
161
pass
162
163
- if envs.VLLM_TORCH_PROFILER_DIR:
163
+ if (
164
+ envs.VLLM_TORCH_PROFILER_DIR
165
+ and not envs.VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM
166
+ ):
167
logger.info(
168
"Torch profiler enabled. AsyncLLM CPU traces will be collected under %s", # noqa: E501
169
envs.VLLM_TORCH_PROFILER_DIR,
170
)
171
+ if envs.VLLM_PROFILER_MAX_ITERS > 0 or envs.VLLM_PROFILER_DELAY_ITERS > 0:
172
+ logger.warning_once(
173
+ "Torch profiler received max_iters or delay_iters setting. These "
174
+ "are not compatible with the AsyncLLM profiler and will be ignored "
175
+ "for the AsyncLLM process. Engine process profiling will still "
176
+ "respect these settings. Consider setting "
177
+ "VLLM_TORCH_PROFILER_DISABLE_ASYNC_LLM=1 to disable "
178
+ "AsyncLLM profiling."
179
+ )
180
worker_name = f"{socket.gethostname()}_{os.getpid()}.async_llm"
181
self.profiler = torch.profiler.profile(
182
activities=[
vllm/v1/worker/gpu_worker.py
+9
-41
@@ -36,7 +36,7 @@ from vllm.model_executor import set_random_seed
36
from vllm.model_executor.models.interfaces import is_mixture_of_experts
37
from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
38
from vllm.platforms import current_platform
39
-from vllm.profiler.gpu_profiler import CudaProfilerWrapper
39
+from vllm.profiler.gpu_profiler import CudaProfilerWrapper, TorchProfilerWrapper
40
from vllm.sequence import IntermediateTensors
41
from vllm.tasks import SupportedTask
42
from vllm.utils.mem_constants import GiB_bytes
@@ -90,32 +90,9 @@ class Worker(WorkerBase):
90
# Torch profiler. Enabled and configured through env vars:
91
# VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
92
if envs.VLLM_TORCH_PROFILER_DIR:
93
- torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
93
worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
95
- logger.info(
96
- "Profiling enabled. Traces will be saved to: %s",
97
- torch_profiler_trace_dir,
98
- )
99
- logger.debug(
100
- "Profiler config: record_shapes=%s,"
101
- "profile_memory=%s,with_stack=%s,with_flops=%s",
102
- envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
103
- envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
104
- envs.VLLM_TORCH_PROFILER_WITH_STACK,
105
- envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
106
- )
107
- self.profiler = torch.profiler.profile(
108
- activities=[
109
- torch.profiler.ProfilerActivity.CPU,
110
- torch.profiler.ProfilerActivity.CUDA,
111
- ],
112
- record_shapes=envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
113
- profile_memory=envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
114
- with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
115
- with_flops=envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
116
- on_trace_ready=torch.profiler.tensorboard_trace_handler(
117
- torch_profiler_trace_dir, worker_name=worker_name, use_gzip=True
118
- ),
94
+ self.profiler = TorchProfilerWrapper(
95
+ worker_name=worker_name, local_rank=self.local_rank
96
)
97
elif envs.VLLM_TORCH_CUDA_PROFILE:
98
self.profiler = CudaProfilerWrapper()
@@ -526,10 +503,12 @@ class Worker(WorkerBase):
503
if not self.profiler:
504
return nullcontext()
505
506
+ self.profiler.step()
507
+
508
num_new = len(scheduler_output.scheduled_new_reqs)
509
num_cached = len(scheduler_output.scheduled_cached_reqs.req_ids)
510
532
- return torch.profiler.record_function(
511
+ return self.profiler.annotate_context_manager(
512
f"execute_new_{num_new}_cached_{num_cached}"
513
)
514
@@ -587,24 +566,11 @@ class Worker(WorkerBase):
566
567
def profile(self, is_start: bool = True):
568
if self.profiler is None:
590
- raise RuntimeError("Profiler is not enabled.")
569
+ raise RuntimeError("Profiling is not enabled.")
570
if is_start:
571
self.profiler.start()
572
else:
573
self.profiler.stop()
595
- if isinstance(self.profiler, torch.profiler.profile):
596
- rank = self.local_rank
597
- profiler_dir = envs.VLLM_TORCH_PROFILER_DIR
598
- profiler_out_file = f"{profiler_dir}/profiler_out_{rank}.txt"
599
- sort_key = "self_cuda_time_total"
600
- table = self.profiler.key_averages().table(sort_by=sort_key)
601
-
602
- with open(profiler_out_file, "w") as f:
603
- print(table, file=f)
604
-
605
- # only print profiler results on rank 0
606
- if rank == 0:
607
- print(table)
574
575
def execute_dummy_batch(self) -> None:
576
self.model_runner._dummy_run(1, uniform_decode=True)
@@ -865,6 +831,8 @@ class Worker(WorkerBase):
831
def shutdown(self) -> None:
832
if runner := getattr(self, "model_runner", None):
833
runner.ensure_kv_transfer_shutdown()
834
+ if self.profiler is not None:
835
+ self.profiler.shutdown()
836
837
838
def init_worker_distributed_environment(