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 9, 00:00 UTC
[torch.compile] Significantly speed up cold start times (#33641)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Feb 3, 2026 at 09:12 UTC
b1bb18de8d12cac63e3bbb0f59b3726fcb68dc80
3 files changed
+41
-21
tests/compile/test_cold_start.py
+5
-4
@@ -37,12 +37,13 @@ def test_moe_compilation_cold_start(monkeypatch, use_fresh_inductor_cache):
37
# The forward pass consists of 32 transformer layers.
38
# Then, we split on the attention operation. This results in
39
# 33 subgraphs (not including the attention operation).
40
- # The 33 subgraphs then get standalone_compile'd.
40
+ # We then standalone_compile the unique subgraphs.
41
#
42
# There are actually only 3 unique subgraphs for this model
43
# (all of its transformer layers are the same modulo weights);
44
# this is true for most vLLM models.
45
- # So we test that during cold start, the aot_autograd cache
46
- # misses for 3 subgraphs and hits for the rest.
45
+ # So we test that during cold start, only 3 subgraphs are compiled
46
+ # These 3 subgraphs should cache miss, and then there should be
47
+ # no other compilation (so no cache hits).
48
assert counters["aot_autograd"]["autograd_cache_miss"] == 3
48
- assert counters["aot_autograd"]["autograd_cache_hit"] == 30
49
+ assert counters["aot_autograd"]["autograd_cache_hit"] == 0
vllm/compilation/backends.py
+36
-14
@@ -121,7 +121,7 @@ class CompilerManager:
121
and compiling the graph.
122
123
The cache is a dict mapping
124
- `(runtime_shape, graph_index, backend_name)`
124
+ `(runtime_shape, graph_hash, backend_name)`
125
to `any_data` returned from the compiler.
126
127
When serializing the cache, we save it to a Python file
@@ -130,7 +130,7 @@ class CompilerManager:
130
"""
131
132
def __init__(self, compilation_config: CompilationConfig) -> None:
133
- self.cache: dict[tuple[Range, int, str], Any] = dict()
133
+ self.cache: dict[tuple[Range, str, str], Any] = dict()
134
self.is_cache_updated = False
135
self.compilation_config = compilation_config
136
self.compiler = make_compiler(compilation_config)
@@ -173,6 +173,7 @@ class CompilerManager:
173
self.disable_cache = disable_cache
174
self.cache_dir = cache_dir
175
self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py")
176
+ self.loaded_cache_entries: dict[tuple[Range, str, str], Any] = {}
177
178
if not disable_cache and os.path.exists(self.cache_file_path):
179
# load the cache from the file
@@ -186,9 +187,9 @@ class CompilerManager:
187
if not isinstance(value, ty):
188
raise TypeError(f"Expected {ty} but got {type(value)} for {value}")
189
189
- def parse_key(key: Any) -> tuple[Range, int, str]:
190
- range_tuple, graph_index, compiler_name = key
191
- check_type(graph_index, int)
190
+ def parse_key(key: Any) -> tuple[Range, str, str]:
191
+ range_tuple, graph_hash, compiler_name = key
192
+ check_type(graph_hash, str)
193
check_type(compiler_name, str)
194
if isinstance(range_tuple, tuple):
195
start, end = range_tuple
@@ -196,7 +197,7 @@ class CompilerManager:
197
check_type(end, int)
198
range_tuple = Range(start=start, end=end)
199
check_type(range_tuple, Range)
199
- return range_tuple, graph_index, compiler_name
200
+ return range_tuple, graph_hash, compiler_name
201
202
self.cache = {parse_key(key): value for key, value in cache.items()}
203
@@ -216,18 +217,25 @@ class CompilerManager:
217
self,
218
graph: fx.GraphModule,
219
example_inputs: list[Any],
219
- graph_index: int,
220
+ graph_hash: str,
221
compile_range: Range,
222
) -> Callable[..., Any] | None:
222
- if (compile_range, graph_index, self.compiler.name) not in self.cache:
223
+ key = (compile_range, graph_hash, self.compiler.name)
224
+ # See if we've already loaded this cache entry
225
+ if key in self.loaded_cache_entries:
226
+ return self.loaded_cache_entries[key]
227
+ # Otherwise, go load it from disk
228
+ if key not in self.cache:
229
return None
224
- handle = self.cache[(compile_range, graph_index, self.compiler.name)]
230
+ handle = self.cache[key]
231
compiled_graph = self.compiler.load(
226
- handle, graph, example_inputs, graph_index, compile_range
232
+ handle, graph, example_inputs, compile_range
233
)
234
+ self.loaded_cache_entries[key] = compiled_graph
235
logger.debug(
229
- "Directly load the %s-th graph for compile range %sfrom %s via handle %s",
230
- graph_index,
236
+ "Directly load the graph (hash %s) for compile range "
237
+ "%sfrom %s via handle %s",
238
+ graph_hash,
239
str(compile_range),
240
self.compiler.name,
241
handle,
@@ -249,12 +257,22 @@ class CompilerManager:
257
global compilation_start_time
258
compilation_start_time = time.time()
259
260
+ from torch._functorch._aot_autograd.autograd_cache import (
261
+ AOTAutogradCachePickler,
262
+ sanitize_gm_for_cache,
263
+ )
264
+
265
+ with sanitize_gm_for_cache(graph):
266
+ pickler = AOTAutogradCachePickler(graph)
267
+ dumped_graph = pickler.dumps(graph)
268
+ graph_hash = hashlib.sha256(dumped_graph).hexdigest()
269
+
270
compilation_counter.num_backend_compilations += 1
271
272
compiled_graph = None
273
274
# try to load from the cache
257
- compiled_graph = self.load(graph, example_inputs, graph_index, compile_range)
275
+ compiled_graph = self.load(graph, example_inputs, graph_hash, compile_range)
276
if compiled_graph is not None:
277
if graph_index == num_graphs - 1:
278
# after loading the last graph for this shape, record the time.
@@ -290,9 +308,13 @@ class CompilerManager:
308
309
assert compiled_graph is not None, "Failed to compile the graph"
310
311
+ self.loaded_cache_entries[(compile_range, graph_hash, self.compiler.name)] = (
312
+ compiled_graph
313
+ )
314
+
315
# store the artifact in the cache
316
if is_compile_cache_enabled(additional_inductor_config) and handle is not None:
295
- self.cache[(compile_range, graph_index, self.compiler.name)] = handle
317
+ self.cache[(compile_range, graph_hash, self.compiler.name)] = handle
318
compilation_counter.num_cache_entries_updated += 1
319
self.is_cache_updated = True
320
if graph_index == 0:
vllm/compilation/compiler_interface.py
-3
@@ -101,7 +101,6 @@ class CompilerInterface:
101
handle: Any,
102
graph: fx.GraphModule,
103
example_inputs: list[Any],
104
- graph_index: int,
104
compile_range: Range,
105
) -> Callable[..., Any]:
106
"""
@@ -302,7 +301,6 @@ class InductorStandaloneAdaptor(CompilerInterface):
301
handle: Any,
302
graph: fx.GraphModule,
303
example_inputs: list[Any],
305
- graph_index: int,
304
compile_range: Range,
305
) -> Callable[..., Any]:
306
assert isinstance(handle, tuple)
@@ -527,7 +525,6 @@ class InductorAdaptor(CompilerInterface):
525
handle: Any,
526
graph: fx.GraphModule,
527
example_inputs: list[Any],
530
- graph_index: int,
528
compile_range: Range,
529
) -> Callable[..., Any]:
530
assert isinstance(handle, tuple)