Exclude the cpu backend from vLLM's active-Triton-driver count (#53530) (#53530)
Signed-off-by: Minjang Kim <minjang@meta.com> Co-authored-by: Autopilot Bot <noreply+1608173377072046@fb.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Minjang Kim committed
Aug 24, 2026 at 21:13 UTC
af119619c4ca2090e0fac310867199f8fadec3a6
2 files changed
+59
-1
tests/test_triton_utils.py
+51
@@ -1,13 +1,46 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import importlib
5
import sys
6
import types
7
from unittest import mock
8
9
+from vllm.triton_utils import importing as triton_importing
10
from vllm.triton_utils.importing import TritonLanguagePlaceholder, TritonPlaceholder
11
12
13
+def _has_triton_for_backends(**drivers: bool) -> bool:
14
+ """Re-evaluate ``HAS_TRITON`` against a synthetic ``triton.backends`` map.
15
+
16
+ ``drivers`` maps backend name to whether its driver reports itself active.
17
+ """
18
+ backends = {}
19
+ for name, is_active in drivers.items():
20
+ driver = mock.Mock()
21
+ driver.is_active.return_value = is_active
22
+ backend = mock.Mock()
23
+ backend.driver = driver
24
+ backends[name] = backend
25
+
26
+ triton_mod = types.ModuleType("triton")
27
+ triton_mod.__spec__ = importlib.machinery.ModuleSpec("triton", None)
28
+ backends_mod = types.ModuleType("triton.backends")
29
+ backends_mod.__dict__["backends"] = backends
30
+ triton_mod.__dict__["backends"] = backends_mod
31
+
32
+ patched_modules = {"triton": triton_mod, "triton.backends": backends_mod}
33
+ try:
34
+ with (
35
+ mock.patch.dict(sys.modules, patched_modules),
36
+ mock.patch.dict("os.environ", {}, clear=True),
37
+ ):
38
+ return importlib.reload(triton_importing).HAS_TRITON
39
+ finally:
40
+ # Restore the module state derived from the real environment.
41
+ importlib.reload(triton_importing)
42
+
43
+
44
def test_triton_placeholder_is_module():
45
triton = TritonPlaceholder()
46
assert isinstance(triton, types.ModuleType)
@@ -77,6 +110,24 @@ def test_triton_placeholder_language_from_parent():
110
assert isinstance(lang, TritonLanguagePlaceholder)
111
112
113
+def test_cpu_backend_does_not_disable_triton():
114
+ # The cpu backend's driver is always active, so counting it alongside a GPU
115
+ # backend used to yield 2 active drivers and disable Triton entirely.
116
+ assert _has_triton_for_backends(amd=True, cpu=True) is True
117
+
118
+
119
+def test_single_gpu_backend_keeps_triton():
120
+ assert _has_triton_for_backends(amd=True) is True
121
+
122
+
123
+def test_multiple_active_gpu_backends_disable_triton():
124
+ assert _has_triton_for_backends(amd=True, nvidia=True) is False
125
+
126
+
127
+def test_cpu_backend_alone_disables_triton():
128
+ assert _has_triton_for_backends(cpu=True) is False
129
+
130
+
131
def test_no_triton_fallback():
132
# clear existing triton modules
133
sys.modules.pop("triton", None)
vllm/triton_utils/importing.py
+8
-1
@@ -23,8 +23,15 @@ if HAS_TRITON:
23
# It's generally expected that x.driver exists and has
24
# an is_active method.
25
# The `x.driver and` check adds a small layer of safety.
26
+ # The "cpu" backend's driver reports itself active unconditionally, so
27
+ # it must be excluded here or every GPU host looks like it has 2 active
28
+ # drivers. Triton's own driver selection skips it for the same reason
29
+ # (see triton.runtime.driver._create_driver): CPU is only ever used
30
+ # when explicitly selected.
31
active_drivers = [
27
- x.driver for x in backends.values() if x.driver and x.driver.is_active()
32
+ x.driver
33
+ for name, x in backends.items()
34
+ if name != "cpu" and x.driver and x.driver.is_active()
35
]
36
37
# Check if we're in a distributed environment where CUDA_VISIBLE_DEVICES