[Feat] Add runtime monitor for post-warmup TileLang compilation (#46718)
Signed-off-by: LopezCastroRoberto <rocastro@redhat.com> Signed-off-by: Roberto L. Castro <38211239+LopezCastroRoberto@users.noreply.github.com>
Roberto L. Castro committed
Jul 9, 2026 at 00:11 UTC
5f8597562467d3b0495e300ebd902ba1aa781914
2 files changed
+367
-1
tests/test_jit_monitor.py
+155
@@ -1,5 +1,6 @@
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
@@ -19,11 +20,15 @@ def _reset_monitor():
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
# ------------------------------------------------------------------
@@ -55,6 +60,45 @@ 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``."""
@@ -66,6 +110,7 @@ def _patch_jit_modules(fake_knobs, *, cute_compile=_fake_cute_compile):
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),
@@ -276,6 +321,116 @@ class TestCuTeDSLHook:
321
cute.compile(lambda: None, "arg", option=True)
322
323
324
+class TestTileLangHook:
325
+ def test_jit_kernel_logs_warning(self):
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
+ def test_jit_impl_logs_warning(self):
340
+ with _patch_jit_modules(_make_fake_knobs()):
341
+ from tilelang.jit import JITImpl
342
+
343
+ def tilelang_fn(
344
+ gemm_out_mul,
345
+ hidden_size: int,
346
+ n_splits: int = 1,
347
+ hc_mult: int = 4,
348
+ ):
349
+ return None
350
+
351
+ class FakeFunc:
352
+ orig_func = tilelang_fn
353
+
354
+ def parse_args(self, *args, **kwargs):
355
+ return (
356
+ (
357
+ "tilelang_key",
358
+ kwargs["hidden_size"],
359
+ kwargs.get("n_splits", 1),
360
+ ),
361
+ {},
362
+ )
363
+
364
+ def set_mode(self, mode):
365
+ self.mode = mode
366
+
367
+ tensor = SimpleNamespace(
368
+ shape=(2, 16, 24),
369
+ dtype="float32",
370
+ device="cuda:0",
371
+ )
372
+ impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn))
373
+
374
+ jit_monitor.activate()
375
+ with (
376
+ mock.patch.object(jit_monitor.logger, "warning_once") as warning_once,
377
+ mock.patch.object(jit_monitor.logger, "warning") as warning,
378
+ ):
379
+ impl(tensor, hidden_size=7168, n_splits=2)
380
+
381
+ warning_once.assert_called_once()
382
+ warning.assert_not_called()
383
+ msg = warning_once.call_args[0][0] % warning_once.call_args[0][1:]
384
+ assert "TileLang JIT compilation during inference" in msg
385
+ assert "tilelang_fn" in msg
386
+
387
+ def test_jit_impl_does_not_log_on_cache_hit(self):
388
+ with _patch_jit_modules(_make_fake_knobs()):
389
+ from tilelang.jit import JITImpl
390
+
391
+ def tilelang_fn(gemm_out_mul, n_splits: int = 1):
392
+ return None
393
+
394
+ class FakeFunc:
395
+ orig_func = tilelang_fn
396
+
397
+ def parse_args(self, *args, **kwargs):
398
+ return (("tilelang_key", kwargs.get("n_splits", 1)), {})
399
+
400
+ def set_mode(self, mode):
401
+ self.mode = mode
402
+
403
+ tensor = SimpleNamespace(shape=(2, 16, 24), dtype="float32")
404
+ impl = JITImpl(FakeFunc(), inspect.signature(tilelang_fn))
405
+
406
+ jit_monitor.activate()
407
+ with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
408
+ impl(tensor, n_splits=2)
409
+ impl(tensor, n_splits=2)
410
+
411
+ warning_once.assert_called_once()
412
+
413
+ def test_from_database_does_not_log(self):
414
+ with _patch_jit_modules(_make_fake_knobs()):
415
+ from tilelang.jit.kernel import JITKernel
416
+
417
+ func = SimpleNamespace(attrs={"global_symbol": "cached_tl_kernel"})
418
+ jit_monitor.activate()
419
+ with mock.patch.object(jit_monitor.logger, "warning_once") as warning_once:
420
+ JITKernel(func=func, from_database=True)
421
+
422
+ warning_once.assert_not_called()
423
+
424
+ def test_error_mode_raises(self):
425
+ with _patch_jit_modules(_make_fake_knobs()):
426
+ from tilelang.jit.kernel import JITKernel
427
+
428
+ func = SimpleNamespace(attrs={"global_symbol": "error_tl_kernel"})
429
+ jit_monitor.activate(mode="error")
430
+ with pytest.raises(RuntimeError, match="TileLang JIT compilation"):
431
+ JITKernel(func=func)
432
+
433
+
434
# ------------------------------------------------------------------
435
# Integration tests (real Triton + GPU)
436
# ------------------------------------------------------------------
vllm/utils/jit_monitor.py
+212
-1
@@ -18,13 +18,16 @@ Currently monitors:
18
- Triton ``@triton.autotune`` cache misses (via ``knobs.autotuning.print``)
19
- Triton ``@triton.jit`` first-time compilations
20
(via ``knobs.runtime.jit_post_compile_hook``)
21
+- TileLang ``@tilelang.jit`` first-time compilations
22
"""
23
24
import contextlib
25
import functools
26
+import importlib
27
import os
28
from collections.abc import Iterator, Mapping
27
-from typing import Literal
29
+from contextlib import suppress
30
+from typing import Any, Literal, cast
31
32
from vllm.logger import init_logger
33
from vllm.triton_utils.importing import HAS_TRITON
@@ -36,6 +39,8 @@ _active: bool = False
39
_mode: JitMonitorMode = "warn"
40
_verbose: bool = False
41
_cutedsl_hook_installed: bool = False
42
+_tilelang_hook_installed: bool = False
43
+_tilelang_jitimpl_compile_depth: int = 0
44
45
46
def is_active() -> bool:
@@ -69,6 +74,7 @@ def activate(*, mode: JitMonitorMode = "warn", verbose: bool = False) -> None:
74
_setup_triton_autotuning_print()
75
_setup_triton_jit_hook()
76
_setup_cutedsl_jit_hook()
77
+ _setup_tilelang_jit_hook()
78
79
logger.info(
80
"Kernel JIT monitor activated; monitored JIT compilations during "
@@ -290,6 +296,211 @@ def _setup_cutedsl_jit_hook() -> None:
296
_cutedsl_hook_installed = True
297
298
299
+# ------------------------------------------------------------------
300
+# TileLang JIT compilation hook
301
+# ------------------------------------------------------------------
302
+
303
+
304
+def _tilelang_arg(
305
+ args: tuple[object, ...],
306
+ kwargs: Mapping[str, object],
307
+ index: int,
308
+ name: str,
309
+ default: object = None,
310
+) -> object:
311
+ if len(args) > index:
312
+ return args[index]
313
+ return kwargs.get(name, default)
314
+
315
+
316
+def _tilelang_kernel_name(func: object) -> str:
317
+ attrs = getattr(func, "attrs", None)
318
+ global_symbol = None
319
+ get_attr = getattr(attrs, "get", None)
320
+ if callable(get_attr):
321
+ try:
322
+ global_symbol = get_attr("global_symbol")
323
+ except Exception:
324
+ global_symbol = None
325
+ elif isinstance(attrs, Mapping):
326
+ global_symbol = attrs.get("global_symbol")
327
+
328
+ if global_symbol is not None:
329
+ return str(global_symbol)
330
+
331
+ name = getattr(func, "__name__", None)
332
+ if name is not None:
333
+ return str(name)
334
+ return func.__class__.__name__ if func is not None else "<unknown>"
335
+
336
+
337
+def _tilelang_call_kwargs(kwargs: Mapping[str, object]) -> dict[str, object]:
338
+ call_kwargs = dict(kwargs)
339
+ tune_params = call_kwargs.pop("__tune_params", {})
340
+ if isinstance(tune_params, Mapping):
341
+ call_kwargs.update(tune_params)
342
+ return call_kwargs
343
+
344
+
345
+def _tilelang_cache_miss_key(
346
+ jit_impl: object,
347
+ args: tuple[object, ...],
348
+ kwargs: Mapping[str, object],
349
+) -> object | None:
350
+ if kwargs.get("__return_compile_arguments", False):
351
+ return None
352
+
353
+ call_kwargs = _tilelang_call_kwargs(kwargs)
354
+ func = getattr(jit_impl, "func", None)
355
+ parse_args = getattr(func, "parse_args", None)
356
+ if not callable(parse_args):
357
+ return None
358
+
359
+ try:
360
+ if getattr(jit_impl, "mode", None) == "auto":
361
+ impl = cast(Any, jit_impl)
362
+ mode = impl._infer_jit_mode(*args, **call_kwargs)
363
+ impl.mode = mode
364
+ if func is not None:
365
+ func.set_mode(mode)
366
+ key, _ = parse_args(*args, **call_kwargs)
367
+ except Exception:
368
+ return None
369
+
370
+ cache = getattr(jit_impl, "_kernel_cache", {})
371
+ if isinstance(cache, Mapping):
372
+ return key if key not in cache else None
373
+
374
+ get = getattr(cache, "get", None)
375
+ if not callable(get):
376
+ return key
377
+ try:
378
+ return key if get(key) is None else None
379
+ except Exception:
380
+ return key
381
+
382
+
383
+def _format_tilelang_runtime_shapes(
384
+ jit_impl: object,
385
+ args: tuple[object, ...],
386
+ kwargs: Mapping[str, object],
387
+) -> str:
388
+ signature = getattr(jit_impl, "signature", None)
389
+ bind_partial = getattr(signature, "bind_partial", None)
390
+ if not callable(bind_partial):
391
+ return "{}"
392
+
393
+ try:
394
+ bound = bind_partial(*args, **_tilelang_call_kwargs(kwargs))
395
+ bound.apply_defaults()
396
+ except Exception:
397
+ return "{}"
398
+
399
+ items = []
400
+ for name, value in bound.arguments.items():
401
+ if str(name).startswith("__"):
402
+ continue
403
+ if not (hasattr(value, "shape") and hasattr(value, "dtype")):
404
+ continue
405
+ shape = getattr(value, "shape", None)
406
+ if shape is not None:
407
+ with suppress(TypeError):
408
+ shape = tuple(shape)
409
+ items.append(f"{name}={_safe_repr(shape, max_len=80)}")
410
+ return "{" + ", ".join(items) + "}"
411
+
412
+
413
+def _format_verbose_tilelang_compile_details(
414
+ jit_impl: object,
415
+ args: tuple[object, ...],
416
+ kwargs: Mapping[str, object],
417
+ cache_key: object,
418
+) -> str:
419
+ return (
420
+ f"cache_key={_safe_repr(cache_key, max_len=400)}; "
421
+ f"runtime_shapes={_format_tilelang_runtime_shapes(jit_impl, args, kwargs)}"
422
+ )
423
+
424
+
425
+def _log_tilelang_jit_compile(
426
+ fn_name: str,
427
+ detail: str | None = None,
428
+) -> None:
429
+ _handle_jit_event(
430
+ backend="TileLang",
431
+ event="JIT compilation",
432
+ fn_name=fn_name,
433
+ detail=detail,
434
+ )
435
+
436
+
437
+def _setup_tilelang_jit_hook() -> None:
438
+ """Wrap TileLang JIT entry points to warn on compilation."""
439
+ global _tilelang_hook_installed
440
+ if _tilelang_hook_installed:
441
+ return
442
+
443
+ try:
444
+ tilelang_kernel = importlib.import_module("tilelang.jit.kernel")
445
+ except Exception:
446
+ logger.debug("TileLang is not available; skipping TileLang JIT monitor.")
447
+ return
448
+
449
+ jit_kernel_cls = getattr(tilelang_kernel, "JITKernel", None)
450
+ if jit_kernel_cls is None:
451
+ logger.debug(
452
+ "TileLang JITKernel is unavailable; skipping TileLang JIT monitor."
453
+ )
454
+ return
455
+
456
+ try:
457
+ tilelang_jit = importlib.import_module("tilelang.jit")
458
+ except Exception:
459
+ tilelang_jit = None
460
+ jit_impl_cls = getattr(tilelang_jit, "JITImpl", None)
461
+ original_init = jit_kernel_cls.__init__
462
+
463
+ @functools.wraps(original_init)
464
+ def _init_with_monitor(self, *args, **kwargs):
465
+ from_database = bool(_tilelang_arg(args, kwargs, 7, "from_database", False))
466
+ if not from_database and _tilelang_jitimpl_compile_depth == 0:
467
+ func = _tilelang_arg(args, kwargs, 0, "func")
468
+ _log_tilelang_jit_compile(_tilelang_kernel_name(func))
469
+ return original_init(self, *args, **kwargs)
470
+
471
+ jit_kernel_cls.__init__ = _init_with_monitor
472
+
473
+ if jit_impl_cls is not None:
474
+ original_call = jit_impl_cls.__call__
475
+
476
+ @functools.wraps(original_call)
477
+ def _call_with_monitor(self, *args, **kwargs):
478
+ global _tilelang_jitimpl_compile_depth
479
+ cache_key = _tilelang_cache_miss_key(self, args, kwargs)
480
+ if cache_key is None:
481
+ return original_call(self, *args, **kwargs)
482
+
483
+ _tilelang_jitimpl_compile_depth += 1
484
+ try:
485
+ detail = None
486
+ if _verbose:
487
+ detail = _format_verbose_tilelang_compile_details(
488
+ self, args, kwargs, cache_key
489
+ )
490
+ func = getattr(self, "func", None)
491
+ orig_func = getattr(func, "orig_func", None)
492
+ _log_tilelang_jit_compile(
493
+ _tilelang_kernel_name(orig_func or func), detail
494
+ )
495
+ return original_call(self, *args, **kwargs)
496
+ finally:
497
+ _tilelang_jitimpl_compile_depth -= 1
498
+
499
+ jit_impl_cls.__call__ = _call_with_monitor
500
+
501
+ _tilelang_hook_installed = True
502
+
503
+
504
@contextlib.contextmanager
505
def numba_workqueue_threading_layer() -> Iterator[None]:
506
"""Force numba's fork-safe `workqueue` threading layer for this block.