[go: up one dir, main page]
Skip to content

vllm.tilelang_utils

Functions:

  • tilelang_jit –

    Apply tilelang.jit, deferring until first use on ROCm.

_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.

Source code in vllm/tilelang_utils/__init__.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
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)

_ensure_tilelang_imported()

Bind the tilelang and T module globals, importing them if needed.

On ROCm, this runs on the first kernel call instead of at import time.

Raises:

Source code in vllm/tilelang_utils/__init__.py 
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def _ensure_tilelang_imported() -> None:
    """Bind the `tilelang` and `T` module globals, importing them if needed.

    On ROCm, this runs on the first kernel call instead of at import time.

    Raises:
        ImportError: If TileLang is not installed.

    """
    global T, tilelang

    if tilelang is not None:
        return
    if not has_tilelang():
        raise ImportError(
            "tilelang is required for mhc but is not installed. Install it with "
            "`pip install tilelang`."
        )
    import tilelang as tilelang_module
    import tilelang.language as tilelang_language

    tilelang = tilelang_module
    T = tilelang_language

tilelang_jit(kernel_function)

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.

The kernel body parsed by TileLang references T as an unqualified global, so on the deferred ROCm path this rebinds T/tilelang in the decorated function's own module globals once they become available.

Source code in vllm/tilelang_utils/__init__.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
def tilelang_jit(kernel_function: Callable[..., Any]) -> Callable[..., Any]:
    """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.

    The kernel body parsed by TileLang references `T` as an unqualified
    global, so on the deferred ROCm path this rebinds `T`/`tilelang` in the
    decorated function's own module globals once they become available.
    """
    if not current_platform.is_rocm():
        _ensure_tilelang_imported()
        return tilelang.jit(pass_configs=_get_pass_configs())(kernel_function)

    return _DeferredTileLangJitKernel(kernel_function)