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 14, 05:30 UTC
Enable --profile in 'vllm bench throughput' (#24575)
Signed-off-by: Tomas Ruiz <tomas.ruiz.te@gmail.com>
Tomas Ruiz committed
Sep 11, 2025 at 08:06 UTC
ee0bc5e1b42c98de83b6364ee32f1a21422bfae7
1 file changed
+34
-4
vllm/benchmarks/throughput.py
+34
-4
@@ -37,6 +37,7 @@ def run_vllm(
37
requests: list[SampleRequest],
38
n: int,
39
engine_args: EngineArgs,
40
+ do_profile: bool,
41
disable_detokenize: bool = False,
42
) -> tuple[float, Optional[list[RequestOutput]]]:
43
from vllm import LLM, SamplingParams
@@ -75,10 +76,14 @@ def run_vllm(
76
outputs = None
77
if not use_beam_search:
78
start = time.perf_counter()
79
+ if do_profile:
80
+ llm.start_profile()
81
outputs = llm.generate(prompts,
82
sampling_params,
83
lora_request=lora_requests,
84
use_tqdm=True)
85
+ if do_profile:
86
+ llm.stop_profile()
87
end = time.perf_counter()
88
else:
89
assert lora_requests is None, "BeamSearch API does not support LoRA"
@@ -88,6 +93,8 @@ def run_vllm(
93
for request in requests:
94
assert request.expected_output_len == output_len
95
start = time.perf_counter()
96
+ if do_profile:
97
+ llm.start_profile()
98
llm.beam_search(
99
prompts,
100
BeamSearchParams(
@@ -95,6 +102,8 @@ def run_vllm(
102
max_tokens=output_len,
103
ignore_eos=True,
104
))
105
+ if do_profile:
106
+ llm.stop_profile()
107
end = time.perf_counter()
108
return end - start, outputs
109
@@ -103,6 +112,7 @@ def run_vllm_chat(
112
requests: list[SampleRequest],
113
n: int,
114
engine_args: EngineArgs,
115
+ do_profile: bool,
116
disable_detokenize: bool = False) -> tuple[float, list[RequestOutput]]:
117
"""
118
Run vLLM chat benchmark. This function is recommended ONLY for benchmarking
@@ -133,7 +143,11 @@ def run_vllm_chat(
143
detokenize=not disable_detokenize,
144
))
145
start = time.perf_counter()
146
+ if do_profile:
147
+ llm.start_profile()
148
outputs = llm.chat(prompts, sampling_params, use_tqdm=True)
149
+ if do_profile:
150
+ llm.stop_profile()
151
end = time.perf_counter()
152
return end - start, outputs
153
@@ -142,6 +156,7 @@ async def run_vllm_async(
156
requests: list[SampleRequest],
157
n: int,
158
engine_args: AsyncEngineArgs,
159
+ do_profile: bool,
160
disable_frontend_multiprocessing: bool = False,
161
disable_detokenize: bool = False,
162
) -> float:
@@ -185,6 +200,8 @@ async def run_vllm_async(
200
201
generators = []
202
start = time.perf_counter()
203
+ if do_profile:
204
+ await llm.start_profile()
205
for i, (prompt, sp,
206
lr) in enumerate(zip(prompts, sampling_params, lora_requests)):
207
generator = llm.generate(prompt,
@@ -195,6 +212,8 @@ async def run_vllm_async(
212
all_gens = merge_async_iterators(*generators)
213
async for i, res in all_gens:
214
pass
215
+ if do_profile:
216
+ await llm.stop_profile()
217
end = time.perf_counter()
218
return end - start
219
@@ -543,6 +562,12 @@ def add_cli_args(parser: argparse.ArgumentParser):
562
type=str,
563
default=None,
564
help="Split of the HF dataset.")
565
+ parser.add_argument(
566
+ "--profile",
567
+ action="store_true",
568
+ default=False,
569
+ help="Use Torch Profiler. The env variable "
570
+ "VLLM_TORCH_PROFILER_DIR must be set to enable profiler.")
571
572
# prefix repetition dataset
573
prefix_repetition_group = parser.add_argument_group(
@@ -600,22 +625,27 @@ def main(args: argparse.Namespace):
625
requests,
626
args.n,
627
AsyncEngineArgs.from_cli_args(args),
603
- args.disable_frontend_multiprocessing,
604
- args.disable_detokenize,
628
+ disable_frontend_multiprocessing=args.disable_frontend_multiprocessing,
629
+ disable_detokenize=args.disable_detokenize,
630
+ do_profile=args.profile,
631
))
632
else:
633
elapsed_time, request_outputs = run_vllm(
634
requests, args.n, EngineArgs.from_cli_args(args),
609
- args.disable_detokenize)
635
+ disable_detokenize=args.disable_detokenize,
636
+ do_profile=args.profile)
637
elif args.backend == "hf":
638
assert args.tensor_parallel_size == 1
639
+ if args.profile:
640
+ raise NotImplementedError(
641
+ "Profiling not implemented yet for backend='hf'.")
642
elapsed_time = run_hf(requests, args.model, tokenizer, args.n,
643
args.hf_max_batch_size, args.trust_remote_code,
644
args.disable_detokenize)
645
elif args.backend == "vllm-chat":
646
elapsed_time, request_outputs = run_vllm_chat(
647
requests, args.n, EngineArgs.from_cli_args(args),
618
- args.disable_detokenize)
648
+ disable_detokenize=args.disable_detokenize, do_profile=args.profile)
649
else:
650
raise ValueError(f"Unknown backend: {args.backend}")
651