Caching cubins during build for reuse

2026-02-13 19:30:28 -08:00
parent 3470345624
commit 4214d4fefe
4 changed files with 34 additions and 2 deletions
--- a/flashinfer_cache.patch
+++ b/flashinfer_cache.patch
@@ -0,0 +1,18 @@
+--- a/flashinfer/artifacts.py
+++ b/flashinfer/artifacts.py
+@@ -203,9 +203,13 @@
+         with ThreadPoolExecutor(num_threads) as pool:
+             futures = []
+-            for name, _ in cubin_files:
+-                source = safe_urljoin(FLASHINFER_CUBINS_REPOSITORY, name)
+-                local_path = FLASHINFER_CUBIN_DIR / name
+            for name, checksum in cubin_files:
+                local_path = FLASHINFER_CUBIN_DIR / name
+                if local_path.exists() and verify_cubin(str(local_path), checksum):
+                    pbar.update(1)
+                    continue
+
+                source = safe_urljoin(FLASHINFER_CUBINS_REPOSITORY, name)
+                 # Ensure parent directory exists
+                 local_path.parent.mkdir(parents=True, exist_ok=True)
+                 fut = pool.submit(