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 7, 02:45 UTC
[compile] Split compile/warmup monitoring (#36098)
Richard Zou committed
Mar 7, 2026 at 16:52 UTC
2dde535df1b736315e56eace0fa1923fe0beffc5
3 files changed
+125
-50
vllm/compilation/caching.py
+25
-1
@@ -189,13 +189,13 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
189
self.shape_env = None
190
self.vllm_backend = vllm_backend
191
self.sym_tensor_indices = sym_tensor_indices
192
+ self._fake_mode: Any | None = None
193
194
import torch._functorch.config as functorch_config
195
196
self.aot_autograd_config = (
197
aot_autograd_config or functorch_config.save_config_portable()
198
)
198
-
199
sym_input = next(
200
(i for i in self.example_inputs if isinstance(i, torch.SymInt)), None
201
)
@@ -217,6 +217,7 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
217
state.pop("optimized_call")
218
state.pop("shape_env")
219
state.pop("vllm_backend", None)
220
+ state.pop("_fake_mode", None)
221
for node in state["graph_module"].graph.nodes:
222
node.meta.pop("source_fn_stack", None)
223
node.meta.pop("nn_module_stack", None)
@@ -351,8 +352,31 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
352
return fn.optimized_call(*example_inputs)
353
354
fn = cls(**state, optimized_call=optimized_call)
355
+ fn._fake_mode = fake_mode
356
return fn
357
358
+ def finalize_loading(self, vllm_config: VllmConfig) -> None:
359
+ """Eagerly initialize the compiled backend and perform all loading.
360
+
361
+ Must be called after _verify_source_unchanged has populated
362
+ compilation_config.traced_files, which is needed for cache dir
363
+ computation.
364
+ """
365
+ if self._fake_mode is None:
366
+ return # Already finalized, or mega path (no _fake_mode set)
367
+
368
+ from torch._guards import TracingContext, tracing
369
+
370
+ from vllm.compilation.backends import VllmBackend
371
+
372
+ vllm_backend = VllmBackend(vllm_config, self.prefix, self.is_encoder)
373
+ with tracing(TracingContext(self._fake_mode)):
374
+ result = vllm_backend(self.graph_module, list(self.example_inputs))
375
+ self.optimized_call = result.optimized_call
376
+ self.vllm_backend = vllm_backend
377
+
378
+ self._fake_mode = None
379
+
380
@property
381
def co_name(self) -> Literal["VllmSerializableFunction"]:
382
"""
vllm/compilation/decorators.py
+41
-27
@@ -30,7 +30,7 @@ from vllm.sequence import IntermediateTensors
30
from vllm.utils.import_utils import resolve_obj_by_qualname
31
from vllm.utils.torch_utils import is_torch_equal_or_newer
32
33
-from .monitor import start_monitoring_torch_compile
33
+from .monitor import monitor_profiling_run, monitor_torch_compile
34
35
if TYPE_CHECKING:
36
# Only added on nightly/2.10 so wrap
@@ -434,17 +434,24 @@ def _support_torch_compile(
434
cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")
435
aot_compilation_path = os.path.join(cache_dir, "model")
436
try:
437
- with (
438
- set_current_vllm_config(self.vllm_config),
439
- open(aot_compilation_path, "rb") as f,
440
- ):
441
- start_monitoring_torch_compile(self.vllm_config)
442
- loaded_fn = torch.compiler.load_compiled_function(
443
- f, f_globals=self.forward.__globals__
444
- )
445
- _verify_source_unchanged(loaded_fn.source_info(), self.vllm_config)
446
- if not self.compilation_config.dynamic_shapes_config.evaluate_guards:
447
- loaded_fn.disable_guard_check()
437
+ with monitor_torch_compile(self.vllm_config):
438
+ with (
439
+ set_current_vllm_config(self.vllm_config),
440
+ open(aot_compilation_path, "rb") as f,
441
+ ):
442
+ loaded_fn = torch.compiler.load_compiled_function(
443
+ f, f_globals=self.forward.__globals__
444
+ )
445
+ _verify_source_unchanged(loaded_fn.source_info(), self.vllm_config)
446
+ ds_config = self.compilation_config.dynamic_shapes_config
447
+ if not ds_config.evaluate_guards:
448
+ loaded_fn.disable_guard_check()
449
+ # Eagerly load compiled artifacts now that traced_files
450
+ # is populated by _verify_source_unchanged.
451
+ with maybe_use_cudagraph_partition_wrapper(self.vllm_config):
452
+ loaded_fn._artifacts.compiled_fn.finalize_loading(
453
+ self.vllm_config
454
+ )
455
self.aot_compiled_fn = loaded_fn
456
self.was_aot_compile_fn_loaded_from_disk = True
457
except Exception as e:
@@ -465,12 +472,11 @@ def _support_torch_compile(
472
logger.info(
473
"Directly load AOT compilation from path %s", aot_compilation_path
474
)
468
- # Apply partition wrapper context for proper CUDA graph capture
469
- from .monitor import end_monitoring_torch_compile
470
-
471
- with maybe_use_cudagraph_partition_wrapper(self.vllm_config):
475
+ with (
476
+ monitor_profiling_run(),
477
+ maybe_use_cudagraph_partition_wrapper(self.vllm_config),
478
+ ):
479
output = self.aot_compiled_fn(self, *args, **kwargs)
473
- end_monitoring_torch_compile(self.vllm_config)
480
return output
481
482
if self.compiled:
@@ -489,8 +495,6 @@ def _support_torch_compile(
495
**kwargs,
496
)
497
492
- # here, it is the starting point of the `torch.compile` process
493
- start_monitoring_torch_compile(self.vllm_config)
498
original_code_object = self.original_code_object()
499
logger.debug("Start compiling function %s", original_code_object)
500
@@ -559,16 +563,26 @@ def _support_torch_compile(
563
# store the path for saving after warmup
564
self._aot_compilation_path = aot_compilation_path
565
self._aot_cache_dir = cache_dir
562
- self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
563
- # All compilation is done at this point, save the AOT artifact.
564
- self.save_aot_compiled_function()
565
- output = self.aot_compiled_fn(self, *args, **kwargs)
566
- else:
567
- output = TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs) # type: ignore[arg-type]
566
+ with monitor_torch_compile(self.vllm_config):
567
+ self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
568
+ # All compilation is done at this point, save the
569
+ # AOT artifact.
570
+ self.save_aot_compiled_function()
571
569
- from .monitor import end_monitoring_torch_compile
572
+ with monitor_profiling_run():
573
+ output = self.aot_compiled_fn(self, *args, **kwargs)
574
+ else:
575
+ with monitor_torch_compile(
576
+ self.vllm_config,
577
+ "torch.compile and initial profiling/warmup "
578
+ "run together took %.2f s in total",
579
+ ):
580
+ output = TorchCompileWithNoGuardsWrapper.__call__(
581
+ self, # type: ignore[arg-type]
582
+ *args,
583
+ **kwargs,
584
+ )
585
571
- end_monitoring_torch_compile(self.vllm_config)
586
self.compiled = True
587
return output
588
vllm/compilation/monitor.py
+59
-22
@@ -1,46 +1,83 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import contextlib
5
import time
6
+from collections.abc import Generator
7
6
-from vllm.config import CompilationConfig, CompilationMode, VllmConfig
8
+from vllm.config import CompilationMode, VllmConfig
9
from vllm.logger import init_logger
10
11
logger = init_logger(__name__)
12
11
-context_manager = None
13
+# Shared global so backends.py can read the start time for Dynamo timing.
14
torch_compile_start_time: float = 0.0
15
16
15
-def start_monitoring_torch_compile(vllm_config: VllmConfig) -> None:
17
+@contextlib.contextmanager
18
+def monitor_torch_compile(
19
+ vllm_config: VllmConfig,
20
+ message: str = "torch.compile took %.2f s in total",
21
+) -> Generator[None, None, None]:
22
+ """Context manager that times torch.compile and manages depyf debugging.
23
+
24
+ On normal exit: logs the compile time and exits depyf.
25
+ On exception: cleans up depyf without logging (compilation failed).
26
+ """
27
global torch_compile_start_time
28
torch_compile_start_time = time.perf_counter()
29
19
- compilation_config: CompilationConfig = vllm_config.compilation_config
30
+ compilation_config = vllm_config.compilation_config
31
+ depyf_cm = None
32
path = vllm_config.compile_debug_dump_path()
33
if compilation_config.mode == CompilationMode.VLLM_COMPILE and path:
34
import depyf
35
36
path.mkdir(parents=True, exist_ok=True)
37
logger.debug("Dumping depyf output to %s", path)
26
- global context_manager
27
- context_manager = depyf.prepare_debug(path.as_posix())
28
- context_manager.__enter__()
29
-
30
-
31
-def end_monitoring_torch_compile(vllm_config: VllmConfig) -> None:
32
- compilation_config: CompilationConfig = vllm_config.compilation_config
33
- total_compile_time: float = time.perf_counter() - torch_compile_start_time
34
- if compilation_config.mode == CompilationMode.VLLM_COMPILE:
35
- logger.info_once(
36
- "torch.compile and initial profiling run took %.2f s in total",
37
- total_compile_time,
38
- scope="local",
39
- )
40
- global context_manager
41
- if context_manager is not None:
42
- context_manager.__exit__(None, None, None)
43
- context_manager = None
38
+ depyf_cm = depyf.prepare_debug(path.as_posix())
39
+ depyf_cm.__enter__()
40
+
41
+ try:
42
+ yield
43
+ except Exception:
44
+ raise
45
+ else:
46
+ total_compile_time = time.perf_counter() - torch_compile_start_time
47
+ if compilation_config.mode == CompilationMode.VLLM_COMPILE:
48
+ logger.info_once(message, total_compile_time, scope="local")
49
+ finally:
50
+ if depyf_cm is not None:
51
+ try:
52
+ depyf_cm.__exit__(None, None, None)
53
+ except Exception:
54
+ logger.warning("Exception during depyf cleanup.", exc_info=True)
55
+
56
+
57
+@contextlib.contextmanager
58
+def monitor_profiling_run() -> Generator[None, None, None]:
59
+ """Context manager that times the initial profiling run.
60
+
61
+ Asserts that no backend compilation occurs during the profiling run
62
+ (all compilation should have completed before this point).
63
+ """
64
+ from vllm.compilation.counter import compilation_counter
65
+
66
+ backend_compilations_before = compilation_counter.num_backend_compilations
67
+ start = time.perf_counter()
68
+ yield
69
+ elapsed = time.perf_counter() - start
70
+ assert (
71
+ compilation_counter.num_backend_compilations == backend_compilations_before
72
+ ), (
73
+ "backend compilation occurred during the initial profiling run; "
74
+ "all compilation should be complete before the profiling run starts."
75
+ )
76
+ logger.info_once(
77
+ "Initial profiling/warmup run took %.2f s",
78
+ elapsed,
79
+ scope="local",
80
+ )
81
82
83
cudagraph_capturing_enabled: bool = True