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

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

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
+<https://ui.perfetto.dev/>. 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.

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

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
 

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(

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. "

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)

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)