Optimize disk cache source loading

Author: sepcnt

testing/python/cache/test_tilelang_kernel_cache_atomic_save.py

@@ -1,7 +1,11 @@
+import builtins
 import errno
 from pathlib import Path
+
+import cloudpickle
 import pytest
 
+import tilelang.cache.kernel_cache as kernel_cache_mod
 from tilelang.cache.kernel_cache import KernelCache
 from tilelang.env import env
 from tilelang.jit.adapter.nvrtc.kernel_cache import NVRTCKernelCache
@@ -46,6 +50,112 @@ def _make_fake_nvrtc_kernel(tmp_path):
     return _FakeKernel(str(lib_path))
 
 
+def _write_complete_kernel_cache_entry(
+    cache: KernelCache,
+    key: str,
+    device_source: str = "// device kernel",
+    host_source: str = "// host kernel",
+) -> Path:
+    cache_path = Path(cache._get_cache_path(key))
+    cache_path.mkdir(parents=True)
+    (cache_path / cache.device_kernel_path).write_text(device_source)
+    (cache_path / cache.host_kernel_path).write_text(host_source)
+    (cache_path / cache.kernel_lib_path).write_bytes(b"fake-so")
+    with (cache_path / cache.params_path).open("wb") as f:
+        cloudpickle.dump(["param"], f)
+    return cache_path
+
+
+def test_kernel_cache_disk_hit_defers_source_loading(cache_dirs, monkeypatch):
+    cache = KernelCache()
+    key = "lazy-source-load"
+    cache_path = _write_complete_kernel_cache_entry(cache, key)
+
+    sentinel = object()
+    captured = {}
+
+    def fail_source_load(*args, **kwargs):
+        raise AssertionError("disk cache hit should pass source paths through for lazy loading")
+
+    def fake_from_database(cls, **kwargs):
+        captured.update(kwargs)
+        return sentinel
+
+    monkeypatch.setattr(cache, "_load_kernel_source", fail_source_load)
+    monkeypatch.setattr(kernel_cache_mod.JITKernel, "from_database", classmethod(fake_from_database))
+
+    loaded = cache._load_kernel_from_disk(
+        key,
+        target="cuda",
+        target_host=None,
+        out_idx=[0],
+        execution_backend="tvm_ffi",
+        pass_configs=None,
+        compile_flags=None,
+        func=None,
+    )
+
+    assert loaded is sentinel
+    assert captured["host_kernel_source"] is None
+    assert captured["device_kernel_source"] is None
+    assert captured["host_kernel_source_path"] == str(cache_path / cache.host_kernel_path)
+    assert captured["device_kernel_source_path"] == str(cache_path / cache.device_kernel_path)
+    assert captured["kernel_lib_path"] == str(cache_path / cache.kernel_lib_path)
+    assert captured["params"] == ["param"]
+
+
+def test_kernel_cache_disk_hit_perf_skips_large_source_file_reads(cache_dirs, monkeypatch):
+    cache = KernelCache()
+    key = "lazy-source-load-perf"
+    large_source = "// source\n" + ("x" * (2 * 1024 * 1024))
+    cache_path = _write_complete_kernel_cache_entry(
+        cache,
+        key,
+        device_source=large_source,
+        host_source=large_source,
+    )
+    source_paths = {
+        (cache_path / cache.device_kernel_path).resolve(),
+        (cache_path / cache.host_kernel_path).resolve(),
+    }
+    source_read_count = 0
+    sentinel = object()
+
+    real_open = builtins.open
+
+    def tracking_open(file, *args, **kwargs):
+        nonlocal source_read_count
+        mode = args[0] if args else kwargs.get("mode", "r")
+        try:
+            path = Path(file).resolve()
+        except TypeError:
+            return real_open(file, *args, **kwargs)
+        if "r" in mode and path in source_paths:
+            source_read_count += 1
+            raise AssertionError("cache perf regression: source file read during disk cache hit")
+        return real_open(file, *args, **kwargs)
+
+    def fake_from_database(cls, **kwargs):
+        return sentinel
+
+    monkeypatch.setattr(builtins, "open", tracking_open)
+    monkeypatch.setattr(kernel_cache_mod.JITKernel, "from_database", classmethod(fake_from_database))
+
+    loaded = cache._load_kernel_from_disk(
+        key,
+        target="cuda",
+        target_host=None,
+        out_idx=[0],
+        execution_backend="tvm_ffi",
+        pass_configs=None,
+        compile_flags=None,
+        func=None,
+    )
+
+    assert loaded is sentinel
+    assert source_read_count == 0
+
+
 def test_kernel_cache_rewrites_incomplete_cache_dir(cache_dirs, tmp_path):
     cache = KernelCache()
     key = "atomic-repair"

tilelang/cache/kernel_cache.py

