| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | from typing import Any, Literal |
| 6 | |
| 7 | from pydantic import Field, model_validator |
| 8 | from typing_extensions import Self |
| 9 | |
| 10 | from vllm.config.utils import config |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.utils.hashing import safe_hash |
| 13 | |
| 14 | logger = init_logger(__name__) |
| 15 | |
| 16 | ProfilerKind = Literal["torch", "cuda", "proton"] |
| 17 | TorchProfilerActivity = Literal["CPU", "CUDA", "PrivateUse1", "XPU"] |
| 18 | ProtonBackend = Literal["cupti"] |
| 19 | ProtonContext = Literal["shadow", "python"] |
| 20 | ProtonData = Literal["tree", "trace"] |
| 21 | ProtonHook = Literal["triton"] |
| 22 | ProtonOutputFormat = Literal["hatchet", "hatchet_msgpack", "chrome_trace"] |
| 23 | |
| 24 | |
| 25 | def _is_uri_path(path: str) -> bool: |
| 26 | """Check if path is a URI (scheme://...), excluding Windows drive letters. |
| 27 | |
| 28 | Supports custom URI schemes like gs://, s3://, hdfs://, etc. |
| 29 | These paths should not be converted to absolute paths. |
| 30 | """ |
| 31 | if "://" in path: |
| 32 | scheme = path.split("://")[0] |
| 33 | # Windows drive letters are single characters (e.g., C://) |
| 34 | # Valid URI schemes have more than one character |
| 35 | return len(scheme) > 1 |
| 36 | return False |
| 37 | |
| 38 | |
| 39 | @config |
| 40 | class ProfilerConfig: |
| 41 | """Dataclass which contains profiler config for the engine.""" |
| 42 | |
| 43 | profiler: ProfilerKind | None = None |
| 44 | """Which profiler to use. Defaults to None. Options are: |
| 45 | |
| 46 | - 'torch': Use PyTorch profiler. |
| 47 | - 'cuda': Use CUDA profiler. |
| 48 | - 'proton': Use Triton Proton profiler.""" |
| 49 | |
| 50 | torch_profiler_dir: str = "" |
| 51 | """Directory to save torch profiler traces. Both AsyncLLM's CPU traces and |
| 52 | worker's traces (CPU & GPU) will be saved under this directory. Note that |
| 53 | it must be an absolute path.""" |
| 54 | |
| 55 | torch_profiler_activities: list[TorchProfilerActivity] | None = Field( |
| 56 | default=None, min_length=1 |
| 57 | ) |
| 58 | """Activities recorded by workers using the torch profiler. When unset, |
| 59 | each worker uses its platform default: CPU; CPU and CUDA; or CPU and XPU.""" |
| 60 | |
| 61 | proton_profiler_dir: str = "" |
| 62 | """Directory to save Triton Proton profiles. Each worker writes a |
| 63 | separate rank-qualified file.""" |
| 64 | |
| 65 | proton_context: ProtonContext = "shadow" |
| 66 | """Proton context source. ``shadow`` records explicit scopes with low |
| 67 | overhead; ``python`` records Python call stacks.""" |
| 68 | |
| 69 | proton_data: ProtonData = "tree" |
| 70 | """Proton output type. ``tree`` produces Hatchet data and ``trace`` |
| 71 | produces a Chrome trace.""" |
| 72 | |
| 73 | proton_backend: ProtonBackend | None = None |
| 74 | """Proton GPU backend. ``None`` lets Proton select CUPTI automatically.""" |
| 75 | |
| 76 | proton_mode: str | None = None |
| 77 | """Optional backend-specific Proton mode string, such as ``pcsampling``.""" |
| 78 | |
| 79 | proton_hook: ProtonHook | None = None |
| 80 | """Optional Proton hook. Use ``triton`` to add Triton launch metadata.""" |
| 81 | |
| 82 | proton_output_format: ProtonOutputFormat | None = None |
| 83 | """Optional format passed to Proton when finalizing a profile. ``None`` |
| 84 | uses the default format for ``proton_data``.""" |
| 85 | |
| 86 | proton_graph_attribution: bool = False |
| 87 | """Observe CUDA graph capture so replayed kernels can be attributed. |
| 88 | Requires Triton >= 3.7 and ``proton_data='tree'``.""" |
| 89 | |
| 90 | torch_profiler_with_stack: bool = True |
| 91 | """If `True`, enables stack tracing in the torch profiler. Enabled by default |
| 92 | as it is useful for debugging. Can be disabled via |
| 93 | --profiler-config.torch_profiler_with_stack=false CLI flag.""" |
| 94 | |
| 95 | torch_profiler_with_flops: bool = False |
| 96 | """If `True`, enables FLOPS counting in the torch profiler. Disabled by default.""" |
| 97 | |
| 98 | torch_profiler_use_gzip: bool = True |
| 99 | """If `True`, saves torch profiler traces in gzip format. Enabled by default""" |
| 100 | |
| 101 | torch_profiler_dump_cuda_time_total: bool = True |
| 102 | """If `True`, dumps total CUDA time in torch profiler traces. Enabled by default.""" |
| 103 | |
| 104 | torch_profiler_record_shapes: bool = False |
| 105 | """If `True`, records tensor shapes in the torch profiler. Disabled by default.""" |
| 106 | |
| 107 | torch_profiler_with_memory: bool = False |
| 108 | """If `True`, enables memory profiling in the torch profiler. |
| 109 | Disabled by default.""" |
| 110 | |
| 111 | capture_torch_profiler: bool = False |
| 112 | """If `True`, enables a torch profiler during CUDA graph capture on rank 0. |
| 113 | Traces are saved to a `capture_traces` subdirectory under `torch_profiler_dir`. |
| 114 | Requires `profiler` to be set to 'torch'.""" |
| 115 | |
| 116 | detailed_trace_annotation: bool = False |
| 117 | """If `True`, uses detailed annotations with roofline metrics (sk, sqsq, |
| 118 | sqsk) in profiler trace events. If `False`, uses simple annotations with |
| 119 | only context/generation request counts and token counts. |
| 120 | Disabled by default.""" |
| 121 | |
| 122 | ignore_frontend: bool = False |
| 123 | """If `True`, disables the front-end profiling of AsyncLLM when using the |
| 124 | 'torch' profiler. This is needed to reduce overhead when using delay/limit options, |
| 125 | since the front-end profiling does not track iterations and will capture the |
| 126 | entire range. |
| 127 | """ |
| 128 | |
| 129 | delay_iterations: int = Field(default=0, ge=0) |
| 130 | """Number of engine iterations to skip before starting profiling. |
| 131 | Defaults to 0, meaning profiling starts immediately after receiving /start_profile. |
| 132 | """ |
| 133 | |
| 134 | max_iterations: int = Field(default=0, ge=0) |
| 135 | """Maximum number of engine iterations to profile after starting profiling. |
| 136 | Defaults to 0, meaning no limit. |
| 137 | """ |
| 138 | |
| 139 | warmup_iterations: int = Field(default=0, ge=0) |
| 140 | """Number of warmup iterations for PyTorch profiler schedule. |
| 141 | During warmup, the profiler runs but data is discarded. This helps reduce |
| 142 | noise from JIT compilation and other one-time costs in the profiled trace. |
| 143 | Defaults to 0 (schedule-based profiling disabled, recording all iterations). |
| 144 | Set to a positive value (e.g., 2) to enable schedule-based profiling. |
| 145 | """ |
| 146 | |
| 147 | active_iterations: int = Field(default=5, ge=1) |
| 148 | """Number of active iterations for PyTorch profiler schedule. |
| 149 | This is the number of iterations where profiling data is actually collected. |
| 150 | Defaults to 5 active iterations. |
| 151 | """ |
| 152 | |
| 153 | wait_iterations: int = Field(default=0, ge=0) |
| 154 | """Number of wait iterations for PyTorch profiler schedule. |
| 155 | During wait, the profiler is completely off with zero overhead. |
| 156 | This allows skipping initial iterations before warmup begins. |
| 157 | Defaults to 0 (no wait period). |
| 158 | """ |
| 159 | |
| 160 | def compute_hash(self) -> str: |
| 161 | """WARNING: Whenever a new field is added to this config, |
| 162 | ensure that it is included in the factors list if |
| 163 | it affects the computation graph. |
| 164 | |
| 165 | Provide a hash that uniquely identifies all the configs |
| 166 | that affect the structure of the computation |
| 167 | graph from input ids/embeddings to the final hidden states, |
| 168 | excluding anything before input ids/embeddings and after |
| 169 | the final hidden states. |
| 170 | """ |
| 171 | # no factors to consider. |
| 172 | # this config will not affect the computation graph. |
| 173 | factors: list[Any] = [] |
| 174 | hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest() |
| 175 | return hash_str |
| 176 | |
| 177 | @model_validator(mode="after") |
| 178 | def _validate_profiler_config(self) -> Self: |
| 179 | has_delay_or_limit = self.delay_iterations > 0 or self.max_iterations > 0 |
| 180 | activities = self.torch_profiler_activities |
| 181 | if activities is not None: |
| 182 | if self.profiler != "torch": |
| 183 | raise ValueError( |
| 184 | "torch_profiler_activities is only applicable when profiler " |
| 185 | "is set to 'torch'" |
| 186 | ) |
| 187 | if len(activities) != len(set(activities)): |
| 188 | raise ValueError( |
| 189 | "torch_profiler_activities must not contain duplicates" |
| 190 | ) |
| 191 | records_cpu_activity = activities is None or "CPU" in activities |
| 192 | if ( |
| 193 | self.profiler == "torch" |
| 194 | and has_delay_or_limit |
| 195 | and not self.ignore_frontend |
| 196 | and records_cpu_activity |
| 197 | ): |
| 198 | logger.warning_once( |
| 199 | "Using 'torch' profiler with delay_iterations or max_iterations " |
| 200 | "while ignore_frontend is False may result in high overhead." |
| 201 | ) |
| 202 | |
| 203 | torch_profiler_dir = self.torch_profiler_dir |
| 204 | if torch_profiler_dir and self.profiler != "torch": |
| 205 | raise ValueError( |
| 206 | "torch_profiler_dir is only applicable when profiler is set to 'torch'" |
| 207 | ) |
| 208 | if self.profiler == "torch" and not torch_profiler_dir: |
| 209 | raise ValueError("torch_profiler_dir must be set when profiler is 'torch'") |
| 210 | |
| 211 | # Support any URI scheme (gs://, s3://, hdfs://, etc.) |
| 212 | # These paths should not be converted to absolute paths |
| 213 | if torch_profiler_dir and not _is_uri_path(torch_profiler_dir): |
| 214 | self.torch_profiler_dir = os.path.abspath( |
| 215 | os.path.expanduser(torch_profiler_dir) |
| 216 | ) |
| 217 | |
| 218 | proton_profiler_dir = self.proton_profiler_dir |
| 219 | non_default_proton_options = [ |
| 220 | name |
| 221 | for name, value, default in ( |
| 222 | ("proton_profiler_dir", proton_profiler_dir, ""), |
| 223 | ("proton_context", self.proton_context, "shadow"), |
| 224 | ("proton_data", self.proton_data, "tree"), |
| 225 | ("proton_backend", self.proton_backend, None), |
| 226 | ("proton_mode", self.proton_mode, None), |
| 227 | ("proton_hook", self.proton_hook, None), |
| 228 | ("proton_output_format", self.proton_output_format, None), |
| 229 | ("proton_graph_attribution", self.proton_graph_attribution, False), |
| 230 | ) |
| 231 | if value != default |
| 232 | ] |
| 233 | if self.profiler != "proton" and non_default_proton_options: |
| 234 | options = ", ".join(non_default_proton_options) |
| 235 | raise ValueError( |
| 236 | f"{options} only applicable when profiler is set to 'proton'" |
| 237 | ) |
| 238 | if self.profiler == "proton" and not proton_profiler_dir: |
| 239 | raise ValueError( |
| 240 | "proton_profiler_dir must be set when profiler is 'proton'" |
| 241 | ) |
| 242 | if proton_profiler_dir: |
| 243 | if _is_uri_path(proton_profiler_dir): |
| 244 | raise ValueError("proton_profiler_dir must be a local directory") |
| 245 | self.proton_profiler_dir = os.path.abspath( |
| 246 | os.path.expanduser(proton_profiler_dir) |
| 247 | ) |
| 248 | |
| 249 | if self.profiler == "proton": |
| 250 | output_format = self.proton_output_format |
| 251 | if self.proton_graph_attribution and self.proton_data != "tree": |
| 252 | raise ValueError("proton_graph_attribution requires proton_data='tree'") |
| 253 | if ( |
| 254 | self.proton_graph_attribution |
| 255 | and self.proton_mode |
| 256 | and self.proton_mode.split(":", 1)[0].lower() == "periodic_flushing" |
| 257 | ): |
| 258 | raise ValueError( |
| 259 | "proton_graph_attribution is incompatible with periodic_flushing: " |
| 260 | "both manage Proton data phases." |
| 261 | ) |
| 262 | if output_format == "chrome_trace" and self.proton_data != "trace": |
| 263 | raise ValueError("chrome_trace output requires proton_data='trace'") |
| 264 | if ( |
| 265 | output_format in ("hatchet", "hatchet_msgpack") |
| 266 | and self.proton_data != "tree" |
| 267 | ): |
| 268 | raise ValueError(f"{output_format} output requires proton_data='tree'") |
| 269 | |
| 270 | if self.capture_torch_profiler and self.profiler != "torch": |
| 271 | raise ValueError( |
| 272 | "capture_torch_profiler is only applicable when profiler is " |
| 273 | "set to 'torch'" |
| 274 | ) |
| 275 | |
| 276 | return self |