✨ 要点🔬 技术摘要
现代计算机的内部架构正呈现出一种日益增长的分裂。一方面,是专门为人工智能设计的强大引擎,它们擅长极速执行数十亿次简单的计算。这些引擎在处理短小且简单的数字时表现最佳,通过牺牲极小部分的细节来换取巨大的速度。另一方面,科学发现的世界——模拟天气模式、模拟原子如何结合或预测流体流动——仍然依赖于长而精确的数字。这些科学计算需要每一分细节来保持稳定性和准确性,但处理它们的标准计算机部件通常比新型 AI 引擎更慢,效率也更低。这产生了一个困境:科学家们既需要新硬件的速度,又无法承受失去其工作所需的精度。
来自中国迈格芬(Maginfra Co., Ltd.)的研究人员探索了一种利用特定类型的计算机芯片(即 Intel AMX)来弥合这一差距的方法。他们的目标是观察能否“欺骗”那些快速、低精度的 AI 引擎,使其执行科学领域所需的缓慢、高精度的数学运算。研究人员并没有直接要求芯片进行复杂的数学运算,而是将问题分解成更小、更简单的部分。想象一下,尝试用一把只标有整英寸刻度的尺子去测量一段很长的距离。你可以先测量整英寸的部分,然后测量剩余的分数,接着测量剩下的微小碎片,最后将它们全部相加,从而得到一个精确的总量。研究人员将这种逻辑应用到了数字上。他们将一个复杂的数字拆分为几个简单的部分,以便快速的 AI 引擎能够轻松处理。然后,他们对这些部分进行了多次快速计算,并仔细地将结果重新组合起来,以重建最终的高精度答案。
该团队在两个不同的精度水平上测试了这种方法。首先,他们处理了单精度数学运算,这是许多科学应用的标准。他们发现,通过将每个数字拆分为三个部分并运行六次特定的计算,他们可以获得与现有最佳软件同样准确的结果,但速度显著提升。在他们测试的计算机芯片上,这种方法比标准的数学运算方式快了 1.14 到 2.56 倍。这种加速在处理大规模数据集时最为明显,因为在这种情况下,拆分和重组数字所带来的额外开销相对于计算本身的巨大速度而言变得微不足道。
当他们转向双精度数学运算——这种运算更加精确,用于最严苛的科学模拟时——挑战也随之增加。在这里,研究人员必须将每个数字拆分为六个部分。由于计算过程需要极其谨慎地进行重组,整个过程变得更加复杂。他们测试了该方法的不同版本,保留的计算小块数量从六个到二十一个不等。他们发现了一个明显的权衡:保留的碎片越多,答案就越精确,但也会使过程变慢。仅使用六个碎片时,该方法对于超大型问题而言足够快,能够比标准软件快达 1.7 倍。然而,随着他们通过增加碎片数量来提高精度,管理这些碎片的额外工作量逐渐抵消了速度优势。最终,当试图保留二十一个碎片时,尽管它更精确,但该方法反而比标准方法更慢。
这项研究还强调,这种技术并非适用于所有情况的万能方案。它在计算的数字保持在特定范围内时效果最好,类似于一把长度有限的尺子在没有特殊调整的情况下,无法测量过长或过短的距离。研究人员指出,他们的方法并不适用于每一种可能的数字类型,特别是那些极大或极小的数字,并且它不能保证与现有软件实现位对位的完美匹配。相反,它为那些需要在其数据符合该方法限制的前提下,同时追求高速度和高精度的科学家提供了一种新的工具。通过展示低精度硬件也可以用来解决高精度问题,这项工作预示着一个未来:即为人工智能而构建的专用引擎,也可以为科学发现的繁重任务提供加速。
技术摘要:基于 Intel AMX 的 FP32 与 FP64 GEMM 之 BF16 组件-乘积仿真
1. 问题陈述
现代 CPU 正日益集成高吞吐量的矩阵引擎(例如 Intel 高级矩阵扩展,简称 AMX),这些引擎针对低精度 AI 工作负载(如 BF16、FP8)进行了优化。然而,计算流体力学和量子化学等领域的许多科学计算应用仍然依赖 FP32 和 FP64 通用矩阵乘法(GEMM),以确保数值稳定性和可重现性。这造成了硬件上的不匹配:专为 AI 设计的高性能低精度引擎无法直接用于高精度的科学计算任务。
本文解决的核心问题是:能否通过算法手段,将旨在处理低精度 AI 的低精度矩阵吞吐量,重新利用于加速 CPU 上的高精度科学 GEMM。挑战在于,如何在不产生过高的分解、数据移动和重构开销的前提下,利用 BF16 组件来仿真 FP32 和 FP64 算术,同时保持与 Intel oneMKL 等标准库相当的数值精度。
2. 方法论
本文提出了一种组件-乘积仿真框架,该框架将高精度操作数分解为多个 BF16 矩阵,使用 AMX 计算选定的低精度乘积,然后将结果重构为目标格式。
2.1 浮点格式与定义域
BF16: 保留了 FP32 的 8 位指数范围,但将尾数精度降低至 8 位。
FP32 仿真: 对标准的 FP32 输入进行操作。
FP64 仿真: 在受限定义域(F B 64 FB_{64} F B 64 )内进行操作,其中 FP64 值位于正常的 BF16 指数范围内。这种限制避免了对每个分量进行缩放因子的需求,从而保留了性能余量。
2.2 分解策略
AMX-FP32 路径: 使用三组件 BF16 残差分解 。每个 FP32 操作数通过直接残差舍入被拆分为三个 BF16 切片(A 0 , A 1 , A 2 A_0, A_1, A_2 A 0 , A 1 , A 2 )。这遵循了“Henry 式”方法,该方法利用共享的指数范围来确保 FP32 中残差的精确减法。
AMX-FP64 路径: 使用简化的固定六切片 Ozaki 分解 。每个 FP64 操作数被拆分为六个 BF16 切片。与使用动态缩放的完整 Ozaki 方案不同,该方法将提取的块直接存储为 BF16 矩阵(假设它们保持为规格化数),从而消除了对显式缩放元数据和重构权重的需求。
2.3 乘积调度与重构
FP32 重构: 该方法评估六个选定的组件乘积 (组件交互矩阵的前三个对角线)。至关重要的是,这些乘积直接在 FP32 AMX 瓷砖(tiles)中累加。实现采用了 操作数复用调度 (例如 A 1 B 1 → A 1 B 0 → … A_1B_1 \to A_1B_0 \to \dots A 1 B 1 → A 1 B 0 → … ),使累加器瓷砖保持驻留,从而最大限度地减少了瓷砖加载/存储的流量。
FP64 重构: 该方法评估可变数量的乘积(6、10、15 或 21 个)。与 FP32 路径不同,每个选定的 BF16 乘积都在 FP32 中计算,存储到内存,提升宽度至 FP64,然后在 FP64 块累加器中进行累加 。这可以防止 FP32 累加误差污染高精度结果,但由于重复的存储和格式转换,会产生显著的开销。
2.4 实现细节
该实现采用分块 GEMM 结构,具有 256 × 256 256 \times 256 256 × 256 的输出块和 32 × 32 32 \times 32 32 × 32 的微型瓷砖(micro-tiles)。它依赖于预计算的、块连续的组件缓冲区,以优化内存访问模式。A 操作数按标准行优先布局存储,而 B 操作数则被预先打包成 AMX BF16 点积指令所需的 VNNI 布局。
3. 主要贡献
算法桥梁: 展示了一种实用的方法,通过使用 Intel AMX BF16 指令来仿真 FP32 和 FP64 GEMM,弥合了 AI 优化硬件与科学计算需求之间的鸿型。
具备 FP32 级精度的 FP32 仿真: 表明六乘积 BF16 调度结合 FP32 累加可以达到与 oneMKL SGEMM 相当的精度(通过范数相对误差和 GEMM 缩放后的分量级误差进行衡量),尽管并未声称达到位级一致(bitwise identity)。
受控的 FP64 仿真: 通过在固定的六切片分解中改变保留的组件乘积数量(6 到 21 个),引入了一种可调的精度-性能权衡机制。
优化的数据移动: 开发了特定的调度策略(操作数复用用于 FP32)和数据布局(VNNI 打包、预计算缓冲区),以减轻 CPU 矩阵引擎上的分解和重构开销。
4. 实验结果
实验在双路 Intel Xeon Platinum 8462Y+ 服务器上进行,将提出的内核与 Intel oneMKL SGEMM 和 DGEMM 进行对比。
FP32 性能: AMX-FP32 实现优于 oneMKL SGEMM ,涵盖了所有测试的方阵规模(N = 256 N=256 N = 256 至 N = 32768 N=32768 N = 32768 )。加速比在 1.14× \times × 到 2.56× \times × 之间,在较大的矩阵规模下观察到最高的增益,因为此时 BF16 吞吐量的优势抵消了分解开销。
FP64 性能: 结果受到更多限制。
AMX-FP64-6 (6 个乘积)对于大矩阵(N ≥ 16384 N \ge 16384 N ≥ 16384 )实现了高达 1.72× \times × 的 DGEMM 加速。
AMX-FP64-10 仅对最大的矩阵有边际收益。
AMX-FP64-15 和 AMX-FP64-21 在测试配置中并未超越 DGEMM,因为存储、转换和累加额外乘积的成本消耗了可用的性能余量。
精度:
AMX-FP32 实现的误差在 10 − 7 10^{-7} 1 0 − 7 量级,与 FP32 级精度一致。
AMX-FP64 各变体显示出随着乘积数量增加而降低的误差,其中 21 乘积变体在接近 DGEMM 级精度的同时,也付出了性能代价。
5. 重要性与声明
论文将这项工作定位为一项早期算法研究 ,而非完整的原生高精度库替代方案。其重要性在于证明了:
通过精心的算法分解和数据移动优化,可以有效地利用低精度 CPU 矩阵引擎进行高精度科学计算。
存在一个明确的性能天花板 :由于在 AMX 瓷砖寄存器之外进行提升宽度和累加的开销,虽然 FP32 仿真非常高效,但 FP64 仿真受到了限制。
该方法提供了一种可调的权衡 ,允许应用程序根据其特定的精度需求与性能需求之间的平衡来选择乘积数量。
作者明确指出了局限性:当前原型仅支持方阵,要求输入处于 BF16 指数范围内(针对 FP64),并且不处理溢出/下溢或异常值。未来的工作确定为扩展该框架以支持矩形矩阵、通过缩放分解支持更广泛的输入定义域,以及适配其他 CPU 矩阵引擎(如 Arm SME)。
每周获取最佳 computer science 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。