Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions vllm/v1/worker/gpu_worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -397,8 +397,10 @@ def profile(self, is_start: bool = True):
self.profiler.start()
else:
self.profiler.stop()
print(self.profiler.key_averages().table(
sort_by="self_cuda_time_total"))
# only print profiler results on rank 0
if self.rank == 0:
print(self.profiler.key_averages().table(
sort_by="self_cuda_time_total"))

def execute_dummy_batch(self) -> None:
self.model_runner._dummy_run(1)
Expand Down
6 changes: 4 additions & 2 deletions vllm/worker/worker.py
Original file line number Diff line number Diff line change
Expand Up @@ -131,8 +131,10 @@ def stop_profile(self):
if self.profiler is None:
raise RuntimeError("Profiler is not enabled.")
self.profiler.stop()
print(
self.profiler.key_averages().table(sort_by="self_cuda_time_total"))
# only print profiler results on rank 0
if self.rank == 0:
print(self.profiler.key_averages().table(
sort_by="self_cuda_time_total"))

def sleep(self, level: int = 1) -> None:
free_bytes_before_sleep = torch.cuda.mem_get_info()[0]
Expand Down