Nvmolkit Usage

nvMolKit usage

What nvMolKit is

GPU-accelerated, batched implementations of common RDKit operations. APIs mirror RDKit where possible but are batch-oriented: they take lists of rdkit.Chem.Mol (or lists of fingerprints) and process them in parallel on one or more GPUs. nvMolKit links against RDKit at build time; inputs and outputs are real RDKit Mol objects.

Where nvMolKit does well

Reach for nvMolKit when:

• The workload is a large batch of molecules processed together (typically thousands or more).
• The metric is throughput / total wall time across the batch, not per-molecule latency.
• The same operation is repeated identically across the batch (fingerprinting a library, embedding/minimizing many conformers, bulk pairwise similarity), so the GPU stays saturated.

Plain RDKit is usually the better choice for single-molecule one-offs or workflows that can't be expressed as a batch. nvMolKit is not meant to replace RDKit for those cases.

Runtime requirements

• An NVIDIA GPU with compute capability 7.0 (V100) or higher
• A CUDA driver compatible with CUDA 12.6+.
• A working torch install with CUDA support (nvMolKit returns GPU tensors via torch's CUDA array interface).

If CUDA is unavailable, nvMolKit calls raise. There is no CPU fallback - if the user needs one, use RDKit directly for that path.

When helping with installation, make the user choose a PyTorch CUDA backend that the host driver supports before installing nvMolKit. nvMolKit's PyPI wheels are built with CUDA Toolkit 12.9 and depend on CUDA 12 runtime packages, but pip/uv can still select a CUDA 13 PyTorch wheel unless the install command says otherwise.

• Conda: prefer conda-forge pytorch-gpu; pin cuda-version=12.6 or another CUDA version supported by the driver.
• pip: send the user to the PyTorch install selector (https://pytorch.org/get-started/locally/) or previous-versions page (https://pytorch.org/get-started/previous-versions/) to install torch for a CUDA 12.x backend before installing nvMolKit.
• uv: install nvMolKit with an explicit backend, e.g. uv pip install --torch-backend=cu128 nvmolkit.

Verify the install before writing real code

Run this once to confirm nvMolKit is importable and a GPU op works end to end:

import nvmolkit
import torch
from rdkit import Chem
from nvmolkit.fingerprints import MorganFingerprintGenerator

print("nvmolkit:", nvmolkit.__version__)
print("cuda available:", torch.cuda.is_available())
print("device count:", torch.cuda.device_count())

mols = [Chem.MolFromSmiles(smi) for smi in ["CCO", "c1ccccc1", "CC(=O)O"]]
fpgen = MorganFingerprintGenerator(radius=2, fpSize=1024)
result = fpgen.GetFingerprints(mols)
torch.cuda.synchronize()
fps = result.torch()
print("fps shape:", tuple(fps.shape), "dtype:", fps.dtype)
# Expected: shape (3, 32), dtype torch.int32  (1024 bits packed into 32 int32s per row)

If this fails, point the user at the install guide on the docs site rather than guessing - see "Going deeper" below.

Entry points

Result types and execution model

Two return shapes carry GPU-resident output, depending on what the operation produces.

AsyncGpuResult

Used by operations that return a single flat tensor (fingerprints, similarity matrices, RMSD/TFD vectors, Butina inputs). Key behaviors:

• Asynchronous. The kernel may not have completed when the call returns.
• result.torch() returns a zero-copy torch.Tensor on the GPU. Caller is responsible for synchronizing before reading values on the host.
• result.numpy() synchronizes and returns a CPU numpy array.
• Exposes __cuda_array_interface__, so it can be passed directly into other nvMolKit functions (e.g. fingerprints → similarity) with no host round-trip.

CUDA stream control

A subset of the AsyncGpuResult-returning APIs accept an optional stream: torch.cuda.Stream | None = None argument so callers can submit nvMolKit work to a non-default stream and overlap it with their own kernels. When omitted, the call uses the current torch stream.

APIs that take a stream argument:

• MorganFingerprintGenerator.GetFingerprints
• crossTanimotoSimilarity, crossCosineSimilarity, and their *MemoryConstrained variants
• butina, fused_butina
• GetConformerRMSMatrix, GetConformerRMSMatrixBatch

Other APIs (ETKDG, MMFF/UFF optimization, TFD, substructure search) are synchronous to the caller — no stream plumbing needed.

Typical pattern:

import torch
from rdkit import Chem
from nvmolkit.fingerprints import MorganFingerprintGenerator
from nvmolkit.similarity import crossTanimotoSimilarity

stream = torch.cuda.Stream()
fpgen = MorganFingerprintGenerator(radius=2, fpSize=1024)
mols = [Chem.MolFromSmiles(smi) for smi in ["CCO", "c1ccccc1", "CC(=O)O"]]

with torch.cuda.stream(stream):
    fps = fpgen.GetFingerprints(mols, stream=stream)
    sim = crossTanimotoSimilarity(fps, stream=stream)
stream.synchronize()
print(sim.torch())

Device3DResult

Used by ETKDG embedding and MMFF/UFF optimization (one-shot and BatchedForcefield) when called with output=CoordinateOutput.DEVICE. The GPU-resident equivalent of writing conformers back to Mol objects. Fields:

• values: AsyncGpuResult of shape (total_atoms, 3) float64. Concatenated conformer coordinates in CSR-style layout.
• atom_starts, mol_indices, conf_indices: AsyncGpuResult int32 buffers describing the layout (values[atom_starts[i]:atom_starts[i+1]] is conformer i's atoms).
• energies, converged: AsyncGpuResult buffers populated only for MMFF/UFF minimization (not for plain ETKDG).
• gpu_id: device the buffers live on. The targetGpu argument on each API picks this; targetGpu=-1 uses the default consolidation device.
• .per_molecule() returns nested list[list[torch.Tensor]] of per-conformer views; .dense(pad_value=nan) materializes a padded (n_mols, max_confs, max_atoms, 3) tensor.

The default mode (CoordinateOutput.RDKIT_CONFORMERS) still writes optimized coordinates back into each Mol and returns Python lists of energies/convergence flags. Reach for CoordinateOutput.DEVICE when chaining downstream GPU work (e.g. ETKDG → MMFF → similarity scoring) without host round-trips.

Configuration

Two configuration objects expose the GPU/CPU knobs.

HardwareOptions (ETKDG, MMFF, UFF)

from nvmolkit.types import HardwareOptions. Passed via hardwareOptions= to EmbedMolecules, MMFFOptimizeMoleculesConfs, UFFOptimizeMoleculesConfs, and the BatchedForcefield constructors. Every field has an "auto" sentinel; the defaults are usually fine.

Passing a gpuIds entry for a device that isn't visible raises RuntimeError: invalid device ordinal. For finding good values automatically across a representative sample, see nvmolkit.autotune (requires the optuna extra); each tune_* function returns a TuneResult whose best_config is a fully-populated HardwareOptions ready to pass back into the real call.

HardwareOptions round-trips through to_dict() / from_dict() for persisting tuned configs to disk.