From c068db0e1d8e302219b9bc4bec72474af2de5f8a Mon Sep 17 00:00:00 2001 From: Taimys <96967844+Taimys@users.noreply.github.com> Date: Wed, 14 Oct 2026 03:17:29 +0800 Subject: [PATCH] [Bugfix][Compile] Replace Transformers AOT caching with graph caching (#58918) Signed-off-by: fangyao fan <2281693949@qq.com> Signed-off-by: Taimys Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Taimys Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> --- tests/compile/test_aot_compile.py | 177 ++++++++++++++++++ tests/compile/test_compiler_interface.py | 85 ++++++++- .../test_rocm_attention_backends_selection.py | 4 +- vllm/compilation/backends.py | 53 ++++-- vllm/compilation/compiler_interface.py | 88 ++++++--- vllm/compilation/decorators.py | 15 +- vllm/compilation/wrapper.py | 9 +- vllm/config/compilation.py | 9 +- .../models/transformers/base.py | 1 + 9 files changed, 387 insertions(+), 54 deletions(-) diff --git a/tests/compile/test_aot_compile.py b/tests/compile/test_aot_compile.py index e5c11b8acc..bbb9edb371 100644 --- a/tests/compile/test_aot_compile.py +++ b/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 diff --git a/tests/compile/test_compiler_interface.py b/tests/compile/test_compiler_interface.py index a54f2a1056..7f7eb79fe3 100644 --- a/tests/compile/test_compiler_interface.py +++ b/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" diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 45580b395f..f103b43650 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/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) diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 434c1af47e..d45be552a7 100644 --- a/vllm/compilation/backends.py +++ b/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 diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py index a297b0ed95..7c2e05cc6b 100644 --- a/vllm/compilation/compiler_interface.py +++ b/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: diff --git a/vllm/compilation/decorators.py b/vllm/compilation/decorators.py index 3aad32cc24..78e166466f 100644 --- a/vllm/compilation/decorators.py +++ b/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 diff --git a/vllm/compilation/wrapper.py b/vllm/compilation/wrapper.py index 78b34f8959..2f8c7ce2ba 100644 --- a/vllm/compilation/wrapper.py +++ b/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: diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index 7a16e3d6a0..395b6e2c8f 100644 --- a/vllm/config/compilation.py +++ b/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 diff --git a/vllm/model_executor/models/transformers/base.py b/vllm/model_executor/models/transformers/base.py index 9f045173f3..6f4cbdc923 100644 --- a/vllm/model_executor/models/transformers/base.py +++ b/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):