| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import copy |
| 5 | from collections import defaultdict |
| 6 | from collections.abc import Callable |
| 7 | from dataclasses import asdict, dataclass, field |
| 8 | from typing import Any, Generic, TypeAlias, TypeVar |
| 9 | |
| 10 | from torch._C._autograd import DeviceType, _KinetoEvent, _ProfilerResult |
| 11 | from torch._C._profiler import _EventType, _ExperimentalConfig, _ProfilerEvent |
| 12 | from torch.autograd.profiler import FunctionEvent |
| 13 | from torch.profiler import ProfilerActivity, profile |
| 14 | |
| 15 | from vllm.profiler.utils import ( |
| 16 | TablePrinter, |
| 17 | event_has_module, |
| 18 | event_is_torch_op, |
| 19 | event_module_repr, |
| 20 | event_torch_op_stack_trace, |
| 21 | indent_string, |
| 22 | ) |
| 23 | from vllm.utils.import_utils import PlaceholderModule |
| 24 | |
| 25 | try: |
| 26 | import pandas as pd |
| 27 | except ImportError: |
| 28 | pd = PlaceholderModule("pandas") |
| 29 | |
| 30 | |
| 31 | @dataclass |
| 32 | class _ModuleTreeNode: |
| 33 | event: _ProfilerEvent |
| 34 | parent: "_ModuleTreeNode | None" = None |
| 35 | children: list["_ModuleTreeNode"] = field(default_factory=list) |
| 36 | trace: str = "" |
| 37 | |
| 38 | @property |
| 39 | def is_leaf(self): |
| 40 | return self.event.children is None or len(self.event.children) == 0 |
| 41 | |
| 42 | @property |
| 43 | def is_torch_op(self): |
| 44 | return event_is_torch_op(self.event) |
| 45 | |
| 46 | @property |
| 47 | def is_cuda(self): |
| 48 | return ( |
| 49 | self.event.tag == _EventType.Kineto |
| 50 | and self.event.typed[1].device_type == DeviceType.CUDA |
| 51 | ) |
| 52 | |
| 53 | |
| 54 | @dataclass |
| 55 | class SummaryStatsEntry: |
| 56 | name: str |
| 57 | cuda_time_us: float |
| 58 | pct_cuda_time: float |
| 59 | invocations: int |
| 60 | |
| 61 | |
| 62 | @dataclass |
| 63 | class ModelStatsEntry: |
| 64 | name: str |
| 65 | cpu_time_us: float |
| 66 | cuda_time_us: float |
| 67 | pct_cuda_time: float |
| 68 | trace: str |
| 69 | |
| 70 | |
| 71 | StatsEntry: TypeAlias = ModelStatsEntry | SummaryStatsEntry |
| 72 | StatsEntryT = TypeVar("StatsEntryT", bound=StatsEntry) |
| 73 | |
| 74 | |
| 75 | @dataclass |
| 76 | class _StatsTreeNode(Generic[StatsEntryT]): |
| 77 | entry: StatsEntryT |
| 78 | children: list["_StatsTreeNode[StatsEntryT]"] = field(default_factory=list) |
| 79 | parent: "_StatsTreeNode[StatsEntryT] | None" = None |
| 80 | |
| 81 | |
| 82 | @dataclass |
| 83 | class LayerwiseProfileResults(profile): |
| 84 | _kineto_results: _ProfilerResult |
| 85 | _kineto_event_correlation_map: dict[int, list[_KinetoEvent]] = field(init=False) |
| 86 | _event_correlation_map: dict[int, list[FunctionEvent]] = field(init=False) |
| 87 | _module_tree: list[_ModuleTreeNode] = field(init=False) |
| 88 | _model_stats_tree: list[_StatsTreeNode[ModelStatsEntry]] = field(init=False) |
| 89 | _summary_stats_tree: list[_StatsTreeNode[SummaryStatsEntry]] = field(init=False) |
| 90 | |
| 91 | # profile metadata |
| 92 | num_running_seqs: int | None = None |
| 93 | |
| 94 | def __post_init__(self): |
| 95 | self._build_correlation_map() |
| 96 | self._build_module_tree() |
| 97 | self._build_stats_trees() |
| 98 | |
| 99 | def print_model_table(self, column_widths: dict[str, int] | None = None): |
| 100 | _column_widths = dict( |
| 101 | name=60, cpu_time_us=12, cuda_time_us=12, pct_cuda_time=12, trace=60 |
| 102 | ) |
| 103 | if column_widths: |
| 104 | _column_widths.update(**column_widths) |
| 105 | filtered_model_table = [ |
| 106 | (depth, row) |
| 107 | for depth, row in self._flatten_stats_tree(self._model_stats_tree) |
| 108 | if row.cuda_time_us > 0 or row.cpu_time_us > 0 |
| 109 | ] |
| 110 | TablePrinter(ModelStatsEntry, _column_widths).print_table( |
| 111 | self._indent_row_names_based_on_depth( |
| 112 | filtered_model_table, |
| 113 | indent_style=lambda indent: "|" + "-" * indent + " ", |
| 114 | ) |
| 115 | ) |
| 116 | |
| 117 | def print_summary_table(self, column_widths: dict[str, int] | None = None): |
| 118 | _column_widths = dict( |
| 119 | name=80, cuda_time_us=12, pct_cuda_time=12, invocations=15 |
| 120 | ) |
| 121 | if column_widths: |
| 122 | _column_widths.update(**column_widths) |
| 123 | filtered_summary_table = [ |
| 124 | (depth, row) |
| 125 | for depth, row in self._flatten_stats_tree(self._summary_stats_tree) |
| 126 | if row.cuda_time_us > 0 |
| 127 | ] |
| 128 | TablePrinter(SummaryStatsEntry, _column_widths).print_table( |
| 129 | self._indent_row_names_based_on_depth( |
| 130 | filtered_summary_table, |
| 131 | indent_style=lambda indent: "|" + "-" * indent + " ", |
| 132 | ) |
| 133 | ) |
| 134 | |
| 135 | def export_model_stats_table_csv(self, filename: str): |
| 136 | df = pd.DataFrame( |
| 137 | [asdict(row) for _, row in self._flatten_stats_tree(self._model_stats_tree)] |
| 138 | ) |
| 139 | df.to_csv(filename) |
| 140 | |
| 141 | def export_summary_stats_table_csv(self, filename: str): |
| 142 | df = pd.DataFrame( |
| 143 | [ |
| 144 | asdict(row) |
| 145 | for _, row in self._flatten_stats_tree(self._summary_stats_tree) |
| 146 | ] |
| 147 | ) |
| 148 | df.to_csv(filename) |
| 149 | |
| 150 | def convert_stats_to_dict(self) -> dict[str, Any]: |
| 151 | return { |
| 152 | "metadata": {"num_running_seqs": self.num_running_seqs}, |
| 153 | "summary_stats": self._convert_stats_tree_to_dict(self._summary_stats_tree), |
| 154 | "model_stats": self._convert_stats_tree_to_dict(self._model_stats_tree), |
| 155 | } |
| 156 | |
| 157 | @staticmethod |
| 158 | def _indent_row_names_based_on_depth( |
| 159 | depths_rows: list[tuple[int, StatsEntryT]], |
| 160 | indent_style: Callable[[int], str] | str = " ", |
| 161 | ): |
| 162 | indented_rows: list[StatsEntryT] = [] |
| 163 | for depth, row in depths_rows: |
| 164 | if row.cuda_time_us == 0: |
| 165 | continue |
| 166 | indented_row: StatsEntryT = copy.deepcopy(row) |
| 167 | indented_row.name = indent_string(indented_row.name, depth, indent_style) |
| 168 | indented_rows.append(indented_row) |
| 169 | return indented_rows |
| 170 | |
| 171 | def _build_correlation_map(self): |
| 172 | self._kineto_event_correlation_map = defaultdict(list) |
| 173 | for event in self._kineto_results.events(): |
| 174 | self._kineto_event_correlation_map[event.correlation_id()].append(event) |
| 175 | |
| 176 | def _build_module_tree(self): |
| 177 | self._module_tree = [] |
| 178 | event_tree = self._kineto_results.experimental_event_tree() |
| 179 | |
| 180 | def _df_traversal( |
| 181 | event: _ProfilerEvent, curr_node: _ModuleTreeNode | None = None |
| 182 | ): |
| 183 | # For the tensor parallel case for now only look at task 1 |
| 184 | if event.start_tid != 1: |
| 185 | return |
| 186 | |
| 187 | if event_has_module(event): |
| 188 | node = _ModuleTreeNode(event=event, parent=curr_node) |
| 189 | if curr_node: |
| 190 | curr_node.children.append(node) |
| 191 | else: |
| 192 | self._module_tree.append(node) |
| 193 | curr_node = node |
| 194 | |
| 195 | is_leaf = event.children is None or len(event.children) == 0 |
| 196 | if is_leaf and curr_node: |
| 197 | node = _ModuleTreeNode( |
| 198 | event=event, |
| 199 | parent=curr_node, |
| 200 | trace=event_torch_op_stack_trace( |
| 201 | event, until=lambda x: event_has_module(x) |
| 202 | ), |
| 203 | ) |
| 204 | curr_node.children.append(node) |
| 205 | curr_node = node |
| 206 | |
| 207 | for child in event.children: |
| 208 | _df_traversal(child, curr_node) |
| 209 | |
| 210 | for root in event_tree: |
| 211 | _df_traversal(root) |
| 212 | |
| 213 | def _get_kineto_gpu_event(self, node: _ModuleTreeNode): |
| 214 | if node.event.tag != _EventType.Kineto: |
| 215 | return None |
| 216 | correlated_kineto_events = self._kineto_event_correlation_map.get( |
| 217 | node.event.correlation_id, [] |
| 218 | ) |
| 219 | iterator = ( |
| 220 | x |
| 221 | for x in correlated_kineto_events |
| 222 | if x.device_type() == DeviceType.CUDA and x.name() == node.event.name |
| 223 | ) |
| 224 | return next(iterator, None) |
| 225 | |
| 226 | def _cumulative_cuda_time(self, node: _ModuleTreeNode): |
| 227 | """Return cuda time in microseconds.""" |
| 228 | |
| 229 | def _cumulative_cuda_time_recursive(node: _ModuleTreeNode): |
| 230 | if node.is_leaf and (gpu_kineto_event := self._get_kineto_gpu_event(node)): |
| 231 | return gpu_kineto_event.duration_ns() / 1000.0 |
| 232 | else: |
| 233 | cumulative_cuda_time = 0 |
| 234 | for child in node.children: |
| 235 | cumulative_cuda_time += _cumulative_cuda_time_recursive(child) |
| 236 | return cumulative_cuda_time |
| 237 | |
| 238 | return _cumulative_cuda_time_recursive(node) |
| 239 | |
| 240 | def _total_cuda_time(self): |
| 241 | return sum([self._cumulative_cuda_time(root) for root in self._module_tree]) |
| 242 | |
| 243 | def _build_stats_trees(self): |
| 244 | summary_dict: dict[tuple[str, ...], _StatsTreeNode[SummaryStatsEntry]] = {} |
| 245 | total_cuda_time = self._total_cuda_time() |
| 246 | |
| 247 | def pct_cuda_time(cuda_time_us): |
| 248 | return (cuda_time_us / total_cuda_time) * 100 |
| 249 | |
| 250 | def build_summary_stats_tree_df( |
| 251 | node: _ModuleTreeNode, |
| 252 | parent: _StatsTreeNode[SummaryStatsEntry] | None = None, |
| 253 | summary_trace: tuple[str, ...] = (), |
| 254 | ) -> _StatsTreeNode[SummaryStatsEntry] | None: |
| 255 | if event_has_module(node.event): |
| 256 | name = event_module_repr(node.event) |
| 257 | cuda_time_us = self._cumulative_cuda_time(node) |
| 258 | elif gpu_kineto_event := self._get_kineto_gpu_event(node): |
| 259 | name = gpu_kineto_event.name() |
| 260 | cuda_time_us = gpu_kineto_event.duration_ns() / 1000.0 |
| 261 | else: |
| 262 | return None |
| 263 | |
| 264 | summary_trace = summary_trace + (name,) |
| 265 | if summary_trace in summary_dict: |
| 266 | entry = summary_dict[summary_trace].entry |
| 267 | entry.cuda_time_us += cuda_time_us |
| 268 | entry.invocations += 1 |
| 269 | entry.pct_cuda_time = pct_cuda_time(entry.cuda_time_us) |
| 270 | else: |
| 271 | new_node = _StatsTreeNode( |
| 272 | entry=SummaryStatsEntry( |
| 273 | name=name, |
| 274 | cuda_time_us=cuda_time_us, |
| 275 | pct_cuda_time=pct_cuda_time(cuda_time_us), |
| 276 | invocations=1, |
| 277 | ), |
| 278 | parent=parent, |
| 279 | ) |
| 280 | if parent: |
| 281 | parent.children.append(new_node) |
| 282 | summary_dict[summary_trace] = new_node |
| 283 | |
| 284 | for child in node.children: |
| 285 | build_summary_stats_tree_df( |
| 286 | child, summary_dict[summary_trace], summary_trace |
| 287 | ) |
| 288 | |
| 289 | return summary_dict[summary_trace] |
| 290 | |
| 291 | self._summary_stats_tree = [] |
| 292 | for root in self._module_tree: |
| 293 | summary_node = build_summary_stats_tree_df(root) |
| 294 | if summary_node is not None: |
| 295 | self._summary_stats_tree.append(summary_node) |
| 296 | |
| 297 | def build_model_stats_tree_df( |
| 298 | node: _ModuleTreeNode, |
| 299 | parent: _StatsTreeNode[ModelStatsEntry] | None = None, |
| 300 | ) -> _StatsTreeNode[ModelStatsEntry] | None: |
| 301 | if event_has_module( |
| 302 | node.event, |
| 303 | ): |
| 304 | name = event_module_repr(node.event) |
| 305 | cuda_time_us = self._cumulative_cuda_time(node) |
| 306 | cpu_time_us = node.event.duration_time_ns / 1000 |
| 307 | trace = "" |
| 308 | elif gpu_kineto_event := self._get_kineto_gpu_event(node): |
| 309 | name = gpu_kineto_event.name() |
| 310 | cuda_time_us = gpu_kineto_event.duration_ns() / 1000.0 |
| 311 | cpu_time_us = 0 |
| 312 | trace = node.trace |
| 313 | else: |
| 314 | return None |
| 315 | |
| 316 | new_node = _StatsTreeNode( |
| 317 | entry=ModelStatsEntry( |
| 318 | name=name, |
| 319 | cpu_time_us=cpu_time_us, |
| 320 | cuda_time_us=cuda_time_us, |
| 321 | pct_cuda_time=pct_cuda_time(cuda_time_us), |
| 322 | trace=trace, |
| 323 | ), |
| 324 | parent=parent, |
| 325 | ) |
| 326 | if parent: |
| 327 | parent.children.append(new_node) |
| 328 | |
| 329 | for child in node.children: |
| 330 | build_model_stats_tree_df(child, new_node) |
| 331 | |
| 332 | return new_node |
| 333 | |
| 334 | self._model_stats_tree = [] |
| 335 | for root in self._module_tree: |
| 336 | model_node = build_model_stats_tree_df(root) |
| 337 | if model_node is not None: |
| 338 | self._model_stats_tree.append(model_node) |
| 339 | |
| 340 | def _flatten_stats_tree( |
| 341 | self, tree: list[_StatsTreeNode[StatsEntryT]] |
| 342 | ) -> list[tuple[int, StatsEntryT]]: |
| 343 | entries: list[tuple[int, StatsEntryT]] = [] |
| 344 | |
| 345 | def df_traversal(node: _StatsTreeNode[StatsEntryT], depth: int = 0): |
| 346 | entries.append((depth, node.entry)) |
| 347 | for child in node.children: |
| 348 | df_traversal(child, depth=depth + 1) |
| 349 | |
| 350 | for root in tree: |
| 351 | df_traversal(root) |
| 352 | |
| 353 | return entries |
| 354 | |
| 355 | def _convert_stats_tree_to_dict( |
| 356 | self, tree: list[_StatsTreeNode[StatsEntryT]] |
| 357 | ) -> list[dict[str, Any]]: |
| 358 | root_dicts: list[dict[str, Any]] = [] |
| 359 | |
| 360 | def df_traversal( |
| 361 | node: _StatsTreeNode[StatsEntryT], curr_json_list: list[dict[str, Any]] |
| 362 | ): |
| 363 | curr_json_list.append({"entry": asdict(node.entry), "children": []}) |
| 364 | for child in node.children: |
| 365 | df_traversal(child, curr_json_list[-1]["children"]) |
| 366 | |
| 367 | for root in tree: |
| 368 | df_traversal(root, root_dicts) |
| 369 | |
| 370 | return root_dicts |
| 371 | |
| 372 | |
| 373 | class layerwise_profile(profile): |
| 374 | def __init__(self, num_running_seqs: int | None = None): |
| 375 | """Layerwise profile constructor. |
| 376 | |
| 377 | Args: |
| 378 | num_running_seqs (Optional[int], optional): When given, |
| 379 | num_running_seqs will be passed to LayerProfileResults |
| 380 | for metadata update. Defaults to None. |
| 381 | |
| 382 | """ |
| 383 | super().__init__( |
| 384 | activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], |
| 385 | record_shapes=True, |
| 386 | with_stack=True, |
| 387 | with_modules=True, |
| 388 | experimental_config=_ExperimentalConfig(verbose=True), |
| 389 | ) |
| 390 | |
| 391 | self.num_running_seqs = num_running_seqs |
| 392 | |
| 393 | def __enter__(self): |
| 394 | return super().__enter__() |
| 395 | |
| 396 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 397 | super().__exit__(exc_type, exc_val, exc_tb) |
| 398 | self.results = LayerwiseProfileResults( |
| 399 | self.profiler.kineto_results, num_running_seqs=self.num_running_seqs |
| 400 | ) |