[Bugfix][ROCm][GLM-5.3-Flash] Apply deferred tilelang.jit already on attribute access (#57192)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Simon Danielsson committed
Sep 17, 2026 at 14:02 UTC
667b26e50bc3287793efd92672c520fa101ecf80
2 files changed
+64
-23
tests/kernels/test_mhc_tilelang_jit.py
+29
-7
@@ -24,19 +24,26 @@ class _PassConfigKey:
24
def _install_tilelang_stub(
25
monkeypatch: pytest.MonkeyPatch,
26
) -> dict[str, int]:
27
- calls = {"jit_decorate": 0, "compiled_call": 0}
27
+ calls = {"jit_decorate": 0, "compiled_call": 0, "compiled_compile": 0}
28
29
tilelang: Any = ModuleType("tilelang")
30
31
+ class _JitImpl:
32
+ def __init__(self, func: Any) -> None:
33
+ self.func = func
34
+
35
+ def __call__(self, *args: Any, **kw: Any) -> Any:
36
+ calls["compiled_call"] += 1
37
+ return self.func.__name__
38
+
39
+ def compile(self, *args: Any, **kw: Any) -> Any:
40
+ calls["compiled_compile"] += 1
41
+ return self.func.__name__
42
+
43
def jit(**kwargs: Any) -> Any:
44
def decorate(func: Any) -> Any:
45
calls["jit_decorate"] += 1
34
-
35
- def compiled(*args: Any, **kw: Any) -> Any:
36
- calls["compiled_call"] += 1
37
- return func.__name__
38
-
39
- return compiled
46
+ return _JitImpl(func)
47
48
return decorate
49
@@ -78,6 +85,21 @@ def test_tilelang_jit_decorator_is_lazy_only_on_rocm(
85
assert calls["compiled_call"] == 1
86
87
88
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="Test requires ROCm")
89
+def test_tilelang_jit_proxies_compile_only_warmup(
90
+ monkeypatch: pytest.MonkeyPatch,
91
+) -> None:
92
+ calls = _install_tilelang_stub(monkeypatch)
93
+ module_name = "vllm.model_executor.kernels.mhc.tilelang_kernels"
94
+ monkeypatch.delitem(sys.modules, module_name, raising=False)
95
+ module = importlib.import_module(module_name)
96
+
97
+ assert module.mhc_post_tilelang.compile() == "mhc_post_tilelang"
98
+ assert calls["jit_decorate"] == 1
99
+ assert calls["compiled_compile"] == 1
100
+ assert calls["compiled_call"] == 0
101
+
102
+
103
@pytest.mark.skipif(not current_platform.is_rocm(), reason="Test requires ROCm")
104
def test_deepseek_v4_import_and_jit_monitor_do_not_hijack_hip_symbols() -> None:
105
if importlib.util.find_spec("tilelang") is None:
vllm/tilelang_utils/__init__.py
+35
-16
@@ -61,8 +61,41 @@ def _get_pass_configs() -> dict[Any, Any]:
61
return pass_configs
62
63
64
+class _DeferredTileLangJitKernel:
65
+ """Stand-in for a `tilelang.jit` kernel that decorates on first use.
66
+
67
+ Both attribute access and calling apply the decoration. Required for
68
+ compile-only JIT warmup on platforms that defer import of tilelang.
69
+ """
70
+
71
+ _kernel_function: Callable[..., Any] | None = None
72
+ _jit_kernel: Any = None
73
+
74
+ def __init__(self, kernel_function: Callable[..., Any]) -> None:
75
+ self._kernel_function = kernel_function
76
+ functools.update_wrapper(self, kernel_function)
77
+
78
+ def _ensure_jit_kernel(self) -> Any:
79
+ if self._jit_kernel is None:
80
+ _ensure_tilelang_imported()
81
+ kernel_function = self._kernel_function
82
+ assert kernel_function is not None
83
+ kernel_function.__globals__["tilelang"] = tilelang
84
+ kernel_function.__globals__["T"] = T
85
+ self._jit_kernel = tilelang.jit(pass_configs=_get_pass_configs())(
86
+ kernel_function
87
+ )
88
+ return self._jit_kernel
89
+
90
+ def __call__(self, *args: Any, **kwargs: Any) -> Any:
91
+ return self._ensure_jit_kernel()(*args, **kwargs)
92
+
93
+ def __getattr__(self, name: str) -> Any:
94
+ return getattr(self._ensure_jit_kernel(), name)
95
+
96
+
97
def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]:
65
- """Apply `tilelang.jit`, deferring until first call on ROCm.
98
+ """Apply `tilelang.jit`, deferring until first use on ROCm.
99
100
ROCm defers JIT decoration so importing the caller's module does not
101
require TileLang immediately. CUDA keeps the eager decoration behavior.
@@ -75,18 +108,4 @@ def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]:
108
_ensure_tilelang_imported()
109
return tilelang.jit(pass_configs=_get_pass_configs())(kernel_function)
110
78
- compiled_kernel: Callable[..., Any] | None = None
79
-
80
- @functools.wraps(kernel_function)
81
- def wrapper(*args: Any, **kwargs: Any) -> Any:
82
- nonlocal compiled_kernel
83
- if compiled_kernel is None:
84
- _ensure_tilelang_imported()
85
- kernel_function.__globals__["tilelang"] = tilelang
86
- kernel_function.__globals__["T"] = T
87
- compiled_kernel = tilelang.jit(pass_configs=_get_pass_configs())(
88
- kernel_function
89
- )
90
- return compiled_kernel(*args, **kwargs)
91
-
92
- return wrapper
111
+ return _DeferredTileLangJitKernel(kernel_function)