[Bugfix][ROCm][GLM-5.3-Flash] Apply deferred tilelang.jit already on attribute access (#57192)

Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>

tests/kernels/test_mhc_tilelang_jit.py

@@ -24,19 +24,26 @@ class _PassConfigKey:
 def _install_tilelang_stub(
     monkeypatch: pytest.MonkeyPatch,
 ) -> dict[str, int]:
-    calls = {"jit_decorate": 0, "compiled_call": 0}
+    calls = {"jit_decorate": 0, "compiled_call": 0, "compiled_compile": 0}
 
     tilelang: Any = ModuleType("tilelang")
 
+    class _JitImpl:
+        def __init__(self, func: Any) -> None:
+            self.func = func
+
+        def __call__(self, *args: Any, **kw: Any) -> Any:
+            calls["compiled_call"] += 1
+            return self.func.__name__
+
+        def compile(self, *args: Any, **kw: Any) -> Any:
+            calls["compiled_compile"] += 1
+            return self.func.__name__
+
     def jit(**kwargs: Any) -> Any:
         def decorate(func: Any) -> Any:
             calls["jit_decorate"] += 1
-
-            def compiled(*args: Any, **kw: Any) -> Any:
-                calls["compiled_call"] += 1
-                return func.__name__
-
-            return compiled
+            return _JitImpl(func)
 
         return decorate
 
@@ -78,6 +85,21 @@ def test_tilelang_jit_decorator_is_lazy_only_on_rocm(
     assert calls["compiled_call"] == 1
 
 
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="Test requires ROCm")
+def test_tilelang_jit_proxies_compile_only_warmup(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    calls = _install_tilelang_stub(monkeypatch)
+    module_name = "vllm.model_executor.kernels.mhc.tilelang_kernels"
+    monkeypatch.delitem(sys.modules, module_name, raising=False)
+    module = importlib.import_module(module_name)
+
+    assert module.mhc_post_tilelang.compile() == "mhc_post_tilelang"
+    assert calls["jit_decorate"] == 1
+    assert calls["compiled_compile"] == 1
+    assert calls["compiled_call"] == 0
+
+
 @pytest.mark.skipif(not current_platform.is_rocm(), reason="Test requires ROCm")
 def test_deepseek_v4_import_and_jit_monitor_do_not_hijack_hip_symbols() -> None:
     if importlib.util.find_spec("tilelang") is None:

vllm/tilelang_utils/__init__.py

@@ -61,8 +61,41 @@ def _get_pass_configs() -> dict[Any, Any]:
     return pass_configs
 
 
+class _DeferredTileLangJitKernel:
+    """Stand-in for a `tilelang.jit` kernel that decorates on first use.
+
+    Both attribute access and calling apply the decoration. Required for
+    compile-only JIT warmup on platforms that defer import of tilelang.
+    """
+
+    _kernel_function: Callable[..., Any] | None = None
+    _jit_kernel: Any = None
+
+    def __init__(self, kernel_function: Callable[..., Any]) -> None:
+        self._kernel_function = kernel_function
+        functools.update_wrapper(self, kernel_function)
+
+    def _ensure_jit_kernel(self) -> Any:
+        if self._jit_kernel is None:
+            _ensure_tilelang_imported()
+            kernel_function = self._kernel_function
+            assert kernel_function is not None
+            kernel_function.__globals__["tilelang"] = tilelang
+            kernel_function.__globals__["T"] = T
+            self._jit_kernel = tilelang.jit(pass_configs=_get_pass_configs())(
+                kernel_function
+            )
+        return self._jit_kernel
+
+    def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        return self._ensure_jit_kernel()(*args, **kwargs)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._ensure_jit_kernel(), name)
+
+
 def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]:
-    """Apply `tilelang.jit`, deferring until first call on ROCm.
+    """Apply `tilelang.jit`, deferring until first use on ROCm.
 
     ROCm defers JIT decoration so importing the caller's module does not
     require TileLang immediately. CUDA keeps the eager decoration behavior.
@@ -75,18 +108,4 @@ def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]:
         _ensure_tilelang_imported()
         return tilelang.jit(pass_configs=_get_pass_configs())(kernel_function)
 
-    compiled_kernel: Callable[..., Any] | None = None
-
-    @functools.wraps(kernel_function)
-    def wrapper(*args: Any, **kwargs: Any) -> Any:
-        nonlocal compiled_kernel
-        if compiled_kernel is None:
-            _ensure_tilelang_imported()
-            kernel_function.__globals__["tilelang"] = tilelang
-            kernel_function.__globals__["T"] = T
-            compiled_kernel = tilelang.jit(pass_configs=_get_pass_configs())(
-                kernel_function
-            )
-        return compiled_kernel(*args, **kwargs)
-
-    return wrapper
+    return _DeferredTileLangJitKernel(kernel_function)