refactor: pull trtllm-gen batch-gemm/gemm headers from artifactory; update tma descriptor shape init

.gitignore

@@ -13,6 +13,8 @@ flashinfer/_build_meta.py
 flashinfer/data/
 flashinfer/jit/aot_config.py
 csrc/aot_default_additional_params.h
+flashinfer/trtllm/gemm/trtllmGen_gemm_export
+flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export
 
 # DS_Store files
 .DS_store

csrc/trtllm_batched_gemm_runner.cu

@@ -214,6 +214,9 @@ void TrtllmGenBatchedGemmRunner::run(
   gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m;
   gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n;
   gemmData.mProblemDimensions.mK = k;
+  gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM;
+  gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN;
+  gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK;
   gemmData.mProblemDimensions.mRank = 0;
   gemmData.mProblemDimensions.mWorldSize = 1;
 
@@ -252,10 +255,6 @@ void TrtllmGenBatchedGemmRunner::run(
   int32_t multiProcessorCount;
   cudaDeviceGetAttribute(&multiProcessorCount, cudaDevAttrMultiProcessorCount, device);
 
-  gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM;
-  gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN;
-  gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK;
-
   // FIXME once we start using all-reduce in the epilogue of the bmm this can be moved elsewhere
   bmm.runInitBeforeWorldSync(config, gemmData, static_cast<void*>(stream));
 
@@ -449,6 +448,9 @@ bool TrtllmGenBatchedGemmRunner::isValidConfigIndex(int32_t configIndex, int32_t
   gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m;
   gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n;
   gemmData.mProblemDimensions.mK = k;
+  gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM;
+  gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN;
+  gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK;
   gemmData.mProblemDimensions.mRank = 0;
   gemmData.mProblemDimensions.mWorldSize = 1;
   gemmData.mProblemDimensions.mMaxNumCtasInTokenDim = maxNumCtasInBatchDim;

csrc/trtllm_gemm_runner.cu

@@ -119,6 +119,9 @@ class TrtllmGenGemmRunner {
     gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m;
     gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n;
     gemmData.mProblemDimensions.mK = k;
+    gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM;
+    gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN;
+    gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK;
     gemmData.mProblemDimensions.mRank = 0;
     gemmData.mProblemDimensions.mWorldSize = 1;
 
@@ -139,6 +142,9 @@ class TrtllmGenGemmRunner {
     gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m;
     gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n;
     gemmData.mProblemDimensions.mK = k;
+    gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM;
+    gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN;
+    gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK;
     gemmData.mProblemDimensions.mRank = 0;
     gemmData.mProblemDimensions.mWorldSize = 1;
 
@@ -187,6 +193,9 @@ class TrtllmGenGemmRunner {
     gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m;
     gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n;
     gemmData.mProblemDimensions.mK = k;
+    gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM;
+    gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN;
+    gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK;
     gemmData.mProblemDimensions.mRank = 0;
     gemmData.mProblemDimensions.mWorldSize = 1;
 

csrc/trtllm_low_latency_gemm_runner.cu

@@ -49,6 +49,9 @@ gemm::gemm::GemmData createGemmData(int64_t m, int64_t n, int64_t k) {
   gemmData.mProblemDimensions.mM = n;
   gemmData.mProblemDimensions.mN = m;
   gemmData.mProblemDimensions.mK = k;
+  gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM;
+  gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN;
+  gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK;
   gemmData.mProblemDimensions.mRank = 0;
   gemmData.mProblemDimensions.mWorldSize = 1;
 

flashinfer/artifacts.py

@@ -79,6 +79,54 @@ def get_available_cubin_files(
     return tuple()
 
 
+def get_available_header_files(
+    source: str, retries: int = 3, delay: int = 5, timeout: int = 10
+) -> tuple[str, ...]:
+    """
+    Recursively navigates through child directories (e.g., include/) and finds
+    all *.h header files, returning them as a tuple of relative paths.
+    """
+    result: list[str] = []
+
+    def fetch_directory(url: str, prefix: str = "") -> None:
+        for attempt in range(1, retries + 1):
+            try:
+                response = requests.get(url, timeout=timeout)
+                response.raise_for_status()
+
+                # Find all .h header files in this directory
+                header_hrefs = re.findall(r'<a href="([^"]+\.h)">', response.text)
+                for h in header_hrefs:
+                    result.append(prefix + h if prefix else h)
+
+                # Find all subdirectories (links ending with /)
+                dir_hrefs = re.findall(r'<a href="([^"]+/)">', response.text)
+                for d in dir_hrefs:
+                    # Skip parent directory links
+                    if d == "../" or d.startswith(".."):
+                        continue
+                    subdir_url = safe_urljoin(url, d)
+                    subdir_prefix = prefix + d if prefix else d
+                    fetch_directory(subdir_url, subdir_prefix)
+
+                return  # Success, exit retry loop
+
+            except requests.exceptions.RequestException as e:
+                logger.warning(
+                    f"Fetching available header files {url}: attempt {attempt} failed: {e}"
+                )
+
+                if attempt < retries:
+                    logger.info(f"Retrying in {delay} seconds...")
+                    time.sleep(delay)
+
+        logger.error(f"Max retries reached for {url}. Fetch failed.")
+
+    fetch_directory(source)
+    logger.info(f"result: {result}")
+    return tuple(result)
+
+
 @dataclass(frozen=True)
 class ArtifactPath:
     """
@@ -182,6 +230,9 @@ def get_subdir_file_list() -> Generator[tuple[str, str], None, None]:
         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])
+        for name in get_available_header_files(safe_urljoin(base, cubin_dir)):
+            full_path = safe_urljoin(cubin_dir, name)
+            yield (full_path, checksums[full_path])
 
 
 def download_artifacts() -> None:
@@ -190,7 +241,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_subdir_file_list())
+    cubin_files = list[tuple[str, str]](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"

flashinfer/jit/cubin_loader.py

@@ -17,13 +17,16 @@
 import ctypes
 import hashlib
 import os
+import pathlib
 from urllib.parse import urljoin
 import shutil
 import time
+from typing import Union
 import uuid
 
 import filelock
 
+from .utils import write_if_different
 from .core import logger
 from .env import FLASHINFER_CUBIN_DIR
 
@@ -136,14 +139,20 @@ def download_file(
         return False
 
 
-def get_meta_hash(checksums_bytes: bytes) -> str:
+def get_meta_hash(
+    checksums_bytes: bytes, target_file: str = "flashinferMetaInfo.h"
+) -> str:
     """
     Parse the checksums.txt file and get the hash of corresponding flashinferMetaInfo.h file
     """
     checksums_lines = checksums_bytes.decode("utf-8").splitlines()
     for line in checksums_lines:
         sha256, filename = line.strip().split()
-        if ".h" in filename:
+        # Match on path segment boundary to avoid substring collisions
+        # (e.g. "Enums.h" must not match "BatchedGemmEnums.h")
+        if filename.lower() == target_file.lower() or filename.lower().endswith(
+            "/" + target_file.lower()
+        ):
             return sha256
     raise ValueError("Invalid checksums.txt, no flashinferMetaInfo.h found")
 
@@ -189,6 +198,27 @@ 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}.
+    """
+
+    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)
+
+
 def get_cubin(file_name: str, sha256: str, session=None) -> bytes:
     """
     Load a cubin from the local cache directory with {file_name} and
@@ -211,6 +241,80 @@ def get_cubin(file_name: str, sha256: str, session=None) -> bytes:
     return load_cubin(cubin_path, sha256)
 
 
+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:
+            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}."
+            )
+
+    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):
     return ctypes.c_char_p(data)
 

flashinfer/jit/fused_moe.py

@@ -26,7 +26,7 @@
     sm89_nvcc_flags,
 )
 from .cpp_ext import is_cuda_version_at_least
-from .cubin_loader import get_cubin, get_meta_hash
+from .cubin_loader import get_cubin, get_meta_hash, download_trtllm_headers
 from .gemm.cutlass.generate_kernels import generate_gemm_operations
 
 
@@ -233,6 +233,19 @@ def gen_trtllm_gen_fused_moe_sm100_module() -> JitSpec:
     # 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 = (
+        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
+    )
+
     # currently only support Blackwell
     nvcc_flags = current_compilation_context.get_nvcc_flags_list(
         supported_major_versions=[10]
@@ -265,7 +278,7 @@ def gen_trtllm_gen_fused_moe_sm100_module() -> JitSpec:
         ]
         + nvcc_flags,
         extra_include_paths=[
-            # link "include" sub-directory in cache
+            jit_env.FLASHINFER_CUBIN_DIR,
             jit_env.FLASHINFER_CUBIN_DIR / include_path,
             jit_env.FLASHINFER_CSRC_DIR / "nv_internal",
             jit_env.FLASHINFER_CSRC_DIR / "nv_internal/include",