Ozaki Scheme II：用 CRT 把高精度 GEMM 映射成少量 INT8 GEMM

论文：Katsuhisa Ozaki, Yuki Uchino, Toshiyuki Imamura, Ozaki Scheme II: A GEMM-oriented emulation of floating-point matrix multiplication using an integer modular technique, arXiv:2504.08009v3, 2025.

如果你还不熟悉 Scheme I 的位切分和 CRT 的构造公式，建议先读：Ozaki Scheme I 与中国剩余定理：从位切分到 Residue Number System。

一句话结论

这篇论文最重要的变化不是“换了一种切 FP64 的方法”，而是把 高精度矩阵乘法的表示方式换掉了。

传统 Ozaki Scheme I 把矩阵按有效位切成多个 slice：

$$A \approx \sum_{i=1}^{k} A_i,\qquad B \approx \sum_{j=1}^{k} B_j,$$

然后需要计算许多交叉项 $A_iB_j$。若保留 $i+j\le k+1$ 的项，GEMM 数量约为

$$\frac{k(k+1)}{2}.$$

例如 8 个 slice 就需要 36 次左右的低精度矩阵乘法。

Ozaki Scheme II 不再做这些交叉项。它先把缩放后的矩阵近似成整数矩阵 $A',B'$，然后选择多个两两互素的模数 $m_1,\dots,m_s$。对每个模数只计算一次

$$C_t=(A'\bmod m_t)(B'\bmod m_t),$$

最后利用 Chinese Remainder Theorem, CRT 把这些余数重新拼回 $A'B'$。

因此，精度提高一档只需要 再增加一个模数、再做一次 GEMM。核心计算数量从 Ozaki I 的近似 $O(k^2)$ 变成 Ozaki II 的 $O(s)$。

---

1. 它到底想解决什么问题？

现代 GPU 的低精度 Tensor Core 极快，但 FP64 尤其在消费级 GPU 上很慢。例如论文列出的 RTX 4090 峰值中，原生 FP64 只有约 1.29 TFLOPS，而 INT8 Tensor Core 峰值约 660 TOPS。

于是自然会有一个想法：

能不能用很多次极快的 INT8 GEMM，模拟一次高精度 FP64 GEMM？

这就是 Ozaki 系列方法的目标。

关键要求有两个：

1. 内部低精度 GEMM 最好能够 无舍入误差；
2. 额外的切分、转换、累加成本不能吃掉 Tensor Core 带来的吞吐优势。

Ozaki Scheme I 已经证明这条路线可行，但问题是 slice 越多，交叉乘法数量增长很快。

---

2. Ozaki Scheme I：为什么 GEMM 数量会爆炸？

设原矩阵为

$$A\in\mathbb{F}^{p\times q},\qquad B\in\mathbb{F}^{q\times r}.$$

传统 Ozaki 会把每个数的 significand 拆成多个低精度部分：

$$A=A_1+A_2+\cdots+A_k,$$ $$B=B_1+B_2+\cdots+B_k.$$

于是

$$AB=\sum_i\sum_j A_iB_j.$$

为了逐渐提高精度，可以只保留“较重要”的三角区域，例如

$$i+j\le k+1.$$

所以所需 GEMM 数量为

$$1+2+\cdots+k=\frac{k(k+1)}2.$$

论文给出的对应关系是：

slice 数	GEMM 数
2	3
3	6
4	10
5	15
6	21
7	28
8	36
10	55

这意味着：精度大致随着 slice 线性提升，但算力开销却近似二次增长。

这正是 Scheme II 要消掉的结构性浪费。

---

3. Scheme II 的第一步：先把 FP 数变成“大整数”

论文首先引入两个对角缩放矩阵 $D,E$，其对角元素都是 2 的幂：

$$C=AB=D^{-1}(DA)(BE)E^{-1}.$$

由于乘以 2 的幂在二进制浮点中可以精确改变指数而不破坏 significand，因此可以把

