misc: Update artifacts docstring and MetaInfoHash (#1967)

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

## 📌 Description

Amendment to [PR
1761](#1761), appending
docstring to two artifactory path classes and deprecating need to update
MetaInfoHash by directly accessing the checksum.txt file.

## 🔍 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

- [ ] I have installed `pre-commit` by running `pip install pre-commit`
(or used your preferred method).
- [ ] I have installed the hooks with `pre-commit install`.
- [ ] 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.
-->


<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added runtime integrity checks for compiled artifacts that verify and
use checksum data during loading to prevent missing or mismatched
artifact headers.

* **Refactor**
* Switched artifact hash resolution to compute hashes dynamically from
provided checksums, improving validation, reliability, and resilience
when loading precompiled components.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Author: jimmyzho

flashinfer/artifacts.py

@@ -79,7 +79,14 @@ def get_available_cubin_files(
     return tuple()
 
 
+@dataclass(frozen=True)
 class ArtifactPath:
+    """
+    This class is used to store the paths of the cubin files in artifactory.
+    The paths are generated in cubin publishing script logs (accessible by codeowners).
+    When compiling new cubins for backend directories, update the corresponding path.
+    """
+
     TRTLLM_GEN_FMHA: str = "463def7494c9fc6792b5aa5b5beef34025e247ac/fmha/trtllm-gen/"
     TRTLLM_GEN_BMM: str = (
         "56fea80cb22f8b2ef2a2c6a822a075fb20b36803/batched_gemm-074aec4-cc00b23"
@@ -106,6 +113,12 @@ class MetaInfoHash:
 
 
 class CheckSumHash:
+    """
+    This class is used to store the checksums of the cubin files in artifactory.
+    The sha256 hashes are generated in cubin publishing script logs (accessible by codeowners).
+    When updating the ArtifactPath for backend directories, update the corresponding hash.
+    """
+
     TRTLLM_GEN_FMHA: str = (
         "639c534614e9fdf5a9cfa91f7ea8f53989613019c0e1f8b755f461e1fcc7546f"
     )

flashinfer/jit/attention/modules.py

@@ -28,7 +28,7 @@
     sm90a_nvcc_flags,
     current_compilation_context,
 )
-from ...jit.cubin_loader import get_cubin
+from ...jit.cubin_loader import get_cubin, get_meta_hash
 from ..utils import (
     dtype_map,
     filename_safe_dtype_map,
@@ -1568,15 +1568,21 @@ def gen_fmha_cutlass_sm100a_module(
 
 
 def gen_trtllm_gen_fmha_module():
-    from ...artifacts import ArtifactPath, MetaInfoHash
+    from ...artifacts import ArtifactPath, CheckSumHash
 
     include_path = f"{ArtifactPath.TRTLLM_GEN_FMHA}/include"
     header_name = "flashInferMetaInfo"
 
+    # Check if checksums.txt exists in the cubin directory
+    checksum_path = f"{ArtifactPath.TRTLLM_GEN_FMHA}/checksums.txt"
+    checksum = get_cubin(checksum_path, CheckSumHash.TRTLLM_GEN_FMHA)
+    assert checksum, f"Failed to get checksums.txt from {checksum_path}"
+
+    meta_hash = get_meta_hash(checksum)
     # use `get_cubin` to get "flashinferMetaInfo.h"
     metainfo = get_cubin(
         f"{include_path}/{header_name}.h",
-        MetaInfoHash.TRTLLM_GEN_FMHA,
+        meta_hash,
     )
 
     # make sure "flashinferMetaInfo.h" is downloaded or cached
@@ -1592,7 +1598,7 @@ def gen_trtllm_gen_fmha_module():
         extra_include_paths=[jit_env.FLASHINFER_CUBIN_DIR / include_path],
         extra_cuda_cflags=[
             f'-DTLLM_GEN_FMHA_CUBIN_PATH=\\"{ArtifactPath.TRTLLM_GEN_FMHA}\\"',
-            f'-DTLLM_GEN_FMHA_METAINFO_HASH=\\"{MetaInfoHash.TRTLLM_GEN_FMHA}\\"',
+            f'-DTLLM_GEN_FMHA_METAINFO_HASH=\\"{meta_hash}\\"',
         ],
     )
 

flashinfer/jit/cubin_loader.py

@@ -136,18 +136,16 @@ def download_file(
         return False
 
 
-def get_meta_hash(checksum_path: str) -> str:
+def get_meta_hash(checksums_bytes: bytes) -> str:
     """
-    Load the file from local cache (checksums.txt)
-    and get the hash of corresponding flashinferMetaInfo.h file
+    Parse the checksums.txt file 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}")
+    checksums_lines = checksums_bytes.decode("utf-8").splitlines()
+    for line in checksums_lines:
+        sha256, filename = line.strip().split()
+        if ".h" in filename:
+            return sha256
+    raise ValueError("Invalid checksums.txt, no flashinferMetaInfo.h found")
 
 
 def verify_cubin(cubin_path: str, expected_sha256: str) -> bool:

flashinfer/jit/fused_moe.py

@@ -17,10 +17,10 @@
 from typing import List
 
 from . import env as jit_env
-from ..artifacts import ArtifactPath, MetaInfoHash
+from ..artifacts import ArtifactPath, CheckSumHash
 from .core import JitSpec, gen_jit_spec, current_compilation_context, sm90a_nvcc_flags
 from .cpp_ext import is_cuda_version_at_least
-from .cubin_loader import get_cubin
+from .cubin_loader import get_cubin, get_meta_hash
 from .gemm.cutlass.generate_kernels import generate_gemm_operations
 
 
@@ -180,10 +180,16 @@ def gen_trtllm_gen_fused_moe_sm100_module() -> JitSpec:
     include_path = f"{ArtifactPath.TRTLLM_GEN_BMM}/include"
     header_name = "flashinferMetaInfo"
 
+    # Check if checksums.txt exists in the cubin directory
+    checksum_path = f"{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt"
+    checksum = get_cubin(checksum_path, CheckSumHash.TRTLLM_GEN_BMM)
+    assert checksum, f"Failed to get checksums.txt from {checksum_path}"
+    meta_hash = get_meta_hash(checksum)
+
     # use `get_cubin` to get "flashinferMetaInfo.h"
     metainfo = get_cubin(
         f"{include_path}/{header_name}.h",
-        MetaInfoHash.TRTLLM_GEN_BMM,
+        meta_hash,
     )
     # make sure "flashinferMetaInfo.h" is downloaded or cached
     assert metainfo, f"{header_name}.h not found"

flashinfer/jit/gemm/core.py

@@ -20,7 +20,7 @@
 import jinja2
 import torch
 
-from ...artifacts import ArtifactPath, MetaInfoHash
+from ...artifacts import ArtifactPath, CheckSumHash
 from .. import env as jit_env
 from ..core import (
     JitSpec,
@@ -30,7 +30,7 @@
     sm100f_nvcc_flags,
     current_compilation_context,
 )
-from ..cubin_loader import get_cubin
+from ..cubin_loader import get_cubin, get_meta_hash
 from ..utils import dtype_cutlass_map, filename_safe_dtype_map, write_if_different
 
 
@@ -361,10 +361,16 @@ def gen_trtllm_gen_gemm_module() -> JitSpec:
     include_path = f"{ArtifactPath.TRTLLM_GEN_GEMM}/include"
     header_name = "flashinferMetaInfo"
 
+    # Check if checksums.txt exists in the cubin directory
+    checksum_path = f"{ArtifactPath.TRTLLM_GEN_GEMM}/checksums.txt"
+    checksum = get_cubin(checksum_path, CheckSumHash.TRTLLM_GEN_GEMM)
+    assert checksum, f"Failed to get checksums.txt from {checksum_path}"
+    meta_hash = get_meta_hash(checksum)
+
     # use `get_cubin` to get "flashinferMetaInfo.h"
     metainfo = get_cubin(
         f"{include_path}/{header_name}.h",
-        MetaInfoHash.TRTLLM_GEN_GEMM,
+        meta_hash,
     )
     # make sure "flashinferMetaInfo.h" is downloaded or cached
     assert metainfo, f"{header_name}.h not found"
@@ -505,10 +511,16 @@ def gen_trtllm_low_latency_gemm_module() -> JitSpec:
     include_path = f"{ArtifactPath.TRTLLM_GEN_GEMM}/include"
     header_name = "flashinferMetaInfo"
 
+    # Check if checksums.txt exists in the cubin directory
+    checksum_path = f"{ArtifactPath.TRTLLM_GEN_GEMM}/checksums.txt"
+    checksum = get_cubin(checksum_path, CheckSumHash.TRTLLM_GEN_GEMM)
+    assert checksum, f"Failed to get checksums.txt from {checksum_path}"
+    meta_hash = get_meta_hash(checksum)
+
     # use `get_cubin` to get "flashinferMetaInfo.h"
     metainfo = get_cubin(
         f"{include_path}/{header_name}.h",
-        MetaInfoHash.TRTLLM_GEN_GEMM,
+        meta_hash,
     )
     # make sure "flashinferMetaInfo.h" is downloaded or cached
     assert metainfo, f"{header_name}.h not found"