Add GPU-accelerated conformer RMSD matrix computation

[volgin]

Summary

Adds GetConformerRMSMatrix — a GPU-accelerated equivalent of RDKit's AllChem.GetConformerRMSMatrix that computes pairwise Kabsch-aligned RMSD for all conformer pairs on the GPU.

• CUDA kernel: One thread-block (128 threads) per conformer pair. Computes centroids, cross-covariance matrix H, eigenvalues of H^T·H via Cardano's analytical formula for 3×3 symmetric matrices, singular values, reflection handling via det(H), and RMSD — all in a single kernel launch.
• Prealigned mode: When prealigned=True, skips Kabsch alignment and computes raw coordinate RMSD (no centering), matching RDKit's behavior.
• Async API: Returns AsyncGpuResult with __cuda_array_interface__ support, consistent with the existing nvMolKit pattern.
• Benchmark included: benchmarks/conformer_rmsd_bench.py with correctness validation against numpy Kabsch SVD.

Benchmark results (RTX 4090, CUDA 12.4)

Molecule	Heavy atoms	Conformers	Pairs	CPU (ms)	GPU (ms)	Speedup
Aspirin	13	100	4,950	26.6	0.28	94x
Aspirin	13	200	19,900	121.5	0.99	122x
Celecoxib	26	100	4,950	43.1	0.29	147x
Celecoxib	26	200	19,900	190.8	3.59	53x
Osimertinib	37	200	19,900	252.7	3.72	68x
Osimertinib	37	500	124,750	1,889	21.9	86x
Cyclic pentapeptide	48	200	19,900	316.1	3.95	80x
Cyclic pentapeptide	48	500	124,750	2,276	19.0	120x

Fits the existing GPU pipeline

The result can be passed directly to butina() for GPU-accelerated conformer clustering, keeping the entire generate → optimize → RMSD → cluster pipeline on the GPU:

from nvmolkit.conformerRmsd import GetConformerRMSMatrix
from nvmolkit.clustering import butina

rmsd = GetConformerRMSMatrix(mol, prealigned=False)
torch.cuda.synchronize()
n = mol.GetNumConformers()
square = torch.zeros(n, n, device='cuda', dtype=torch.float64)
idx = torch.tril_indices(n, n, offset=-1)
square[idx[0], idx[1]] = rmsd.torch()
square = square + square.T
clusters = butina(square, cutoff=0.5)

Note on RDKit compatibility

Results may differ slightly from RDKit's GetConformerRMSMatrix(prealigned=False) because RDKit modifies conformer coordinates in-place during sequential alignment (each AlignMol call updates the probe conformer, affecting subsequent pairs). This kernel computes each pair independently from the original coordinates, which is the mathematically correct pairwise Kabsch RMSD. Tests validate against an independent numpy SVD-based Kabsch reference.

Test plan

• 12 tests covering correctness (vs numpy Kabsch SVD), prealigned mode, large conformer sets, edge cases, streams, and invalid inputs
• Tested on RTX 4090 (CUDA 12.4)
• Benchmark validates correctness across 4 drug-like molecules (13–48 heavy atoms) and 10–500 conformers
• CI validation on NVIDIA Blossom

🤖 Generated with Claude Code

[greptile-apps]

Greptile Summary

This PR adds GetConformerRMSMatrix — a GPU-accelerated replacement for RDKit's AllChem.GetConformerRMSMatrix — along with a batch variant and benchmark. The CUDA kernel is mathematically sound: it uses Cardano's analytical formula to compute eigenvalues of H^T·H (avoiding a full SVD on-device), handles reflections via det(H) < 0, and properly guards against numerical edge cases. The host plumbing (async allocation, pinned-memory H2D transfer, stream threading) follows established nvMolKit patterns.

Key findings:

• Signed integer overflow in batch binding (nvmolkit/conformerRmsd.cpp:67): nc * (nc - 1) / 2 uses int arithmetic; for nc = 65536 the intermediate product nc * (nc - 1) = 4,294,901,760 exceeds INT_MAX, causing undefined behavior and producing a corrupt PyArray shape. The single-mol path on line 88 of the same file already applies the correct int64_t cast — the batch path should match.
• cccl_interface not linked in src/CMakeLists.txt for the new conformer_rmsd target, unlike every other CUDA kernel target in the repo that uses CUB / cuda::std::span. This may work via transitive dependencies through device_vector today but is inconsistent and fragile.
• Fixed block size of 128 is suboptimal for small molecules (< 128 heavy atoms, which is the majority of drug-like inputs). Most benchmarked molecules have 13–48 atoms, so many threads are idle and all 11 BlockReduce calls still incur full synchronization overhead regardless.

Confidence Score: 3/5

