Splitting FP32 matrix multiplies into low-precision pieces let a molecular-dynamics model run 1.7x faster on Blackwell and 1.5x on Ampere without retraining, and the one-pass shortcut still drifted
AspireAI community · NVIDIA Technical Blog
Machine-learning interatomic potentials are usually forced to run inference in FP32, because forces are gradients of the predicted energy and small numerical errors accumulate over a molecular-dynamics trajectory. That leaves most of a modern GPU idle: Tensor Cores deliver far more throughput for BF16 and TF32 than for FP32. NVIDIA's evaluation of the Point Edge Transformer from LabCOSMO at EPFL asks whether FP32 accuracy can be kept while using Tensor Cores anyway.
The answer is to decompose the FP32 operands. On Blackwell, cuBLAS's BF16x9 path splits each FP32 operand into three BF16 components and recombines nine pairwise products in FP32, which needs no model change at all - just an environment variable. It accelerated PET's large GEMMs by 2.2 to 2.5 times on a 17,280-atom supercell and sped up the molecular-dynamics step 1.7 times, with error below the model's own level of non-determinism. On Ampere, a custom CUTLASS TF32x3 layer reached roughly 1.7 to 1.9 times on GEMMs and 1.5 times on the full MD step for the larger system. The physical behaviour held up: TF32x3 followed the smooth FP32 energy curve on a CO/Pt(111) surface where one-pass TF32 produced visible point-to-point wobble, and in NVE simulation FP32 and TF32x3 showed comparable bounded energy fluctuations while one-pass TF32 drifted systematically.
The interesting result is what one-pass TF32 could not do. It was faster per step - 2.0 times on structural relaxation - but all three replicates of a 539-atom Cantor-alloy relaxation failed to converge within 500 steps with forces plateaued near 1.2e-3 eV/A, and even at 1,000 steps they did not converge, while FP32 and TF32x3 both converged in about 218 to 222 steps to a 1.0e-3 eV/A tolerance. The limits are explicit: this is one third-party model on specific hardware, the combined gain with graph construction and custom kernels was not measured, and end-to-end speedup depends on system size and on GEMMs being a large share of runtime.
Read the story and explore more from the AspireAI community.
Two sparse-autoencoder analyses have resisted scaling: matching semantically similar features across layers, and compressing large feature circuits into interpretable supernodes. This paper argues both are the same underlying problem, estimating semantic distances between SAE features that sit on different activation manifolds. Its representation treats each feature not as a single decoder vector but as an activation-weighted distribution over the hidden states expressing it, projects those distributions into a shared reference space, and compares them with Wasserstein distance. The representation is proved invariant to activation rescaling and stable under perturbation, and recovers true matches under finite-sample margin conditions. Empirically it outperformed decoder-vector and LLM-based baselines and captured subtle functional distinctions between related features, and it compressed large feature circuits into interpretable supernodes automatically. The limitation is that this is interpretability tooling: the paper demonstrates better feature matching and circuit compression, not that downstream interpretability claims become more accurate.