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, 05:15 UTC
[compile] Fix extra cache save on warm start. (#35921)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Mar 4, 2026 at 23:56 UTC
dd6dbd93f8d299ee1e0fdbdd7cd0d41f47a4093f
2 files changed
+27
-6
tests/compile/test_startup.py
+4
-4
@@ -61,11 +61,11 @@ def test_moe_startup(monkeypatch, vllm_runner, fresh_vllm_cache):
61
counters.clear()
62
with compilation_counter.expect(
63
num_compiled_artifacts_loaded=3,
64
- # TODO: warm start should not save any artifacts
65
- # https://github.com/vllm-project/vllm/issues/35708
66
- num_compiled_artifacts_saved=1,
64
+ num_compiled_artifacts_saved=0,
65
):
66
_run_vllm(vllm_runner)
67
assert counters["aot_autograd"]["total"] == 30
68
assert counters["aot_autograd"]["autograd_cache_miss"] == 0
71
- assert counters["aot_autograd"]["autograd_cache_hit"] == 1
69
+ assert (
70
+ counters["aot_autograd"]["autograd_cache_hit"] == 0
71
+ ) # No miss at aot_autograd level causing disk I/O.
vllm/compilation/backends.py
+23
-2
@@ -221,10 +221,28 @@ class CompilerManager:
221
) -> Callable[..., Any] | None:
222
if (compile_range, graph_index, self.compiler.name) not in self.cache:
223
return None
224
- handle = self.cache[(compile_range, graph_index, self.compiler.name)]
224
+
225
+ def parse_value(value: Any) -> tuple[tuple[str, str], str]:
226
+ assert isinstance(value, dict)
227
+ handle = value["graph_handle"]
228
+ assert isinstance(handle[0], str)
229
+ assert isinstance(handle[1], str)
230
+ cache_key = value["cache_key"]
231
+ return handle, cache_key
232
+
233
+ try:
234
+ handle, cache_key = parse_value(
235
+ self.cache[(compile_range, graph_index, self.compiler.name)]
236
+ )
237
+ except Exception:
238
+ # When the cache is outdated, we should ignore the existing file.
239
+ # This should cause the correct cache to be generated again.
240
+ return None
241
+
242
compiled_graph = self.compiler.load(
243
handle, graph, example_inputs, graph_index, compile_range
244
)
245
+ self.loaded_artifacts[cache_key] = compiled_graph
246
logger.debug(
247
"Directly load the %s-th graph for compile range %sfrom %s via handle %s",
248
graph_index,
@@ -341,7 +359,10 @@ class CompilerManager:
359
360
# store the artifact in the cache
361
if is_compile_cache_enabled(additional_inductor_config) and handle is not None:
344
- self.cache[(compile_range, graph_index, self.compiler.name)] = handle
362
+ self.cache[(compile_range, graph_index, self.compiler.name)] = {
363
+ "graph_handle": handle,
364
+ "cache_key": cache_key,
365
+ }
366
compilation_counter.num_cache_entries_updated += 1
367
self.is_cache_updated = True
368
if graph_index == 0: