A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 00:00 UTC
[Misc] Fix mypy errors in vllm/profiler and remove from exclude list (#34959)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Feb 20, 2026 at 21:56 UTC
d38cd3dde549b8e421d7d0390799b985e13bd8ab
2 files changed
+37
-30
tools/pre_commit/mypy.py
-1
@@ -43,7 +43,6 @@ EXCLUDE = [
43
"vllm/benchmarks",
44
"vllm/config",
45
"vllm/device_allocator",
46
- "vllm/profiler",
46
"vllm/reasoning",
47
"vllm/tool_parser",
48
]
vllm/profiler/layerwise_profile.py
+37
-29
@@ -5,7 +5,7 @@ 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, TypeAlias
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
@@ -69,13 +69,14 @@ class ModelStatsEntry:
69
70
71
StatsEntry: TypeAlias = ModelStatsEntry | SummaryStatsEntry
72
+StatsEntryT = TypeVar("StatsEntryT", bound=StatsEntry)
73
74
75
@dataclass
75
-class _StatsTreeNode:
76
- entry: StatsEntry
77
- children: list[StatsEntry]
78
- parent: StatsEntry | None
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
@@ -84,8 +85,8 @@ class LayerwiseProfileResults(profile):
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)
87
- _model_stats_tree: list[_StatsTreeNode] = field(init=False)
88
- _summary_stats_tree: list[_StatsTreeNode] = 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
@@ -95,7 +96,7 @@ class LayerwiseProfileResults(profile):
96
self._build_module_tree()
97
self._build_stats_trees()
98
98
- def print_model_table(self, column_widths: dict[str, int] = None):
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
)
@@ -113,7 +114,7 @@ class LayerwiseProfileResults(profile):
114
)
115
)
116
116
- def print_summary_table(self, column_widths: dict[str, int] = None):
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
)
@@ -155,14 +156,14 @@ class LayerwiseProfileResults(profile):
156
157
@staticmethod
158
def _indent_row_names_based_on_depth(
158
- depths_rows: list[tuple[int, StatsEntry]],
159
+ depths_rows: list[tuple[int, StatsEntryT]],
160
indent_style: Callable[[int], str] | str = " ",
161
):
161
- indented_rows = []
162
+ indented_rows: list[StatsEntryT] = []
163
for depth, row in depths_rows:
164
if row.cuda_time_us == 0:
165
continue
165
- indented_row = copy.deepcopy(row)
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
@@ -240,7 +241,7 @@ class LayerwiseProfileResults(profile):
241
return sum([self._cumulative_cuda_time(root) for root in self._module_tree])
242
243
def _build_stats_trees(self):
243
- summary_dict: dict[str, _StatsTreeNode] = {}
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,9 +249,9 @@ class LayerwiseProfileResults(profile):
249
250
def build_summary_stats_tree_df(
251
node: _ModuleTreeNode,
251
- parent: _StatsTreeNode | None = None,
252
- summary_trace: tuple[str] = (),
253
- ):
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)
@@ -274,7 +275,6 @@ class LayerwiseProfileResults(profile):
275
pct_cuda_time=pct_cuda_time(cuda_time_us),
276
invocations=1,
277
),
277
- children=[],
278
parent=parent,
279
)
280
if parent:
@@ -290,11 +290,14 @@ class LayerwiseProfileResults(profile):
290
291
self._summary_stats_tree = []
292
for root in self._module_tree:
293
- self._summary_stats_tree.append(build_summary_stats_tree_df(root))
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(
296
- node: _ModuleTreeNode, parent: _StatsTreeNode | None = None
297
- ):
298
+ node: _ModuleTreeNode,
299
+ parent: _StatsTreeNode[ModelStatsEntry] | None = None,
300
+ ) -> _StatsTreeNode[ModelStatsEntry] | None:
301
if event_has_module(
302
node.event,
303
):
@@ -319,7 +322,6 @@ class LayerwiseProfileResults(profile):
322
trace=trace,
323
),
324
parent=parent,
322
- children=[],
325
)
326
if parent:
327
parent.children.append(new_node)
@@ -331,14 +333,16 @@ class LayerwiseProfileResults(profile):
333
334
self._model_stats_tree = []
335
for root in self._module_tree:
334
- self._model_stats_tree.append(build_model_stats_tree_df(root))
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(
337
- self, tree: list[_StatsTreeNode]
338
- ) -> list[tuple[int, StatsEntry]]:
339
- entries: list[tuple[int, StatsEntry]] = []
341
+ self, tree: list[_StatsTreeNode[StatsEntryT]]
342
+ ) -> list[tuple[int, StatsEntryT]]:
343
+ entries: list[tuple[int, StatsEntryT]] = []
344
341
- def df_traversal(node: _StatsTreeNode, depth=0):
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)
@@ -348,10 +352,14 @@ class LayerwiseProfileResults(profile):
352
353
return entries
354
351
- def _convert_stats_tree_to_dict(self, tree: list[_StatsTreeNode]) -> list[dict]:
352
- root_dicts: list[dict] = []
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
354
- def df_traversal(node: _StatsTreeNode, curr_json_list: list[dict]):
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"])