Improved Scaling for Fast Mode of Ozaki Scheme II

Technical Summary: Improved Scaling for Fast Mode of Ozaki Scheme II

Problem Statement
Ozaki Scheme II is a method for emulating high-precision matrix multiplication (e.g., FP64/FP32) using low-precision integer hardware (e.g., INT8 Tensor Cores) based on the Chinese Remainder Theorem (CRT). The scheme requires scaling high-precision input matrices into integer matrices such that the product of these integers, when reconstructed via CRT, does not exceed a specific threshold ($P/2$). If this threshold is exceeded, the CRT recovery fails, leading to catastrophic accuracy loss.

The scheme offers two modes for determining scaling factors ($\mu$ and $\nu$):

1. Accurate Mode: Uses an auxiliary INT8 matrix multiplication (GEMM) to estimate scaling factors, ensuring high accuracy but incurring significant computational overhead.
2. Fast Mode: Uses the Cauchy–Schwarz inequality to estimate scaling factors without the auxiliary GEMM, offering higher throughput but historically lower accuracy.

This paper identifies a critical flaw in the existing Fast Mode scaling formula: it lacks scale invariance. When input matrices are multiplied by a scalar constant $\alpha$, the effective bit width of the resulting scaled integer matrices changes unpredictably.

• For large $\alpha$, the effective bit width decreases, causing accuracy degradation.
• For small $\alpha$, the scaled integer product may exceed the CRT threshold ($P/2$), causing complete recovery failure.
  Numerical experiments on all-ones matrices and random matrices demonstrated that for certain scalar ranges, the Fast Mode fails to recover the correct result, whereas the Accurate Mode remains stable.

Methodology
The authors propose a revised scaling formula for the Fast Mode that is derived directly from the CRT uniqueness condition via the Cauchy–Schwarz inequality.

1. Theoretical Derivation:
   The authors reformulate the condition for the scaling factors $\mu_i$ and $\nu_j$. Instead of the previous heuristic, they define a sufficient condition where:
   $$2\mu_i \|a_{i,:}\|_2 \|b_{:,j}\|_2 \nu_j \leq P - 1$$
   By solving for $\mu_i$ and $\nu_j$ using base-2 logarithms and accounting for the absolute error bounds of the log2f function provided by GPU math APIs, they derive new formulas:
   $$\mu_i := 2^{\lfloor P'_{prop} - 0.5 \cdot (\widehat{\log_2}(\sum a_{ih}^2) + 4u_{32}) \rfloor}$$
   where $P'_{prop}$ is a precomputed constant and $\widehat{\log_2}$ is a bounded approximation of the logarithm.

2. Scale Invariance Proof:
   The authors analytically demonstrate that the new formula is scale-invariant. If the input matrix $A$ is scaled by $\alpha = 2^s$, the term $\log_2(\alpha)$ cancels out in the calculation of the effective bit width of the scaled integer matrix. Consequently, the bit width remains constant regardless of the input magnitude, preventing the accuracy degradation and CRT failure observed in the original Fast Mode.

3. Implementation:
   The proposed method (OS II-prop) replaces the scaling formula in the Fast Mode of the GEMMul8 library. It retains the $O(mk)$ and $O(kn)$ computational complexity of the original Fast Mode, introducing no additional overhead compared to the original Fast Mode and remaining faster than the Accurate Mode (which requires an extra GEMM).

Key Contributions

• Identification of Limitation: The paper demonstrates that the existing Fast Mode scaling formula lacks scale invariance, leading to accuracy degradation and CRT recovery failure for specific input scalings.
• Revised Formula: A new scaling formula is proposed that guarantees scale invariance and ensures the CRT uniqueness condition is always satisfied.
• Zero Overhead: The new formula introduces no additional computational cost over the original Fast Mode.
• Comprehensive Evaluation: The method is evaluated on an NVIDIA GH200 GPU for both Double-Precision (DGEMM) and Single-Precision (SGEMM) matrix multiplications.

Results
Experiments were conducted using random matrices with varying element magnitude spreads ($\phi$) and matrix dimensions ($k \in \{1024, 16384\}$).

• Accuracy: The proposed method (OS II-prop) achieves accuracy comparable to the Accurate Mode (OS II-accu) across all values of $\phi$. In contrast, the original Fast Mode (OS II-fast) showed significant accuracy degradation and failure for large $\phi$ values (ill-conditioned cases).
• Throughput: OS II-prop maintains throughput comparable to the original Fast Mode. While OS II-prop is slightly slower than the original Fast Mode in specific scenarios with small moduli counts (due to fewer zero elements in the scaled integers), it significantly outperforms the Accurate Mode in throughput.
• Trade-off: In the accuracy-throughput trade-off analysis, OS II-prop dominates OS II-accu (higher throughput for the same accuracy) and overcomes the accuracy limitations of OS II-fast. It is the only method capable of simultaneously outperforming cuBLAS in both accuracy and throughput for a wide range of conditions.

Significance
The paper claims that the proposed method resolves the fundamental accuracy limitation of the Fast Mode in Ozaki Scheme II without sacrificing its performance advantage. By ensuring scale invariance, the method guarantees that CRT recovery succeeds regardless of input scaling, offering a superior balance of accuracy and performance. This allows users to achieve high-precision results on low-precision hardware with the speed of the Fast Mode and the reliability of the Accurate Mode, effectively overcoming the throughput constraints of the Accurate Mode and the accuracy constraints of the original Fast Mode.

The authors note that while the current work focuses on INT8 matrix engines for DGEMM and SGEMM, the approach is applicable to FP8 engines and complex matrix multiplications (ZGEMM/CGEMM), which are left for future work.