$$DA\approx A',\qquad BE\approx B'$$

截断成整数矩阵：

$$A'\in\mathbb{Z}^{p\times q},\qquad B'\in\mathbb{Z}^{q\times r}.$$

问题就变成：

如何高效计算一个“位数很大”的整数矩阵乘积 $A'B'$？

如果直接把大整数拆成 base-$2^7$ digit，还是会回到大量 digit 交叉乘法。

Scheme II 的答案是：不要 digit decomposition，改用 residue representation。

---

4. 核心：CRT 把一个大整数乘法变成多个小整数乘法

选取两两互素的模数

$$m_1,m_2,\dots,m_s,$$

并记

$$M=\prod_{t=1}^{s}m_t.$$

对于整数矩阵 $A',B'$，对每个模数分别构造

$$A'_t=A'\bmod m_t,$$ $$B'_t=B'\bmod m_t.$$

那么

$$A'B'\bmod m_t = (A'\bmod m_t)(B'\bmod m_t)\bmod m_t.$$

所以第 $t$ 个 residue domain 只需要一发 GEMM：

$$C_t=A'_tB'_t.$$

然后再取

$$C_t\bmod m_t.$$

只要知道同一个整数在所有模数下的余数，CRT 就能恢复它在模 $M$ 意义下的唯一表示：

$$X\equiv\sum_{t=1}^{s} C_t M_t y_t\pmod M,$$

其中

$$M_t=M/m_t,$$

而 $y_t$ 是 $M_t$ 在模 $m_t$ 下的逆元。

这些 $M_t y_t$ 都可以提前预计算。

为什么这一下就省掉大量 GEMM？

因为每个模数是一个 完整的大整数表示通道，不是一个 positional digit。

• digit decomposition：不同 digit 之间必须做交叉项；
• residue decomposition：每个 residue channel 都是封闭的，乘法后仍然只属于自己的 modulo ring。

也就是说

$$(A\bmod m)(B\bmod m)$$

自己就包含了原始大整数乘法在这个模数下的全部信息，没有 $i\times j$ 的 cross term。

因此：

$$\boxed{\text{Scheme I: }O(k^2)\text{ GEMMs}}$$ $$\boxed{\text{Scheme II: }O(s)\text{ GEMMs}}$$

这是整篇论文最重要的结构变化。

---

5. 但 CRT 不是只能恢复“模 $M$”吗？如何得到真正的 $A'B'$？

这是理解这篇论文最关键的数学条件。

CRT 告诉我们的是

$$X\equiv A'B'\pmod M.$$

也就是说可能的答案有

$$\dots,X-2M,X-M,X,X+M,X+2M,\dots$$

怎么知道到底是哪一个？

论文要求真实整数乘积的范围完全落在一个宽度小于 $M$ 的区间里。

若

$$-c_{\max}\le (A'B')_{ij}\le c_{\max},$$

并保证

$$2c_{\max}<M,$$

那么所有真实答案都落在

$$(-M/2,M/2)$$

中，因此 CRT 的 symmetric representative 就是唯一真实答案。

论文用一个保守上界：

$$c_{\max}\le q\,2^{k_A+k_B}.$$

于是需要满足

$$q2^{k_A+k_B}<\frac M2.$$

所以可以选择

$$k_A+k_B= \left\lfloor \log_2\frac{M/2-1}{q} \right\rfloor.$$

当 $A,B$ 精度相同时，令

$$k_A=k_B= \left\lfloor \frac12\log_2\frac{M/2-1}{q} \right\rfloor.$$

这条式子非常重要：

模数越多 $\Rightarrow M$ 越大 $\Rightarrow$ 可以保留的整数有效位 $k_A,k_B$ 越多 $\Rightarrow$ 对 FP 输入的截断误差越小。

