DOCSEARCHbyAlgolia
Submit Docs

NVIDIA ALCHEMI Toolkit-Ops

repository·main·Indexed 42 days ago

https://github.com/nvidia/nvalchemi-toolkit-ops

High-performance, GPU-optimized, and batched primitives for computational chemistry and atomistic simulation workflows. Built with NVIDIA Warp, it provides scalable kernels for neighbor lists, molecular dynamics (NVE, NVT, NPT, NPH), geometry optimization (FIRE, FIRE2), and interatomic interactions including DFT-D3(BJ) dispersion and PME electrostatics. Includes native bindings for PyTorch and JAX.

Tokens
78.5K
Snippets
146
Records
276
Agent score
51%

What's inside NVIDIA ALCHEMI Toolkit-Ops

  1. Overview of NVIDIA ALCHEMI Toolkit-Ops

    main

    NVIDIA ALCHEMI Toolkit-Ops is a collection of GPU-optimized, batched primitives designed to accelerate atomistic simulations. The core compute kernels are written in NVIDIA warp-lang and are intended to be highly scalable (supporting >100,000 atoms) and optimized for high throughput on GPUs.

    Key capabilities include:

    • Neighbor lists: Naive, cell-list, and tiled cluster-pair methods with support for Matrix and COO formats.
    • Molecular dynamics: Support for NVE, NVT, NPT, and NPH ensembles using various thermostats (Langevin, Nosé-Hoover Chain, velocity-rescaling).
    • Geometry optimization: FIRE and FIRE2 methods supporting coordinate and lattice relaxation.
    • Interatomic interactions: DFT-D3(BJ) dispersion, DSF, Ewald, and PME electrostatics.
    • Differentiable electrostatics: Integration with PyTorch for automatic differentiation of forces, charge gradients, virials, and stress.
    • Framework Bindings: Native support for PyTorch and JAX.

  2. Overview of Neighbor List Benchmarks

    main

    The benchmarks/neighborlist suite provides performance evaluations for various neighbor list construction algorithms. These benchmarks are used to measure how different algorithms scale with system size (number of particles) and batch dimensions.

    Supported algorithms and scaling patterns include:

    • Cell list algorithm: Optimized for $O(N)$ scaling.
    • Naive algorithm: Standard $O(N^2)$ scaling.
    • Batch processing variants: Evaluates performance when processing multiple systems simultaneously.

  3. Overview of ALCHEMI Toolkit-Ops

    main

    ALCHEMI Toolkit-Ops provides GPU-accelerated algorithms for atomistic simulations, computational chemistry, and graph neural networks. Built on NVIDIA Warp, it delivers high-performance primitives for neighbor list construction, dispersion corrections, and related operations on both single systems and batched datasets.

    Key Features:

    • Density Functional Theory: DFT-D3(BJ) dispersion corrections and high-throughput batch calculations.
    • Graph Neural Networks: Generation of edge connectivity via neighbor lists in COO format.
    • Molecular Dynamics: Neighbor list construction with skin distance optimization and rebuild detection.
    • High-Throughput Screening: Batch processing of large molecular databases.
    • Method Development: Access to low-level Warp kernels for custom algorithm development.

  4. Overview of Neighbor List capabilities

    main

    The nvalchemi-toolkit-ops neighbor list functionality provides efficient $O(N)$ neighbor list construction using GPU-accelerated cell list algorithms. It is designed for high-performance computational chemistry workflows and supports several advanced features:

    • Flexible Output Formats: Supports both dense (padded matrix) and sparse (COO/edge_index) output formats.
    • Batching: Capability to build neighbor lists for single systems or batches of systems.
    • Differentiable Geometry: Ability to return per-neighbor vectors and distances, enabling differentiable physics simulations.
    • Integrated Physics: Compute inline Lennard-Jones pair energies and forces using the pair_fn parameter.
    • Optimization: Supports torch.compile for performance optimization and provides Torch cost-model dispatch estimates.
    • Workflow Integration: Includes mechanisms to detect when neighbor lists need rebuilding and integrates with molecular dynamics workflows.

  5. Overview of Electrostatics capabilities

    main

    The Electrostatics module provides GPU-accelerated computation of long-range electrostatic interactions in periodic systems. It supports multiple methods including direct Coulomb (damped and undamped), Ewald summation with automatic parameter estimation, and Particle Mesh Ewald (PME) for $O(N \log N)$ scaling.

    Key capabilities include:

    • Interfacial Systems: Applying two-dimensional slab corrections for Ewald and PME.
    • Advanced Multipoles: Computing multipole Ewald and PME totals using charges, dipoles, and quadrupoles ($l_{max} = 0, 1, 2$), including stress tensors via cell gradients.
    • Feature Extraction: Extracting atom-centered multipole features by projecting periodic potential onto receiver GTOs.
    • Optimization & Training:
      • Leveraging autograd for forces and gradients.
      • Using the energy-derivative contract for training on forces, stress, and charge gradients (replaces deprecated direct-output flags).
      • Amortizing reciprocal-space state using the multipole SCF cache.
      • Using second-order backward Warp kernels for force-loss-style training.

  6. Overview of Molecular Dynamics Integrators

    main

    ALCHEMI Toolkit-Ops provides GPU-accelerated MD integrators and thermostats built on NVIDIA Warp, featuring full PyTorch autograd support. Integrators are categorized by the statistical ensemble they simulate:

    IntegratorEnsembleConservationBest For
    Velocity VerletNVEEnergyTesting, production NVE runs
    Langevin (BAOAB)NVTTemperatureCanonical sampling, equilibration
    Nosé-Hoover ChainNVTTemperatureDeterministic thermostat, long runs
    NPT (MTK)NPTTemperature, PressureConstant pressure simulations
    NPH (MTK)NPHEnthalpy, PressureAdiabatic constant pressure
    Velocity Rescaling--Quick equilibration (non-canonical)

    All integrators support:

    • Single-system and batched calculations (via batch_idx or atom_ptr).
    • Automatic differentiation (positions, velocities, forces).
    • Both mutating (in-place) and non-mutating (out) variants.
    • Float32 and float64 precision.

  7. Use NVIDIA Warp kernels for electrostatic interactions

    main

    The nvalchemiops.interactions.electrostatics module provides framework-agnostic NVIDIA Warp kernels for calculating electrostatic interactions. These kernels operate directly on warp.array objects, making them suitable for building custom integrators or high-performance bindings.

    Note: If you require PyTorch compatibility (accepting torch.Tensor inputs and supporting autograd), use the PyTorch-specific electrostatic module instead of these Warp kernels.

  8. Explore the NVIDIA ALCHEMI Toolkit-Ops API Reference

    main

    The NVIDIA ALCHEMI Toolkit-Ops package provides high-performance NVIDIA Warp primitives for computational chemistry. The API is organized into three primary backend references depending on your preferred compute framework:

    • PyTorch Reference: Includes modules for neighbors, segment operations, dispersion, electrostatics, dynamics, and autograd.
    • JAX Reference: Includes modules for neighbors, segment operations, dispersion, and electrostatics.
    • Warp Reference: Includes modules for neighbors, segment operations, dispersion, electrostatics, dynamics, and utilities.

    For detailed usage instructions and implementation patterns, refer to the full user guide and examples, as the API reference provides technical signatures but may lack contextual guidance.

  9. Understand ALCHEMI Toolkit-Ops ecosystem integration

    main

    ALCHEMI Toolkit-Ops is designed to work seamlessly within the scientific Python ecosystem and NVIDIA hardware stack:

    • PyTorch: Supports automatic CPU/GPU handling and torch.compile compatibility.
    • JAX: Provides a functional API using jax.Array that mirrors the PyTorch interface.
    • NVIDIA Warp: Kernels are implemented in NVIDIA Warp, utilizing GPU memory layouts optimized for NVIDIA architectures. This allows for standard profiling using tools like Nsight Systems.
    • Computational Chemistry: DFT-D3 dispersion correction parameters follow the standard format from the Grimme group, ensuring compatibility with established DFT workflows.

  10. Use segmented operations in PyTorch

    main

    The nvalchemiops.torch.segment_ops module provides differentiable PyTorch bindings for GPU-accelerated segmented reductions and per-segment algebra. These operations allow you to perform computations (like sums, means, or matrix-vector products) on groups of data defined by segment indices, which is highly efficient for computational chemistry workloads on GPUs.

    For advanced usage guidance and patterns for CUDA graph capture, refer to the segment_ops_userguide.

  11. Use segmented operations in JAX with nvalchemiops.jax.segment_ops

    main

    The nvalchemiops.jax.segment_ops module provides differentiable JAX bindings designed for GPU-accelerated segmented reductions and per-segment algebra. These operations allow you to perform computations (like sums, means, or matrix-vector products) on groups of data defined by segment indices, which is highly efficient for computational chemistry workloads on GPUs.

    Available public operations include:

    • segmented_sum: Segmented summation.
    • segmented_dot: Segmented dot product.
    • segmented_mul: Segmented multiplication.
    • segmented_mean: Segmented mean.
    • segmented_rms_norm: Segmented Root Mean Square (RMS) normalization.
    • segmented_matvec: Segmented matrix-vector multiplication.

  12. Understand Neighbor List performance benchmarks

    main

    The Neighbor List benchmarks evaluate the performance of different neighbor list algorithms across two primary backends: Torch and JAX. Benchmarks are categorized by chemical systems (e.g., CsCl, NH₃) and analyzed through three scaling lenses:

    1. System Size Scaling: Measures mean execution time, throughput ($10^6$ atoms/s), and peak GPU memory as the number of atoms in a system increases.
    2. Constant Workload: Evaluates execution time, throughput, and memory near a configured total-atom target while varying the batch size.
    3. Batch Scaling: Measures execution time, throughput, and peak GPU memory relative to the batch size (with fixed atoms per system).

    Note on Memory Reporting:

    • Torch: Reports memory using the Torch allocator.
    • JAX: Memory plots are omitted. Because XLA's allocator pool and buffer reuse make per-call allocation deltas unreliable, the suite does not measure JAX memory. Torch memory reports are not a proxy for JAX memory due to differences in framework wrappers and allocators.