deprecate metainfohash

Author: jimmyzho

flashinfer/artifacts.py

@@ -85,7 +85,7 @@ class ArtifactPath:
     DEEPGEMM: str = "a72d85b019dc125b9f711300cb989430f762f5a6/deep-gemm/"
 
 
-# TODO: Should be deprecated
+# TODO (jimmyzho): Should be deprecated except DEEPGEMM
 @dataclass(frozen=True)
 class MetaInfoHash:
     TRTLLM_GEN_FMHA: str = (

flashinfer/jit/attention/modules.py

@@ -20,7 +20,7 @@
 import jinja2
 import torch
 
-from ...artifacts import ArtifactPath, MetaInfoHash
+from ...artifacts import ArtifactPath
 from .. import env as jit_env
 from ..core import (
     JitSpec,
@@ -29,7 +29,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,
@@ -1574,7 +1574,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",
+        get_meta_hash(f"{ArtifactPath.TRTLLM_GEN_FMHA}/checksums.txt"),
     )
 
     # make sure "flashinferMetaInfo.h" is downloaded or cached
@@ -1590,7 +1591,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=\\"{get_meta_hash(f"{ArtifactPath.TRTLLM_GEN_FMHA}/checksums.txt")}\\"',
         ],
     )
 

flashinfer/jit/cubin_loader.py

@@ -120,6 +120,19 @@ 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
+    """
+    with open(checksum_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 load_cubin(cubin_path: str, sha256: str) -> bytes:
     """
     Load a cubin from the provide local path and

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
 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
 
 
@@ -179,7 +179,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",
+        get_meta_hash(f"{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt"),
+    )
     # 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
 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
 
 
@@ -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",
-        MetaInfoHash.TRTLLM_GEN_GEMM,
+        get_meta_hash(f"{ArtifactPath.TRTLLM_GEN_GEMM}/checksums.txt"),
     )
     # make sure "flashinferMetaInfo.h" is downloaded or cached
     assert metainfo, f"{header_name}.h not found"