[Compile][CI] Honor Triton cache overrides and add AMD timeout headroom (#58474)

Signed-off-by: Andreas Karatzas <andreas.karatzas@protonmail.com>
Co-authored-by: Codex <noreply@openai.com>

.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

.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:

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

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