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, 14:00 UTC
[Bugfix] Propagate compilation_time from workers to main process for TP>1 (#35503)
Signed-off-by: Huy Do <huydhn@gmail.com>
Huy Do committed
Feb 27, 2026 at 21:03 UTC
7b346ba8ed546ad274a7905ccfd9f3c0c242226a
4 files changed
+20
-5
vllm/v1/executor/abstract.py
+9
-1
@@ -115,7 +115,15 @@ class Executor(ABC):
115
underlying workers.
116
"""
117
self.collective_rpc("initialize_from_config", args=(kv_cache_configs,))
118
- self.collective_rpc("compile_or_warm_up_model")
118
+ compilation_times: list[float] = self.collective_rpc("compile_or_warm_up_model")
119
+ # Propagate compilation time from workers back to the main process.
120
+ # With TP>1, compilation happens in worker processes, so the main
121
+ # process config is never updated. Use max across workers since they
122
+ # compile in parallel.
123
+ if compilation_times:
124
+ self.vllm_config.compilation_config.compilation_time = max(
125
+ compilation_times
126
+ )
127
128
def register_failure_callback(self, callback: FailureCallback): # noqa: B027
129
"""
vllm/v1/worker/cpu_worker.py
+2
-1
@@ -118,11 +118,12 @@ class CPUWorker(Worker):
118
def determine_available_memory(self) -> int:
119
return self.cache_config.cpu_kvcache_space_bytes or 0
120
121
- def compile_or_warm_up_model(self) -> None:
121
+ def compile_or_warm_up_model(self) -> float:
122
# Reset the seed to ensure that the random state is not affected by
123
# the model initialization and profiling.
124
set_random_seed(self.model_config.seed)
125
self.model_runner.warming_up_model()
126
+ return self.compilation_config.compilation_time
127
128
def _get_autobind_cpu_ids(
129
self, cpu_selector: Callable[[list[LogicalCPUInfo]], list[LogicalCPUInfo]]
vllm/v1/worker/gpu_worker.py
+3
-1
@@ -480,7 +480,7 @@ class Worker(WorkerBase):
480
self.model_runner.initialize_kv_cache(kv_cache_config)
481
482
@instrument(span_name="Warmup (GPU)")
483
- def compile_or_warm_up_model(self) -> None:
483
+ def compile_or_warm_up_model(self) -> float:
484
warmup_sizes = []
485
486
if self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE:
@@ -605,6 +605,8 @@ class Worker(WorkerBase):
605
# the model initialization and profiling.
606
set_random_seed(self.model_config.seed)
607
608
+ return self.compilation_config.compilation_time
609
+
610
def reset_mm_cache(self) -> None:
611
self.model_runner.reset_mm_cache()
612
vllm/v1/worker/worker_base.py
+6
-2
@@ -87,8 +87,12 @@ class WorkerBase:
87
"""Get specifications for KV cache implementation."""
88
raise NotImplementedError
89
90
- def compile_or_warm_up_model(self) -> None:
91
- """Prepare model for execution through compilation/warmup."""
90
+ def compile_or_warm_up_model(self) -> float:
91
+ """Prepare model for execution through compilation/warmup.
92
+
93
+ Returns:
94
+ The accumulated compilation time in seconds.
95
+ """
96
raise NotImplementedError
97
98
def check_health(self) -> None: