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 6, 21:45 UTC
[compile] aot_compile should respect VLLM_DISABLE_COMPILE_CACHE (#36358)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Mar 11, 2026 at 06:12 UTC
09b6f9985225109fbe2c30bc3956501433128aa4
4 files changed
+182
-43
tests/compile/test_aot_compile.py
+114
@@ -4,6 +4,7 @@
4
import functools
5
import hashlib
6
import multiprocessing
7
+import os
8
import pickle
9
import tempfile
10
from contextlib import contextmanager
@@ -19,6 +20,7 @@ from vllm.compilation.caching import (
20
StandaloneCompiledArtifacts,
21
VllmSerializableFunction,
22
)
23
+from vllm.compilation.counter import compilation_counter
24
from vllm.compilation.decorators import support_torch_compile
25
from vllm.config import (
26
CompilationConfig,
@@ -763,3 +765,115 @@ class TestStandaloneCompiledArtifactsIntegration:
765
assert isinstance(config, dict)
766
assert "bundled_autograd_cache" in config
767
assert config["bundled_autograd_cache"] is True
768
+
769
+
770
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
771
+def test_disable_compile_cache_skips_aot_save(
772
+ monkeypatch: pytest.MonkeyPatch, fresh_vllm_cache: str
773
+):
774
+ """When VLLM_DISABLE_COMPILE_CACHE=1, AOT artifacts must not be saved."""
775
+ monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
776
+ monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1")
777
+ disable_envs_cache()
778
+
779
+ args = (torch.randn(10, 10),)
780
+ expected = reference_fn(*args)
781
+ vllm_config = make_vllm_config()
782
+
783
+ with (
784
+ use_vllm_config(vllm_config),
785
+ compilation_counter.expect(
786
+ num_aot_compiles=1,
787
+ num_aot_artifacts_saved=0,
788
+ num_aot_artifacts_loaded=0,
789
+ ),
790
+ ):
791
+ mod = CompiledMod(vllm_config=vllm_config)
792
+ actual = mod(*args)
793
+
794
+ assert torch.allclose(actual, expected)
795
+
796
+ # No cached artifact should exist on disk
797
+ aot_dir = os.path.join(fresh_vllm_cache, "torch_compile_cache", "torch_aot_compile")
798
+ if os.path.isdir(aot_dir):
799
+ for root, _dirs, files in os.walk(aot_dir):
800
+ for f in files:
801
+ assert f != "model", (
802
+ f"AOT artifact unexpectedly saved at {os.path.join(root, f)}"
803
+ )
804
+
805
+
806
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
807
+def test_disable_compile_cache_skips_aot_load(
808
+ monkeypatch: pytest.MonkeyPatch, fresh_vllm_cache: str
809
+):
810
+ """When VLLM_DISABLE_COMPILE_CACHE=1, AOT artifacts must not be loaded."""
811
+ # Phase 1: compile and save with cache enabled
812
+ monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1")
813
+ disable_envs_cache()
814
+
815
+ args = (torch.randn(10, 10),)
816
+ vllm_config = make_vllm_config()
817
+
818
+ with (
819
+ use_vllm_config(vllm_config),
820
+ compilation_counter.expect(num_aot_artifacts_saved=1),
821
+ ):
822
+ CompiledMod(vllm_config=vllm_config)(*args)
823
+
824
+ # Phase 2: disable cache, compile again — should NOT load from disk
825
+ monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
826
+ disable_envs_cache()
827
+ torch._dynamo.reset()
828
+
829
+ vllm_config = make_vllm_config()
830
+ with (
831
+ use_vllm_config(vllm_config),
832
+ compilation_counter.expect(
833
+ num_aot_compiles=1,
834
+ num_aot_artifacts_saved=0,
835
+ num_aot_artifacts_loaded=0,
836
+ ),
837
+ ):
838
+ mod = CompiledMod(vllm_config=vllm_config)
839
+ mod(*args)
840
+
841
+ assert not mod.was_aot_compile_fn_loaded_from_disk
842
+
843
+
844
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
845
+def test_aot_counters_on_save_and_load(
846
+ monkeypatch: pytest.MonkeyPatch, fresh_vllm_cache: str
847
+):
848
+ """Verify AOT counters are incremented correctly on save and load."""
849
+ monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1")
850
+ disable_envs_cache()
851
+
852
+ args = (torch.randn(10, 10),)
853
+
854
+ # Phase 1: fresh compile + save
855
+ vllm_config = make_vllm_config()
856
+ with (
857
+ use_vllm_config(vllm_config),
858
+ compilation_counter.expect(
859
+ num_aot_compiles=1,
860
+ num_aot_artifacts_saved=1,
861
+ num_aot_artifacts_loaded=0,
862
+ ),
863
+ ):
864
+ CompiledMod(vllm_config=vllm_config)(*args)
865
+
866
+ # Phase 2: load from cache
867
+ monkeypatch.setenv("VLLM_FORCE_AOT_LOAD", "1")
868
+ disable_envs_cache()
869
+
870
+ vllm_config = make_vllm_config()
871
+ with (
872
+ use_vllm_config(vllm_config),
873
+ compilation_counter.expect(
874
+ num_aot_compiles=0,
875
+ num_aot_artifacts_saved=0,
876
+ num_aot_artifacts_loaded=1,
877
+ ),
878
+ ):
879
+ CompiledMod(vllm_config=vllm_config)(*args)
vllm/compilation/counter.py
+6
@@ -31,6 +31,12 @@ class CompilationCounter:
31
num_compiled_artifacts_saved: int = 0
32
# The number of standalone_compile compiled artifacts loaded from cache
33
num_compiled_artifacts_loaded: int = 0
34
+ # The number of AOT compile invocations
35
+ num_aot_compiles: int = 0
36
+ # The number of AOT compiled artifacts saved to disk
37
+ num_aot_artifacts_saved: int = 0
38
+ # The number of AOT compiled artifacts loaded from disk
39
+ num_aot_artifacts_loaded: int = 0
40
# Number of times a model was loaded with CompilationMode.STOCK_TORCH_COMPILE
41
stock_torch_compile_count: int = 0
42
vllm/compilation/decorators.py
+59
-43
@@ -266,6 +266,51 @@ def _verify_source_unchanged(
266
)
267
268
269
+def _try_load_aot_compiled_fn(
270
+ model: Any,
271
+ aot_compilation_path: str,
272
+) -> Any | None:
273
+ """Try to load an AOT-compiled function from disk.
274
+
275
+ Returns the loaded callable on success, or None on failure.
276
+ Re-raises on failure when ``VLLM_FORCE_AOT_LOAD`` is set.
277
+ """
278
+ try:
279
+ with monitor_torch_compile(model.vllm_config):
280
+ with (
281
+ set_current_vllm_config(model.vllm_config),
282
+ open(aot_compilation_path, "rb") as f,
283
+ ):
284
+ loaded_fn = torch.compiler.load_compiled_function(
285
+ f, f_globals=model.forward.__globals__
286
+ )
287
+ _verify_source_unchanged(loaded_fn.source_info(), model.vllm_config)
288
+ ds_config = model.compilation_config.dynamic_shapes_config
289
+ if not ds_config.evaluate_guards:
290
+ loaded_fn.disable_guard_check()
291
+ # Eagerly load compiled artifacts now that traced_files
292
+ # is populated by _verify_source_unchanged.
293
+ with maybe_use_cudagraph_partition_wrapper(model.vllm_config):
294
+ loaded_fn._artifacts.compiled_fn.finalize_loading(model.vllm_config)
295
+ compilation_counter.num_aot_artifacts_loaded += 1
296
+ logger.info("Directly load AOT compilation from path %s", aot_compilation_path)
297
+ return loaded_fn
298
+ except Exception as e:
299
+ if os.path.exists(aot_compilation_path):
300
+ if isinstance(e, EOFError):
301
+ message = "Compile cache file corrupted."
302
+ else:
303
+ message = str(e)
304
+ logger.warning(
305
+ "Compiling model again due to a load failure from %s, reason: %s",
306
+ aot_compilation_path,
307
+ message,
308
+ )
309
+ if envs.VLLM_FORCE_AOT_LOAD:
310
+ raise e
311
+ return None
312
+
313
+
314
def _support_torch_compile(
315
cls: type[_T],
316
dynamic_arg_dims: dict[str, int | list[int]],
@@ -438,51 +483,17 @@ def _support_torch_compile(
483
dp_rank = self.vllm_config.parallel_config.data_parallel_index
484
cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")
485
aot_compilation_path = os.path.join(cache_dir, "model")
441
- try:
442
- with monitor_torch_compile(self.vllm_config):
486
+ if not envs.VLLM_DISABLE_COMPILE_CACHE:
487
+ loaded_fn = _try_load_aot_compiled_fn(self, aot_compilation_path)
488
+ if loaded_fn is not None:
489
+ self.aot_compiled_fn = loaded_fn
490
+ self.was_aot_compile_fn_loaded_from_disk = True
491
with (
444
- set_current_vllm_config(self.vllm_config),
445
- open(aot_compilation_path, "rb") as f,
492
+ monitor_profiling_run(),
493
+ maybe_use_cudagraph_partition_wrapper(self.vllm_config),
494
):
447
- loaded_fn = torch.compiler.load_compiled_function(
448
- f, f_globals=self.forward.__globals__
449
- )
450
- _verify_source_unchanged(loaded_fn.source_info(), self.vllm_config)
451
- ds_config = self.compilation_config.dynamic_shapes_config
452
- if not ds_config.evaluate_guards:
453
- loaded_fn.disable_guard_check()
454
- # Eagerly load compiled artifacts now that traced_files
455
- # is populated by _verify_source_unchanged.
456
- with maybe_use_cudagraph_partition_wrapper(self.vllm_config):
457
- loaded_fn._artifacts.compiled_fn.finalize_loading(
458
- self.vllm_config
459
- )
460
- self.aot_compiled_fn = loaded_fn
461
- self.was_aot_compile_fn_loaded_from_disk = True
462
- except Exception as e:
463
- if os.path.exists(aot_compilation_path):
464
- if isinstance(e, EOFError):
465
- message = "Compile cache file corrupted."
466
- else:
467
- message = str(e)
468
- logger.warning(
469
- "Compiling model again due to a load failure from %s, "
470
- "reason: %s",
471
- aot_compilation_path,
472
- message,
473
- )
474
- if envs.VLLM_FORCE_AOT_LOAD:
475
- raise e
476
- if getattr(self, "aot_compiled_fn", None) is not None:
477
- logger.info(
478
- "Directly load AOT compilation from path %s", aot_compilation_path
479
- )
480
- with (
481
- monitor_profiling_run(),
482
- maybe_use_cudagraph_partition_wrapper(self.vllm_config),
483
- ):
484
- output = self.aot_compiled_fn(self, *args, **kwargs)
485
- return output
495
+ output = self.aot_compiled_fn(self, *args, **kwargs)
496
+ return output
497
498
if self.compiled:
499
assert (
@@ -570,6 +581,7 @@ def _support_torch_compile(
581
self._aot_cache_dir = cache_dir
582
with monitor_torch_compile(self.vllm_config):
583
self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
584
+ compilation_counter.num_aot_compiles += 1
585
# All compilation is done at this point, save the
586
# AOT artifact.
587
self.save_aot_compiled_function()
@@ -593,6 +605,9 @@ def _support_torch_compile(
605
606
# triggers VllmSerializableFunction.serialize()
607
def save_aot_compiled_function(self: type[_T]) -> None:
608
+ if envs.VLLM_DISABLE_COMPILE_CACHE:
609
+ return
610
+
611
if self.was_aot_compile_fn_loaded_from_disk:
612
logger.debug("AOT compiled function was loaded from cache, skipping save")
613
return
@@ -608,6 +623,7 @@ def _support_torch_compile(
623
tmp_file = f"{self._aot_compilation_path}.{os.getpid()}.tmp"
624
self.aot_compiled_fn.save_compiled_function(tmp_file)
625
os.replace(tmp_file, self._aot_compilation_path)
626
+ compilation_counter.num_aot_artifacts_saved += 1
627
logger.info_once(
628
"saved AOT compiled function to %s",
629
self._aot_compilation_path,
vllm/compilation/wrapper.py
+3
@@ -349,6 +349,9 @@ def reset_compile_wrapper(model: torch.nn.Module) -> None:
349
compilation_counter.num_cache_entries_updated = 0
350
compilation_counter.num_compiled_artifacts_saved = 0
351
compilation_counter.stock_torch_compile_count = 0
352
+ compilation_counter.num_aot_compiles = 0
353
+ compilation_counter.num_aot_artifacts_saved = 0
354
+ compilation_counter.num_aot_artifacts_loaded = 0
355
356
# Clear the AOT compiled function so the model is forced to
357
# recompile on the next call. Without this, decorators.py