@@ -330,19 +330,23 @@ def cached(
                 )
                 return self._memory_cache[key]
 
-            if verbose:
-                self.logger.debug(f"Checking disk cache for kernel {get_prim_func_name(func, '<unknown>')}")
+        if verbose:
+            self.logger.debug(f"Checking disk cache for kernel {get_prim_func_name(func, '<unknown>')}")
 
-            # Then check disk cache
-            kernel = self._load_kernel_from_disk(
-                key, target, target_host, out_idx, execution_backend, pass_configs, compile_flags, func, verbose
-            )
-            if kernel is not None:
-                if verbose:
-                    self.logger.debug(f"Found kernel in disk cache for {get_prim_func_name(func, '<unknown>')}")
-                # Populate memory cache with disk result
+        # Disk loads can be expensive for large kernel sets; keep them outside
+        # the global cache lock so independent cache hits can proceed in parallel.
+        kernel = self._load_kernel_from_disk(
+            key, target, target_host, out_idx, execution_backend, pass_configs, compile_flags, func, verbose
+        )
+        if kernel is not None:
+            if verbose:
+                self.logger.debug(f"Found kernel in disk cache for {get_prim_func_name(func, '<unknown>')}")
+            with self._lock:
+                existing = self._memory_cache.get(key)
+                if existing is not None:
+                    return existing
                 self._memory_cache[key] = kernel
-                return kernel
+            return kernel
 
         if verbose:
             self.logger.debug(f"No cached kernel for {get_prim_func_name(func, '<unknown>')}")
