[UX] Tag torch.compile log lines with the component being compiled (#48133)
Signed-off-by: Francesco Fusco <ffu@zurich.ibm.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Francesco Fusco committed
Sep 30, 2026 at 12:34 UTC
d5ee309314c78a52a0e154cb15d53e34844d623f
3 files changed
+79
-27
vllm/compilation/backends.py
+25
-12
@@ -142,6 +142,10 @@ class CompilerManager:
142
self.compilation_config = compilation_config
143
self.compiler = make_compiler(compilation_config)
144
self.loaded_artifacts: dict[str, Any] = {}
145
+ self.prefix: str = ""
146
+
147
+ def _log_prefix(self) -> str:
148
+ return f"[{self.prefix}] " if self.prefix else ""
149
150
def compute_hash(self, vllm_config: VllmConfig) -> str:
151
return self.compiler.compute_hash(vllm_config)
@@ -179,6 +183,7 @@ class CompilerManager:
183
self.disable_cache = disable_cache
184
self.cache_dir = cache_dir
185
self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py")
186
+ self.prefix = prefix
187
188
if not disable_cache and os.path.exists(self.cache_file_path):
189
# load the cache from the file
@@ -287,9 +292,10 @@ class CompilerManager:
292
# after loading the last graph for this shape, record the time.
293
# there can be multiple graphs due to piecewise compilation.
294
elapsed = time.perf_counter() - compilation_start_time
290
- logger.info_once(
291
- "Directly load the compiled graph(s) for compile range %s "
292
- "from the cache, took %.3f s",
295
+ logger.info(
296
+ "%sDirectly load the compiled graph(s) for compile "
297
+ "range %s from the cache, took %.3f s",
298
+ self._log_prefix(),
299
str(compile_range),
300
elapsed,
301
)
@@ -373,8 +379,9 @@ class CompilerManager:
379
self.is_cache_updated = True
380
if graph_index == 0:
381
# adds some info logging for the first graph
376
- logger.info_once(
377
- "Cache the graph of compile range %s for later use",
382
+ logger.info(
383
+ "%sCache the graph of compile range %s for later use",
384
+ self._log_prefix(),
385
str(compile_range),
386
)
387
logger.debug_once(
@@ -388,8 +395,9 @@ class CompilerManager:
395
# after compiling the last graph, record the end time
396
if graph_index == num_graphs - 1:
397
elapsed = time.perf_counter() - compilation_start_time
391
- logger.info_once(
392
- "Compiling a graph for compile range %s takes %.2f s",
398
+ logger.info(
399
+ "%sCompiling a graph for compile range %s takes %.2f s",
400
+ self._log_prefix(),
401
str(compile_range),
402
elapsed,
403
)
@@ -865,6 +873,9 @@ class VllmBackend:
873
# `torch.compile` is JIT compiled, so we don't need to
874
# do anything here
875
876
+ def _log_prefix(self) -> str:
877
+ return f"[{self.prefix}] " if self.prefix else ""
878
+
879
def collect_standalone_compile_artifacts(
880
self,
881
) -> tuple[Any, dict[str, list[int]] | None, dict[str, bool] | None]:
@@ -1085,10 +1096,11 @@ class VllmBackend:
1096
disable_cache = disable_cache or is_ngram_gpu_enabled
1097
1098
if disable_cache:
1088
- logger.info_once("vLLM's torch.compile cache is disabled.")
1099
+ logger.info("%svLLM's torch.compile cache is disabled.", self._log_prefix())
1100
else:
1090
- logger.info_once(
1091
- "Using cache directory: %s for vLLM's torch.compile",
1101
+ logger.info(
1102
+ "%sUsing cache directory: %s for vLLM's torch.compile",
1103
+ self._log_prefix(),
1104
local_cache_dir,
1105
)
1106
@@ -1148,8 +1160,9 @@ class VllmBackend:
1160
current_perf = time.perf_counter()
1161
current_epoch = time.time()
1162
dynamo_time = current_perf - torch_compile_start_time
1151
- logger.info_once(
1152
- "Dynamo bytecode transform time: %.2f s",
1163
+ logger.info(
1164
+ "%sDynamo bytecode transform time: %.2f s",
1165
+ self._log_prefix(),
1166
dynamo_time,
1167
)
1168
vllm/compilation/decorators.py
+38
-10
@@ -292,7 +292,12 @@ def _try_load_aot_compiled_fn(
292
Re-raises on failure when ``VLLM_FORCE_AOT_LOAD`` is set.
293
"""
294
try:
295
- with monitor_torch_compile(model.vllm_config, is_encoder=model._is_encoder):
295
+ with monitor_torch_compile(
296
+ model.vllm_config,
297
+ is_encoder=model._is_encoder,
298
+ tag=getattr(model, "_compile_tag", ""),
299
+ announce_start=False,
300
+ ):
301
with (
302
set_current_vllm_config(model.vllm_config),
303
open(aot_compilation_path, "rb") as f,
@@ -309,8 +314,12 @@ def _try_load_aot_compiled_fn(
314
with maybe_use_cudagraph_partition_wrapper(model.vllm_config):
315
loaded_fn._artifacts.compiled_fn.finalize_loading(model.vllm_config)
316
compilation_counter.num_aot_artifacts_loaded += 1
317
+ tag = getattr(model, "_compile_tag", "")
318
+ log_prefix = f"[{tag}] " if tag else ""
319
logger.info(
313
- "Directly load AOT compilation from path %s", aot_compilation_path
320
+ "%sDirectly load AOT compilation from path %s",
321
+ log_prefix,
322
+ aot_compilation_path,
323
)
324
return loaded_fn
325
except Exception as e:
@@ -319,8 +328,11 @@ def _try_load_aot_compiled_fn(
328
message = "Compile cache file corrupted."
329
else:
330
message = str(e)
331
+ tag = getattr(model, "_compile_tag", "")
332
+ log_prefix = f"[{tag}] " if tag else ""
333
logger.warning(
323
- "Compiling model again due to a load failure from %s, reason: %s",
334
+ "%sCompiling model again due to a load failure from %s, reason: %s",
335
+ log_prefix,
336
aot_compilation_path,
337
message,
338
)
@@ -402,6 +414,16 @@ def _support_torch_compile(
414
compilation_counter.num_models_seen += 1
415
self.compiled = False
416
417
+ # Capture the active model tag (e.g. "backbone", "eagle_head") so
418
+ # later compilation/warmup log lines can identify which component
419
+ # they belong to. Mirror VllmBackend's precedence rule
420
+ # (`compile_prefix or model_tag`) so log tags line up with the
421
+ # cache-dir tag: encoders use the class name, everything else
422
+ # falls through to `model_tag`.
423
+ from vllm.compilation.backends import model_tag as _current_model_tag
424
+
425
+ self._compile_tag = (cls.__name__ if is_encoder else "") or _current_model_tag
426
+
427
# Handled by monkeypatching `TorchCompileWithNoGuardsWrapper` into base class
428
TorchCompileWithNoGuardsWrapper.__init__(
429
self,
@@ -567,7 +589,7 @@ def _support_torch_compile(
589
self.aot_compiled_fn = loaded_fn
590
self.was_aot_compile_fn_loaded_from_disk = True
591
with (
570
- monitor_profiling_run(),
592
+ monitor_profiling_run(tag=self._compile_tag),
593
maybe_use_cudagraph_partition_wrapper(self.vllm_config),
594
):
595
output = self.aot_compiled_fn(self, *args, **kwargs)
@@ -667,7 +689,9 @@ def _support_torch_compile(
689
self._aot_compilation_path = aot_compilation_path
690
self._aot_cache_dir = cache_dir
691
with monitor_torch_compile(
670
- self.vllm_config, is_encoder=self._is_encoder
692
+ self.vllm_config,
693
+ is_encoder=self._is_encoder,
694
+ tag=self._compile_tag,
695
):
696
self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
697
compilation_counter.num_aot_compiles += 1
@@ -675,14 +699,15 @@ def _support_torch_compile(
699
# AOT artifact.
700
self.save_aot_compiled_function()
701
678
- with monitor_profiling_run():
702
+ with monitor_profiling_run(tag=self._compile_tag):
703
output = self.aot_compiled_fn(self, *args, **kwargs)
704
else:
705
with monitor_torch_compile(
706
self.vllm_config,
683
- "torch.compile and initial profiling/warmup "
707
+ "%storch.compile and initial profiling/warmup "
708
"run together took %.2f s in total",
709
is_encoder=self._is_encoder,
710
+ tag=self._compile_tag,
711
):
712
output = TorchCompileWithNoGuardsWrapper.__call__(
713
self, # type: ignore[arg-type]
@@ -706,6 +731,7 @@ def _support_torch_compile(
731
self.aot_compiled_fn and self._aot_compilation_path and self._aot_cache_dir
732
)
733
734
+ log_prefix = f"[{self._compile_tag}] " if self._compile_tag else ""
735
try:
736
os.makedirs(self._aot_cache_dir, exist_ok=True)
737
# File saving should be atomic, so we will save to a temporary location
@@ -714,13 +740,15 @@ def _support_torch_compile(
740
self.aot_compiled_fn.save_compiled_function(tmp_file)
741
os.replace(tmp_file, self._aot_compilation_path)
742
compilation_counter.num_aot_artifacts_saved += 1
717
- logger.info_once(
718
- "saved AOT compiled function to %s",
743
+ logger.info(
744
+ "%ssaved AOT compiled function to %s",
745
+ log_prefix,
746
self._aot_compilation_path,
747
)
748
except Exception as e:
749
logger.warning(
723
- "unable to save AOT compiled function to %s: %s",
750
+ "%sunable to save AOT compiled function to %s: %s",
751
+ log_prefix,
752
self._aot_compilation_path,
753
e,
754
)
vllm/compilation/monitor.py
+16
-5
@@ -17,18 +17,27 @@ torch_compile_start_time: float = 0.0
17
@contextlib.contextmanager
18
def monitor_torch_compile(
19
vllm_config: VllmConfig,
20
- message: str = "torch.compile took %.2f s in total",
20
+ message: str = "%storch.compile took %.2f s in total",
21
is_encoder: bool = False,
22
+ tag: str = "",
23
+ announce_start: bool = True,
24
) -> Generator[None, None, None]:
25
"""Context manager that times torch.compile and manages depyf debugging.
26
27
On normal exit: logs the compile time and exits depyf.
28
On exception: cleans up depyf without logging (compilation failed).
29
+
30
+ ``announce_start=False`` suppresses the "Starting torch.compile" line;
31
+ used by the AOT-load probe so a failed load attempt doesn't announce a
32
+ compile it didn't perform.
33
"""
34
global torch_compile_start_time
35
torch_compile_start_time = time.perf_counter()
36
37
compilation_config = vllm_config.compilation_config
38
+ log_prefix = f"[{tag}] " if tag else ""
39
+ if announce_start and compilation_config.mode == CompilationMode.VLLM_COMPILE:
40
+ logger.info("%sStarting torch.compile", log_prefix)
41
depyf_cm = None
42
path = vllm_config.compile_debug_dump_path()
43
if compilation_config.mode == CompilationMode.VLLM_COMPILE and path:
@@ -50,7 +59,7 @@ def monitor_torch_compile(
59
compilation_config.encoder_compilation_time += total_compile_time
60
else:
61
compilation_config.compilation_time += total_compile_time
53
- logger.info_once(message, total_compile_time)
62
+ logger.info(message, log_prefix, total_compile_time)
63
finally:
64
if depyf_cm is not None:
65
try:
@@ -60,7 +69,7 @@ def monitor_torch_compile(
69
70
71
@contextlib.contextmanager
63
-def monitor_profiling_run() -> Generator[None, None, None]:
72
+def monitor_profiling_run(tag: str = "") -> Generator[None, None, None]:
73
"""Context manager that times the initial profiling run.
74
75
Asserts that no backend compilation occurs during the profiling run
@@ -78,8 +87,10 @@ def monitor_profiling_run() -> Generator[None, None, None]:
87
"backend compilation occurred during the initial profiling run; "
88
"all compilation should be complete before the profiling run starts."
89
)
81
- logger.info_once(
82
- "Initial profiling/warmup run took %.2f s",
90
+ log_prefix = f"[{tag}] " if tag else ""
91
+ logger.info(
92
+ "%sInitial profiling/warmup run took %.2f s",
93
+ log_prefix,
94
elapsed,
95
)
96