flashinfer-python 0.6.7__tar.gz → 0.6.7.post1__tar.gz

{flashinfer_python-0.6.7 → flashinfer_python-0.6.7.post1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashinfer-python
-Version: 0.6.7
+Version: 0.6.7.post1
 Summary: FlashInfer: Kernel Library for LLM Serving
 Author: FlashInfer team
 License-Expression: Apache-2.0

flashinfer_python-0.6.7.post1/flashinfer/_build_meta.py

@@ -0,0 +1,3 @@
+"""Build metadata for flashinfer package."""
+__version__ = "0.6.7.post1"
+__git_version__ = "d6129c7ab41e2e2427e6fbf6d3d8001e5b0d5449"

{flashinfer_python-0.6.7 → flashinfer_python-0.6.7.post1}/flashinfer/artifacts.py

@@ -284,7 +284,7 @@ def get_artifacts_status() -> tuple[tuple[str, bool], ...]:
     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
+        # get_artifact stores files in FLASHINFER_CUBIN_DIR with the same relative path
         # Remove any leading slashes from name
         local_path = os.path.join(FLASHINFER_CUBIN_DIR, file_name)
         exists = os.path.isfile(local_path)

{flashinfer_python-0.6.7 → flashinfer_python-0.6.7.post1}/flashinfer/deep_gemm.py

@@ -46,7 +46,7 @@ logger = logging.getLogger(__name__)
 
 from .artifacts import ArtifactPath
 from .cuda_utils import checkCudaErrors
-from .jit.cubin_loader import get_cubin
+from .jit.cubin_loader import get_artifact
 from .jit.env import FLASHINFER_CUBIN_DIR
 from .utils import (
     ceil_div,
@@ -945,7 +945,7 @@ def load_all():
             continue
         symbol, sha256 = KERNEL_MAP[cubin_name]
         cubin_name = cubin_name + ".cubin"
-        get_cubin(ArtifactPath.DEEPGEMM + "/" + cubin_name, sha256)
+        get_artifact(ArtifactPath.DEEPGEMM + "/" + cubin_name, sha256)
         path = FLASHINFER_CUBIN_DIR / ArtifactPath.DEEPGEMM / cubin_name
         assert path.exists()
         RUNTIME_CACHE[cubin_name] = SM100FP8GemmRuntime(str(path), symbol)
@@ -960,7 +960,7 @@ def load(name: str, code: str) -> SM100FP8GemmRuntime:
         return RUNTIME_CACHE[cubin_name]
     symbol, sha256 = KERNEL_MAP[cubin_name]
     cubin_name = cubin_name + ".cubin"
-    get_cubin(ArtifactPath.DEEPGEMM + "/" + cubin_name, sha256)
+    get_artifact(ArtifactPath.DEEPGEMM + "/" + cubin_name, sha256)
     path = FLASHINFER_CUBIN_DIR / ArtifactPath.DEEPGEMM / cubin_name
     assert path.exists()
     RUNTIME_CACHE[cubin_name] = SM100FP8GemmRuntime(str(path), symbol)
@@ -1593,7 +1593,7 @@ class KernelMap:
 
     def init_indices(self):
         indice_path = ArtifactPath.DEEPGEMM + "/" + "kernel_map.json"
-        assert get_cubin(indice_path, self.KERNEL_MAP_HASH), (
+        assert get_artifact(indice_path, self.KERNEL_MAP_HASH), (
             "cubin kernel map file not found, nor downloaded with matched sha256"
         )
         path = FLASHINFER_CUBIN_DIR / indice_path

{flashinfer_python-0.6.7 → flashinfer_python-0.6.7.post1}/flashinfer/jit/attention/modules.py

@@ -28,7 +28,7 @@ from ..core import (
     sm90a_nvcc_flags,
     current_compilation_context,
 )
-from ...jit.cubin_loader import get_cubin, get_meta_hash
+from ...jit.cubin_loader import get_artifact, get_meta_hash
 from ..utils import (
     dtype_map,
     filename_safe_dtype_map,
@@ -1769,12 +1769,12 @@ def gen_trtllm_gen_fmha_module():
 
     # 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)
+    checksum = get_artifact(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(
+    # use `get_artifact` to get "flashinferMetaInfo.h"
+    metainfo = get_artifact(
         f"{include_path}/{header_name}.h",
         meta_hash,
     )

{flashinfer_python-0.6.7 → flashinfer_python-0.6.7.post1}/flashinfer/jit/cubin_loader.py

@@ -26,7 +26,6 @@ import uuid
 
 import filelock
 
-from .utils import write_if_different
 from .core import logger
 from .env import FLASHINFER_CUBIN_DIR
 
@@ -198,121 +197,90 @@ def load_cubin(cubin_path: str, sha256: str) -> bytes:
     return b""
 
 
-def get_file(
-    uri_path: str,
-    sha256: str,
-    file_path: str,
-    session=None,
-) -> bytes:
-    """
-    Load a file from local cache directory {file_path}, ensure that the sha256 signature matches.
-    Otherwise, download the file from {uri_path} and write to {file_path}.
-    """
+def get_artifact(file_name: str, sha256: str, session=None) -> bytes:
+    """Load an artifact (cubin, header, checksum, etc.) from the local cache.
 
-    file = load_cubin(file_path, sha256)
-    if file:
-        return file
-    os.makedirs(os.path.dirname(file_path), exist_ok=True)
-    uri = safe_urljoin(FLASHINFER_CUBINS_REPOSITORY, uri_path)
-    logger.info(f"Fetching file from {uri}")
-    download_file(uri, file_path, session=session)
-    return load_cubin(file_path, sha256)
+    Checks ``FLASHINFER_CUBIN_DIR / file_name`` first.  If the file is missing
+    or its SHA-256 doesn't match, it is downloaded from
+    ``FLASHINFER_CUBINS_REPOSITORY``.
 
-
-def get_cubin(file_name: str, sha256: str, session=None) -> bytes:
+    Returns the file contents as bytes, or empty bytes on failure.
     """
-    Load a cubin from the local cache directory with {file_name} and
-    ensure that the sha256 signature matches.
+    local_path = str(FLASHINFER_CUBIN_DIR / file_name)
+    data = load_cubin(local_path, sha256)
+    if data:
+        return data
+
+    if os.getenv("FLASHINFER_NO_DOWNLOAD"):
+        raise RuntimeError(
+            f"Artifact not found locally: {file_name} "
+            f"(looked at {local_path}). "
+            f"FLASHINFER_NO_DOWNLOAD is set — refusing to download. "
+            f"This means flashinfer-cubin is missing this file."
+        )
 
-    If the kernel does not exist in the cache, it will downloaded.
+    os.makedirs(os.path.dirname(local_path), exist_ok=True)
+    uri = safe_urljoin(FLASHINFER_CUBINS_REPOSITORY, file_name)
+    logger.info(f"Fetching {file_name} from {uri}")
+    download_file(uri, local_path, session=session)
+    return load_cubin(local_path, sha256)
 
-    Returns:
-    None on failure.
-    """
-    cubin_path = str(FLASHINFER_CUBIN_DIR / file_name)
-    cubin = load_cubin(cubin_path, sha256)
-    if cubin:
-        return cubin
-    # either the file does not exist or it is corrupted, we'll download a new one.
 
-    uri = safe_urljoin(FLASHINFER_CUBINS_REPOSITORY, file_name)
-    logger.info(f"Fetching cubin {file_name} from {uri}")
-    download_file(uri, cubin_path, session=session)
-    return load_cubin(cubin_path, sha256)
+# Backward-compatible alias
+get_cubin = get_artifact
 
 
-def download_trtllm_headers(
-    op: str,
-    header_dest_dir: Union[str, pathlib.Path],
-    header_path: str,
-    artifact_path: str,
-    checksum: bytes,
-):
-    header_dest_dir = pathlib.Path(header_dest_dir)
-
-    if op == "bmm":
-        header_files = [
-            "BatchedGemmEnums.h",
-            "BatchedGemmInterface.h",
-            "BatchedGemmOptions.h",
-            "Enums.h",
-            "GemmGatedActOptions.h",
-            "GemmOptions.h",
-            "KernelParams.h",
-            "KernelParamsDecl.h",
-            "KernelTraits.h",
-            "TmaDescriptor.h",
-            "trtllm/gen/CommonUtils.h",
-            "trtllm/gen/CudaArchDecl.h",
-            "trtllm/gen/CudaKernelLauncher.h",
-            "trtllm/gen/DtypeDecl.h",
-            "trtllm/gen/MmaDecl.h",
-            "trtllm/gen/SfLayoutDecl.h",
-            "trtllm/gen/SparsityDecl.h",
-        ]
-
-    else:
-        header_files = [
-            "GemmInterface.h",
-            "GemmOptions.h",
-            "Enums.h",
-            "KernelTraits.h",
-            "KernelParams.h",
-            "KernelParamsDecl.h",
-            "TmaDescriptor.h",
-            "trtllm/gen/CommonUtils.h",
-            "trtllm/gen/CudaKernelLauncher.h",
-            "trtllm/gen/DtypeDecl.h",
-            "trtllm/gen/MmaDecl.h",
-            "trtllm/gen/SfLayoutDecl.h",
-            "trtllm/gen/CudaArchDecl.h",
-        ]
-
-    artifact_hash_path = header_dest_dir / ".artifact_hash"
-
-    # Check if cached headers are from a different artifact version (e.g. after git checkout)
-    if artifact_hash_path.exists():
-        with open(artifact_hash_path, "r") as f:
-            cached_hash = f.read().strip()
-        if cached_hash != artifact_path:
+def ensure_symlink(
+    link: Union[str, pathlib.Path], target: Union[str, pathlib.Path]
+) -> None:
+    """Create or update a symlink, removing any stale file/directory at *link*.
+
+    This is used to map C++ include paths (e.g.
+    ``CUBIN_DIR/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export``) to the
+    canonical artifact directory where ``get_artifact()`` stores downloaded files.
+    """
+    link = pathlib.Path(link)
+    target = pathlib.Path(target)
+    if link.is_symlink() or link.exists():
+        if link.is_symlink() and link.resolve() == target.resolve():
+            return  # already correct
+        # Stale symlink or directory from a previous version; remove it.
+        if link.is_symlink() or link.is_file():
+            link.unlink()
+        else:
+            shutil.rmtree(link)
+    link.parent.mkdir(parents=True, exist_ok=True)
+    link.symlink_to(target)
+
+
+def verify_symlinked_headers(
+    symlink_path: Union[str, pathlib.Path],
+    headers: list,
+    checksums: bytes,
+) -> None:
+    """Verify that headers accessible through the symlink match expected checksums.
+
+    This catches stale cached headers after e.g. ``git checkout`` to a branch
+    with a different artifact version.  Each header file is read through the
+    symlink and its SHA-256 is compared against the hash from ``checksums.txt``.
+    """
+    symlink_path = pathlib.Path(symlink_path)
+    for header in headers:
+        header_path = symlink_path / header
+        expected_hash = get_meta_hash(checksums, header)
+        if not header_path.exists():
             raise RuntimeError(
-                f"Detected inconsistent cached artifacts. "
-                f"(Cached trtllm headers were downloaded for artifact "
-                f"'{cached_hash}', but current code expects "
-                f"'{artifact_path}'). "
-                f"Please clear the cache to confirm and allow the new headers to be downloaded: "
-                f"rm -rf {header_dest_dir}."
+                f"Header {header} not found at {header_path}. "
+                f"Try clearing the cache: rm -rf {FLASHINFER_CUBIN_DIR}"
+            )
+        actual_hash = hashlib.sha256(header_path.read_bytes()).hexdigest()
+        if actual_hash != expected_hash:
+            raise RuntimeError(
+                f"Header {header} at {header_path} has wrong checksum "
+                f"(expected {expected_hash}, got {actual_hash}). "
+                f"This can happen after switching branches. "
+                f"Try clearing the cache: rm -rf {FLASHINFER_CUBIN_DIR}"
             )
-
-    for file in header_files:
-        uri_path = f"{header_path}/{file}"
-        file_hash = get_meta_hash(checksum, file)
-        file_path = str(header_dest_dir / file)
-        result = get_file(uri_path, file_hash, file_path)
-        assert result, f"{file} not found"
-
-    # Record which artifact version these headers belong to
-    write_if_different(artifact_hash_path, artifact_path)
 
 
 def convert_to_ctypes_char_p(data: bytes):
@@ -334,7 +302,7 @@ def setup_cubin_loader(dll_path: str) -> None:
 
     def get_cubin_callback(name: bytes, sha256: bytes):
         # Both name and sha256 are bytes (c_char_p)
-        cubin = get_cubin(name.decode("utf-8"), sha256.decode("utf-8"))
+        cubin = get_artifact(name.decode("utf-8"), sha256.decode("utf-8"))
         _LIB.FlashInferSetCurrentCubin(
             convert_to_ctypes_char_p(cubin), ctypes.c_int(len(cubin))
         )

{flashinfer_python-0.6.7 → flashinfer_python-0.6.7.post1}/flashinfer/jit/fused_moe.py

@@ -26,9 +26,34 @@ from .core import (
     sm89_nvcc_flags,
 )
 from .cpp_ext import is_cuda_version_at_least
-from .cubin_loader import get_cubin, get_meta_hash, download_trtllm_headers
+from .cubin_loader import (
+    get_artifact,
+    get_meta_hash,
+    ensure_symlink,
+    verify_symlinked_headers,
+)
 from .gemm.cutlass.generate_kernels import generate_gemm_operations
 
+BMM_EXPORT_HEADERS = [
+    "BatchedGemmEnums.h",
+    "BatchedGemmInterface.h",
+    "BatchedGemmOptions.h",
+    "Enums.h",
+    "GemmGatedActOptions.h",
+    "GemmOptions.h",
+    "KernelParams.h",
+    "KernelParamsDecl.h",
+    "KernelTraits.h",
+    "TmaDescriptor.h",
+    "trtllm/gen/CommonUtils.h",
+    "trtllm/gen/CudaArchDecl.h",
+    "trtllm/gen/CudaKernelLauncher.h",
+    "trtllm/gen/DtypeDecl.h",
+    "trtllm/gen/MmaDecl.h",
+    "trtllm/gen/SfLayoutDecl.h",
+    "trtllm/gen/SparsityDecl.h",
+]
+
 
 def gen_cutlass_fused_moe_sm120_module(use_fast_build: bool = False) -> JitSpec:
     nvcc_flags = [
@@ -221,30 +246,32 @@ def gen_trtllm_gen_fused_moe_sm100_module() -> JitSpec:
 
     # 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)
+    checksum = get_artifact(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(
+    # use `get_artifact` to get "flashinferMetaInfo.h"
+    metainfo = get_artifact(
         f"{include_path}/{header_name}.h",
         meta_hash,
     )
     # make sure "flashinferMetaInfo.h" is downloaded or cached
     assert metainfo, f"{header_name}.h not found"
 
-    header_path = f"{include_path}/trtllmGen_bmm_export"
-    header_dest_dir = (
+    # Fetch BMM export headers via get_artifact() and symlink for C++ includes.
+    bmm_export_path = f"{include_path}/trtllmGen_bmm_export"
+    for header in BMM_EXPORT_HEADERS:
+        h = get_artifact(f"{bmm_export_path}/{header}", get_meta_hash(checksum, header))
+        assert h, f"{header} not found"
+    symlink_path = (
         jit_env.FLASHINFER_CUBIN_DIR
         / "flashinfer"
         / "trtllm"
         / "batched_gemm"
         / "trtllmGen_bmm_export"
     )
-
-    download_trtllm_headers(
-        "bmm", header_dest_dir, header_path, ArtifactPath.TRTLLM_GEN_BMM, checksum
-    )
+    ensure_symlink(symlink_path, jit_env.FLASHINFER_CUBIN_DIR / bmm_export_path)
+    verify_symlinked_headers(symlink_path, BMM_EXPORT_HEADERS, checksum)
 
     # currently only support Blackwell
     nvcc_flags = current_compilation_context.get_nvcc_flags_list(

{flashinfer_python-0.6.7 → flashinfer_python-0.6.7.post1}/flashinfer/jit/gemm/core.py

@@ -32,7 +32,7 @@ from ..core import (
 )
 from ..cubin_loader import (
     # download_trtllm_headers,
-    get_cubin,
+    get_artifact,
     get_meta_hash,
 )
 from ..utils import dtype_cutlass_map, filename_safe_dtype_map, write_if_different
@@ -540,12 +540,12 @@ def gen_trtllm_gen_gemm_module() -> JitSpec:
 
     # 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)
+    checksum = get_artifact(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(
+    # use `get_artifact` to get "flashinferMetaInfo.h"
+    metainfo = get_artifact(
         f"{include_path}/{header_name}.h",
         meta_hash,
     )
@@ -709,12 +709,12 @@ def gen_trtllm_low_latency_gemm_module() -> JitSpec:
 
     # 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)
+    checksum = get_artifact(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(
+    # use `get_artifact` to get "flashinferMetaInfo.h"
+    metainfo = get_artifact(
         f"{include_path}/{header_name}.h",
         meta_hash,
     )

{flashinfer_python-0.6.7 → flashinfer_python-0.6.7.post1}/flashinfer/jit/moe_utils.py

@@ -37,22 +37,31 @@ def gen_moe_utils_module() -> JitSpec:
     # artifacts.py imports from jit.cubin_loader, which triggers jit/__init__.py,
     # which imports this module — so ArtifactPath/CheckSumHash aren't defined yet
     # at module load time if these imports are at the top level.
-    from .cubin_loader import download_trtllm_headers, get_cubin
+    from .cubin_loader import (
+        get_artifact,
+        get_meta_hash,
+        ensure_symlink,
+        verify_symlinked_headers,
+    )
+    from .fused_moe import BMM_EXPORT_HEADERS
     from ..artifacts import ArtifactPath, CheckSumHash
 
-    download_trtllm_headers(
-        "bmm",
+    checksum = get_artifact(
+        f"{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt", CheckSumHash.TRTLLM_GEN_BMM
+    )
+    bmm_export_path = f"{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export"
+    for header in BMM_EXPORT_HEADERS:
+        h = get_artifact(f"{bmm_export_path}/{header}", get_meta_hash(checksum, header))
+        assert h, f"{header} not found"
+    symlink_path = (
         jit_env.FLASHINFER_CUBIN_DIR
         / "flashinfer"
         / "trtllm"
         / "batched_gemm"
-        / "trtllmGen_bmm_export",
-        f"{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export",
-        ArtifactPath.TRTLLM_GEN_BMM,
-        get_cubin(
-            f"{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt", CheckSumHash.TRTLLM_GEN_BMM
-        ),
+        / "trtllmGen_bmm_export"
     )
+    ensure_symlink(symlink_path, jit_env.FLASHINFER_CUBIN_DIR / bmm_export_path)
+    verify_symlinked_headers(symlink_path, BMM_EXPORT_HEADERS, checksum)
     nvcc_flags = [
         "-DTLLM_GEN_EXPORT_INTERFACE",  # Use relative includes in downloaded headers
         "-DENABLE_BF16",

{flashinfer_python-0.6.7 → flashinfer_python-0.6.7.post1}/flashinfer_python.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashinfer-python
-Version: 0.6.7
+Version: 0.6.7.post1
 Summary: FlashInfer: Kernel Library for LLM Serving
 Author: FlashInfer team
 License-Expression: Apache-2.0

flashinfer_python-0.6.7/flashinfer/_build_meta.py

@@ -1,3 +0,0 @@
-"""Build metadata for flashinfer package."""
-__version__ = "0.6.7"
-__git_version__ = "56ed54098411946525da17eca2bcfd7ec72a515e"