[Bugfix][Compile] Replace Transformers AOT caching with graph caching (#58918)

Signed-off-by: fangyao fan <2281693949@qq.com>
Signed-off-by: Taimys <ffy2281693949@gmail.com>
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Co-authored-by: Taimys <ffy2281693949@gmail.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

tests/compile/test_aot_compile.py

@@ -1,9 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import functools
 import hashlib
+import json
 import os
 import pickle
+import subprocess
+import sys
 import tempfile
 from contextlib import contextmanager
 from pathlib import Path
@@ -890,3 +894,176 @@ def test_disable_compile_cache_skips_aot_load(
         mod(*args)
 
     assert not mod.was_aot_compile_fn_loaded_from_disk
+
+
+_CACHE_OFFSET = int(os.environ.get("VLLM_TEST_CACHE_OFFSET", "0"))
+
+
+def _capture_cache_offset(fn):
+    offset = _CACHE_OFFSET
+
+    @functools.wraps(fn)
+    def wrapped(self, x):
+        return fn(self, x) + offset
+
+    return wrapped
+
+
+def _run_graph_cache_probe():
+    """Fresh-process probe; only the vLLM cache root survives between runs."""
+    from torch._inductor.runtime.cache_dir_utils import cache_dir
+
+    @support_torch_compile(dynamic_arg_dims={"x": 0}, use_aot_compile=False)
+    class GraphCacheMod(torch.nn.Module):
+        def __init__(self, **kwargs):
+            super().__init__()
+
+        @_capture_cache_offset
+        def forward(self, x):
+            return x * 2 + 1
+
+    cache_dir()  # Exercise a tmp default populated before backend initialization.
+    config = make_vllm_config()
+    x = torch.arange(64, device=current_platform.device_type).reshape(8, 8).float()
+    expected = x * 2 + 1 + int(os.environ["VLLM_TEST_CACHE_OFFSET"])
+    with use_vllm_config(config):
+        mod = GraphCacheMod(vllm_config=config)
+        assert torch.equal(mod(x), expected)
+        traced = torch._dynamo.utils.counters["stats"]["unique_graphs"]
+        assert traced > 0
+        aot_compiles = compilation_counter.num_aot_compiles
+        assert torch.equal(mod(x), expected)
+        assert torch._dynamo.utils.counters["stats"]["unique_graphs"] == traced
+        assert compilation_counter.num_aot_compiles == aot_compiles == 0
+    for name in (
+        "num_aot_artifacts_saved",
+        "num_aot_artifacts_loaded",
+        "num_compiled_artifacts_saved",
+        "num_compiled_artifacts_loaded",
+    ):
+        assert getattr(compilation_counter, name) == 0
+    assert Path(cache_dir()).is_relative_to(config.compilation_config.cache_dir)
+    assert Path(os.environ["TRITON_CACHE_DIR"]).is_relative_to(
+        config.compilation_config.cache_dir
+    )
+    result = {
+        "inductor_dir": cache_dir(),
+        "cache_dir": config.compilation_config.cache_dir,
+        "fx_hits": torch._dynamo.utils.counters["inductor"]["fxgraph_cache_hit"],
+        "autograd_hits": torch._dynamo.utils.counters["aot_autograd"][
+            "autograd_cache_hit"
+        ],
+    }
+    native = os.environ.get("VLLM_TEST_NATIVE_AOT")
+    if native:
+        if native == "warm":
+            os.environ["VLLM_FORCE_AOT_LOAD"] = "1"
+            disable_envs_cache()
+        with use_vllm_config(config):
+            native_mod = CompiledMod(vllm_config=config)
+            assert torch.equal(native_mod(x), reference_fn(x))
+        assert compilation_counter.num_aot_artifacts_loaded == int(native == "warm")
+        assert compilation_counter.num_aot_artifacts_saved == int(native == "cold")
+        assert compilation_counter.num_aot_compiles == int(native == "cold")
+    Path(os.environ["VLLM_TEST_CACHE_RESULT"]).write_text(json.dumps(result))
+
+
+@pytest.fixture
+def graph_cache_probe(tmp_path):
+    def run(offset, *, standalone="1", native=""):
+        private = Path(tempfile.mkdtemp(dir=tmp_path))
+        env = os.environ.copy()
+        for key in ("TORCHINDUCTOR_CACHE_DIR", "TRITON_CACHE_DIR"):
+            env.pop(key, None)
+        for key in ("HOME", "TMPDIR", "XDG_CACHE_HOME"):
+            directory = private / key
+            directory.mkdir()
+            env[key] = str(directory)
+        result = private / "result.json"
+        env.update(
+            VLLM_CACHE_ROOT=str(tmp_path / "shared"),
+            VLLM_USE_AOT_COMPILE="1",
+            VLLM_USE_MEGA_AOT_ARTIFACT="1",
+            VLLM_USE_STANDALONE_COMPILE=standalone,
+            VLLM_FORCE_AOT_LOAD="0",
+            VLLM_DISABLE_COMPILE_CACHE="0",
+            VLLM_USE_BYTECODE_HOOK="0",
+            VLLM_TEST_CACHE_OFFSET=str(offset),
+            VLLM_TEST_CACHE_RESULT=str(result),
+            VLLM_TEST_NATIVE_AOT=native,
+            TORCHINDUCTOR_FX_GRAPH_REMOTE_CACHE="0",
+            TORCHINDUCTOR_AUTOGRAD_REMOTE_CACHE="0",
+        )
+        proc = subprocess.run(
+            [
+                sys.executable,
+                "-c",
+                "from tests.compile.test_aot_compile import "
+                "_run_graph_cache_probe; _run_graph_cache_probe()",
+            ],
+            cwd=Path(__file__).resolve().parents[2],
+            env=env,
+            capture_output=True,
+            text=True,
+            timeout=180,
+        )
+        assert proc.returncode == 0, proc.stdout + proc.stderr
+        return json.loads(result.read_text())
+
+    return run
+
+
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
+@pytest.mark.parametrize("standalone", ["0", "1"])
+def test_compile_opt_out_retraces_changed_closure(graph_cache_probe, standalone):
+    """A warm graph is reused, but a changed closure never loads a stale handle."""
+    results = [
+        graph_cache_probe(offset, standalone=standalone) for offset in (0, 0, 7, 7)
+    ]
+    assert len({r["cache_dir"] for r in results}) == 1
+    assert len({r["inductor_dir"] for r in results}) == 1
+    for warm in (results[1], results[3]):
+        assert warm["fx_hits"] > 0
+        if standalone == "1":
+            assert warm["autograd_hits"] > 0
+
+
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
+def test_compile_opt_out_preserves_native_aot_roundtrip(graph_cache_probe):
+    """Native AOT survives an opt-out module compiled first under the same config."""
+    graph_cache_probe(0, native="cold")
+    graph_cache_probe(0, native="warm")
+
+
+@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
+def test_compile_opt_out_respected_in_dynamo_trace_once(monkeypatch, fresh_vllm_cache):
+    """Opt-out skips model AOT without disabling tracing or causing retracing."""
+    monkeypatch.setenv("VLLM_USE_AOT_COMPILE", "1")
+    monkeypatch.setenv("VLLM_FORCE_AOT_LOAD", "0")
+    monkeypatch.setenv("VLLM_USE_BYTECODE_HOOK", "0")
+    disable_envs_cache()
+
+    @support_torch_compile(dynamic_arg_dims={"x": 0}, use_aot_compile=False)
+    class OptOutMod(torch.nn.Module):
+        def forward(self, x):
+            return x * 2 + 1
+
+    config = VllmConfig(
+        compilation_config=CompilationConfig(
+            mode=CompilationMode.DYNAMO_TRACE_ONCE, backend="inductor"
+        )
+    )
+    x = torch.arange(64, device=current_platform.device_type).reshape(8, 8).float()
+    counters = torch._dynamo.utils.counters["stats"]
+    with (
+        use_vllm_config(config),
+        patch.object(OptOutMod, "aot_compile", side_effect=AssertionError("Model AOT")),
+    ):
+        assert envs.VLLM_USE_AOT_COMPILE
+        mod = OptOutMod(vllm_config=config)
+        before = counters["unique_graphs"]
+        assert torch.equal(mod(x), x * 2 + 1)
+        traced = counters["unique_graphs"]
+        assert traced > before
+        assert torch.equal(mod(x + 10), (x + 10) * 2 + 1)
+        assert counters["unique_graphs"] == traced

tests/compile/test_compiler_interface.py

@@ -18,7 +18,7 @@ def test_inductor_standalone_load_uses_current_cache_dir(
     new_cache_dir = tmp_path / "new"
     key = "artifact_shape_None_subgraph_0"
 
-    adaptor = InductorStandaloneAdaptor(save_format=save_format)
+    adaptor = InductorStandaloneAdaptor(use_aot_compile=True, save_format=save_format)
     adaptor.initialize_cache(str(new_cache_dir))
 
     # Simulate a handle persisted before the cache directory was relocated.
@@ -46,3 +46,86 @@ def test_inductor_standalone_load_uses_current_cache_dir(
         path=str(new_cache_dir / key),
         format=save_format,
     )
+
+
+@pytest.mark.parametrize("raises", [False, True])
+def test_opt_out_functorch_config_is_scoped(monkeypatch, raises):
+    from contextlib import nullcontext
+
+    import torch
+
+    from vllm.compilation.compiler_interface import (
+        get_inductor_factors,
+        set_functorch_config,
+    )
+
+    config = torch._functorch.config
+    monkeypatch.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
+    from vllm.envs import disable_envs_cache
+
+    disable_envs_cache()
+    with config.patch(bundled_autograd_cache=True):
+        error = pytest.raises(RuntimeError, match="probe") if raises else nullcontext()
+        with error, set_functorch_config(use_aot_compile=False):
+            assert config.bundled_autograd_cache is False
+            if raises:
+                raise RuntimeError("probe")
+        assert config.bundled_autograd_cache is True
+        original = type(config).save_config_portable
+        observed = []
+
+        def observe(module):
+            if module is config:
+                observed.append(config.bundled_autograd_cache)
+            return original(module)
+
+        with patch.object(type(config), "save_config_portable", new=observe):
+            get_inductor_factors(use_aot_compile=False)
+        assert observed == [False]
+        assert config.bundled_autograd_cache is True
+
+
+def test_opt_out_ignores_existing_direct_cache_index(tmp_path, monkeypatch):
+    import builtins
+
+    from vllm.compilation.backends import CompilerManager
+    from vllm.config import CompilationConfig, CompilationMode
+    from vllm.config.utils import Range
+    from vllm.envs import disable_envs_cache
+
+    monkeypatch.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
+    disable_envs_cache()
+    cache = tmp_path / "rank_0_0" / "decoder"
+    cache.mkdir(parents=True)
+    index = cache / "vllm_compile_cache.py"
+    index.write_text("old index must not be read")
+    manager = CompilerManager(
+        CompilationConfig(mode=CompilationMode.VLLM_COMPILE), use_aot_compile=False
+    )
+    real_open = builtins.open
+    reads = []
+
+    def check_open(path, *args, **kwargs):
+        if str(path) == str(index):
+            reads.append(path)
+            raise AssertionError("read old index")
+        return real_open(path, *args, **kwargs)
+
+    with (
+        patch.object(manager.compiler, "initialize_cache"),
+        patch("builtins.open", side_effect=check_open),
+    ):
+        manager.initialize_cache(str(cache), prefix="decoder")
+    assert Path(manager.cache_file_path) == index
+    assert reads == []
+    compile_range = Range(1, 8)
+    manager.cache[(compile_range, 0, manager.compiler.name)] = {
+        "graph_handle": ("old", "unused"),
+        "cache_key": "old",
+    }
+    with patch.object(
+        manager.compiler, "load", side_effect=AssertionError("direct load")
+    ) as load:
+        assert manager.load(MagicMock(), [], 0, compile_range) is None
+        load.assert_not_called()
+    assert index.read_text() == "old index must not be read"

tests/v1/attention/test_rocm_attention_backends_selection.py

@@ -206,7 +206,9 @@ def test_segmented_triton_cache_survives_standalone_compilation(
         RocmSegmentedAttentionImpl(8, 128, 128**-0.5, 2, None, None, "auto")
         assert os.environ["TRITON_CACHE_DIR"] == str(expected)
         assert expected.is_dir()
-        InductorStandaloneAdaptor("binary").initialize_cache(str(tmp_path / model))
+        InductorStandaloneAdaptor(
+            use_aot_compile=True, save_format="binary"
+        ).initialize_cache(str(tmp_path / model))
         assert os.environ["TRITON_CACHE_DIR"] == str(expected)
 
 

vllm/compilation/backends.py

@@ -94,10 +94,13 @@ def make_copy_and_call(
     return copy_and_call
 
 
-def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface:
-    assert not envs.VLLM_USE_MEGA_AOT_ARTIFACT or envs.VLLM_USE_STANDALONE_COMPILE, (
-        "VLLM_USE_MEGA_AOT_ARTIFACT=1 requires VLLM_USE_STANDALONE_COMPILE=1"
-    )
+def make_compiler(
+    compilation_config: CompilationConfig, *, use_aot_compile: bool
+) -> CompilerInterface:
+    assert (
+        not (use_aot_compile and envs.VLLM_USE_MEGA_AOT_ARTIFACT)
+        or envs.VLLM_USE_STANDALONE_COMPILE
+    ), "VLLM_USE_MEGA_AOT_ARTIFACT=1 requires VLLM_USE_STANDALONE_COMPILE=1"
 
     if compilation_config.backend == "inductor":
         # Use standalone compile only if requested, version is new enough,
@@ -107,11 +110,12 @@ def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface:
         ):
             logger.debug("Using InductorStandaloneAdaptor")
             return InductorStandaloneAdaptor(
-                compilation_config.compile_cache_save_format
+                compilation_config.compile_cache_save_format,
+                use_aot_compile=use_aot_compile,
             )
         else:
             logger.debug("Using InductorAdaptor")
-            return InductorAdaptor()
+            return InductorAdaptor(use_aot_compile=use_aot_compile)
     elif compilation_config.backend == "eager":
         logger.debug("Using EagerAdaptor")
         return EagerAdaptor()
@@ -136,11 +140,16 @@ class CompilerManager:
     support int as key.
     """
 
-    def __init__(self, compilation_config: CompilationConfig) -> None:
+    def __init__(
+        self, compilation_config: CompilationConfig, *, use_aot_compile: bool
+    ) -> None:
         self.cache: dict[tuple[Range, int, str], Any] = dict()
         self.is_cache_updated = False
         self.compilation_config = compilation_config
-        self.compiler = make_compiler(compilation_config)
+        self.use_aot_compile = use_aot_compile
+        self.compiler = make_compiler(
+            compilation_config, use_aot_compile=use_aot_compile
+        )
         self.loaded_artifacts: dict[str, Any] = {}
         self.prefix: str = ""
 
@@ -185,7 +194,11 @@ class CompilerManager:
         self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py")
         self.prefix = prefix
 
-        if not disable_cache and os.path.exists(self.cache_file_path):
+        if (
+            self.use_aot_compile
+            and not disable_cache
+            and os.path.exists(self.cache_file_path)
+        ):
             # load the cache from the file
             with open(self.cache_file_path) as f:
                 # we use ast.literal_eval to parse the data
@@ -216,7 +229,7 @@ class CompilerManager:
         )
 
     def save_to_file(self) -> None:
-        if self.disable_cache or not self.is_cache_updated:
+        if not self.use_aot_compile or self.disable_cache or not self.is_cache_updated:
             return
         printer = pprint.PrettyPrinter(indent=4)
         data = printer.pformat(self.cache)
@@ -230,6 +243,8 @@ class CompilerManager:
         graph_index: int,
         compile_range: Range,
     ) -> Callable[..., Any] | None:
+        if not self.use_aot_compile:
+            return None
         if (compile_range, graph_index, self.compiler.name) not in self.cache:
             return None
 
@@ -370,7 +385,11 @@ class CompilerManager:
         assert compiled_graph is not None, "Failed to compile the graph"
 
         # store the artifact in the cache
-        if is_compile_cache_enabled(additional_inductor_config) and handle is not None:
+        if (
+            self.use_aot_compile
+            and is_compile_cache_enabled(additional_inductor_config)
+            and handle is not None
+        ):
             self.cache[(compile_range, graph_index, self.compiler.name)] = {
                 "graph_handle": handle,
                 "cache_key": cache_key,
@@ -839,6 +858,7 @@ class VllmBackend:
         vllm_config: VllmConfig,
         prefix: str = "",
         is_encoder: bool = False,
+        use_aot_compile: bool = True,
     ) -> None:
         # if the model is initialized with a non-empty prefix,
         # then usually it's enough to use that prefix,
@@ -859,9 +879,10 @@ class VllmBackend:
 
         self.vllm_config = vllm_config
         self.compilation_config = vllm_config.compilation_config
+        self.use_mega_artifact = use_aot_compile and envs.VLLM_USE_MEGA_AOT_ARTIFACT
 
         self.compiler_manager: CompilerManager = CompilerManager(
-            self.compilation_config
+            self.compilation_config, use_aot_compile=use_aot_compile
         )
 
         # Deepcopy the inductor config to detach the post-grad custom pass
@@ -892,7 +913,7 @@ class VllmBackend:
                   returns_tuple
 
         """
-        if not envs.VLLM_USE_MEGA_AOT_ARTIFACT:
+        if not self.use_mega_artifact:
             return None, None, None
 
         from .caching import StandaloneCompiledArtifacts
@@ -1193,7 +1214,7 @@ class VllmBackend:
         # keep a split_gm copy from BEFORE the interpreter replaces
         # submodules with PiecewiseBackend -- used for serialization
         original_split_gm = None
-        if envs.VLLM_USE_MEGA_AOT_ARTIFACT:
+        if self.use_mega_artifact:
             original_split_gm = deepcopy(self.split_gm)
 
         from torch._dynamo.utils import lazy_format_graph_code
@@ -1285,9 +1306,7 @@ class VllmBackend:
             logger.debug_once("Computation graph saved to %s", graph_path)
 
         self._called = True
-        graph_to_serialize = (
-            original_split_gm if envs.VLLM_USE_MEGA_AOT_ARTIFACT else self.graph
-        )
+        graph_to_serialize = original_split_gm if self.use_mega_artifact else self.graph
 
         execution_code, submod_names, consts = generate_execution_code(self.split_gm)
         # Use getattr to get correct callables: __dict__ has PiecewiseBackend

vllm/compilation/compiler_interface.py

@@ -3,7 +3,7 @@
 import contextlib
 import copy
 import os
-from collections.abc import Callable
+from collections.abc import Callable, Generator
 from contextlib import ExitStack
 from typing import Any, Literal
 from unittest.mock import patch
@@ -26,6 +26,22 @@ logger = init_logger(__name__)
 _last_inductor_triton_cache_dir: str | None = None
 
 
+def initialize_inductor_cache(base_cache_dir: str) -> None:
+    global _last_inductor_triton_cache_dir
+    # Keep artifacts together unless the caller supplies a Triton cache.
+    inductor_cache = os.path.join(base_cache_dir, "inductor_cache")
+    os.makedirs(inductor_cache, exist_ok=True)
+    os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache
+    triton_cache = os.environ.get("TRITON_CACHE_DIR")
+    # A directory assigned by an earlier adaptor is still a model-local
+    # default, rather than a caller override for subsequent models.
+    if triton_cache is None or triton_cache == _last_inductor_triton_cache_dir:
+        triton_cache = os.path.join(base_cache_dir, "triton_cache")
+        _last_inductor_triton_cache_dir = triton_cache
+    os.makedirs(triton_cache, exist_ok=True)
+    os.environ["TRITON_CACHE_DIR"] = triton_cache
+
+
 class CompilerInterface:
     """The interface for a compiler that can be used by vLLM."""
 
@@ -151,18 +167,18 @@ class AlwaysHitShapeEnv:
         return ""
 
 
-def _get_vllm_functorch_config() -> dict[str, Any]:
+def _get_vllm_functorch_config(use_aot_compile: bool = True) -> dict[str, Any]:
     """Return the functorch config overrides that vLLM applies at compile time.
 
     Used by both set_functorch_config() and get_inductor_factors() to ensure
     the compile-time config and cache key are always consistent."""
     cfg: dict[str, Any] = {}
-    if not envs.VLLM_USE_MEGA_AOT_ARTIFACT:
+    if not (use_aot_compile and envs.VLLM_USE_MEGA_AOT_ARTIFACT):
         cfg["bundled_autograd_cache"] = False
     return cfg
 
 
-def get_inductor_factors() -> list[Any]:
+def get_inductor_factors(use_aot_compile: bool = True) -> list[Any]:
     factors: list[Any] = []
     # summarize system state
     from torch._inductor.codecache import CacheBase
@@ -180,7 +196,7 @@ def get_inductor_factors() -> list[Any]:
     from torch._inductor import config as inductor_config
 
     factors.append(inductor_config.save_config_portable())
-    with functorch_config.patch(_get_vllm_functorch_config()):
+    with functorch_config.patch(_get_vllm_functorch_config(use_aot_compile)):
         factors.append(functorch_config.save_config_portable())
     return factors
 
@@ -254,13 +270,16 @@ class InductorStandaloneAdaptor(CompilerInterface):
 
     name = "inductor_standalone"
 
-    def __init__(self, save_format: Literal["binary", "unpacked"]) -> None:
+    def __init__(
+        self, save_format: Literal["binary", "unpacked"], *, use_aot_compile: bool
+    ) -> None:
         if not is_torch_equal_or_newer("2.10.0"):
             _patch_standalone_compile_atomic_save()
         self.save_format = save_format
+        self.use_aot_compile = use_aot_compile
 
     def compute_hash(self, vllm_config: VllmConfig) -> str:
-        factors = get_inductor_factors()
+        factors = get_inductor_factors(self.use_aot_compile)
         hash_str: str = safe_hash(
             str(factors).encode(), usedforsecurity=False
         ).hexdigest()[:10]
@@ -270,6 +289,10 @@ class InductorStandaloneAdaptor(CompilerInterface):
         self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
     ) -> None:
         self.cache_dir = cache_dir
+        if not self.use_aot_compile and not disable_cache:
+            initialize_inductor_cache(
+                cache_dir[: -len(prefix)] if prefix else cache_dir
+            )
 
     def compile(
         self,
@@ -285,7 +308,6 @@ class InductorStandaloneAdaptor(CompilerInterface):
         if compiler_config is not None:
             current_config.update(compiler_config)
         set_inductor_config(current_config, compile_range)
-        set_functorch_config()
 
         if compile_range.is_single_size():
             dynamic_shapes = "from_example_inputs"
@@ -296,7 +318,11 @@ class InductorStandaloneAdaptor(CompilerInterface):
 
         supports_aot = is_torch_equal_or_newer("2.10.0")
 
-        if not supports_aot and envs.VLLM_USE_MEGA_AOT_ARTIFACT:
+        if (
+            not supports_aot
+            and self.use_aot_compile
+            and envs.VLLM_USE_MEGA_AOT_ARTIFACT
+        ):
             logger.error(
                 "CRITICAL: VLLM_USE_MEGA_AOT_ARTIFACT "
                 "is enabled but PyTorch version does not support 'aot' "
@@ -314,7 +340,9 @@ class InductorStandaloneAdaptor(CompilerInterface):
         if is_torch_equal_or_newer("2.13.0.dev"):
             compile_kwargs["donate_graph_module"] = True  # type: ignore[assignment]
 
-        use_aot: bool = supports_aot and envs.VLLM_USE_MEGA_AOT_ARTIFACT
+        use_aot: bool = (
+            supports_aot and self.use_aot_compile and envs.VLLM_USE_MEGA_AOT_ARTIFACT
+        )
         # only add 'aot' parameter if both supported and enabled...
         # this will set bundled_autograd_cache
         # https://github.com/pytorch/pytorch/blob/9bbc5b2905c260adf41bc866a732f9c121a2828a/torch/_inductor/standalone_compile.py#L359 # noqa
@@ -366,9 +394,12 @@ class InductorStandaloneAdaptor(CompilerInterface):
         else:
             fake_mode_ctx = contextlib.nullcontext()
 
-        with pregrad_ctx, fake_mode_ctx:
+        with pregrad_ctx, fake_mode_ctx, set_functorch_config(self.use_aot_compile):
             compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)
 
+        if not self.use_aot_compile:
+            return compiled_graph, None
+
         if use_aot:
             from torch._inductor.standalone_compile import AOTCompiledArtifact
 
@@ -445,8 +476,11 @@ class InductorAdaptor(CompilerInterface):
 
     name = "inductor"
 
+    def __init__(self, *, use_aot_compile: bool) -> None:
+        self.use_aot_compile = use_aot_compile
+
     def compute_hash(self, vllm_config: VllmConfig) -> str:
-        factors = get_inductor_factors()
+        factors = get_inductor_factors(self.use_aot_compile)
         hash_str: str = safe_hash(
             str(factors).encode(), usedforsecurity=False
         ).hexdigest()[:10]
@@ -455,24 +489,12 @@ class InductorAdaptor(CompilerInterface):
     def initialize_cache(
         self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
     ) -> None:
-        global _last_inductor_triton_cache_dir
         self.cache_dir = cache_dir
         self.prefix = prefix
         self.base_cache_dir = cache_dir[: -len(prefix)] if prefix else cache_dir
         if disable_cache:
             return
-        # Keep artifacts together unless the caller supplies a Triton cache.
-        inductor_cache = os.path.join(self.base_cache_dir, "inductor_cache")
-        os.makedirs(inductor_cache, exist_ok=True)
-        os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache
-        triton_cache = os.environ.get("TRITON_CACHE_DIR")
-        # A directory assigned by an earlier adaptor is still a model-local
-        # default, rather than a caller override for subsequent models.
-        if triton_cache is None or triton_cache == _last_inductor_triton_cache_dir:
-            triton_cache = os.path.join(self.base_cache_dir, "triton_cache")
-            _last_inductor_triton_cache_dir = triton_cache
-        os.makedirs(triton_cache, exist_ok=True)
-        os.environ["TRITON_CACHE_DIR"] = triton_cache
+        initialize_inductor_cache(self.base_cache_dir)
 
     def compile(
         self,
@@ -495,7 +517,6 @@ class InductorAdaptor(CompilerInterface):
         current_config["fx_graph_remote_cache"] = False
 
         set_inductor_config(current_config, compile_range)
-        set_functorch_config()
 
         # inductor can inplace modify the graph, so we need to copy it
         # see https://github.com/pytorch/pytorch/issues/138980
@@ -554,6 +575,7 @@ class InductorAdaptor(CompilerInterface):
             return AlwaysHitShapeEnv()
 
         with ExitStack() as stack:
+            stack.enter_context(set_functorch_config(self.use_aot_compile))
             # for hijacking the hash of the compiled graph
             stack.enter_context(
                 patch(
@@ -754,9 +776,17 @@ def set_inductor_config(config: dict[str, Any], compile_range: Range) -> None:
         )
 
 
-def set_functorch_config() -> None:
-    for k, v in _get_vllm_functorch_config().items():
-        setattr(torch._functorch.config, k, v)
+@contextlib.contextmanager
+def set_functorch_config(use_aot_compile: bool) -> Generator[None, None, None]:
+    overrides = _get_vllm_functorch_config(use_aot_compile)
+    if use_aot_compile:
+        # Preserve the existing native configuration lifetime.
+        for key, value in overrides.items():
+            setattr(torch._functorch.config, key, value)
+        yield
+    else:
+        with torch._functorch.config.patch(overrides):
+            yield
 
 
 def trigger_inductor_lazy_init(device: torch.device | None = None) -> None:

vllm/compilation/decorators.py

@@ -110,6 +110,7 @@ def support_torch_compile(
     mark_unbacked_dims: dict[str, int | list[int]] | None = None,
     enable_if: Callable[[VllmConfig], bool] | None = None,
     is_encoder: bool = False,
+    use_aot_compile: bool = True,
 ) -> Callable[[type[_T]], type[_T]]: ...
 
 
@@ -124,6 +125,7 @@ def support_torch_compile(
     mark_unbacked_dims: dict[str, int | list[int]] | None = None,
     enable_if: Callable[[VllmConfig], bool] | None = None,
     is_encoder: bool = False,
+    use_aot_compile: bool = True,
 ) -> Callable[[type[_T]], type[_T]] | type[_T]:
     """A decorator to add support for compiling the forward method of a class.
 
@@ -175,6 +177,10 @@ def support_torch_compile(
     the lifetime of the model, otherwise, it cannot be captured as a single
     computation graph.
 
+    `use_aot_compile=False` disables model AOT in DYNAMO_TRACE_ONCE and
+    VLLM_COMPILE modes. In VLLM_COMPILE mode, it also skips mega and direct
+    compiled artifacts while retaining the underlying PyTorch graph caches.
+
     `enable_if` is a function that takes a `VllmConfig` object as input and
     returns a boolean value indicating whether to compile the model or not.
     This is useful if you want to compile the model only when certain
@@ -243,6 +249,7 @@ def support_torch_compile(
             mark_unbacked_dims,
             enable_if,
             is_encoder,
+            use_aot_compile,
         )
 
     if cls is not None:
@@ -347,6 +354,7 @@ def _support_torch_compile(
     mark_unbacked_dims: dict[str, int | list[int]] | None = None,
     enable_if: Callable[[VllmConfig], bool] | None = None,
     is_encoder: bool = False,
+    use_aot_compile: bool = True,
 ) -> type[_T]:
     """Internal implementation of support_torch_compile decorator."""
     if TorchCompileWithNoGuardsWrapper in cls.__bases__:
@@ -396,6 +404,7 @@ def _support_torch_compile(
 
         self.vllm_config = vllm_config
         self.compilation_config = self.vllm_config.compilation_config
+        self._allow_aot_compile = use_aot_compile
         enable_compile = enable_if is None or enable_if(vllm_config)
         # for CompilationMode.STOCK_TORCH_COMPILE , the upper level model runner
         # will handle the compilation, so we don't need to do anything here.
@@ -544,7 +553,7 @@ def _support_torch_compile(
         ds_type = self.compilation_config.dynamic_shapes_config.type
         cache_dir = None
         aot_compilation_path = None
-        if envs.VLLM_USE_AOT_COMPILE:
+        if self._allow_aot_compile and envs.VLLM_USE_AOT_COMPILE:
             """
             When using torch.compile in AOT mode, we store the cache artifacts
             under VLLM_CACHE_ROOT/torch_compile_cache/torch_aot_compile/{hash}
@@ -597,7 +606,7 @@ def _support_torch_compile(
 
         if self.compiled:
             assert (
-                not envs.VLLM_USE_AOT_COMPILE
+                not (self._allow_aot_compile and envs.VLLM_USE_AOT_COMPILE)
                 or self.vllm_config.compilation_config.backend == "eager"
             )
             return TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs)  # type: ignore[arg-type]
@@ -680,7 +689,7 @@ def _support_torch_compile(
             torch.fx.experimental._config.patch(**fx_config_patches),
             torch._inductor.config.patch(**inductor_config_patches),
         ):
-            use_aot_compile = envs.VLLM_USE_AOT_COMPILE
+            use_aot_compile = self._allow_aot_compile and envs.VLLM_USE_AOT_COMPILE
             if self.vllm_config.compilation_config.backend == "eager":
                 logger.warning("Detected eager backend, disabling AOT compile.")
                 use_aot_compile = False

vllm/compilation/wrapper.py

@@ -85,6 +85,8 @@ class TorchCompileWithNoGuardsWrapper:
             return ctx.result
         return callable_fn(*args, **kwargs)
 
+    _allow_aot_compile: bool = True
+
     def __init__(
         self,
         compile_prefix: str = "",
@@ -104,7 +106,10 @@ class TorchCompileWithNoGuardsWrapper:
             raise RuntimeError("Compilation mode cannot be NO_COMPILATION")
 
         backend = vllm_config.compilation_config.init_backend(
-            vllm_config, prefix=compile_prefix, is_encoder=is_encoder
+            vllm_config,
+            prefix=compile_prefix,
+            is_encoder=is_encoder,
+            use_aot_compile=self._allow_aot_compile,
         )
         options = {}
 
@@ -151,7 +156,7 @@ class TorchCompileWithNoGuardsWrapper:
         _apply_constrain_to_fx_strides_patch()
 
         aot_context = nullcontext()
-        if envs.VLLM_USE_AOT_COMPILE:
+        if self._allow_aot_compile and envs.VLLM_USE_AOT_COMPILE:
             if hasattr(torch._dynamo.config, "enable_aot_compile"):
                 aot_context = torch._dynamo.config.patch(enable_aot_compile=True)
             else:

vllm/config/compilation.py

@@ -1085,6 +1085,7 @@ class CompilationConfig:
         vllm_config: "VllmConfig",
         prefix: str = "",
         is_encoder: bool = False,
+        use_aot_compile: bool = True,
     ) -> str | Callable:
         """Initialize the backend for the compilation config from a vllm config.
 
@@ -1093,6 +1094,7 @@ class CompilationConfig:
             prefix: Cache directory prefix for this compiled module.
             is_encoder: Whether this module is used in an encoder (as
                 opposed to a text backbone).
+            use_aot_compile: Whether the module permits AOT and direct artifacts.
 
         Returns:
             The backend for the compilation config.
@@ -1124,7 +1126,12 @@ class CompilationConfig:
 
         from vllm.compilation.backends import VllmBackend
 
-        return VllmBackend(vllm_config, prefix=prefix, is_encoder=is_encoder)
+        return VllmBackend(
+            vllm_config,
+            prefix=prefix,
+            is_encoder=is_encoder,
+            use_aot_compile=use_aot_compile,
+        )
 
     def post_init_cudagraph_sizes(self) -> None:
         """To complete the initialization after cudagraph related

vllm/model_executor/models/transformers/base.py

@@ -282,6 +282,7 @@ class Base(
             dynamic_arg_dims=dynamic_arg_dims,
             enable_if=enable_if,
             is_encoder=is_encoder,
+            use_aot_compile=False,
         )(cls)
 
     def _decorate_for_torch_compile(self):