Unlocking Tensor Core Performance for Machine Learning Interatomic Potentials

Machine learning interatomic potentials let researchers simulate molecules and materials far faster than quantum methods like density functional theory. But many still run in full FP32 precision, because small numerical errors pile up over thousands of simulation steps and can wreck a trajectory. That leaves a lot of GPU throughput unused.

NVIDIA and collaborators tested two methods that emulate FP32 math on Tensor Cores on PET, a transformer-based potential from EPFL's Lab COSMO. On Blackwell, the BF16x9 path sped up the full molecular dynamics step by up to 1.7x. On an A100, TF32x3 sped it up by up to 1.5x. Neither required retraining or changing the model.

This is the kind of unglamorous engineering we want to see more of. Faster materials and chemistry simulation means more candidate catalysts, alloys, and molecules screened per dollar and per hour of compute. A free 1.5x on existing models, enabled in the Blackwell case by an environment variable, is a real gift to labs with fixed budgets.

What makes this post credible is that it does not stop at kernel speedups. The team checked whether the physics still behaved. TF32x3 tracked the smooth FP32 energy curve as a CO molecule was pulled off a platinum surface, and it held energy conservation in a 30-picosecond simulation of water. In a 539-atom alloy relaxation, it converged in about the same number of steps as FP32.

The cautionary tale is the cheaper option. Plain one-pass TF32 was faster, about 2x per step, but it introduced roughness, energy drift, and failed to converge even after the step limit was doubled to 1,000. That is the general lesson I'd underline: a faster number that quietly breaks the science is worse than no speedup, because it produces results that look fine and aren't.

The gains are also not universal. Small systems benefit less because non-matrix work dominates, and TF32x3 needs tuning for each model's matrix shapes. BF16x9 matched or beat native FP32 accuracy in matrix tests, while TF32x3 has a slightly different accuracy profile and should be validated per workload. And yes, this is a vendor blog promoting its own libraries and hardware, so we would like to see other groups reproduce it on other potentials.

Our take: worth trying, with discipline. Labs should adopt these paths with the same physical checks NVIDIA ran here, not just a speed test.

Both methods split each FP32 number into lower-precision pieces, multiply those pieces on Tensor Cores, and recombine the result in FP32. BF16x9 is available through cuBLAS on supported Blackwell GPUs; TF32x3 was implemented on Ampere through a CUTLASS kernel wrapped as a drop-in replacement for PyTorch linear layers.