Improved Scaling for Fast Mode of Ozaki Scheme II

技术摘要：改进 Ozaki Scheme II 中 Fast Mode 的缩放机制

问题陈述
Ozaki Scheme II 是一种基于中国剩余定理（CRT）的方法，旨在利用低精度整数硬件（如 INT8 Tensor Cores）来模拟高精度矩阵乘法（如 FP64/FP32）。该方案要求将高精度输入矩阵缩放为整数矩阵，使得这些整数的乘积在通过 CRT 重构时不会超过特定阈值（$P/2$）。如果超过此阈值，CRT 恢复将会失败，导致灾难性的精度损失。

该方案提供了两种确定缩放因子（$\mu$ 和 $\nu$）的模式：

1. 准确模式（Accurate Mode）： 使用辅助的 INT8 矩阵乘法（GEMM）来估计缩放因子，确保高精度，但会产生显著的计算开销。
2. 快速模式（Fast Mode）： 使用柯西-施瓦茨不等式（Cauchy–Schwarz inequality）来估计缩放因子，无需辅助 GEMM，具有更高的吞吐量，但在历史上精度较低。

本文指出了现有 Fast Mode 缩放公式中的一个关键缺陷：它缺乏尺度不变性（scale invariance）。当输入矩阵乘以标量常数 $\alpha$ 时，所得缩放整数矩阵的有效位宽会发生不可预测的变化。

• 对于较大的 $\alpha$，有效位宽会减小，导致精度下降。
• 对于较小的 $\alpha$，缩放后的整数乘积可能会超过 CRT 阈值（$P/2$），从而导致完全的恢复失败。
  在全 1 矩阵和随机矩阵上的数值实验表明，对于某些标量范围，Fast Mode 无法恢复正确结果，而 Accurate Mode 则保持稳定。

方法论
作者提出了一个修订后的 Fast Mode 缩放公式，该公式直接通过柯西-施瓦茨不等式从 CRT 唯一性条件中推导而来。

1. 理论推导：
   作者重新构建了缩放因子 $\mu_i$ 和 $\nu_j$ 的条件。不同于之前的启发式方法，他们定义了一个充分条件，其中：
   $$2\mu_i \|a_{i,:}\|_2 \|b_{:,j}\|_2 \nu_j \leq P - 1$$
   通过使用基 2 对数并考虑 GPU 数学 API 提供的 log2f 函数的绝对误差界限，他们推导出了新的公式：
   $$\mu_i := 2^{\lfloor P'_{prop} - 0.5 \cdot (\widehat{\log_2}(\sum a_{ih}^2) + 4u_{32}) \rfloor}$$
   其中 $P'_{prop}$ 是预计算常数，$\widehat{\log_2}$ 是对数函数的有界近似。

2. 尺度不变性证明：
   作者从解析上证明了新公式具有尺度不变性。如果输入矩阵 $A$ 被缩放 $\alpha = 2^s$，则 $\log_2(\alpha)$ 项会在计算缩放整数矩阵的有效位宽时相互抵消。因此，无论输入幅度如何，位宽都保持不变，从而防止了在原始 Fast Mode 中观察到的精度下降和 CRT 失败问题。

3. 实现：
   所提方法（OS II-prop）替换了 GEMMul8 库中 Fast Mode 的缩放公式。它保留了原始 Fast Mode 的 $O(mk)$和$O(kn)$ 计算复杂度，与原始 Fast Mode 相比未引入额外开销，并且仍然比 Accurate Mode（需要额外的 GEMM）更快。

核心贡献

• 识别局限性： 本文证明了现有的 Fast Mode 缩放公式缺乏尺度不变性，导致在特定的输入缩放情况下会出现精度下降和 CRT 恢复失败。
• 修订公式： 提出了一种新的缩放公式，该公式保证了尺度不变性，并确保始终满足 CRT 唯一性条件。
• 零开销： 新公式相对于原始 Fast Mode 引入了零额外的计算成本。
• 全面评估： 该方法在 NVIDIA GH200 GPU 上针对双精度（DGEMM）和单精度（SGEMM）矩阵乘法进行了评估。

结果
实验使用具有不同元素量级跨度（$\phi$）的随机矩阵以及矩阵维度（$k \in \{1024, 16384\}$）进行。

• 精度： 所提方法（OS II-prop）在所有 $\phi$ 值下均实现了与 Accurate Mode（OS II-accu）相当的精度。相比之下，原始 Fast Mode（OS II-fast）在较大的 $\phi$ 值（病态情况）下表现出明显的精度下降和失效。
• 吞吐量： OS II-prop 保持了与原始 Fast Mode 相当的吞吐量。虽然在模数较少（由于缩放整数中的零元素较少）的特定场景下，OS II-prop 比原始 Fast Mode 略慢，但其吞吐量显著优于 Accurate Mode。
• 权衡： 在精度-吞吐量权衡分析中，OS II-prop 优于 OS II-accu（在相同精度下具有更高吞吐量），并克服了 OS II-fast 的精度限制。它是唯一能够在广泛条件下同时在精度和吞吐量方面超越 cuBLAS 的方法。

意义
论文声称，所提方法解决了 Ozaki Scheme II 中 Fast Mode 的根本精度局限性，且没有牺牲其性能优势。通过确保尺度不变性，该方法保证了无论输入缩放如何，CRT 恢复都能成功，从而提供了精度与性能的卓越平衡。这使得用户能够利用 Fast Mode 的速度和 Accurate Mode 的可靠性，在低精度硬件上获得高精度结果，有效地克服了 Accurate Mode 的吞吐量约束和原始 Fast Mode 的精度约束。

作者指出，虽然目前的工作侧重于用于 DGEMM 和 SGEMM 的 INT8 矩阵引擎，但该方法同样适用于 FP8 引擎和复数矩阵乘法（ZGEMM/CGEMM），这些内容已作为未来工作予以保留。