• Safe to merge after fixing the integer overflow in the batch binding; the overflow is unreachable for typical drug-discovery workloads but is undefined behavior and would produce wrong results at ~65k conformers.
• The CUDA math is correct and well-tested against an independent numpy Kabsch reference. The single-mol Python API is clean. However, the signed integer overflow in the batch binding path is a real correctness bug (UB in C++, wrong PyArray shape at the boundary case) that should be fixed before merging. The missing cccl_interface link is a build-consistency issue.
• nvmolkit/conformerRmsd.cpp (integer overflow at line 67) and src/CMakeLists.txt (missing cccl_interface for conformer_rmsd target).

Important Files Changed

Filename	Overview
nvmolkit/conformerRmsd.cpp	Python/C++ binding layer — contains a signed integer overflow (nc * (nc - 1) / 2 in int) at line 67 for large conformer counts (nc ≈ 65 536) that corrupts the PyArray shape passed back to Python; line 88 of the same file correctly uses int64_t for the equivalent single-mol path.
src/conformer_rmsd.cu	Core CUDA kernel implementing per-pair Kabsch RMSD via Cardano eigenvalue decomposition of H^T·H; mathematics is sound, reflection handling via det(H) < 0 is correct, numerical guards (fmax, clamping) are appropriate, though fixed block size of 128 is suboptimal for small molecules.
src/conformer_rmsd_mol.cpp	Host-side mol-to-GPU pipeline; properly uses int64_t for overflow detection, initiates async device allocation before CPU coordinate packing, and handles edge cases (null mols, zero atoms, < 2 conformers) cleanly.
nvmolkit/tests/test_conformer_rmsd.py	12 tests covering correctness vs numpy Kabsch, prealigned mode, edge cases (zero atoms, < 2 conformers, invalid types), streams, and both single/batch APIs; good coverage.
src/CMakeLists.txt	New conformer_rmsd and conformer_rmsd_mol targets are correctly structured, but cccl_interface (providing CUB / cuda::std::span) is not explicitly linked unlike every other CUDA kernel target in the repo that uses CUB.

Comments Outside Diff (3)

1. nvmolkit/conformerRmsd.cpp, line 67 (link)

   Integer overflow in numPairs computation

   nc * (nc - 1) is evaluated as a 32-bit signed int multiplication. For nc = 65536, the product is 4,294,901,760, which exceeds INT_MAX (2,147,483,647), causing signed integer overflow — undefined behavior in C++. On two's complement hardware this typically wraps to -65,536, so numPairs becomes -32,768, corrupting the shape passed to makePyArray.

   Note that the upstream conformerRmsdBatchMatrixMol does validate that the final pair count fits in int (using int64_t arithmetic in conformer_rmsd_mol.cpp), but it allows nc = 65536 through because numPairs = 2,147,450,880 ≤ INT_MAX. The overflow therefore occurs silently in the binding code.

   Line 88 in the same file already applies the correct fix with static_cast<int64_t>. Apply the same pattern here:

2. src/CMakeLists.txt, line 132-134 (link)

   Missing cccl_interface dependency

   conformer_rmsd.cu includes <cub/cub.cuh> and <cuda/std/span> from CCCL. Every other target in this repo that uses CUB explicitly links cccl_interface (e.g. similarity_kernels, morganFingerprintKernels), but conformer_rmsd does not:

   add_library(conformer_rmsd conformer_rmsd.cu)
   target_link_libraries(conformer_rmsd PUBLIC device_vector PRIVATE cuda_error_check)

   If device_vector does not transitively expose cccl_interface, builds may fail or silently pick up a different CCCL version from the system CUDA install. Consider explicitly adding cccl_interface consistent with the rest of the codebase:

3. src/conformer_rmsd.cu, line 84-85 (link)

   Performance degrades severely for very small molecules

   kRmsdBlockSize = 128 means every conformer pair dispatches exactly 128 threads regardless of atom count. For molecules with fewer than 128 heavy atoms (most drug-like molecules: aspirin has 13, celecoxib 26), most threads are idle on every loop iteration. All 11 cub::BlockReduce calls are still executed, and each requires a __syncthreads(), meaning the inter-thread synchronization cost is disproportionate to the actual work.

   The benchmark results confirm this pattern (Celecoxib 26 atoms, 200 conformers: 3.59 ms vs Aspirin 13 atoms, 200 conformers: 0.99 ms — 3.6× more pairs but 3.6× slower despite having only 2× more atoms, suggesting overhead dominates).

   A smaller block size (e.g., 32) for small molecules would improve occupancy and reduce the per-pair reduction overhead. Consider exposing kRmsdBlockSize as a template parameter or choosing it adaptively at the launch site based on atom count.

