vLLM/vllm/profiler/wrapper.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import importlib
import inspect
import json
import os
import threading
from abc import ABC, abstractmethod
from collections.abc import Callable, Iterator, Sequence
from contextlib import AbstractContextManager, contextmanager, nullcontext, suppress
from glob import glob
from uuid import uuid4

import torch
from packaging.version import InvalidVersion, Version
from typing_extensions import override

import vllm.version
from vllm.config import ProfilerConfig
from vllm.config.profiler import (
    ProfilerKind,
    TorchProfilerActivity,
    _is_uri_path,
    validate_profile_iteration_bounds,
)
from vllm.config.utils import replace
from vllm.logger import init_logger
from vllm.platforms import current_platform

logger = init_logger(__name__)

_TRITON_PROTON_3_7_VERSION = Version("3.7.0")


class WorkerProfiler(ABC):
    def __init__(self, profiler_config: ProfilerConfig) -> None:
        self._default_delay_iters = profiler_config.delay_iterations
        self._delay_iters = self._default_delay_iters
        if self._delay_iters > 0:
            logger.info_once(
                "GPU profiling will start "
                f"{self._delay_iters} steps after start_profile."
            )

        self._default_max_iters = profiler_config.max_iterations
        self._max_iters = self._default_max_iters
        if self._max_iters > 0:
            logger.info_once(
                "GPU profiling will stop "
                f"after {self._max_iters} worker steps, "
                "or when stop_profile is received."
            )

        # Serializes start()/stop(), which AsyncLLM may call from different threads.
        self._lifecycle_lock = threading.Lock()

        # Track when the profiler gets triggered by start_profile
        self._active_iteration_count = 0
        self._active = False

        # Track when the profiler is actually running
        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

    @property
    def should_annotate(self) -> bool:
        """Whether worker iterations should receive profiler annotations."""
        return self._running

    @abstractmethod
    def _start(self) -> None:
        """Start the profiler."""
        pass

    @abstractmethod
    def _stop(self) -> None:
        """Stop the profiler."""
        pass

    def _call_start(self) -> None:
        """Call _start with error handling but no safeguards."""
        try:
            self._start()
            self._running = True  # Only mark as running if start succeeds
        except Exception as e:
            logger.warning("Failed to start profiler: %s", e)

    def _call_stop(self) -> None:
        """Call _stop with error handling but no safeguards."""
        try:
            self._stop()
            logger.info_once("Profiler stopped successfully.")
        except Exception as e:
            logger.warning("Failed to stop profiler: %s", e)
        self._running = False  # Always mark as not running, assume stop worked

    def start(
        self,
        *,
        delay_iterations: int | None = None,
        max_iterations: int | None = None,
    ) -> None:
        """Attempt to start the profiler, accounting for delayed starts."""
        with self._lifecycle_lock:
            if self._active:
                logger.debug(
                    "start_profile received when profiler is already active. "
                    "Ignoring request."
                )
                return
            validate_profile_iteration_bounds(delay_iterations, max_iterations)
            self._delay_iters = (
                self._default_delay_iters
                if delay_iterations is None
                else delay_iterations
            )
            self._max_iters = (
                self._default_max_iters if max_iterations is None else max_iterations
            )
            self._active = True
            if self._delay_iters == 0:
                self._call_start()

    def step(self) -> None:
        """Update the profiler state at each worker step,
        to handle delayed starts and max iteration limits."""
        if not self._active:
            return

        self._active_iteration_count += 1

        if (
            not self._running
            and self._delay_iters > 0
            and self._active_iteration_count == self._delay_iters
        ):
            logger.info_once("Starting profiler after delay...")
            self._call_start()

        # Call profiler step for schedule-based profiling
        # Only count iterations where data is actually recorded (not warmup)
        if self._running and self._profiler_step():
            self._profiling_for_iters += 1

        if (
            self._max_iters > 0
            and self._running
            and self._profiling_for_iters > self._max_iters
        ):
            # Automatically stop the profiler after max iters. Go through the
            # public stop() (not _call_stop() directly) so _active and the
            # iteration counters reset too -- otherwise a later start_profile
            # is silently ignored forever, since start() bails out early
            # whenever _active is still True.
            logger.info_once("Max profiling iterations reached. Stopping profiler...")
            self.stop()
            return

    def _profiler_step(self) -> bool:
        """Called each step when profiler is running.
        Override in subclasses to handle schedule-based profiling.

        Returns:
            True if the step was an active profiling step (data recorded),
            False if the step was a warmup step (data discarded).

        """
        return True

    def stop(self) -> None:
        """Attempt to stop the profiler, accounting for overlapped calls."""
        with self._lifecycle_lock:
            if not self._active:
                logger.debug(
                    "stop_profile received when profiler is not active. "
                    "Ignoring request."
                )
                return
            self._active = False
            self._active_iteration_count = 0
            self._profiling_for_iters = 0

            if self._running:
                self._call_stop()

    def shutdown(self) -> None:
        """Ensure profiler is stopped when shutting down."""
        logger.info_once("Shutting down profiler")
        if self._running:
            self.stop()

    @property
    def has_cuda_graph_session(self) -> bool:
        """Whether a capture-time session is retained to attribute replays."""
        return False

    def capture_cuda_graphs(self) -> AbstractContextManager[None]:
        """Observe graph creation for backends that attribute replay activity."""
        return nullcontext()

    def set_output_name(self, worker_name: str) -> None:
        """Set the output name for the next profiling session."""
        return None

    def annotate_context_manager(self, name: str):
        """Return a context manager to annotate profiler traces."""
        return nullcontext()


