From dc89fdfb0ede2918a884faae771a49bbbe7b6cf9 Mon Sep 17 00:00:00 2001 From: Chinmay-Kulkarni-AMD Date: Mon, 14 Sep 2026 14:46:48 +0530 Subject: [PATCH] [CPU][Profiler] Group torch profiler tables by input shape when record_shapes is on (#56016) Signed-off-by: Chinmay Kulkarni Co-authored-by: Li, Jiang --- vllm/profiler/wrapper.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm/profiler/wrapper.py b/vllm/profiler/wrapper.py index 877a390d6d..fb89631986 100644 --- a/vllm/profiler/wrapper.py +++ b/vllm/profiler/wrapper.py @@ -19,6 +19,7 @@ import vllm.version from vllm.config import ProfilerConfig from vllm.config.profiler import _is_uri_path from vllm.logger import init_logger +from vllm.platforms import current_platform logger = init_logger(__name__) @@ -259,9 +260,14 @@ class TorchProfilerWrapper(WorkerProfiler): sort_key: str, row_limit: int | None = None, ) -> str: + group_by_input_shape = ( + current_platform.is_cpu() + and self.profiler_config.torch_profiler_record_shapes + ) + averages = self.profiler.key_averages(group_by_input_shape=group_by_input_shape) if row_limit is None: # use profiler default row limit of 100 - return self.profiler.key_averages().table(sort_by=sort_key) - return self.profiler.key_averages().table( + return averages.table(sort_by=sort_key) + return averages.table( sort_by=sort_key, row_limit=row_limit, )