checksum check

Author: jimmyzho

flashinfer/artifacts.py

@@ -22,13 +22,15 @@
 from typing import Generator
 import requests  # type: ignore[import-untyped]
 import shutil
+import hashlib
 
 from .jit.core import logger
 from .jit.cubin_loader import (
     FLASHINFER_CUBINS_REPOSITORY,
     get_cubin,
     safe_urljoin,
     FLASHINFER_CUBIN_DIR,
+    download_file,
 )
 
 
@@ -72,50 +74,88 @@ def get_available_cubin_files(
     return tuple()
 
 
-@dataclass(frozen=True)
 class ArtifactPath:
-    TRTLLM_GEN_FMHA: str = "7206d64e67f4c8949286246d6e2e07706af5d223/fmha/trtllm-gen"
+    TRTLLM_GEN_FMHA: str = "a72d85b019dc125b9f711300cb989430f762f5a6/fmha/trtllm-gen/"
     TRTLLM_GEN_BMM: str = (
-        "e6f22dcc3fdeb29ff87af2f4a2cb3d30b8d273e0/batched_gemm-45beda1-ee6a802"
+        "a72d85b019dc125b9f711300cb989430f762f5a6/batched_gemm-145d1b1-9e1d49a/"
     )
     TRTLLM_GEN_GEMM: str = (
-        "037e528e719ec3456a7d7d654f26b805e44c63b1/gemm-8704aa4-f91dc9e"
+        "a72d85b019dc125b9f711300cb989430f762f5a6/gemm-145d1b1-f91dc9e/"
     )
-    CUDNN_SDPA: str = "4c623163877c8fef5751c9c7a59940cd2baae02e/fmha/cudnn"
-    DEEPGEMM: str = "51d730202c9eef782f06ecc950005331d85c5d4b/deep-gemm"
-
+    CUDNN_SDPA: str = "a72d85b019dc125b9f711300cb989430f762f5a6/fmha/cudnn/"
+    DEEPGEMM: str = "a72d85b019dc125b9f711300cb989430f762f5a6/deep-gemm/"
 
+# TODO: Should be deprecated
 @dataclass(frozen=True)
 class MetaInfoHash:
     TRTLLM_GEN_FMHA: str = (
         "2f605255e71d673768f5bece66dde9e2e9f4c873347bfe8fefcffbf86a3c847d"
     )
     TRTLLM_GEN_BMM: str = (
-        "c98b4ce69a39fd41556d67033c30ea814ef76b0a2fe16e798e55baf0104acc34"
+        "9490085267aed30a387bfff024a0605e1ca4d39dfe06a5abc159d7d7e129bdf4"
     )
     DEEPGEMM: str = "b4374f857c3066089c4ec6b5e79e785559fa2c05ce2623710b0b04bf86414a48"
     TRTLLM_GEN_GEMM: str = (
-        "0345358c916d990709f9670e113e93f35c76aa22715e2d5128ec2ca8740be5ba"
+        "7d8ef4e6d89b6990e3e90a3d3a21e96918824d819f8f897a9bfd994925b9ea67"
     )
 
 
-def get_cubin_file_list() -> Generator[str, None, None]:
-    base = FLASHINFER_CUBINS_REPOSITORY
+# @dataclass(frozen=True)
+class CheckSumHash:
+    TRTLLM_GEN_FMHA: str = "b2d9d40db550ef85585e980bee651ac19d3e416f10b0c8bf9de0a7f9d0bee3d4"
+    TRTLLM_GEN_BMM: str = "efb9379c924193f6d3cb792bafb12b0811cab8eaa12bf324c7c410636c7769cd"
+    DEEPGEMM: str = "1a2a166839042dbd2a57f48051c82cd1ad032815927c753db269a4ed10d0ffbf"  
+    TRTLLM_GEN_GEMM: str = "e475e37989eed16418e0e858e2868ff07cb4b650cc48759cc23012f1afea310a"
 
-    # The meta info header files first.
-    yield safe_urljoin(ArtifactPath.TRTLLM_GEN_FMHA, "include/flashInferMetaInfo.h")
-    yield safe_urljoin(ArtifactPath.TRTLLM_GEN_GEMM, "include/flashinferMetaInfo.h")
-    yield safe_urljoin(ArtifactPath.TRTLLM_GEN_BMM, "include/flashinferMetaInfo.h")
+    map_checksums: [dict[str, str]] = {
+        safe_urljoin(ArtifactPath.TRTLLM_GEN_FMHA, "checksums.txt"): TRTLLM_GEN_FMHA,
+        safe_urljoin(ArtifactPath.TRTLLM_GEN_BMM, "checksums.txt"): TRTLLM_GEN_BMM,
+        safe_urljoin(ArtifactPath.DEEPGEMM, "checksums.txt"): DEEPGEMM,
+        safe_urljoin(ArtifactPath.TRTLLM_GEN_GEMM, "checksums.txt"): TRTLLM_GEN_GEMM,
+    }
 
-    # All the actual kernel cubin's.
-    for kernel in [
+
+def get_checksums(subdirs):
+    checksums = {}
+    for subdir in subdirs:
+        uri = safe_urljoin(FLASHINFER_CUBINS_REPOSITORY, safe_urljoin(subdir, "checksums.txt"))
+        checksum_path = FLASHINFER_CUBIN_DIR / safe_urljoin(subdir, "checksums.txt")
+        download_file(uri, checksum_path)
+        with open(checksum_path, "r") as f:
+            for line in f:
+                sha256, filename = line.strip().split()
+
+                # Distinguish between all meta info header files
+                if ".h" in filename:
+                    filename = safe_urljoin(subdir, filename)
+                checksums[filename] = sha256
+    return checksums
+
+
+def get_subdir_file_list():
+    base = FLASHINFER_CUBINS_REPOSITORY
+
+    cubin_dirs = [
         ArtifactPath.TRTLLM_GEN_FMHA,
         ArtifactPath.TRTLLM_GEN_BMM,
         ArtifactPath.TRTLLM_GEN_GEMM,
         ArtifactPath.DEEPGEMM,
-    ]:
-        for name in get_available_cubin_files(safe_urljoin(base, kernel)):
-            yield safe_urljoin(kernel, name)
+    ]
+
+    # Get checksums of all files
+    checksums = get_checksums(cubin_dirs)
+
+    # The meta info header files first.
+    yield (safe_urljoin(ArtifactPath.TRTLLM_GEN_FMHA, "include/flashInferMetaInfo.h"), checksums[safe_urljoin(ArtifactPath.TRTLLM_GEN_FMHA, "include/flashInferMetaInfo.h")])
+    yield (safe_urljoin(ArtifactPath.TRTLLM_GEN_GEMM, "include/flashinferMetaInfo.h"), checksums[safe_urljoin(ArtifactPath.TRTLLM_GEN_GEMM, "include/flashinferMetaInfo.h")])
+    yield (safe_urljoin(ArtifactPath.TRTLLM_GEN_BMM, "include/flashinferMetaInfo.h"), checksums[safe_urljoin(ArtifactPath.TRTLLM_GEN_BMM, "include/flashinferMetaInfo.h")])
+
+    # All the actual kernel cubin's.
+    for cubin_dir in cubin_dirs:
+        checksum_path = safe_urljoin(cubin_dir, "checksums.txt")
+        yield (checksum_path, CheckSumHash.map_checksums[checksum_path])
+        for name in get_available_cubin_files(safe_urljoin(base, cubin_dir)):
+            yield (safe_urljoin(cubin_dir, name), checksums[name])
 
 
 def download_artifacts() -> None:
@@ -124,27 +164,25 @@ def download_artifacts() -> None:
     # use a shared session to make use of HTTP keep-alive and reuse of
     # HTTPS connections.
     session = requests.Session()
+    cubin_files = list(get_subdir_file_list())
+    num_threads = int(os.environ.get("FLASHINFER_CUBIN_DOWNLOAD_THREADS", "4"))
+    with tqdm_logging_redirect(
+        total=len(cubin_files), desc="Downloading cubins"
+    ) as pbar:
 
-    with temp_env_var("FLASHINFER_CUBIN_CHECKSUM_DISABLED", "1"):
-        cubin_files = list(get_cubin_file_list())
-        num_threads = int(os.environ.get("FLASHINFER_CUBIN_DOWNLOAD_THREADS", "4"))
-        with tqdm_logging_redirect(
-            total=len(cubin_files), desc="Downloading cubins"
-        ) as pbar:
-
-            def update_pbar_cb(_) -> None:
-                pbar.update(1)
+        def update_pbar_cb(_) -> None:
+            pbar.update(1)
 
-            with ThreadPoolExecutor(num_threads) as pool:
-                futures = []
-                for name in cubin_files:
-                    fut = pool.submit(get_cubin, name, "", session)
-                    fut.add_done_callback(update_pbar_cb)
-                    futures.append(fut)
+        with ThreadPoolExecutor(num_threads) as pool:
+            futures = []
+            for name, checksum in cubin_files:
+                fut = pool.submit(get_cubin, name, checksum, session)
+                fut.add_done_callback(update_pbar_cb)
+                futures.append(fut)
 
-                results = [fut.result() for fut in as_completed(futures)]
+            results = [fut.result() for fut in as_completed(futures)]
 
-        all_success = all(results)
+    all_success = all(results)
     if not all_success:
         raise RuntimeError("Failed to download cubins")
 
@@ -154,7 +192,7 @@ def get_artifacts_status() -> tuple[tuple[str, bool], ...]:
     Check which cubins are already downloaded and return (num_downloaded, total).
     Does not download any cubins.
     """
-    cubin_files = get_cubin_file_list()
+    cubin_files = get_subdir_file_list()
 
     def _check_file_status(file_name: str) -> tuple[str, bool]:
         # get_cubin stores cubins in FLASHINFER_CUBIN_DIR with the same relative path
@@ -163,7 +201,7 @@ def _check_file_status(file_name: str) -> tuple[str, bool]:
         exists = os.path.isfile(local_path)
         return (file_name, exists)
 
-    return tuple(_check_file_status(file_name) for file_name in cubin_files)
+    return tuple(_check_file_status(file_name) for file_name, _ in cubin_files)
 
 
 def clear_cubin():