misc: checksum check when downloading artifacts (#1761)

<!-- .github/pull_request_template.md -->

## 📌 Description

checks the sha256 hash when downloading cubins from artifactory, using
the generated `checksum.txt` in each cubin directory.

## 🔍 Related Issues

<!-- Link any related issues here -->

## 🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull
request, please make sure the following items are complete.

### ✅ Pre-commit Checks

- [x] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [x] I have installed the hooks with `pre-commit install`.
- [x] I have run the hooks manually with `pre-commit run --all-files`
and fixed any reported issues.

> If you are unsure about how to set up `pre-commit`, see [the
pre-commit documentation](https://pre-commit.com/).

## 🧪 Tests

- [ ] Tests have been added or updated as needed.
- [ ] All tests are passing (`unittest`, etc.).

## Reviewer Notes

<!-- Optional: anything you'd like reviewers to focus on, concerns, etc.
-->

Author: jimmyzho

flashinfer/artifacts.py

@@ -32,9 +32,10 @@
 
 from .jit.cubin_loader import (
     FLASHINFER_CUBINS_REPOSITORY,
-    download_file,
     safe_urljoin,
     FLASHINFER_CUBIN_DIR,
+    download_file,
+    verify_cubin,
 )
 
 
@@ -78,50 +79,109 @@ 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 = (
         "56fea80cb22f8b2ef2a2c6a822a075fb20b36803/batched_gemm-074aec4-cc00b23"
     )
     TRTLLM_GEN_GEMM: str = (
         "1fddc48b7b48af33914d040051b3e2ee9ba4701e/gemm-145d1b1-9b113e3"
     )
-    CUDNN_SDPA: str = "4c623163877c8fef5751c9c7a59940cd2baae02e/fmha/cudnn"
-    DEEPGEMM: str = "51d730202c9eef782f06ecc950005331d85c5d4b/deep-gemm"
+    CUDNN_SDPA: str = "a72d85b019dc125b9f711300cb989430f762f5a6/fmha/cudnn/"
+    DEEPGEMM: str = "a72d85b019dc125b9f711300cb989430f762f5a6/deep-gemm/"
 
 
 @dataclass(frozen=True)
 class MetaInfoHash:
+    DEEPGEMM: str = "b4374f857c3066089c4ec6b5e79e785559fa2c05ce2623710b0b04bf86414a48"
     TRTLLM_GEN_FMHA: str = (
-        "2f605255e71d673768f5bece66dde9e2e9f4c873347bfe8fefcffbf86a3c847d"
+        "d26dbf837f40ff2dcd964094ab6e1b3f2424edda5979c313f5262655161fce98"
     )
     TRTLLM_GEN_BMM: str = (
         "4a8ceeb356fc5339021acf884061e97e49e01da5c75dbf0f7cf4932c37a70152"
     )
-    DEEPGEMM: str = "b4374f857c3066089c4ec6b5e79e785559fa2c05ce2623710b0b04bf86414a48"
     TRTLLM_GEN_GEMM: str = (
         "bd5c3227bec4f8d7a7d3a27fd7628e010d99a5c42651d0a6b97e146803e63340"
     )
 
 
-def get_cubin_file_list() -> Generator[str, None, None]:
+class CheckSumHash:
+    TRTLLM_GEN_FMHA: str = (
+        "b2d9d40db550ef85585e980bee651ac19d3e416f10b0c8bf9de0a7f9d0bee3d4"
+    )
+    TRTLLM_GEN_BMM: str = (
+        "8df2aae8f3aa39d64d2c723e775640beb4ac602a6cbb02e497c2a7316e349934"
+    )
+    DEEPGEMM: str = "1a2a166839042dbd2a57f48051c82cd1ad032815927c753db269a4ed10d0ffbf"
+    TRTLLM_GEN_GEMM: str = (
+        "15cb8c85dfb5eddd4f121d64cb5a718321fb55b85aa19df10ddc1329d4a726b9"
+    )
+    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,
+    }
+
+
+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() -> Generator[tuple[str, str], None, None]:
     base = FLASHINFER_CUBINS_REPOSITORY
 
-    # 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")
-
-    # All the actual kernel cubin's.
-    for kernel in [
+    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:
@@ -130,8 +190,7 @@ 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_cubin_file_list())
+    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"
@@ -142,7 +201,7 @@ def update_pbar_cb(_) -> None:
 
         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
@@ -159,13 +218,19 @@ 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], ...]:
     """
     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
@@ -174,7 +239,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():

flashinfer/jit/attention/modules.py

@@ -1575,7 +1575,8 @@ def gen_trtllm_gen_fmha_module():
 
     # use `get_cubin` to get "flashinferMetaInfo.h"
     metainfo = get_cubin(
-        f"{include_path}/{header_name}.h", MetaInfoHash.TRTLLM_GEN_FMHA
+        f"{include_path}/{header_name}.h",
+        MetaInfoHash.TRTLLM_GEN_FMHA,
     )
 
     # make sure "flashinferMetaInfo.h" is downloaded or cached

flashinfer/jit/cubin_loader.py

@@ -136,6 +136,35 @@ def download_file(
         return False
 
 
+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
+    """
+    local_path = FLASHINFER_CUBIN_DIR / safe_urljoin(checksum_path, "checksums.txt")
+    with open(local_path, "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,10 @@ def gen_trtllm_gen_fused_moe_sm100_module() -> JitSpec:
     header_name = "flashinferMetaInfo"
 
     # use `get_cubin` to get "flashinferMetaInfo.h"
-    metainfo = get_cubin(f"{include_path}/{header_name}.h", MetaInfoHash.TRTLLM_GEN_BMM)
+    metainfo = get_cubin(
+        f"{include_path}/{header_name}.h",
+        MetaInfoHash.TRTLLM_GEN_BMM,
+    )
     # 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

@@ -35,17 +35,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.