From 76f6036f70e064d375974d66801a76ad0bd903f2 Mon Sep 17 00:00:00 2001 From: Azra Bano Date: Sun, 16 Aug 2026 00:23:00 -0700 Subject: [PATCH 1/3] fix: enumerate artifact downloads from checksums.txt so cute-dsl .so kernels are pre-fetched download_artifacts() enumerated files by scraping the artifactory HTML index with a .cubin-only regex, so the DSL_FMHA (cute-dsl) kernels, which ship as .so files, were never downloaded even though their directories are explicitly listed in get_subdir_file_list(). Each directory's checksums.txt was downloaded regardless, so the cache looked complete while every kernel it lists was absent, and the missing kernels were then fetched over HTTP lazily inside the first model forward pass that needed them. Drive the enumeration from the per-directory checksums.txt manifests instead: they are authoritative for the directory contents, are already fetched by get_checksums(), and are themselves SHA-256 pinned via CheckSumHash.map_checksums. This also means every manifest entry is downloaded and checksum-verified by download_artifacts(), so a listed-but-missing file now fails loudly instead of silently, and the per-directory HTML index round-trips (and their retry stalls) are gone. Also widen the retained get_available_cubin_files() helper to match .so hrefs in addition to .cubin. Fixes #4432 Signed-off-by: Azra Bano --- flashinfer/artifacts.py | 63 +++++++++------------- tests/test_artifacts.py | 115 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 138 insertions(+), 40 deletions(-) diff --git a/flashinfer/artifacts.py b/flashinfer/artifacts.py index 9effb776659..4f367ebde77 100644 --- a/flashinfer/artifacts.py +++ b/flashinfer/artifacts.py @@ -62,8 +62,10 @@ def get_available_cubin_files( try: response = requests.get(source, timeout=timeout) response.raise_for_status() - hrefs = re.findall(r'\', 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'', response.text) + return tuple(hrefs) except requests.exceptions.RequestException as e: logger.warning( @@ -267,7 +269,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 = [ @@ -289,53 +290,39 @@ 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") - ], - ) - yield ( safe_urljoin( ArtifactPath.TRTLLM_GEN_GEMM_RUBIN, "include/flashinferMetaInfo.h" ), - checksums[ - safe_urljoin( - ArtifactPath.TRTLLM_GEN_GEMM_RUBIN, "include/flashinferMetaInfo.h" - ) - ], - ) - yield ( safe_urljoin(ArtifactPath.TRTLLM_GEN_BMM_RUBIN, "include/flashinferMetaInfo.h"), - checksums[ - safe_urljoin( - ArtifactPath.TRTLLM_GEN_BMM_RUBIN, "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: diff --git a/tests/test_artifacts.py b/tests/test_artifacts.py index 6a20ea63acc..ee2e6d33cd2 100644 --- a/tests/test_artifacts.py +++ b/tests/test_artifacts.py @@ -270,6 +270,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 = """ + + + + + Index of sw-kernelinferencelibrary-public-generic-local/5b34f84266cbc2135066ce96885b664992535670/fmha/cute-dsl/x86_64/sm_103a + + +

Index of sw-kernelinferencelibrary-public-generic-local/5b34f84266cbc2135066ce96885b664992535670/fmha/cute-dsl/x86_64/sm_103a

