From a58bdd0d4ec62764a8d6c5d4f9ed6a1dd4bf9b80 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Nicol=C3=B2=20Lucchesi?= Date: Thu, 8 Oct 2026 02:11:59 +0200 Subject: [PATCH] [Perf] JIT monitor: hook TileLang JITImpl.compile instead of __call__ (#58685) Signed-off-by: NickLucche --- tests/jit_monitor/test_hooks.py | 5 ++++- vllm/utils/jit_monitor.py | 18 ++++++++---------- 2 files changed, 12 insertions(+), 11 deletions(-) diff --git a/tests/jit_monitor/test_hooks.py b/tests/jit_monitor/test_hooks.py index 452fe96995..6d4f25314c 100644 --- a/tests/jit_monitor/test_hooks.py +++ b/tests/jit_monitor/test_hooks.py @@ -65,10 +65,13 @@ def _fake_tilelang_import_modules(): key, _ = self.func.parse_args(*args, **kwargs) kernel = self._kernel_cache.get(key) if kernel is None: - kernel = "compiled" + kernel = self.compile(*args, **kwargs) self._kernel_cache[key] = kernel return kernel + def compile(self, *args, **kwargs): + return "compiled" + fake_kernel = cast(Any, ModuleType("tilelang.jit.kernel")) fake_kernel.JITKernel = FakeJITKernel diff --git a/vllm/utils/jit_monitor.py b/vllm/utils/jit_monitor.py index ffb63bc493..387189c793 100644 --- a/vllm/utils/jit_monitor.py +++ b/vllm/utils/jit_monitor.py @@ -499,32 +499,30 @@ def _setup_tilelang_jit_hook() -> None: jit_kernel_cls.__init__ = _init_with_monitor if jit_impl_cls is not None: - original_call = jit_impl_cls.__call__ + # JITImpl.__call__ only calls compile() on a kernel-cache miss, so + # hooking compile() keeps the per-call hot path untouched. + original_compile = jit_impl_cls.compile - @functools.wraps(original_call) - def _call_with_monitor(self, *args, **kwargs): + @functools.wraps(original_compile) + def _compile_with_monitor(self, *args, **kwargs): global _tilelang_jitimpl_compile_depth - cache_key = _tilelang_cache_miss_key(self, args, kwargs) - if cache_key is None: - return original_call(self, *args, **kwargs) - _tilelang_jitimpl_compile_depth += 1 try: detail = None if _verbose: detail = _format_verbose_tilelang_compile_details( - self, args, kwargs, cache_key + self, args, kwargs, _tilelang_cache_miss_key(self, args, kwargs) ) func = getattr(self, "func", None) orig_func = getattr(func, "orig_func", None) _log_tilelang_jit_compile( _tilelang_kernel_name(orig_func or func), detail ) - return original_call(self, *args, **kwargs) + return original_compile(self, *args, **kwargs) finally: _tilelang_jitimpl_compile_depth -= 1 - jit_impl_cls.__call__ = _call_with_monitor + jit_impl_cls.compile = _compile_with_monitor _tilelang_hook_installed = True