机器学习原子间势不再只跑FP32:PET换用张量核心低精度分解,MD步最高快1.7倍
机器学习原子间势(MLIP)的推理长期停留在FP32,因为分子动力学对力的小误差很敏感。针对EPFL Lab COSMO开发的PET模型,在Blackwell GPU上用BF16x9、在Ampere GPU上用TF32x3把FP32矩阵乘分解为低精度分量在张量核心上计算,无需重训练或改架构。在17,280原子水盒上,BF16x9使PET的大矩阵乘快2.2–2.5倍、完整MD步最高快1.7倍;TF32x3在A100上使主要GEMM快约1.7–1.9倍、MD步最高快1.5倍,并保持与FP32接近的能量守恒与结构弛豫行为。