update checksum check

Author: jimmyzho

flashinfer/artifacts.py

@@ -32,10 +32,10 @@
 
 from .jit.cubin_loader import (
     FLASHINFER_CUBINS_REPOSITORY,
-    download_file,
     safe_urljoin,
     FLASHINFER_CUBIN_DIR,
     download_file,
+    verify_cubin,
 )
 
 
@@ -91,22 +91,11 @@ class ArtifactPath:
     DEEPGEMM: str = "a72d85b019dc125b9f711300cb989430f762f5a6/deep-gemm/"
 
 
-# TODO (jimmyzho): Should be deprecated except DEEPGEMM
 @dataclass(frozen=True)
 class MetaInfoHash:
-    TRTLLM_GEN_FMHA: str = (
-        "2f605255e71d673768f5bece66dde9e2e9f4c873347bfe8fefcffbf86a3c847d"
-    )
-    TRTLLM_GEN_BMM: str = (
-        "9490085267aed30a387bfff024a0605e1ca4d39dfe06a5abc159d7d7e129bdf4"
-    )
     DEEPGEMM: str = "b4374f857c3066089c4ec6b5e79e785559fa2c05ce2623710b0b04bf86414a48"
-    TRTLLM_GEN_GEMM: str = (
-        "7d8ef4e6d89b6990e3e90a3d3a21e96918824d819f8f897a9bfd994925b9ea67"
-    )
 
 
-# @dataclass(frozen=True)
 class CheckSumHash:
     TRTLLM_GEN_FMHA: str = (
         "b2d9d40db550ef85585e980bee651ac19d3e416f10b0c8bf9de0a7f9d0bee3d4"
@@ -197,12 +186,13 @@ def download_artifacts() -> None:
     with tqdm_logging_redirect(
         total=len(cubin_files), desc="Downloading cubins"
     ) as pbar:
+
         def update_pbar_cb(_) -> None:
             pbar.update(1)
 
         with ThreadPoolExecutor(num_threads) as pool:
             futures = []
-            for name in cubin_files:
+            for name, _ in cubin_files:
                 source = safe_urljoin(FLASHINFER_CUBINS_REPOSITORY, name)
                 local_path = FLASHINFER_CUBIN_DIR / name
                 # Ensure parent directory exists
@@ -219,6 +209,12 @@ def update_pbar_cb(_) -> None:
     if not all_success:
         raise RuntimeError("Failed to download cubins")
 
+    # Check checksums of all downloaded cubins
+    for name, checksum in cubin_files:
+        local_path = FLASHINFER_CUBIN_DIR / name
+        if not verify_cubin(str(local_path), checksum):
+            raise RuntimeError("Failed to download cubins: checksum mismatch")
+
 
 def get_artifacts_status() -> tuple[tuple[str, bool], ...]:
     """

flashinfer/jit/attention/modules.py

@@ -1568,15 +1568,15 @@ def gen_fmha_cutlass_sm100a_module(
 
 
 def gen_trtllm_gen_fmha_module():
-    from ...artifacts import ArtifactPath, MetaInfoHash
+    from ...artifacts import ArtifactPath
 
     include_path = f"{ArtifactPath.TRTLLM_GEN_FMHA}/include"
     header_name = "flashInferMetaInfo"
 
     # use `get_cubin` to get "flashinferMetaInfo.h"
     metainfo = get_cubin(
         f"{include_path}/{header_name}.h",
-        get_meta_hash(f"{ArtifactPath.TRTLLM_GEN_FMHA}/checksums.txt"),
+        get_meta_hash(ArtifactPath.TRTLLM_GEN_FMHA),
     )
 
     # make sure "flashinferMetaInfo.h" is downloaded or cached

flashinfer/jit/cubin_loader.py

@@ -132,14 +132,29 @@ def get_meta_hash(checksum_path: str) -> str:
     Load the file from local cache (checksums.txt)
     and get the hash of corresponding flashinferMetaInfo.h file
     """
-    with open(checksum_path, "r") as f:
+    with open(checksum_path + "/checksums.txt", "r") as f:
         for line in f:
             sha256, filename = line.strip().split()
             if ".h" in filename:
                 return sha256
     raise ValueError(f"Invalid path: checksums.txt not found in {checksum_path}")
 
 
+def verify_cubin(cubin_path: str, expected_sha256: str) -> bool:
+    """
+    Verify the cubin file against the sha256 checksum.
+    """
+    with open(cubin_path, "rb") as f:
+        data = f.read()
+    actual_sha256 = hashlib.sha256(data).hexdigest()
+    if actual_sha256 != expected_sha256:
+        logger.warning(
+            f"sha256 mismatch (expected {expected_sha256} actual {actual_sha256}) for {cubin_path}"
+        )
+        return False
+    return True
+
+
 def load_cubin(cubin_path: str, sha256: str) -> bytes:
     """
     Load a cubin from the provide local path and

flashinfer/jit/fused_moe.py

@@ -181,7 +181,7 @@ def gen_trtllm_gen_fused_moe_sm100_module() -> JitSpec:
     # use `get_cubin` to get "flashinferMetaInfo.h"
     metainfo = get_cubin(
         f"{include_path}/{header_name}.h",
-        get_meta_hash(f"{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt"),
+        get_meta_hash(ArtifactPath.TRTLLM_GEN_BMM),
     )
     # make sure "flashinferMetaInfo.h" is downloaded or cached
     assert metainfo, f"{header_name}.h not found"

flashinfer/jit/gemm/core.py

@@ -364,7 +364,7 @@ def gen_trtllm_gen_gemm_module() -> JitSpec:
     # use `get_cubin` to get "flashinferMetaInfo.h"
     metainfo = get_cubin(
         f"{include_path}/{header_name}.h",
-        get_meta_hash(f"{ArtifactPath.TRTLLM_GEN_GEMM}/checksums.txt"),
+        get_meta_hash(ArtifactPath.TRTLLM_GEN_GEMM),
     )
     # make sure "flashinferMetaInfo.h" is downloaded or cached
     assert metainfo, f"{header_name}.h not found"

tests/utils/test_load_cubin_compile_race_condition.py

@@ -34,17 +34,12 @@ def worker_process(temp_dir):
     os.environ["FLASHINFER_CUBIN_DIR"] = temp_dir
 
     # Import here to ensure FLASHINFER_CUBIN_DIR is set before module loads
-    from flashinfer.artifacts import ArtifactPath, MetaInfoHash
-    from flashinfer.jit.cubin_loader import get_cubin
+    from flashinfer.artifacts import ArtifactPath
 
     # Define the target file - same for all processes
     include_path = f"{ArtifactPath.TRTLLM_GEN_BMM}/include"
     header_name = "flashinferMetaInfo"
 
-    # Use get_cubin to get "flashinferMetaInfo.h"
-    # Note: all processes target the same file name
-    metainfo = get_cubin(f"{include_path}/{header_name}.h", MetaInfoHash.TRTLLM_GEN_BMM)  # noqa: F841
-
     # Read the file from FLASHINFER_CUBIN_DIR
     # NOTE(Zihao): instead of using metainfo, we directly read from the file path,
     # that aligns with how we compile the kernel.