@@ -518,9 +522,6 @@ def _load_kernel_from_disk(
         if not all([os.path.exists(file) for file in required_files]):
             return None
 
-        # Load the kernel source file (optional)
-        device_kernel_source, host_kernel_source = self._load_kernel_source(device_kernel_path, host_kernel_path, verbose)
-
         # Load kernel parameters
         kernel_params: list[KernelParam] | None = None
         try:
@@ -533,8 +534,10 @@ def _load_kernel_from_disk(
 
         return self._build_kernel(
             func=func,
-            host_kernel_source=host_kernel_source,
-            device_kernel_source=device_kernel_source,
+            host_kernel_source=None,
+            device_kernel_source=None,
+            host_kernel_path=host_kernel_path,
+            device_kernel_path=device_kernel_path,
             kernel_lib_path=kernel_lib_path,
             kernel_params=kernel_params,
             target=target,
@@ -638,8 +641,10 @@ def _set_adapter_cache_path(self, kernel: JITKernel, cache_path: str):
     def _build_kernel(
         self,
         func: Callable | None,
-        host_kernel_source: str,
-        device_kernel_source: str,
+        host_kernel_source: str | None,
+        device_kernel_source: str | None,
+        host_kernel_path: str | None,
+        device_kernel_path: str | None,
         kernel_lib_path: str,
         kernel_params: list[KernelParam] | None,
         target: str | Target,
@@ -651,10 +656,6 @@ def _build_kernel(
     ) -> JITKernel | None:
         # Check all required components and report specific failures
         missing_components = []
-        if not host_kernel_source:
-            missing_components.append("host_kernel_source")
-        if not device_kernel_source:
-            missing_components.append("device_kernel_source")
         if not kernel_params:
             missing_components.append("kernel_params")
 
@@ -674,4 +675,6 @@ def _build_kernel(
             execution_backend=execution_backend,
             pass_configs=pass_configs,
             compile_flags=compile_flags,
+            host_kernel_source_path=host_kernel_path,
+            device_kernel_source_path=device_kernel_path,
         )

tilelang/jit/adapter/base.py

@@ -94,3 +94,17 @@ def get_kernel_source(self, kernel_only: bool = True) -> str:
 
     def _post_init(self):
         self.func = self._convert_torch_func()
+
+    def _load_cached_text_source(self, source_attr: str, path_attr: str) -> str | None:
+        source = getattr(self, source_attr, None)
+        if source is not None:
+            return source
+
+        path = getattr(self, path_attr, None)
+        if path is None:
+            return None
+
+        with open(path, encoding="utf-8") as file:
+            source = file.read()
+        setattr(self, source_attr, source)
+        return source

tilelang/jit/adapter/cutedsl/adapter.py

@@ -104,18 +104,22 @@ def from_database(
         result_idx: list[int],
         target: str,
         func_or_mod: tir.PrimFunc | tvm.IRModule,
-        host_kernel_source: str,
-        device_kernel_source: str,
+        host_kernel_source: str | None,
+        device_kernel_source: str | None,
         kernel_lib_path: str,
         verbose: bool = False,
         pass_configs: dict[str, Any] | None = None,
         compile_flags: list[str] | None = None,
+        host_kernel_source_path: str | None = None,
+        device_kernel_source_path: str | None = None,
     ):
         adapter = cls.__new__(cls)
         adapter.params = params
         adapter.result_idx = adapter._legalize_result_idx(result_idx)
         adapter.host_kernel_source = host_kernel_source
         adapter.device_kernel_source = device_kernel_source
+        adapter._host_kernel_source_path = host_kernel_source_path
+        adapter._device_kernel_source_path = device_kernel_source_path
 
         if isinstance(func_or_mod, tir.PrimFunc):
             gsym = func_or_mod.attrs.get("global_symbol")
@@ -228,7 +232,10 @@ def get_kernel_source(self, kernel_only: bool = True) -> str | None:
         str | None
             The kernel source code, or None if not available
         """
-        return self.device_kernel_source
+        source = self._load_cached_text_source("device_kernel_source", "_device_kernel_source_path")
+        if source is not None:
+            self.kernel_global_source = source
+        return source
 
     def _forward_from_prebuild_lib(self, *args, stream: int | None = None, device_id: int = 0):
         """Low-level function to call the compiled CUDA kernel.

tilelang/jit/adapter/cython/adapter.py

@@ -155,18 +155,22 @@ def from_database(
         result_idx: list[int],
         target: str,
         func_or_mod: tir.PrimFunc | tvm.IRModule,
-        host_kernel_source: str,
-        device_kernel_source: str,
+        host_kernel_source: str | None,
+        device_kernel_source: str | None,
         kernel_lib_path: str,
         verbose: bool = False,
         pass_configs: dict[str, Any] | None = None,
         compile_flags: list[str] | None = None,
+        host_kernel_source_path: str | None = None,
+        device_kernel_source_path: str | None = None,
     ):
         adapter = cls.__new__(cls)
         adapter.params = params
         adapter.result_idx = adapter._legalize_result_idx(result_idx)
         adapter.host_kernel_source = host_kernel_source
         adapter.device_kernel_source = device_kernel_source
+        adapter._host_kernel_source_path = host_kernel_source_path
+        adapter._device_kernel_source_path = device_kernel_source_path
         adapter.kernel_global_source = device_kernel_source  # Set alias for compatibility
         adapter.pass_configs = pass_configs
 
@@ -386,12 +390,16 @@ def is_dynamic(self):
     def get_kernel_source(self, kernel_only: bool = False):
         """Returns the source code of the compiled kernel."""
         if kernel_only:
-            return self.device_kernel_source
+            source = self._load_cached_text_source("device_kernel_source", "_device_kernel_source_path")
+            if source is not None:
+                self.kernel_global_source = source
+            return source
         else:
             # Wrapper only has host kernel source
-            assert self.host_kernel_source is not None, "Wrapped source is not available"
-            return self.host_kernel_source
+            source = self._load_cached_text_source("host_kernel_source", "_host_kernel_source_path")
+            assert source is not None, "Wrapped source is not available"
+            return source
 
     def get_host_source(self):
         """Returns the source code of the host function."""
-        return self.host_kernel_source
+        return self._load_cached_text_source("host_kernel_source", "_host_kernel_source_path")

tilelang/jit/adapter/nvrtc/adapter.py

@@ -102,18 +102,23 @@ def from_database(
         result_idx: list[int],
         target: str,
         func_or_mod: tir.PrimFunc | tvm.IRModule,
-        host_kernel_source: str,
-        device_kernel_source: str,
+        host_kernel_source: str | None,
+        device_kernel_source: str | None,
         kernel_lib_path: str,
         verbose: bool = False,
         pass_configs: dict[str, Any] | None = None,
         compile_flags: list[str] | None = None,
+        host_kernel_source_path: str | None = None,
+        device_kernel_source_path: str | None = None,
     ):
         adapter = cls.__new__(cls)
         adapter.params = params
         adapter.result_idx = adapter._legalize_result_idx(result_idx)
         adapter.host_kernel_source = host_kernel_source
         adapter.device_kernel_source = device_kernel_source
+        adapter.host_func = host_kernel_source
+        adapter._host_kernel_source_path = host_kernel_source_path
+        adapter._device_kernel_source_path = device_kernel_source_path
 
         if isinstance(func_or_mod, tir.PrimFunc):
             adapter.ir_module = tvm.IRModule({func_or_mod.attrs["global_symbol"]: func_or_mod})
@@ -201,9 +206,9 @@ def get_kernel_source(self, kernel_only: bool = True) -> str | None:
             The kernel source code, or None if not available
         """
         if kernel_only:
-            return self.device_kernel_source
+            return self._load_cached_text_source("device_kernel_source", "_device_kernel_source_path")
         else:
-            return self.host_func
+            return self._load_cached_text_source("host_func", "_host_kernel_source_path")
 
     def _forward_from_prebuild_lib(self, *args, stream: int | None = None):
         """Low-level function to call the compiled CUDA kernel."""