| 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.platforms import current_platform |
| 16 | from vllm.utils import jit_monitor |
| 17 | |
| 18 | pytestmark = pytest.mark.cpu_test |
| 19 | |
| 20 | |
| 21 | # ------------------------------------------------------------------ |
| 22 | # Helpers — lightweight stand-ins for the modules ``activate()`` patches |
| 23 | # ------------------------------------------------------------------ |
| 24 | |
| 25 | |
| 26 | def _make_fake_knobs(*, autotuning_print=False, jit_hook=None): |
| 27 | """Build a minimal fake ``triton.knobs`` namespace.""" |
| 28 | autotuning = SimpleNamespace(print=autotuning_print) |
| 29 | runtime = SimpleNamespace(jit_post_compile_hook=jit_hook) |
| 30 | return SimpleNamespace(autotuning=autotuning, runtime=runtime) |
| 31 | |
| 32 | |
| 33 | def _fake_cute_import_modules(compile_fn): |
| 34 | """Fake Python's parent package + submodule for ``import cutlass.cute``.""" |
| 35 | fake_cute = cast(Any, ModuleType("cutlass.cute")) |
| 36 | fake_cute.compile = compile_fn |
| 37 | fake_parent_package = cast(Any, ModuleType("cutlass")) |
| 38 | fake_parent_package.__path__ = [] |
| 39 | fake_parent_package.cute = fake_cute |
| 40 | return { |
| 41 | "cutlass": fake_parent_package, |
| 42 | "cutlass.cute": fake_cute, |
| 43 | } |
| 44 | |
| 45 | |
| 46 | def _fake_cute_compile(*args, **kwargs): |
| 47 | return "compiled" |
| 48 | |
| 49 | |
| 50 | def _fake_tilelang_import_modules(): |
| 51 | """Fake Python's TileLang modules touched by ``jit_monitor.activate``.""" |
| 52 | |
| 53 | class FakeJITKernel: |
| 54 | def __init__(self, *args, **kwargs): |
| 55 | pass |
| 56 | |
| 57 | class FakeJITImpl: |
| 58 | def __init__(self, func, signature): |
| 59 | self.func = func |
| 60 | self.signature = signature |
| 61 | self.mode = "lazy" |
| 62 | self._kernel_cache = {} |
| 63 | |
| 64 | def __call__(self, *args, **kwargs): |
| 65 | key, _ = self.func.parse_args(*args, **kwargs) |
| 66 | kernel = self._kernel_cache.get(key) |
| 67 | if kernel is None: |
| 68 | kernel = "compiled" |
| 69 | self._kernel_cache[key] = kernel |
| 70 | return kernel |
| 71 | |
| 72 | fake_kernel = cast(Any, ModuleType("tilelang.jit.kernel")) |
| 73 | fake_kernel.JITKernel = FakeJITKernel |
| 74 | |
| 75 | fake_jit = cast(Any, ModuleType("tilelang.jit")) |
| 76 | fake_jit.JITImpl = FakeJITImpl |
| 77 | fake_jit.kernel = fake_kernel |
| 78 | |
| 79 | fake_tilelang = cast(Any, ModuleType("tilelang")) |
| 80 | fake_tilelang.jit = fake_jit |
| 81 | |
| 82 | return { |
| 83 | "tilelang": fake_tilelang, |
| 84 | "tilelang.jit": fake_jit, |
| 85 | "tilelang.jit.kernel": fake_kernel, |
| 86 | } |
| 87 | |
| 88 | |
| 89 | @contextmanager |
| 90 | def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile): |
| 91 | """Patch the Triton and CuTeDSL imports touched by ``jit_monitor.activate``.""" |
| 92 | fake_triton = cast(Any, ModuleType("triton")) |
| 93 | fake_triton.knobs = fake_knobs |
| 94 | with ( |
| 95 | mock.patch.dict( |
| 96 | sys.modules, |
| 97 | { |
| 98 | "triton": fake_triton, |
| 99 | **_fake_cute_import_modules(cute_compile), |
| 100 | **_fake_tilelang_import_modules(), |
| 101 | }, |
| 102 | ), |
| 103 | mock.patch.object(jit_monitor, "HAS_TRITON", True), |
| 104 | ): |
| 105 | yield |
| 106 | |
| 107 | |
| 108 | def _triton_hook_kwargs(name: str): |
| 109 | return dict( |
| 110 | key="k", |
| 111 | repr="r", |
| 112 | fn=SimpleNamespace(name=name), |
| 113 | compile=lambda: None, |
| 114 | is_manual_warmup=False, |
| 115 | already_compiled=False, |
| 116 | ) |
| 117 | |
| 118 | |
| 119 | # ------------------------------------------------------------------ |
| 120 | # activate() |
| 121 | # ------------------------------------------------------------------ |
| 122 | |
| 123 | |
| 124 | def test_activate_sets_active(): |
| 125 | assert not jit_monitor.is_active() |
| 126 | with _patch_jit_modules(_make_fake_knobs()): |
| 127 | jit_monitor.activate() |
| 128 | assert jit_monitor.is_active() |
| 129 | |
| 130 | |
| 131 | def test_activate_is_idempotent(): |
| 132 | fake = _make_fake_knobs() |
| 133 | with _patch_jit_modules(fake): |
| 134 | jit_monitor.activate() |
| 135 | first_hook = fake.runtime.jit_post_compile_hook |
| 136 | jit_monitor.activate() |
| 137 | assert fake.runtime.jit_post_compile_hook is first_hook |
| 138 | |
| 139 | |
| 140 | def test_activate_logs_info(): |
| 141 | with ( |
| 142 | mock.patch.object(jit_monitor.logger, "info") as m, |
| 143 | _patch_jit_modules(_make_fake_knobs()), |
| 144 | ): |
| 145 | jit_monitor.activate() |
| 146 | m.assert_called_once() |
| 147 | assert "Kernel JIT monitor activated" in m.call_args[0][0] |
| 148 | |
| 149 | |
| 150 | def test_activate_rejects_unknown_mode(): |
| 151 | with pytest.raises(ValueError, match="Unsupported JIT monitor mode"): |
| 152 | jit_monitor.activate(mode="panic") |
| 153 | |
| 154 | |
| 155 | def test_activate_without_triton(): |
| 156 | with mock.patch.object(jit_monitor, "HAS_TRITON", False): |
| 157 | jit_monitor.activate() |
| 158 | assert jit_monitor.is_active() |
| 159 | |
| 160 | |
| 161 | # ------------------------------------------------------------------ |
| 162 | # Triton autotuning print |
| 163 | # ------------------------------------------------------------------ |
| 164 | |
| 165 | |
| 166 | def test_autotuning_print_is_enabled(): |
| 167 | fake = _make_fake_knobs(autotuning_print=False) |
| 168 | with _patch_jit_modules(fake): |
| 169 | jit_monitor.activate() |
| 170 | assert fake.autotuning.print is True |
| 171 | |
| 172 | |
| 173 | def test_autotuning_print_respects_user_opt_out(): |
| 174 | fake = _make_fake_knobs(autotuning_print=False) |
| 175 | with ( |
| 176 | mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "0"}), |
| 177 | _patch_jit_modules(fake), |
| 178 | ): |
| 179 | jit_monitor.activate() |
| 180 | assert fake.autotuning.print is False |
| 181 | |
| 182 | |
| 183 | def test_autotuning_print_noop_when_user_already_enabled(): |
| 184 | fake = _make_fake_knobs(autotuning_print=True) |
| 185 | with ( |
| 186 | mock.patch.dict(os.environ, {"TRITON_PRINT_AUTOTUNING": "1"}), |
| 187 | _patch_jit_modules(fake), |
| 188 | ): |
| 189 | jit_monitor.activate() |
| 190 | assert fake.autotuning.print is True |
| 191 | |
| 192 | |
| 193 | # ------------------------------------------------------------------ |
| 194 | # Triton JIT hook |
| 195 | # ------------------------------------------------------------------ |
| 196 | |
| 197 | |
| 198 | def test_triton_hook_is_registered(): |
| 199 | fake = _make_fake_knobs() |
| 200 | assert fake.runtime.jit_post_compile_hook is None |
| 201 | with _patch_jit_modules(fake): |
| 202 | jit_monitor.activate() |
| 203 | assert fake.runtime.jit_post_compile_hook is not None |
| 204 | |
| 205 | |
| 206 | def test_triton_hook_logs_warning(): |
| 207 | fake = _make_fake_knobs() |
| 208 | with _patch_jit_modules(fake): |
| 209 | jit_monitor.activate() |
| 210 | |
| 211 | hook = fake.runtime.jit_post_compile_hook |
| 212 | |
| 213 | with ( |
| 214 | mock.patch.object(jit_monitor.logger, "warning_once") as m, |
| 215 | mock.patch.object(jit_monitor.logger, "warning") as warning, |
| 216 | ): |
| 217 | hook(**_triton_hook_kwargs("test_kernel")) |
| 218 | |
| 219 | m.assert_called_once() |
| 220 | warning.assert_not_called() |
| 221 | msg = m.call_args[0][0] % m.call_args[0][1:] |
| 222 | assert "Triton kernel JIT compilation during inference" in msg |
| 223 | assert "test_kernel" in msg |
| 224 | |
| 225 | |
| 226 | def test_triton_hook_chains_existing_hook(): |
| 227 | existing = mock.MagicMock(return_value="existing_result") |
| 228 | fake = _make_fake_knobs(jit_hook=existing) |
| 229 | with _patch_jit_modules(fake): |
| 230 | jit_monitor.activate() |
| 231 | |
| 232 | hook = fake.runtime.jit_post_compile_hook |
| 233 | result = hook(**_triton_hook_kwargs("chained_kernel")) |
| 234 | |
| 235 | existing.assert_called_once() |
| 236 | assert result == "existing_result" |
| 237 | |
| 238 | |
| 239 | def test_triton_hook_works_without_existing_hook(): |
| 240 | fake = _make_fake_knobs(jit_hook=None) |
| 241 | with _patch_jit_modules(fake): |
| 242 | jit_monitor.activate() |
| 243 | |
| 244 | hook = fake.runtime.jit_post_compile_hook |
| 245 | assert hook(**_triton_hook_kwargs("solo_kernel")) is None |
| 246 | |
| 247 | |
| 248 | def test_triton_hook_error_mode_raises(): |
| 249 | fake = _make_fake_knobs() |
| 250 | with _patch_jit_modules(fake): |
| 251 | jit_monitor.activate(mode="error") |
| 252 | |
| 253 | hook = fake.runtime.jit_post_compile_hook |
| 254 | with pytest.raises(RuntimeError, match="Triton kernel JIT compilation"): |
| 255 | hook(**_triton_hook_kwargs("error_kernel")) |
| 256 | |
| 257 | |
| 258 | # ------------------------------------------------------------------ |
| 259 | # CuTeDSL hook |
| 260 | # ------------------------------------------------------------------ |
| 261 | |
| 262 | |
| 263 | def test_cutedsl_compile_logs_warning(): |
| 264 | with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile): |
| 265 | import cutlass.cute as cute |
| 266 | |
| 267 | jit_monitor.activate() |
| 268 | with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: |
| 269 | result = cute.compile(lambda: None, "arg", option=True) |
| 270 | |
| 271 | assert result == "compiled" |
| 272 | warning_once.assert_called_once() |
| 273 | msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] |
| 274 | assert "CuTeDSL JIT compilation during inference" in msg |
| 275 | |
| 276 | |
| 277 | def test_cutedsl_compile_logs_verbose_warning(): |
| 278 | with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile): |
| 279 | import cutlass.cute as cute |
| 280 | |
| 281 | jit_monitor.activate(verbose=True) |
| 282 | with mock.patch.object(jit_monitor.logger, "warning") as warning: |
| 283 | result = cute.compile(lambda: None, "arg", option=True) |
| 284 | |
| 285 | assert result == "compiled" |
| 286 | warning.assert_called_once() |
| 287 | msg = warning.call_args[0][0] % warning.call_args[0][1:] |
| 288 | assert "CuTeDSL JIT compilation during inference" in msg |
| 289 | |
| 290 | |
| 291 | def test_cutedsl_error_mode_raises(): |
| 292 | with _patch_jit_modules(_make_fake_knobs(), cute_compile=_fake_cute_compile): |
| 293 | import cutlass.cute as cute |
| 294 | |
| 295 | jit_monitor.activate(mode="error") |
| 296 | with pytest.raises(RuntimeError, match="CuTeDSL JIT compilation"): |
| 297 | cute.compile(lambda: None, "arg", option=True) |
| 298 | |
| 299 | |
| 300 | def test_cutedsl_subscripted_compile_is_monitored(): |
| 301 | """``cute.compile[options](...)`` (flashinfer >= 0.6.14) must work.""" |
| 302 | |
| 303 | class FakeCompileCallable: |
| 304 | def __getitem__(self, options): |
| 305 | return self |
| 306 | |
| 307 | def __call__(self, *args, **kwargs): |
| 308 | return "compiled" |
| 309 | |
| 310 | with _patch_jit_modules(_make_fake_knobs(), cute_compile=FakeCompileCallable()): |
| 311 | import cutlass.cute as cute |
| 312 | |
| 313 | jit_monitor.activate() |
| 314 | with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: |
| 315 | result = cute.compile[("opt_level", 3)](lambda: None, "arg") |
| 316 | |
| 317 | assert result == "compiled" |
| 318 | warning_once.assert_called_once() |
| 319 | |
| 320 | |
| 321 | # ------------------------------------------------------------------ |
| 322 | # TileLang hook |
| 323 | # ------------------------------------------------------------------ |
| 324 | |
| 325 | |
| 326 | @pytest.mark.skipif( |
| 327 | current_platform.is_rocm(), |
| 328 | reason="TileLang JIT monitoring is disabled on ROCm", |
| 329 | ) |
| 330 | def test_tilelang_jit_kernel_logs_warning(): |
| 331 | with _patch_jit_modules(_make_fake_knobs()): |
| 332 | from tilelang.jit.kernel import JITKernel |
| 333 | |
| 334 | func = SimpleNamespace(attrs={"global_symbol": "tl_kernel"}) |
| 335 | jit_monitor.activate() |
| 336 | with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: |
| 337 | JITKernel(func=func, out_idx=None, execution_backend="tvm_ffi") |
| 338 | |
| 339 | warning_once.assert_called_once() |
| 340 | msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] |
| 341 | assert "TileLang JIT compilation during inference" in msg |
| 342 | assert "tl_kernel" in msg |
| 343 | |
| 344 | |
| 345 | @pytest.mark.skipif( |
| 346 | current_platform.is_rocm(), |
| 347 | reason="TileLang JIT monitoring is disabled on ROCm", |
| 348 | ) |
| 349 | def test_tilelang_jit_impl_logs_warning(): |
| 350 | with _patch_jit_modules(_make_fake_knobs()): |
| 351 | from tilelang.jit import JITImpl |
| 352 | |
| 353 | def tilelang_fn( |
| 354 | gemm_out_mul, |
| 355 | hidden_size: int, |
| 356 | n_splits: int = 1, |
| 357 | hc_mult: int = 4, |
| 358 | ): |
| 359 | return None |
| 360 | |
| 361 | class FakeFunc: |
| 362 | orig_func = tilelang_fn |
| 363 | |
| 364 | def parse_args(self, *args, **kwargs): |
| 365 | return ( |
| 366 | ( |
| 367 | "tilelang_key", |
| 368 | kwargs["hidden_size"], |
| 369 | kwargs.get("n_splits", 1), |
| 370 | ), |
| 371 | {}, |
| 372 | ) |
| 373 | |
| 374 | def set_mode(self, mode): |
| 375 | self.mode = mode |
| 376 | |
| 377 | tensor = SimpleNamespace( |
| 378 | shape=(2, 16, 24), |
| 379 | dtype="float32", |
| 380 | device="cuda:0", |
| 381 | ) |
| 382 | impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn)) |
| 383 | |
| 384 | jit_monitor.activate() |
| 385 | with ( |
| 386 | mock.patch.object(jit_monitor.logger, "warning_once") as warning_once, |
| 387 | mock.patch.object(jit_monitor.logger, "warning") as warning, |
| 388 | ): |
| 389 | impl(tensor, hidden_size=7168, n_splits=2) |
| 390 | |
| 391 | warning_once.assert_called_once() |
| 392 | warning.assert_not_called() |
| 393 | msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:] |
| 394 | assert "TileLang JIT compilation during inference" in msg |
| 395 | assert "tilelang_fn" in msg |
| 396 | |
| 397 | |
| 398 | @pytest.mark.skipif( |
| 399 | current_platform.is_rocm(), |
| 400 | reason="TileLang JIT monitoring is disabled on ROCm", |
| 401 | ) |
| 402 | def test_tilelang_jit_impl_does_not_log_on_cache_hit(): |
| 403 | with _patch_jit_modules(_make_fake_knobs()): |
| 404 | from tilelang.jit import JITImpl |
| 405 | |
| 406 | def tilelang_fn(gemm_out_mul, n_splits: int = 1): |
| 407 | return None |
| 408 | |
| 409 | class FakeFunc: |
| 410 | orig_func = tilelang_fn |
| 411 | |
| 412 | def parse_args(self, *args, **kwargs): |
| 413 | return (("tilelang_key", kwargs.get("n_splits", 1)), {}) |
| 414 | |
| 415 | def set_mode(self, mode): |
| 416 | self.mode = mode |
| 417 | |
| 418 | tensor = SimpleNamespace(shape=(2, 16, 24), dtype="float32") |
| 419 | impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn)) |
| 420 | |
| 421 | jit_monitor.activate() |
| 422 | with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once: |
| 423 | impl(tensor, n_splits=2) |
| 424 | impl(tensor, n_splits=2) |
| 425 | |
| 426 | warning_once.assert_called_once() |
| 427 | |
| 428 | |
| 429 | @pytest.mark.skipif( |
| 430 | current_platform.is_rocm(), |
| 431 | reason="TileLang JIT monitoring is disabled on ROCm", |
| 432 | ) |
| 433 | def test_tilelang_from_database_does_not_log(): |
| 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 | |
| 445 | @pytest.mark.skipif( |
| 446 | current_platform.is_rocm(), |
| 447 | reason="TileLang JIT monitoring is disabled on ROCm", |
| 448 | ) |
| 449 | def test_tilelang_error_mode_raises(): |
| 450 | with _patch_jit_modules(_make_fake_knobs()): |
| 451 | from tilelang.jit.kernel import JITKernel |
| 452 | |
| 453 | func = SimpleNamespace(attrs={"global_symbol": "error_tl_kernel"}) |
| 454 | jit_monitor.activate(mode="error") |
| 455 | with pytest.raises(RuntimeError, match="TileLang JIT compilation"): |
| 456 | JITKernel(func=func) |