iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[UX] Tag torch.compile log lines with the component being compiled (#48133)

Browse Source 
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>
This commit is contained in:
Francesco Fusco
2026-09-30 10:34:31 +00:00
committed by GitHub
co-authored by Claude Luka Govedič
parent d71f662606
commit d5ee309314
3 changed files with 79 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/compilation/backends.py
+25 -12
View File
@@ -142,6 +142,10 @@ class CompilerManager:
self.compilation_config = compilation_config
self.compiler = make_compiler(compilation_config)
self.loaded_artifacts: dict[str, Any] = {}
self.prefix: str = ""
def _log_prefix(self) -> str:
return f"[{self.prefix}] " if self.prefix else ""
def compute_hash(self, vllm_config: VllmConfig) -> str:
return self.compiler.compute_hash(vllm_config)
@@ -179,6 +183,7 @@ class CompilerManager:
self.disable_cache = disable_cache
self.cache_dir = cache_dir
self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py")
self.prefix = prefix
if not disable_cache and os.path.exists(self.cache_file_path):
# load the cache from the file
@@ -287,9 +292,10 @@ class CompilerManager:
# after loading the last graph for this shape, record the time.
# there can be multiple graphs due to piecewise compilation.
elapsed = time.perf_counter() - compilation_start_time
logger.info_once(
"Directly load the compiled graph(s) for compile range %s "
"from the cache, took %.3f s",
logger.info(
"%sDirectly load the compiled graph(s) for compile "
"range %s from the cache, took %.3f s",
self._log_prefix(),
str(compile_range),
elapsed,
)
@@ -373,8 +379,9 @@ class CompilerManager:
self.is_cache_updated = True
if graph_index == 0:
# adds some info logging for the first graph
logger.info_once(
"Cache the graph of compile range %s for later use",
logger.info(
"%sCache the graph of compile range %s for later use",
self._log_prefix(),
str(compile_range),
)
logger.debug_once(
@@ -388,8 +395,9 @@ class CompilerManager:
# after compiling the last graph, record the end time
if graph_index == num_graphs - 1:
elapsed = time.perf_counter() - compilation_start_time
logger.info_once(
"Compiling a graph for compile range %s takes %.2f s",
logger.info(
"%sCompiling a graph for compile range %s takes %.2f s",
self._log_prefix(),
str(compile_range),
elapsed,
)
@@ -865,6 +873,9 @@ class VllmBackend:
# `torch.compile` is JIT compiled, so we don't need to
# do anything here
def _log_prefix(self) -> str:
return f"[{self.prefix}] " if self.prefix else ""
def collect_standalone_compile_artifacts(
self,
) -> tuple[Any, dict[str, list[int]] | None, dict[str, bool] | None]:
@@ -1085,10 +1096,11 @@ class VllmBackend:
disable_cache = disable_cache or is_ngram_gpu_enabled
if disable_cache:
logger.info_once("vLLM's torch.compile cache is disabled.")
logger.info("%svLLM's torch.compile cache is disabled.", self._log_prefix())
else:
logger.info_once(
"Using cache directory: %s for vLLM's torch.compile",
logger.info(
"%sUsing cache directory: %s for vLLM's torch.compile",
self._log_prefix(),
local_cache_dir,
)
@@ -1148,8 +1160,9 @@ class VllmBackend:
current_perf = time.perf_counter()
current_epoch = time.time()
dynamo_time = current_perf - torch_compile_start_time
logger.info_once(
"Dynamo bytecode transform time: %.2f s",
logger.info(
"%sDynamo bytecode transform time: %.2f s",
self._log_prefix(),
dynamo_time,
)
vllm/compilation/decorators.py
+38 -10
View File
@@ -292,7 +292,12 @@ def _try_load_aot_compiled_fn(
Re-raises on failure when ``VLLM_FORCE_AOT_LOAD`` is set.
"""
try:
with monitor_torch_compile(model.vllm_config, is_encoder=model._is_encoder):
with monitor_torch_compile(
model.vllm_config,
is_encoder=model._is_encoder,
tag=getattr(model, "_compile_tag", ""),
announce_start=False,
):
with (
set_current_vllm_config(model.vllm_config),
open(aot_compilation_path, "rb") as f,
@@ -309,8 +314,12 @@ def _try_load_aot_compiled_fn(
with maybe_use_cudagraph_partition_wrapper(model.vllm_config):
loaded_fn._artifacts.compiled_fn.finalize_loading(model.vllm_config)
compilation_counter.num_aot_artifacts_loaded += 1
tag = getattr(model, "_compile_tag", "")
log_prefix = f"[{tag}] " if tag else ""
logger.info(
"Directly load AOT compilation from path %s", aot_compilation_path
"%sDirectly load AOT compilation from path %s",
log_prefix,
aot_compilation_path,
)
return loaded_fn
except Exception as e:
@@ -319,8 +328,11 @@ def _try_load_aot_compiled_fn(
message = "Compile cache file corrupted."
else:
message = str(e)
tag = getattr(model, "_compile_tag", "")
log_prefix = f"[{tag}] " if tag else ""
logger.warning(
"Compiling model again due to a load failure from %s, reason: %s",
"%sCompiling model again due to a load failure from %s, reason: %s",
log_prefix,
aot_compilation_path,
message,
)
@@ -402,6 +414,16 @@ def _support_torch_compile(
compilation_counter.num_models_seen += 1
self.compiled = False
# Capture the active model tag (e.g. "backbone", "eagle_head") so
# later compilation/warmup log lines can identify which component
# they belong to. Mirror VllmBackend's precedence rule
# (`compile_prefix or model_tag`) so log tags line up with the
# cache-dir tag: encoders use the class name, everything else
# falls through to `model_tag`.
from vllm.compilation.backends import model_tag as _current_model_tag
self._compile_tag = (cls.__name__ if is_encoder else "") or _current_model_tag
# Handled by monkeypatching `TorchCompileWithNoGuardsWrapper` into base class
TorchCompileWithNoGuardsWrapper.__init__(
self,
@@ -567,7 +589,7 @@ def _support_torch_compile(
self.aot_compiled_fn = loaded_fn
self.was_aot_compile_fn_loaded_from_disk = True
with (
monitor_profiling_run(),
monitor_profiling_run(tag=self._compile_tag),
maybe_use_cudagraph_partition_wrapper(self.vllm_config),
):
output = self.aot_compiled_fn(self, *args, **kwargs)
@@ -667,7 +689,9 @@ def _support_torch_compile(
self._aot_compilation_path = aot_compilation_path
self._aot_cache_dir = cache_dir
with monitor_torch_compile(
self.vllm_config, is_encoder=self._is_encoder
self.vllm_config,
is_encoder=self._is_encoder,
tag=self._compile_tag,
):
self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
compilation_counter.num_aot_compiles += 1
@@ -675,14 +699,15 @@ def _support_torch_compile(
# AOT artifact.
self.save_aot_compiled_function()
with monitor_profiling_run():
with monitor_profiling_run(tag=self._compile_tag):
output = self.aot_compiled_fn(self, *args, **kwargs)
else:
with monitor_torch_compile(
self.vllm_config,
"torch.compile and initial profiling/warmup "
"%storch.compile and initial profiling/warmup "
"run together took %.2f s in total",
is_encoder=self._is_encoder,
tag=self._compile_tag,
):
output = TorchCompileWithNoGuardsWrapper.__call__(
self, # type: ignore[arg-type]
@@ -706,6 +731,7 @@ def _support_torch_compile(
self.aot_compiled_fn and self._aot_compilation_path and self._aot_cache_dir
)
log_prefix = f"[{self._compile_tag}] " if self._compile_tag else ""
try:
os.makedirs(self._aot_cache_dir, exist_ok=True)
# File saving should be atomic, so we will save to a temporary location
@@ -714,13 +740,15 @@ def _support_torch_compile(
self.aot_compiled_fn.save_compiled_function(tmp_file)
os.replace(tmp_file, self._aot_compilation_path)
compilation_counter.num_aot_artifacts_saved += 1
logger.info_once(
"saved AOT compiled function to %s",
logger.info(
"%ssaved AOT compiled function to %s",
log_prefix,
self._aot_compilation_path,
)
except Exception as e:
logger.warning(
"unable to save AOT compiled function to %s: %s",
"%sunable to save AOT compiled function to %s: %s",
log_prefix,
self._aot_compilation_path,
e,
)
vllm/compilation/monitor.py
+16 -5
View File
@@ -17,18 +17,27 @@ torch_compile_start_time: float = 0.0
@contextlib.contextmanager
def monitor_torch_compile(
vllm_config: VllmConfig,
message: str = "torch.compile took %.2f s in total",
message: str = "%storch.compile took %.2f s in total",
is_encoder: bool = False,
tag: str = "",
announce_start: bool = True,
) -> Generator[None, None, None]:
"""Context manager that times torch.compile and manages depyf debugging.
On normal exit: logs the compile time and exits depyf.
On exception: cleans up depyf without logging (compilation failed).
``announce_start=False`` suppresses the "Starting torch.compile" line;
used by the AOT-load probe so a failed load attempt doesn't announce a
compile it didn't perform.
"""
global torch_compile_start_time
torch_compile_start_time = time.perf_counter()
compilation_config = vllm_config.compilation_config
log_prefix = f"[{tag}] " if tag else ""
if announce_start and compilation_config.mode == CompilationMode.VLLM_COMPILE:
logger.info("%sStarting torch.compile", log_prefix)
depyf_cm = None
path = vllm_config.compile_debug_dump_path()
if compilation_config.mode == CompilationMode.VLLM_COMPILE and path:
@@ -50,7 +59,7 @@ def monitor_torch_compile(
compilation_config.encoder_compilation_time += total_compile_time
else:
compilation_config.compilation_time += total_compile_time
logger.info_once(message, total_compile_time)
logger.info(message, log_prefix, total_compile_time)
finally:
if depyf_cm is not None:
try:
@@ -60,7 +69,7 @@ def monitor_torch_compile(
@contextlib.contextmanager
def monitor_profiling_run() -> Generator[None, None, None]:
def monitor_profiling_run(tag: str = "") -> Generator[None, None, None]:
"""Context manager that times the initial profiling run.
Asserts that no backend compilation occurs during the profiling run
@@ -78,8 +87,10 @@ def monitor_profiling_run() -> Generator[None, None, None]:
"backend compilation occurred during the initial profiling run; "
"all compilation should be complete before the profiling run starts."
)
logger.info_once(
"Initial profiling/warmup run took %.2f s",
log_prefix = f"[{tag}] " if tag else ""
logger.info(
"%sInitial profiling/warmup run took %.2f s",
log_prefix,
elapsed,
)