Replace fused Butina Python/Triton path with CUDA backend

[Matthew-Neba]

Summary

• Replaced the fused Butina Python/Triton implementation with a CUDA backend.
• Keeps the full reordering clustering loop on the GPU with CUDA Graphs
• Avoids CPU-GPU round trips, Python loop/object overhead and materializing an N x N hit matrix.
• Supports Tanimoto and cosine fingerprint similarity metrics, just like the previous version. Although, it can be extended for any fingerprint similarity metric
• Fixed a CUDA stream synchronization bug in CPU-returning similarity operations that caused flaky results and tests.
• Unified regular and fused Butina outputs to return device-side cluster IDs and optional centroids through a shared result type.
• Consolidated Butina code and added shared helpers for CUDA graphs, fingerprint preparation, and Tanimoto/cosine calculations.
• Standardized packed fingerprints on uint32 while retaining int32 input support

Performance

About 12-15x faster than the previous Python implementation on my machine.

[greptile-apps]

Greptile Summary

This PR replaces the Python/Triton fused Butina path with a new CUDA backend (src/fused_butina.cu) that keeps the entire reordering clustering loop on-device using CUDA Graphs conditional while-nodes. It also fixes a CPU-returning similarity stream-synchronization bug by threading an input-stream event through crossSimilarityImpl, unifies Butina result types into a shared ButinaResult struct, and standardizes packed fingerprints on uint32.

• New CUDA fused Butina (src/fused_butina.cu): computes an initial O(N²) neighbour scan with a tiled initialNeighborCountKernel, then drives a graph-captured WHILE loop (extract → record → update → argmax) that avoids any CPU–GPU round trips or materialising an N×N hit matrix; all O(N) result buffers stay on-device.
• Stream-sync fix (src/similarity.cpp): crossSimilarityImpl now records a cudaEvent_t on the caller's input stream and makes every internal compute stream call cudaStreamWaitEvent before touching fingerprint data, eliminating the previously flaky CPU-result tests.
• Unified API: both butina() and fused_butina() now return AsyncGpuResult(cluster_ids) (optionally + AsyncGpuResult(centroids)), replacing the old fused path's Python-side (list[tuple], list[int]) output; _fingerprint_inputs is consolidated into _prepare_packed_fingerprints.

Confidence Score: 4/5

The new CUDA graph conditional while-loop requires sm_90+ hardware and CUDA 12.4; any user running fused_butina on Ampere or older GPUs will get a runtime error at graph instantiation.

The new fused Butina path builds its iteration loop around cudaGraphConditionalHandleCreate and cudaGraphCondTypeWhile, both introduced in CUDA 12.4 and executable only on sm_90+ devices. The previous Triton path had no such hardware restriction. The stream-sync fix, fingerprint index arithmetic (consistently using static_caststd::size_t), and unified ButinaResult API are all correct.

src/fused_butina.cu and src/butina_common.cuh — the ConditionalLoopGraph class and its use of cudaGraphConditionalHandleCreate / cudaGraphCondTypeWhile are the surface that limits device compatibility.

Important Files Changed

Filename	Overview
src/fused_butina.cu	New file implementing the CUDA fused Butina pipeline; uses cudaGraphConditionalHandleCreate / cudaGraphCondTypeWhile (sm_90+ / CUDA 12.4 requirement); fingerprint index arithmetic correctly uses static_caststd::size_t throughout, addressing the prior overflow concern.
src/butina_common.cuh	New header owning ConditionalLoopGraph RAII wrapper; well-structured with proper move/copy suppression, but cudaGraphConditionalHandleCreate and cudaGraphCondTypeWhile require sm_90+.
src/fingerprint_similarity_device.cuh	New device-side header consolidating Tanimoto/cosine similarity computation; bounds checks, pruning predicates, and denominator-zero handling are all correct.
src/similarity.cpp	Fixes the CPU-result stream-sync bug: records a cudaEvent_t on the caller's input stream and makes every ScopedStream wait on it before accessing fingerprint buffers; correctly threaded through both single-shot and chunked paths.
nvmolkit/clustering.py	fused_butina now delegates entirely to the C++ backend; metric/stream parameter order swapped vs old signature (pre-existing outside-diff note); both butina() and fused_butina() share _wrap_result for consistent AsyncGpuResult output.
nvmolkit/_fingerprint_inputs.py	New shared helper for packed fingerprint preparation; correctly enqueues tensor moves on the resolved CUDA stream, validates dtype (int32/uint32) and 2D shape, and reinterprets int32 as uint32 in-place.
nvmolkit/similarity.py	Wraps CrossTanimoto/Cosine calls inside with torch.cuda.stream(active_stream) so that kernel dispatch lands on the right stream; CPU-returning variants pass the active stream to C++ for the event-wait fix.
nvmolkit/DataStructs.cpp	CrossTanimotoSimilarityCPURawBuffers / CrossCosineSimilarityCPURawBuffers now accept a stream pointer and forward it to the C++ compute path; duplicate boilerplate eliminated via crossSimilarityCPUFromRawBuffers / crossSimilarityGPUFromRawBuffers templates.
src/butina.h	Old in-place butinaGpu overloads replaced by butinaFromDistanceMatrix / butinaFromHitMatrix / fusedButinaGpu returning the shared ButinaResult struct; well-documented with parameter semantics clarified.
nvmolkit/types.py	Adds np.ascontiguousarray guard before torch.as_tensor for non-C-contiguous NumPy arrays; small but correct correctness fix.
nvmolkit/fingerprints.py	unpack_fingerprint now casts to int64 before shifting (avoids signed-shift UB on uint32 MSB), pack_fingerprint accumulates in int64 and converts to uint32; both functions now validate ndim and dtype explicitly.

