| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Skip Triton autotuning under VLLM_TRITON_FORCE_FIRST_CONFIG.""" |
| 4 | |
| 5 | from vllm.logger import init_logger |
| 6 | from vllm.triton_utils.importing import HAS_TRITON |
| 7 | |
| 8 | logger = init_logger(__name__) |
| 9 | |
| 10 | _installed: bool = False |
| 11 | |
| 12 | |
| 13 | def is_installed() -> bool: |
| 14 | """Return whether the first-valid-config patch is currently installed.""" |
| 15 | return _installed |
| 16 | |
| 17 | |
| 18 | def install() -> None: |
| 19 | """Install the Autotuner.run replacement.""" |
| 20 | global _installed |
| 21 | if _installed: |
| 22 | return |
| 23 | if not HAS_TRITON: |
| 24 | return |
| 25 | |
| 26 | import importlib |
| 27 | |
| 28 | autotuner_mod = importlib.import_module("triton.runtime.autotuner") |
| 29 | Autotuner = autotuner_mod.Autotuner |
| 30 | from triton.compiler.errors import CompileTimeAssertionFailure |
| 31 | from triton.runtime.errors import OutOfResources, PTXASError |
| 32 | |
| 33 | _invalid_config_errors = (OutOfResources, CompileTimeAssertionFailure, PTXASError) |
| 34 | _picked_cache: dict[tuple, int] = {} |
| 35 | seen_kernels: set[str] = set() |
| 36 | |
| 37 | def _run_first_valid_config(self, *args, **kwargs): |
| 38 | if not self.configs: |
| 39 | return self.fn(*args, **kwargs) |
| 40 | |
| 41 | key_vals = tuple(kwargs[name] for name in self.keys if name in kwargs) |
| 42 | cache_key = (id(self), key_vals) |
| 43 | kernel_name = getattr(self.base_fn, "__name__", repr(self.fn)) |
| 44 | |
| 45 | cached_idx = _picked_cache.get(cache_key) |
| 46 | candidate_indices = ( |
| 47 | [cached_idx] if cached_idx is not None else list(range(len(self.configs))) |
| 48 | ) |
| 49 | |
| 50 | last_exc: Exception | None = None |
| 51 | for idx in candidate_indices: |
| 52 | config = self.configs[idx] |
| 53 | if config.pre_hook is not None: |
| 54 | full_nargs = { |
| 55 | **dict(zip(self.arg_names, args)), |
| 56 | **kwargs, |
| 57 | **config.all_kwargs(), |
| 58 | } |
| 59 | config.pre_hook(full_nargs) |
| 60 | # Prefer self.fn.run(...) — the kernel-launch entrypoint for both |
| 61 | # JITFunction and Heuristics. Calling JITFunction(...) directly |
| 62 | # raises "Cannot call @triton.jit'd outside of the scope of a |
| 63 | # kernel". Fall back to plain call only if .run is missing. |
| 64 | launch = getattr(self.fn, "run", self.fn) |
| 65 | try: |
| 66 | result = launch(*args, **kwargs, **config.all_kwargs()) |
| 67 | except _invalid_config_errors as e: |
| 68 | last_exc = e |
| 69 | continue |
| 70 | |
| 71 | if cached_idx is None: |
| 72 | _picked_cache[cache_key] = idx |
| 73 | self.best_config = config |
| 74 | if kernel_name not in seen_kernels: |
| 75 | seen_kernels.add(kernel_name) |
| 76 | logger.info( |
| 77 | "[triton-autotune-disabled] kernel=%s configs=%d " |
| 78 | "picked_index=%d picked=%s", |
| 79 | kernel_name, |
| 80 | len(self.configs), |
| 81 | idx, |
| 82 | config, |
| 83 | ) |
| 84 | return result |
| 85 | |
| 86 | raise RuntimeError( |
| 87 | f"No valid config for kernel " |
| 88 | f"{kernel_name} key={key_vals} (tried {len(self.configs)} configs)" |
| 89 | ) from last_exc |
| 90 | |
| 91 | Autotuner.run = _run_first_valid_config |
| 92 | _installed = True |