[CPU][Profiler] Group torch profiler tables by input shape when record_shapes is on (#56016)
Signed-off-by: Chinmay Kulkarni <Chinmay.Kulkarni@amd.com> Co-authored-by: Li, Jiang <jiang1.li@intel.com>
Chinmay-Kulkarni-AMD committed
Sep 14, 2026 at 14:46 UTC
dc89fdfb0ede2918a884faae771a49bbbe7b6cf9
1 file changed
+8
-2
vllm/profiler/wrapper.py
+8
-2
@@ -19,6 +19,7 @@ import vllm.version
19
from vllm.config import ProfilerConfig
20
from vllm.config.profiler import _is_uri_path
21
from vllm.logger import init_logger
22
+from vllm.platforms import current_platform
23
24
logger = init_logger(__name__)
25
@@ -259,9 +260,14 @@ class TorchProfilerWrapper(WorkerProfiler):
260
sort_key: str,
261
row_limit: int | None = None,
262
) -> str:
263
+ group_by_input_shape = (
264
+ current_platform.is_cpu()
265
+ and self.profiler_config.torch_profiler_record_shapes
266
+ )
267
+ averages = self.profiler.key_averages(group_by_input_shape=group_by_input_shape)
268
if row_limit is None: # use profiler default row limit of 100
263
- return self.profiler.key_averages().table(sort_by=sort_key)
264
- return self.profiler.key_averages().table(
269
+ return averages.table(sort_by=sort_key)
270
+ return averages.table(
271
sort_by=sort_key,
272
row_limit=row_limit,
273
)