所以 Scheme II 可以通过增加模数数量 $s$ 来平滑控制精度，而且 每增加一个模数只增加一次 GEMM。

---

6. 为什么 INT8 Tensor Core 可以“精确”算这些 GEMM？

GPU 版本要求

$$2\le m_t\le 256.$$

使用 symmetric modulo 后，余数可以放进 signed INT8：

$$-128\le a'_{t,ij},b'_{t,ij}\le 127.$$

INT8 Tensor Core 的乘积会累加到 INT32。

一个 dot product 的最坏情况量级大约为

$$q\cdot 128^2=q\cdot 2^{14}.$$

只要它不会溢出 INT32，整个 dot product 就是 整数域精确运算，不存在浮点 rounding。

论文给出的条件是当

$$q<2^{17}$$

时，可以直接使用 cublasGemmEx 的 INT8 输入、INT32 累加而保持精确；更大的 $q$ 可以进一步使用 block matrix multiplication。

这也是这个方法特别适合 Tensor Core 的地方：

Tensor Core 虽然“低精度”，但这里恰恰利用的是 整数乘法 + 足够宽的整数累加器，因此低精度输入并不意味着矩阵乘积有数值舍入误差。

真正的精度损失主要来自最开始把 $DA,BE$ 截断到 $A',B'$。

---

7. Algorithm 1 可以压缩成这条流水线

对于给定的模数数量 $s$：

Step 1：决定可保留多少 bit

根据 $M=\prod m_t$ 和内积长度 $q$，计算 $k_A,k_B$。

Step 2：缩放并整数化

$$A'=\operatorname{trunc}(DA),$$ $$B'=\operatorname{trunc}(BE).$$

Step 3：对每个模数执行一次低精度 GEMM

对 $t=1,\dots,s$：

$$A'_t=A'\bmod m_t,$$ $$B'_t=B'\bmod m_t,$$ $$C'_t=A'_tB'_t.$$

然后再对输出取模：

$$C''_t=C'_t\bmod m_t.$$

Step 4：CRT accumulation

$$Z=\sum_t C''_t M_t y_t.$$

Step 5：range reduction

$$C'''=Z\bmod M$$

并选择落在合法范围内的 symmetric representative。

Step 6：反缩放

$$C=D^{-1}C'''E^{-1}.$$

最终得到对 $AB$ 的高精度近似。

---

8. Scheme II 相比 Scheme I 还有一个容易忽略的优点：显存

Scheme I 需要保存多个 slice：

$$A_1,\dots,A_k,\qquad B_1,\dots,B_k.$$

而 Scheme II 可以按模数流式处理：

1. 生成 $A'_t,B'_t$；
2. 做 $C_t=A'_tB'_t$；
3. 累加到 CRT reconstruction；
4. 立即丢掉 $A'_t,B'_t,C_t$。

所以它天然适合做 streaming / pipelining，也更容易降低临时 workspace。

从 GPU kernel 设计角度看，这一点甚至可能和 GEMM 数量减少同样重要。

---

9. 为什么论文有 OS II-fast 和 OS II-accu？

前面为了保证 CRT reconstruction 唯一，需要估计

$$c_{\max}\ge\max_{ij}(|A'||B'|)_{ij}.$$

估计越保守，就必须选择越小的 $k_A,k_B$，也就是更早截断输入，导致精度下降。

论文实现了两种方案：

OS II-fast

用 Cauchy-Schwarz 等便宜上界来估计 $c_{\max}$。

优点：快。

缺点：上界可能过松，浪费 CRT dynamic range，因此 truncation error 更大。

OS II-accu

额外用低精度 GEMM 来更准确地估计范围。

优点：能把 $k_A,k_B$ 开得更大，保留更多有效位，因此面对 exponent distribution 更宽的数据时精度明显更好。

代价：多一些前处理计算。

这说明 Scheme II 的核心 trade-off 不只是“GEMM 次数”，而是：

