A high-throughput and memory-efficient inference and serving engine for LLMs
[Perf] Add VLLM_TRITON_FORCE_FIRST_CONFIG to skip Triton autotuning (#42425)
Signed-off-by: Francesco Fusco <ffu@zurich.ibm.com> Co-authored-by: Claude <noreply@anthropic.com>
Francesco Fusco committed
Jun 16, 2026 at 15:16 UTC
ced32bb474feaed7e02723544013298fb1401b31
4 files changed
+209
tests/test_force_first_config.py
new
+93
@@ -0,0 +1,93 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Targeted unit tests for VLLM_TRITON_FORCE_FIRST_CONFIG.
4
+
5
+These tests exercise only the patched `Autotuner.run` logic installed by
6
+`vllm.triton_utils.force_first_config.install`. The wrapped kernel is a
7
+plain callable so the tests run on CPU-only hosts (no GPU, no actual
8
+kernel launch) as long as the `triton` package is importable.
9
+"""
10
+
11
+from types import SimpleNamespace
12
+
13
+import pytest
14
+
15
+from vllm.triton_utils import HAS_TRITON, triton
16
+
17
+if not HAS_TRITON:
18
+ pytest.skip("triton not available", allow_module_level=True)
19
+
20
+from vllm.triton_utils import force_first_config # noqa: E402
21
+
22
+OutOfResources = triton.runtime.errors.OutOfResources
23
+
24
+
25
+@pytest.fixture
26
+def patched_autotuner(monkeypatch: pytest.MonkeyPatch):
27
+ """Install the first-valid-config patch and restore after.
28
+
29
+ The env-var gate lives in vllm.env_override; install() itself does not
30
+ read the environment, so the test calls it directly.
31
+ """
32
+ Autotuner = triton.runtime.autotuner.Autotuner
33
+ original_run = Autotuner.run
34
+ # Reset the once-only guard so install() re-runs for each test.
35
+ monkeypatch.setattr(force_first_config, "_installed", False)
36
+ force_first_config.install()
37
+ yield Autotuner
38
+ Autotuner.run = original_run
39
+
40
+
41
+def _make_fake_self(configs, fn):
42
+ """Minimal stand-in for an Autotuner instance."""
43
+ return SimpleNamespace(
44
+ configs=configs,
45
+ keys=[],
46
+ arg_names=[],
47
+ base_fn=fn,
48
+ fn=fn,
49
+ best_config=None,
50
+ )
51
+
52
+
53
+def test_skips_invalid_first_config_and_caches_second(patched_autotuner):
54
+ bad = triton.Config({"BLOCK": 1024})
55
+ good = triton.Config({"BLOCK": 64})
56
+ calls = []
57
+
58
+ def fake_fn(*args, **kwargs):
59
+ calls.append(kwargs["BLOCK"])
60
+ if kwargs["BLOCK"] == 1024:
61
+ raise OutOfResources(required=99999, limit=1, name="shared memory")
62
+ return "ok"
63
+
64
+ fake_self = _make_fake_self([bad, good], fake_fn)
65
+
66
+ # First call: walks past the invalid config, picks the second.
67
+ assert patched_autotuner.run(fake_self) == "ok"
68
+ assert calls == [1024, 64]
69
+ assert fake_self.best_config is good
70
+
71
+ # Second call: cached index is reused, invalid config is NOT retried.
72
+ calls.clear()
73
+ assert patched_autotuner.run(fake_self) == "ok"
74
+ assert calls == [64]
75
+
76
+
77
+def test_empty_configs_falls_back_to_direct_call(patched_autotuner):
78
+ def fake_fn(*args, **kwargs):
79
+ return "direct"
80
+
81
+ fake_self = _make_fake_self([], fake_fn)
82
+ assert patched_autotuner.run(fake_self) == "direct"
83
+
84
+
85
+def test_all_configs_invalid_raises_runtime_error(patched_autotuner):
86
+ cfgs = [triton.Config({"BLOCK": 1024}), triton.Config({"BLOCK": 2048})]
87
+
88
+ def always_oor(*args, **kwargs):
89
+ raise OutOfResources(required=99999, limit=1, name="shared memory")
90
+
91
+ fake_self = _make_fake_self(cfgs, always_oor)
92
+ with pytest.raises(RuntimeError, match="[Nn]o valid config"):
93
+ patched_autotuner.run(fake_self)
vllm/env_override.py
+15
@@ -856,3 +856,18 @@ def _patch_inductor_fallback_allow_list() -> None:
856
857
858
_patch_inductor_fallback_allow_list()
859
+
860
+# ============================================================
861
+# Triton Autotuner determinism
862
+# ============================================================
863
+# Replace the Autotuner.run so it always pick the first running configuration.
864
+# Useful to eliminate autotune variability leading to non determinism.
865
+if os.environ.get("VLLM_TRITON_FORCE_FIRST_CONFIG", "0").strip().lower() in (
866
+ "1",
867
+ "true",
868
+):
869
+ from vllm.triton_utils.force_first_config import (
870
+ install as _install_force_first_config,
871
+ )
872
+
873
+ _install_force_first_config()
vllm/envs.py
+9
@@ -108,6 +108,7 @@ if TYPE_CHECKING:
108
VLLM_USE_MEGA_AOT_ARTIFACT: bool = False
109
VLLM_USE_TRITON_AWQ: bool = False
110
VLLM_FASTSAFETENSORS_QUEUE_SIZE: int = 0
111
+ VLLM_TRITON_FORCE_FIRST_CONFIG: bool = False
112
VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
113
VLLM_SKIP_P2P_CHECK: bool = False
114
VLLM_DISABLED_KERNELS: list[str] = []
@@ -1072,6 +1073,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
1073
),
1074
# If set, vLLM will use Triton implementations of AWQ.
1075
"VLLM_USE_TRITON_AWQ": lambda: bool(int(os.getenv("VLLM_USE_TRITON_AWQ", "0"))),
1076
+ # If set, monkey-patch triton.runtime.autotuner.Autotuner.run to skip
1077
+ # benchmarking and select the first valid config (walking past invalid
1078
+ # ones). Used to eliminate autotuning variability when measuring kernel
1079
+ # performance and applied before running any kernel.
1080
+ "VLLM_TRITON_FORCE_FIRST_CONFIG": lambda: (
1081
+ os.environ.get("VLLM_TRITON_FORCE_FIRST_CONFIG", "0").strip().lower()
1082
+ in ("1", "true")
1083
+ ),
1084
# If set, allow loading or unloading lora adapters in runtime,
1085
"VLLM_ALLOW_RUNTIME_LORA_UPDATING": lambda: (
1086
os.environ.get("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "0").strip().lower()
vllm/triton_utils/force_first_config.py
new
+92
@@ -0,0 +1,92 @@
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