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 13, 17:00 UTC
[CPU] Enable torch profiling (#28130)
Signed-off-by: Aditya Tewari <aditya.tewari@arm.com>
Aditya Tewari committed
Nov 6, 2025 at 07:32 UTC
3755c14532ae7278ada06c02aa90aa92f70556fa
1 file changed
+36
vllm/v1/worker/cpu_worker.py
+36
@@ -37,6 +37,28 @@ class CPUWorker(Worker):
37
38
self.parallel_config.disable_custom_all_reduce = True
39
40
+ if envs.VLLM_TORCH_PROFILER_DIR:
41
+ torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
42
+ worker_name = f"{vllm_config.instance_id}-rank-{self.rank}"
43
+ logger.info(
44
+ "Profiling enabled. Traces will be saved to: %s",
45
+ torch_profiler_trace_dir,
46
+ )
47
+ self.profiler = torch.profiler.profile(
48
+ activities=[
49
+ torch.profiler.ProfilerActivity.CPU,
50
+ ],
51
+ record_shapes=envs.VLLM_TORCH_PROFILER_RECORD_SHAPES,
52
+ profile_memory=envs.VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY,
53
+ with_stack=envs.VLLM_TORCH_PROFILER_WITH_STACK,
54
+ with_flops=envs.VLLM_TORCH_PROFILER_WITH_FLOPS,
55
+ on_trace_ready=torch.profiler.tensorboard_trace_handler(
56
+ torch_profiler_trace_dir, worker_name=worker_name, use_gzip=False
57
+ ),
58
+ )
59
+ else:
60
+ self.profiler = None
61
+
62
def init_device(self):
63
# Setup OpenMP threads affinity.
64
omp_cpuids = envs.VLLM_CPU_OMP_THREADS_BIND
@@ -166,3 +188,17 @@ class CPUWorker(Worker):
188
[(x.id, x.physical_core) for x in logical_cpu_list],
189
)
190
return ",".join([str(x.id) for x in logical_cpu_list])
191
+
192
+ def profile(self, is_start: bool = True):
193
+ if self.profiler is None:
194
+ raise RuntimeError("Profiler is not enabled.")
195
+ if is_start:
196
+ self.profiler.start()
197
+ else:
198
+ self.profiler.stop()
199
+ if self.local_rank == 0:
200
+ logger.info(
201
+ self.profiler.key_averages().table(
202
+ sort_by="self_cpu_time_total", row_limit=50
203
+ )
204
+ )