$$\text{更精确的 range estimation} \Longleftrightarrow \text{更多可保留 significand bits}.$$

---

10. 实验结果应该怎么看？

精度

在论文用的 HPL-like 数据分布（$\phi=0.5$）下，Ozaki Scheme II 大约需要 14–15 个模数 才能达到与 DGEMM 相近的精度。

也就是说大约是 14–15 次 INT8 GEMM。

相比之下，Ozaki Scheme I 通常需要 7–8 个 slice，而这对应约 28–36 次矩阵乘法。

这基本验证了论文最重要的目标：

达到 FP64-like 精度时，把低精度 GEMM 数量大约砍半。

GPU 性能

论文报告：

• RTX 4090：FP64-equivalent 约 7.4–9.8 TFLOPS；
• GH200：约 56.6–80.2 TFLOPS；
• GH200 上 $16384^3$、$s=14$ 时达到 80.2 TFLOPS。

而论文测得 GH200 原生 FP64 Tensor Core DGEMM 约为 61.9 TFLOPS，因此在足够大的矩阵上，模拟 FP64 可以达到甚至超过原生 FP64 的吞吐。

这里要注意：这不是说 14 次 INT8 GEMM 比 1 次 FP64 GEMM 的 FLOP 数少，而是 INT8 Tensor Core 的硬件吞吐高太多，因此即使执行很多次仍然可能赢。

---

11. 真正瓶颈居然不是 INT8 GEMM

论文的 timing breakdown 很值得看。

当矩阵足够大时，INT8 Tensor Core GEMM 本身已经非常快，真正长期占主导的是：

$$\boxed{\text{FP64} \rightarrow \text{scaled integer} \rightarrow \text{INT8 residue conversion}}$$

也就是把原始 double 数据转换成每个模数下的 INT8 矩阵。

论文观察到：

• 小矩阵：kernel launch overhead 明显；
• 中等矩阵：FP64 → INT8 conversion 成为主要瓶颈；
• 大矩阵：GEMM 与 accumulation 的相对开销进一步下降，但 conversion 仍然占主导。

这个结果对 GPU kernel 优化非常有启发：

Scheme II 下一步的优化重点很可能不是继续优化 GEMM，而是把 scaling、truncation、modulo、layout conversion 融合成高吞吐 preprocessing kernel，甚至考虑和 GEMM pipeline 做 producer-consumer overlap。

从算子工程视角，它已经从“GEMM 优化问题”转化成了“数据表示转换问题”。

---

12. 为什么 Blackwell 时代这类方法反而更有意义？

论文给出的硬件趋势很直观：低精度 Tensor Core 的峰值增长远快于 FP64。

例如它列出的 B200 规格中：

• INT8 Tensor Core：约 4500 TOPS；
• FP64 Tensor Core：约 40 TFLOPS。

因此 INT8 / FP64 的峰值比越来越夸张。

Scheme II 的价值就在于利用这个 precision throughput gap：

$$\text{many cheap low-precision GEMMs} < \text{one expensive high-precision GEMM}$$

只要额外的数据转换与 CRT reconstruction 足够便宜，架构越偏向低精度，这种 emulation 越可能有优势。

---

13. CPU 上为什么也有价值？

论文不仅做 INT8 → FP64 emulation，还展示了另一种用法：

用 FP64 GEMM 模拟 quad-word 高精度矩阵乘法。

这里余数矩阵仍然被设计成足够小，使得 FP64 GEMM 中每个整数 dot product 都能 exactly represent。

若 unit roundoff 为

$$u=2^{-53},$$

论文选择模数满足

$$qm_i^2\le 4u^{-1},$$

从而保证

$$(|A'_t||B'_t|)_{ij}\le u^{-1},$$

这样 FP64 的 53-bit significand 就可以精确表达内部整数累加结果。

实验中，Scheme II 在 quad-word matrix multiplication 上相对 Scheme I 最多取得约 2.3× speedup。

