From 117f5277df719a26c2a53b6b141737e1dbcae72a Mon Sep 17 00:00:00 2001 From: Alex Yang Date: Tue, 28 Jul 2026 01:56:08 -0700 Subject: [PATCH] Key cubin checksums by full path to prevent cross-pin collisions get_checksums() flattens every subdir's checksums.txt into one dict keyed by bare filename, namespacing only ".h" entries. Bare filenames are not unique across subdirs: per-arch cubin pins publish the same sm100f/sm103a kernel names built from different sources, so whichever subdir is processed last silently overwrites the earlier one's hashes. Every shared name is then verified against the wrong pin's hash and download_artifacts() fails with "Failed to download cubins: checksum mismatch". main cannot trigger this today because it defines only a single BMM/GEMM pin. It did trigger on release-v0.6.16, which carries a second, Rubin-specific pin: all 2534 shared BMM filenames and 105 of 108 shared GEMM filenames resolved to the wrong hash, taking down build-flashinfer-cubin for the whole release. This is a latent landmine on main that arms itself the moment a second pin lands. Key every entry by full path, and look cubins up the same way headers already were. Add a pin-independent regression test that feeds get_checksums() two subdirs publishing the same kernel name with different hashes; against the unfixed version it fails on the bare filename still being present as a key. Co-Authored-By: Claude Opus 5 --- flashinfer/artifacts.py | 14 +++++---- tests/test_artifacts.py | 64 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 5 deletions(-) diff --git a/flashinfer/artifacts.py b/flashinfer/artifacts.py index 0d0c9602bc2..6e3fec1be82 100644 --- a/flashinfer/artifacts.py +++ b/flashinfer/artifacts.py @@ -207,10 +207,13 @@ def get_checksums(subdirs): for line in f: sha256, filename = line.strip().split() - # Distinguish between all meta info header files - if ".h" in filename: - filename = safe_urljoin(subdir, filename) - checksums[filename] = sha256 + # Key every entry by its full path. Bare filenames are not + # unique across subdirs: per-arch cubin pins publish the same + # kernel names built from different sources, so a flat dict + # lets whichever subdir is processed last silently overwrite + # the earlier one's hashes, failing verification for every + # shared name. + checksums[safe_urljoin(subdir, filename)] = sha256 return checksums @@ -268,7 +271,8 @@ def get_subdir_file_list() -> Generator[tuple[str, str], None, None]: 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)): - yield (safe_urljoin(cubin_dir, name), checksums[name]) + 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]) diff --git a/tests/test_artifacts.py b/tests/test_artifacts.py index 6dabf1884bc..5d433bfc3fa 100644 --- a/tests/test_artifacts.py +++ b/tests/test_artifacts.py @@ -1,6 +1,7 @@ from flashinfer.artifacts import ( ArtifactPath, get_available_cubin_files, + get_checksums, get_subdir_file_list, ) @@ -413,3 +414,66 @@ def test_get_subdir_file_list(monkeypatch, tmp_path): assert len(meta_info_headers) == 3, ( f"Meta info headers count mismatch. Expected 3, got {len(meta_info_headers)}. Headers found: {meta_info_headers}" ) + + +@responses.activate +def test_get_checksums_does_not_collide_across_subdirs(monkeypatch, tmp_path): + """Two pins publishing the same kernel filename must keep separate hashes. + + Per-arch cubin pins ship identically-named sm100f/sm103a kernels built from + different sources. Keying the checksum map by bare filename let whichever + subdir was processed last overwrite the earlier one, so every shared name + was then verified against the wrong pin's hash. + """ + from flashinfer import artifacts + + temp_cubin_dir = tmp_path / "cubins" + temp_cubin_dir.mkdir(exist_ok=True) + monkeypatch.setattr( + artifacts, "FLASHINFER_CUBINS_REPOSITORY", test_cubin_repository + ) + monkeypatch.setattr(artifacts, "FLASHINFER_CUBIN_DIR", temp_cubin_dir) + + shared_kernel = "Bmm_Bfloat16_E2m1E2m1_Fp32_shared_name_sm100f.cubin" + subdir_a = "pin-a/batched_gemm-aaaaaaa-bbbbbbb/" + subdir_b = "pin-b/batched_gemm-ccccccc-ddddddd/" + hash_a = "1111111111111111111111111111111111111111111111111111111111111111" + hash_b = "2222222222222222222222222222222222222222222222222222222222222222" + header_a = "3333333333333333333333333333333333333333333333333333333333333333" + header_b = "4444444444444444444444444444444444444444444444444444444444444444" + + for subdir, kernel_hash, header_hash in ( + (subdir_a, hash_a, header_a), + (subdir_b, hash_b, header_b), + ): + responses.add( + responses.GET, + safe_urljoin(test_cubin_repository, safe_urljoin(subdir, "checksums.txt")), + body=( + f"{header_hash} include/flashinferMetaInfo.h\n" + f"{kernel_hash} {shared_kernel}\n" + ), + status=200, + ) + + checksums = get_checksums([subdir_a, subdir_b]) + + key_a = safe_urljoin(subdir_a, shared_kernel) + key_b = safe_urljoin(subdir_b, shared_kernel) + + # The bare filename must not be a key at all -- that is what collided. + assert shared_kernel not in checksums + + assert checksums[key_a] == hash_a, ( + f"{subdir_a} kernel resolved to {checksums[key_a]}, expected {hash_a}" + ) + assert checksums[key_b] == hash_b, ( + f"{subdir_b} kernel resolved to {checksums[key_b]}, expected {hash_b}" + ) + assert checksums[key_a] != checksums[key_b], ( + "both pins resolved to the same checksum -- the per-pin hashes collided" + ) + + # Headers were already namespaced before this fix; keep them covered. + assert checksums[safe_urljoin(subdir_a, "include/flashinferMetaInfo.h")] == header_a + assert checksums[safe_urljoin(subdir_b, "include/flashinferMetaInfo.h")] == header_b