fix: enumerate artifact downloads from checksums.txt so cute-dsl .so kernels are pre-fetched

flashinfer/artifacts.py

@@ -19,6 +19,7 @@
 import os
 import re
 import time
+from pathlib import PureWindowsPath
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import Generator
 import requests  # type: ignore[import-untyped]
@@ -62,8 +63,10 @@ def get_available_cubin_files(
         try:
             response = requests.get(source, timeout=timeout)
             response.raise_for_status()
-            hrefs = re.findall(r'\<a href=".*\.cubin">', response.text)
-            return tuple((h[9:-8] + ".cubin") for h in hrefs)
+            # Kernel binaries are shipped as .cubin (trtllm-gen, deepgemm)
+            # or .so (cute-dsl, exported via TVM-FFI).
+            hrefs = re.findall(r'<a href="([^"]+\.(?:cubin|so))">', response.text)
+            return tuple(hrefs)
 
         except requests.exceptions.RequestException as e:
             logger.warning(
@@ -220,10 +223,40 @@ def get_checksums(subdirs):
                 f"from {uri}. Check that the pin exists in "
                 f"{FLASHINFER_CUBINS_REPOSITORY} and is reachable."
             )
+        # Verify the manifest against its pinned SHA-256 *before* parsing it:
+        # its entries become download paths and per-file checksums, so a
+        # tampered manifest must be rejected up front, not discovered after
+        # files derived from it have already been written.
+        pinned_sha = CheckSumHash.map_checksums.get(
+            safe_urljoin(subdir, "checksums.txt")
+        )
+        if pinned_sha is not None and not verify_cubin(str(checksum_path), pinned_sha):
+            raise RuntimeError(
+                f"Checksum manifest for artifact pin '{subdir}' does not match "
+                f"its pinned SHA-256; refusing to parse it. Delete "
+                f"'{checksum_path}' and retry."
+            )
         with open(checksum_path, "r") as f:
             for line in f:
                 sha256, filename = line.strip().split()
 
+                # Manifest entries are joined onto FLASHINFER_CUBIN_DIR and
+                # downloaded to; never accept a name that could escape it
+                # (absolute, parent-relative, backslash-separated, or
+                # drive-qualified like ``C:/x.so``, which would replace the
+                # cache root when joined on Windows).
+                if (
+                    "\\" in filename
+                    or filename.startswith("/")
+                    or ".." in filename.split("/")
+                    or PureWindowsPath(filename).drive
+                ):
+                    raise RuntimeError(
+                        f"Unsafe filename {filename!r} in checksum manifest for "
+                        f"artifact pin '{subdir}'; refusing to use it as a "
+                        f"download path."
+                    )
+
                 # Key every entry by its full path. Bare filenames are not
                 # unique across subdirs: two pins built from different sources
                 # can ship identically named kernels, so a flat dict would let
@@ -244,7 +277,6 @@ def _get_host_cpu_arch() -> str:
 
 
 def get_subdir_file_list() -> Generator[tuple[str, str], None, None]:
-    base = FLASHINFER_CUBINS_REPOSITORY
     cpu_arch = _get_host_cpu_arch()
 
     cubin_dirs = [
@@ -263,35 +295,35 @@ def get_subdir_file_list() -> Generator[tuple[str, str], None, None]:
     checksums = get_checksums(cubin_dirs)
 
     # The meta info header files first.
-    yield (
+    meta_info_headers = (
         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")
-        ],
     )
+    for header_path in meta_info_headers:
+        yield (header_path, checksums[header_path])
 
-    # All the actual kernel cubin's.
+    # The checksum manifests themselves, pinned by CheckSumHash.
     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)):
-            full_path = safe_urljoin(cubin_dir, name)
-            yield (full_path, checksums[full_path])
-        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])
+
+    # Everything else each directory's checksums.txt manifest lists.
+    #
+    # The manifest is authoritative for a directory's contents: it is generated
+    # by the cubin publishing pipeline, already fetched by get_checksums(), and
+    # pinned by SHA-256 via CheckSumHash.map_checksums. Enumerating downloads
+    # from it (instead of scraping the artifactory HTML index for hard-coded
+    # extensions) guarantees every published artifact is downloaded and
+    # checksum-verified by download_artifacts(). Scraping silently dropped any
+    # file the regex did not anticipate -- the cute-dsl FMHA kernels ship as
+    # .so, so none of them were ever pre-fetched and each was lazily fetched
+    # over HTTP inside the first forward pass that needed it (#4432). It also
+    # made a failed/unparseable index listing indistinguishable from an empty
+    # directory, silently skipping the whole directory.
+    for file_path, checksum in checksums.items():
+        if file_path not in meta_info_headers:
+            yield (file_path, checksum)
 
 
 def download_artifacts() -> None:

