From a367cdbca0824c8c428fe15b4552fbc4b824d718 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=90=98=E5=A4=A9=E6=A5=BD?= Date: Tue, 11 Aug 2026 14:51:57 -0700 Subject: [PATCH] [Profiler] Add minimal Triton Proton profiling backend (#48789) Signed-off-by: Luossu --- docs/contributing/profiling.md | 64 ++++++ tests/v1/worker/test_gpu_profiler.py | 327 ++++++++++++++++++++++++++- vllm/benchmarks/latency.py | 5 + vllm/config/profiler.py | 87 ++++++- vllm/config/vllm.py | 12 + vllm/profiler/wrapper.py | 129 +++++++++++ vllm/v1/worker/gpu_worker.py | 33 ++- 7 files changed, 637 insertions(+), 20 deletions(-) diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index c9bd0e5bdd..e420a3aa68 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -84,6 +84,70 @@ curl -X POST http://localhost:8000/v1/chat/completions \ $ curl -X POST http://localhost:8000/stop_profile ``` +## Profile with Triton Proton + +[Proton](https://github.com/triton-lang/triton/tree/main/third_party/proton) +is Triton's GPU profiler. It can collect a low-overhead aggregate tree or a +Chrome trace and works through the same vLLM profiling controls as the PyTorch +and CUDA profilers. Proton currently supports NVIDIA GPUs through CUPTI and +requires eager execution. + +Start a server with a local output directory: + +```bash +vllm serve meta-llama/Llama-3.1-8B-Instruct \ + --enforce-eager \ + --profiler-config '{ + "profiler": "proton", + "proton_profiler_dir": "./proton_profile", + "proton_output_format": "hatchet", + "proton_hook": "triton" + }' +``` + +Then use `/start_profile` and `/stop_profile` as shown above, or pass +`--profile` to a vLLM benchmark. Each worker uses a topology- and +rank-qualified output name, such as +`proton_dp0_pp0_tp0_dcp0_ep0_rank0_pid1234_0123456789abcdef0123456789abcdef_run0.hatchet`, +so distributed workers, restarted servers, and repeated profiling runs do not +overwrite one another. A `profile_prefix` is included when supplied. Each +profile is finalized by `/stop_profile` and is ready to inspect immediately. + +The Proton-specific options are: + +- `proton_context`: `shadow` (default) or `python` +- `proton_data`: `tree` (default) or `trace` +- `proton_backend`: `cupti` or automatic +- `proton_mode`: an optional backend mode string +- `proton_hook`: `triton` to record Triton launch metadata, or unset +- `proton_output_format`: `hatchet`, `hatchet_msgpack`, `chrome_trace`, or unset + +`hatchet` and `hatchet_msgpack` require `proton_data: "tree"`, while +`chrome_trace` requires `proton_data: "trace"`. CUDA graph profiling is not yet +supported, so `--enforce-eager` is required for every Proton mode. + +Automatic backend selection is recommended. vLLM currently supports Proton's +`cupti` backend on NVIDIA GPUs. ROCm support is not yet available. vLLM does not +expose Proton's experimental instrumentation backend because current upstream +Triton builds can produce profiles without timing metrics. Backend-specific +modes, including `pcsampling`, can be selected with `proton_mode`. + +Triton 3.6 supports explicit `hatchet` and `chrome_trace` output. The +`hatchet_msgpack` format and `periodic_flushing` mode require Triton 3.7 or +newer; vLLM rejects these options with the detected version before starting a +session. + +Inspect tree profiles with: + +```bash +proton-viewer -m time/ns \ + proton_profile/proton_dp0_pp0_tp0_dcp0_ep0_rank0_pid1234_0123456789abcdef0123456789abcdef_run0.hatchet +``` + +Chrome traces (`proton_data: "trace"`) can be opened in +. Proton is imported lazily, so selecting another +profiler does not require a Proton-capable Triton installation. + ## Profile with NVIDIA Nsight Systems 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. diff --git a/tests/v1/worker/test_gpu_profiler.py b/tests/v1/worker/test_gpu_profiler.py index 8ff89354ea..3e236a73b1 100644 --- a/tests/v1/worker/test_gpu_profiler.py +++ b/tests/v1/worker/test_gpu_profiler.py @@ -1,12 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from unittest.mock import MagicMock +import os +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import MagicMock, Mock, call, patch +from uuid import UUID import pytest +from pydantic import ValidationError -from vllm.config import CUDAGraphMode, ProfilerConfig +from vllm.config import ( + CompilationConfig, + CUDAGraphMode, + ProfilerConfig, + VllmConfig, +) from vllm.config.profiler import _is_uri_path -from vllm.profiler.wrapper import WorkerProfiler +from vllm.platforms import current_platform +from vllm.profiler.wrapper import ProtonProfilerWrapper, WorkerProfiler from vllm.v1.core.sched.output import CachedRequestData from vllm.v1.worker.gpu_model_runner import GPUModelRunner from vllm.v1.worker.gpu_worker import Worker @@ -281,6 +292,16 @@ class TestAnnotateProfile: "execute_5_context_1(sq4sk4sqsq16sqsk16)_generation_1(sq1sk11sqsq1sqsk11)" ) + def test_skips_annotation_work_after_profiler_stops(self): + worker = MagicMock() + worker.profiler.is_running = False + + context = Worker.annotate_profile(worker, scheduler_output=None) + + worker.profiler.step.assert_called_once_with() + worker.profiler.annotate_context_manager.assert_not_called() + assert isinstance(context, nullcontext) + def test_profiler_entered_during_capture(): """Profiler is used as a context manager in _warmup_and_capture, @@ -298,3 +319,303 @@ def test_profiler_entered_during_capture(): mock_profiler.__enter__.assert_called_once() mock_profiler.__exit__.assert_called_once() + + +def make_proton(session_id: int | None = 7): + return SimpleNamespace( + start=Mock(return_value=session_id), + activate=Mock(), + deactivate=Mock(), + finalize=Mock(), + scope=Mock(return_value=nullcontext()), + ) + + +def make_proton_wrapper( + tmp_path, proton=None, triton_version="3.6.0", **config_overrides +): + proton = proton or make_proton() + config = ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + **config_overrides, + ) + + def import_module(name): + if name == "triton.profiler": + return proton + assert name == "triton" + return SimpleNamespace(__version__=triton_version) + + with patch( + "vllm.profiler.wrapper.importlib.import_module", side_effect=import_module + ): + wrapper = ProtonProfilerWrapper(config, worker_name="rank_3") + return wrapper, proton + + +_requires_cuda_for_proton = pytest.mark.skipif( + not current_platform.is_cuda(), + reason="Proton profiling tests require an NVIDIA CUDA platform.", +) + + +@_requires_cuda_for_proton +class TestProtonConfig: + def test_normalizes_local_output_directory(self, tmp_path, monkeypatch): + monkeypatch.chdir(tmp_path) + config = ProfilerConfig(profiler="proton", proton_profiler_dir="profiles") + assert config.proton_profiler_dir == os.path.join(tmp_path, "profiles") + + @pytest.mark.parametrize( + ("options", "message"), + [ + ({"proton_profiler_dir": ""}, "must be set"), + ({"proton_profiler_dir": "s3://bucket/profiles"}, "local directory"), + ( + {"proton_data": "tree", "proton_output_format": "chrome_trace"}, + "requires proton_data", + ), + ( + {"proton_data": "trace", "proton_output_format": "hatchet"}, + "requires proton_data", + ), + ( + { + "proton_data": "trace", + "proton_output_format": "hatchet_msgpack", + }, + "requires proton_data", + ), + ], + ) + def test_rejects_invalid_option_combinations(self, tmp_path, options, message): + kwargs = {"proton_profiler_dir": str(tmp_path), **options} + with pytest.raises(ValueError, match=message): + ProfilerConfig(profiler="proton", **kwargs) + + @pytest.mark.parametrize( + "field", + [ + "proton_context", + "proton_data", + "proton_backend", + "proton_hook", + "proton_output_format", + ], + ) + def test_rejects_invalid_typed_options(self, field, tmp_path): + with pytest.raises(ValidationError): + ProfilerConfig( + profiler="proton", + proton_profiler_dir=str(tmp_path), + **{field: "invalid"}, + ) + + @pytest.mark.parametrize( + ("field", "value"), + [ + ("proton_profiler_dir", "profiles"), + ("proton_context", "python"), + ("proton_data", "trace"), + ("proton_backend", "cupti"), + ("proton_mode", "pcsampling"), + ("proton_hook", "triton"), + ("proton_output_format", "chrome_trace"), + ], + ) + def test_rejects_proton_options_for_other_profilers(self, field, value): + with pytest.raises(ValueError, match=f"{field} only applicable"): + ProfilerConfig(**{field: value}) + + def test_allows_proton_when_cuda_graphs_are_disabled(self, tmp_path): + config = VllmConfig( + profiler_config=ProfilerConfig( + profiler="proton", proton_profiler_dir=str(tmp_path) + ), + compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE), + ) + + assert config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE + + def test_rejects_proton_on_non_cuda_platforms(self, tmp_path): + with ( + patch("vllm.platforms.current_platform.is_cuda", return_value=False), + pytest.raises(ValueError, match="supports NVIDIA CUDA only"), + ): + VllmConfig( + profiler_config=ProfilerConfig( + profiler="proton", proton_profiler_dir=str(tmp_path) + ), + compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE), + ) + + def test_rejects_proton_when_cuda_graphs_are_enabled(self, tmp_path): + with pytest.raises(ValueError, match="requires CUDA graphs to be disabled"): + VllmConfig( + profiler_config=ProfilerConfig( + profiler="proton", proton_profiler_dir=str(tmp_path) + ), + compilation_config=CompilationConfig( + cudagraph_mode=CUDAGraphMode.PIECEWISE + ), + ) + + +@_requires_cuda_for_proton +class TestProtonProfilerWrapper: + def test_passes_config_and_global_rank_name_to_proton(self, tmp_path): + wrapper, proton = make_proton_wrapper( + tmp_path, + proton_context="python", + proton_data="trace", + proton_backend="cupti", + proton_mode="pcsampling", + proton_hook="triton", + proton_output_format="chrome_trace", + ) + + wrapper.start() + + start_args = proton.start.call_args.kwargs + assert start_args["name"].startswith(os.path.join(tmp_path, "proton_rank_3_")) + assert start_args["name"].endswith("_run0") + assert start_args | {"name": None} == { + "name": None, + "context": "python", + "data": "trace", + "backend": "cupti", + "mode": "pcsampling", + "hook": "triton", + } + wrapper.stop() + proton.finalize.assert_called_once_with(session=7, output_format="chrome_trace") + assert tmp_path.is_dir() + + def test_finalizes_each_profile_with_unique_output_names(self, tmp_path): + proton = make_proton() + proton.start.side_effect = [7, 8] + wrapper, proton = make_proton_wrapper(tmp_path, proton) + + wrapper.start() + wrapper.start() + wrapper.stop() + wrapper.start() + wrapper.stop() + + names = [c.kwargs["name"] for c in proton.start.call_args_list] + assert len(names) == len(set(names)) == 2 + assert names[0].endswith("_run0") + assert names[1].endswith("_run1") + assert proton.deactivate.call_count == 2 + assert proton.finalize.call_args_list == [call(session=7), call(session=8)] + + def test_output_names_are_unique_across_worker_restarts(self, tmp_path): + with patch( + "vllm.profiler.wrapper.uuid4", + side_effect=[UUID(int=1), UUID(int=2)], + ): + first, first_proton = make_proton_wrapper(tmp_path) + second, second_proton = make_proton_wrapper(tmp_path) + + first.start() + second.start() + + first_name = first_proton.start.call_args.kwargs["name"] + second_name = second_proton.start.call_args.kwargs["name"] + assert first_name != second_name + assert first_name.endswith(f"_{UUID(int=1).hex}_run0") + assert second_name.endswith(f"_{UUID(int=2).hex}_run0") + + @pytest.mark.parametrize( + ("option", "value", "feature"), + [ + ("proton_output_format", "hatchet_msgpack", "hatchet_msgpack"), + ("proton_mode", "periodic_flushing", "periodic flushing"), + ], + ) + def test_newer_features_reject_triton_3_6(self, tmp_path, option, value, feature): + with pytest.raises(RuntimeError, match=feature): + make_proton_wrapper(tmp_path, **{option: value}) + + @pytest.mark.parametrize( + ("option", "value"), + [ + ("proton_output_format", "hatchet_msgpack"), + ("proton_mode", "periodic_flushing"), + ], + ) + def test_triton_3_7_features(self, tmp_path, option, value): + make_proton_wrapper(tmp_path, triton_version="3.7.0", **{option: value}) + + def test_rejects_output_format_when_finalize_lacks_capability(self, tmp_path): + proton = make_proton() + proton.finalize = lambda session=None: None + + with pytest.raises(RuntimeError, match="does not support selecting"): + make_proton_wrapper(tmp_path, proton, proton_output_format="hatchet") + + def test_missing_proton_has_actionable_error(self, tmp_path): + config = ProfilerConfig(profiler="proton", proton_profiler_dir=str(tmp_path)) + with ( + patch( + "vllm.profiler.wrapper.importlib.import_module", + side_effect=ImportError, + ), + pytest.raises(RuntimeError, match="requires a Triton installation"), + ): + ProtonProfilerWrapper(config, worker_name="rank_0") + + def test_scope_annotations_delegate_to_proton(self, tmp_path): + wrapper, proton = make_proton_wrapper(tmp_path) + wrapper.start() + + context = wrapper.annotate_context_manager("decode") + + proton.scope.assert_called_once_with("decode") + assert context is not None + + +@_requires_cuda_for_proton +def test_gpu_worker_creates_proton_profiler(): + worker = MagicMock() + worker.rank = 1 + worker.profiler = None + worker.profiler_config.profiler = "proton" + + with ( + patch( + "vllm.distributed.utils.get_worker_rank_suffix", + return_value="rank1", + ), + patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper, + ): + Worker.profile(worker) + + wrapper.assert_called_once_with(worker.profiler_config, worker_name="rank1") + worker.profiler.start.assert_called_once_with() + + +@_requires_cuda_for_proton +def test_gpu_worker_recreates_proton_profiler_for_each_run(): + worker = MagicMock() + worker.rank = 1 + worker.profiler = None + worker.profiler_config.profiler = "proton" + + with ( + patch( + "vllm.distributed.utils.get_worker_rank_suffix", + return_value="rank1", + ), + patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper, + ): + Worker.profile(worker, profile_prefix="first") + Worker.profile(worker, is_start=False) + Worker.profile(worker, profile_prefix="second") + + assert wrapper.call_args_list == [ + call(worker.profiler_config, worker_name="first_rank1"), + call(worker.profiler_config, worker_name="second_rank1"), + ] + assert wrapper.return_value.start.call_count == 2 diff --git a/vllm/benchmarks/latency.py b/vllm/benchmarks/latency.py index cce531629f..6209f5f2d8 100644 --- a/vllm/benchmarks/latency.py +++ b/vllm/benchmarks/latency.py @@ -147,6 +147,11 @@ def main(args: argparse.Namespace): ) elif profiler_config.profiler == "cuda": print("Profiling with cuda profiler ...") + elif profiler_config.profiler == "proton": + print( + "Profiling with Proton (results will be saved under" + f" {profiler_config.proton_profiler_dir})..." + ) run_to_completion(do_profile=True) return diff --git a/vllm/config/profiler.py b/vllm/config/profiler.py index f0e29d08f3..b3b232490f 100644 --- a/vllm/config/profiler.py +++ b/vllm/config/profiler.py @@ -13,7 +13,12 @@ from vllm.utils.hashing import safe_hash logger = init_logger(__name__) -ProfilerKind = Literal["torch", "cuda"] +ProfilerKind = Literal["torch", "cuda", "proton"] +ProtonBackend = Literal["cupti"] +ProtonContext = Literal["shadow", "python"] +ProtonData = Literal["tree", "trace"] +ProtonHook = Literal["triton"] +ProtonOutputFormat = Literal["hatchet", "hatchet_msgpack", "chrome_trace"] def _is_uri_path(path: str) -> bool: @@ -38,13 +43,39 @@ class ProfilerConfig: """Which profiler to use. Defaults to None. Options are: - 'torch': Use PyTorch profiler. - - 'cuda': Use CUDA profiler.""" + - 'cuda': Use CUDA profiler. + - 'proton': Use Triton Proton profiler.""" torch_profiler_dir: str = "" """Directory to save torch profiler traces. Both AsyncLLM's CPU traces and worker's traces (CPU & GPU) will be saved under this directory. Note that it must be an absolute path.""" + proton_profiler_dir: str = "" + """Directory to save Triton Proton profiles. Each worker writes a + separate rank-qualified file.""" + + proton_context: ProtonContext = "shadow" + """Proton context source. ``shadow`` records explicit scopes with low + overhead; ``python`` records Python call stacks.""" + + proton_data: ProtonData = "tree" + """Proton output type. ``tree`` produces Hatchet data and ``trace`` + produces a Chrome trace.""" + + proton_backend: ProtonBackend | None = None + """Proton GPU backend. ``None`` lets Proton select CUPTI automatically.""" + + proton_mode: str | None = None + """Optional backend-specific Proton mode string, such as ``pcsampling``.""" + + proton_hook: ProtonHook | None = None + """Optional Proton hook. Use ``triton`` to add Triton launch metadata.""" + + proton_output_format: ProtonOutputFormat | None = None + """Optional format passed to Proton when finalizing a profile. ``None`` + uses the default format for ``proton_data``.""" + torch_profiler_with_stack: bool = True """If `True`, enables stack tracing in the torch profiler. Enabled by default as it is useful for debugging. Can be disabled via @@ -142,18 +173,60 @@ class ProfilerConfig: "while ignore_frontend is False may result in high overhead." ) - profiler_dir = self.torch_profiler_dir - if profiler_dir and self.profiler != "torch": + torch_profiler_dir = self.torch_profiler_dir + if torch_profiler_dir and self.profiler != "torch": raise ValueError( "torch_profiler_dir is only applicable when profiler is set to 'torch'" ) - if self.profiler == "torch" and not profiler_dir: + if self.profiler == "torch" and not torch_profiler_dir: raise ValueError("torch_profiler_dir must be set when profiler is 'torch'") # Support any URI scheme (gs://, s3://, hdfs://, etc.) # These paths should not be converted to absolute paths - if profiler_dir and not _is_uri_path(profiler_dir): - self.torch_profiler_dir = os.path.abspath(os.path.expanduser(profiler_dir)) + if torch_profiler_dir and not _is_uri_path(torch_profiler_dir): + self.torch_profiler_dir = os.path.abspath( + os.path.expanduser(torch_profiler_dir) + ) + + proton_profiler_dir = self.proton_profiler_dir + non_default_proton_options = [ + name + for name, value, default in ( + ("proton_profiler_dir", proton_profiler_dir, ""), + ("proton_context", self.proton_context, "shadow"), + ("proton_data", self.proton_data, "tree"), + ("proton_backend", self.proton_backend, None), + ("proton_mode", self.proton_mode, None), + ("proton_hook", self.proton_hook, None), + ("proton_output_format", self.proton_output_format, None), + ) + if value != default + ] + if self.profiler != "proton" and non_default_proton_options: + options = ", ".join(non_default_proton_options) + raise ValueError( + f"{options} only applicable when profiler is set to 'proton'" + ) + if self.profiler == "proton" and not proton_profiler_dir: + raise ValueError( + "proton_profiler_dir must be set when profiler is 'proton'" + ) + if proton_profiler_dir: + if _is_uri_path(proton_profiler_dir): + raise ValueError("proton_profiler_dir must be a local directory") + self.proton_profiler_dir = os.path.abspath( + os.path.expanduser(proton_profiler_dir) + ) + + if self.profiler == "proton": + output_format = self.proton_output_format + if output_format == "chrome_trace" and self.proton_data != "trace": + raise ValueError("chrome_trace output requires proton_data='trace'") + if ( + output_format in ("hatchet", "hatchet_msgpack") + and self.proton_data != "tree" + ): + raise ValueError(f"{output_format} output requires proton_data='tree'") if self.capture_torch_profiler and self.profiler != "torch": raise ValueError( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index e03f70a295..2bbfd57330 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1269,6 +1269,18 @@ class VllmConfig: self.compilation_config.mode = CompilationMode.NONE self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE + if self.profiler_config.profiler == "proton": + if not current_platform.is_cuda(): + raise ValueError( + "The Proton profiler currently supports NVIDIA CUDA only" + ) + if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE: + raise ValueError( + "The Proton profiler requires CUDA graphs to be disabled. " + "Use --enforce-eager or set " + "--compilation-config.cudagraph_mode=none." + ) + if os.environ.get("TORCH_COMPILE_DISABLE") == "1": logger.warning_once( "TORCH_COMPILE_DISABLE is set, disabling torch.compile. " diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index bbdec955a0..2f56d0fc7c 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -1,13 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib +import inspect import json +import os from abc import ABC, abstractmethod from collections.abc import Callable from contextlib import nullcontext from typing import Literal +from uuid import uuid4 import torch +from packaging.version import InvalidVersion, Version from typing_extensions import override import vllm.version @@ -17,6 +22,8 @@ from vllm.logger import init_logger logger = init_logger(__name__) +_TRITON_PROTON_3_7_VERSION = Version("3.7.0") + class WorkerProfiler(ABC): def __init__(self, profiler_config: ProfilerConfig) -> None: @@ -43,6 +50,11 @@ class WorkerProfiler(ABC): self._profiling_for_iters = 0 self._running = False + @property + def is_running(self) -> bool: + """Whether the underlying profiler is currently collecting data.""" + return self._running + @abstractmethod def _start(self) -> None: """Start the profiler.""" @@ -340,6 +352,123 @@ class TorchProfilerWrapper(WorkerProfiler): return torch.profiler.record_function(name) +class ProtonProfilerWrapper(WorkerProfiler): + """Worker profiler backed by :mod:`triton.profiler` (Proton).""" + + def __init__( + self, + profiler_config: ProfilerConfig, + worker_name: str, + ) -> None: + super().__init__(profiler_config) + + try: + self._proton = importlib.import_module("triton.profiler") + triton = importlib.import_module("triton") + except ImportError as exc: + raise RuntimeError( + "The Proton profiler requires a Triton installation with " + "triton.profiler support." + ) from exc + + self._output_dir = profiler_config.proton_profiler_dir + self._output_path = os.path.join(self._output_dir, f"proton_{worker_name}") + self._context = profiler_config.proton_context + self._data = profiler_config.proton_data + self._backend = profiler_config.proton_backend + self._mode = profiler_config.proton_mode + self._hook = profiler_config.proton_hook + self._output_format = profiler_config.proton_output_format + self._triton_version_string = getattr(triton, "__version__", "unknown") + try: + self._triton_version = Version(self._triton_version_string) + except InvalidVersion: + self._triton_version = None + self._validate_capabilities() + self._session_id: int | None = None + # Qualify output names by process and wrapper instance so a new + # worker cannot overwrite profiles left by an earlier server process. + self._instance_id = f"pid{os.getpid()}_{uuid4().hex}" + self._run_id = 0 + + logger.info_once( + "Proton profiling enabled. Output will be saved under: %s", + self._output_dir, + ) + + def _require_triton_version(self, feature: str, minimum: Version) -> None: + if self._triton_version is None or self._triton_version < minimum: + raise RuntimeError( + f"Proton {feature} requires Triton >= {minimum}; found " + f"{self._triton_version_string}." + ) + + def _validate_capabilities(self) -> None: + if self._output_format is not None: + parameters = inspect.signature(self._proton.finalize).parameters + supports_output_format = "output_format" in parameters or any( + parameter.kind == inspect.Parameter.VAR_KEYWORD + for parameter in parameters.values() + ) + if not supports_output_format: + raise RuntimeError( + "The installed Triton Proton does not support selecting " + "an output format during finalize." + ) + + if self._output_format == "hatchet_msgpack": + self._require_triton_version( + "hatchet_msgpack output", _TRITON_PROTON_3_7_VERSION + ) + if self._mode and self._mode.split(":", 1)[0] == "periodic_flushing": + self._require_triton_version( + "periodic flushing", _TRITON_PROTON_3_7_VERSION + ) + + def _create_session(self, output_path: str) -> int: + os.makedirs(self._output_dir, exist_ok=True) + session_id = self._proton.start( + name=output_path, + context=self._context, + data=self._data, + backend=self._backend, + mode=self._mode, + hook=self._hook, + ) + if session_id is None: + raise RuntimeError("Proton did not create a profiling session") + return session_id + + @override + def _start(self) -> None: + output_path = f"{self._output_path}_{self._instance_id}_run{self._run_id}" + self._session_id = self._create_session(output_path) + self._run_id += 1 + + @override + def _stop(self) -> None: + assert self._session_id is not None + session_id = self._session_id + try: + self._proton.deactivate(session=session_id) + finally: + try: + if self._output_format is None: + self._proton.finalize(session=session_id) + else: + self._proton.finalize( + session=session_id, output_format=self._output_format + ) + finally: + self._session_id = None + + @override + def annotate_context_manager(self, name: str): + if not self._running: + return nullcontext() + return self._proton.scope(name) + + class CudaProfilerWrapper(WorkerProfiler): def __init__(self, profiler_config: ProfilerConfig) -> None: super().__init__(profiler_config) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index bee85c6eca..2340b64aba 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -55,7 +55,11 @@ from vllm.lora.request import LoRARequest from vllm.model_executor.warmup.kernel_warmup import kernel_warmup from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory from vllm.platforms import current_platform -from vllm.profiler.wrapper import CudaProfilerWrapper, TorchProfilerWrapper +from vllm.profiler.wrapper import ( + CudaProfilerWrapper, + ProtonProfilerWrapper, + TorchProfilerWrapper, +) from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask from vllm.tracing import instrument @@ -162,16 +166,12 @@ class Worker(WorkerBase): self._weight_update_active = False self._weight_update_is_draft = False - # Torch/CUDA profiler. Enabled and configured through profiler_config. + # Worker profiler. Enabled and configured through profiler_config. # Profiler wrapper is created lazily in profile() when start is called, # so we have all the information needed for proper trace naming. self.profiler: Any | None = None self.profiler_config = vllm_config.profiler_config - # Only validate profiler config is valid, don't instantiate yet - if self.profiler_config.profiler not in ("torch", "cuda", None): - raise ValueError(f"Unknown profiler type: {self.profiler_config.profiler}") - self.use_v2_model_runner = vllm_config.use_v2_model_runner # pending non-blocking PP send work from the previous iteration self._pp_send_work: list[Handle] = [] @@ -909,6 +909,8 @@ class Worker(WorkerBase): return nullcontext() self.profiler.step() + if not self.profiler.is_running: + return nullcontext() iteration_details = compute_iteration_details(scheduler_output) @@ -1125,6 +1127,7 @@ class Worker(WorkerBase): ) if is_start: + profiler_type = self.profiler_config.profiler # Generate the trace name by combining prefix with comprehensive rank suffix from vllm.distributed.utils import get_worker_rank_suffix @@ -1138,7 +1141,6 @@ class Worker(WorkerBase): # Create the profiler wrapper only on the first start call if self.profiler is None: - profiler_type = self.profiler_config.profiler if profiler_type == "torch": self.profiler = TorchProfilerWrapper( self.profiler_config, @@ -1152,20 +1154,31 @@ class Worker(WorkerBase): elif profiler_type == "cuda": self.profiler = CudaProfilerWrapper(self.profiler_config) logger.debug("Starting CUDA profiler") + elif profiler_type == "proton": + self.profiler = ProtonProfilerWrapper( + self.profiler_config, worker_name=trace_name + ) + logger.debug( + "Starting Proton profiler with trace name: %s", trace_name + ) else: # Config validation should prevent this code being reached raise ValueError( f"Invalid profiler value of {self.profiler_config.profiler}" ) - # If profiler already initialized, restart profiling but keep - # the original trace name from the first initialization. self.profiler.start() else: if self.profiler is None: logger.warning("Profiler was not started, nothing to stop.") return - self.profiler.stop() + try: + self.profiler.stop() + finally: + if self.profiler_config.profiler == "proton": + # Proton output names are fixed when the wrapper is constructed. + # Recreate it so the next profile_prefix is honored. + self.profiler = None def execute_dummy_batch(self) -> None: num_tokens = getattr(self.model_runner, "uniform_decode_query_len", 1)