TorchProfilerActivityMap = {
    "CPU": torch.profiler.ProfilerActivity.CPU,
    "CUDA": torch.profiler.ProfilerActivity.CUDA,
    "PrivateUse1": torch.profiler.ProfilerActivity.PrivateUse1,
    "XPU": torch.profiler.ProfilerActivity.XPU,
}


class TorchProfilerWrapper(WorkerProfiler):
    def __init__(
        self,
        profiler_config: ProfilerConfig,
        worker_name: str,
        local_rank: int,
        activities: Sequence[TorchProfilerActivity],
        on_trace_ready: Callable[[torch.profiler.profile], None] | None = None,
    ) -> None:
        super().__init__(profiler_config)

        self.local_rank = local_rank
        self.profiler_config = profiler_config
        torch_profiler_trace_dir = profiler_config.torch_profiler_dir
        self._torch_profiler_trace_dir = torch_profiler_trace_dir
        self._torch_profiler_use_gzip = profiler_config.torch_profiler_use_gzip
        self._worker_name = worker_name
        self._custom_trace_handler = on_trace_ready is not None
        if local_rank in (None, 0):
            logger.info_once(
                "Torch profiling enabled. Traces will be saved to: %s",
                torch_profiler_trace_dir,
            )
            logger.debug(
                "Profiler config: record_shapes=%s,"
                "profile_memory=%s,with_stack=%s,with_flops=%s",
                profiler_config.torch_profiler_record_shapes,
                profiler_config.torch_profiler_with_memory,
                profiler_config.torch_profiler_with_stack,
                profiler_config.torch_profiler_with_flops,
            )

        # Determine trace handler: use custom handler if provided,
        # otherwise default to tensorboard trace handler
        if on_trace_ready is not None:
            trace_handler = on_trace_ready
        else:
            trace_handler = torch.profiler.tensorboard_trace_handler(
                torch_profiler_trace_dir,
                worker_name=worker_name,
                use_gzip=profiler_config.torch_profiler_use_gzip,
            )

        self._records_cpu_activity = "CPU" in activities
        self.dump_device_time_total = (
            any(activity != "CPU" for activity in activities)
            and profiler_config.torch_profiler_dump_cuda_time_total
        )
        self.dump_cpu_time_total = self._records_cpu_activity and len(activities) == 1

        # Create profiler schedule if warmup or wait iterations are configured
        profiler_schedule = None
        if profiler_config.warmup_iterations > 0 or profiler_config.wait_iterations > 0:
            profiler_schedule = torch.profiler.schedule(
                skip_first=0,
                wait=profiler_config.wait_iterations,
                warmup=profiler_config.warmup_iterations,
                active=profiler_config.active_iterations,
                repeat=1,
            )
            if local_rank in (None, 0):
                logger.info_once(
                    "Profiler schedule configured: wait=%d, warmup=%d, active=%d",
                    profiler_config.wait_iterations,
                    profiler_config.warmup_iterations,
                    profiler_config.active_iterations,
                )

        self._profiler_kwargs = dict(
            activities=[TorchProfilerActivityMap[activity] for activity in activities],
            schedule=profiler_schedule,
            record_shapes=profiler_config.torch_profiler_record_shapes,
            profile_memory=profiler_config.torch_profiler_with_memory,
            with_stack=profiler_config.torch_profiler_with_stack,
            with_flops=profiler_config.torch_profiler_with_flops,
            on_trace_ready=trace_handler,
        )
        self.profiler: torch.profiler.profile

        # Track if we're using a schedule (need to call step())
        self._uses_schedule = profiler_schedule is not None
        self._warmup_iterations = profiler_config.warmup_iterations
        # Subtract 1 because profiler.start() already consumes step 0
        # (WAIT or WARMUP), so only wait + warmup - 1 non-active steps
        # remain to be advanced through via profiler.step() calls.
        self._initial_warmup_steps_remaining = max(
            profiler_config.wait_iterations + profiler_config.warmup_iterations - 1,
            0,
        )
        self._warmup_steps_remaining = self._initial_warmup_steps_remaining
        self._version_metadata_added = False

    @override
    def set_output_name(self, worker_name: str) -> None:
        if self._active:
            return
        self._worker_name = worker_name
        if self._custom_trace_handler:
            return
        self._profiler_kwargs["on_trace_ready"] = (
            torch.profiler.tensorboard_trace_handler(
                self._torch_profiler_trace_dir,
                worker_name=worker_name,
                use_gzip=self._torch_profiler_use_gzip,
            )
        )

    def _build_profiler_table(
        self,
        sort_key: str,
        row_limit: int | None = None,
    ) -> str:
        group_by_input_shape = (
            current_platform.is_cpu()
            and self.profiler_config.torch_profiler_record_shapes
        )
        averages = self.profiler.key_averages(group_by_input_shape=group_by_input_shape)
        if row_limit is None:  # use profiler default row limit of 100
            return averages.table(sort_by=sort_key)
        return averages.table(
            sort_by=sort_key,
            row_limit=row_limit,
        )

    def _write_profiler_table(self, table: str) -> None:
        profiler_dir = self.profiler_config.torch_profiler_dir

        # Skip file write for URI paths (gs://, s3://, etc.)
        # as standard file I/O doesn't work with URI schemes
        if not _is_uri_path(profiler_dir):
            profiler_out_file = os.path.join(
                profiler_dir, f"{self._worker_name}.profiler_out.txt"
            )
            with open(profiler_out_file, "w") as f:
                print(table, file=f)

    def _maybe_add_version_metadata(self) -> None:
        """Stamp the vLLM version (which embeds the git commit) into the trace.

        add_metadata_json is a no-op until Kineto is initialized, which with a
        schedule only happens after the WAIT phase, so stamp once it's live.
        """
        if self._version_metadata_added:
            return
        # None while the schedule is still in the WAIT phase.
        if self.profiler.profiler is None:
            return
        try:
            self.profiler.add_metadata_json(
                "vllm_version", json.dumps(vllm.version.__version__)
            )
            self.profiler.add_metadata_json(
                "vllm_version_tuple",
                json.dumps([str(p) for p in vllm.version.__version_tuple__]),
            )
        except Exception as e:
            logger.warning("Failed to add vLLM version to profiler metadata: %s", e)
        # Mark done even on failure, to avoid retrying every step.
        self._version_metadata_added = True

    @override
    def _start(self) -> None:
        self.profiler = torch.profiler.profile(**self._profiler_kwargs)
        self._warmup_steps_remaining = self._initial_warmup_steps_remaining
        self._version_metadata_added = False
        self.profiler.start()
        # No-schedule case: Kineto is live immediately. With a schedule this
        # no-ops and _profiler_step stamps it once WAIT ends.
        self._maybe_add_version_metadata()

    @override
    def _stop(self) -> None:
        self.profiler.stop()

        rank = self.local_rank
        if self.dump_device_time_total:
            table = self._build_profiler_table(sort_key="self_device_time_total")
            self._write_profiler_table(table)

            # only print profiler results on rank 0
            if rank == 0:
                print(table)

        if self.dump_cpu_time_total:
            table = self._build_profiler_table(
                sort_key="self_cpu_time_total", row_limit=50
            )
            self._write_profiler_table(table)

            # only print profiler results on rank 0
            if rank == 0:
                print(table)

    @override
    def _profiler_step(self) -> bool:
        """Call profiler.step() when using schedule-based profiling.

        Returns:
            True if the step was an active profiling step (data recorded),
            False if the step was a warmup step (data discarded).

        """
        if self._uses_schedule:
            self.profiler.step()
            # Stamp once the schedule leaves WAIT and Kineto is live.
            self._maybe_add_version_metadata()
            # Track warmup steps - only count active steps toward max_iterations
            if self._warmup_steps_remaining > 0:
                self._warmup_steps_remaining -= 1
                return False
        return True

    @property
    @override
    def should_annotate(self) -> bool:
        return self._running and self._records_cpu_activity

    @override
    def annotate_context_manager(self, name: str):
        if not self.should_annotate:
            return nullcontext()
        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._graph_attribution = profiler_config.proton_graph_attribution
        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
        self._graph_session = False
        self._phase = 0
        self._active_output_path: str | None = None
        self._session_storage_path = os.path.join(
            self._output_dir,
            f".proton_cuda_graph_session_{worker_name}_{self._instance_id}",
        )

        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._graph_attribution:
            self._require_triton_version(
                "CUDA graph attribution", _TRITON_PROTON_3_7_VERSION
            )
        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].lower() == "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

    @property
    def has_cuda_graph_session(self) -> bool:
        return self._graph_session

    @override
    def set_output_name(self, worker_name: str) -> None:
        """Set the next run's output name after startup graph capture."""
        if self._active:
            return
        self._output_path = os.path.join(self._output_dir, f"proton_{worker_name}")

    @override
    @contextmanager
    def capture_cuda_graphs(self) -> Iterator[None]:
        """Keep a Proton session active while vLLM captures CUDA graphs."""
        if not self._graph_attribution:
            yield
            return
        if self._session_id is None:
            self._session_id = self._create_session(self._session_storage_path)
        else:
            self._proton.activate(session=self._session_id)
        self._graph_session = True

        try:
            yield
        finally:
            captured_phase = self._phase
            try:
                self._phase = self._proton.data.advance_phase(self._session_id)
            finally:
                self._proton.deactivate(session=self._session_id, flushing=True)
            self._proton.data.clear(self._session_id, captured_phase)

    @override
    def _start(self) -> None:
        self._active_output_path = (
            f"{self._output_path}_{self._instance_id}_run{self._run_id}"
        )
        self._run_id += 1
        if self._graph_session:
            assert self._session_id is not None
            self._proton.activate(session=self._session_id)
        else:
            self._session_id = self._create_session(self._active_output_path)

    def _write_graph_phase(self, phase: int) -> None:
        assert self._active_output_path is not None
        output_format = self._output_format or "hatchet"
        output_path = f"{self._active_output_path}.{output_format}"
        if output_format == "hatchet_msgpack":
            with open(output_path, "wb") as output_file:
                output_file.write(
                    self._proton.data.get_msgpack(self._session_id, phase)
                )
        else:
            with open(output_path, "w", encoding="utf-8") as output_file:
                json.dump(self._proton.data.get(self._session_id, phase), output_file)

    def _finalize_session(self, session_id: int) -> None:
        if self._output_format is None:
            self._proton.finalize(session=session_id)
        else:
            self._proton.finalize(session=session_id, output_format=self._output_format)

    @override
    def _stop(self) -> None:
        assert self._session_id is not None
        session_id = self._session_id
        if self._graph_session:
            completed_phase = self._phase
            try:
                self._phase = self._proton.data.advance_phase(session_id)
            finally:
                self._proton.deactivate(session=session_id, flushing=True)
            try:
                self._write_graph_phase(completed_phase)
            finally:
                self._proton.data.clear(session_id, completed_phase)
                self._active_output_path = None
            return

        try:
            self._proton.deactivate(session=session_id)
        finally:
            try:
                self._finalize_session(session_id)
            finally:
                self._session_id = None
                self._active_output_path = None

    @override
    def shutdown(self) -> None:
        super().shutdown()
        if self._graph_session and self._session_id is not None:
            session_id = self._session_id
            self._session_id = None
            try:
                self._finalize_session(session_id)
            except Exception:
                logger.exception(
                    "Failed to finalize Proton CUDA graph session during shutdown."
                )
            finally:
                for output_path in glob(f"{self._session_storage_path}.*"):
                    with suppress(FileNotFoundError):
                        os.remove(output_path)

    @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)
        # Note: lazy import to avoid dependency issues if CUDA is not available.
        import torch.cuda.profiler as cuda_profiler

        self._cuda_profiler = cuda_profiler

    @override
    def _start(self) -> None:
        self._cuda_profiler.start()

    @override
    def _stop(self) -> None:
        self._cuda_profiler.stop()

    @override
    def annotate_context_manager(self, name: str):
        return torch.cuda.nvtx.range(name)


