layerwise_profile.py

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

import copy
from collections import defaultdict
from dataclasses import asdict, dataclass, field
from typing import Any, Callable, Optional, TypeAlias, Union

import pandas as pd
from torch._C._autograd import DeviceType, _KinetoEvent, _ProfilerResult
from torch._C._profiler import _EventType, _ExperimentalConfig, _ProfilerEvent
from torch.autograd.profiler import FunctionEvent
from torch.profiler import ProfilerActivity, profile

from vllm.profiler.utils import (
    TablePrinter,
    event_has_module,
    event_is_torch_op,
    event_module_repr,
    event_torch_op_stack_trace,
    indent_string,
)


@dataclass
class _ModuleTreeNode:
    event: _ProfilerEvent
    parent: Optional["_ModuleTreeNode"] = None
    children: list["_ModuleTreeNode"] = field(default_factory=list)
    trace: str = ""

    @property
    def is_leaf(self):
        return self.event.children is None or len(self.event.children) == 0

    @property
    def is_torch_op(self):
        return event_is_torch_op(self.event)

    @property
    def is_cuda(self):
        return (
            self.event.tag == _EventType.Kineto
            and self.event.typed[1].device_type == DeviceType.CUDA
        )


@dataclass
class SummaryStatsEntry:
    name: str
    cuda_time_us: float
    pct_cuda_time: float
    invocations: int


@dataclass
class ModelStatsEntry:
    name: str
    cpu_time_us: float
    cuda_time_us: float
    pct_cuda_time: float
    trace: str


StatsEntry: TypeAlias = Union[ModelStatsEntry, SummaryStatsEntry]


@dataclass
class _StatsTreeNode:
    entry: StatsEntry
    children: list[StatsEntry]
    parent: Optional[StatsEntry]


