@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 222 lines 7.32 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the apply-time Triton kernel override registry.
4
5 The registry must let out-of-tree platforms replace kernels defined as
6 plain ``@triton.jit`` functions, patching every hold-site (defining
7 module, from-import copies, and JIT warmup owners) so that existing call
8 sites dispatch to the override without any core-side decoration.
9 """
10
11 import sys
12 import types
13 from typing import Any
14
15 import pytest
16
17 from vllm.model_executor.warmup.jit_warmup_triton_helper import (
18 _DecoratedTritonJitKernel,
19 )
20 from vllm.triton_utils import tl, triton
21 from vllm.triton_utils.dispatcher import (
22 KernelOverride,
23 register_kernels,
24 )
25
26
27 class _FakeTritonKernel:
28 """Stands in for a Triton JITFunction (works without Triton installed)."""
29
30 arg_names = ("output_ptr", "input_ptr", "replace_from", "MAX_NUM_TOKENS")
31
32 def __getitem__(self, grid: Any) -> Any:
33 def launch(*args: Any, **kwargs: Any) -> Any:
34 return ("default", grid, kwargs)
35
36 return launch
37
38
39 @pytest.fixture
40 def _fake_module(monkeypatch: pytest.MonkeyPatch) -> Any:
41 """Expose a fake defining module holding the kernel and its warmup owner."""
42 import types
43
44 kernel = _FakeTritonKernel()
45 owner = _DecoratedTritonJitKernel(
46 kernel,
47 lambda: {},
48 lambda output, input, replace_from: ( # noqa: ARG005
49 (1,),
50 dict(MAX_NUM_TOKENS=8),
51 ),
52 )
53 module = types.ModuleType("vllm.tests.fake_kernel_home")
54 module.expand_kernel = kernel # type: ignore[attr-defined]
55 module._expand = owner # type: ignore[attr-defined]
56 monkeypatch.setitem(sys.modules, module.__name__, module)
57 return module
58
59
60 def test_override_patches_defining_module_and_warmup_owner(
61 _fake_module: Any,
62 ) -> None:
63 calls: list[tuple[Any, ...]] = []
64
65 def cpu_impl(output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS=None):
66 calls.append((output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS))
67 return "cpu"
68
69 register_kernels(
70 {
71 "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl,
72 }
73 )
74
75 overridden = _fake_module.expand_kernel
76 assert isinstance(overridden, KernelOverride)
77 # Launch with kernel-arg names like the warmup launch path does.
78 assert overridden[(4,)](output_ptr=1, input_ptr=2, replace_from=0) == "cpu"
79 assert calls == [(1, 2, 0, None)]
80
81 # The warmup owner's captured kernel is patched too, and its cached
82 # argument-name introspection is invalidated.
83 assert _fake_module._expand.kernel is overridden
84
85
86 def test_override_forwards_positionally_for_mismatched_names(
87 _fake_module: Any,
88 ) -> None:
89 calls: list[tuple[Any, ...]] = []
90
91 def cpu_impl(output, input_val, replace_from, MAX_NUM_TOKENS=None):
92 calls.append((output, input_val, replace_from, MAX_NUM_TOKENS))
93 return "cpu"
94
95 register_kernels(
96 {
97 "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl,
98 }
99 )
100 overridden = _fake_module.expand_kernel
101 assert overridden[(4,)](output_ptr=1, input_ptr=2, replace_from=0) == "cpu"
102 assert calls == [(1, 2, 0, None)]
103
104
105 def test_override_covers_from_import_copies(_fake_module: Any) -> None:
106 import types
107
108 importer = types.ModuleType("vllm.tests.fake_kernel_importer")
109 importer.expand_kernel = _fake_module.expand_kernel # type: ignore[attr-defined]
110 sys.modules["vllm.tests.fake_kernel_importer"] = importer
111
112 def cpu_impl(output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS=None):
113 return "cpu"
114
115 register_kernels(
116 {
117 "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl,
118 }
119 )
120 assert isinstance(importer.expand_kernel, KernelOverride)
121 del sys.modules["vllm.tests.fake_kernel_importer"]
122
123
124 def test_override_owner_launch_goes_through_impl(_fake_module: Any) -> None:
125 launched: list[Any] = []
126
127 def cpu_impl(output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS=None):
128 launched.append((output_ptr, input_ptr, replace_from, MAX_NUM_TOKENS))
129 return "cpu"
130
131 register_kernels(
132 {
133 "vllm.tests.fake_kernel_home.expand_kernel": cpu_impl,
134 }
135 )
136 owner = _fake_module._expand
137 # In-place kernels return None from launch; assert on the recorded call.
138 owner(1, 2, 0)
139 assert launched == [(1, 2, 0, 8)]
140
141
142 def test_register_kernel_auto_imports_defining_module(
143 monkeypatch: pytest.MonkeyPatch,
144 ) -> None:
145 """Registering a kernel whose defining module is not imported yet must
146 import it on the spot and apply the override."""
147 kernel_name = "vllm.v1.sample.rejection_sampler.expand_kernel"
148 module_name = kernel_name.rpartition(".")[0]
149 calls: list[tuple[Any, ...]] = []
150
151 def cpu_impl(
152 output_ptr,
153 input_ptr,
154 cu_num_tokens_ptr,
155 replace_from,
156 replace_to,
157 MAX_NUM_TOKENS=None,
158 ):
159 calls.append((output_ptr, replace_from, MAX_NUM_TOKENS))
160 return "cpu"
161
162 # Ensure the module is NOT imported before registration. Drop it from
163 # sys.modules; a fresh import creates new kernel/owner objects, which
164 # is what a cold-start platform process would see.
165 saved = sys.modules.pop(module_name, None)
166 # Re-run the registry patch against the original module afterwards so
167 # other tests that already imported it keep working.
168 try:
169 assert module_name not in sys.modules
170 register_kernels({kernel_name: cpu_impl})
171 assert module_name in sys.modules
172 finally:
173 if saved is None:
174 # First import in the whole session: keep the fresh module.
175 saved = sys.modules[module_name]
176 else:
177 # Restore the original module object and re-apply the override
178 # to its (pre-existing) kernel so both generations agree.
179 sys.modules[module_name] = saved
180 register_kernels({kernel_name: cpu_impl})
181
182 overridden = saved.expand_kernel
183 assert isinstance(overridden, KernelOverride)
184 overridden[(2,)](
185 output_ptr=1,
186 input_ptr=2,
187 cu_num_tokens_ptr=3,
188 replace_from=0,
189 replace_to=0,
190 MAX_NUM_TOKENS=8,
191 )
192 assert calls == [(1, 0, 8)]
193 # The warmup owner captured at the fresh import must be patched too
194 # (when the module existed before, there are two module generations).
195 fresh = sys.modules[module_name]
196 if fresh is not saved:
197 assert fresh._expand.kernel is overridden
198
199
200 def test_real_triton_kernel_override_end_to_end() -> None:
201 """Override a real @triton.jit kernel through the registry."""
202 if not triton.__class__.__module__.startswith("triton"):
203 pytest.skip("Triton is not installed")
204
205 @triton.jit
206 def _add_one_kernel(x_ptr, out_ptr, BLOCK: tl.constexpr): # pragma: no cover
207 offsets = tl.arange(0, BLOCK)
208 tl.store(out_ptr + offsets, tl.load(x_ptr + offsets) + 1)
209
210 module = types.ModuleType("vllm.tests.real_jit_kernel_home")
211 module._add_one_kernel = _add_one_kernel # type: ignore[attr-defined]
212 sys.modules[module.__name__] = module
213
214 def cpu_impl(x_ptr, out_ptr, BLOCK=None): # noqa: ARG001
215 return "cpu"
216
217 register_kernels({"vllm.tests.real_jit_kernel_home._add_one_kernel": cpu_impl})
218 assert module._add_one_kernel.arg_names == ( # type: ignore[attr-defined]
219 "x_ptr",
220 "out_ptr",
221 "BLOCK",
222 )