@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 864a87febb

[Kernel][CI] `--jit-monitor-mode error` e2e tests for kernel warmup infra (#50109)

Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>

Nicolò Lucchesi committed Jul 31, 2026 at 14:53 UTC 864a87febb5b4f3fe15b50520a6857b241d5bf63
10 files changed +716 -539
.buildkite/intel_jobs/engine_intel.yaml
+2 -2
@@ -33,12 +33,12 @@ steps:
33 - tests/test_config
34 - tests/test_logger
35 - tests/test_vllm_port
36 - - tests/test_jit_monitor.py
36 + - tests/jit_monitor/test_hooks.py
37 commands:
38 - >-
39 bash .buildkite/scripts/hardware_ci/run-intel-test.sh
40 'cd tests &&
41 - pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py test_jit_monitor.py'
41 + pytest -v -s engine/test_arg_utils.py test_sequence.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py'
42
43 - label: Engine (1 GPU)
44 timeout_in_minutes: 30
.buildkite/test-amd.yaml
+1 -1
@@ -971,7 +971,7 @@ steps:
971 - tests/test_logger
972 - tests/test_vllm_port
973 commands:
974 - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
974 + - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py
975
976 #-------------------------------------------------------- mi300 · entrypoints --------------------------------------------------------#
977
.buildkite/test_areas/engine.yaml
+3 -2
@@ -23,9 +23,10 @@ steps:
23 - tests/test_config
24 - tests/test_logger
25 - tests/test_vllm_port
26 - - tests/test_jit_monitor.py
26 + - tests/jit_monitor/test_hooks.py
27 + - tests/jit_monitor/test_hooks_gpu.py
28 commands:
28 - - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py test_jit_monitor.py
29 + - pytest -v -s engine test_sequence.py test_config.py test_logger.py test_vllm_port.py jit_monitor/test_hooks.py jit_monitor/test_hooks_gpu.py
30 mirror:
31 amd:
32 dind: false
.buildkite/test_areas/jit_monitor.yaml new
+22
@@ -0,0 +1,22 @@
1 +group: JIT Monitor
2 +depends_on:
3 + - image-build
4 +steps:
5 +- label: No Runtime JITs e2e tests
6 + key: jit-monitor-no-runtime-jit
7 + device: h200_35gb
8 + timeout_in_minutes: 45
9 + source_file_dependencies:
10 + - vllm/utils/jit_monitor.py
11 + - vllm/v1/worker/gpu_worker.py
12 + - vllm/model_executor/warmup/
13 + - vllm/config/observability.py
14 + - tests/jit_monitor/test_no_runtime_jit.py
15 + - tests/models/registry.py
16 + commands:
17 + # Boot a curated JIT-heavy model set with the JIT monitor in "error" mode
18 + # and run generation; any post-warmup JIT compilation fails the test.
19 + # Per-test watchdog so a wedged engine/CUDA init fails with a traceback
20 + # instead of running until the build timeout.
21 + - export PYTHONFAULTHANDLER=1
22 + - pytest -v -s jit_monitor/test_no_runtime_jit.py --timeout=900 --timeout-method=thread
tests/jit_monitor/__init__.py
tests/jit_monitor/conftest.py new
+27
@@ -0,0 +1,27 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from vllm.utils import jit_monitor
7 +
8 +
9 +@pytest.fixture(autouse=True)
10 +def _reset_monitor():
11 + """Reset global monitor state between tests.
12 +
13 + ``activate()`` installs process-global hooks and flips module globals, so
14 + without this every test would inherit the previous test's monitor state.
15 + """
16 +
17 + def reset():
18 + jit_monitor._active = False
19 + jit_monitor._mode = "warn"
20 + jit_monitor._verbose = False
21 + jit_monitor._cutedsl_hook_installed = False
22 + jit_monitor._tilelang_hook_installed = False
23 + jit_monitor._tilelang_jitimpl_compile_depth = 0
24 +
25 + reset()
26 + yield
27 + reset()
tests/jit_monitor/test_hooks.py new
+435
@@ -0,0 +1,435 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the JIT monitor hooks. Backends are mocked, so no GPU."""
4 +
5 +import inspect
6 +import os
7 +import sys
8 +from contextlib import contextmanager
9 +from types import ModuleType, SimpleNamespace
10 +from typing import Any, cast
11 +from unittest import mock
12 +
13 +import pytest
14 +
15 +from vllm.utils import jit_monitor
16 +
17 +pytestmark = pytest.mark.cpu_test
18 +
19 +
20 +# ------------------------------------------------------------------
21 +# Helpers — lightweight stand-ins for the modules ``activate()`` patches
22 +# ------------------------------------------------------------------
23 +
24 +
25 +def _make_fake_knobs(*, autotuning_print=False, jit_hook=None):
26 + """Build a minimal fake ``triton.knobs`` namespace."""
27 + autotuning = SimpleNamespace(print=autotuning_print)
28 + runtime = SimpleNamespace(jit_post_compile_hook=jit_hook)
29 + return SimpleNamespace(autotuning=autotuning, runtime=runtime)
30 +
31 +
32 +def _fake_cute_import_modules(compile_fn):
33 + """Fake Python's parent package + submodule for ``import cutlass.cute``."""
34 + fake_cute = cast(Any, ModuleType("cutlass.cute"))
35 + fake_cute.compile = compile_fn
36 + fake_parent_package = cast(Any, ModuleType("cutlass"))
37 + fake_parent_package.__path__ = []
38 + fake_parent_package.cute = fake_cute
39 + return {
40 + "cutlass": fake_parent_package,
41 + "cutlass.cute": fake_cute,
42 + }
43 +
44 +
45 +def _fake_cute_compile(*args, **kwargs):
46 + return "compiled"
47 +
48 +
49 +def _fake_tilelang_import_modules():
50 + """Fake Python's TileLang modules touched by ``jit_monitor.activate``."""
51 +
52 + class FakeJITKernel:
53 + def __init__(self, *args, **kwargs):
54 + pass
55 +
56 + class FakeJITImpl:
57 + def __init__(self, func, signature):
58 + self.func = func
59 + self.signature = signature
60 + self.mode = "lazy"
61 + self._kernel_cache = {}
62 +
63 + def __call__(self, *args, **kwargs):
64 + key, _ = self.func.parse_args(*args, **kwargs)
65 + kernel = self._kernel_cache.get(key)
66 + if kernel is None:
67 + kernel = "compiled"
68 + self._kernel_cache[key] = kernel
69 + return kernel
70 +
71 + fake_kernel = cast(Any, ModuleType("tilelang.jit.kernel"))
72 + fake_kernel.JITKernel = FakeJITKernel
73 +
74 + fake_jit = cast(Any, ModuleType("tilelang.jit"))
75 + fake_jit.JITImpl = FakeJITImpl
76 + fake_jit.kernel = fake_kernel
77 +
78 + fake_tilelang = cast(Any, ModuleType("tilelang"))
79 + fake_tilelang.jit = fake_jit
80 +
81 + return {
82 + "tilelang": fake_tilelang,
83 + "tilelang.jit": fake_jit,
84 + "tilelang.jit.kernel": fake_kernel,
85 + }
86 +
87 +
88 +@contextmanager
89 +def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile):
90 + """Patch the Triton and CuTeDSL imports touched by ``jit_monitor.activate``."""
91 + fake_triton = cast(Any, ModuleType("triton"))
92 + fake_triton.knobs = fake_knobs
93 + with (
94 + mock.patch.dict(
95 + sys.modules,
96 + {
97 + "triton": fake_triton,
98 + **_fake_cute_import_modules(cute_compile),
99 + **_fake_tilelang_import_modules(),
100 + },
101 + ),
102 + mock.patch.object(jit_monitor, "HAS_TRITON", True),
103 + ):
104 + yield
105 +
106 +
107 +def _triton_hook_kwargs(name: str):
108 + return dict(
109 + key="k",
110 + repr="r",
111 + fn=SimpleNamespace(name=name),
112 + compile=lambda: None,
113 + is_manual_warmup=False,
114 + already_compiled=False,
115 + )
116 +
117 +
118 +# ------------------------------------------------------------------
119 +# activate()
120 +# ------------------------------------------------------------------
121 +
122 +
123 +def test_activate_sets_active():
124 + assert not jit_monitor.is_active()
125 + with _patch_jit_modules(_make_fake_knobs()):
126 + jit_monitor.activate()
127 + assert jit_monitor.is_active()
128 +
129 +
130 +def test_activate_is_idempotent():
131 + fake = _make_fake_knobs()
132 + with _patch_jit_modules(fake):
133 + jit_monitor.activate()
134 + first_hook = fake.runtime.jit_post_compile_hook
135 + jit_monitor.activate()
136 + assert fake.runtime.jit_post_compile_hook is first_hook
137 +
138 +
139 +def test_activate_logs_info():
140 + with (
141 + mock.patch.object(jit_monitor.logger, "info") as m,
142 + _patch_jit_modules(_make_fake_knobs()),
143 + ):
144 + jit_monitor.activate()
145 + m.assert_called_once()
146 + assert "Kernel JIT monitor activated" in m.call_args[0][0]
147 +
148 +
149 +def test_activate_rejects_unknown_mode():
150 + with pytest.raises(ValueError, match="Unsupported JIT monitor mode"):
151 + jit_monitor.activate(mode="panic") # type: ignore[arg-type]
152 +
153 +
154 +def test_activate_without_triton():
155 + with mock.patch.object(jit_monitor, "HAS_TRITON", False):
156 + jit_monitor.activate()
157 + assert jit_monitor.is_active()
158 +
159 +
160 +# ------------------------------------------------------------------
161 +# Triton autotuning print
162 +# ------------------------------------------------------------------
163 +
164 +
165 +def test_autotuning_print_is_enabled():
166 + fake = _make_fake_knobs(autotuning_print=False)
167 + with _patch_jit_modules(fake):
168 + jit_monitor.activate()
169 + assert fake.autotuning.print is True
170 +
171 +
172 +def test_autotuning_print_respects_user_opt_out():
173 + fake = _make_fake_knobs(autotuning_print=False)
174 + with (
175 + mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "0"}),
176 + _patch_jit_modules(fake),
177 + ):
178 + jit_monitor.activate()
179 + assert fake.autotuning.print is False
180 +
181 +
182 +def test_autotuning_print_noop_when_user_already_enabled():
183 + fake = _make_fake_knobs(autotuning_print=True)
184 + with (
185 + mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "1"}),
186 + _patch_jit_modules(fake),
187 + ):
188 + jit_monitor.activate()
189 + assert fake.autotuning.print is True
190 +
191 +
192 +# ------------------------------------------------------------------
193 +# Triton JIT hook
194 +# ------------------------------------------------------------------
195 +
196 +
197 +def test_triton_hook_is_registered():
198 + fake = _make_fake_knobs()
199 + assert fake.runtime.jit_post_compile_hook is None
200 + with _patch_jit_modules(fake):
201 + jit_monitor.activate()
202 + assert fake.runtime.jit_post_compile_hook is not None
203 +
204 +
205 +def test_triton_hook_logs_warning():
206 + fake = _make_fake_knobs()
207 + with _patch_jit_modules(fake):
208 + jit_monitor.activate()
209 +
210 + hook = fake.runtime.jit_post_compile_hook
211 +
212 + with (
213 + mock.patch.object(jit_monitor.logger, "warning_once") as m,
214 + mock.patch.object(jit_monitor.logger, "warning") as warning,
215 + ):
216 + hook(**_triton_hook_kwargs("test_kernel"))
217 +
218 + m.assert_called_once()
219 + warning.assert_not_called()
220 + msg = m.call_args[0][0] % m.call_args[0][1:]
221 + assert "Triton kernel JIT compilation during inference" in msg
222 + assert "test_kernel" in msg
223 +
224 +
225 +def test_triton_hook_chains_existing_hook():
226 + existing = mock.MagicMock(return_value="existing_result")
227 + fake = _make_fake_knobs(jit_hook=existing)
228 + with _patch_jit_modules(fake):
229 + jit_monitor.activate()
230 +
231 + hook = fake.runtime.jit_post_compile_hook
232 + result = hook(**_triton_hook_kwargs("chained_kernel"))
233 +
234 + existing.assert_called_once()
235 + assert result == "existing_result"
236 +
237 +
238 +def test_triton_hook_works_without_existing_hook():
239 + fake = _make_fake_knobs(jit_hook=None)
240 + with _patch_jit_modules(fake):
241 + jit_monitor.activate()
242 +
243 + hook = fake.runtime.jit_post_compile_hook
244 + assert hook(**_triton_hook_kwargs("solo_kernel")) is None
245 +
246 +
247 +def test_triton_hook_error_mode_raises():
248 + fake = _make_fake_knobs()
249 + with _patch_jit_modules(fake):
250 + jit_monitor.activate(mode="error")
251 +
252 + hook = fake.runtime.jit_post_compile_hook
253 + with pytest.raises(RuntimeError, match="Triton kernel JIT compilation"):
254 + hook(**_triton_hook_kwargs("error_kernel"))
255 +
256 +
257 +# ------------------------------------------------------------------
258 +# CuTeDSL hook
259 +# ------------------------------------------------------------------
260 +
261 +
262 +def test_cutedsl_compile_logs_warning():
263 + with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile):
264 + import cutlass.cute as cute
265 +
266 + jit_monitor.activate()
267 + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
268 + result = cute.compile(lambda: None, "arg", option=True)
269 +
270 + assert result == "compiled"
271 + warning_once.assert_called_once()
272 + msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
273 + assert "CuTeDSL JIT compilation during inference" in msg
274 +
275 +
276 +def test_cutedsl_compile_logs_verbose_warning():
277 + with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile):
278 + import cutlass.cute as cute
279 +
280 + jit_monitor.activate(verbose=True)
281 + with mock.patch.object(jit_monitor.logger, "warning") as warning:
282 + result = cute.compile(lambda: None, "arg", option=True)
283 +
284 + assert result == "compiled"
285 + warning.assert_called_once()
286 + msg = warning.call_args[0][0] % warning.call_args[0][1:]
287 + assert "CuTeDSL JIT compilation during inference" in msg
288 +
289 +
290 +def test_cutedsl_error_mode_raises():
291 + with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile):
292 + import cutlass.cute as cute
293 +
294 + jit_monitor.activate(mode="error")
295 + with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"):
296 + cute.compile(lambda: None, "arg", option=True)
297 +
298 +
299 +def test_cutedsl_subscripted_compile_is_monitored():
300 + """``cute.compile[options](...)`` (flashinfer >= 0.6.14) must work."""
301 +
302 + class FakeCompileCallable:
303 + def __getitem__(self, options):
304 + return self
305 +
306 + def __call__(self, *args, **kwargs):
307 + return "compiled"
308 +
309 + with _patch_jit_modules(_make_fake_knobs(), cute_compile=FakeCompileCallable()):
310 + import cutlass.cute as cute
311 +
312 + jit_monitor.activate()
313 + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
314 + result = cute.compile[("opt_level", 3)](lambda: None, "arg")
315 +
316 + assert result == "compiled"
317 + warning_once.assert_called_once()
318 +
319 +
320 +# ------------------------------------------------------------------
321 +# TileLang hook
322 +# ------------------------------------------------------------------
323 +
324 +
325 +def test_tilelang_jit_kernel_logs_warning():
326 + with _patch_jit_modules(_make_fake_knobs()):
327 + from tilelang.jit.kernel import JITKernel
328 +
329 + func = SimpleNamespace(attrs={"global_symbol": "tl_kernel"})
330 + jit_monitor.activate()
331 + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
332 + JITKernel(func=func, out_idx=None, execution_backend="tvm_ffi")
333 +
334 + warning_once.assert_called_once()
335 + msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
336 + assert "TileLang JIT compilation during inference" in msg
337 + assert "tl_kernel" in msg
338 +
339 +
340 +def test_tilelang_jit_impl_logs_warning():
341 + with _patch_jit_modules(_make_fake_knobs()):
342 + from tilelang.jit import JITImpl
343 +
344 + def tilelang_fn(
345 + gemm_out_mul,
346 + hidden_size: int,
347 + n_splits: int = 1,
348 + hc_mult: int = 4,
349 + ):
350 + return None
351 +
352 + class FakeFunc:
353 + orig_func = tilelang_fn
354 +
355 + def parse_args(self, *args, **kwargs):
356 + return (
357 + (
358 + "tilelang_key",
359 + kwargs["hidden_size"],
360 + kwargs.get("n_splits", 1),
361 + ),
362 + {},
363 + )
364 +
365 + def set_mode(self, mode):
366 + self.mode = mode
367 +
368 + tensor = SimpleNamespace(
369 + shape=(2, 16, 24),
370 + dtype="float32",
371 + device="cuda:0",
372 + )
373 + impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn))
374 +
375 + jit_monitor.activate()
376 + with (
377 + mock.patch.object(jit_monitor.logger, "warning_once") as warning_once,
378 + mock.patch.object(jit_monitor.logger, "warning") as warning,
379 + ):
380 + impl(tensor, hidden_size=7168, n_splits=2)
381 +
382 + warning_once.assert_called_once()
383 + warning.assert_not_called()
384 + msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
385 + assert "TileLang JIT compilation during inference" in msg
386 + assert "tilelang_fn" in msg
387 +
388 +
389 +def test_tilelang_jit_impl_does_not_log_on_cache_hit():
390 + with _patch_jit_modules(_make_fake_knobs()):
391 + from tilelang.jit import JITImpl
392 +
393 + def tilelang_fn(gemm_out_mul, n_splits: int = 1):
394 + return None
395 +
396 + class FakeFunc:
397 + orig_func = tilelang_fn
398 +
399 + def parse_args(self, *args, **kwargs):
400 + return (("tilelang_key", kwargs.get("n_splits", 1)), {})
401 +
402 + def set_mode(self, mode):
403 + self.mode = mode
404 +
405 + tensor = SimpleNamespace(shape=(2, 16, 24), dtype="float32")
406 + impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn))
407 +
408 + jit_monitor.activate()
409 + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
410 + impl(tensor, n_splits=2)
411 + impl(tensor, n_splits=2)
412 +
413 + warning_once.assert_called_once()
414 +
415 +
416 +def test_tilelang_from_database_does_not_log():
417 + with _patch_jit_modules(_make_fake_knobs()):
418 + from tilelang.jit.kernel import JITKernel
419 +
420 + func = SimpleNamespace(attrs={"global_symbol": "cached_tl_kernel"})
421 + jit_monitor.activate()
422 + with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
423 + JITKernel(func=func, from_database=True)
424 +
425 + warning_once.assert_not_called()
426 +
427 +
428 +def test_tilelang_error_mode_raises():
429 + with _patch_jit_modules(_make_fake_knobs()):
430 + from tilelang.jit.kernel import JITKernel
431 +
432 + func = SimpleNamespace(attrs={"global_symbol": "error_tl_kernel"})
433 + jit_monitor.activate(mode="error")
434 + with pytest.raises(RuntimeError, match="TileLang JIT compilation"):
435 + JITKernel(func=func)
tests/jit_monitor/test_hooks_gpu.py new
+85
@@ -0,0 +1,85 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""End-to-end JIT monitor tests: real Triton kernel, real GPU, real hook."""
4 +
5 +from unittest import mock
6 +
7 +import pytest
8 +
9 +from vllm.utils import jit_monitor
10 +
11 +try:
12 + import torch
13 +
14 + _HAS_CUDA = torch.cuda.is_available()
15 +except ImportError:
16 + _HAS_CUDA = False
17 +
18 +try:
19 + import triton
20 + import triton.language as tl
21 +
22 + _HAS_TRITON = True
23 +except ImportError:
24 + _HAS_TRITON = False
25 +
26 +pytestmark = pytest.mark.skipif(
27 + not (_HAS_CUDA and _HAS_TRITON),
28 + reason="Requires CUDA GPU and Triton",
29 +)
30 +
31 +
32 +if _HAS_TRITON:
33 +
34 + @triton.jit
35 + def _add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
36 + pid = tl.program_id(0)
37 + offs = pid * BLOCK + tl.arange(0, BLOCK)
38 + mask = offs < n
39 + x = tl.load(x_ptr + offs, mask=mask)
40 + y = tl.load(y_ptr + offs, mask=mask)
41 + tl.store(out_ptr + offs, x + y, mask=mask)
42 +
43 +
44 +def _run_add_kernel(n: int, block: int = 256, offset: int = 0) -> None:
45 + """Launch ``_add_kernel`` with vectors of length *n*."""
46 + x = torch.randn(n + offset, device="cuda")[offset:] # affect alignment
47 + y = torch.randn(n, device="cuda")
48 + out = torch.empty(n, device="cuda")
49 + grid = ((n + block - 1) // block,)
50 + _add_kernel[grid](x, y, out, n, BLOCK=block)
51 + torch.accelerator.synchronize()
52 +
53 +
54 +def test_no_warning_on_cached_shape():
55 + _run_add_kernel(1024)
56 +
57 + jit_monitor.activate()
58 + with mock.patch.object(jit_monitor.logger, "warning_once") as w:
59 + _run_add_kernel(1024)
60 + w.assert_not_called()
61 +
62 +
63 +def test_warning_on_new_constexpr():
64 + _run_add_kernel(1024, block=256)
65 +
66 + jit_monitor.activate()
67 + with mock.patch.object(jit_monitor.logger, "warning_once") as w:
68 + # Different BLOCK (a tl.constexpr) forces recompilation.
69 + _run_add_kernel(1024, block=512)
70 + w.assert_called()
71 + msg = w.call_args[0][0] % w.call_args[0][1:]
72 + assert "_add_kernel" in msg
73 +
74 +
75 +def test_verbose_warning_on_each_new_pointer_alignment():
76 + _run_add_kernel(1024)
77 +
78 + jit_monitor.activate(verbose=True)
79 + with (
80 + mock.patch.object(jit_monitor.logger, "warning") as w,
81 + mock.patch.object(jit_monitor.logger, "warning_once") as w_once,
82 + ):
83 + _run_add_kernel(1024, offset=1)
84 + assert w.called
85 + w_once.assert_not_called()
tests/jit_monitor/test_no_runtime_jit.py new
+141
@@ -0,0 +1,141 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Catch runtime (post-warmup) JIT compilations for JIT-heavy backends running
4 +e2e tests on popular models, for which we only load a few blocks for performance.
5 +
6 +NOTE(NickLucche) With cuda graphs on, kernels fully covered by graphs captured during
7 +warmup do not re-trigger the Python JIT hooks. The targeted paths (prefill
8 +MoE/MLA/SSM and the sampler) run mixed, so they are unaffected.
9 +"""
10 +
11 +from dataclasses import dataclass
12 +
13 +import pytest
14 +
15 +from vllm import LLM, SamplingParams
16 +from vllm.inputs import TokensPrompt
17 +
18 +from ..models.utils import dummy_hf_overrides
19 +from ..utils import create_new_process_for_each_test
20 +
21 +# Warmup coverage is still incomplete for these backends, so the monitor fires
22 +# during inference. Tracked in https://github.com/vllm-project/vllm/issues/49349;
23 +# drop this once the warmup contract migrations land.
24 +pytestmark = pytest.mark.skip(reason="Kernel warmup coverage is still incomplete")
25 +
26 +
27 +@dataclass(frozen=True)
28 +class JitModel:
29 + model: str
30 + draft: str | None = None
31 + trust_remote_code: bool = False
32 +
33 +
34 +JIT_MONITOR_MODELS = [
35 + JitModel("Qwen/Qwen3-0.6B"),
36 + JitModel("deepseek-ai/DeepSeek-V2-Lite-Chat", trust_remote_code=True),
37 + JitModel("deepseek-ai/DeepSeek-V3", trust_remote_code=True),
38 + JitModel("ibm-granite/granite-4.0-tiny-preview"),
39 + JitModel(
40 + "luccafong/deepseek_mtp_main_random",
41 + draft="luccafong/deepseek_mtp_draft_random",
42 + trust_remote_code=True,
43 + ),
44 + JitModel(
45 + "eagle618/deepseek-v3-random",
46 + draft="eagle618/eagle-deepseek-v3-random",
47 + trust_remote_code=True,
48 + ),
49 +]
50 +
51 +
52 +def _run_shape_battery(llm: LLM) -> None:
53 + """Exercise diverse compile keys so missing warmup keys surface.
54 +
55 + Token-id prompts keep shapes exact and avoid depending on a tokenizer.
56 + Outputs are meaningless under dummy weights; we assert only that no JIT
57 + fired.
58 + """
59 + short = TokensPrompt(prompt_token_ids=[1, 2, 3, 4])
60 + medium = TokensPrompt(prompt_token_ids=list(range(1, 33)))
61 + long = TokensPrompt(prompt_token_ids=list(range(1, 129)))
62 +
63 + # Greedy single-sequence multi-step decode: prefill + autoregressive decode
64 + # + greedy sampler.
65 + llm.generate(medium, SamplingParams(temperature=0.0, max_tokens=16))
66 +
67 + # Batched prefill with mixed lengths: varlen prefill + padded decode.
68 + llm.generate([short, medium, long], SamplingParams(temperature=0.0, max_tokens=8))
69 +
70 + # Triton sampler kernels: top_k / top_p / min_p each specialize.
71 + for sampling_params in (
72 + SamplingParams(temperature=0.8, top_k=20, max_tokens=8, seed=0),
73 + SamplingParams(temperature=0.8, top_p=0.9, max_tokens=8, seed=0),
74 + SamplingParams(temperature=0.8, min_p=0.1, max_tokens=8, seed=0),
75 + SamplingParams(
76 + temperature=0.8, top_k=20, top_p=0.9, min_p=0.1, max_tokens=8, seed=0
77 + ),
78 + ):
79 + llm.generate(medium, sampling_params)
80 +
81 + # Heterogeneous SamplingParams in one step, where missing sampler warmup
82 + # keys most often hide.
83 + llm.generate(
84 + [medium] * 4,
85 + [
86 + SamplingParams(temperature=0.0, max_tokens=8),
87 + SamplingParams(temperature=0.8, top_k=20, max_tokens=8, seed=0),
88 + SamplingParams(temperature=0.8, top_p=0.9, max_tokens=8, seed=0),
89 + SamplingParams(temperature=0.8, min_p=0.1, max_tokens=8, seed=0),
90 + ],
91 + )
92 +
93 +
94 +@create_new_process_for_each_test("spawn")
95 +def can_run_without_jit(spec: JitModel):
96 + """Boot ``spec`` with the monitor armed and run the shape battery.
97 +
98 + A subprocess per model is required: the monitor's hooks are process-global
99 + and, once armed in ``error`` mode, stay armed. It must be spawned rather
100 + than forked, since forking a pytest process that already initialized CUDA
101 + poisons the child.
102 + """
103 + llm = LLM(
104 + spec.model,
105 + trust_remote_code=spec.trust_remote_code,
106 + max_model_len=2048,
107 + max_num_seqs=8,
108 + gpu_memory_utilization=0.80,
109 + load_format="dummy",
110 + hf_overrides=dummy_hf_overrides,
111 + # cuda graphs cover captured decode shapes, run eager.
112 + enforce_eager=False,
113 + jit_monitor_mode="error",
114 + speculative_config={
115 + "model": spec.draft,
116 + "num_speculative_tokens": 2,
117 + }
118 + if spec.draft
119 + else None,
120 + )
121 +
122 + try:
123 + _run_shape_battery(llm)
124 + except Exception as e:
125 + # The monitor's message contains "during inference"; distinguish a real
126 + # JIT miss from an unrelated crash.
127 + if "during inference" in str(e):
128 + pytest.fail(
129 + f"{spec.model}: post-warmup JIT compilation detected - a warmup "
130 + f"key is missing for a shape in the battery.\n{e}"
131 + )
132 + raise
133 +
134 +
135 +@pytest.mark.parametrize("spec", JIT_MONITOR_MODELS, ids=lambda s: s.model)
136 +def test_no_runtime_jit(spec: JitModel, monkeypatch: pytest.MonkeyPatch):
137 + """Assert JIT-heavy backends do not JIT-compile during inference."""
138 + # Set here rather than in the child so the spawned process inherits it:
139 + # the engine core must not be forked once the test process has CUDA up.
140 + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
141 + can_run_without_jit(spec)
tests/test_jit_monitor.py deleted
-534
@@ -1,534 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -import inspect
4 -import os
5 -import sys
6 -from contextlib import contextmanager
7 -from types import ModuleType, SimpleNamespace
8 -from typing import Any, cast
9 -from unittest import mock
10 -
11 -import pytest
12 -
13 -from vllm.utils import jit_monitor
14 -
15 -
16 -@pytest.fixture(autouse=True)
17 -def _reset_monitor():
18 - """Reset global monitor state between tests."""
19 - jit_monitor._active = False
20 - jit_monitor._mode = "warn"
21 - jit_monitor._verbose = False
22 - jit_monitor._cutedsl_hook_installed = False
23 - jit_monitor._tilelang_hook_installed = False
24 - jit_monitor._tilelang_jitimpl_compile_depth = 0
25 - yield
26 - jit_monitor._active = False
27 - jit_monitor._mode = "warn"
28 - jit_monitor._verbose = False
29 - jit_monitor._cutedsl_hook_installed = False
30 - jit_monitor._tilelang_hook_installed = False
31 - jit_monitor._tilelang_jitimpl_compile_depth = 0
32 -
33 -
34 -# ------------------------------------------------------------------
35 -# Helpers — lightweight stand-ins for the modules ``activate()`` patches
36 -# ------------------------------------------------------------------
37 -
38 -
39 -def _make_fake_knobs(*, autotuning_print=False, jit_hook=None):
40 - """Build a minimal fake ``triton.knobs`` namespace."""
41 - autotuning = SimpleNamespace(print=autotuning_print)
42 - runtime = SimpleNamespace(jit_post_compile_hook=jit_hook)
43 - return SimpleNamespace(autotuning=autotuning, runtime=runtime)
44 -
45 -
46 -def _fake_cute_import_modules(compile_fn):
47 - """Fake Python's parent package + submodule for ``import cutlass.cute``."""
48 - fake_cute = cast(Any, ModuleType("cutlass.cute"))
49 - fake_cute.compile = compile_fn
50 - fake_parent_package = cast(Any, ModuleType("cutlass"))
51 - fake_parent_package.__path__ = []
52 - fake_parent_package.cute = fake_cute
53 - return {
54 - "cutlass": fake_parent_package,
55 - "cutlass.cute": fake_cute,
56 - }
57 -
58 -
59 -def _fake_cute_compile(*args, **kwargs):
60 - return "compiled"
61 -
62 -
63 -def _fake_tilelang_import_modules():
64 - """Fake Python's TileLang modules touched by ``jit_monitor.activate``."""
65 -
66 - class FakeJITKernel:
67 - def __init__(self, *args, **kwargs):
68 - pass
69 -
70 - class FakeJITImpl:
71 - def __init__(self, func, signature):
72 - self.func = func
73 - self.signature = signature
74 - self.mode = "lazy"
75 - self._kernel_cache = {}
76 -
77 - def __call__(self, *args, **kwargs):
78 - key, _ = self.func.parse_args(*args, **kwargs)
79 - kernel = self._kernel_cache.get(key)
80 - if kernel is None:
81 - kernel = "compiled"
82 - self._kernel_cache[key] = kernel
83 - return kernel
84 -
85 - fake_kernel = cast(Any, ModuleType("tilelang.jit.kernel"))
86 - fake_kernel.JITKernel = FakeJITKernel
87 -
88 - fake_jit = cast(Any, ModuleType("tilelang.jit"))
89 - fake_jit.JITImpl = FakeJITImpl
90 - fake_jit.kernel = fake_kernel
91 -
92 - fake_tilelang = cast(Any, ModuleType("tilelang"))
93 - fake_tilelang.jit = fake_jit
94 -
95 - return {
96 - "tilelang": fake_tilelang,
97 - "tilelang.jit": fake_jit,
98 - "tilelang.jit.kernel": fake_kernel,
99 - }
100 -
101 -
102 -@contextmanager
103 -def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile):
104 - """Patch the Triton and CuTeDSL imports touched by ``jit_monitor.activate``."""
105 - fake_triton = cast(Any, ModuleType("triton"))
106 - fake_triton.knobs = fake_knobs
107 - with (
108 - mock.patch.dict(
109 - sys.modules,
110 - {
111 - "triton": fake_triton,
112 - **_fake_cute_import_modules(cute_compile),
113 - **_fake_tilelang_import_modules(),
114 - },
115 - ),
116 - mock.patch.object(jit_monitor, "HAS_TRITON", True),
117 - ):
118 - yield
119 -
120 -
121 -# ------------------------------------------------------------------
122 -# Unit tests (no GPU required, triton is mocked)
123 -# ------------------------------------------------------------------
124 -
125 -
126 -class TestActivateBasic:
127 - def test_sets_active(self):
128 - assert not jit_monitor.is_active()
129 - with _patch_jit_modules(_make_fake_knobs()):
130 - jit_monitor.activate()
131 - assert jit_monitor.is_active()
132 -
133 - def test_idempotent(self):
134 - fake = _make_fake_knobs()
135 - with _patch_jit_modules(fake):
136 - jit_monitor.activate()
137 - first_hook = fake.runtime.jit_post_compile_hook
138 - jit_monitor.activate()
139 - assert fake.runtime.jit_post_compile_hook is first_hook
140 -
141 - def test_logs_info_on_activation(self):
142 - with (
143 - mock.patch.object(jit_monitor.logger, "info") as m,
144 - _patch_jit_modules(_make_fake_knobs()),
145 - ):
146 - jit_monitor.activate()
147 - m.assert_called_once()
148 - assert "Kernel JIT monitor activated" in m.call_args[0][0]
149 -
150 - def test_rejects_unknown_mode(self):
151 - with pytest.raises(ValueError, match="Unsupported JIT monitor mode"):
152 - jit_monitor.activate(mode="panic") # type: ignore[arg-type]
153 -
154 -
155 -class TestAutotuningPrint:
156 - def test_enables_autotuning_print(self):
157 - fake = _make_fake_knobs(autotuning_print=False)
158 - with _patch_jit_modules(fake):
159 - jit_monitor.activate()
160 - assert fake.autotuning.print is True
161 -
162 - def test_respects_user_opt_out(self):
163 - fake = _make_fake_knobs(autotuning_print=False)
164 - with (
165 - mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "0"}),
166 - _patch_jit_modules(fake),
167 - ):
168 - jit_monitor.activate()
169 - assert fake.autotuning.print is False
170 -
171 - def test_noop_when_user_already_enabled(self):
172 - fake = _make_fake_knobs(autotuning_print=True)
173 - with (
174 - mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "1"}),
175 - _patch_jit_modules(fake),
176 - ):
177 - jit_monitor.activate()
178 - assert fake.autotuning.print is True
179 -
180 -
181 -class TestTritonJitHook:
182 - def test_hook_registered(self):
183 - fake = _make_fake_knobs()
184 - assert fake.runtime.jit_post_compile_hook is None
185 - with _patch_jit_modules(fake):
186 - jit_monitor.activate()
187 - assert fake.runtime.jit_post_compile_hook is not None
188 -
189 - def test_hook_logs_warning(self):
190 - fake = _make_fake_knobs()
191 - with _patch_jit_modules(fake):
192 - jit_monitor.activate()
193 -
194 - hook = fake.runtime.jit_post_compile_hook
195 - mock_fn = SimpleNamespace(name="test_kernel")
196 -
197 - with (
198 - mock.patch.object(jit_monitor.logger, "warning_once") as m,
199 - mock.patch.object(jit_monitor.logger, "warning") as warning,
200 - ):
201 - hook(
202 - key="some_key",
203 - repr="some_repr",
204 - fn=mock_fn,
205 - compile=lambda: None,
206 - is_manual_warmup=False,
207 - already_compiled=False,
208 - )
209 -
210 - m.assert_called_once()
211 - warning.assert_not_called()
212 - msg = m.call_args[0][0] % m.call_args[0][1:]
213 - assert "Triton kernel JIT compilation during inference" in msg
214 - assert "test_kernel" in msg
215 -
216 - def test_hook_chains_existing_hook(self):
217 - existing = mock.MagicMock(return_value="existing_result")
218 - fake = _make_fake_knobs(jit_hook=existing)
219 - with _patch_jit_modules(fake):
220 - jit_monitor.activate()
221 -
222 - hook = fake.runtime.jit_post_compile_hook
223 - mock_fn = SimpleNamespace(name="chained_kernel")
224 - kwargs = dict(
225 - key="k",
226 - repr="r",
227 - fn=mock_fn,
228 - compile=lambda: None,
229 - is_manual_warmup=False,
230 - already_compiled=False,
231 - )
232 - result = hook(**kwargs)
233 -
234 - existing.assert_called_once()
235 - assert result == "existing_result"
236 -
237 - def test_hook_works_without_existing_hook(self):
238 - fake = _make_fake_knobs(jit_hook=None)
239 - with _patch_jit_modules(fake):
240 - jit_monitor.activate()
241 -
242 - hook = fake.runtime.jit_post_compile_hook
243 - mock_fn = SimpleNamespace(name="solo_kernel")
244 - result = hook(
245 - key="k",
246 - repr="r",
247 - fn=mock_fn,
248 - compile=lambda: None,
249 - is_manual_warmup=False,
250 - already_compiled=False,
251 - )
252 - assert result is None
253 -
254 - def test_error_mode_raises(self):
255 - fake = _make_fake_knobs()
256 - with _patch_jit_modules(fake):
257 - jit_monitor.activate(mode="error")
258 -
259 - hook = fake.runtime.jit_post_compile_hook
260 - mock_fn = SimpleNamespace(name="error_kernel")
261 - with pytest.raises(RuntimeError, match="Triton kernel JIT compilation"):
262 - hook(
263 - key="k",
264 - repr="r",
265 - fn=mock_fn,
266 - compile=lambda: None,
267 - is_manual_warmup=False,
268 - already_compiled=False,
269 - )
270 -
271 -
272 -class TestNoTritonFallback:
273 - def test_activate_without_triton(self):
274 - with mock.patch.object(jit_monitor, "HAS_TRITON", False):
275 - jit_monitor.activate()
276 - assert jit_monitor.is_active()
277 -
278 -
279 -class TestCuTeDSLHook:
280 - def test_compile_logs_warning(self):
281 - def compile_fn(*args, **kwargs):
282 - return "compiled"
283 -
284 - with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn):
285 - import cutlass.cute as cute
286 -
287 - jit_monitor.activate()
288 - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
289 - result = cute.compile(lambda: None, "arg", option=True)
290 -
291 - assert result == "compiled"
292 - warning_once.assert_called_once()
293 - msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
294 - assert "CuTeDSL JIT compilation during inference" in msg
295 -
296 - def test_compile_logs_verbose_warning(self):
297 - def compile_fn(*args, **kwargs):
298 - return "compiled"
299 -
300 - with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn):
301 - import cutlass.cute as cute
302 -
303 - jit_monitor.activate(verbose=True)
304 - with mock.patch.object(jit_monitor.logger, "warning") as warning:
305 - result = cute.compile(lambda: None, "arg", option=True)
306 -
307 - assert result == "compiled"
308 - warning.assert_called_once()
309 - msg = warning.call_args[0][0] % warning.call_args[0][1:]
310 - assert "CuTeDSL JIT compilation during inference" in msg
311 -
312 - def test_error_mode_raises(self):
313 - def compile_fn(*args, **kwargs):
314 - return "compiled"
315 -
316 - with _patch_jit_modules(_make_fake_knobs(), cute_compile=compile_fn):
317 - import cutlass.cute as cute
318 -
319 - jit_monitor.activate(mode="error")
320 - with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"):
321 - cute.compile(lambda: None, "arg", option=True)
322 -
323 - def test_subscripted_compile_is_monitored(self):
324 - """``cute.compile[options](...)`` (flashinfer >= 0.6.14) must work."""
325 -
326 - class FakeCompileCallable:
327 - def __getitem__(self, options):
328 - return self
329 -
330 - def __call__(self, *args, **kwargs):
331 - return "compiled"
332 -
333 - with _patch_jit_modules(_make_fake_knobs(), cute_compile=FakeCompileCallable()):
334 - import cutlass.cute as cute
335 -
336 - jit_monitor.activate()
337 - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
338 - result = cute.compile[("opt_level", 3)](lambda: None, "arg")
339 -
340 - assert result == "compiled"
341 - warning_once.assert_called_once()
342 -
343 -
344 -class TestTileLangHook:
345 - def test_jit_kernel_logs_warning(self):
346 - with _patch_jit_modules(_make_fake_knobs()):
347 - from tilelang.jit.kernel import JITKernel
348 -
349 - func = SimpleNamespace(attrs={"global_symbol": "tl_kernel"})
350 - jit_monitor.activate()
351 - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
352 - JITKernel(func=func, out_idx=None, execution_backend="tvm_ffi")
353 -
354 - warning_once.assert_called_once()
355 - msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
356 - assert "TileLang JIT compilation during inference" in msg
357 - assert "tl_kernel" in msg
358 -
359 - def test_jit_impl_logs_warning(self):
360 - with _patch_jit_modules(_make_fake_knobs()):
361 - from tilelang.jit import JITImpl
362 -
363 - def tilelang_fn(
364 - gemm_out_mul,
365 - hidden_size: int,
366 - n_splits: int = 1,
367 - hc_mult: int = 4,
368 - ):
369 - return None
370 -
371 - class FakeFunc:
372 - orig_func = tilelang_fn
373 -
374 - def parse_args(self, *args, **kwargs):
375 - return (
376 - (
377 - "tilelang_key",
378 - kwargs["hidden_size"],
379 - kwargs.get("n_splits", 1),
380 - ),
381 - {},
382 - )
383 -
384 - def set_mode(self, mode):
385 - self.mode = mode
386 -
387 - tensor = SimpleNamespace(
388 - shape=(2, 16, 24),
389 - dtype="float32",
390 - device="cuda:0",
391 - )
392 - impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn))
393 -
394 - jit_monitor.activate()
395 - with (
396 - mock.patch.object(jit_monitor.logger, "warning_once") as warning_once,
397 - mock.patch.object(jit_monitor.logger, "warning") as warning,
398 - ):
399 - impl(tensor, hidden_size=7168, n_splits=2)
400 -
401 - warning_once.assert_called_once()
402 - warning.assert_not_called()
403 - msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
404 - assert "TileLang JIT compilation during inference" in msg
405 - assert "tilelang_fn" in msg
406 -
407 - def test_jit_impl_does_not_log_on_cache_hit(self):
408 - with _patch_jit_modules(_make_fake_knobs()):
409 - from tilelang.jit import JITImpl
410 -
411 - def tilelang_fn(gemm_out_mul, n_splits: int = 1):
412 - return None
413 -
414 - class FakeFunc:
415 - orig_func = tilelang_fn
416 -
417 - def parse_args(self, *args, **kwargs):
418 - return (("tilelang_key", kwargs.get("n_splits", 1)), {})
419 -
420 - def set_mode(self, mode):
421 - self.mode = mode
422 -
423 - tensor = SimpleNamespace(shape=(2, 16, 24), dtype="float32")
424 - impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn))
425 -
426 - jit_monitor.activate()
427 - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
428 - impl(tensor, n_splits=2)
429 - impl(tensor, n_splits=2)
430 -
431 - warning_once.assert_called_once()
432 -
433 - def test_from_database_does_not_log(self):
434 - with _patch_jit_modules(_make_fake_knobs()):
435 - from tilelang.jit.kernel import JITKernel
436 -
437 - func = SimpleNamespace(attrs={"global_symbol": "cached_tl_kernel"})
438 - jit_monitor.activate()
439 - with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
440 - JITKernel(func=func, from_database=True)
441 -
442 - warning_once.assert_not_called()
443 -
444 - def test_error_mode_raises(self):
445 - with _patch_jit_modules(_make_fake_knobs()):
446 - from tilelang.jit.kernel import JITKernel
447 -
448 - func = SimpleNamespace(attrs={"global_symbol": "error_tl_kernel"})
449 - jit_monitor.activate(mode="error")
450 - with pytest.raises(RuntimeError, match="TileLang JIT compilation"):
451 - JITKernel(func=func)
452 -
453 -
454 -# ------------------------------------------------------------------
455 -# Integration tests (real Triton + GPU)
456 -# ------------------------------------------------------------------
457 -
458 -try:
459 - import torch
460 -
461 - _HAS_CUDA = torch.cuda.is_available()
462 -except ImportError:
463 - _HAS_CUDA = False
464 -
465 -try:
466 - import triton
467 - import triton.language as tl
468 -
469 - _HAS_TRITON = True
470 -except ImportError:
471 - _HAS_TRITON = False
472 -
473 -_skip_no_gpu = pytest.mark.skipif(
474 - not (_HAS_CUDA and _HAS_TRITON),
475 - reason="Requires CUDA GPU and Triton",
476 -)
477 -
478 -
479 -if _HAS_TRITON:
480 -
481 - @triton.jit
482 - def _add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr):
483 - pid = tl.program_id(0)
484 - offs = pid * BLOCK + tl.arange(0, BLOCK)
485 - mask = offs < n
486 - x = tl.load(x_ptr + offs, mask=mask)
487 - y = tl.load(y_ptr + offs, mask=mask)
488 - tl.store(out_ptr + offs, x + y, mask=mask)
489 -
490 -
491 -def _run_add_kernel(n: int, block: int = 256, offset: int = 0) -> None:
492 - """Launch ``_add_kernel`` with vectors of length *n*."""
493 - x = torch.randn(n + offset, device="cuda")[offset:] # affect alignment
494 - y = torch.randn(n, device="cuda")
495 - out = torch.empty(n, device="cuda")
496 - grid = ((n + block - 1) // block,)
497 - _add_kernel[grid](x, y, out, n, BLOCK=block)
498 - torch.accelerator.synchronize()
499 -
500 -
501 -@_skip_no_gpu
502 -class TestTritonJitHookIntegration:
503 - """End-to-end: real Triton kernel, real GPU, real hook."""
504 -
505 - def test_no_warning_on_cached_shape(self):
506 - _run_add_kernel(1024)
507 -
508 - jit_monitor.activate()
509 - with mock.patch.object(jit_monitor.logger, "warning_once") as w:
510 - _run_add_kernel(1024)
511 - w.assert_not_called()
512 -
513 - def test_warning_on_new_constexpr(self):
514 - _run_add_kernel(1024, block=256)
515 -
516 - jit_monitor.activate()
517 - with mock.patch.object(jit_monitor.logger, "warning_once") as w:
518 - # Different BLOCK (a tl.constexpr) forces recompilation.
519 - _run_add_kernel(1024, block=512)
520 - w.assert_called()
521 - msg = w.call_args[0][0] % w.call_args[0][1:]
522 - assert "_add_kernel" in msg
523 -
524 - def test_verbose_warning_on_each_new_pointer_alignment(self):
525 - _run_add_kernel(1024)
526 -
527 - jit_monitor.activate(verbose=True)
528 - with (
529 - mock.patch.object(jit_monitor.logger, "warning") as w,
530 - mock.patch.object(jit_monitor.logger, "warning_once") as w_once,
531 - ):
532 - _run_add_kernel(1024, offset=1)
533 - assert w.called
534 - w_once.assert_not_called()