@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 182 lines 5.94 KB
Raw
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(*, vllm_version: str = "0.0.0", **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 ``vllm_version`` selects the build type: importing.py takes its Triton-CPU
18 branch only when "cpu" is in the installed vLLM version, so pin it rather
19 than depending on the real install (the CPU CI image is a "+cpu" build).
20 """
21 backends = {}
22 for name, is_active in drivers.items():
23 driver = mock.Mock()
24 driver.is_active.return_value = is_active
25 backend = mock.Mock()
26 backend.driver = driver
27 backends[name] = backend
28
29 triton_mod = types.ModuleType("triton")
30 triton_mod.__spec__ = importlib.machinery.ModuleSpec("triton", None)
31 backends_mod = types.ModuleType("triton.backends")
32 backends_mod.__dict__["backends"] = backends
33 triton_mod.__dict__["backends"] = backends_mod
34
35 patched_modules = {"triton": triton_mod, "triton.backends": backends_mod}
36 try:
37 with (
38 mock.patch.dict(sys.modules, patched_modules),
39 mock.patch.dict("os.environ", {}, clear=True),
40 mock.patch("importlib.metadata.version", return_value=vllm_version),
41 ):
42 return importlib.reload(triton_importing).HAS_TRITON
43 finally:
44 # Restore the module state derived from the real environment.
45 importlib.reload(triton_importing)
46
47
48 def test_triton_placeholder_is_module():
49 triton = TritonPlaceholder()
50 assert isinstance(triton, types.ModuleType)
51 assert triton.__name__ == "triton"
52
53
54 def test_triton_language_placeholder_is_module():
55 triton_language = TritonLanguagePlaceholder()
56 assert isinstance(triton_language, types.ModuleType)
57 assert triton_language.__name__ == "triton.language"
58
59
60 def test_triton_placeholder_decorators():
61 triton = TritonPlaceholder()
62
63 @triton.jit
64 def foo(x):
65 return x
66
67 @triton.autotune
68 def bar(x):
69 return x
70
71 @triton.heuristics
72 def baz(x):
73 return x
74
75 assert foo(1) == 1
76 assert bar(2) == 2
77 assert baz(3) == 3
78
79
80 def test_triton_placeholder_decorators_with_args():
81 triton = TritonPlaceholder()
82
83 @triton.jit(debug=True)
84 def foo(x):
85 return x
86
87 @triton.autotune(configs=[], key="x")
88 def bar(x):
89 return x
90
91 @triton.heuristics({"BLOCK_SIZE": lambda args: 128 if args["x"] > 1024 else 64})
92 def baz(x):
93 return x
94
95 assert foo(1) == 1
96 assert bar(2) == 2
97 assert baz(3) == 3
98
99
100 def test_triton_placeholder_language():
101 lang = TritonLanguagePlaceholder()
102 assert isinstance(lang, types.ModuleType)
103 assert lang.__name__ == "triton.language"
104 assert lang.constexpr(2**31 - 1) == 2**31 - 1
105 assert lang.constexpr(1.5) == 1.5
106 assert lang.dtype is None
107 assert lang.int64 is None
108 assert lang.int32 is None
109 assert lang.tensor is None
110
111
112 def test_triton_placeholder_language_from_parent():
113 triton = TritonPlaceholder()
114 lang = triton.language
115 assert isinstance(lang, TritonLanguagePlaceholder)
116
117
118 def test_cpu_backend_does_not_disable_triton():
119 # The cpu backend's driver is always active, so counting it alongside a GPU
120 # backend used to yield 2 active drivers and disable Triton entirely.
121 assert _has_triton_for_backends(amd=True, cpu=True) is True
122
123
124 def test_single_gpu_backend_keeps_triton():
125 assert _has_triton_for_backends(amd=True) is True
126
127
128 def test_multiple_active_gpu_backends_disable_triton():
129 assert _has_triton_for_backends(amd=True, nvidia=True) is False
130
131
132 def test_cpu_backend_alone_disables_triton():
133 assert _has_triton_for_backends(cpu=True) is False
134
135
136 def test_cpu_build_with_cpu_backend_keeps_triton():
137 assert _has_triton_for_backends(vllm_version="0.0.0+cpu", cpu=True) is True
138
139
140 def test_cpu_build_without_cpu_backend_disables_triton():
141 assert _has_triton_for_backends(vllm_version="0.0.0+cpu", amd=True) is False
142
143
144 def test_no_triton_fallback():
145 # Save the real modules so they can be put back afterward - popping them
146 # (rather than mock.patch.dict alone) is what forces vllm.triton_utils to
147 # be freshly, unmemoized re-imported under the "triton is absent"
148 # condition below.
149 saved_modules = {
150 name: sys.modules.get(name)
151 for name in (
152 "triton",
153 "triton.language",
154 "vllm.triton_utils",
155 "vllm.triton_utils.importing",
156 )
157 }
158 for name in saved_modules:
159 sys.modules.pop(name, None)
160
161 try:
162 # mock triton not being installed
163 with mock.patch.dict(sys.modules, {"triton": None}):
164 from vllm.triton_utils import HAS_TRITON, tl, triton
165
166 assert HAS_TRITON is False
167 assert triton.__class__.__name__ == "TritonPlaceholder"
168 assert triton.language.__class__.__name__ == "TritonLanguagePlaceholder"
169 assert tl.__class__.__name__ == "TritonLanguagePlaceholder"
170 assert tl.constexpr(2**31 - 1) == 2**31 - 1
171 finally:
172 # The pops above are outside mock.patch.dict's scope, so exiting the
173 # `with` restores "triton" to absent, not to the real module - unlike
174 # _has_triton_for_backends above, nothing here puts it back. Left
175 # alone, every subsequent import of vllm.triton_utils in this process
176 # (including other test files, and other tests' cleanup fixtures that
177 # need real triton) would get this fake, HAS_TRITON=False module.
178 for name, module in saved_modules.items():
179 if module is not None:
180 sys.modules[name] = module
181 else:
182 sys.modules.pop(name, None)