NVIDIA Tensor Core FP32 methods speed up Point Edge Transformer MLIP inference

Why precision matters in atomistic simulation

Atomistic simulations help researchers understand how molecules and materials behave, from chemical reactions and phase transitions to mechanical failure. Quantum-mechanical methods such as density functional theory (DFT) give accurate energies and forces, but their cost often limits them to small systems or short simulations. Classical interatomic potentials are much faster, but their fixed functional forms limit accuracy and generalizability.

Machine learning interatomic potentials (MLIPs) sit between the two. They use neural networks to predict potential energies and interatomic forces at a fraction of the cost. Running them efficiently is still hard, according to NVIDIA. Forces are obtained by differentiating the predicted energy, and during molecular dynamics they guide updates to atomic velocities and positions over many timesteps. Small errors can accumulate, so numerical precision matters for both physical accuracy and trajectory stability.

For this reason, many MLIPs still run inference in FP32, even though modern NVIDIA GPUs offer much higher matrix-multiplication throughput through low-precision Tensor Core operations.

How the FP32 emulation methods work

The approach NVIDIA examines decomposes each FP32 operand into lower-precision components, evaluates their cross-products on Tensor Cores and recombines the results with FP32 accumulation. According to NVIDIA, this can raise throughput while keeping accuracy close to native FP32, without changing the model's weights or architecture.

On Blackwell GPUs, NVIDIA cuBLAS offers a BF16x9 path. Each FP32 operand is split into three BF16 components, and nine pairwise products are recombined with the required scaling. It can be enabled with the CUBLAS_EMULATE_SINGLE_PRECISION=1 environment variable, so compatible PyTorch GEMMs can use it without a custom GEMM. cuBLAS supports it on Blackwell GPUs with compute capability 10.0 or 10.3.

On Ampere GPUs, TF32x3 uses a leading TF32 component and a residual TF32 component per operand. NVIDIA's implementation uses no scaling factors, evaluates three cross-products and omits the residual-by-residual term. NVIDIA CUTLASS provides it through the OpMultiplyAddFastF32 operator, with inputs and outputs remaining in FP32.

Results on the Point Edge Transformer

PET is a transformer-based MLIP developed by Lab COSMO at EPFL. NVIDIA chose it as a demanding test because it already spends much of its runtime in highly optimized GEMMs, leaving limited headroom. The tests used the 26M-parameter PET-MAD small checkpoint on a 192-atom water box and a 17,280-atom supercell, plus the 255M-parameter large checkpoint on the 192-atom system.

• BF16x9 on Blackwell: large GEMMs ran 2.2, 2.5x faster for the 17,280-atom workload and 1.2, 1.4x faster for the 192-atom workload.
• MD step, 192 atoms: the performant strategy gave 1.2x and the eager strategy 1.1x, relative to compiled FP32.
• MD step, 17,280 atoms: the performant strategy gave 1.7x and the eager strategy 1.4x.
• TF32x3 on Ampere: major GEMMs sped up by 1.7, 1.9x and the full molecular dynamics step by up to 1.5x, with force error near the FP32 baseline.

The performant strategy uses BF16x9 when cuBLAS expects it to be faster, while the eager strategy uses it whenever the operation supports it. NVIDIA says the error from BF16x9 is below the model's own level of non-determinism.

Energy conservation and structural relaxation

NVIDIA also tested TF32x3 for potential-energy smoothness, energy conservation and structural relaxation. In NVE trajectories, TF32x3 matched FP32 behavior, while one-pass TF32 showed drift and roughness.

In relaxation of a 539-atom Cantor-alloy slab, TF32x3 converged in 218, 222 steps with a 1.5x per-step speedup. One-pass TF32 failed to reach the force tolerance even after 1,000 steps.

What it means for users

The results suggest that MLIP users on Blackwell or Ampere GPUs may gain speed without retraining, changing architecture or writing custom kernels. NVIDIA notes that the methods may add further gains when combined with optimized graph construction, fused operations, custom kernels and CUDA graphs, but it did not measure that combined gain and says it should be established end to end for each workload.

For next steps, NVIDIA points to the cuBLAS documentation for BF16x9, the NVIDIA ALCHEMI Toolkit documentation for TF32x3 integration examples, and the UPET repository for PET model checkpoints.