_{Last reviewed commit: "conformerRmsd: rever..."}

[volgin]

For context on the real-world impact: we're processing ~200M molecule states in our conformer generation pipeline (embed → optimize → RMSD cluster → select). The RMSD matrix is currently the only stage that falls back to CPU. At ~2s saved per state, this kernel alone would save over 12 years of cumulative CPU time across our workload — and keep the entire pipeline on GPU.

[evasnow1992]

Thank you for the contribution! This is a substantial addition — we'll review it thoroughly and follow up with any comments. We'll keep you updated as we work through it and determine the best approach for integration.

[scal444]

Thanks for the contribution! Overall this is a useful acceleration that aligns well with nvMolKit's functionality even though it wasn't initially on our rader. It's a very clean implementation and my code comments are minor.

One significant comment - I wonder if we're getting good GPU saturation. With many (100+ conformers) there's probably saturation but for smaller cases we typically need to send batches of molecules at a time, and calling singleApi(mol) repeatedly is less efficient that calling batchAPI([mol1, mol2, mol3]) multiple times. Adapting this to multiple molecules would require some minor preprocessing up front, and potentially an additional layer of indexing to match blocks to molecule ID and then to conformer ID. The output would be a list of results tensors.

Would your use cases match the multiple molecule multiple conformer case?

[volgin]

Addressed the latest round of feedback:

Batch API (scal444) — Added GetConformerRMSMatrixBatch for processing multiple molecules in a single kernel launch. Follows existing nvMolKit batch patterns: flat coordinate buffer with per-molecule offset arrays, one block per conformer pair across all molecules, array of per-molecule device output pointers. Python wrapper validates inputs, accepts an optional CUDA stream, and returns list[AsyncGpuResult]. Tests cover: batch vs. single-molecule parity, mixed conformer counts, prealigned=True, empty input, None molecule, and explicit stream.

Dead disc variable (Greptile) — Removed. The variable was computed but never used after the else fallback branch was deleted in the previous round. Near-degenerate eigenvalue inputs are handled entirely by the existing fmax guards on sqrtPP and the acos argument.

Coordinate offset overflow (Greptile) — Promoted coordOffsets from int to size_t end-to-end: PinnedHostVector, the host packing index (base), AsyncDeviceVector, the kernel parameter, and the header declaration. Also fixed a narrowing bug in the host packing loop where the index expression coordOffsetsArr[m] + confIdx * na * 3 + a * 3 was being computed in int despite coordOffsetsArr[m] being size_t — the intermediate products are now widened via static_cast<size_t> before addition.

O(N) linear scan (Greptile) — Replaced with an O(log N) binary search in conformerRmsdBatchKernel.

Cumulative pairOffsets overflow (scal444) — pairOffsets and totalPairs are intentionally 32-bit: this kernel launches one block per pair on the grid x-dimension, whose hardware limit is 2^31 - 1 (INT_MAX). The int64_t running sum exists solely to detect prefix-sum overflow before materializing into that launch/index type — an unchecked overflow would silently route blocks to the wrong molecule or write out of bounds. Added a clarifying comment in the source.

[volgin]

Fixed three issues from the latest Greptile pass: removed a premature #undef BLOCK_REDUCE_SUM inside the prealigned branch of conformerRmsdBatchKernel (the preprocessor processes it unconditionally, breaking compilation of the non-prealigned path); replaced nullptr output pointers for zero-pair molecules with a guaranteed valid device allocation; added precision comments at both inverse triangular-number formula sites.

[volgin]

The Kabsch computation duplication between the two kernels is noted. We're happy to extract a __device__ helper in this PR if you'd prefer to keep it together, but can also defer it to a follow-up PR to keep this review focused. Let us know which you'd prefer.

[volgin]

Went ahead with the refactoring rather than deferring — the benefits are clear enough to include here: computePairRmsd() is now a single __device__ __forceinline__ helper that contains the full RMSD computation path (prealigned fast-path, centroid reduction, cross-covariance accumulation, Cardano eigenvalue decomposition, reflection correction). Both kernels resolve their coordinate pointers and output slot, then delegate to it — reducing each kernel body from ~80 lines to ~15. BLOCK_REDUCE_SUM is defined once at file scope and #undef'd after the helper, removing the duplicate macro definition. Also added the if (numMols <= 0) return; guard at the top of conformerRmsdBatchKernel.

[scal444]

Thanks Andrei. Please see my comments from a few days ago. The only significant ones are moving business logic outside of the pybind layer where possible, and using cub::blockreduce as it's the nvidia standard for in-kernel sum reductions. The others are minor and optional, overall the change looks great.

[scal444]

Thanks for the contribution Andrei!