| |
| |
|
|
| """Persistent AOT cache for CuTe DSL compiled kernels. |
| |
| Saves compiled TVM FFI kernels as .o files on first compile, |
| loads them on subsequent runs to skip JIT compilation. |
| |
| Environment variables: |
| MM_SPARSE_ATTN_AOT_CACHE: Override cache directory |
| (default: ~/.cache/minfer/mm_sparse_attn) |
| MM_SPARSE_ATTN_AOT_DISABLE=1: Disable AOT cache entirely |
| """ |
|
|
| import hashlib |
| import os |
| import time |
|
|
| import cutlass.cute as cute |
|
|
| _AOT_CACHE_DIR = os.environ.get( |
| "MM_SPARSE_ATTN_AOT_CACHE", |
| os.path.expanduser("~/.cache/minfer/mm_sparse_attn"), |
| ) |
| _AOT_DISABLE = os.environ.get("MM_SPARSE_ATTN_AOT_DISABLE", "0") == "1" |
|
|
| _loaded_modules: dict[str, object] = {} |
|
|
|
|
| def _key_to_path(key: tuple) -> str: |
| h = hashlib.sha256(repr(key).encode()).hexdigest()[:16] |
| name = str(key[0]).replace("/", "_") |
| return os.path.join(_AOT_CACHE_DIR, f"{name}_{h}") |
|
|
|
|
| def try_load_aot(key: tuple): |
| if _AOT_DISABLE: |
| return None |
| obj_path = _key_to_path(key) + ".o" |
| if not os.path.isfile(obj_path): |
| return None |
| func_name = str(key[0]) |
| try: |
| if obj_path not in _loaded_modules: |
| _loaded_modules[obj_path] = cute.runtime.load_module( |
| obj_path, enable_tvm_ffi=True |
| ) |
| return getattr(_loaded_modules[obj_path], func_name) |
| except Exception as e: |
| print(f"[aot_cache] Failed to load {obj_path}: {e}") |
| return None |
|
|
|
|
| def save_aot(key: tuple, compiled) -> None: |
| if _AOT_DISABLE: |
| return |
| if not hasattr(compiled, "export_to_c"): |
| return |
| obj_path = _key_to_path(key) + ".o" |
| os.makedirs(_AOT_CACHE_DIR, exist_ok=True) |
| tmp_path = obj_path + f".tmp.{os.getpid()}" |
| func_name = str(key[0]) |
| try: |
| t0 = time.time() |
| compiled.export_to_c(tmp_path, function_name=func_name) |
| os.replace(tmp_path, obj_path) |
| dt = time.time() - t0 |
| print(f"[aot_cache] Saved {func_name} -> {obj_path} ({dt:.1f}s)") |
| except Exception as e: |
| print(f"[aot_cache] Failed to save {func_name}: {e}") |
| if os.path.exists(tmp_path): |
| os.remove(tmp_path) |
|
|
