238 lines
9.7 KiB
Python
238 lines
9.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import os
|
|
from typing import Any, Literal
|
|
|
|
from pydantic import Field, model_validator
|
|
from typing_extensions import Self
|
|
|
|
from vllm.config.utils import config
|
|
from vllm.logger import init_logger
|
|
from vllm.utils.hashing import safe_hash
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
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:
|
|
"""Check if path is a URI (scheme://...), excluding Windows drive letters.
|
|
|
|
Supports custom URI schemes like gs://, s3://, hdfs://, etc.
|
|
These paths should not be converted to absolute paths.
|
|
"""
|
|
if "://" in path:
|
|
scheme = path.split("://")[0]
|
|
# Windows drive letters are single characters (e.g., C://)
|
|
# Valid URI schemes have more than one character
|
|
return len(scheme) > 1
|
|
return False
|
|
|
|
|
|
@config
|
|
class ProfilerConfig:
|
|
"""Dataclass which contains profiler config for the engine."""
|
|
|
|
profiler: ProfilerKind | None = None
|
|
"""Which profiler to use. Defaults to None. Options are:
|
|
|
|
- 'torch': Use PyTorch 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
|
|
--profiler-config.torch_profiler_with_stack=false CLI flag."""
|
|
|
|
torch_profiler_with_flops: bool = False
|
|
"""If `True`, enables FLOPS counting in the torch profiler. Disabled by default."""
|
|
|
|
torch_profiler_use_gzip: bool = True
|
|
"""If `True`, saves torch profiler traces in gzip format. Enabled by default"""
|
|
|
|
torch_profiler_dump_cuda_time_total: bool = True
|
|
"""If `True`, dumps total CUDA time in torch profiler traces. Enabled by default."""
|
|
|
|
torch_profiler_record_shapes: bool = False
|
|
"""If `True`, records tensor shapes in the torch profiler. Disabled by default."""
|
|
|
|
torch_profiler_with_memory: bool = False
|
|
"""If `True`, enables memory profiling in the torch profiler.
|
|
Disabled by default."""
|
|
|
|
capture_torch_profiler: bool = False
|
|
"""If `True`, enables a torch profiler during CUDA graph capture on rank 0.
|
|
Traces are saved to a `capture_traces` subdirectory under `torch_profiler_dir`.
|
|
Requires `profiler` to be set to 'torch'."""
|
|
|
|
detailed_trace_annotation: bool = False
|
|
"""If `True`, uses detailed annotations with roofline metrics (sk, sqsq,
|
|
sqsk) in profiler trace events. If `False`, uses simple annotations with
|
|
only context/generation request counts and token counts.
|
|
Disabled by default."""
|
|
|
|
ignore_frontend: bool = False
|
|
"""If `True`, disables the front-end profiling of AsyncLLM when using the
|
|
'torch' profiler. This is needed to reduce overhead when using delay/limit options,
|
|
since the front-end profiling does not track iterations and will capture the
|
|
entire range.
|
|
"""
|
|
|
|
delay_iterations: int = Field(default=0, ge=0)
|
|
"""Number of engine iterations to skip before starting profiling.
|
|
Defaults to 0, meaning profiling starts immediately after receiving /start_profile.
|
|
"""
|
|
|
|
max_iterations: int = Field(default=0, ge=0)
|
|
"""Maximum number of engine iterations to profile after starting profiling.
|
|
Defaults to 0, meaning no limit.
|
|
"""
|
|
|
|
warmup_iterations: int = Field(default=0, ge=0)
|
|
"""Number of warmup iterations for PyTorch profiler schedule.
|
|
During warmup, the profiler runs but data is discarded. This helps reduce
|
|
noise from JIT compilation and other one-time costs in the profiled trace.
|
|
Defaults to 0 (schedule-based profiling disabled, recording all iterations).
|
|
Set to a positive value (e.g., 2) to enable schedule-based profiling.
|
|
"""
|
|
|
|
active_iterations: int = Field(default=5, ge=1)
|
|
"""Number of active iterations for PyTorch profiler schedule.
|
|
This is the number of iterations where profiling data is actually collected.
|
|
Defaults to 5 active iterations.
|
|
"""
|
|
|
|
wait_iterations: int = Field(default=0, ge=0)
|
|
"""Number of wait iterations for PyTorch profiler schedule.
|
|
During wait, the profiler is completely off with zero overhead.
|
|
This allows skipping initial iterations before warmup begins.
|
|
Defaults to 0 (no wait period).
|
|
"""
|
|
|
|
def compute_hash(self) -> str:
|
|
"""
|
|
WARNING: Whenever a new field is added to this config,
|
|
ensure that it is included in the factors list if
|
|
it affects the computation graph.
|
|
|
|
Provide a hash that uniquely identifies all the configs
|
|
that affect the structure of the computation
|
|
graph from input ids/embeddings to the final hidden states,
|
|
excluding anything before input ids/embeddings and after
|
|
the final hidden states.
|
|
"""
|
|
# no factors to consider.
|
|
# this config will not affect the computation graph.
|
|
factors: list[Any] = []
|
|
hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()
|
|
return hash_str
|
|
|
|
@model_validator(mode="after")
|
|
def _validate_profiler_config(self) -> Self:
|
|
has_delay_or_limit = self.delay_iterations > 0 or self.max_iterations > 0
|
|
if self.profiler == "torch" and has_delay_or_limit and not self.ignore_frontend:
|
|
logger.warning_once(
|
|
"Using 'torch' profiler with delay_iterations or max_iterations "
|
|
"while ignore_frontend is False may result in high overhead."
|
|
)
|
|
|
|
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 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 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(
|
|
"capture_torch_profiler is only applicable when profiler is "
|
|
"set to 'torch'"
|
|
)
|
|
|
|
return self
|