@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a367cdbca0

[Profiler] Add minimal Triton Proton profiling backend (#48789)

Signed-off-by: Luossu <zhluosuu@outlook.com>

鐘天楽 committed Aug 11, 2026 at 14:51 UTC a367cdbca0824c8c428fe15b4552fbc4b824d718
7 files changed +638 -21
docs/contributing/profiling.md
+64
@@ -84,6 +84,70 @@ curl -X POST http://localhost:8000/v1/chat/completions \
84 $ curl -X POST http://localhost:8000/stop_profile
85 ```
86
87 +## Profile with Triton Proton
88 +
89 +[Proton](https://github.com/triton-lang/triton/tree/main/third_party/proton)
90 +is Triton's GPU profiler. It can collect a low-overhead aggregate tree or a
91 +Chrome trace and works through the same vLLM profiling controls as the PyTorch
92 +and CUDA profilers. Proton currently supports NVIDIA GPUs through CUPTI and
93 +requires eager execution.
94 +
95 +Start a server with a local output directory:
96 +
97 +```bash
98 +vllm serve meta-llama/Llama-3.1-8B-Instruct \
99 + --enforce-eager \
100 + --profiler-config '{
101 + "profiler": "proton",
102 + "proton_profiler_dir": "./proton_profile",
103 + "proton_output_format": "hatchet",
104 + "proton_hook": "triton"
105 + }'
106 +```
107 +
108 +Then use `/start_profile` and `/stop_profile` as shown above, or pass
109 +`--profile` to a vLLM benchmark. Each worker uses a topology- and
110 +rank-qualified output name, such as
111 +`proton_dp0_pp0_tp0_dcp0_ep0_rank0_pid1234_0123456789abcdef0123456789abcdef_run0.hatchet`,
112 +so distributed workers, restarted servers, and repeated profiling runs do not
113 +overwrite one another. A `profile_prefix` is included when supplied. Each
114 +profile is finalized by `/stop_profile` and is ready to inspect immediately.
115 +
116 +The Proton-specific options are:
117 +
118 +- `proton_context`: `shadow` (default) or `python`
119 +- `proton_data`: `tree` (default) or `trace`
120 +- `proton_backend`: `cupti` or automatic
121 +- `proton_mode`: an optional backend mode string
122 +- `proton_hook`: `triton` to record Triton launch metadata, or unset
123 +- `proton_output_format`: `hatchet`, `hatchet_msgpack`, `chrome_trace`, or unset
124 +
125 +`hatchet` and `hatchet_msgpack` require `proton_data: "tree"`, while
126 +`chrome_trace` requires `proton_data: "trace"`. CUDA graph profiling is not yet
127 +supported, so `--enforce-eager` is required for every Proton mode.
128 +
129 +Automatic backend selection is recommended. vLLM currently supports Proton's
130 +`cupti` backend on NVIDIA GPUs. ROCm support is not yet available. vLLM does not
131 +expose Proton's experimental instrumentation backend because current upstream
132 +Triton builds can produce profiles without timing metrics. Backend-specific
133 +modes, including `pcsampling`, can be selected with `proton_mode`.
134 +
135 +Triton 3.6 supports explicit `hatchet` and `chrome_trace` output. The
136 +`hatchet_msgpack` format and `periodic_flushing` mode require Triton 3.7 or
137 +newer; vLLM rejects these options with the detected version before starting a
138 +session.
139 +
140 +Inspect tree profiles with:
141 +
142 +```bash
143 +proton-viewer -m time/ns \
144 + proton_profile/proton_dp0_pp0_tp0_dcp0_ep0_rank0_pid1234_0123456789abcdef0123456789abcdef_run0.hatchet
145 +```
146 +
147 +Chrome traces (`proton_data: "trace"`) can be opened in
148 +<https://ui.perfetto.dev/>. Proton is imported lazily, so selecting another
149 +profiler does not require a Proton-capable Triton installation.
150 +
151 ## Profile with NVIDIA Nsight Systems
152
153 Nsight systems is an advanced tool that exposes more profiling details, such as register and shared memory usage, annotated code regions and low-level CUDA APIs and events.
tests/v1/worker/test_gpu_profiler.py
+325 -4
@@ -1,12 +1,23 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from unittest.mock import MagicMock
3 +import os
4 +from contextlib import nullcontext
5 +from types import SimpleNamespace
6 +from unittest.mock import MagicMock, Mock, call, patch
7 +from uuid import UUID
8
9 import pytest
6 -
7 -from vllm.config import CUDAGraphMode, ProfilerConfig
10 +from pydantic import ValidationError
11 +
12 +from vllm.config import (
13 + CompilationConfig,
14 + CUDAGraphMode,
15 + ProfilerConfig,
16 + VllmConfig,
17 +)
18 from vllm.config.profiler import _is_uri_path
9 -from vllm.profiler.wrapper import WorkerProfiler
19 +from vllm.platforms import current_platform
20 +from vllm.profiler.wrapper import ProtonProfilerWrapper, WorkerProfiler
21 from vllm.v1.core.sched.output import CachedRequestData
22 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
23 from vllm.v1.worker.gpu_worker import Worker
@@ -281,6 +292,16 @@ class TestAnnotateProfile:
292 "execute_5_context_1(sq4sk4sqsq16sqsk16)_generation_1(sq1sk11sqsq1sqsk11)"
293 )
294
295 + def test_skips_annotation_work_after_profiler_stops(self):
296 + worker = MagicMock()
297 + worker.profiler.is_running = False
298 +
299 + context = Worker.annotate_profile(worker, scheduler_output=None)
300 +
301 + worker.profiler.step.assert_called_once_with()
302 + worker.profiler.annotate_context_manager.assert_not_called()
303 + assert isinstance(context, nullcontext)
304 +
305
306 def test_profiler_entered_during_capture():
307 """Profiler is used as a context manager in _warmup_and_capture,
@@ -298,3 +319,303 @@ def test_profiler_entered_during_capture():
319
320 mock_profiler.__enter__.assert_called_once()
321 mock_profiler.__exit__.assert_called_once()
322 +
323 +
324 +def make_proton(session_id: int | None = 7):
325 + return SimpleNamespace(
326 + start=Mock(return_value=session_id),
327 + activate=Mock(),
328 + deactivate=Mock(),
329 + finalize=Mock(),
330 + scope=Mock(return_value=nullcontext()),
331 + )
332 +
333 +
334 +def make_proton_wrapper(
335 + tmp_path, proton=None, triton_version="3.6.0", **config_overrides
336 +):
337 + proton = proton or make_proton()
338 + config = ProfilerConfig(
339 + profiler="proton",
340 + proton_profiler_dir=str(tmp_path),
341 + **config_overrides,
342 + )
343 +
344 + def import_module(name):
345 + if name == "triton.profiler":
346 + return proton
347 + assert name == "triton"
348 + return SimpleNamespace(__version__=triton_version)
349 +
350 + with patch(
351 + "vllm.profiler.wrapper.importlib.import_module", side_effect=import_module
352 + ):
353 + wrapper = ProtonProfilerWrapper(config, worker_name="rank_3")
354 + return wrapper, proton
355 +
356 +
357 +_requires_cuda_for_proton = pytest.mark.skipif(
358 + not current_platform.is_cuda(),
359 + reason="Proton profiling tests require an NVIDIA CUDA platform.",
360 +)
361 +
362 +
363 +@_requires_cuda_for_proton
364 +class TestProtonConfig:
365 + def test_normalizes_local_output_directory(self, tmp_path, monkeypatch):
366 + monkeypatch.chdir(tmp_path)
367 + config = ProfilerConfig(profiler="proton", proton_profiler_dir="profiles")
368 + assert config.proton_profiler_dir == os.path.join(tmp_path, "profiles")
369 +
370 + @pytest.mark.parametrize(
371 + ("options", "message"),
372 + [
373 + ({"proton_profiler_dir": ""}, "must be set"),
374 + ({"proton_profiler_dir": "s3://bucket/profiles"}, "local directory"),
375 + (
376 + {"proton_data": "tree", "proton_output_format": "chrome_trace"},
377 + "requires proton_data",
378 + ),
379 + (
380 + {"proton_data": "trace", "proton_output_format": "hatchet"},
381 + "requires proton_data",
382 + ),
383 + (
384 + {
385 + "proton_data": "trace",
386 + "proton_output_format": "hatchet_msgpack",
387 + },
388 + "requires proton_data",
389 + ),
390 + ],
391 + )
392 + def test_rejects_invalid_option_combinations(self, tmp_path, options, message):
393 + kwargs = {"proton_profiler_dir": str(tmp_path), **options}
394 + with pytest.raises(ValueError, match=message):
395 + ProfilerConfig(profiler="proton", **kwargs)
396 +
397 + @pytest.mark.parametrize(
398 + "field",
399 + [
400 + "proton_context",
401 + "proton_data",
402 + "proton_backend",
403 + "proton_hook",
404 + "proton_output_format",
405 + ],
406 + )
407 + def test_rejects_invalid_typed_options(self, field, tmp_path):
408 + with pytest.raises(ValidationError):
409 + ProfilerConfig(
410 + profiler="proton",
411 + proton_profiler_dir=str(tmp_path),
412 + **{field: "invalid"},
413 + )
414 +
415 + @pytest.mark.parametrize(
416 + ("field", "value"),
417 + [
418 + ("proton_profiler_dir", "profiles"),
419 + ("proton_context", "python"),
420 + ("proton_data", "trace"),
421 + ("proton_backend", "cupti"),
422 + ("proton_mode", "pcsampling"),
423 + ("proton_hook", "triton"),
424 + ("proton_output_format", "chrome_trace"),
425 + ],
426 + )
427 + def test_rejects_proton_options_for_other_profilers(self, field, value):
428 + with pytest.raises(ValueError, match=f"{field} only applicable"):
429 + ProfilerConfig(**{field: value})
430 +
431 + def test_allows_proton_when_cuda_graphs_are_disabled(self, tmp_path):
432 + config = VllmConfig(
433 + profiler_config=ProfilerConfig(
434 + profiler="proton", proton_profiler_dir=str(tmp_path)
435 + ),
436 + compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE),
437 + )
438 +
439 + assert config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE
440 +
441 + def test_rejects_proton_on_non_cuda_platforms(self, tmp_path):
442 + with (
443 + patch("vllm.platforms.current_platform.is_cuda", return_value=False),
444 + pytest.raises(ValueError, match="supports NVIDIA CUDA only"),
445 + ):
446 + VllmConfig(
447 + profiler_config=ProfilerConfig(
448 + profiler="proton", proton_profiler_dir=str(tmp_path)
449 + ),
450 + compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE),
451 + )
452 +
453 + def test_rejects_proton_when_cuda_graphs_are_enabled(self, tmp_path):
454 + with pytest.raises(ValueError, match="requires CUDA graphs to be disabled"):
455 + VllmConfig(
456 + profiler_config=ProfilerConfig(
457 + profiler="proton", proton_profiler_dir=str(tmp_path)
458 + ),
459 + compilation_config=CompilationConfig(
460 + cudagraph_mode=CUDAGraphMode.PIECEWISE
461 + ),
462 + )
463 +
464 +
465 +@_requires_cuda_for_proton
466 +class TestProtonProfilerWrapper:
467 + def test_passes_config_and_global_rank_name_to_proton(self, tmp_path):
468 + wrapper, proton = make_proton_wrapper(
469 + tmp_path,
470 + proton_context="python",
471 + proton_data="trace",
472 + proton_backend="cupti",
473 + proton_mode="pcsampling",
474 + proton_hook="triton",
475 + proton_output_format="chrome_trace",
476 + )
477 +
478 + wrapper.start()
479 +
480 + start_args = proton.start.call_args.kwargs
481 + assert start_args["name"].startswith(os.path.join(tmp_path, "proton_rank_3_"))
482 + assert start_args["name"].endswith("_run0")
483 + assert start_args | {"name": None} == {
484 + "name": None,
485 + "context": "python",
486 + "data": "trace",
487 + "backend": "cupti",
488 + "mode": "pcsampling",
489 + "hook": "triton",
490 + }
491 + wrapper.stop()
492 + proton.finalize.assert_called_once_with(session=7, output_format="chrome_trace")
493 + assert tmp_path.is_dir()
494 +
495 + def test_finalizes_each_profile_with_unique_output_names(self, tmp_path):
496 + proton = make_proton()
497 + proton.start.side_effect = [7, 8]
498 + wrapper, proton = make_proton_wrapper(tmp_path, proton)
499 +
500 + wrapper.start()
501 + wrapper.start()
502 + wrapper.stop()
503 + wrapper.start()
504 + wrapper.stop()
505 +
506 + names = [c.kwargs["name"] for c in proton.start.call_args_list]
507 + assert len(names) == len(set(names)) == 2
508 + assert names[0].endswith("_run0")
509 + assert names[1].endswith("_run1")
510 + assert proton.deactivate.call_count == 2
511 + assert proton.finalize.call_args_list == [call(session=7), call(session=8)]
512 +
513 + def test_output_names_are_unique_across_worker_restarts(self, tmp_path):
514 + with patch(
515 + "vllm.profiler.wrapper.uuid4",
516 + side_effect=[UUID(int=1), UUID(int=2)],
517 + ):
518 + first, first_proton = make_proton_wrapper(tmp_path)
519 + second, second_proton = make_proton_wrapper(tmp_path)
520 +
521 + first.start()
522 + second.start()
523 +
524 + first_name = first_proton.start.call_args.kwargs["name"]
525 + second_name = second_proton.start.call_args.kwargs["name"]
526 + assert first_name != second_name
527 + assert first_name.endswith(f"_{UUID(int=1).hex}_run0")
528 + assert second_name.endswith(f"_{UUID(int=2).hex}_run0")
529 +
530 + @pytest.mark.parametrize(
531 + ("option", "value", "feature"),
532 + [
533 + ("proton_output_format", "hatchet_msgpack", "hatchet_msgpack"),
534 + ("proton_mode", "periodic_flushing", "periodic flushing"),
535 + ],
536 + )
537 + def test_newer_features_reject_triton_3_6(self, tmp_path, option, value, feature):
538 + with pytest.raises(RuntimeError, match=feature):
539 + make_proton_wrapper(tmp_path, **{option: value})
540 +
541 + @pytest.mark.parametrize(
542 + ("option", "value"),
543 + [
544 + ("proton_output_format", "hatchet_msgpack"),
545 + ("proton_mode", "periodic_flushing"),
546 + ],
547 + )
548 + def test_triton_3_7_features(self, tmp_path, option, value):
549 + make_proton_wrapper(tmp_path, triton_version="3.7.0", **{option: value})
550 +
551 + def test_rejects_output_format_when_finalize_lacks_capability(self, tmp_path):
552 + proton = make_proton()
553 + proton.finalize = lambda session=None: None
554 +
555 + with pytest.raises(RuntimeError, match="does not support selecting"):
556 + make_proton_wrapper(tmp_path, proton, proton_output_format="hatchet")
557 +
558 + def test_missing_proton_has_actionable_error(self, tmp_path):
559 + config = ProfilerConfig(profiler="proton", proton_profiler_dir=str(tmp_path))
560 + with (
561 + patch(
562 + "vllm.profiler.wrapper.importlib.import_module",
563 + side_effect=ImportError,
564 + ),
565 + pytest.raises(RuntimeError, match="requires a Triton installation"),
566 + ):
567 + ProtonProfilerWrapper(config, worker_name="rank_0")
568 +
569 + def test_scope_annotations_delegate_to_proton(self, tmp_path):
570 + wrapper, proton = make_proton_wrapper(tmp_path)
571 + wrapper.start()
572 +
573 + context = wrapper.annotate_context_manager("decode")
574 +
575 + proton.scope.assert_called_once_with("decode")
576 + assert context is not None
577 +
578 +
579 +@_requires_cuda_for_proton
580 +def test_gpu_worker_creates_proton_profiler():
581 + worker = MagicMock()
582 + worker.rank = 1
583 + worker.profiler = None
584 + worker.profiler_config.profiler = "proton"
585 +
586 + with (
587 + patch(
588 + "vllm.distributed.utils.get_worker_rank_suffix",
589 + return_value="rank1",
590 + ),
591 + patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper,
592 + ):
593 + Worker.profile(worker)
594 +
595 + wrapper.assert_called_once_with(worker.profiler_config, worker_name="rank1")
596 + worker.profiler.start.assert_called_once_with()
597 +
598 +
599 +@_requires_cuda_for_proton
600 +def test_gpu_worker_recreates_proton_profiler_for_each_run():
601 + worker = MagicMock()
602 + worker.rank = 1
603 + worker.profiler = None
604 + worker.profiler_config.profiler = "proton"
605 +
606 + with (
607 + patch(
608 + "vllm.distributed.utils.get_worker_rank_suffix",
609 + return_value="rank1",
610 + ),
611 + patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper,
612 + ):
613 + Worker.profile(worker, profile_prefix="first")
614 + Worker.profile(worker, is_start=False)
615 + Worker.profile(worker, profile_prefix="second")
616 +
617 + assert wrapper.call_args_list == [
618 + call(worker.profiler_config, worker_name="first_rank1"),
619 + call(worker.profiler_config, worker_name="second_rank1"),
620 + ]
621 + assert wrapper.return_value.start.call_count == 2
vllm/benchmarks/latency.py
+5
@@ -147,6 +147,11 @@ def main(args: argparse.Namespace):
147 )
148 elif profiler_config.profiler == "cuda":
149 print("Profiling with cuda profiler ...")
150 + elif profiler_config.profiler == "proton":
151 + print(
152 + "Profiling with Proton (results will be saved under"
153 + f" {profiler_config.proton_profiler_dir})..."
154 + )
155 run_to_completion(do_profile=True)
156 return
157
vllm/config/profiler.py
+80 -7
@@ -13,7 +13,12 @@ from vllm.utils.hashing import safe_hash
13
14 logger = init_logger(__name__)
15
16 -ProfilerKind = Literal["torch", "cuda"]
16 +ProfilerKind = Literal["torch", "cuda", "proton"]
17 +ProtonBackend = Literal["cupti"]
18 +ProtonContext = Literal["shadow", "python"]
19 +ProtonData = Literal["tree", "trace"]
20 +ProtonHook = Literal["triton"]
21 +ProtonOutputFormat = Literal["hatchet", "hatchet_msgpack", "chrome_trace"]
22
23
24 def _is_uri_path(path: str) -> bool:
@@ -38,13 +43,39 @@ class ProfilerConfig:
43 """Which profiler to use. Defaults to None. Options are:
44
45 - 'torch': Use PyTorch profiler.
41 - - 'cuda': Use CUDA profiler."""
46 + - 'cuda': Use CUDA profiler.
47 + - 'proton': Use Triton Proton profiler."""
48
49 torch_profiler_dir: str = ""
50 """Directory to save torch profiler traces. Both AsyncLLM's CPU traces and
51 worker's traces (CPU & GPU) will be saved under this directory. Note that
52 it must be an absolute path."""
53
54 + proton_profiler_dir: str = ""
55 + """Directory to save Triton Proton profiles. Each worker writes a
56 + separate rank-qualified file."""
57 +
58 + proton_context: ProtonContext = "shadow"
59 + """Proton context source. ``shadow`` records explicit scopes with low
60 + overhead; ``python`` records Python call stacks."""
61 +
62 + proton_data: ProtonData = "tree"
63 + """Proton output type. ``tree`` produces Hatchet data and ``trace``
64 + produces a Chrome trace."""
65 +
66 + proton_backend: ProtonBackend | None = None
67 + """Proton GPU backend. ``None`` lets Proton select CUPTI automatically."""
68 +
69 + proton_mode: str | None = None
70 + """Optional backend-specific Proton mode string, such as ``pcsampling``."""
71 +
72 + proton_hook: ProtonHook | None = None
73 + """Optional Proton hook. Use ``triton`` to add Triton launch metadata."""
74 +
75 + proton_output_format: ProtonOutputFormat | None = None
76 + """Optional format passed to Proton when finalizing a profile. ``None``
77 + uses the default format for ``proton_data``."""
78 +
79 torch_profiler_with_stack: bool = True
80 """If `True`, enables stack tracing in the torch profiler. Enabled by default
81 as it is useful for debugging. Can be disabled via
@@ -142,18 +173,60 @@ class ProfilerConfig:
173 "while ignore_frontend is False may result in high overhead."
174 )
175
145 - profiler_dir = self.torch_profiler_dir
146 - if profiler_dir and self.profiler != "torch":
176 + torch_profiler_dir = self.torch_profiler_dir
177 + if torch_profiler_dir and self.profiler != "torch":
178 raise ValueError(
179 "torch_profiler_dir is only applicable when profiler is set to 'torch'"
180 )
150 - if self.profiler == "torch" and not profiler_dir:
181 + if self.profiler == "torch" and not torch_profiler_dir:
182 raise ValueError("torch_profiler_dir must be set when profiler is 'torch'")
183
184 # Support any URI scheme (gs://, s3://, hdfs://, etc.)
185 # These paths should not be converted to absolute paths
155 - if profiler_dir and not _is_uri_path(profiler_dir):
156 - self.torch_profiler_dir = os.path.abspath(os.path.expanduser(profiler_dir))
186 + if torch_profiler_dir and not _is_uri_path(torch_profiler_dir):
187 + self.torch_profiler_dir = os.path.abspath(
188 + os.path.expanduser(torch_profiler_dir)
189 + )
190 +
191 + proton_profiler_dir = self.proton_profiler_dir
192 + non_default_proton_options = [
193 + name
194 + for name, value, default in (
195 + ("proton_profiler_dir", proton_profiler_dir, ""),
196 + ("proton_context", self.proton_context, "shadow"),
197 + ("proton_data", self.proton_data, "tree"),
198 + ("proton_backend", self.proton_backend, None),
199 + ("proton_mode", self.proton_mode, None),
200 + ("proton_hook", self.proton_hook, None),
201 + ("proton_output_format", self.proton_output_format, None),
202 + )
203 + if value != default
204 + ]
205 + if self.profiler != "proton" and non_default_proton_options:
206 + options = ", ".join(non_default_proton_options)
207 + raise ValueError(
208 + f"{options} only applicable when profiler is set to 'proton'"
209 + )
210 + if self.profiler == "proton" and not proton_profiler_dir:
211 + raise ValueError(
212 + "proton_profiler_dir must be set when profiler is 'proton'"
213 + )
214 + if proton_profiler_dir:
215 + if _is_uri_path(proton_profiler_dir):
216 + raise ValueError("proton_profiler_dir must be a local directory")
217 + self.proton_profiler_dir = os.path.abspath(
218 + os.path.expanduser(proton_profiler_dir)
219 + )
220 +
221 + if self.profiler == "proton":
222 + output_format = self.proton_output_format
223 + if output_format == "chrome_trace" and self.proton_data != "trace":
224 + raise ValueError("chrome_trace output requires proton_data='trace'")
225 + if (
226 + output_format in ("hatchet", "hatchet_msgpack")
227 + and self.proton_data != "tree"
228 + ):
229 + raise ValueError(f"{output_format} output requires proton_data='tree'")
230
231 if self.capture_torch_profiler and self.profiler != "torch":
232 raise ValueError(
vllm/config/vllm.py
+12
@@ -1269,6 +1269,18 @@ class VllmConfig:
1269 self.compilation_config.mode = CompilationMode.NONE
1270 self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE
1271
1272 + if self.profiler_config.profiler == "proton":
1273 + if not current_platform.is_cuda():
1274 + raise ValueError(
1275 + "The Proton profiler currently supports NVIDIA CUDA only"
1276 + )
1277 + if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
1278 + raise ValueError(
1279 + "The Proton profiler requires CUDA graphs to be disabled. "
1280 + "Use --enforce-eager or set "
1281 + "--compilation-config.cudagraph_mode=none."
1282 + )
1283 +
1284 if os.environ.get("TORCH_COMPILE_DISABLE") == "1":
1285 logger.warning_once(
1286 "TORCH_COMPILE_DISABLE is set, disabling torch.compile. "
vllm/profiler/wrapper.py
+129
@@ -1,13 +1,18 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import importlib
5 +import inspect
6 import json
7 +import os
8 from abc import ABC, abstractmethod
9 from collections.abc import Callable
10 from contextlib import nullcontext
11 from typing import Literal
12 +from uuid import uuid4
13
14 import torch
15 +from packaging.version import InvalidVersion, Version
16 from typing_extensions import override
17
18 import vllm.version
@@ -17,6 +22,8 @@ from vllm.logger import init_logger
22
23 logger = init_logger(__name__)
24
25 +_TRITON_PROTON_3_7_VERSION = Version("3.7.0")
26 +
27
28 class WorkerProfiler(ABC):
29 def __init__(self, profiler_config: ProfilerConfig) -> None:
@@ -43,6 +50,11 @@ class WorkerProfiler(ABC):
50 self._profiling_for_iters = 0
51 self._running = False
52
53 + @property
54 + def is_running(self) -> bool:
55 + """Whether the underlying profiler is currently collecting data."""
56 + return self._running
57 +
58 @abstractmethod
59 def _start(self) -> None:
60 """Start the profiler."""
@@ -340,6 +352,123 @@ class TorchProfilerWrapper(WorkerProfiler):
352 return torch.profiler.record_function(name)
353
354
355 +class ProtonProfilerWrapper(WorkerProfiler):
356 + """Worker profiler backed by :mod:`triton.profiler` (Proton)."""
357 +
358 + def __init__(
359 + self,
360 + profiler_config: ProfilerConfig,
361 + worker_name: str,
362 + ) -> None:
363 + super().__init__(profiler_config)
364 +
365 + try:
366 + self._proton = importlib.import_module("triton.profiler")
367 + triton = importlib.import_module("triton")
368 + except ImportError as exc:
369 + raise RuntimeError(
370 + "The Proton profiler requires a Triton installation with "
371 + "triton.profiler support."
372 + ) from exc
373 +
374 + self._output_dir = profiler_config.proton_profiler_dir
375 + self._output_path = os.path.join(self._output_dir, f"proton_{worker_name}")
376 + self._context = profiler_config.proton_context
377 + self._data = profiler_config.proton_data
378 + self._backend = profiler_config.proton_backend
379 + self._mode = profiler_config.proton_mode
380 + self._hook = profiler_config.proton_hook
381 + self._output_format = profiler_config.proton_output_format
382 + self._triton_version_string = getattr(triton, "__version__", "unknown")
383 + try:
384 + self._triton_version = Version(self._triton_version_string)
385 + except InvalidVersion:
386 + self._triton_version = None
387 + self._validate_capabilities()
388 + self._session_id: int | None = None
389 + # Qualify output names by process and wrapper instance so a new
390 + # worker cannot overwrite profiles left by an earlier server process.
391 + self._instance_id = f"pid{os.getpid()}_{uuid4().hex}"
392 + self._run_id = 0
393 +
394 + logger.info_once(
395 + "Proton profiling enabled. Output will be saved under: %s",
396 + self._output_dir,
397 + )
398 +
399 + def _require_triton_version(self, feature: str, minimum: Version) -> None:
400 + if self._triton_version is None or self._triton_version < minimum:
401 + raise RuntimeError(
402 + f"Proton {feature} requires Triton >= {minimum}; found "
403 + f"{self._triton_version_string}."
404 + )
405 +
406 + def _validate_capabilities(self) -> None:
407 + if self._output_format is not None:
408 + parameters = inspect.signature(self._proton.finalize).parameters
409 + supports_output_format = "output_format" in parameters or any(
410 + parameter.kind == inspect.Parameter.VAR_KEYWORD
411 + for parameter in parameters.values()
412 + )
413 + if not supports_output_format:
414 + raise RuntimeError(
415 + "The installed Triton Proton does not support selecting "
416 + "an output format during finalize."
417 + )
418 +
419 + if self._output_format == "hatchet_msgpack":
420 + self._require_triton_version(
421 + "hatchet_msgpack output", _TRITON_PROTON_3_7_VERSION
422 + )
423 + if self._mode and self._mode.split(":", 1)[0] == "periodic_flushing":
424 + self._require_triton_version(
425 + "periodic flushing", _TRITON_PROTON_3_7_VERSION
426 + )
427 +
428 + def _create_session(self, output_path: str) -> int:
429 + os.makedirs(self._output_dir, exist_ok=True)
430 + session_id = self._proton.start(
431 + name=output_path,
432 + context=self._context,
433 + data=self._data,
434 + backend=self._backend,
435 + mode=self._mode,
436 + hook=self._hook,
437 + )
438 + if session_id is None:
439 + raise RuntimeError("Proton did not create a profiling session")
440 + return session_id
441 +
442 + @override
443 + def _start(self) -> None:
444 + output_path = f"{self._output_path}_{self._instance_id}_run{self._run_id}"
445 + self._session_id = self._create_session(output_path)
446 + self._run_id += 1
447 +
448 + @override
449 + def _stop(self) -> None:
450 + assert self._session_id is not None
451 + session_id = self._session_id
452 + try:
453 + self._proton.deactivate(session=session_id)
454 + finally:
455 + try:
456 + if self._output_format is None:
457 + self._proton.finalize(session=session_id)
458 + else:
459 + self._proton.finalize(
460 + session=session_id, output_format=self._output_format
461 + )
462 + finally:
463 + self._session_id = None
464 +
465 + @override
466 + def annotate_context_manager(self, name: str):
467 + if not self._running:
468 + return nullcontext()
469 + return self._proton.scope(name)
470 +
471 +
472 class CudaProfilerWrapper(WorkerProfiler):
473 def __init__(self, profiler_config: ProfilerConfig) -> None:
474 super().__init__(profiler_config)
vllm/v1/worker/gpu_worker.py
+23 -10
@@ -55,7 +55,11 @@ from vllm.lora.request import LoRARequest
55 from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
56 from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory
57 from vllm.platforms import current_platform
58 -from vllm.profiler.wrapper import CudaProfilerWrapper, TorchProfilerWrapper
58 +from vllm.profiler.wrapper import (
59 + CudaProfilerWrapper,
60 + ProtonProfilerWrapper,
61 + TorchProfilerWrapper,
62 +)
63 from vllm.sequence import IntermediateTensors
64 from vllm.tasks import SupportedTask
65 from vllm.tracing import instrument
@@ -162,16 +166,12 @@ class Worker(WorkerBase):
166 self._weight_update_active = False
167 self._weight_update_is_draft = False
168
165 - # Torch/CUDA profiler. Enabled and configured through profiler_config.
169 + # Worker profiler. Enabled and configured through profiler_config.
170 # Profiler wrapper is created lazily in profile() when start is called,
171 # so we have all the information needed for proper trace naming.
172 self.profiler: Any | None = None
173 self.profiler_config = vllm_config.profiler_config
174
171 - # Only validate profiler config is valid, don't instantiate yet
172 - if self.profiler_config.profiler not in ("torch", "cuda", None):
173 - raise ValueError(f"Unknown profiler type: {self.profiler_config.profiler}")
174 -
175 self.use_v2_model_runner = vllm_config.use_v2_model_runner
176 # pending non-blocking PP send work from the previous iteration
177 self._pp_send_work: list[Handle] = []
@@ -909,6 +909,8 @@ class Worker(WorkerBase):
909 return nullcontext()
910
911 self.profiler.step()
912 + if not self.profiler.is_running:
913 + return nullcontext()
914
915 iteration_details = compute_iteration_details(scheduler_output)
916
@@ -1125,6 +1127,7 @@ class Worker(WorkerBase):
1127 )
1128
1129 if is_start:
1130 + profiler_type = self.profiler_config.profiler
1131 # Generate the trace name by combining prefix with comprehensive rank suffix
1132 from vllm.distributed.utils import get_worker_rank_suffix
1133
@@ -1138,7 +1141,6 @@ class Worker(WorkerBase):
1141
1142 # Create the profiler wrapper only on the first start call
1143 if self.profiler is None:
1141 - profiler_type = self.profiler_config.profiler
1144 if profiler_type == "torch":
1145 self.profiler = TorchProfilerWrapper(
1146 self.profiler_config,
@@ -1152,20 +1154,31 @@ class Worker(WorkerBase):
1154 elif profiler_type == "cuda":
1155 self.profiler = CudaProfilerWrapper(self.profiler_config)
1156 logger.debug("Starting CUDA profiler")
1157 + elif profiler_type == "proton":
1158 + self.profiler = ProtonProfilerWrapper(
1159 + self.profiler_config, worker_name=trace_name
1160 + )
1161 + logger.debug(
1162 + "Starting Proton profiler with trace name: %s", trace_name
1163 + )
1164 else:
1165 # Config validation should prevent this code being reached
1166 raise ValueError(
1167 f"Invalid profiler value of {self.profiler_config.profiler}"
1168 )
1169
1161 - # If profiler already initialized, restart profiling but keep
1162 - # the original trace name from the first initialization.
1170 self.profiler.start()
1171 else:
1172 if self.profiler is None:
1173 logger.warning("Profiler was not started, nothing to stop.")
1174 return
1168 - self.profiler.stop()
1175 + try:
1176 + self.profiler.stop()
1177 + finally:
1178 + if self.profiler_config.profiler == "proton":
1179 + # Proton output names are fixed when the wrapper is constructed.
1180 + # Recreate it so the next profile_prefix is honored.
1181 + self.profiler = None
1182
1183 def execute_dummy_batch(self) -> None:
1184 num_tokens = getattr(self.model_runner, "uniform_decode_query_len", 1)