_DEFAULT_TORCH_PROFILER_ACTIVITIES: dict[str, tuple[TorchProfilerActivity, ...]] = {
    "cpu": ("CPU",),
    "cuda": ("CPU", "CUDA"),
    "xpu": ("CPU", "XPU"),
}
_SUPPORTED_TORCH_PROFILER_ACTIVITIES = {
    device: frozenset(activities)
    for device, activities in _DEFAULT_TORCH_PROFILER_ACTIVITIES.items()
}
_SUPPORTED_PROFILER_KINDS: dict[str, frozenset[ProfilerKind]] = {
    "cpu": frozenset(("torch",)),
    "cuda": frozenset(("torch", "cuda", "proton")),
    "xpu": frozenset(("torch",)),
}


def validate_worker_profiler_config(profiler_config: ProfilerConfig) -> None:
    """Validate profiler selections against the worker's capabilities."""
    profiler_type = profiler_config.profiler
    if profiler_type is None:
        return
    device_type = current_platform.device_type
    if device_type not in _SUPPORTED_PROFILER_KINDS:
        raise ValueError(f"Unsupported profiler device type: {device_type}")
    supported_kinds = _SUPPORTED_PROFILER_KINDS[device_type]
    if profiler_type not in supported_kinds:
        supported_names = ", ".join(sorted(supported_kinds))
        raise ValueError(
            f"Unsupported profiler type for {device_type}: "
            f"{profiler_type}. Supported profiler types: {supported_names}."
        )
    if profiler_type == "torch":
        default_activities = _DEFAULT_TORCH_PROFILER_ACTIVITIES[device_type]
        supported_activities = _SUPPORTED_TORCH_PROFILER_ACTIVITIES[device_type]
        configured = profiler_config.torch_profiler_activities
        activities = default_activities if configured is None else configured
        unsupported = set(activities) - supported_activities
        if unsupported:
            unsupported_names = ", ".join(sorted(unsupported))
            supported_names = ", ".join(sorted(supported_activities))
            raise ValueError(
                f"Unsupported torch profiler activities for "
                f"{device_type}: {unsupported_names}. "
                f"Supported activities: {supported_names}."
            )


