tilelang.cache.kernel_cache

The cache utils with class and database persistence - KernelCache Class

Classes

KernelCache

Caches compiled kernels using a class and database persistence to avoid redundant compilation.

Module Contents

class tilelang.cache.kernel_cache.KernelCache

Caches compiled kernels using a class and database persistence to avoid redundant compilation. Cache files:

kernel.cu: The compiled kernel source code wrapped_kernel.cu: The compiled wrapped kernel source code kernel_lib.so: The compiled kernel library params.pkl: The compiled kernel parameters

execution_backend: Literal['tvm_ffi', 'cython', 'nvrtc', 'torch', 'cutedsl'] = 'tvm_ffi'
device_kernel_path = 'device_kernel.cu'
host_kernel_path = 'host_kernel.cu'
kernel_lib_path = 'kernel_lib.so'
params_path = 'params.pkl'
cached(func=None, out_idx=None, *args, target, target_host=None, execution_backend='tvm_ffi', verbose, pass_configs=None, compile_flags=None)

Caches and reuses compiled kernels to avoid redundant compilation.

This is the ONLY place where environment variable processing, target normalization, and execution backend resolution should happen. All compilation paths go through here.

Parameters:

  • func (tvm.tir.PrimFunc) – Function to be compiled or a prepared PrimFunc

  • out_idx (list[int]) – Indices specifying which outputs to return

  • target (str | tvm.target.Target) – Compilation target platform (None = read from TILELANG_TARGET env var)

  • target_host (str | tvm.target.Target | None) – Host target platform

  • execution_backend (Literal['tvm_ffi', 'cython', 'nvrtc', 'torch', 'cutedsl']) – Execution backend (None = read from TILELANG_EXECUTION_BACKEND)

  • verbose (bool) – Enable verbose output (None = read from TILELANG_VERBOSE)

  • *args – Arguments passed to func

  • pass_configs (dict | None)

  • compile_flags (list[str] | str | None)

Returns:

The compiled kernel, either freshly compiled or from cache

Return type:

JITKernel

Environment Variables

TILELANG_TARGETstr

Default compilation target (e.g., “cuda”, “llvm”). Defaults to “auto”.

TILELANG_EXECUTION_BACKENDstr

Default execution backend. Defaults to “auto”.

TILELANG_VERBOSEstr

Set to “1”, “true”, “yes”, or “on” to enable verbose compilation by default.

clear_cache()

Clears the entire kernel cache, including both in-memory and disk cache.