所以这套方法本质上并不绑定 INT8 Tensor Core：

它是一种 用高吞吐 GEMM 数据类型作为“精确的小整数乘法器”，再用 CRT 扩大有效整数位宽 的通用框架。

---

14. 我认为这篇论文最值得记住的三个点

1. 从 positional representation 切换成 residue representation

Ozaki I 是“按位切片”；Ozaki II 是“按模数投影”。

前者乘法产生 cross terms，后者每个 modulo channel 独立封闭。

这是 GEMM 数从二次降到线性的根本原因。

2. 低精度 Tensor Core 不等于低数值精度

只要输入被限制成小整数、accumulator 足够宽，INT8 × INT8 → INT32 可以作为 exact arithmetic primitive。

低精度只描述 operand 的位宽，不代表最终算法必须有低精度 rounding error。

3. 算法优化最终会把瓶颈推到 representation conversion

当 GEMM 快到一定程度后，真正决定性能的是

• scaling；
• truncation；
• modulo；
• FP64 → INT8 conversion；
• CRT accumulation / reduction。

这和很多现代 GPU 算法的演化路径一样：核心 Tensor Core 算子被优化到极致后，瓶颈迁移到 surrounding kernels 和 data movement。

---

15. 和你熟悉的 GPU/NLA 视角怎么连接？

可以把这篇论文理解成一种 “算术格式层面的 communication-avoiding”：它不是减少节点通信，而是减少不同 precision slice 之间必须发生的交叉计算。

传统 slice expansion 中，表示本身制造了 $A_iB_j$ 的组合爆炸；CRT 改变 representation 后，每个 residue channel 都能独立完成 GEMM，因此直接把计算 DAG 从三角形压成一条长度为 $s$ 的链/并行集合。

如果未来自己实现 GPU 版本，我会优先关注：

1. mod conversion fusion：把 scaling + trunc + symmetric modulo + INT8 pack 融为一个 kernel；
2. 多模数并行：不同 $m_t$ 天然独立，可以 stream / batched 调度；
3. workspace streaming：一组 residue GEMM 完成立即回收输入；
4. CRT reconstruction fusion：尽量避免把每个 $C_t$ 全部落 HBM 后再单独 reduce；
5. Blackwell INT8 pipeline：研究转换 kernel 是否会彻底成为 roofline 上的新瓶颈。

换句话说，论文已经证明数学上“少做 GEMM”是可行的，而真正有 GPU systems 味道的后续工作，是让 representation conversion 和 reconstruction 不拖后腿。

---

16. 读原文时要避开的一个公式笔误

原文第 3.2 节在讨论 FP64 模数时写道

$$M \approx s m_s$$

并据此说 $M$ 与 $s$ 成正比。按前文定义

$$M=\prod_{i=1}^{s}m_i,$$

而该段的前提是各个模数大小非常接近，即 $m_i\approx m_s$，因此正确的数量级应是

$$M\approx m_s^s, \qquad \log_2 M\approx s\log_2m_s.$$

所以真正近似随模数数量 $s$ 线性增长的是 $\log_2M$，进而是可保留的位数 $k_A+k_B$，而不是 $M$ 本身。这个笔误不影响算法和实验：式 (17)、(18) 实际使用的正是 $\log_2M$。

当前开源实现是 RIKEN-RCCS/GEMMul8。其代码仍按四阶段组织：scaling/quantization、低精度矩阵乘法、乘积重新取模、CRT reduction + undo scaling；较新版本还扩展到了 CUDA/HIP、INT8/FP8 和多种 BLAS Level-3 操作。

---

Reference

• Katsuhisa Ozaki, Yuki Uchino, Toshiyuki Imamura. Ozaki Scheme II: A GEMM-oriented emulation of floating-point matrix multiplication using an integer modular technique. arXiv:2504.08009v3, 2025. https://arxiv.org/abs/2504.08009