# Unlocking Tensor Core Performance for Machine Learning Interatomic Potentials > Source: > Published: 2026-10-29 15:00:00+00:00 Atomistic simulations help researchers understand how molecules and materials behave, from chemical reactions and phase transitions to mechanical failure. These simulations come with a cost-accuracy trade-off. Quantum-mechanical methods such as density functional theory (DFT) can provide accurate energies and forces, but their computational cost often limits their use for large systems or long simulations. Classical interatomic potentials are much faster, but their fixed functional forms limit accuracy and generalizability across different atomic environments. Machine learning interatomic potentials (MLIPs) help bridge this gap. They use neural networks to predict potential energies and interatomic forces at a fraction of the cost, enabling researchers to study larger systems and longer timescales while retaining much of the required accuracy. Running an MLIP efficiently, however, presents its own challenge. Atomistic simulations are sensitive to small errors that can accumulate over time and are most visible when computing forces, which are obtained by differentiating the predicted energy. During molecular dynamics simulations, these forces guide updates to atomic velocities and positions over many timesteps, making numerical precision important not only for physical accuracy but also for the stability of the resulting trajectory. For this reason, many MLIPs continue to run inference in FP32, even though modern NVIDIA GPUs offer substantially higher matrix-multiplication throughput through low-precision Tensor Core operations. Tensor Core methods for FP32 matrix multiplication offer a way to narrow that gap. They decompose each FP32 operand into lower-precision components, evaluate their cross-products on Tensor Cores, and recombine the results in FP32. This approach can increase matrix-multiplication throughput while retaining accuracy close to native FP32, without retraining or changing the model’s weights or architecture. In this post, we evaluate how Tensor Core FP32 methods perform in practice on the Point Edge Transformer (PET), a transformer-based MLIP developed by Lab COSMO at EPFL. PET spends much of its runtime in highly optimized general matrix multiplications (GEMMs), leaving limited headroom for further optimization and making it a demanding test case of whether these methods can deliver additional performance gains. We evaluate BF16x9 on NVIDIA Blackwell GPUs and TF32x3 on NVIDIA Ampere GPUs, comparing kernel and end-to-end molecular dynamics performance and trajectory-level behavior. We further test the TF32x3 method for potential-energy smoothness, energy conservation, and structural relaxation. We show how these methods can accelerate MLIP inference while retaining behavior close to the FP32 baseline and outline how to evaluate them on other precision-sensitive scientific machine learning workloads. ## How Tensor Core FP32 closes the hardware gap Modern NVIDIA GPUs deliver substantially higher Tensor Core throughput for BF16 and TF32 than for FP32 matrix multiplication, creating a gap between the available throughput at lower precision and the precision requirements of many scientific workloads. Tensor Core FP32 methods help narrow this gap by decomposing FP32 operands into lower-precision components, evaluating selected products on Tensor Cores, and recombining them with FP32 accumulation. Depending on the method and workload, this can deliver accuracy at or near native FP32 while taking advantage of higher Tensor Core throughput. On Blackwell GPUs, NVIDIA cuBLAS provides a [BF16x9 Tensor Core path](https://arxiv.org/abs/2605.16617) (Figure 1). Each FP32 operand is decomposed into three BF16 components, and nine pairwise products are recombined with the required scaling. cuBLAS can enable this path through the `CUBLAS_EMULATE_SINGLE_PRECISION=1` environment variable, so compatible PyTorch GEMMs can use it without changing the model architecture or implementing a custom GEMM. Because this approach changes how GEMMs are evaluated, it may provide additional gains when combined with optimized graph construction, fused operations, custom kernels, and NVIDIA CUDA graphs. We did not measure the combined gain in this evaluation; it should be established end to end for each workload. The [TF32x3](https://arxiv.org/abs/2203.03341) [method](https://arxiv.org/abs/2203.03341) applies the same principle on Ampere GPUs (Figure 2). TF32x3 represents each FP32 operand using a leading TF32 component and a residual TF32 component. Unlike BF16x9, our TF32x3 implementation does not use scaling factors. It evaluates three cross-products and omits the residual-by-residual term, as shown in the following equation NVIDIA CUTLASS provides this approach through the `OpMultiplyAddFastF32` operator. Inputs and outputs remain in FP32. Tensor Cores execute the three TF32 products with accumulation in FP32. ## Benchmarking Tensor Core FP32 matrix math on PET We evaluated the 26M-parameter PET-MAD small checkpoint (`pet-mad-s-v1.6.0-nn-32`) on a 192-atom water box and a 17,280-atom supercell. We also evaluated the 255M-parameter large checkpoint (`pet-mad-m-v1.6.0`) on the 192-atom system. We compared predicted energies, forces, and stresses with those from the native FP32 path. ### BF16x9 method on NVIDIA Blackwell GPUs On NVIDIA Blackwell GPUs with compute capability 10.0 or 10.3, cuBLAS supports enabling BF16x9 through an environment variable, math mode, or compute type (requested with `CUBLAS_COMPUTE_32F_EMULATED_16BFX9`, `CUBLAS_FP32_EMULATED_BF16X9_MATH`, or `CUBLAS_EMULATE_SINGLE_PRECISION=1`). No model-code changes are required to use this path for single-precision general matrix multiplication (SGEMM). With BF16x9 on Blackwell, large GEMMs in PET were accelerated by 2.2–2.5x for the 17,280-atom workload (Figure 3). Figure 4 compares MD-step speedup and single-point energy and force accuracy for cuBLAS’s performant and eager strategies against compiled FP32. - **192-atom system:** The performant strategy sped up the MD step by 1.2x, while the eager strategy achieved 1.1x. - **17,280-atom system:** The performant strategy achieved a 1.7x speedup, while the eager strategy achieved 1.4x. The performant strategy uses BF16x9 when cuBLAS expects it to be faster; the eager strategy uses it whenever the operation supports it. Importantly, the error due to the BF16x9 Tensor Core method is below the model’s own level of non-determinism. This result is significant because PET already expresses much of its runtime as optimized GEMM work, and yet the BF16x9 path exposed performance headroom without retraining the model, changing its transformer architecture, or writing custom kernels. ### TF32x3 method on Ampere GPUs On Ampere GPUs, we implemented TF32x3 by replacing PET’s standard `nn.Linear` layers with a custom `TF32Linear` layer that calls a CUTLASS C++ TF32x3 kernel. CUTLASS C++ provides a reference implementation of this approach in [Example 27](https://github.com/NVIDIA/cutlass/blob/main/examples/27_ampere_3xtf32_fast_accurate_tensorop_gemm/27_ampere_3xtf32_fast_accurate_tensorop_gemm.cu). The key setting in the PET implementation is: C++ ``` cutlass::arch::OpMultiplyAddFastF32 ``` This selects CUTLASS’s fast, accurate FP32 multiplication mode. The kernel represents each FP32 operand as leading and residual TF32 components. Tensor Cores then evaluate three products—leading-by-leading and the two leading-by-residual products—and accumulate them in FP32. The smallest residual-by-residual product is omitted. The CUTLASS C++ operator is declared with `float` inputs and outputs. PET therefore continues to supply ordinary FP32 activations and weights and receives an FP32 result; from the rest of the model’s perspective, the layer behaves like a standard `nn.Linear`. Only the implementation of the underlying matrix multiplication changes. The C++ example is written for `sm80` (Ampere) and can also be compiled for `sm90` (Hopper). The custom `TF32Linear` module replaces compatible PyTorch linear layers and performs their matrix multiplications using the CUTLASS C++ TF32x3 kernel while preserving the trained weights, bias, and input shape. Because PET calculates forces by differentiating predicted energies, TF32x3 is also used to propagate gradients to the layer inputs. This ensures that both energy and force calculations use TF32x3 Tensor Core arithmetic. Weight gradients use the standard PyTorch path, but they are not required during fixed-weight inference. Across PET’s major linear-layer shapes on A100, TF32x3 accelerated GEMMs by approximately 1.7–1.9x (Figure 5). Figure 6 compares MD-step speedup and single-point accuracy for one-pass TF32 and TF32x3, each measured against its corresponding FP32 baseline. One-pass TF32 achieved higher speedups, but at a greater cost to accuracy and trajectory stability (Figures 6–7). - **Small checkpoint (26M parameters):** Results are shown for a 192-atom water box and its 17,280-atom supercell. TF32x3 kept the force error to roughly three times the difference between FP32 runs after recompilation and sped up the MD step by 1.5x on the larger system. - **Large checkpoint (255M parameters):** On the 192-atom system, TF32x3 achieved a 1.4x speedup, compared with 1.2x for the small checkpoint. Its force error was 3.9 × 10⁻⁶ eV/Å, essentially unchanged from the small checkpoint. ## Validating physical consistency For molecular simulation, matching isolated energy and force predictions is only the first test of numerical fidelity. Forces are gradients of the learned potential-energy surface and are integrated over many timesteps, so even small numerical irregularities can roughen the surface, destabilize trajectories, or produce unphysical energy drift. The promise of the Tensor Core FP32 method is therefore especially compelling: it can unlock Tensor Core performance for matrix multiplications while preserving the smooth and stable behavior expected from FP32. We tested potential-energy smoothness and energy conservation using two separate tests (Figure 7). To test energy smoothness, we dragged a single CO molecule rigidly away from an on-top site of a Pt(111) slab and evaluated the total energy at 60 carbon heights from 1.2 to 6.0 Å in steps of 8.1 × 10⁻² Å. TF32x3 followed the smooth FP32 CO/Pt(111) energy curve, while one-pass TF32 introduced visible point-to-point variations. In a separate NVE simulation of a 192-atom water box, FP32 and TF32x3 showed comparable bounded energy fluctuations, while one-pass TF32 showed a larger systematic drift. These results show that TF32x3 preserved potential-energy smoothness and maintained energy conservation in the tested systems, not just single-point accuracy. ## Testing structural relaxation Structural relaxation also tests numerical fidelity as the optimizer repeatedly uses predicted forces to update the atomic geometry. Small force inconsistencies can alter the optimization path, increase the number of steps required, or prevent the system from reaching the desired force tolerance. We relaxed a 539-atom [Cantor-alloy](https://www.nature.com/articles/s41467-025-65662-7) slab using PET and the FIRE2 optimizer with the NVIDIA ALCHEMI Toolkit. Across three replicates, FP32 converged in 218 steps and TF32x3 in 218–222 steps; every run reached the 1.0 × 10⁻³ eV/Å force tolerance. With one-pass TF32, all three replicates failed to converge within the 500-step limit, and forces plateaued around 1.2 × 10⁻³ eV/Å. TF32x3 accelerated each optimization step by 1.5x while matching the FP32 convergence rate and criterion (Figure 8). One-pass TF32 achieved a 2.0x per-step speedup but failed to converge even after doubling the optimization step limit to 1,000 steps. These results highlight the practical value of the Tensor Core path: substantial acceleration without sacrificing the convergence behavior of the FP32 workflow. ## Limitations and decision criteria End-to-end gains depend on system size and the share of runtime spent in GEMMs. Small systems generally benefit less because launch overhead, neighbor processing, graph operations, and other non-GEMM work make up a larger portion of the runtime. Larger systems benefit more when GEMMs account for more of the work. TF32x3’s performance also depends on kernel tuning for the model’s matrix shapes and hardware. Numerical behavior depends on the Tensor Core strategy used. In matrix-multiplication tests, [BF16x9 matched or exceeded native FP32 accuracy](https://arxiv.org/abs/2605.16617) against an FP64 reference, supporting its use as a drop-in replacement for FP32 GEMMs on supported hardware. TF32x3 has a slightly different accuracy profile and should be validated against the model’s accuracy and physical-stability requirements. ## Conclusion Tensor Core FP32 methods improved PET performance without retraining or changing the model architecture. TF32x3 accelerated major GEMMs on A100 by up to 1.9x and the complete molecular dynamics step by up to 1.5x. BF16x9 accelerated major GEMMs on an NVIDIA B300 Tensor Core GPU by up to 2.5x and the complete step by up to 1.7x. As demonstrated for TF32x3, Tensor Core methods preserve behavior close to FP32 in the tested energy-conservation and structural-relaxation workflows. These results support evaluating the methods in other MLIPs and precision-sensitive scientific machine learning models whose runtime is dominated by large FP32 GEMMs. Neither of the Tensor Core FP32 strategies requires retraining or an architectural redesign. Each can therefore be evaluated as an independent performance layer and combined with other software optimizations when the workload supports it. Try TF32x3 or BF16x9 Tensor Core methods with your model today—no retraining or architectural changes needed. Learn more in the [cuBLAS documentation](https://docs.nvidia.com/cuda/cublas/) and the [NVIDIA ALCHEMI Toolkit documentation](https://developer.nvidia.com/cuda/cuda-x-libraries/alchemi). Explore the PET models in the [UPET repository](https://github.com/lab-cosmo/upet). ## Acknowledgments We’d like to thank Nikita Fedik, Harun Bayraktar, Gregory Thomsen, and Vamsi Sripathi for their contributions to this post.