DGEMM with Ozaki Scheme I/II on FP4 Tensor Cores: A Base-13 E2M1 Limb Representation
Shun-ichiro Hayashi, Daichi Mukunoki, Tetsuya Hoshino, Takahiro Katagiri
cs.DC
2026-08-07
Blackwell 砍掉九成 FP64 单元后,作者用 FP4 张量核心配 13 进制拆分模拟 FP64 矩阵乘,在 RTX PRO 6000 上比 FP8 版 Ozaki 方案更快、比原生 FP64 快 8 倍。
科学计算要 FP64,但 GPU 厂商在为 AI 砍 FP64。NVIDIA B300(Blackwell Ultra,2025)把 FP64 单元砍掉了九成以上,下一代 Rubin 还会继续降。FP64 矩阵乘(DGEMM)是高性能计算的基础算子,硬件支持一缩水,整类计算就面临停滞。
低精度张量核心(Tensor Core)倒是越来越猛:Blackwell 第五代张量核心的 FP4 吞吐是 FP8 的两倍。问题在于 FP4 只有 4 位,直接算 FP64 误差大到没法用。此前已有工作用 INT8 或 FP8 张量核心去「模拟」高精度矩阵乘,思路叫 Ozaki 方案(Ozaki scheme):把高精度数拆成多段低精度「枝」,在低精度单元上分别乘,再精确累加重组。FP4 更快,但一直没人做出来,因为部分和一旦写回 FP4 格式就被舍入,精度全毁。这篇的突破是让 FP4 也能用上 Ozaki 方案。
关键观察藏在 FP4 的取值里。E2M1 格式(FP4,2 位指数 1 位尾数)能表示的值是 {0, ±0.5, ±1, ±1.5, ±2, ±3, ±4, ±6}。把这些值全部翻倍,得到的正好是一组整数 S = {0, ±1, ±2, ±3, ±4, ±6, ±8, ±12}。
作者证明:任何整数 N 都能写成 N = c + 13n(c 属于 S,n 为整数)。S 在以 13 的倍数平移之后能覆盖所有整数。于是任意整数可以被贪心拆成一串「13 进制枝」:N = Σ 13ⁱ × cᵢ,p 个枝能无空隙表示 |N| ≤ (13ᵖ−1)/3(p=1 时是 4,p=2 时 56,p=3 时 732,依此类推)。
真正解决舍入的,是不让中间和回到 FP4。两个枝相乘最大是 144,只要内积长度 K 满足 144K ≤ 2²⁴(K ≤ 116,508),和就可以完整地存在 FP32 累加器里,一个 bit 都不丢。FP4 张量核心就此成了一个精确的整数矩阵乘机器。同一套原理还能在 FP4 上逐位精确地模拟 INT8 张量核心的整数矩阵乘。
在此之上搭两套 Ozaki 方案:
实测在 NVIDIA RTX PRO 6000 Blackwell 上:FP4 峰值 2000 TFLOPS,FP8 与 INT8 各 1000,原生 FP64 只有 1.85 TFLOPS。
| 方法 | 16384³ 算力 TFLOPS | 16384³ 端到端 TFLOPS |
| cuBLAS FP64 | 1.85 | 1.85 |
| OzII-FP4(本文) | 17.31 | 15.26 |
| GEMMul8-FP8(FP8 版 Ozaki) | 15.68 | 14.65 |
| GEMMul8-INT8 | 40.94 | 36.39 |
算力阶段 OzII-FP4 比 FP8 版方案快 1.10 到 1.19 倍。到 16384³ 这个最大规模,端到端时间反过来只有 GEMMul8-FP8 的 0.96(更快);中小规模(4096³/8192³)略慢,是 1.08/1.03 倍。对照原生 FP64,OzII-FP4 算力阶段快 8.8/9.2/9.4 倍,端到端快 5.8/7.3/8.2 倍。峰值利用率 OzII-FP4 是 6165%,GEMMul8-FP8 是 5361%。预处理开销随规模增大被摊薄:N=4096/8192/16384 时分别占 14.1/7.3/3.9%。
精度上,这套方法对「转成整数后的输入」完全精确,只在最后舍入一次。当输入数量级均匀时,它甚至比 FP8 版还准。但一旦一行里数量级相差变大(共享指数量化是唯一误差源),精度掉得比 FP8 版快:从 φ=0 的 58.2 位掉到 φ=32 的 52.3 位(掉 5.9 位),FP8 版只掉 3.5 位。
硬件走向很明确:AI 的需求让低精度张量核心越来越便宜、FP64 越来越稀缺。这篇证明了给 AI 训练用的那批 FP4 单元,顺手就能拿去做高精度科学计算,不必等厂商发还 FP64 的型号。
但要泼一盆冷水:只要 INT8 单元还在,GEMMul8-INT8 就仍然最快(上表 40.94 对 17.31)。OzII-FP4 的优势真正发威,是在连 INT8 单元也被砍掉的 B300 这一代,以及之后 FP4 一家独大的 Rubin 上。它买的是一个「未来 FP4 为主」的期权,不是当下的最快解。
作者自己承认:FP4 每枝信息量少,需要的模数/枝数比 FP8 多(达到 FP64 精度要 19 个模数,FP8 版只要 13 个),矩阵乘次数更多;数量级跨度大时精度掉得更快;部分 Karatsuba 优化在 epilogue 融合后反而因为占满共享内存而变慢,没用上。
读下来的存疑:验证只在一块卡(RTX PRO 6000 Blackwell)、最大 16384³、内积 K ≤ 16384 上做;与 INT8 方案的对照其实说明这篇方法在当下硬件上未必最优,更多是面向下一代硬件的占位。端到端在中小规模仍略慢于 FP8 方案,只有足够大的矩阵才反超,这对实际 HPC 工作负载(往往分块成中等大小)收益多大,论文没展开。