@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 3918f3c5a3

[Profiler] Unify platform-aware torch profiling (#57460)

czhu-cohere committed Sep 20, 2026 at 19:17 UTC 3918f3c5a3e389fb17d9c0d629811f3745014db0
9 files changed +408 -110
docs/contributing/profiling.md
+4
@@ -22,6 +22,10 @@ To use the `torch.profiler` module, set the `profiler` entry to `'torch'` and `t
22 - `torch_profiler_with_flops` to enable recording FLOPs, off by default
23 - `torch_profiler_use_gzip` to control gzip-compressing profiling files, on by default
24 - `torch_profiler_dump_cuda_time_total` to control dumping and printing the aggregated CUDA self time table, on by default
25 +- `torch_profiler_activities` to select worker activities. Defaults are
26 + platform-specific: `["CPU"]` on CPU, `["CPU", "CUDA"]` on NVIDIA GPUs, and
27 + `["CPU", "XPU"]` on Intel GPUs. Selecting `["CUDA"]` omits CPU annotations
28 + and the AsyncLLM CPU trace to reduce profiling overhead and trace size.
29
30 When using `vllm bench serve`, you can enable profiling by passing the `--profile` flag.
31
tests/v1/engine/test_async_llm.py
+30 -7
@@ -60,13 +60,7 @@ VISION_PROMPT = {
60 }
61
62
63 -def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch):
64 - vllm_config = MagicMock()
65 - vllm_config.observability_config.otlp_traces_endpoint = None
66 - vllm_config.scheduler_config.stream_interval = 1
67 - vllm_config.profiler_config.profiler = "cuda"
68 - vllm_config.profiler_config.ignore_frontend = False
69 -
63 +def _mock_async_llm_dependencies(monkeypatch: pytest.MonkeyPatch):
64 renderer = MagicMock()
65 engine_core = MagicMock()
66 engine_core.profile_async = AsyncMock()
@@ -90,6 +84,16 @@ def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatc
84 "make_async_mp_client",
85 MagicMock(return_value=engine_core),
86 )
87 + return engine_core
88 +
89 +
90 +def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatch):
91 + vllm_config = MagicMock()
92 + vllm_config.observability_config.otlp_traces_endpoint = None
93 + vllm_config.scheduler_config.stream_interval = 1
94 + vllm_config.profiler_config.profiler = "cuda"
95 + vllm_config.profiler_config.ignore_frontend = False
96 + engine_core = _mock_async_llm_dependencies(monkeypatch)
97
98 engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False)
99
@@ -103,6 +107,25 @@ def test_cuda_profiler_requests_reach_engine_core(monkeypatch: pytest.MonkeyPatc
107 engine_core.profile_async.assert_has_awaits([call(True, None), call(False)])
108
109
110 +def test_cuda_only_torch_profiler_skips_frontend_cpu_trace(
111 + monkeypatch: pytest.MonkeyPatch,
112 +):
113 + vllm_config = MagicMock()
114 + vllm_config.observability_config.otlp_traces_endpoint = None
115 + vllm_config.scheduler_config.stream_interval = 1
116 + vllm_config.profiler_config.profiler = "torch"
117 + vllm_config.profiler_config.ignore_frontend = False
118 + vllm_config.profiler_config.torch_profiler_activities = ["CUDA"]
119 + _mock_async_llm_dependencies(monkeypatch)
120 + profiler = MagicMock()
121 + monkeypatch.setattr(async_llm_module, "TorchProfilerWrapper", profiler)
122 +
123 + engine = AsyncLLM(vllm_config, MagicMock(), log_stats=False)
124 +
125 + assert engine.profiler is None
126 + profiler.assert_not_called()
127 +
128 +
129 async def generate(
130 engine: AsyncLLM,
131 request_id: str,
tests/v1/worker/test_gpu_profiler.py
+231 -10
@@ -7,6 +7,7 @@ from unittest.mock import MagicMock, Mock, call, patch
7 from uuid import UUID
8
9 import pytest
10 +import torch
11 from pydantic import ValidationError
12
13 from vllm.config import (
@@ -17,10 +18,17 @@ from vllm.config import (
18 )
19 from vllm.config.profiler import _is_uri_path
20 from vllm.platforms import current_platform
20 -from vllm.profiler.wrapper import ProtonProfilerWrapper, WorkerProfiler
21 +from vllm.profiler.wrapper import (
22 + ProtonProfilerWrapper,
23 + TorchProfilerWrapper,
24 + WorkerProfiler,
25 + create_worker_profiler,
26 + validate_worker_profiler_config,
27 +)
28 from vllm.v1.core.sched.output import CachedRequestData
29 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
30 from vllm.v1.worker.gpu_worker import Worker
31 +from vllm.v1.worker.xpu_worker import XPUWorker
32
33
34 class ConcreteWorkerProfiler(WorkerProfiler):
@@ -51,6 +59,213 @@ def default_profiler_config():
59 )
60
61
62 +def test_torch_profiler_rebuilds_one_shot_profiler_each_round(tmp_path):
63 + config = ProfilerConfig(
64 + profiler="torch",
65 + torch_profiler_dir=str(tmp_path),
66 + torch_profiler_activities=["CUDA"],
67 + torch_profiler_dump_cuda_time_total=False,
68 + warmup_iterations=2,
69 + )
70 + profilers = [MagicMock(), MagicMock()]
71 + for profiler in profilers:
72 + profiler.profiler = MagicMock()
73 +
74 + with patch(
75 + "vllm.profiler.wrapper.torch.profiler.profile", side_effect=profilers
76 + ) as profile:
77 + wrapper = TorchProfilerWrapper(
78 + config, worker_name="worker", local_rank=0, activities=["CUDA"]
79 + )
80 + profile.assert_not_called()
81 +
82 + wrapper.start()
83 + assert not wrapper.should_annotate
84 + assert not wrapper._profiler_step()
85 + wrapper.stop()
86 + wrapper.start()
87 + assert not wrapper._profiler_step()
88 + wrapper.stop()
89 +
90 + assert profile.call_count == 2
91 + assert profile.call_args.kwargs["activities"] == [
92 + torch.profiler.ProfilerActivity.CUDA
93 + ]
94 + for profiler in profilers:
95 + profiler.start.assert_called_once_with()
96 + profiler.step.assert_called_once_with()
97 + profiler.stop.assert_called_once_with()
98 +
99 +
100 +def test_torch_profiler_records_each_profile_round(tmp_path):
101 + traces: list[torch.profiler.profile] = []
102 + wrapper = TorchProfilerWrapper(
103 + ProfilerConfig(
104 + profiler="torch",
105 + torch_profiler_dir=str(tmp_path),
106 + torch_profiler_dump_cuda_time_total=False,
107 + ),
108 + worker_name="worker",
109 + local_rank=1,
110 + activities=["CPU"],
111 + on_trace_ready=traces.append,
112 + )
113 +
114 + for run in range(2):
115 + wrapper.start()
116 + with torch.profiler.record_function(f"run_{run}"):
117 + pass
118 + wrapper.stop()
119 +
120 + assert len(traces) == 2
121 + for run, trace in enumerate(traces):
122 + assert {
123 + event.name for event in trace.events() if event.name.startswith("run_")
124 + } == {f"run_{run}"}
125 +
126 +
127 +@pytest.mark.parametrize(
128 + "activities", [["CPU", "CUDA"], ["CUDA"], ["CPU", "XPU"], ["XPU"]]
129 +)
130 +@pytest.mark.parametrize("dump_device_time", [True, False])
131 +def test_torch_profiler_device_summary(tmp_path, capsys, activities, dump_device_time):
132 + """Device summaries honor the dump option on both CUDA and XPU."""
133 + config = ProfilerConfig(
134 + profiler="torch",
135 + torch_profiler_dir=str(tmp_path),
136 + torch_profiler_dump_cuda_time_total=dump_device_time,
137 + )
138 + with patch("vllm.profiler.wrapper.torch.profiler.profile") as profile:
139 + profile.return_value.key_averages.return_value.table.return_value = (
140 + "device times"
141 + )
142 + wrapper = TorchProfilerWrapper(
143 + config, worker_name="worker", local_rank=0, activities=activities
144 + )
145 + wrapper.start()
146 + wrapper.stop()
147 +
148 + summary = tmp_path / "profiler_out_0.txt"
149 + assert summary.exists() == dump_device_time
150 + assert ("device times" in capsys.readouterr().out) == dump_device_time
151 + if dump_device_time:
152 + assert summary.read_text() == "device times\n"
153 +
154 +
155 +@pytest.mark.parametrize(
156 + "activities",
157 + [[], ["CPU", "CPU"], ["INVALID"]],
158 +)
159 +def test_torch_profiler_activities_reject_invalid_values(activities):
160 + with pytest.raises(ValueError, match="torch_profiler_activities"):
161 + ProfilerConfig(
162 + profiler="torch",
163 + torch_profiler_dir="/tmp/mock",
164 + torch_profiler_activities=activities,
165 + )
166 +
167 +
168 +def test_torch_profiler_activities_require_torch_profiler():
169 + with pytest.raises(ValueError, match="only applicable"):
170 + ProfilerConfig(torch_profiler_activities=["CPU"])
171 +
172 +
173 +@pytest.mark.parametrize(
174 + ("device_type", "activities"),
175 + [
176 + ("cuda", ["XPU"]),
177 + ("xpu", ["CUDA"]),
178 + ("cpu", ["CUDA"]),
179 + ("cpu", ["XPU"]),
180 + ],
181 +)
182 +def test_worker_rejects_unsupported_activities_at_startup(device_type, activities):
183 + config = ProfilerConfig(
184 + profiler="torch",
185 + torch_profiler_dir="/tmp/mock",
186 + torch_profiler_activities=activities,
187 + )
188 +
189 + with (
190 + patch.object(current_platform, "device_type", device_type),
191 + pytest.raises(ValueError, match="Unsupported torch profiler activities"),
192 + ):
193 + validate_worker_profiler_config(config)
194 +
195 +
196 +@pytest.mark.parametrize("device_type", ["cpu", "xpu"])
197 +def test_worker_rejects_cuda_profiler_on_other_devices(device_type):
198 + with (
199 + patch.object(current_platform, "device_type", device_type),
200 + pytest.raises(ValueError, match="Unsupported profiler type"),
201 + ):
202 + validate_worker_profiler_config(ProfilerConfig(profiler="cuda"))
203 +
204 +
205 +@pytest.mark.parametrize(
206 + ("device_type", "activities", "expected"),
207 + [
208 + ("cuda", None, ("CPU", "CUDA")),
209 + ("xpu", None, ("CPU", "XPU")),
210 + ("cpu", None, ("CPU",)),
211 + ("cuda", ["CUDA"], ("CUDA",)),
212 + ("xpu", ["XPU"], ("XPU",)),
213 + ("cuda", ["CPU"], ("CPU",)),
214 + ("xpu", ["CPU"], ("CPU",)),
215 + ("cpu", ["CPU"], ("CPU",)),
216 + ],
217 +)
218 +def test_worker_creates_platform_torch_profiler(device_type, activities, expected):
219 + config = ProfilerConfig(
220 + profiler="torch",
221 + torch_profiler_dir="/tmp/mock",
222 + torch_profiler_activities=activities,
223 + )
224 +
225 + with (
226 + patch.object(current_platform, "device_type", device_type),
227 + patch("vllm.profiler.wrapper.TorchProfilerWrapper") as wrapper,
228 + ):
229 + validate_worker_profiler_config(config)
230 + profiler = create_worker_profiler(
231 + config,
232 + worker_name="rank0",
233 + local_rank=0,
234 + )
235 +
236 + assert profiler is wrapper.return_value
237 + wrapper.assert_called_once_with(
238 + config,
239 + worker_name="rank0",
240 + local_rank=0,
241 + activities=expected,
242 + )
243 +
244 +
245 +@pytest.mark.parametrize("worker_type", [Worker, XPUWorker])
246 +def test_worker_reuses_torch_wrapper_across_profile_rounds(worker_type):
247 + worker = object.__new__(worker_type)
248 + worker.rank = 0
249 + worker.local_rank = 0
250 + worker.profiler = None
251 + worker.profiler_config = ProfilerConfig(
252 + profiler="torch", torch_profiler_dir="/tmp/mock"
253 + )
254 +
255 + with (
256 + patch("vllm.distributed.utils.get_worker_rank_suffix", return_value="rank0"),
257 + patch("vllm.profiler.wrapper.TorchProfilerWrapper") as wrapper,
258 + ):
259 + worker.profile()
260 + worker.profile(is_start=False)
261 + worker.profile()
262 +
263 + assert worker.profiler is wrapper.return_value
264 + wrapper.assert_called_once()
265 + assert wrapper.return_value.start.call_count == 2
266 + wrapper.return_value.stop.assert_called_once_with()
267 +
268 +
269 def test_immediate_start_stop(default_profiler_config):
270 """Test standard start without delay."""
271 profiler = ConcreteWorkerProfiler(default_profiler_config)
@@ -318,13 +533,17 @@ class TestAnnotateProfile:
533 "execute_5_context_1(sq4sk4sqsq16sqsk16)_generation_1(sq1sk11sqsq1sqsk11)"
534 )
535
321 - def test_skips_annotation_work_after_profiler_stops(self):
536 + def test_skips_annotation_work_when_profiler_does_not_annotate(self):
537 worker = MagicMock()
323 - worker.profiler.is_running = False
538 + worker.profiler.should_annotate = False
539
325 - context = Worker.annotate_profile(worker, scheduler_output=None)
540 + with patch(
541 + "vllm.v1.worker.gpu_worker.compute_iteration_details"
542 + ) as compute_iteration_details:
543 + context = Worker.annotate_profile(worker, scheduler_output=None)
544
545 worker.profiler.step.assert_called_once_with()
546 + compute_iteration_details.assert_not_called()
547 worker.profiler.annotate_context_manager.assert_not_called()
548 assert isinstance(context, nullcontext)
549
@@ -837,17 +1056,18 @@ class TestProtonProfilerWrapper:
1056
1057 @_requires_cuda_for_proton
1058 def test_gpu_worker_creates_proton_profiler():
840 - worker = MagicMock()
1059 + worker = object.__new__(Worker)
1060 worker.rank = 1
1061 + worker.local_rank = 1
1062 worker.profiler = None
843 - worker.profiler_config.profiler = "proton"
1063 + worker.profiler_config = MagicMock(profiler="proton")
1064
1065 with (
1066 patch(
1067 "vllm.distributed.utils.get_worker_rank_suffix",
1068 return_value="rank1",
1069 ),
850 - patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper,
1070 + patch("vllm.profiler.wrapper.ProtonProfilerWrapper") as wrapper,
1071 ):
1072 Worker.profile(worker)
1073
@@ -857,17 +1077,18 @@ def test_gpu_worker_creates_proton_profiler():
1077
1078 @_requires_cuda_for_proton
1079 def test_gpu_worker_recreates_proton_profiler_for_each_run():
860 - worker = MagicMock()
1080 + worker = object.__new__(Worker)
1081 worker.rank = 1
1082 + worker.local_rank = 1
1083 worker.profiler = None
863 - worker.profiler_config.profiler = "proton"
1084 + worker.profiler_config = MagicMock(profiler="proton")
1085
1086 with (
1087 patch(
1088 "vllm.distributed.utils.get_worker_rank_suffix",
1089 return_value="rank1",
1090 ),
870 - patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper,
1091 + patch("vllm.profiler.wrapper.ProtonProfilerWrapper") as wrapper,
1092 ):
1093 wrapper.return_value.has_cuda_graph_session = False
1094 Worker.profile(worker, profile_prefix="first")
vllm/config/profiler.py
+25 -1
@@ -14,6 +14,7 @@ from vllm.utils.hashing import safe_hash
14 logger = init_logger(__name__)
15
16 ProfilerKind = Literal["torch", "cuda", "proton"]
17 +TorchProfilerActivity = Literal["CPU", "CUDA", "PrivateUse1", "XPU"]
18 ProtonBackend = Literal["cupti"]
19 ProtonContext = Literal["shadow", "python"]
20 ProtonData = Literal["tree", "trace"]
@@ -51,6 +52,12 @@ class ProfilerConfig:
52 worker's traces (CPU & GPU) will be saved under this directory. Note that
53 it must be an absolute path."""
54
55 + torch_profiler_activities: list[TorchProfilerActivity] | None = Field(
56 + default=None, min_length=1
57 + )
58 + """Activities recorded by workers using the torch profiler. When unset,
59 + each worker uses its platform default: CPU; CPU and CUDA; or CPU and XPU."""
60 +
61 proton_profiler_dir: str = ""
62 """Directory to save Triton Proton profiles. Each worker writes a
63 separate rank-qualified file."""
@@ -170,7 +177,24 @@ class ProfilerConfig:
177 @model_validator(mode="after")
178 def _validate_profiler_config(self) -> Self:
179 has_delay_or_limit = self.delay_iterations > 0 or self.max_iterations > 0
173 - if self.profiler == "torch" and has_delay_or_limit and not self.ignore_frontend:
180 + activities = self.torch_profiler_activities
181 + if activities is not None:
182 + if self.profiler != "torch":
183 + raise ValueError(
184 + "torch_profiler_activities is only applicable when profiler "
185 + "is set to 'torch'"
186 + )
187 + if len(activities) != len(set(activities)):
188 + raise ValueError(
189 + "torch_profiler_activities must not contain duplicates"
190 + )
191 + records_cpu_activity = activities is None or "CPU" in activities
192 + if (
193 + self.profiler == "torch"
194 + and has_delay_or_limit
195 + and not self.ignore_frontend
196 + and records_cpu_activity
197 + ):
198 logger.warning_once(
199 "Using 'torch' profiler with delay_iterations or max_iterations "
200 "while ignore_frontend is False may result in high overhead."
vllm/profiler/wrapper.py
+106 -11
@@ -6,10 +6,9 @@ import inspect
6 import json
7 import os
8 from abc import ABC, abstractmethod
9 -from collections.abc import Callable, Iterator
9 +from collections.abc import Callable, Iterator, Sequence
10 from contextlib import AbstractContextManager, contextmanager, nullcontext, suppress
11 from glob import glob
12 -from typing import Literal
12 from uuid import uuid4
13
14 import torch
@@ -18,7 +17,7 @@ from typing_extensions import override
17
18 import vllm.version
19 from vllm.config import ProfilerConfig
21 -from vllm.config.profiler import _is_uri_path
20 +from vllm.config.profiler import ProfilerKind, TorchProfilerActivity, _is_uri_path
21 from vllm.logger import init_logger
22 from vllm.platforms import current_platform
23
@@ -57,6 +56,11 @@ class WorkerProfiler(ABC):
56 """Whether the underlying profiler is currently collecting data."""
57 return self._running
58
59 + @property
60 + def should_annotate(self) -> bool:
61 + """Whether worker iterations should receive profiler annotations."""
62 + return self._running
63 +
64 @abstractmethod
65 def _start(self) -> None:
66 """Start the profiler."""
@@ -176,7 +180,6 @@ class WorkerProfiler(ABC):
180 return nullcontext()
181
182
179 -TorchProfilerActivity = Literal["CPU", "CUDA", "PrivateUse1", "XPU"]
183 TorchProfilerActivityMap = {
184 "CPU": torch.profiler.ProfilerActivity.CPU,
185 "CUDA": torch.profiler.ProfilerActivity.CUDA,
@@ -191,7 +194,7 @@ class TorchProfilerWrapper(WorkerProfiler):
194 profiler_config: ProfilerConfig,
195 worker_name: str,
196 local_rank: int,
194 - activities: list[TorchProfilerActivity],
197 + activities: Sequence[TorchProfilerActivity],
198 on_trace_ready: Callable[[torch.profiler.profile], None] | None = None,
199 ) -> None:
200 super().__init__(profiler_config)
@@ -224,7 +227,12 @@ class TorchProfilerWrapper(WorkerProfiler):
227 use_gzip=profiler_config.torch_profiler_use_gzip,
228 )
229
227 - self.dump_cpu_time_total = "CPU" in activities and len(activities) == 1
230 + self._records_cpu_activity = "CPU" in activities
231 + self.dump_device_time_total = (
232 + any(activity != "CPU" for activity in activities)
233 + and profiler_config.torch_profiler_dump_cuda_time_total
234 + )
235 + self.dump_cpu_time_total = self._records_cpu_activity and len(activities) == 1
236
237 # Create profiler schedule if warmup or wait iterations are configured
238 profiler_schedule = None
@@ -244,7 +252,7 @@ class TorchProfilerWrapper(WorkerProfiler):
252 profiler_config.active_iterations,
253 )
254
247 - self.profiler = torch.profiler.profile(
255 + self._profiler_kwargs = dict(
256 activities=[TorchProfilerActivityMap[activity] for activity in activities],
257 schedule=profiler_schedule,
258 record_shapes=profiler_config.torch_profiler_record_shapes,
@@ -253,6 +261,7 @@ class TorchProfilerWrapper(WorkerProfiler):
261 with_flops=profiler_config.torch_profiler_with_flops,
262 on_trace_ready=trace_handler,
263 )
264 + self.profiler: torch.profiler.profile
265
266 # Track if we're using a schedule (need to call step())
267 self._uses_schedule = profiler_schedule is not None
@@ -260,10 +269,11 @@ class TorchProfilerWrapper(WorkerProfiler):
269 # Subtract 1 because profiler.start() already consumes step 0
270 # (WAIT or WARMUP), so only wait + warmup - 1 non-active steps
271 # remain to be advanced through via profiler.step() calls.
263 - self._warmup_steps_remaining = max(
272 + self._initial_warmup_steps_remaining = max(
273 profiler_config.wait_iterations + profiler_config.warmup_iterations - 1,
274 0,
275 )
276 + self._warmup_steps_remaining = self._initial_warmup_steps_remaining
277 self._version_metadata_added = False
278
279 def _build_profiler_table(
@@ -319,6 +329,9 @@ class TorchProfilerWrapper(WorkerProfiler):
329
330 @override
331 def _start(self) -> None:
332 + self.profiler = torch.profiler.profile(**self._profiler_kwargs)
333 + self._warmup_steps_remaining = self._initial_warmup_steps_remaining
334 + self._version_metadata_added = False
335 self.profiler.start()
336 # No-schedule case: Kineto is live immediately. With a schedule this
337 # no-ops and _profiler_step stamps it once WAIT ends.
@@ -328,10 +341,9 @@ class TorchProfilerWrapper(WorkerProfiler):
341 def _stop(self) -> None:
342 self.profiler.stop()
343
331 - profiler_config = self.profiler_config
344 rank = self.local_rank
333 - if profiler_config.torch_profiler_dump_cuda_time_total:
334 - table = self._build_profiler_table(sort_key="self_cuda_time_total")
345 + if self.dump_device_time_total:
346 + table = self._build_profiler_table(sort_key="self_device_time_total")
347 self._write_profiler_table(rank, table)
348
349 # only print profiler results on rank 0
@@ -367,8 +379,15 @@ class TorchProfilerWrapper(WorkerProfiler):
379 return False
380 return True
381
382 + @property
383 + @override
384 + def should_annotate(self) -> bool:
385 + return self._running and self._records_cpu_activity
386 +
387 @override
388 def annotate_context_manager(self, name: str):
389 + if not self.should_annotate:
390 + return nullcontext()
391 return torch.profiler.record_function(name)
392
393
@@ -606,6 +625,82 @@ class CudaProfilerWrapper(WorkerProfiler):
625 return torch.cuda.nvtx.range(name)
626
627
628 +_DEFAULT_TORCH_PROFILER_ACTIVITIES: dict[str, tuple[TorchProfilerActivity, ...]] = {
629 + "cpu": ("CPU",),
630 + "cuda": ("CPU", "CUDA"),
631 + "xpu": ("CPU", "XPU"),
632 +}
633 +_SUPPORTED_TORCH_PROFILER_ACTIVITIES = {
634 + device: frozenset(activities)
635 + for device, activities in _DEFAULT_TORCH_PROFILER_ACTIVITIES.items()
636 +}
637 +_SUPPORTED_PROFILER_KINDS: dict[str, frozenset[ProfilerKind]] = {
638 + "cpu": frozenset(("torch",)),
639 + "cuda": frozenset(("torch", "cuda", "proton")),
640 + "xpu": frozenset(("torch",)),
641 +}
642 +
643 +
644 +def validate_worker_profiler_config(profiler_config: ProfilerConfig) -> None:
645 + """Validate profiler selections against the worker's capabilities."""
646 + profiler_type = profiler_config.profiler
647 + if profiler_type is None:
648 + return
649 + device_type = current_platform.device_type
650 + if device_type not in _SUPPORTED_PROFILER_KINDS:
651 + raise ValueError(f"Unsupported profiler device type: {device_type}")
652 + supported_kinds = _SUPPORTED_PROFILER_KINDS[device_type]
653 + if profiler_type not in supported_kinds:
654 + supported_names = ", ".join(sorted(supported_kinds))
655 + raise ValueError(
656 + f"Unsupported profiler type for {device_type}: "
657 + f"{profiler_type}. Supported profiler types: {supported_names}."
658 + )
659 + if profiler_type == "torch":
660 + default_activities = _DEFAULT_TORCH_PROFILER_ACTIVITIES[device_type]
661 + supported_activities = _SUPPORTED_TORCH_PROFILER_ACTIVITIES[device_type]
662 + configured = profiler_config.torch_profiler_activities
663 + activities = default_activities if configured is None else configured
664 + unsupported = set(activities) - supported_activities
665 + if unsupported:
666 + unsupported_names = ", ".join(sorted(unsupported))
667 + supported_names = ", ".join(sorted(supported_activities))
668 + raise ValueError(
669 + f"Unsupported torch profiler activities for "
670 + f"{device_type}: {unsupported_names}. "
671 + f"Supported activities: {supported_names}."
672 + )
673 +
674 +
675 +def create_worker_profiler(
676 + profiler_config: ProfilerConfig,
677 + *,
678 + worker_name: str,
679 + local_rank: int,
680 +) -> WorkerProfiler:
681 + """Create a profiler using a validated config and platform defaults."""
682 + profiler_type = profiler_config.profiler
683 + if profiler_type == "torch":
684 + default_activities = _DEFAULT_TORCH_PROFILER_ACTIVITIES[
685 + current_platform.device_type
686 + ]
687 + configured = profiler_config.torch_profiler_activities
688 + logger.debug("Starting torch profiler with trace name: %s", worker_name)
689 + return TorchProfilerWrapper(
690 + profiler_config,
691 + worker_name=worker_name,
692 + local_rank=local_rank,
693 + activities=default_activities if configured is None else tuple(configured),
694 + )
695 + if profiler_type == "cuda":
696 + logger.debug("Starting CUDA profiler")
697 + return CudaProfilerWrapper(profiler_config)
698 +
699 + assert profiler_type == "proton", f"Unknown profiler type: {profiler_type}"
700 + logger.debug("Starting Proton profiler with trace name: %s", worker_name)
701 + return ProtonProfilerWrapper(profiler_config, worker_name=worker_name)
702 +
703 +
704 def create_graph_capture_profiler(
705 profiler_config: ProfilerConfig, global_rank: int
706 ) -> WorkerProfiler | None:
vllm/v1/engine/async_llm.py
+2
@@ -208,9 +208,11 @@ class AsyncLLM(EngineClient):
208 pass
209
210 self.profiler = profiler
211 + configured_activities = vllm_config.profiler_config.torch_profiler_activities
212 if (
213 vllm_config.profiler_config.profiler == "torch"
214 and not vllm_config.profiler_config.ignore_frontend
215 + and (configured_activities is None or "CPU" in configured_activities)
216 ):
217 profiler_dir = vllm_config.profiler_config.torch_profiler_dir
218 logger.info(
vllm/v1/worker/cpu_worker.py
-22
@@ -7,7 +7,6 @@ import vllm.v1.worker.cpu.shm # noqa # isort: skip
7 import math
8 import os
9 import sys
10 -from typing import Any
10
11 import psutil
12 import torch
@@ -16,7 +15,6 @@ from vllm import envs
15 from vllm.config import CompilationMode, VllmConfig
16 from vllm.logger import init_logger
17 from vllm.platforms import CpuArchEnum, current_platform
19 -from vllm.profiler.wrapper import TorchProfilerWrapper
18 from vllm.utils.cpu_resource_utils import (
19 get_allowed_cpu_list,
20 get_memory_node_info,
@@ -106,18 +104,6 @@ class CPUWorker(Worker):
104
105 self.parallel_config.disable_custom_all_reduce = True
106
109 - # Torch profiler. Enabled and configured through profiler_config.
110 - self.profiler: Any | None = None
111 - profiler_config = vllm_config.profiler_config
112 - if profiler_config.profiler == "torch":
113 - worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
114 - self.profiler = TorchProfilerWrapper(
115 - profiler_config,
116 - worker_name=worker_name,
117 - local_rank=self.local_rank,
118 - activities=["CPU"],
119 - )
120 -
107 def init_device(self):
108 self.device = torch.device("cpu")
109
@@ -287,11 +273,3 @@ class CPUWorker(Worker):
273 language_model=self.compilation_config.compilation_time,
274 encoder=self.compilation_config.encoder_compilation_time,
275 )
290 -
291 - def profile(self, is_start: bool = True, profile_prefix: str | None = None):
292 - if self.profiler is None:
293 - raise RuntimeError("Profiler is not enabled.")
294 - if is_start:
295 - self.profiler.start()
296 - else:
297 - self.profiler.stop()
vllm/v1/worker/gpu_worker.py
+10 -31
@@ -61,10 +61,9 @@ from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
61 from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory
62 from vllm.platforms import current_platform
63 from vllm.profiler.wrapper import (
64 - CudaProfilerWrapper,
65 - ProtonProfilerWrapper,
66 - TorchProfilerWrapper,
64 create_graph_capture_profiler,
65 + create_worker_profiler,
66 + validate_worker_profiler_config,
67 )
68 from vllm.sequence import IntermediateTensors
69 from vllm.tasks import SupportedTask
@@ -222,6 +221,7 @@ class Worker(WorkerBase):
221 # so we have all the information needed for proper trace naming.
222 self.profiler: Any | None = None
223 self.profiler_config = vllm_config.profiler_config
224 + validate_worker_profiler_config(self.profiler_config)
225
226 self.use_v2_model_runner = vllm_config.use_v2_model_runner
227
@@ -1038,7 +1038,7 @@ class Worker(WorkerBase):
1038 return nullcontext()
1039
1040 self.profiler.step()
1041 - if not self.profiler.is_running:
1041 + if not self.profiler.should_annotate:
1042 return nullcontext()
1043
1044 iteration_details = compute_iteration_details(scheduler_output)
@@ -1262,7 +1262,6 @@ class Worker(WorkerBase):
1262 )
1263
1264 if is_start:
1265 - profiler_type = self.profiler_config.profiler
1265 # Generate the trace name by combining prefix with comprehensive rank suffix
1266 from vllm.distributed.utils import get_worker_rank_suffix
1267
@@ -1274,36 +1273,16 @@ class Worker(WorkerBase):
1273 else:
1274 trace_name = rank_suffix
1275
1277 - if profiler_type == "proton" and self.profiler is not None:
1276 + if self.profiler_config.profiler == "proton" and self.profiler is not None:
1277 self.profiler.set_output_name(trace_name)
1278
1279 # Create the profiler wrapper only on the first start call
1280 if self.profiler is None:
1282 - if profiler_type == "torch":
1283 - self.profiler = TorchProfilerWrapper(
1284 - self.profiler_config,
1285 - worker_name=trace_name,
1286 - local_rank=self.local_rank,
1287 - activities=["CPU", "CUDA"],
1288 - )
1289 - logger.debug(
1290 - "Starting torch profiler with trace name: %s", trace_name
1291 - )
1292 - elif profiler_type == "cuda":
1293 - self.profiler = CudaProfilerWrapper(self.profiler_config)
1294 - logger.debug("Starting CUDA profiler")
1295 - elif profiler_type == "proton":
1296 - self.profiler = ProtonProfilerWrapper(
1297 - self.profiler_config, worker_name=trace_name
1298 - )
1299 - logger.debug(
1300 - "Starting Proton profiler with trace name: %s", trace_name
1301 - )
1302 - else:
1303 - # Config validation should prevent this code being reached
1304 - raise ValueError(
1305 - f"Invalid profiler value of {self.profiler_config.profiler}"
1306 - )
1281 + self.profiler = create_worker_profiler(
1282 + self.profiler_config,
1283 + worker_name=trace_name,
1284 + local_rank=self.local_rank,
1285 + )
1286
1287 self.profiler.start()
1288 else:
vllm/v1/worker/xpu_worker.py
-28
@@ -8,7 +8,6 @@ import torch
8 from vllm.config import VllmConfig
9 from vllm.logger import init_logger
10 from vllm.platforms import current_platform
11 -from vllm.profiler.wrapper import TorchProfilerWrapper
11 from vllm.utils.mem_utils import MemorySnapshot, format_gib
12 from vllm.utils.torch_utils import set_random_seed
13 from vllm.v1.utils import report_usage_stats
@@ -225,33 +224,6 @@ class XPUWorker(Worker):
224 # If usage stat is enabled, collect relevant info.
225 report_usage_stats(self.vllm_config)
226
228 - def profile(self, is_start: bool = True, profile_prefix: str | None = None):
229 - if self.profiler_config is None or self.profiler_config.profiler is None:
230 - raise RuntimeError(
231 - "Profiling is not enabled. Please set --profiler-config to enable "
232 - "profiling. Example: "
233 - "'--profiler-config.profiler=torch --profiler-config.torch_profiler_dir"
234 - "=YOUR_DIR_PATH_TO_DUMP_TRACE'"
235 - )
236 -
237 - if is_start and self.profiler is None:
238 - from vllm.distributed.utils import get_worker_rank_suffix
239 -
240 - rank_suffix = get_worker_rank_suffix(global_rank=self.rank)
241 - trace_name = (
242 - f"{profile_prefix}_{rank_suffix}" if profile_prefix else rank_suffix
243 - )
244 -
245 - self.profiler = TorchProfilerWrapper(
246 - self.profiler_config,
247 - worker_name=trace_name,
248 - local_rank=self.local_rank,
249 - activities=["CPU", "XPU"],
250 - )
251 - logger.debug("Starting torch profiler with trace name: %s", trace_name)
252 -
253 - super().profile(is_start=is_start, profile_prefix=profile_prefix)
254 -
227 def shutdown(self) -> None:
228 logger.info(
229 "XPUWorker shutdown: cleaning up (rank=%d, local_rank=%d)",