From 467e8a0cd7b1ed9a69b7f6fccaeab75141838a35 Mon Sep 17 00:00:00 2001 From: Yuan Tang Date: Sun, 19 Jan 2025 20:04:25 -0500 Subject: [PATCH 1/3] [Bugfix] Fix incorrect types in LayerwiseProfileResults Signed-off-by: Yuan Tang --- vllm/profiler/layerwise_profile.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/profiler/layerwise_profile.py b/vllm/profiler/layerwise_profile.py index 33babfebdca1..b51473816b33 100644 --- a/vllm/profiler/layerwise_profile.py +++ b/vllm/profiler/layerwise_profile.py @@ -128,7 +128,7 @@ def export_summary_stats_table_csv(self, filename: str): ]) df.to_csv(filename) - def convert_stats_to_dict(self) -> str: + def convert_stats_to_dict(self) -> dict: return { "metadata": { "num_running_seqs": self.num_running_seqs @@ -227,7 +227,7 @@ def _total_cuda_time(self): [self._cumulative_cuda_time(root) for root in self._module_tree]) def _build_stats_trees(self): - summary_dict: Dict[str, self.StatsTreeNode] = {} + summary_dict: Dict[str, _StatsTreeNode] = {} total_cuda_time = self._total_cuda_time() def pct_cuda_time(cuda_time_us): From e706e9d98b782eac1834b3860704667996f186b7 Mon Sep 17 00:00:00 2001 From: Yuan Tang Date: Sun, 19 Jan 2025 21:35:51 -0500 Subject: [PATCH 2/3] more precise type Signed-off-by: Yuan Tang --- vllm/profiler/layerwise_profile.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/profiler/layerwise_profile.py b/vllm/profiler/layerwise_profile.py index b51473816b33..db251a1b8030 100644 --- a/vllm/profiler/layerwise_profile.py +++ b/vllm/profiler/layerwise_profile.py @@ -1,7 +1,7 @@ import copy from collections import defaultdict from dataclasses import asdict, dataclass, field -from typing import Callable, Dict, List, Optional, Tuple, TypeAlias, Union +from typing import Callable, Dict, List, Optional, Tuple, TypeAlias, Union, Any import pandas as pd from torch._C._autograd import DeviceType, _KinetoEvent, _ProfilerResult @@ -128,7 +128,7 @@ def export_summary_stats_table_csv(self, filename: str): ]) df.to_csv(filename) - def convert_stats_to_dict(self) -> dict: + def convert_stats_to_dict(self) -> dict[str, Any]: return { "metadata": { "num_running_seqs": self.num_running_seqs From 121036f7b7411bda46d8a5e40f7b40666c4b05c1 Mon Sep 17 00:00:00 2001 From: Yuan Tang Date: Sun, 19 Jan 2025 21:46:57 -0500 Subject: [PATCH 3/3] Fix lint Signed-off-by: Yuan Tang --- vllm/profiler/layerwise_profile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/profiler/layerwise_profile.py b/vllm/profiler/layerwise_profile.py index db251a1b8030..29c0edd0ee53 100644 --- a/vllm/profiler/layerwise_profile.py +++ b/vllm/profiler/layerwise_profile.py @@ -1,7 +1,7 @@ import copy from collections import defaultdict from dataclasses import asdict, dataclass, field -from typing import Callable, Dict, List, Optional, Tuple, TypeAlias, Union, Any +from typing import Any, Callable, Dict, List, Optional, Tuple, TypeAlias, Union import pandas as pd from torch._C._autograd import DeviceType, _KinetoEvent, _ProfilerResult