perf(fpindex): GPU fingerprint queries + RapidFuzz edit-distance#2
Open
trvachov wants to merge 1 commit into
Open
perf(fpindex): GPU fingerprint queries + RapidFuzz edit-distance#2trvachov wants to merge 1 commit into
trvachov wants to merge 1 commit into
Conversation
Run fingerprint cdist on the model-assigned GPU and accelerate invalid-SMILES Levenshtein lookup with RapidFuzz. Thread explicit multi-GPU device routing, canonicalize CUDA cache keys, restore legacy CPU topk tie identity, pin RapidFuzz in both manifests, and use workers=1. Add tests/test_fpindex.py for GPU/CPU parity, device routing, cache behavior, fallback paths, and the real index. Reviewed-and-tested: adversarially verified on the 211k fingerprint index with an A100, on CPU-only paths, and against both model checkpoints.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Accelerates the building-block fingerprint lookup in
get_reactants— the single largest cost in ReaSyn inference (~60% of wall time in profiling).Two changes to
FingerprintIndex.query_cuda:cdiston the model's GPU. The query fingerprint arrives as a CPU tensor, so despite thequery_cudaname the distance search was executing on the CPU over the 211k×2048 matrix (~53 ms/query). Routed to the model's assignedcuda:<id>→ ~53 → 1.4 ms.editdistanceloop (~335–460 ms/call — the true dominant cost). Replaced withrapidfuzz.process.cdist(batched C++/SIMD) → ~9–15×, identical distances.Correctness / safety
cdistdistances are bit-identical to CPU; to also preserve the selected indices at tied distances, the distance vector is copied back to CPU fortopk(GPU/CPUtopkbreak ties differently). RapidFuzz distances are identical toeditdistance.cuda:<gpu_id>(not the process-default GPU 0);fp_cudacache key canonicalized.pyproject.tomlandenv.yml;workers=1.Tests
tests/test_fpindex.py— GPU/CPU parity (incl. the real 211k index), device routing, cache behavior, fallback paths, RapidFuzz↔editdistance identity. 9 passed on A100.Toggles
REASYN_GPU_QUERY(default on),REASYN_RAPIDFUZZ(default on); set=0to restore legacy behavior.Independently reviewed and fixed (multi-GPU device routing, topk tie identity, manifest pinning) before opening.
Testing environment⚠️
All test/benchmark results above were produced on the review container: Python 3.12 / PyTorch 2.11 / CUDA 13.2 — not the repo-pinned stack (Python 3.10 / PyTorch 2.7 / CUDA 11.8 per
env.yml/pyproject.toml). Kernel selection, SDPA backends, and tie-breaking are version-sensitive, so these results should be re-confirmed on the pinned stack / CI before merge.