Node-level Parallelisation of Ozaki Scheme II

Conor McLeod

Abstract. Driven by the AI boom, NVIDIAs newest accelerators increasingly prioritise low precision GEMMs ahead of FP64 GEMMs. To keep FP64 GEMMs fast, we look to emulation techniques. The Ozaki Scheme II is an efficient FP64 GEMM emulation technique which uses the Chinese Remainder Theorem to perform double precision GEMMs on the INT8 tensor cores. This summer, I worked on a node-level Ozaki scheme library which runs on JUPITER.

Keywords: GEMM, Ozaki Scheme, Chinese Remainder Theorem

1. Introduction

The General Matrix Multiply (GEMM) operation is perhaps the most ubiquitous mathematical operation we perform. On GPUs, we can perform these operations incredibly fast.

Science and engineering codes have classically operated in the high precision FP64 regime.

But today, NVIDIA is designing their GPUs for a different crowd. Unlike scientific computing, mainstream deep learning rarely uses FP64. AI workloads are The max precision you will see in deep learning is FP32, and with care, much lower precisions can be used. A range of low precision float formats have proliferated in the last few years.

So the need for the Ozaki scheme arises out of economic factors just as much as scientific ones. Nvidia is a profit making, publically-listed company, and they will go where the money is. Traditional HPC cannot compete with the historically unprecedented level of capital expenditure behind the AI boom. So if science and engineering want to pursue the best high precision GEMM performance, they must adapt to the way the hardware being designed today.^[1]

The Ozaki Scheme II is a scheme for emulating FP64 matrix multiplication using the Chinese Remainder Theorem.

Nvidia added FP64 emulation on INT8 tensor cores to cuBLAS in CUDA Toolkit 13.0 Update 2.

2. Background

2.1 Chinese Remainder Theroem

Let $x\in \mathbb{Z}$, and let $p_1,\dots ,p_N\in {\mathbb{N}}_{\ge 2}$ be pairwise coprime with $\mathcal{P}:=\prod _{i=1}^N{p}_i$.

Let $q_i\in \mathbb{N}$ be the modular inverse of ${\displaystyle \frac{\mathcal{P}}{p_i}}$, such that:

$$\frac{\mathcal{P}}{p_i}{q}_i\equiv 1(\mathrm{mod}{p}_i)$$

Suppose $x$ is known only through its residues:

$$\{\begin{array}{l}x\equiv y_1(\mathrm{mod}{p}_1),\\ ⋮\\ x\equiv y_N(\mathrm{mod}{p}_N).\end{array}$$

Then $x$ is recovered modulo $\mathcal{P}$:

$$x\equiv \sum _{i=1}^N\frac{\mathcal{P}}{p_i}{q}_i{y}_i(\mathrm{mod}\mathcal{P})$$

This is a weighted sum of the residues: each $y_i$ gets a fixed weight $\frac{\mathcal{P}}{p_i}{q}_i$ that depends only on the moduli

To better communicate the intution behind this

$0$$\pm \mathcal{P}/2$$\mathbb{Z}/\mathcal{P}\mathbb{Z}$integers mod $\mathcal{P}$

target $x$running sum $\hat{x}$

$\mathcal{P} = 1.374 \times 10^{19}$

$\hat{x} = \textstyle\sum_{i=1}^{8} \frac{\mathcal{P}}{p_i} q_i y_i \bmod \mathcal{P}$

$\phantom{\hat{x}} = 123\,456\,789$

✓ $\hat{x} = x$: recovered exactly

$x =$

$N = 8$

$\mathcal{P}/2 \approx 2^{62.6}$$|x| \approx 2^{26.9}$

0326496128160

bits

2.2 Ozaki Scheme II

Armed with an understanding of the Chinese Remainder Theorem. The idea is we can take input FP64 matrices, scale and truncate the values to convert them to large integers, then form matrices of residues for given moduli. We perform the matrix multiplication on these residues

$A$

×2²⁷×2³⁴×2²¹×2²⁶

2.40e−1

−2.08e−3

−1.50e−3

1.28e1

9.88e−2

−7.43e−3

1.05e−3

−7.80e−2

−5.39e2

2.02e−3

−3.83e−1

−1.23e1

4.21e0

3.83e−1

−1.62e1

6.18e0

×

$B$

×2²¹×2²⁵×2²²

6.36e2

−1.11e−2

3.80e2

8.88e−1

−2.24e2

3.88e−1

−3.64e−3

−2.71e−3

1.07e2

1.28e0

−1.14e−1

4.84e2

=

$C$

?

?

?

?

?

?

?

?

?

?

?

?

slice

1/6 FP64 inputs

Entries span six orders of magnitude.

Hover a cell to see where its value comes from.

N = 8

2.3 INT8 Tensor Cores

Tensor cores are specialised hardware units on Nvidia GPUs for performing matrix multiplication. Tensor cores support different number formats, and the lower the format, the better the performance.

3. Results

4. Discussion

5. Conclusion

Acknowledgements

I would like to thank Clément Richefort for supervising my project.

References

1. K. Ozaki, T. Ogita, S. Oishi, and S. M. Rump, ‘Error-free transformations of matrix multiplication by using fast routines of matrix multiplication and its applications’, Numer Algor, vol. 59, no. 1, pp. 95–118, Jan. 2012, doi: 10.1007/s11075-011-9478-1.
2. Y. Uchino et al., ‘High-Performance and Power-Efficient Emulation of Matrix Multiplication using INT8 Matrix Engines’, in Proceedings of the SC ’25 Workshops of the International Conference for High Performance Computing, Networking, Storage and Analysis, in ACM Conferences, 2025, pp. 1824–1831. doi: 10.1145/3731599.3767539.

---

1. Nvidia have said they are not abandoning the HPC community ↩︎