[Refactor] Move compilation outside critical section (#260)

* move compilation outside critical section

* lint fix

---------
Co-authored-by: LeiWang1999 <leiwang1999@outlook.com>

tilelang/cache/kernel_cache.py

@@ -38,7 +38,6 @@ class KernelCache:
         with cls._lock:
             if cls._instance is None:
                 cls._instance = super(KernelCache, cls).__new__(cls)
-                cls._instance._cache = {}  # In-memory cache
                 cls._instance.cache_dir = cache_dir  # Cache directory
                 os.makedirs(cls._instance.cache_dir, exist_ok=True)  # Ensure cache directory exists
                 cls._instance.logger = logging.getLogger(__name__)  # Initialize logger
@@ -91,18 +90,14 @@ class KernelCache:
             JITKernel: The compiled kernel, either freshly compiled or from cache
         """
         key = self._generate_key(func, out_idx, execution_backend, args, target, target_host)
-        with self._lock:  # Thread-safe access to cache
-            if key in self._cache:
-                return self._cache[key]
-
+        with self._lock:  # TODO: use filelock
             # Attempt to load from disk
             kernel = self._load_kernel_from_disk(key, target, target_host, out_idx,
                                                  execution_backend, pass_configs, func)
-            if kernel:
-                self._cache[key] = kernel  # Load to in-memory cache
+            if kernel is not None:
                 return kernel
 
-            # Compile kernel if cache miss
+        # Compile kernel if cache miss; leave critical section
         kernel = JITKernel(
             func,
             out_idx=out_idx,
@@ -112,10 +107,13 @@ class KernelCache:
             verbose=verbose,
             pass_configs=pass_configs,
         )
-            self._cache[key] = kernel  # Store in in-memory cache
         if execution_backend == "dlpack":
             self.logger.warning("DLPack backend does not support cache saving to disk.")
         else:
+            with self._lock:  # enter critical section again to check and update disk cache
+                disk_kernel = self._load_kernel_from_disk(key, target, target_host, out_idx,
+                                                          execution_backend, pass_configs, func)
+                if disk_kernel is None:
                     self._save_kernel_to_disk(key, kernel, func)
         return kernel
 
@@ -123,8 +121,7 @@ class KernelCache:
         """
         Clears the entire kernel cache, including both in-memory and disk cache.
         """
-        with self._lock:  # Thread-safe operation
-            self._cache.clear()  # Clear in-memory cache
+        with self._lock:
             self._clear_disk_cache()  # Clear disk cache
 
     def _get_cache_path(self, key: str) -> str: