@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a8ec486592
Fork 0

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 11, 12:30 UTC

[Misc] Add a script to benchmark compilation time (#29919)

Signed-off-by: Bin Bao <binbao@meta.com>

Bin Bao committed Dec 14, 2025 at 08:02 UTC a8ec486592fd44db67a7390fb91f032ce69f80e1
3 files changed +349
vllm/benchmarks/startup.py new
+326
@@ -0,0 +1,326 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Benchmark the cold and warm startup time of vLLM models.
4 +
5 +This script measures total startup time (including model loading, compilation,
6 +and cache operations) for both cold and warm scenarios:
7 +- Cold startup: Fresh start with no caches (temporary cache directories)
8 +- Warm startup: Using cached compilation and model info
9 +"""
10 +
11 +import argparse
12 +import dataclasses
13 +import json
14 +import multiprocessing
15 +import os
16 +import shutil
17 +import tempfile
18 +import time
19 +from contextlib import contextmanager
20 +from typing import Any
21 +
22 +import numpy as np
23 +from tqdm import tqdm
24 +
25 +from vllm.benchmarks.lib.utils import (
26 + convert_to_pytorch_benchmark_format,
27 + write_to_json,
28 +)
29 +from vllm.engine.arg_utils import EngineArgs
30 +
31 +
32 +@contextmanager
33 +def cold_startup():
34 + """
35 + Context manager to measure cold startup time:
36 + 1. Uses a temporary directory for vLLM cache to avoid any pollution
37 + between cold startup iterations.
38 + 2. Uses inductor's fresh_cache to clear torch.compile caches.
39 + """
40 + from torch._inductor.utils import fresh_cache
41 +
42 + # Use temporary directory for caching to avoid any pollution between cold startups
43 + original_cache_root = os.environ.get("VLLM_CACHE_ROOT")
44 + temp_cache_dir = tempfile.mkdtemp(prefix="vllm_startup_bench_cold_")
45 + try:
46 + os.environ["VLLM_CACHE_ROOT"] = temp_cache_dir
47 + with fresh_cache():
48 + yield
49 + finally:
50 + # Clean up temporary cache directory
51 + shutil.rmtree(temp_cache_dir, ignore_errors=True)
52 + if original_cache_root:
53 + os.environ["VLLM_CACHE_ROOT"] = original_cache_root
54 + else:
55 + os.environ.pop("VLLM_CACHE_ROOT", None)
56 +
57 +
58 +def run_startup_in_subprocess(engine_args_dict, result_queue):
59 + """
60 + Run LLM startup in a subprocess and return timing metrics via a queue.
61 + This ensures complete isolation between iterations.
62 + """
63 + try:
64 + # Import inside the subprocess to avoid issues with forking
65 + from vllm import LLM
66 + from vllm.engine.arg_utils import EngineArgs
67 +
68 + engine_args = EngineArgs(**engine_args_dict)
69 +
70 + # Measure total startup time
71 + start_time = time.perf_counter()
72 +
73 + llm = LLM(**dataclasses.asdict(engine_args))
74 +
75 + total_startup_time = time.perf_counter() - start_time
76 +
77 + # Extract compilation time if available
78 + compilation_time = 0.0
79 + if hasattr(llm.llm_engine, "vllm_config"):
80 + vllm_config = llm.llm_engine.vllm_config
81 + if (
82 + hasattr(vllm_config, "compilation_config")
83 + and vllm_config.compilation_config is not None
84 + ):
85 + compilation_time = vllm_config.compilation_config.compilation_time
86 +
87 + result_queue.put(
88 + {
89 + "total_startup_time": total_startup_time,
90 + "compilation_time": compilation_time,
91 + }
92 + )
93 +
94 + except Exception as e:
95 + result_queue.put(None)
96 + result_queue.put(str(e))
97 +
98 +
99 +def save_to_pytorch_benchmark_format(
100 + args: argparse.Namespace, results: dict[str, Any]
101 +) -> None:
102 + base_name = os.path.splitext(args.output_json)[0]
103 +
104 + cold_startup_records = convert_to_pytorch_benchmark_format(
105 + args=args,
106 + metrics={
107 + "avg_cold_startup_time": results["avg_cold_startup_time"],
108 + },
109 + extra_info={
110 + "cold_startup_times": results["cold_startup_times"],
111 + "cold_startup_percentiles": results["cold_startup_percentiles"],
112 + },
113 + )
114 + if cold_startup_records:
115 + write_to_json(f"{base_name}.cold_startup.pytorch.json", cold_startup_records)
116 +
117 + cold_compilation_records = convert_to_pytorch_benchmark_format(
118 + args=args,
119 + metrics={
120 + "avg_cold_compilation_time": results["avg_cold_compilation_time"],
121 + },
122 + extra_info={
123 + "cold_compilation_times": results["cold_compilation_times"],
124 + "cold_compilation_percentiles": results["cold_compilation_percentiles"],
125 + },
126 + )
127 + if cold_compilation_records:
128 + write_to_json(
129 + f"{base_name}.cold_compilation.pytorch.json", cold_compilation_records
130 + )
131 +
132 + warm_startup_records = convert_to_pytorch_benchmark_format(
133 + args=args,
134 + metrics={
135 + "avg_warm_startup_time": results["avg_warm_startup_time"],
136 + },
137 + extra_info={
138 + "warm_startup_times": results["warm_startup_times"],
139 + "warm_startup_percentiles": results["warm_startup_percentiles"],
140 + },
141 + )
142 + if warm_startup_records:
143 + write_to_json(f"{base_name}.warm_startup.pytorch.json", warm_startup_records)
144 +
145 + warm_compilation_records = convert_to_pytorch_benchmark_format(
146 + args=args,
147 + metrics={
148 + "avg_warm_compilation_time": results["avg_warm_compilation_time"],
149 + },
150 + extra_info={
151 + "warm_compilation_times": results["warm_compilation_times"],
152 + "warm_compilation_percentiles": results["warm_compilation_percentiles"],
153 + },
154 + )
155 + if warm_compilation_records:
156 + write_to_json(
157 + f"{base_name}.warm_compilation.pytorch.json", warm_compilation_records
158 + )
159 +
160 +
161 +def add_cli_args(parser: argparse.ArgumentParser):
162 + parser.add_argument(
163 + "--num-iters-cold",
164 + type=int,
165 + default=5,
166 + help="Number of cold startup iterations.",
167 + )
168 + parser.add_argument(
169 + "--num-iters-warmup",
170 + type=int,
171 + default=3,
172 + help="Number of warmup iterations before benchmarking warm startups.",
173 + )
174 + parser.add_argument(
175 + "--num-iters-warm",
176 + type=int,
177 + default=5,
178 + help="Number of warm startup iterations.",
179 + )
180 + parser.add_argument(
181 + "--output-json",
182 + type=str,
183 + default=None,
184 + help="Path to save the startup time results in JSON format.",
185 + )
186 +
187 + parser = EngineArgs.add_cli_args(parser)
188 + return parser
189 +
190 +
191 +def main(args: argparse.Namespace):
192 + # Set multiprocessing start method to 'spawn' for clean process isolation
193 + # This ensures each subprocess starts fresh without inheriting state
194 + multiprocessing.set_start_method("spawn", force=True)
195 +
196 + engine_args = EngineArgs.from_cli_args(args)
197 +
198 + def create_llm_and_measure_startup():
199 + """
200 + Create LLM instance in a subprocess and measure startup time.
201 + Returns timing metrics, using subprocess for complete isolation.
202 + """
203 + # Convert engine_args to dictionary for pickling
204 + engine_args_dict = dataclasses.asdict(engine_args)
205 +
206 + # Create a queue for inter-process communication
207 + result_queue = multiprocessing.Queue()
208 + process = multiprocessing.Process(
209 + target=run_startup_in_subprocess,
210 + args=(
211 + engine_args_dict,
212 + result_queue,
213 + ),
214 + )
215 + process.start()
216 + process.join()
217 +
218 + if not result_queue.empty():
219 + result = result_queue.get()
220 + if result is None:
221 + if not result_queue.empty():
222 + error_msg = result_queue.get()
223 + raise RuntimeError(f"Subprocess failed: {error_msg}")
224 + else:
225 + raise RuntimeError("Subprocess failed with unknown error")
226 + return result
227 + else:
228 + raise RuntimeError("Subprocess did not return a result")
229 +
230 + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
231 + print("Setting VLLM_ENABLE_V1_MULTIPROCESSING=0 to collect startup metrics.\n")
232 +
233 + print("Measuring cold startup time...\n")
234 + cold_startup_times = []
235 + cold_compilation_times = []
236 + for i in tqdm(range(args.num_iters_cold), desc="Cold startup iterations"):
237 + with cold_startup():
238 + metrics = create_llm_and_measure_startup()
239 + cold_startup_times.append(metrics["total_startup_time"])
240 + cold_compilation_times.append(metrics["compilation_time"])
241 +
242 + # Warmup for warm startup
243 + print("\nWarming up for warm startup measurement...\n")
244 + for _ in tqdm(range(args.num_iters_warmup), desc="Warmup iterations"):
245 + create_llm_and_measure_startup()
246 +
247 + print("\nMeasuring warm startup time...\n")
248 + warm_startup_times = []
249 + warm_compilation_times = []
250 + for i in tqdm(range(args.num_iters_warm), desc="Warm startup iterations"):
251 + metrics = create_llm_and_measure_startup()
252 + warm_startup_times.append(metrics["total_startup_time"])
253 + warm_compilation_times.append(metrics["compilation_time"])
254 +
255 + # Calculate statistics
256 + cold_startup_array = np.array(cold_startup_times)
257 + cold_compilation_array = np.array(cold_compilation_times)
258 + warm_startup_array = np.array(warm_startup_times)
259 + warm_compilation_array = np.array(warm_compilation_times)
260 +
261 + avg_cold_startup = np.mean(cold_startup_array)
262 + avg_cold_compilation = np.mean(cold_compilation_array)
263 + avg_warm_startup = np.mean(warm_startup_array)
264 + avg_warm_compilation = np.mean(warm_compilation_array)
265 +
266 + percentages = [10, 25, 50, 75, 90, 99]
267 + cold_startup_percentiles = np.percentile(cold_startup_array, percentages)
268 + cold_compilation_percentiles = np.percentile(cold_compilation_array, percentages)
269 + warm_startup_percentiles = np.percentile(warm_startup_array, percentages)
270 + warm_compilation_percentiles = np.percentile(warm_compilation_array, percentages)
271 +
272 + print("\n" + "=" * 60)
273 + print("STARTUP TIME BENCHMARK RESULTS")
274 + print("=" * 60)
275 +
276 + # Cold startup statistics
277 + print("\nCOLD STARTUP:")
278 + print(f"Avg total startup time: {avg_cold_startup:.2f} seconds")
279 + print(f"Avg compilation time: {avg_cold_compilation:.2f} seconds")
280 + print("Startup time percentiles:")
281 + for percentage, percentile in zip(percentages, cold_startup_percentiles):
282 + print(f" {percentage}%: {percentile:.2f} seconds")
283 + print("Compilation time percentiles:")
284 + for percentage, percentile in zip(percentages, cold_compilation_percentiles):
285 + print(f" {percentage}%: {percentile:.2f} seconds")
286 +
287 + # Warm startup statistics
288 + print("\nWARM STARTUP:")
289 + print(f"Avg total startup time: {avg_warm_startup:.2f} seconds")
290 + print(f"Avg compilation time: {avg_warm_compilation:.2f} seconds")
291 + print("Startup time percentiles:")
292 + for percentage, percentile in zip(percentages, warm_startup_percentiles):
293 + print(f" {percentage}%: {percentile:.2f} seconds")
294 + print("Compilation time percentiles:")
295 + for percentage, percentile in zip(percentages, warm_compilation_percentiles):
296 + print(f" {percentage}%: {percentile:.2f} seconds")
297 +
298 + print("=" * 60)
299 +
300 + # Output JSON results if specified
301 + if args.output_json:
302 + results = {
303 + "avg_cold_startup_time": float(avg_cold_startup),
304 + "avg_cold_compilation_time": float(avg_cold_compilation),
305 + "cold_startup_times": cold_startup_times,
306 + "cold_compilation_times": cold_compilation_times,
307 + "cold_startup_percentiles": dict(
308 + zip(percentages, cold_startup_percentiles.tolist())
309 + ),
310 + "cold_compilation_percentiles": dict(
311 + zip(percentages, cold_compilation_percentiles.tolist())
312 + ),
313 + "avg_warm_startup_time": float(avg_warm_startup),
314 + "avg_warm_compilation_time": float(avg_warm_compilation),
315 + "warm_startup_times": warm_startup_times,
316 + "warm_compilation_times": warm_compilation_times,
317 + "warm_startup_percentiles": dict(
318 + zip(percentages, warm_startup_percentiles.tolist())
319 + ),
320 + "warm_compilation_percentiles": dict(
321 + zip(percentages, warm_compilation_percentiles.tolist())
322 + ),
323 + }
324 + with open(args.output_json, "w") as f:
325 + json.dump(results, f, indent=4)
326 + save_to_pytorch_benchmark_format(args, results)
vllm/entrypoints/cli/__init__.py
+2
@@ -2,12 +2,14 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 from vllm.entrypoints.cli.benchmark.latency import BenchmarkLatencySubcommand
4 from vllm.entrypoints.cli.benchmark.serve import BenchmarkServingSubcommand
5 +from vllm.entrypoints.cli.benchmark.startup import BenchmarkStartupSubcommand
6 from vllm.entrypoints.cli.benchmark.sweep import BenchmarkSweepSubcommand
7 from vllm.entrypoints.cli.benchmark.throughput import BenchmarkThroughputSubcommand
8
9 __all__: list[str] = [
10 "BenchmarkLatencySubcommand",
11 "BenchmarkServingSubcommand",
12 + "BenchmarkStartupSubcommand",
13 "BenchmarkSweepSubcommand",
14 "BenchmarkThroughputSubcommand",
15 ]
vllm/entrypoints/cli/benchmark/startup.py new
+21
@@ -0,0 +1,21 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import argparse
4 +
5 +from vllm.benchmarks.startup import add_cli_args, main
6 +from vllm.entrypoints.cli.benchmark.base import BenchmarkSubcommandBase
7 +
8 +
9 +class BenchmarkStartupSubcommand(BenchmarkSubcommandBase):
10 + """The `startup` subcommand for `vllm bench`."""
11 +
12 + name = "startup"
13 + help = "Benchmark the startup time of vLLM models."
14 +
15 + @classmethod
16 + def add_cli_args(cls, parser: argparse.ArgumentParser) -> None:
17 + add_cli_args(parser)
18 +
19 + @staticmethod
20 + def cmd(args: argparse.Namespace) -> None:
21 + main(args)