tests/test_artifacts.py

@@ -4,6 +4,8 @@
     get_subdir_file_list,
 )
 
+import hashlib
+
 import pytest
 import responses
 
@@ -249,6 +251,61 @@ def test_get_available_cubin_files():
         )
 
 
+# Directory index of a cute-dsl DSL_FMHA arch directory: the kernels there are
+# TVM-FFI shared objects (.so), not .cubin files (#4432). Mixed with a stray
+# .cubin and non-kernel files to check the enumerator keeps both kernel
+# extensions and nothing else.
+success_dsl_fmha_response = """
+<!DOCTYPE html>
+<html>
+    <head>
+        <meta name="robots" content="noindex"/>
+        <title>Index of sw-kernelinferencelibrary-public-generic-local/5b34f84266cbc2135066ce96885b664992535670/fmha/cute-dsl/x86_64/sm_103a</title>
+    </head>
+    <body>
+        <h1>Index of sw-kernelinferencelibrary-public-generic-local/5b34f84266cbc2135066ce96885b664992535670/fmha/cute-dsl/x86_64/sm_103a</h1>
+        <pre>Name                                          Last modified      Size</pre>
+        <hr/>
+        <pre>
+            <a href="../">../</a>
+            <a href="cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_pdl_tvmffi.so">cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_pdl_tvmffi.so</a>
+            03-Sep-2025 03:45  1.2 MB
+<a href="cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so">cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so</a>
+            03-Sep-2025 03:45  1.2 MB
+<a href="some_kernel.cubin">some_kernel.cubin</a>
+            03-Sep-2025 03:45  60.79 KB
+<a href="checksums.txt">checksums.txt</a>
+            03-Sep-2025 03:45  40.12 KB
+<a href="LICENSE">LICENSE</a>
+            03-Sep-2025 03:45  11.09 KB
+
+        </pre>
+        <hr/>
+        <address style="font-size:small;">Artifactory/7.117.14 Server</address>
+    </body>
+</html>
+"""
+
+expected_dsl_fmha_kernel_files = {
+    "cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_pdl_tvmffi.so",
+    "cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so",
+    "some_kernel.cubin",
+}
+
+
+@responses.activate
+def test_get_available_cubin_files_matches_so():
+    """Regression for #4432: kernel .so files (cute-dsl) must be enumerated
+    alongside .cubin files; non-kernel files must still be excluded."""
+    source = safe_urljoin(
+        test_cubin_repository,
+        safe_urljoin(artifact_paths.DSL_FMHA, "x86_64/sm_103a/"),
+    )
+    responses.add(responses.GET, source, body=success_dsl_fmha_response, status=200)
+    available_files = get_available_cubin_files(source, retries=1, delay=0, timeout=5)
+    assert set(available_files) == expected_dsl_fmha_kernel_files
+
+
 @responses.activate
 def test_get_available_cubin_files_non_200_response():
     """Test that non-200 response codes return an empty tuple."""
@@ -310,9 +367,15 @@ def test_get_checksums_falls_back_to_cached_manifest(monkeypatch, tmp_path):
     monkeypatch.setattr(artifacts, "FLASHINFER_CUBIN_DIR", cubin_dir)
     monkeypatch.setattr(artifacts, "download_file", lambda *args, **kwargs: False)
 
+    manifest_body = "abc123 kernel.fp8_m_grouped_gemm.007d9ebdca7e.cubin\n"
     cached = cubin_dir / safe_urljoin(artifact_paths.DEEPGEMM, "checksums.txt")
     cached.parent.mkdir(parents=True)