def create_worker_profiler(
    profiler_config: ProfilerConfig,
    *,
    worker_name: str,
    local_rank: int,
) -> WorkerProfiler:
    """Create a profiler using a validated config and platform defaults."""
    profiler_type = profiler_config.profiler
    if profiler_type == "torch":
        default_activities = _DEFAULT_TORCH_PROFILER_ACTIVITIES[
            current_platform.device_type
        ]
        configured = profiler_config.torch_profiler_activities
        logger.debug("Starting torch profiler with trace name: %s", worker_name)
        return TorchProfilerWrapper(
            profiler_config,
            worker_name=worker_name,
            local_rank=local_rank,
            activities=default_activities if configured is None else tuple(configured),
        )
    if profiler_type == "cuda":
        logger.debug("Starting CUDA profiler")
        return CudaProfilerWrapper(profiler_config)

    assert profiler_type == "proton", f"Unknown profiler type: {profiler_type}"
    logger.debug("Starting Proton profiler with trace name: %s", worker_name)
    return ProtonProfilerWrapper(profiler_config, worker_name=worker_name)


def create_frontend_profiler(
    profiler_config: ProfilerConfig, *, worker_name: str
) -> TorchProfilerWrapper:
    """Create the CPU-only frontend profiler, which runs for the whole session
    rather than following engine iterations."""
    logger.info(
        "Torch profiler enabled. AsyncLLM CPU traces will be collected under %s",
        profiler_config.torch_profiler_dir,
    )
    return TorchProfilerWrapper(
        replace(
            profiler_config,
            delay_iterations=0,
            max_iterations=0,
            wait_iterations=0,
            warmup_iterations=0,
        ),
        worker_name=worker_name,
        local_rank=0,
        activities=["CPU"],
    )


def create_graph_capture_profiler(
    profiler_config: ProfilerConfig, global_rank: int
) -> WorkerProfiler | None:
    """Create a profiler to observe CUDA graph capture, if configured.

    Applies only to backends that attribute graph replay activity which
    need a session around capture (Proton with ``proton_graph_attribution``).
    """
    if (
        profiler_config.profiler == "proton"
        and profiler_config.proton_graph_attribution
    ):
        from vllm.distributed.utils import get_worker_rank_suffix

        return ProtonProfilerWrapper(
            profiler_config,
            worker_name=get_worker_rank_suffix(global_rank=global_rank),
        )
    return None