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 5, 04:45 UTC
Measure encoder compile time seperate from llm backbone (#39240)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
Lucas Kabela committed
Apr 14, 2026 at 07:52 UTC
c08f3b2a62fd8ddcb644ae91697a7981f2db4097
9 files changed
+188
-149
vllm/benchmarks/startup.py
+115
-130
@@ -16,7 +16,7 @@ import shutil
16
import tempfile
17
import time
18
from contextlib import contextmanager
19
-from typing import Any
19
+from typing import Any, NamedTuple
20
21
import numpy as np
22
from tqdm import tqdm
@@ -27,6 +27,82 @@ from vllm.benchmarks.lib.utils import (
27
)
28
from vllm.engine.arg_utils import EngineArgs
29
30
+PERCENTAGES = [10, 25, 50, 75, 90, 99]
31
+
32
+
33
+class MetricDesc(NamedTuple):
34
+ """Descriptor for a metric to collect from each iteration."""
35
+
36
+ iter_key: str # key in the iteration result dict
37
+ suffix: str # result key suffix, e.g. "startup", "compilation"
38
+ display_name: str
39
+
40
+
41
+class MetricStats(NamedTuple):
42
+ """Aggregated statistics for a single benchmark metric."""
43
+
44
+ key: str # e.g. "cold_startup", "warm_encoder_compilation"
45
+ display_name: str
46
+ values: list[float]
47
+ avg: float
48
+ percentiles: dict[int, float]
49
+
50
+
51
+_BASE_METRICS = [
52
+ MetricDesc("total_startup_time", "startup", "Startup time"),
53
+ MetricDesc("compilation_time", "compilation", "Compilation time"),
54
+]
55
+_ENCODER_METRIC = MetricDesc(
56
+ "encoder_compilation_time",
57
+ "encoder_compilation",
58
+ "Encoder compilation time",
59
+)
60
+
61
+
62
+def _compute_metric(
63
+ phase: str,
64
+ desc: MetricDesc,
65
+ iterations: list[dict[str, float]],
66
+) -> MetricStats:
67
+ values = [m[desc.iter_key] for m in iterations]
68
+ arr = np.array(values)
69
+ return MetricStats(
70
+ key=f"{phase}_{desc.suffix}",
71
+ display_name=desc.display_name,
72
+ values=values,
73
+ avg=float(np.mean(arr)),
74
+ percentiles=dict(zip(PERCENTAGES, np.percentile(arr, PERCENTAGES).tolist())),
75
+ )
76
+
77
+
78
+def _collect_phase_metrics(
79
+ phase: str,
80
+ iterations: list[dict[str, float]],
81
+ has_encoder: bool,
82
+) -> list[MetricStats]:
83
+ metrics = [_compute_metric(phase, desc, iterations) for desc in _BASE_METRICS]
84
+ if has_encoder:
85
+ metrics.append(_compute_metric(phase, _ENCODER_METRIC, iterations))
86
+ return metrics
87
+
88
+
89
+def _print_phase(phase_name: str, metrics: list[MetricStats]) -> None:
90
+ print(f"\n{phase_name}:")
91
+ for m in metrics:
92
+ print(f"Avg {m.display_name.lower()}: {m.avg:.2f} seconds")
93
+ for m in metrics:
94
+ print(f"{m.display_name} percentiles:")
95
+ for pct, val in m.percentiles.items():
96
+ print(f" {pct}%: {val:.2f} seconds")
97
+
98
+
99
+def _metric_to_json(m: MetricStats) -> dict[str, Any]:
100
+ return {
101
+ f"avg_{m.key}_time": m.avg,
102
+ f"{m.key}_times": m.values,
103
+ f"{m.key}_percentiles": m.percentiles,
104
+ }
105
+
106
107
@contextmanager
108
def cold_startup():
@@ -72,6 +148,7 @@ def run_startup_in_subprocess(engine_args, result_queue):
148
149
# Extract compilation time if available
150
compilation_time = 0.0
151
+ encoder_compilation_time = 0.0
152
if hasattr(llm.llm_engine, "vllm_config"):
153
vllm_config = llm.llm_engine.vllm_config
154
if (
@@ -79,11 +156,15 @@ def run_startup_in_subprocess(engine_args, result_queue):
156
and vllm_config.compilation_config is not None
157
):
158
compilation_time = vllm_config.compilation_config.compilation_time
159
+ encoder_compilation_time = (
160
+ vllm_config.compilation_config.encoder_compilation_time
161
+ )
162
163
result_queue.put(
164
{
165
"total_startup_time": total_startup_time,
166
"compilation_time": compilation_time,
167
+ "encoder_compilation_time": encoder_compilation_time,
168
}
169
)
170
@@ -93,65 +174,20 @@ def run_startup_in_subprocess(engine_args, result_queue):
174
175
176
def save_to_pytorch_benchmark_format(
96
- args: argparse.Namespace, results: dict[str, Any]
177
+ args: argparse.Namespace, metrics: list[MetricStats]
178
) -> None:
179
base_name = os.path.splitext(args.output_json)[0]
99
-
100
- cold_startup_records = convert_to_pytorch_benchmark_format(
101
- args=args,
102
- metrics={
103
- "avg_cold_startup_time": [results["avg_cold_startup_time"]],
104
- },
105
- extra_info={
106
- "cold_startup_times": results["cold_startup_times"],
107
- "cold_startup_percentiles": results["cold_startup_percentiles"],
108
- },
109
- )
110
- if cold_startup_records:
111
- write_to_json(f"{base_name}.cold_startup.pytorch.json", cold_startup_records)
112
-
113
- cold_compilation_records = convert_to_pytorch_benchmark_format(
114
- args=args,
115
- metrics={
116
- "avg_cold_compilation_time": [results["avg_cold_compilation_time"]],
117
- },
118
- extra_info={
119
- "cold_compilation_times": results["cold_compilation_times"],
120
- "cold_compilation_percentiles": results["cold_compilation_percentiles"],
121
- },
122
- )
123
- if cold_compilation_records:
124
- write_to_json(
125
- f"{base_name}.cold_compilation.pytorch.json", cold_compilation_records
126
- )
127
-
128
- warm_startup_records = convert_to_pytorch_benchmark_format(
129
- args=args,
130
- metrics={
131
- "avg_warm_startup_time": [results["avg_warm_startup_time"]],
132
- },
133
- extra_info={
134
- "warm_startup_times": results["warm_startup_times"],
135
- "warm_startup_percentiles": results["warm_startup_percentiles"],
136
- },
137
- )
138
- if warm_startup_records:
139
- write_to_json(f"{base_name}.warm_startup.pytorch.json", warm_startup_records)
140
-
141
- warm_compilation_records = convert_to_pytorch_benchmark_format(
142
- args=args,
143
- metrics={
144
- "avg_warm_compilation_time": [results["avg_warm_compilation_time"]],
145
- },
146
- extra_info={
147
- "warm_compilation_times": results["warm_compilation_times"],
148
- "warm_compilation_percentiles": results["warm_compilation_percentiles"],
149
- },
150
- )
151
- if warm_compilation_records:
152
- write_to_json(
153
- f"{base_name}.warm_compilation.pytorch.json", warm_compilation_records
180
+ for m in metrics:
181
+ records = convert_to_pytorch_benchmark_format(
182
+ args=args,
183
+ metrics={f"avg_{m.key}_time": [m.avg]},
184
+ extra_info={
185
+ f"{m.key}_times": m.values,
186
+ f"{m.key}_percentiles": m.percentiles,
187
+ },
188
)
189
+ if records:
190
+ write_to_json(f"{base_name}.{m.key}.pytorch.json", records)
191
192
193
def add_cli_args(parser: argparse.ArgumentParser):
@@ -224,97 +260,46 @@ def main(args: argparse.Namespace):
260
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
261
print("Setting VLLM_ENABLE_V1_MULTIPROCESSING=0 to collect startup metrics.\n")
262
263
+ # Collect cold startup iterations
264
print("Measuring cold startup time...\n")
228
- cold_startup_times = []
229
- cold_compilation_times = []
265
+ cold_iterations = []
266
for i in tqdm(range(args.num_iters_cold), desc="Cold startup iterations"):
267
with cold_startup():
232
- metrics = create_llm_and_measure_startup()
233
- cold_startup_times.append(metrics["total_startup_time"])
234
- cold_compilation_times.append(metrics["compilation_time"])
268
+ cold_iterations.append(create_llm_and_measure_startup())
269
270
# Warmup for warm startup
271
print("\nWarming up for warm startup measurement...\n")
272
for _ in tqdm(range(args.num_iters_warmup), desc="Warmup iterations"):
273
create_llm_and_measure_startup()
274
275
+ # Collect warm startup iterations
276
print("\nMeasuring warm startup time...\n")
242
- warm_startup_times = []
243
- warm_compilation_times = []
277
+ warm_iterations = []
278
for i in tqdm(range(args.num_iters_warm), desc="Warm startup iterations"):
245
- metrics = create_llm_and_measure_startup()
246
- warm_startup_times.append(metrics["total_startup_time"])
247
- warm_compilation_times.append(metrics["compilation_time"])
248
-
249
- # Calculate statistics
250
- cold_startup_array = np.array(cold_startup_times)
251
- cold_compilation_array = np.array(cold_compilation_times)
252
- warm_startup_array = np.array(warm_startup_times)
253
- warm_compilation_array = np.array(warm_compilation_times)
254
-
255
- avg_cold_startup = np.mean(cold_startup_array)
256
- avg_cold_compilation = np.mean(cold_compilation_array)
257
- avg_warm_startup = np.mean(warm_startup_array)
258
- avg_warm_compilation = np.mean(warm_compilation_array)
259
-
260
- percentages = [10, 25, 50, 75, 90, 99]
261
- cold_startup_percentiles = np.percentile(cold_startup_array, percentages)
262
- cold_compilation_percentiles = np.percentile(cold_compilation_array, percentages)
263
- warm_startup_percentiles = np.percentile(warm_startup_array, percentages)
264
- warm_compilation_percentiles = np.percentile(warm_compilation_array, percentages)
279
+ warm_iterations.append(create_llm_and_measure_startup())
280
281
+ # Determine if encoder compilation occurred in any iteration
282
+ has_encoder = any(
283
+ m["encoder_compilation_time"] > 0 for m in cold_iterations + warm_iterations
284
+ )
285
+
286
+ cold_metrics = _collect_phase_metrics("cold", cold_iterations, has_encoder)
287
+ warm_metrics = _collect_phase_metrics("warm", warm_iterations, has_encoder)
288
+ all_metrics = cold_metrics + warm_metrics
289
+
290
+ # Print results
291
print("\n" + "=" * 60)
292
print("STARTUP TIME BENCHMARK RESULTS")
293
print("=" * 60)
269
-
270
- # Cold startup statistics
271
- print("\nCOLD STARTUP:")
272
- print(f"Avg total startup time: {avg_cold_startup:.2f} seconds")
273
- print(f"Avg compilation time: {avg_cold_compilation:.2f} seconds")
274
- print("Startup time percentiles:")
275
- for percentage, percentile in zip(percentages, cold_startup_percentiles):
276
- print(f" {percentage}%: {percentile:.2f} seconds")
277
- print("Compilation time percentiles:")
278
- for percentage, percentile in zip(percentages, cold_compilation_percentiles):
279
- print(f" {percentage}%: {percentile:.2f} seconds")
280
-
281
- # Warm startup statistics
282
- print("\nWARM STARTUP:")
283
- print(f"Avg total startup time: {avg_warm_startup:.2f} seconds")
284
- print(f"Avg compilation time: {avg_warm_compilation:.2f} seconds")
285
- print("Startup time percentiles:")
286
- for percentage, percentile in zip(percentages, warm_startup_percentiles):
287
- print(f" {percentage}%: {percentile:.2f} seconds")
288
- print("Compilation time percentiles:")
289
- for percentage, percentile in zip(percentages, warm_compilation_percentiles):
290
- print(f" {percentage}%: {percentile:.2f} seconds")
291
-
294
+ _print_phase("COLD STARTUP", cold_metrics)
295
+ _print_phase("WARM STARTUP", warm_metrics)
296
print("=" * 60)
297
298
# Output JSON results if specified
299
if args.output_json:
296
- results = {
297
- "avg_cold_startup_time": float(avg_cold_startup),
298
- "avg_cold_compilation_time": float(avg_cold_compilation),
299
- "cold_startup_times": cold_startup_times,
300
- "cold_compilation_times": cold_compilation_times,
301
- "cold_startup_percentiles": dict(
302
- zip(percentages, cold_startup_percentiles.tolist())
303
- ),
304
- "cold_compilation_percentiles": dict(
305
- zip(percentages, cold_compilation_percentiles.tolist())
306
- ),
307
- "avg_warm_startup_time": float(avg_warm_startup),
308
- "avg_warm_compilation_time": float(avg_warm_compilation),
309
- "warm_startup_times": warm_startup_times,
310
- "warm_compilation_times": warm_compilation_times,
311
- "warm_startup_percentiles": dict(
312
- zip(percentages, warm_startup_percentiles.tolist())
313
- ),
314
- "warm_compilation_percentiles": dict(
315
- zip(percentages, warm_compilation_percentiles.tolist())
316
- ),
317
- }
300
+ results: dict[str, Any] = {}
301
+ for m in all_metrics:
302
+ results.update(_metric_to_json(m))
303
with open(args.output_json, "w") as f:
304
json.dump(results, f, indent=4)
320
- save_to_pytorch_benchmark_format(args, results)
305
+ save_to_pytorch_benchmark_format(args, all_metrics)
vllm/compilation/backends.py
+13
-3
@@ -265,6 +265,7 @@ class CompilerManager:
265
compile_range: Range,
266
graph_index: int = 0,
267
num_graphs: int = 1,
268
+ is_encoder: bool = False,
269
) -> Any:
270
if graph_index == 0:
271
# before compiling the first graph, record the start time
@@ -282,7 +283,10 @@ class CompilerManager:
283
# after loading the last graph for this shape, record the time.
284
# there can be multiple graphs due to piecewise compilation.
285
elapsed = time.perf_counter() - compilation_start_time
285
- compilation_config.compilation_time += elapsed
286
+ if is_encoder:
287
+ compilation_config.encoder_compilation_time += elapsed
288
+ else:
289
+ compilation_config.compilation_time += elapsed
290
logger.info_once(
291
"Directly load the compiled graph(s) for compile range %s "
292
"from the cache, took %.3f s",
@@ -387,7 +391,10 @@ class CompilerManager:
391
# after compiling the last graph, record the end time
392
if graph_index == num_graphs - 1:
393
elapsed = time.perf_counter() - compilation_start_time
390
- compilation_config.compilation_time += elapsed
394
+ if is_encoder:
395
+ compilation_config.encoder_compilation_time += elapsed
396
+ else:
397
+ compilation_config.compilation_time += elapsed
398
logger.info_once(
399
"Compiling a graph for compile range %s takes %.2f s",
400
str(compile_range),
@@ -1130,7 +1137,10 @@ class VllmBackend:
1137
logger.info_once(
1138
"Dynamo bytecode transform time: %.2f s", dynamo_time, scope="local"
1139
)
1133
- self.compilation_config.compilation_time += dynamo_time
1140
+ if self.is_encoder:
1141
+ self.compilation_config.encoder_compilation_time += dynamo_time
1142
+ else:
1143
+ self.compilation_config.compilation_time += dynamo_time
1144
1145
# Record Dynamo time in tracing if available
1146
start_time = int(torch_compile_start_time * 1e9)
vllm/compilation/piecewise_backend.py
+1
@@ -270,6 +270,7 @@ class PiecewiseBackend:
270
compile_range=range_entry.compile_range,
271
graph_index=self.piecewise_compile_index,
272
num_graphs=self.total_piecewise_compiles,
273
+ is_encoder=self.vllm_backend.is_encoder,
274
)
275
276
range_entry.compiled = True
vllm/config/compilation.py
+4
@@ -710,6 +710,8 @@ class CompilationConfig:
710
"""files that are traced for compilation"""
711
compilation_time: float = field(default=0.0, init=False)
712
"""time taken for compilation"""
713
+ encoder_compilation_time: float = field(default=0.0, init=False)
714
+ """time taken for multimodal encoder compilation"""
715
716
static_forward_context: dict[str, Any] = field(default_factory=dict, init=False)
717
"""Per-model forward context
@@ -756,6 +758,7 @@ class CompilationConfig:
758
"local_cache_dir",
759
"traced_files",
760
"compilation_time",
761
+ "encoder_compilation_time",
762
"static_forward_context",
763
"pass_config", # handled separately below
764
"dynamic_shapes_config", # handled separately below
@@ -775,6 +778,7 @@ class CompilationConfig:
778
"enabled_custom_ops": True,
779
"disabled_custom_ops": True,
780
"compilation_time": True,
781
+ "encoder_compilation_time": True,
782
"traced_files": True,
783
"inductor_compile_config": {
784
"post_grad_custom_post_pass": True,
vllm/v1/engine/core.py
+27
-5
@@ -282,11 +282,33 @@ class EngineCore:
282
self.model_executor.initialize_from_config(kv_cache_configs)
283
284
elapsed = time.time() - start
285
- logger.info_once(
286
- "init engine (profile, create kv cache, warmup model) took %.2f seconds",
287
- elapsed,
288
- scope="local",
289
- )
285
+ compile_time = vllm_config.compilation_config.compilation_time
286
+ encoder_compile_time = vllm_config.compilation_config.encoder_compilation_time
287
+ if encoder_compile_time > 0:
288
+ logger.info_once(
289
+ "init engine (profile, create kv cache, warmup model) took "
290
+ "%.2f s (compilation: %.2f s — language_model: %.2f s, "
291
+ "encoder: %.2f s)",
292
+ elapsed,
293
+ compile_time + encoder_compile_time,
294
+ compile_time,
295
+ encoder_compile_time,
296
+ scope="local",
297
+ )
298
+ elif compile_time > 0:
299
+ logger.info_once(
300
+ "init engine (profile, create kv cache, warmup model) took "
301
+ "%.2f s (compilation: %.2f s)",
302
+ elapsed,
303
+ compile_time,
304
+ scope="local",
305
+ )
306
+ else:
307
+ logger.info_once(
308
+ "init engine (profile, create kv cache, warmup model) took %.2f s",
309
+ elapsed,
310
+ scope="local",
311
+ )
312
return scheduler_kv_cache_config
313
314
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
vllm/v1/executor/abstract.py
+8
-3
@@ -22,7 +22,7 @@ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
22
from vllm.v1.engine import ReconfigureDistributedRequest
23
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
24
from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
25
-from vllm.v1.worker.worker_base import WorkerBase
25
+from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase
26
27
if TYPE_CHECKING:
28
from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase
@@ -121,14 +121,19 @@ class Executor(ABC):
121
underlying workers.
122
"""
123
self.collective_rpc("initialize_from_config", args=(kv_cache_configs,))
124
- compilation_times: list[float] = self.collective_rpc("compile_or_warm_up_model")
124
+ compilation_times: list[CompilationTimes] = self.collective_rpc(
125
+ "compile_or_warm_up_model"
126
+ )
127
# Propagate compilation time from workers back to the main process.
128
# With TP>1, compilation happens in worker processes, so the main
129
# process config is never updated. Use max across workers since they
130
# compile in parallel.
131
if compilation_times:
132
self.vllm_config.compilation_config.compilation_time = max(
131
- compilation_times
133
+ t.language_model for t in compilation_times
134
+ )
135
+ self.vllm_config.compilation_config.encoder_compilation_time = max(
136
+ t.encoder for t in compilation_times
137
)
138
139
def register_failure_callback(self, callback: FailureCallback): # noqa: B027
vllm/v1/worker/cpu_worker.py
+6
-2
@@ -13,6 +13,7 @@ from vllm.profiler.wrapper import TorchProfilerWrapper
13
from vllm.utils.torch_utils import set_random_seed
14
from vllm.v1.worker.cpu_model_runner import CPUModelRunner
15
from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
16
+from vllm.v1.worker.worker_base import CompilationTimes
17
18
logger = init_logger(__name__)
19
@@ -104,12 +105,15 @@ class CPUWorker(Worker):
105
def determine_available_memory(self) -> int:
106
return self.cache_config.cpu_kvcache_space_bytes or 0
107
107
- def compile_or_warm_up_model(self) -> float:
108
+ def compile_or_warm_up_model(self) -> CompilationTimes:
109
# Reset the seed to ensure that the random state is not affected by
110
# the model initialization and profiling.
111
set_random_seed(self.model_config.seed)
112
self.model_runner.warming_up_model()
112
- return self.compilation_config.compilation_time
113
+ return CompilationTimes(
114
+ language_model=self.compilation_config.compilation_time,
115
+ encoder=self.compilation_config.encoder_compilation_time,
116
+ )
117
118
def profile(self, is_start: bool = True, profile_prefix: str | None = None):
119
if self.profiler is None:
vllm/v1/worker/gpu_worker.py
+6
-3
@@ -56,7 +56,7 @@ from vllm.v1.outputs import (
56
)
57
from vllm.v1.utils import compute_iteration_details, report_usage_stats
58
from vllm.v1.worker.utils import is_residual_scattered_for_sp
59
-from vllm.v1.worker.worker_base import WorkerBase
59
+from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase
60
from vllm.v1.worker.workspace import init_workspace_manager
61
62
from ...model_executor.model_loader import TensorizerLoader
@@ -547,7 +547,7 @@ class Worker(WorkerBase):
547
self.model_runner._init_kv_zero_meta()
548
549
@instrument(span_name="Warmup (GPU)")
550
- def compile_or_warm_up_model(self) -> float:
550
+ def compile_or_warm_up_model(self) -> CompilationTimes:
551
warmup_sizes: list[int] = []
552
553
if self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE:
@@ -689,7 +689,10 @@ class Worker(WorkerBase):
689
# the model initialization and profiling.
690
set_random_seed(self.model_config.seed)
691
692
- return self.compilation_config.compilation_time
692
+ return CompilationTimes(
693
+ language_model=self.compilation_config.compilation_time,
694
+ encoder=self.compilation_config.encoder_compilation_time,
695
+ )
696
697
def reset_mm_cache(self) -> None:
698
self.model_runner.reset_mm_cache()
vllm/v1/worker/worker_base.py
+8
-3
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Callable
5
-from typing import TYPE_CHECKING, Any, TypeVar
5
+from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar
6
7
import torch
8
import torch.nn as nn
@@ -30,6 +30,11 @@ logger = init_logger(__name__)
30
_R = TypeVar("_R")
31
32
33
+class CompilationTimes(NamedTuple):
34
+ language_model: float
35
+ encoder: float
36
+
37
+
38
class WorkerBase:
39
"""Worker interface that allows vLLM to cleanly separate implementations for
40
different hardware. Also abstracts control plane communication, e.g., to
@@ -86,11 +91,11 @@ class WorkerBase:
91
"""Get specifications for KV cache implementation."""
92
raise NotImplementedError
93
89
- def compile_or_warm_up_model(self) -> float:
94
+ def compile_or_warm_up_model(self) -> CompilationTimes:
95
"""Prepare model for execution through compilation/warmup.
96
97
Returns:
93
- The accumulated compilation time in seconds.
98
+ Compilation times (language_model, encoder) in seconds.
99
"""
100
raise NotImplementedError
101