-    cached.write_text("abc123 kernel.fp8_m_grouped_gemm.007d9ebdca7e.cubin\n")
+    cached.write_text(manifest_body)
+    monkeypatch.setitem(
+        artifacts.CheckSumHash.map_checksums,
+        safe_urljoin(artifact_paths.DEEPGEMM, "checksums.txt"),
+        hashlib.sha256(manifest_body.encode()).hexdigest(),
+    )
 
     checksums = artifacts.get_checksums([artifact_paths.DEEPGEMM])
     assert checksums == {
@@ -323,6 +386,63 @@ def test_get_checksums_falls_back_to_cached_manifest(monkeypatch, tmp_path):
     }
 
 
+def test_get_checksums_rejects_tampered_manifest(monkeypatch, tmp_path):
+    """A manifest that does not match its pinned SHA-256 must not be parsed.
+
+    Its entries become download paths and per-file checksums, so a tampered
+    manifest has to be rejected before parsing, not discovered afterwards.
+    """
+    from flashinfer import artifacts
+
+    cubin_dir = tmp_path / "cubins"
+    monkeypatch.setattr(artifacts, "FLASHINFER_CUBIN_DIR", cubin_dir)
+    monkeypatch.setattr(artifacts, "download_file", lambda *args, **kwargs: False)
+
+    cached = cubin_dir / safe_urljoin(artifact_paths.DEEPGEMM, "checksums.txt")
+    cached.parent.mkdir(parents=True)
+    cached.write_text("abc123 kernel.fp8_m_grouped_gemm.007d9ebdca7e.cubin\n")
+    monkeypatch.setitem(
+        artifacts.CheckSumHash.map_checksums,
+        safe_urljoin(artifact_paths.DEEPGEMM, "checksums.txt"),
+        "0" * 64,
+    )
+
+    with pytest.raises(RuntimeError) as excinfo:
+        artifacts.get_checksums([artifact_paths.DEEPGEMM])
+    assert "pinned SHA-256" in str(excinfo.value)
+
+
+def test_get_checksums_rejects_traversal_filenames(monkeypatch, tmp_path):
+    """Manifest entries are joined onto FLASHINFER_CUBIN_DIR; absolute paths,
+    ``..`` segments and drive-qualified names must be rejected so a manifest
+    can never direct a write outside the cubin cache."""
+    from flashinfer import artifacts
+
+    cubin_dir = tmp_path / "cubins"
+    monkeypatch.setattr(artifacts, "FLASHINFER_CUBIN_DIR", cubin_dir)
+    monkeypatch.setattr(artifacts, "download_file", lambda *args, **kwargs: False)
+
+    for bad_name in (
+        "../../outside.so",
+        "/etc/evil.so",
+        "a\\..\\b.cubin",
+        "C:/outside.so",
+    ):
+        manifest_body = f"abc123 {bad_name}\n"
+        cached = cubin_dir / safe_urljoin(artifact_paths.DEEPGEMM, "checksums.txt")
+        cached.parent.mkdir(parents=True, exist_ok=True)
+        cached.write_text(manifest_body)
+        monkeypatch.setitem(
+            artifacts.CheckSumHash.map_checksums,
+            safe_urljoin(artifact_paths.DEEPGEMM, "checksums.txt"),
+            hashlib.sha256(manifest_body.encode()).hexdigest(),
+        )
+
+        with pytest.raises(RuntimeError) as excinfo:
+            artifacts.get_checksums([artifact_paths.DEEPGEMM])
+        assert "Unsafe filename" in str(excinfo.value)
+
+
 @responses.activate
 def test_get_subdir_file_list(monkeypatch, tmp_path):
     _mock_file_index_responses()
@@ -391,10 +511,15 @@ def test_get_subdir_file_list(monkeypatch, tmp_path):
 
     # Mock DSL_FMHA checksums + directory index for the host cpu_arch.
     # Pin to x86_64 so the test is deterministic regardless of the runner arch.
+    # The cute-dsl kernels ship as .so, not .cubin (#4432).
     monkeypatch.setattr(artifacts, "_get_host_cpu_arch", lambda: "x86_64")
-    checksums_dsl_fmha = "aabbccdd11223344 cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so\n"
+    checksums_dsl_fmha = """aabbccdd11223344 cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so
+bbccddee22334455 cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_pdl_tvmffi.so
+"""
     # Minimal directory index: an empty HTML page with no cubin/header hrefs.