+ 
Name                                          Last modified      Size
+
+ 
+            ../
+            cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_lse_pdl_tvmffi.so
+            03-Sep-2025 03:45  1.2 MB
+cute_dsl_fmha_bf16_h128_causal_nonpersistent_varlen_tvmffi.so
+            03-Sep-2025 03:45  1.2 MB
+some_kernel.cubin
+            03-Sep-2025 03:45  60.79 KB
+checksums.txt
+            03-Sep-2025 03:45  40.12 KB
+LICENSE
+            03-Sep-2025 03:45  11.09 KB
+
+
+
+
Artifactory/7.117.14 Server
+ + +""" + +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.""" @@ -462,10 +517,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 = '
../
' for sm_arch in artifact_paths.DSL_FMHA_ARCHS: subdir = safe_urljoin(artifact_paths.DSL_FMHA, f"x86_64/{sm_arch}/") @@ -563,3 +623,54 @@ def test_get_subdir_file_list(monkeypatch, tmp_path): f"{shared_name} resolved to the same checksum for both pins " f"({by_path[plain_path]}) -- the per-pin hashes collided" ) + + # 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.TRTLLM_GEN_BMM_RUBIN, + artifact_paths.TRTLLM_GEN_GEMM_RUBIN, + artifact_paths.DEEPGEMM, + artifact_paths.DEEPGEMM_RUBIN, + ] + + [ + 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}" From acad497160bb088ef61cbebd31e85488f638f41c Mon Sep 17 00:00:00 2001 From: Azra Bano Date: Sun, 23 Aug 2026 22:27:05 -0400 Subject: [PATCH 2/3] fix: verify manifest pin before parsing and reject unsafe manifest filenames The checksums.txt manifests are now the enumeration source for downloads, so harden that trust boundary: refuse to parse a manifest that does not match its SHA-256 pin in CheckSumHash.map_checksums, and reject manifest entries that are absolute or contain '..' segments so an entry can never direct a write outside FLASHINFER_CUBIN_DIR. --- flashinfer/artifacts.py | 26 ++++++++++++ tests/test_artifacts.py | 88 ++++++++++++++++++++++++++++++++++++++++- 2 files changed, 113 insertions(+), 1 deletion(-) diff --git a/flashinfer/artifacts.py b/flashinfer/artifacts.py index 4f367ebde77..47018949aaf 100644 --- a/flashinfer/artifacts.py +++ b/flashinfer/artifacts.py @@ -244,10 +244,36 @@ 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. + if ( + "\\" in filename + or filename.startswith("/") + or ".." in filename.split("/") + ): + 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: the per-arch pins (e.g. TRTLLM_GEN_BMM # and TRTLLM_GEN_BMM_RUBIN) ship the same sm100f/sm103a kernel diff --git a/tests/test_artifacts.py b/tests/test_artifacts.py index ee2e6d33cd2..d4514265b81 100644 --- a/tests/test_artifacts.py +++ b/tests/test_artifacts.py @@ -4,6 +4,8 @@ get_subdir_file_list, ) +import hashlib + import pytest import responses @@ -386,9 +388,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_RUBIN, "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_RUBIN, "checksums.txt"), + hashlib.sha256(manifest_body.encode()).hexdigest(), + ) checksums = artifacts.get_checksums([artifact_paths.DEEPGEMM_RUBIN]) assert checksums == { @@ -399,6 +407,60 @@ 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_RUBIN, "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_RUBIN, "checksums.txt"), + "0" * 64, + ) + + with pytest.raises(RuntimeError) as excinfo: + artifacts.get_checksums([artifact_paths.DEEPGEMM_RUBIN]) + 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 + and ``..`` segments 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"): + manifest_body = f"abc123 {bad_name}\n" + cached = cubin_dir / safe_urljoin( + artifact_paths.DEEPGEMM_RUBIN, "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_RUBIN, "checksums.txt"), + hashlib.sha256(manifest_body.encode()).hexdigest(), + ) + + with pytest.raises(RuntimeError) as excinfo: + artifacts.get_checksums([artifact_paths.DEEPGEMM_RUBIN]) + assert "Unsafe filename" in str(excinfo.value) + + @responses.activate def test_get_subdir_file_list(monkeypatch, tmp_path): _mock_file_index_responses() @@ -542,6 +604,30 @@ 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.TRTLLM_GEN_BMM_RUBIN: checksums_bmm_rubin, + artifact_paths.TRTLLM_GEN_GEMM_RUBIN: checksums_gemm_rubin, + artifact_paths.DEEPGEMM: checksums_deepgemm, + artifact_paths.DEEPGEMM_RUBIN: checksums_deepgemm_rubin, + **{ + 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 From 3206fd6de5b7dc81fbb9d5b44e855b035b53924a Mon Sep 17 00:00:00 2001 From: Azra Bano Date: Fri, 4 Sep 2026 11:20:25 -0400 Subject: [PATCH 3/3] fix: reject drive-qualified manifest filenames too `C:/outside.so` passed the existing absolute / `..` / backslash checks, and on Windows joining a drive-qualified name onto FLASHINFER_CUBIN_DIR replaces the cache root instead of nesting under it. Reject any entry with a non-empty PureWindowsPath(...).drive alongside the other checks, and add the case to the traversal regression test. Signed-off-by: Azra Bano Co-Authored-By: Claude Fable 5.1 --- flashinfer/artifacts.py | 7 ++++++- tests/test_artifacts.py | 13 +++++++++---- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/flashinfer/artifacts.py b/flashinfer/artifacts.py index 5d2059ba130..96cdcf8139b 100644 --- a/flashinfer/artifacts.py +++ b/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] @@ -240,11 +241,15 @@ def get_checksums(subdirs): sha256, filename = line.strip().split() # Manifest entries are joined onto FLASHINFER_CUBIN_DIR and - # downloaded to; never accept a name that could escape it. + # 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 " diff --git a/tests/test_artifacts.py b/tests/test_artifacts.py index 3834407a62e..36661c37a33 100644 --- a/tests/test_artifacts.py +++ b/tests/test_artifacts.py @@ -413,16 +413,21 @@ def test_get_checksums_rejects_tampered_manifest(monkeypatch, tmp_path): def test_get_checksums_rejects_traversal_filenames(monkeypatch, tmp_path): - """Manifest entries are joined onto FLASHINFER_CUBIN_DIR; absolute paths - and ``..`` segments must be rejected so a manifest can never direct a - write outside the cubin cache.""" + """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"): + 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)