_{Reviews (9): Last reviewed commit: "Merge branch 'main' into main" | Re-trigger Greptile}

[scal444]

Thanks for the contribution, I'll be sure to benchmark it!

My main concern is code duplication, but I know the fused/non-fused versions are subtly different enough so that consolidating all the kernels might not make sense. Could you generally take a look and see if anything is worth pulling out into a butina_common.cuh/.cu file?

[scal444]

Hi Matthew - I benchmarked the current PR on my local 5080. I'm seeing good speedups for the smaller cases but perf regressions as it grows.

   Molecules     Main avg.       PR avg.           PR result
  ━━━━━━━━━━━  ━━━━━━━━━━━━  ━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━
          1k      0.563 ms      0.158 ms        3.56× faster
  ───────────  ────────────  ────────────  ──────────────────
          5k      2.398 ms      1.596 ms        1.50× faster
  ───────────  ────────────  ────────────  ──────────────────
         10k      5.951 ms      5.815 ms    effectively tied
  ───────────  ────────────  ────────────  ──────────────────
         20k     17.824 ms     22.064 ms        1.24× slower
  ───────────  ────────────  ────────────  ──────────────────
         30k     31.562 ms     48.922 ms        1.55× slower
  ───────────  ────────────  ────────────  ──────────────────
         40k     48.762 ms     86.065 ms        1.77× slower
  ───────────  ────────────  ────────────  ──────────────────
        100k    226.471 ms    532.878 ms        2.35× slower
  ───────────  ────────────  ────────────  ──────────────────
        500k       4.788 s      27.581 s        5.76× slower

I traced it to the initial neighbors kernel (~95%+ of runtime). - it's scaling quadratically and is less efficient that the original triton. We may need to rework it to be more efficient. If you have time to take a look at it, that would be great, otherwise I can give it a stab later this week.

[Matthew-Neba]

Ok, thanks for the benchmark results. I’ll take a crack at it today

[scal444]

Thanks for all the extra cleanup! Perf now looks great -

   Size          Main    Previous PR    Updated PR         vs main
  ━━━━━━  ━━━━━━━━━━━━  ━━━━━━━━━━━━━  ━━━━━━━━━━━━  ━━━━━━━━━━━━━━
     1k      0.563 ms       0.158 ms      0.148 ms    3.81× faster
  ──────  ────────────  ─────────────  ────────────  ──────────────
     5k      2.398 ms       1.596 ms      0.453 ms    5.30× faster
  ──────  ────────────  ─────────────  ────────────  ──────────────
    10k      5.951 ms       5.815 ms      0.997 ms    5.97× faster
  ──────  ────────────  ─────────────  ────────────  ──────────────
    20k     17.824 ms      22.064 ms      2.905 ms    6.14× faster
  ──────  ────────────  ─────────────  ────────────  ──────────────
    30k     31.562 ms      48.922 ms      6.088 ms    5.18× faster
  ──────  ────────────  ─────────────  ────────────  ──────────────
    40k     48.762 ms      86.065 ms     10.312 ms    4.73× faster
  ──────  ────────────  ─────────────  ────────────  ──────────────
   100k    226.471 ms     532.878 ms     57.488 ms    3.94× faster
  ──────  ────────────  ─────────────  ────────────  ──────────────
   500k       4.788 s       27.581 s       1.531 s    3.13× faster

A few small comments then this can go in.

[scal444]

Again, should not mess with the changelog.

[scal444]

I generally would like to avoid polluting .h with .cuh includes. Any .h file should be include-able by a .cpp file without worrying about if it's bringing in something the requires nvcc. Another thing we haven't been great about throughout the codebase, but I'm trying to enforce it going forward.

[scal444]

I think we'll run into a grid limit crash here on large enough sizes. We had a similar issue with the fused version #194