@dataclass
class LayerwiseProfileResults(profile):
    _kineto_results: _ProfilerResult
    _kineto_event_correlation_map: dict[int, list[_KinetoEvent]] = field(init=False)
    _event_correlation_map: dict[int, list[FunctionEvent]] = field(init=False)
    _module_tree: list[_ModuleTreeNode] = field(init=False)
    _model_stats_tree: list[_StatsTreeNode] = field(init=False)
    _summary_stats_tree: list[_StatsTreeNode] = field(init=False)

    # profile metadata
    num_running_seqs: Optional[int] = None

    def __post_init__(self):
        self._build_correlation_map()
        self._build_module_tree()
        self._build_stats_trees()

    def print_model_table(self, column_widths: dict[str, int] = None):
        _column_widths = dict(
            name=60, cpu_time_us=12, cuda_time_us=12, pct_cuda_time=12, trace=60
        )
        if column_widths:
            _column_widths.update(**column_widths)
        filtered_model_table = [
            (depth, row)
            for depth, row in self._flatten_stats_tree(self._model_stats_tree)
            if row.cuda_time_us > 0 or row.cpu_time_us > 0
        ]
        TablePrinter(ModelStatsEntry, _column_widths).print_table(
            self._indent_row_names_based_on_depth(
                filtered_model_table,
                indent_style=lambda indent: "|" + "-" * indent + " ",
            )
        )

    def print_summary_table(self, column_widths: dict[str, int] = None):
        _column_widths = dict(
            name=80, cuda_time_us=12, pct_cuda_time=12, invocations=15
        )
        if column_widths:
            _column_widths.update(**column_widths)
        filtered_summary_table = [
            (depth, row)
            for depth, row in self._flatten_stats_tree(self._summary_stats_tree)
            if row.cuda_time_us > 0
        ]
        TablePrinter(SummaryStatsEntry, _column_widths).print_table(
            self._indent_row_names_based_on_depth(
                filtered_summary_table,
                indent_style=lambda indent: "|" + "-" * indent + " ",
            )
        )

    def export_model_stats_table_csv(self, filename: str):
        df = pd.DataFrame(
            [asdict(row) for _, row in self._flatten_stats_tree(self._model_stats_tree)]
        )
        df.to_csv(filename)

    def export_summary_stats_table_csv(self, filename: str):
        df = pd.DataFrame(
            [
                asdict(row)
                for _, row in self._flatten_stats_tree(self._summary_stats_tree)
            ]
        )
        df.to_csv(filename)

    def convert_stats_to_dict(self) -> dict[str, Any]:
        return {
            "metadata": {"num_running_seqs": self.num_running_seqs},
            "summary_stats": self._convert_stats_tree_to_dict(self._summary_stats_tree),
            "model_stats": self._convert_stats_tree_to_dict(self._model_stats_tree),
        }

    @staticmethod
    def _indent_row_names_based_on_depth(
        depths_rows: list[tuple[int, StatsEntry]],
        indent_style: Union[Callable[[int], str], str] = " ",
    ):
        indented_rows = []
        for depth, row in depths_rows:
            if row.cuda_time_us == 0:
                continue
            indented_row = copy.deepcopy(row)
            indented_row.name = indent_string(indented_row.name, depth, indent_style)
            indented_rows.append(indented_row)
        return indented_rows

    def _build_correlation_map(self):
        self._kineto_event_correlation_map = defaultdict(list)
        for event in self._kineto_results.events():
            self._kineto_event_correlation_map[event.correlation_id()].append(event)

    def _build_module_tree(self):
        self._module_tree = []
        event_tree = self._kineto_results.experimental_event_tree()

        def _df_traversal(
            event: _ProfilerEvent, curr_node: Optional[_ModuleTreeNode] = None
        ):
            # For the tensor parallel case for now only look at task 1
            if event.start_tid != 1:
                return

            if event_has_module(event):
                node = _ModuleTreeNode(event=event, parent=curr_node)
                if curr_node:
                    curr_node.children.append(node)
                else:
                    self._module_tree.append(node)
                curr_node = node

            is_leaf = event.children is None or len(event.children) == 0
            if is_leaf and curr_node:
                node = _ModuleTreeNode(
                    event=event,
                    parent=curr_node,
                    trace=event_torch_op_stack_trace(
                        event, until=lambda x: event_has_module(x)
                    ),
                )
                curr_node.children.append(node)
                curr_node = node

            for child in event.children:
                _df_traversal(child, curr_node)

        for root in event_tree:
            _df_traversal(root)

    def _get_kineto_gpu_event(self, node: _ModuleTreeNode):
        if node.event.tag != _EventType.Kineto:
            return None
        correlated_kineto_events = self._kineto_event_correlation_map.get(
            node.event.correlation_id, []
        )
        iterator = (
            x
            for x in correlated_kineto_events
            if x.device_type() == DeviceType.CUDA and x.name() == node.event.name
        )
        return next(iterator, None)

    def _cumulative_cuda_time(self, node: _ModuleTreeNode):
        "Return cuda time in microseconds"

        def _cumulative_cuda_time_recursive(node: _ModuleTreeNode):
            if node.is_leaf and (gpu_kineto_event := self._get_kineto_gpu_event(node)):
                return gpu_kineto_event.duration_ns() / 1000.0
            else:
                cumulative_cuda_time = 0
                for child in node.children:
                    cumulative_cuda_time += _cumulative_cuda_time_recursive(child)
                return cumulative_cuda_time

        return _cumulative_cuda_time_recursive(node)

    def _total_cuda_time(self):
        return sum([self._cumulative_cuda_time(root) for root in self._module_tree])

    def _build_stats_trees(self):
        summary_dict: dict[str, _StatsTreeNode] = {}
        total_cuda_time = self._total_cuda_time()

        def pct_cuda_time(cuda_time_us):
            return (cuda_time_us / total_cuda_time) * 100

        def build_summary_stats_tree_df(
            node: _ModuleTreeNode,
            parent: Optional[_StatsTreeNode] = None,
            summary_trace: tuple[str] = (),
        ):
            if event_has_module(node.event):
                name = event_module_repr(node.event)
                cuda_time_us = self._cumulative_cuda_time(node)
            elif gpu_kineto_event := self._get_kineto_gpu_event(node):
                name = gpu_kineto_event.name()
                cuda_time_us = gpu_kineto_event.duration_ns() / 1000.0
            else:
                return None

            summary_trace = summary_trace + (name,)
            if summary_trace in summary_dict:
                entry = summary_dict[summary_trace].entry
                entry.cuda_time_us += cuda_time_us
                entry.invocations += 1
                entry.pct_cuda_time = pct_cuda_time(entry.cuda_time_us)
            else:
                new_node = _StatsTreeNode(
                    entry=SummaryStatsEntry(
                        name=name,
                        cuda_time_us=cuda_time_us,
                        pct_cuda_time=pct_cuda_time(cuda_time_us),
                        invocations=1,
                    ),
                    children=[],
                    parent=parent,
                )
                if parent:
                    parent.children.append(new_node)
                summary_dict[summary_trace] = new_node

            for child in node.children:
                build_summary_stats_tree_df(
                    child, summary_dict[summary_trace], summary_trace
                )

            return summary_dict[summary_trace]

        self._summary_stats_tree = []
        for root in self._module_tree:
            self._summary_stats_tree.append(build_summary_stats_tree_df(root))

        def build_model_stats_tree_df(
            node: _ModuleTreeNode, parent: Optional[_StatsTreeNode] = None
        ):
            if event_has_module(
                node.event,
            ):
                name = event_module_repr(node.event)
                cuda_time_us = self._cumulative_cuda_time(node)
                cpu_time_us = node.event.duration_time_ns / 1000
                trace = ""
            elif gpu_kineto_event := self._get_kineto_gpu_event(node):
                name = gpu_kineto_event.name()
                cuda_time_us = gpu_kineto_event.duration_ns() / 1000.0
                cpu_time_us = 0
                trace = node.trace
            else:
                return None

            new_node = _StatsTreeNode(
                entry=ModelStatsEntry(
                    name=name,
                    cpu_time_us=cpu_time_us,
                    cuda_time_us=cuda_time_us,
                    pct_cuda_time=pct_cuda_time(cuda_time_us),
                    trace=trace,
                ),
                parent=parent,
                children=[],
            )
            if parent:
                parent.children.append(new_node)

            for child in node.children:
                build_model_stats_tree_df(child, new_node)

            return new_node

        self._model_stats_tree = []
        for root in self._module_tree:
            self._model_stats_tree.append(build_model_stats_tree_df(root))

    def _flatten_stats_tree(
        self, tree: list[_StatsTreeNode]
    ) -> list[tuple[int, StatsEntry]]:
        entries: list[tuple[int, StatsEntry]] = []

        def df_traversal(node: _StatsTreeNode, depth=0):
            entries.append((depth, node.entry))
            for child in node.children:
                df_traversal(child, depth=depth + 1)

        for root in tree:
            df_traversal(root)

        return entries

    def _convert_stats_tree_to_dict(self, tree: list[_StatsTreeNode]) -> list[dict]:
        root_dicts: list[dict] = []

        def df_traversal(node: _StatsTreeNode, curr_json_list: list[dict]):
            curr_json_list.append({"entry": asdict(node.entry), "children": []})
            for child in node.children:
                df_traversal(child, curr_json_list[-1]["children"])

        for root in tree:
            df_traversal(root, root_dicts)

        return root_dicts


class layerwise_profile(profile):
    def __init__(self, num_running_seqs: Optional[int] = None):
        """
        layerwise profile constructor.

        Args:
            num_running_seqs (Optional[int], optional): When given,
                num_running_seqs will be passed to LayerProfileResults
                for metadata update. Defaults to None.
        """
        super().__init__(
            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
            record_shapes=True,
            with_stack=True,
            with_modules=True,
            experimental_config=_ExperimentalConfig(verbose=True),
        )

        self.num_running_seqs = num_running_seqs

    def __enter__(self):
        return super().__enter__()

    def __exit__(self, exc_type, exc_val, exc_tb):
        super().__exit__(exc_type, exc_val, exc_tb)
        self.results = LayerwiseProfileResults(
            self.profiler.kineto_results, num_running_seqs=self.num_running_seqs
        )