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 8, 22:45 UTC
Revert "[torch.compile] Significantly speed up cold start times" (#33820)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Feb 4, 2026 at 13:59 UTC
9f14c9224d3d6664e2f5a2e7fecd012fd048fcb1
3 files changed
+21
-41
tests/compile/test_cold_start.py
+4
-5
@@ -37,13 +37,12 @@ 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
- # We then standalone_compile the unique subgraphs.
40
+ # The 33 subgraphs then get standalone_compile'd.
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, 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).
45
+ # So we test that during cold start, the aot_autograd cache
46
+ # misses for 3 subgraphs and hits for the rest.
47
assert counters["aot_autograd"]["autograd_cache_miss"] == 3
49
- assert counters["aot_autograd"]["autograd_cache_hit"] == 0
48
+ assert counters["aot_autograd"]["autograd_cache_hit"] == 30
vllm/compilation/backends.py
+14
-36
@@ -121,7 +121,7 @@ class CompilerManager:
121
and compiling the graph.
122
123
The cache is a dict mapping
124
- `(runtime_shape, graph_hash, backend_name)`
124
+ `(runtime_shape, graph_index, 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, str, str], Any] = dict()
133
+ self.cache: dict[tuple[Range, int, str], Any] = dict()
134
self.is_cache_updated = False
135
self.compilation_config = compilation_config
136
self.compiler = make_compiler(compilation_config)
@@ -173,7 +173,6 @@ 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] = {}
176
177
if not disable_cache and os.path.exists(self.cache_file_path):
178
# load the cache from the file
@@ -187,9 +186,9 @@ class CompilerManager:
186
if not isinstance(value, ty):
187
raise TypeError(f"Expected {ty} but got {type(value)} for {value}")
188
190
- def parse_key(key: Any) -> tuple[Range, str, str]:
191
- range_tuple, graph_hash, compiler_name = key
192
- check_type(graph_hash, str)
189
+ def parse_key(key: Any) -> tuple[Range, int, str]:
190
+ range_tuple, graph_index, compiler_name = key
191
+ check_type(graph_index, int)
192
check_type(compiler_name, str)
193
if isinstance(range_tuple, tuple):
194
start, end = range_tuple
@@ -197,7 +196,7 @@ class CompilerManager:
196
check_type(end, int)
197
range_tuple = Range(start=start, end=end)
198
check_type(range_tuple, Range)
200
- return range_tuple, graph_hash, compiler_name
199
+ return range_tuple, graph_index, compiler_name
200
201
self.cache = {parse_key(key): value for key, value in cache.items()}
202
@@ -217,25 +216,18 @@ class CompilerManager:
216
self,
217
graph: fx.GraphModule,
218
example_inputs: list[Any],
220
- graph_hash: str,
219
+ graph_index: int,
220
compile_range: Range,
221
) -> Callable[..., Any] | None:
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:
222
+ if (compile_range, graph_index, self.compiler.name) not in self.cache:
223
return None
230
- handle = self.cache[key]
224
+ handle = self.cache[(compile_range, graph_index, self.compiler.name)]
225
compiled_graph = self.compiler.load(
232
- handle, graph, example_inputs, compile_range
226
+ handle, graph, example_inputs, graph_index, compile_range
227
)
234
- self.loaded_cache_entries[key] = compiled_graph
228
logger.debug(
236
- "Directly load the graph (hash %s) for compile range "
237
- "%sfrom %s via handle %s",
238
- graph_hash,
229
+ "Directly load the %s-th graph for compile range %sfrom %s via handle %s",
230
+ graph_index,
231
str(compile_range),
232
self.compiler.name,
233
handle,
@@ -257,22 +249,12 @@ class CompilerManager:
249
global compilation_start_time
250
compilation_start_time = time.time()
251
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
-
252
compilation_counter.num_backend_compilations += 1
253
254
compiled_graph = None
255
256
# try to load from the cache
275
- compiled_graph = self.load(graph, example_inputs, graph_hash, compile_range)
257
+ compiled_graph = self.load(graph, example_inputs, graph_index, compile_range)
258
if compiled_graph is not None:
259
if graph_index == num_graphs - 1:
260
# after loading the last graph for this shape, record the time.
@@ -308,13 +290,9 @@ class CompilerManager:
290
291
assert compiled_graph is not None, "Failed to compile the graph"
292
311
- self.loaded_cache_entries[(compile_range, graph_hash, self.compiler.name)] = (
312
- compiled_graph
313
- )
314
-
293
# store the artifact in the cache
294
if is_compile_cache_enabled(additional_inductor_config) and handle is not None:
317
- self.cache[(compile_range, graph_hash, self.compiler.name)] = handle
295
+ self.cache[(compile_range, graph_index, self.compiler.name)] = handle
296
compilation_counter.num_cache_entries_updated += 1
297
self.is_cache_updated = True
298
if graph_index == 0:
vllm/compilation/compiler_interface.py
+3
@@ -101,6 +101,7 @@ class CompilerInterface:
101
handle: Any,
102
graph: fx.GraphModule,
103
example_inputs: list[Any],
104
+ graph_index: int,
105
compile_range: Range,
106
) -> Callable[..., Any]:
107
"""
@@ -301,6 +302,7 @@ class InductorStandaloneAdaptor(CompilerInterface):
302
handle: Any,
303
graph: fx.GraphModule,
304
example_inputs: list[Any],
305
+ graph_index: int,
306
compile_range: Range,
307
) -> Callable[..., Any]:
308
assert isinstance(handle, tuple)
@@ -525,6 +527,7 @@ class InductorAdaptor(CompilerInterface):
527
handle: Any,
528
graph: fx.GraphModule,
529
example_inputs: list[Any],
530
+ graph_index: int,
531
compile_range: Range,
532
) -> Callable[..., Any]:
533
assert isinstance(handle, tuple)