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, 17:45 UTC
[torch.compile] Stop compiling identical artifacts (#34003)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Feb 7, 2026 at 08:24 UTC
81fe69cae5cb0d3aff23934064c639e9c9815a81
2 files changed
+64
-11
tests/compile/test_cold_start.py
+4
-4
@@ -37,12 +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
- # The 33 subgraphs then get standalone_compile'd.
40
+ # We then generate compiled artifacts for 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, we are only compling
46
+ # for 3 unique subgraphs.
47
assert counters["aot_autograd"]["autograd_cache_miss"] == 3
48
- assert counters["aot_autograd"]["autograd_cache_hit"] == 30
48
+ assert counters["aot_autograd"]["autograd_cache_hit"] == 0
vllm/compilation/backends.py
+60
-7
@@ -134,6 +134,7 @@ class CompilerManager:
134
self.is_cache_updated = False
135
self.compilation_config = compilation_config
136
self.compiler = make_compiler(compilation_config)
137
+ self.loaded_artifacts: dict[str, Any] = {}
138
139
def compute_hash(self, vllm_config: VllmConfig) -> str:
140
return self.compiler.compute_hash(vllm_config)
@@ -282,13 +283,61 @@ class CompilerManager:
283
maybe_key += f"{compile_range.start}_{compile_range.end}"
284
maybe_key += f"_subgraph_{graph_index}"
285
with self.compile_context(compile_range):
285
- compiled_graph, handle = self.compiler.compile(
286
- graph,
287
- example_inputs,
288
- additional_inductor_config,
289
- compile_range,
290
- maybe_key,
291
- )
286
+ # There is a compilation time optimization here.
287
+ #
288
+ # If the (input metdata, graph, compiler config) are the same, then
289
+ # we want to avoid compiling the same artifact again. If we didn't
290
+ # do this optimization, the backend compilation (InductorAdaptor or
291
+ # InductorStandaloneAdaptor)
292
+ # is able to cache hit and produce an artifact faster if it was
293
+ # already created, but it is still a duplicate artifact that
294
+ # requires unnecessary things e.g. disk IO.
295
+ #
296
+ # The optimization is: If the backend compilation cache hits,
297
+ # then do an early return from the backend compilation and look up
298
+ # which of the previous in-memory artifacts we created to reuse.
299
+ #
300
+ # We implemented this by monkey-patching torch (torch does not
301
+ # easily expose the cache_key function), but in the future torch
302
+ # should expose the cache_key function that we can just call
303
+ # directly before invoking backend compilation.
304
+ cache_key = None
305
+ orig = torch._functorch._aot_autograd.autograd_cache.autograd_cache_key
306
+
307
+ def autograd_cache_key(*args, **kwargs):
308
+ result = orig(*args, **kwargs)
309
+ if result is None:
310
+ return None
311
+ nonlocal cache_key
312
+ cache_key = result[0]
313
+ if cache_key in self.loaded_artifacts:
314
+ raise StopCompiling()
315
+ return result
316
+
317
+ from unittest.mock import patch
318
+
319
+ with (
320
+ # Graphs that are isometric (different node names but same
321
+ # structure) should be treated as the same.
322
+ torch._functorch.config.patch(autograd_cache_normalize_inputs=True),
323
+ patch(
324
+ "torch._functorch._aot_autograd.autograd_cache.autograd_cache_key",
325
+ autograd_cache_key,
326
+ ),
327
+ ):
328
+ try:
329
+ compiled_graph, handle = self.compiler.compile(
330
+ graph,
331
+ example_inputs,
332
+ additional_inductor_config,
333
+ compile_range,
334
+ maybe_key,
335
+ )
336
+ except StopCompiling:
337
+ assert cache_key is not None
338
+ return self.loaded_artifacts[cache_key]
339
+ if cache_key is not None and compiled_graph is not None:
340
+ self.loaded_artifacts[cache_key] = compiled_graph
341
342
assert compiled_graph is not None, "Failed to compile the graph"
343
@@ -326,6 +375,10 @@ class CompilerManager:
375
return compiled_graph
376
377
378
+class StopCompiling(BaseException):
379
+ pass
380
+
381
+
382
@dataclasses.dataclass
383
class SplitItem:
384
submod_name: str