-    # This avoids 404 retry overhead while still exercising the code path.
+    # Enumeration is driven by the checksums.txt manifest, so the index body
+    # must not matter; this one is registered only so a regression back to
+    # HTML scraping fails fast (empty listing) instead of retrying 404s.
     empty_dir_index = '<html><body><pre><a href="../">../</a></pre></body></html>'
     for sm_arch in artifact_paths.DSL_FMHA_ARCHS:
         subdir = safe_urljoin(artifact_paths.DSL_FMHA, f"x86_64/{sm_arch}/")
@@ -411,6 +536,27 @@ def test_get_subdir_file_list(monkeypatch, tmp_path):
             status=200,
         )
 
+    # get_checksums() refuses to parse a manifest that does not match its
+    # pinned SHA-256, so pin every mocked manifest body for this test.
+    mocked_manifests = {
+        artifact_paths.TRTLLM_GEN_FMHA: checksums_fmha,
+        artifact_paths.TRTLLM_GEN_GEMM: checksums_gemm,
+        artifact_paths.TRTLLM_GEN_BMM: checksums_bmm,
+        artifact_paths.DEEPGEMM: checksums_deepgemm,
+        **{
+            safe_urljoin(
+                artifact_paths.DSL_FMHA, f"x86_64/{sm_arch}/"
+            ): checksums_dsl_fmha
+            for sm_arch in artifact_paths.DSL_FMHA_ARCHS
+        },
+    }
+    for manifest_subdir, manifest_body in mocked_manifests.items():
+        monkeypatch.setitem(
+            artifacts.CheckSumHash.map_checksums,
+            safe_urljoin(manifest_subdir, "checksums.txt"),
+            hashlib.sha256(manifest_body.encode()).hexdigest(),
+        )
+
     cubin_files = list(get_subdir_file_list())
 
     # Extract just the file paths from the (path, checksum) tuples
@@ -463,6 +609,54 @@ def test_get_subdir_file_list(monkeypatch, tmp_path):
     by_path = dict(cubin_files)
     assert len(by_path) == len(cubin_files), "duplicate paths in cubin file list"
 
+    # Regression for #4432: the cute-dsl FMHA kernels are .so files, which the
+    # old HTML-scraping enumerator (cubin/header regexes only) silently
+    # skipped, so download_artifacts() reported success while every DSL kernel
+    # was missing from the cache. Every manifest-listed .so must be enumerated
+    # for every arch, carrying the checksum from its own manifest.
+    for sm_arch in artifact_paths.DSL_FMHA_ARCHS:
+        subdir = safe_urljoin(artifact_paths.DSL_FMHA, f"x86_64/{sm_arch}/")
+        for so_name, so_sha in (
+            (
+                "cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so",
+                "aabbccdd11223344",
+            ),
+            (
+                "cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_pdl_tvmffi.so",
+                "bbccddee22334455",
+            ),
+        ):
+            so_path = safe_urljoin(subdir, so_name)
+            assert so_path in by_path, (
+                f"DSL FMHA kernel '{so_path}' not enumerated -- .so artifacts "
+                f"would be silently skipped by download_artifacts() (#4432)"
+            )
+            assert by_path[so_path] == so_sha
+
+    # Mixed-content directory: enumeration is manifest-driven, so files with
+    # extensions the old scraper never anticipated (e.g. deepgemm's
+    # kernel_map.json) must be enumerated too, not only .cubin/.h files.
+    kernel_map_path = safe_urljoin(artifact_paths.DEEPGEMM, "kernel_map.json")
+    assert kernel_map_path in by_path
+
+    # Every entry of every manifest must be enumerated, so a file that is
+    # listed but missing on the server now fails download_artifacts() loudly
+    # instead of being silently skipped.
+    manifest_entries = artifacts.get_checksums(
+        [
+            artifact_paths.TRTLLM_GEN_FMHA,
+            artifact_paths.TRTLLM_GEN_BMM,
+            artifact_paths.TRTLLM_GEN_GEMM,
+            artifact_paths.DEEPGEMM,
+        ]
+        + [
+            safe_urljoin(artifact_paths.DSL_FMHA, f"x86_64/{sm_arch}/")
+            for sm_arch in artifact_paths.DSL_FMHA_ARCHS
+        ]
+    )
+    missing = set(manifest_entries) - set(by_path)
+    assert not missing, f"manifest entries not enumerated for download: {missing}"
+
 
 def test_get_checksums_keys_by_full_path(monkeypatch, tmp_path):
     """Two pins shipping the same kernel filename must keep separate hashes.