@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c5b7c069a4

[ROCm] Defer `tilelang` import through its import `from vllm.tilelang_utils import tilelang` and relaxed `has_tilelang` (#51159)

Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

fxmarty-amd committed Aug 13, 2026 at 20:12 UTC c5b7c069a43e88dc1fb67e91286887cc664e2af9
8 files changed +313 -52
docs/assets/contributing/dockerfile-stages-dependency.png
Binary files a/docs/assets/contributing/dockerfile-stages-dependency.png and b/docs/assets/contributing/dockerfile-stages-dependency.png differ
tests/jit_monitor/test_hooks.py
+21
@@ -12,6 +12,7 @@ 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
@@ -322,6 +323,10 @@ def test_cutedsl_subscripted_compile_is_monitored():
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
@@ -337,6 +342,10 @@ def test_tilelang_jit_kernel_logs_warning():
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
@@ -386,6 +395,10 @@ def test_tilelang_jit_impl_logs_warning():
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
@@ -413,6 +426,10 @@ def test_tilelang_jit_impl_does_not_log_on_cache_hit():
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
@@ -425,6 +442,10 @@ def test_tilelang_from_database_does_not_log():
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
tests/kernels/test_mhc_tilelang_jit.py new
+128
@@ -0,0 +1,128 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import ctypes
5 +import importlib
6 +import importlib.util
7 +import sys
8 +from types import ModuleType
9 +from typing import Any
10 +
11 +import pytest
12 +
13 +from vllm.platforms import current_platform
14 +from vllm.utils import import_utils
15 +
16 +
17 +class _PassConfigKey:
18 + TL_DISABLE_WARP_SPECIALIZED = "disable_warp_specialized"
19 + TL_DISABLE_TMA_LOWER = "disable_tma_lower"
20 + TL_PTXAS_REGISTER_USAGE_LEVEL = "ptxas_register_usage_level"
21 +
22 +
23 +def _install_tilelang_stub(
24 + monkeypatch: pytest.MonkeyPatch,
25 +) -> dict[str, int]:
26 + calls = {"jit_decorate": 0, "compiled_call": 0}
27 +
28 + tilelang: Any = ModuleType("tilelang")
29 +
30 + def jit(**kwargs: Any) -> Any:
31 + def decorate(func: Any) -> Any:
32 + calls["jit_decorate"] += 1
33 +
34 + def compiled(*args: Any, **kw: Any) -> Any:
35 + calls["compiled_call"] += 1
36 + return func.__name__
37 +
38 + return compiled
39 +
40 + return decorate
41 +
42 + tilelang.PassConfigKey = _PassConfigKey
43 + tilelang.jit = jit
44 +
45 + monkeypatch.setattr(import_utils, "has_tilelang", lambda: True)
46 + monkeypatch.setitem(sys.modules, "tilelang", tilelang)
47 + monkeypatch.setitem(
48 + sys.modules, "tilelang.language", ModuleType("tilelang.language")
49 + )
50 + monkeypatch.delitem(sys.modules, "vllm.tilelang_utils", raising=False)
51 +
52 + return calls
53 +
54 +
55 +def test_tilelang_jit_decorator_is_lazy_only_on_rocm(
56 + monkeypatch: pytest.MonkeyPatch,
57 +) -> None:
58 + if not (current_platform.is_cuda() or current_platform.is_rocm()):
59 + pytest.skip("Test requires CUDA or ROCm")
60 +
61 + calls = _install_tilelang_stub(monkeypatch)
62 + module_name = "vllm.model_executor.kernels.mhc.tilelang_kernels"
63 + monkeypatch.delitem(sys.modules, module_name, raising=False)
64 + module = importlib.import_module(module_name)
65 +
66 + if current_platform.is_rocm():
67 + assert calls["jit_decorate"] == 0
68 + else:
69 + assert calls["jit_decorate"] > 0
70 +
71 + decorated_calls = calls["jit_decorate"]
72 + assert module.mhc_post_tilelang() == "mhc_post_tilelang"
73 + if current_platform.is_rocm():
74 + assert calls["jit_decorate"] == 1
75 + else:
76 + assert calls["jit_decorate"] == decorated_calls
77 + assert calls["compiled_call"] == 1
78 +
79 +
80 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="Test requires ROCm")
81 +def test_deepseek_v4_import_and_jit_monitor_do_not_hijack_hip_symbols(
82 + monkeypatch: pytest.MonkeyPatch,
83 +) -> None:
84 + if importlib.util.find_spec("tilelang") is None:
85 + pytest.skip("Test requires TileLang to be installed")
86 +
87 + class DlInfo(ctypes.Structure):
88 + _fields_ = [
89 + ("dli_fname", ctypes.c_char_p),
90 + ("dli_fbase", ctypes.c_void_p),
91 + ("dli_sname", ctypes.c_char_p),
92 + ("dli_saddr", ctypes.c_void_p),
93 + ]
94 +
95 + libdl = ctypes.CDLL("libdl.so.2")
96 + dlsym = libdl.dlsym
97 + dlsym.argtypes = [ctypes.c_void_p, ctypes.c_char_p]
98 + dlsym.restype = ctypes.c_void_p
99 + dladdr = libdl.dladdr
100 + dladdr.argtypes = [ctypes.c_void_p, ctypes.POINTER(DlInfo)]
101 + dladdr.restype = ctypes.c_int
102 +
103 + from vllm.model_executor.layers import mhc # noqa: F401
104 + from vllm.models import deepseek_v4 # noqa: F401
105 + from vllm.utils import jit_monitor
106 +
107 + monkeypatch.setattr(jit_monitor, "_active", False)
108 + monkeypatch.setattr(jit_monitor, "_setup_triton_autotuning_print", lambda: None)
109 + monkeypatch.setattr(jit_monitor, "_setup_triton_jit_hook", lambda: None)
110 + monkeypatch.setattr(jit_monitor, "_setup_cutedsl_jit_hook", lambda: None)
111 + monkeypatch.setattr(
112 + jit_monitor,
113 + "_setup_tilelang_jit_hook",
114 + lambda: pytest.fail("TileLang JIT monitor must not run on ROCm"),
115 + )
116 + jit_monitor.activate()
117 +
118 + assert not any(
119 + name == "tilelang" or name.startswith("tilelang.") for name in sys.modules
120 + )
121 +
122 + address = dlsym(None, b"hipFree")
123 + assert address is not None, "hipFree is not available in the global symbol table"
124 + info = DlInfo()
125 + assert dladdr(address, ctypes.byref(info))
126 + source = info.dli_fname.decode() if info.dli_fname else "<unknown>"
127 + assert "libhip_stub.so" not in source, source
128 + assert "libamdhip64.so" in source, source
tools/pre_commit/check_forbidden_imports.py
+32
@@ -92,6 +92,15 @@ CHECK_IMPORTS = {
92 ),
93 allowed_files={"vllm/triton_utils/importing.py"},
94 ),
95 + "tilelang": ForbiddenImport(
96 + pattern=r"^(from|import)\s+tilelang(\s|\.|$)",
97 + tip="Use 'from vllm.tilelang_utils import tilelang, T' instead.",
98 + allowed_pattern=re.compile(
99 + r"from\s+vllm\.tilelang_utils\s+import\s+"
100 + r"(tilelang|T|T, tilelang|tilelang, T)\b"
101 + ),
102 + allowed_files={"vllm/tilelang_utils/__init__.py"},
103 + ),
104 "huggingface_hub repo API": ForbiddenImport(
105 # Catch `from huggingface_hub import <fn>`, including parenthesized,
106 # multi-line imports.
@@ -182,6 +191,29 @@ def test_regex():
191 f"(expected {should_match}, got {result})"
192 )
193
194 + tilelang_cases = [
195 + # Should match
196 + ("import tilelang", True),
197 + ("import tilelang.language as T", True),
198 + ("from tilelang.jit import JITImpl", True),
199 + ("from tilelang.jit.kernel import JITKernel", True),
200 + # Should not match: indented (local) imports are allowed, mirroring
201 + # the "triton" rule, so mocked-module test imports are not flagged.
202 + (" import tilelang", False),
203 + (" from tilelang.jit import JITImpl", False),
204 + ("from vllm.tilelang_utils import tilelang", False),
205 + ("from vllm.tilelang_utils import T", False),
206 + ("from vllm.tilelang_utils import T, tilelang", False),
207 + ("from vllm.tilelang_utils import tilelang, T", False),
208 + ("import tilelang_kernels", False),
209 + ]
210 + for i, (content, should_match) in enumerate(tilelang_cases):
211 + result = matches("tilelang", content)
212 + assert result == should_match, (
213 + f"tilelang case {i} failed: {content!r} "
214 + f"(expected {should_match}, got {result})"
215 + )
216 +
217 hf_cases = [
218 # Should match
219 ("from huggingface_hub import snapshot_download", True),
vllm/model_executor/kernels/mhc/tilelang_kernels.py
+12 -49
@@ -1,29 +1,17 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 import math
7 from functools import cache
5 -from typing import TYPE_CHECKING, Any
8
9 import torch
10
11 from vllm.platforms import current_platform
10 -from vllm.utils.import_utils import has_tilelang
12 +from vllm.tilelang_utils import T, tilelang, tilelang_jit
13 from vllm.utils.math_utils import cdiv
14
13 -# TileLang is used for MHC on CUDA and ROCm. Keep non-GPU imports cheap so
14 -# registering the Python wrapper modules does not require TileLang everywhere.
15 -if TYPE_CHECKING or current_platform.is_cuda_alike():
16 - if not has_tilelang():
17 - raise ImportError(
18 - "tilelang is required for mhc but is not installed. Install it with "
19 - "`pip install tilelang`."
20 - )
21 - import tilelang
22 - import tilelang.language as T
23 -else:
24 - tilelang = None # type: ignore[assignment]
25 - T = None # type: ignore[assignment]
26 -
15 ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda()
16
17
@@ -40,18 +28,7 @@ def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
28 return split_k
29
30
43 -pass_configs: dict[tilelang.PassConfigKey, Any] = {
44 - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
45 - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
46 -}
47 -
48 -if current_platform.is_cuda():
49 - pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10
50 -
51 -
52 -@tilelang.jit(
53 - pass_configs=pass_configs,
54 -)
31 +@tilelang_jit
32 def mhc_pre_big_fuse_tilelang(
33 gemm_out_mul,
34 gemm_out_sqrsum,
@@ -191,9 +168,7 @@ def mhc_pre_big_fuse_tilelang(
168 # Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478
169
170
194 -@tilelang.jit(
195 - pass_configs=pass_configs,
196 -)
171 +@tilelang_jit
172 def mhc_pre_big_fuse_with_norm_tilelang(
173 gemm_out_mul,
174 gemm_out_sqrsum,
@@ -353,9 +328,7 @@ def mhc_pre_big_fuse_with_norm_tilelang(
328 T.pdl_trigger()
329
330
356 -@tilelang.jit(
357 - pass_configs=pass_configs,
358 -)
331 +@tilelang_jit
332 def mhc_pre_big_fuse_broadcast_with_norm_tilelang(
333 gemm_out_mul,
334 gemm_out_sqrsum,
@@ -517,9 +490,7 @@ def mhc_pre_big_fuse_broadcast_with_norm_tilelang(
490 T.pdl_trigger()
491
492
520 -@tilelang.jit(
521 - pass_configs=pass_configs,
522 -)
493 +@tilelang_jit
494 def mhc_fused_tilelang(
495 comb_mix,
496 residual_in,
@@ -640,9 +611,7 @@ def mhc_fused_tilelang(
611 T.pdl_trigger()
612
613
643 -@tilelang.jit(
644 - pass_configs=pass_configs,
645 -)
614 +@tilelang_jit
615 def mhc_post_tilelang(
616 a,
617 b,
@@ -695,9 +664,7 @@ def mhc_post_tilelang(
664 T.pdl_trigger()
665
666
698 -@tilelang.jit(
699 - pass_configs=pass_configs,
700 -)
667 +@tilelang_jit
668 def hc_prenorm_gemm_tilelang(
669 x,
670 fn,
@@ -781,9 +748,7 @@ def hc_prenorm_gemm_tilelang(
748 T.pdl_trigger()
749
750
784 -@tilelang.jit(
785 - pass_configs=pass_configs,
786 -)
751 +@tilelang_jit
752 def hc_prenorm_gemm_block_m_tilelang(
753 x,
754 fn,
@@ -878,9 +843,7 @@ def hc_prenorm_gemm_block_m_tilelang(
843 T.pdl_trigger()
844
845
881 -@tilelang.jit(
882 - pass_configs=pass_configs,
883 -)
846 +@tilelang_jit
847 def hc_head_fuse_tilelang(
848 residual,
849 fn,
vllm/tilelang_utils/__init__.py new
+91
@@ -0,0 +1,91 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from __future__ import annotations
5 +
6 +import functools
7 +from collections.abc import Callable
8 +from functools import cache
9 +from typing import TYPE_CHECKING, Any
10 +
11 +from vllm.platforms import current_platform
12 +from vllm.utils.import_utils import has_tilelang
13 +
14 +if TYPE_CHECKING or current_platform.is_cuda():
15 + if not has_tilelang():
16 + raise ImportError(
17 + "tilelang is required for mhc but is not installed. Install it with "
18 + "`pip install tilelang`."
19 + )
20 + import tilelang
21 + import tilelang.language as T
22 +else:
23 + tilelang = None # type: ignore[assignment]
24 + T = None # type: ignore[assignment]
25 +
26 +
27 +def _ensure_tilelang_imported() -> None:
28 + """Bind the `tilelang` and `T` module globals, importing them if needed.
29 +
30 + On ROCm, this runs on the first kernel call instead of at import time.
31 +
32 + Raises:
33 + ImportError: If TileLang is not installed.
34 + """
35 + global T, tilelang
36 +
37 + if tilelang is not None:
38 + return
39 + if not has_tilelang():
40 + raise ImportError(
41 + "tilelang is required for mhc but is not installed. Install it with "
42 + "`pip install tilelang`."
43 + )
44 + import tilelang as tilelang_module
45 + import tilelang.language as tilelang_language
46 +
47 + tilelang = tilelang_module
48 + T = tilelang_language
49 +
50 +
51 +@cache
52 +def _get_pass_configs() -> dict[Any, Any]:
53 + _ensure_tilelang_imported()
54 + pass_configs: dict[Any, Any] = {
55 + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
56 + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
57 + }
58 + if current_platform.is_cuda():
59 + pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10
60 + return pass_configs
61 +
62 +
63 +def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]:
64 + """Apply `tilelang.jit`, deferring until first call on ROCm.
65 +
66 + ROCm defers JIT decoration so importing the caller's module does not
67 + require TileLang immediately. CUDA keeps the eager decoration behavior.
68 +
69 + The kernel body parsed by TileLang references `T` as an unqualified
70 + global, so on the deferred ROCm path this rebinds `T`/`tilelang` in the
71 + decorated function's own module globals once they become available.
72 + """
73 + if not current_platform.is_rocm():
74 + _ensure_tilelang_imported()
75 + return tilelang.jit(pass_configs=_get_pass_configs())(kernel_function)
76 +
77 + compiled_kernel: Callable[..., Any] | None = None
78 +
79 + @functools.wraps(kernel_function)
80 + def wrapper(*args: Any, **kwargs: Any) -> Any:
81 + nonlocal compiled_kernel
82 + if compiled_kernel is None:
83 + _ensure_tilelang_imported()
84 + kernel_function.__globals__["tilelang"] = tilelang
85 + kernel_function.__globals__["T"] = T
86 + compiled_kernel = tilelang.jit(pass_configs=_get_pass_configs())(
87 + kernel_function
88 + )
89 + return compiled_kernel(*args, **kwargs)
90 +
91 + return wrapper
vllm/utils/import_utils.py
+22 -2
@@ -412,6 +412,21 @@ def _has_module(module_name: str) -> bool:
412 return True
413
414
415 +@cache
416 +def _has_module_spec(module_name: str) -> bool:
417 + """Return True if *module_name* is installed, without importing it.
418 +
419 + Unlike [`_has_module`][vllm.utils.import_utils._has_module], this only
420 + resolves the import spec. It therefore does not pay the import cost of
421 + heavyweight modules, at the price of not verifying that native
422 + dependencies (shared libraries, etc.) are satisfied. The result is cached.
423 + """
424 + try:
425 + return importlib.util.find_spec(module_name) is not None
426 + except Exception:
427 + return False
428 +
429 +
430 def has_deep_ep() -> bool:
431 """Whether the optional `deep_ep` package is available."""
432 return _has_module("deep_ep")
@@ -504,8 +519,13 @@ def has_triton_kernels() -> bool:
519
520 @cache
521 def has_tilelang() -> bool:
507 - """Whether the optional `tilelang` package is available."""
508 - if not _has_module("tilelang"):
522 + """Whether the optional `tilelang` package is available.
523 +
524 + Only the import spec is checked: importing `tilelang` is expensive, so
525 + callers must import it lazily at their point of use rather than relying
526 + on this function to have imported it already.
527 + """
528 + if not _has_module_spec("tilelang"):
529 return False
530 # ROCm-only guard, imported lazily to avoid loading rocm on CUDA.
531 from vllm.platforms import current_platform
vllm/utils/jit_monitor.py
+7 -1
@@ -30,6 +30,7 @@ from contextlib import suppress
30 from typing import Any, Literal, cast
31
32 from vllm.logger import init_logger
33 +from vllm.platforms import current_platform
34 from vllm.triton_utils.importing import HAS_TRITON
35
36 logger = init_logger(__name__)
@@ -74,7 +75,12 @@ def activate(*, mode: JitMonitorMode = "warn", verbose: bool = False) -> None:
75 _setup_triton_autotuning_print()
76 _setup_triton_jit_hook()
77 _setup_cutedsl_jit_hook()
77 - _setup_tilelang_jit_hook()
78 +
79 + # Refer to #51159. tilelang ships broken symbols
80 + # on rocm.
81 + # TODO: Remove the guard once tilelang upstream is fixed.
82 + if not current_platform.is_rocm():
83 + _setup_tilelang_jit_hook()
84
85 logger.info(
86 "Kernel JIT monitor activated; monitored JIT compilations during "