From 153f2ba1b986fb461d84ae7b4f9644a608f36bc8 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Thu, 24 Sep 2026 01:27:07 -0500 Subject: [PATCH] [Compile][CI] Honor Triton cache overrides and add AMD timeout headroom (#58474) Signed-off-by: Andreas Karatzas Co-authored-by: Codex --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/model_executor.yaml | 2 +- docs/design/torch_compile.md | 6 ++++++ vllm/compilation/compiler_interface.py | 15 ++++++++++----- 4 files changed, 18 insertions(+), 7 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index a349f8a58e..771a177298 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1995,7 +1995,7 @@ steps: #------------------------------------------------------ mi300 ยท model_executor -------------------------------------------------------# - label: ":amd: (MI300) Model Executor" # TBD - timeout_in_minutes: 60 + timeout_in_minutes: 75 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] dind: false agent_pool: mi300_1 diff --git a/.buildkite/test_areas/model_executor.yaml b/.buildkite/test_areas/model_executor.yaml index 17818cc152..f6905f98ef 100644 --- a/.buildkite/test_areas/model_executor.yaml +++ b/.buildkite/test_areas/model_executor.yaml @@ -33,7 +33,7 @@ steps: label: ":amd: (MI300) Model Executor" dind: false device: mi300_1 - timeout_in_minutes: 60 + timeout_in_minutes: 75 depends_on: - image-build-amd source_file_dependencies: diff --git a/docs/design/torch_compile.md b/docs/design/torch_compile.md index 4dc0da0c7d..1ffc0ac53a 100644 --- a/docs/design/torch_compile.md +++ b/docs/design/torch_compile.md @@ -27,6 +27,12 @@ With all these factors taken into consideration, usually we can guarantee that t A unique aspect of vLLM's `torch.compile` integration, is that we guarantee all the compilation finishes before we serve any requests. No requests will trigger new compilations. Otherwise, the engine would be blocked on that request, and the response time will have unexpected spikes. +An explicit `TRITON_CACHE_DIR` is preserved, allowing processes and model +configurations to reuse Triton artifacts. The directory can start empty: the +first compilation populates it and later matching specializations reuse it. +When copying compilation caches, also copy this directory if set. Cold-cache +tests should set it to a fresh directory. + By default, the cache saves compiled artifacts as binary files. If you would like to interact with the generated code for debugging purposes, set the field `compile_cache_save_format=unpacked` in the compilation config, or omit this and set the env variable `VLLM_COMPILE_CACHE_SAVE_FORMAT=unpacked`. ## Dynamic shapes and vllm guard dropping diff --git a/vllm/compilation/compiler_interface.py b/vllm/compilation/compiler_interface.py index 67aff0fe1e..394799ab87 100644 --- a/vllm/compilation/compiler_interface.py +++ b/vllm/compilation/compiler_interface.py @@ -23,6 +23,8 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer logger = init_logger(__name__) +_last_inductor_triton_cache_dir: str | None = None + class CompilerInterface: """The interface for a compiler that can be used by vLLM.""" @@ -453,19 +455,22 @@ 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 - # redirect the cache directory to a subdirectory - # set flags so that Inductor and Triton store their cache - # in the cache_dir, then users only need to copy the cache_dir - # to another machine to reuse the cache. + # 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.path.join(self.base_cache_dir, "triton_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