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 15, 08:30 UTC
[V1] fix torch profiling for V1 offline scenarios (#18445)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
May 27, 2025 at 23:16 UTC
774c5fde307dcadb8e0d2a58bf2d429eb6cb648d
4 files changed
+23
-51
benchmarks/benchmark_latency.py
+10
-27
@@ -6,13 +6,12 @@ import dataclasses
6
import json
7
import os
8
import time
9
-from pathlib import Path
9
from typing import Any, Optional
10
11
import numpy as np
13
-import torch
12
from tqdm import tqdm
13
14
+import vllm.envs as envs
15
from benchmark_utils import convert_to_pytorch_benchmark_format, write_to_json
16
from vllm import LLM, SamplingParams
17
from vllm.engine.arg_utils import EngineArgs
@@ -80,17 +79,9 @@ def main(args: argparse.Namespace):
79
80
def run_to_completion(profile_dir: Optional[str] = None):
81
if profile_dir:
83
- with torch.profiler.profile(
84
- activities=[
85
- torch.profiler.ProfilerActivity.CPU,
86
- torch.profiler.ProfilerActivity.CUDA,
87
- ],
88
- on_trace_ready=torch.profiler.tensorboard_trace_handler(
89
- str(profile_dir)
90
- ),
91
- ) as p:
92
- llm_generate()
93
- print(p.key_averages().table(sort_by="self_cuda_time_total"))
82
+ llm.start_profile()
83
+ llm_generate()
84
+ llm.stop_profile()
85
else:
86
start_time = time.perf_counter()
87
llm_generate()
@@ -103,11 +94,7 @@ def main(args: argparse.Namespace):
94
run_to_completion(profile_dir=None)
95
96
if args.profile:
106
- profile_dir = args.profile_result_dir
107
- if not profile_dir:
108
- profile_dir = (
109
- Path(".") / "vllm_benchmark_result" / f"latency_result_{time.time()}"
110
- )
97
+ profile_dir = envs.VLLM_TORCH_PROFILER_DIR
98
print(f"Profiling (results will be saved to '{profile_dir}')...")
99
run_to_completion(profile_dir=profile_dir)
100
return
@@ -164,15 +151,6 @@ if __name__ == "__main__":
151
action="store_true",
152
help="profile the generation process of a single batch",
153
)
167
- parser.add_argument(
168
- "--profile-result-dir",
169
- type=str,
170
- default=None,
171
- help=(
172
- "path to save the pytorch profiler output. Can be visualized "
173
- "with ui.perfetto.dev or Tensorboard."
174
- ),
175
- )
154
parser.add_argument(
155
"--output-json",
156
type=str,
@@ -193,4 +171,9 @@ if __name__ == "__main__":
171
# numbers. We need to disable prefix caching by default.
172
parser.set_defaults(enable_prefix_caching=False)
173
args = parser.parse_args()
174
+ if args.profile and not envs.VLLM_TORCH_PROFILER_DIR:
175
+ raise OSError(
176
+ "The environment variable 'VLLM_TORCH_PROFILER_DIR' is not set. "
177
+ "Please set it to a valid path to use torch profiler."
178
+ )
179
main(args)
vllm/benchmarks/latency.py
+9
-24
@@ -6,13 +6,12 @@ import dataclasses
6
import json
7
import os
8
import time
9
-from pathlib import Path
9
from typing import Any, Optional
10
11
import numpy as np
13
-import torch
12
from tqdm import tqdm
13
14
+import vllm.envs as envs
15
from vllm import LLM, SamplingParams
16
from vllm.benchmarks.utils import (convert_to_pytorch_benchmark_format,
17
write_to_json)
@@ -59,13 +58,6 @@ def add_cli_args(parser: argparse.ArgumentParser):
58
action="store_true",
59
help="profile the generation process of a single batch",
60
)
62
- parser.add_argument(
63
- "--profile-result-dir",
64
- type=str,
65
- default=None,
66
- help=("path to save the pytorch profiler output. Can be visualized "
67
- "with ui.perfetto.dev or Tensorboard."),
68
- )
61
parser.add_argument(
62
"--output-json",
63
type=str,
@@ -87,7 +79,10 @@ def add_cli_args(parser: argparse.ArgumentParser):
79
80
def main(args: argparse.Namespace):
81
print(args)
90
-
82
+ if args.profile and not envs.VLLM_TORCH_PROFILER_DIR:
83
+ raise OSError(
84
+ "The environment variable 'VLLM_TORCH_PROFILER_DIR' is not set. "
85
+ "Please set it to a valid path to use torch profiler.")
86
engine_args = EngineArgs.from_cli_args(args)
87
88
# NOTE(woosuk): If the request cannot be processed in a single batch,
@@ -131,16 +126,9 @@ def main(args: argparse.Namespace):
126
127
def run_to_completion(profile_dir: Optional[str] = None):
128
if profile_dir:
134
- with torch.profiler.profile(
135
- activities=[
136
- torch.profiler.ProfilerActivity.CPU,
137
- torch.profiler.ProfilerActivity.CUDA,
138
- ],
139
- on_trace_ready=torch.profiler.tensorboard_trace_handler(
140
- str(profile_dir)),
141
- ) as p:
142
- llm_generate()
143
- print(p.key_averages().table(sort_by="self_cuda_time_total"))
129
+ llm.start_profile()
130
+ llm_generate()
131
+ llm.stop_profile()
132
else:
133
start_time = time.perf_counter()
134
llm_generate()
@@ -153,10 +141,7 @@ def main(args: argparse.Namespace):
141
run_to_completion(profile_dir=None)
142
143
if args.profile:
156
- profile_dir = args.profile_result_dir
157
- if not profile_dir:
158
- profile_dir = (Path(".") / "vllm_benchmark_result" /
159
- f"latency_result_{time.time()}")
144
+ profile_dir = envs.VLLM_TORCH_PROFILER_DIR
145
print(f"Profiling (results will be saved to '{profile_dir}')...")
146
run_to_completion(profile_dir=profile_dir)
147
return
vllm/v1/worker/gpu_worker.py
+2
@@ -292,6 +292,8 @@ class Worker(WorkerBase):
292
self.profiler.start()
293
else:
294
self.profiler.stop()
295
+ print(self.profiler.key_averages().table(
296
+ sort_by="self_cuda_time_total"))
297
298
def execute_dummy_batch(self) -> None:
299
self.model_runner._dummy_run(1)
vllm/worker/worker.py
+2
@@ -128,6 +128,8 @@ class Worker(LocalOrDistributedWorkerBase):
128
if self.profiler is None:
129
raise RuntimeError("Profiler is not enabled.")
130
self.profiler.stop()
131
+ print(
132
+ self.profiler.key_averages().table(sort_by="self_cuda_time_total"))
133
134
def sleep(self, level: int = 1) -> None:
135
free_bytes_before_sleep = torch.cuda.mem_get_info()[0]