Day 6｜先把 Ozaki / CRT 拆成 pipeline

昨天我們把原理講完了。

今天先不急著寫到 PTX、GPU Assembly，或是 mma.sync 這種東西。

我們現在懂了數學，或至少知道它大概在幹嘛：Ozaki 負責把 FP64 切開，中國餘數定理負責把小餘數拼回大整數。可是光是算一個數字再重建，就已經有夠麻煩了。如果一開始就把切片、取餘數、GEMM、CRT、scaling 全部混在同一顆 kernel 裡，答案錯的時候根本不知道該從哪裡查。

所以 Day 6 的目標不是快。

今天只做設計拆解：

先確定每個地方都做對
後面再來想怎麼加速

我們要把 Day 5 的東西拆成三個步驟的 pipeline：

1. 建構 CRT，定義大整數空間。
2. 求餘數，得到 INT8 餘數矩陣。
3. GEMM 之後重建，做 scaling 回到原本的數字尺度。

這三段分開看，才有辦法對答案。
(後面我再來細修文章AI感太重 或是 其他問題)
這幾天專案使用的 代碼連結adaptive_GEMM
有任何問題 或是發現我寫的怪怪的 都非常歡迎跟感謝願意跟我說><
可以使用這個問題回報
或是直接email我信箱 謝謝!

先對一下 repo 裡的主線

目前專案裡的命名還叫 Phase22、Phase24。這些名字聽起來很像 AI 自己亂生出來的代號，之後應該會再改；現在寫文章先將就一下，之後如果 repo 名稱整理過，我再回來改文。

主線要看的是 Phase 22。

它在 OzakiConfig 裡預設是：

ExecutionMode mode = ExecutionMode::Phase22;
bool enable_fp64 = false;
bool enable_residual = false;

這代表它走的是 INT8 + CRT 的快路徑，不會偷偷在後面補一段 naive FP64 residual GEMM。

為什麼要特別講這句？

因為 Phase 24 會開 residual。它看起來比較準，但那是因為後面又用 FP64 重算了一段，不是 Ozaki / CRT 自己真的做到 lossless double。這系列接下來要講的是 Phase 22 這條路。 同時 這個案例也不是完全精度fp64 註解1 會說明到

回到 code 狀態，Phase 22 fast path 大概長這樣：

double sh = 2^17;
double sl = 2^34;

precompute_modulo_hl_A(..., sh, sl);
precompute_modulo_hl_B(..., sh, sl);

for pass in {hi×hi, hi×lo, lo×hi, lo×lo}:
    crt_fast_kernel(..., inv[pass]);

Double 展開後，會對應到 high / low 的四種組合：

A_hi × B_hi
A_hi × B_lo
A_lo × B_hi
A_lo × B_lo

縮放比例分別是：

hi × hi  →  2^-34
hi × lo  →  2^-51
lo × hi  →  2^-51
lo × lo  →  2^-68

融合過後的版本會直接把這些處理掉。

但今天我們先不要站在 fused 版本看事情。今天要把它拆回三段，確定每段都能自己對答案。

Pipeline 長什麼樣

把 Day 5 的數學收成工程流程，大概是這樣：

(1) 建構 CRT
        選定質數 p_i
        算 M、M_i、逆元、重建係數 f_i
        定義 INT8 餘數可以重建多大的整數

(2) 求餘數
        FP64 → high / low 整數薄片
        整數薄片 → x mod p_i
        得到 N 張餘數矩陣

(3) 重建
        對每個 p_i 做完 GEMM
        用 CRT 把結果拼回大整數
        乘上 Ozaki scaling，回到 double 尺度

我們的案例裡 N = 7，因為用了 7 個質數。

但要注意一件事：不是整個系統只有 7 張矩陣。

比較精確地說，每一個 slice 都會對 7 個質數取餘數。A 有 high / low，B 也有 high / low，所以會出現：

A_high mod p_i
A_low  mod p_i
B_high mod p_i
B_low  mod p_i

而 GEMM pass 會用這四組組合出：

high × high
high × low
low  × high
low  × low

所以「7 張餘數矩陣」是指某一個 slice 對 7 個質數的視角。完整 pipeline 會把 A/B 與 high/low 都算進來。

第一步：建構 CRT，定義大整數空間

CRT 不是每個元素都臨時計算一次。

質數、乘積、轉換係數，對整個 GEMM 來說都是常數。這些東西比較像是「這套數字系統的規格表」。

以我們的專案為例，質數必須小於 128：

127, 113, 109, 107, 103, 101, 97

原因很現實：餘數最後要塞進 int8_t，再送去 INT8 GEMM。int8_t 的正數上限是 127，所以我們會挑小於 128、又盡可能大的互質數。

這些質數乘起來，就是 CRT 的大模數 M：

M = 127 × 113 × 109 × 107 × 103 × 101 × 97
  = 168897325606883

大約 47 bits。

這個 M 定義了我們可以重建多大的整數。直覺上，只要真正的整數結果沒有超過 M/2 的範圍，我們就可以把它正確拼回來。

如果重建結果落在：

M/2 ~ M

那它其實代表負數，要折回來：

if (x > M / 2) x -= M;

所以 CRT 這一步不是只是在列質數。

它同時定義了三件事：

1. INT8 能裝哪些餘數
2. 這組質數乘起來可以表示多大的整數空間
3. 後面重建時每個餘數要乘哪個係數

程式裡會把這些表放起來，例如：

d_primes
d_M_arr
d_coeffs_flat
d_coeff_offsets

為了速度和簡化，Phase 22 fast path 固定使用這 7 個質數。

如果走 dynamic 的狀態，則可以根據 tile 的數值範圍決定要用幾個質數。概念上是估：

threshold ≈ K × max|A_slice| × max|B_slice|

然後找一個足夠大的前綴乘積 M_nl。

但教學上先不用把這件事想太複雜。今天只需要記得兩點：

餘數要對每個 p_i 算
拼回來時要乘上對應的 f_i

如果忘了 f_i 是什麼，可以回看 Day 5 怎麼用 3, 5, 7 把 23 重建回來。那個例子裡每個積木的角色，跟這裡的 f_i 是同一件事。

第二步：求餘數矩陣

有了 CRT 規格表，下一步才是把 FP64 轉成 INT8 世界能處理的資料。

Phase 22 會先把 double 拆成 high / low 的整數薄片：

hi = round(v * 2^17)
lo = round((v - hi / 2^17) * 2^34)

注意，這裡的 hi 和 lo 是整數薄片，不是浮點數。

以 v 在 [-1, 1] 附近來看：

round(v × 2^17) 可能接近 ±131072

這已經不是 INT8 裝得下的數字。

所以要再做：

rem = iv % p;
if (rem < 0) rem += p;

取完餘數以後，才會得到可以放進 int8_t 的東西。

對 A 和 B 來說，會產生四組資料：

A_high mod p
A_low  mod p
B_high mod p
B_low  mod p

之所以有四個，是因為：

A × B
= (A_high + A_low) × (B_high + B_low)
= A_high × B_high
 + A_high × B_low
 + A_low  × B_high
 + A_low  × B_low

A、B、high、low，合起來就是 AB 熱滷。

這些餘數矩陣裝的不是原始數值。

例如：

原本整數薄片是 131072
mod 127 之後可能只剩 40

INT8 GEMM 看到的是 40，不是 131072。背後的大整數意義，要等 GEMM 做完，再用 CRT 拼回來。

專案裡後面也會考慮硬體比較適合的 tile size。Phase 22 的 fast path 會把 A 整理成 128×32，B 整理成 64×32，讓後面的 C tile 變成 128×64。

但今天先不要管 Tensor Core、PTX 或 GPU layout。

先用 Python 或 C++，正常用質數把這件事重建出來：

double
→ high / low integer
→ mod p_i
→ 用 CRT 拼回 integer
→ scaling 回 double

數學先算對，確定能重現，再轉移到下一步。

第三步：GEMM 後重建

求餘數只是準備資料。

真正的 GEMM 會對每個質數各做一次：

C_p = A_p × B_p

因為 modular arithmetic 對加法、乘法都成立，所以：

(A × B) mod p = (A mod p) × (B mod p) mod p

也就是說，我們可以先在每個小質數的世界裡做 GEMM，再把 7 個世界的答案拼回大整數。

重建時做的是：

x = Σ (C_p[i,j] mod p) × f_p  mod M

接著把 [0, M) 折成有號整數：

if (x > M / 2) x -= M;

最後才做 scaling：

C[i,j] += x × inv

inv 取決於現在是哪一個 high / low pass：

high × high  →  2^-34
high × low   →  2^-51
low  × high  →  2^-51
low  × low   →  2^-68

這一點很重要：

重建出來的東西不是最後答案
它是一個 Ozaki 薄片乘積
還要乘上對應的 scale 才會回到 double 尺度

所以不要在 INT8 GEMM 做完的瞬間就把它當成 C。

它只是拼圖的一片。

Fused 版本差在哪

如果我們照剛剛的流程分開寫，會慢在哪？

答案是 global memory。

分開寫很容易變成：

precompute_modulo 寫出 A8 / B8
INT8 GEMM 讀 A8 / B8，寫出 C_p
CRT reconstruct 讀 C_p，寫回 C

這樣每一步都在搬大矩陣。

7 個質數已經代表 7 份餘數。A/B 又有 high/low。再加上四種 pass，如果每個中間結果都落地，global memory 的 overhead 會非常明顯。

所以後續 fused 版本會盡量把事情留在比較近的地方：

constant memory：放質數、M、係數
shared memory：放 tile
register / accumulator：放 GEMM 中間結果

Phase 22 fast path 做的事情就是：

1. precompute_modulo_hl_A/B 同時處理 high / low 和 7 個質數。
2. crt_fast_kernel 在 INT8 GEMM accumulator 裡直接做 CRT，不把每個 C_p 都寫回 global memory。

後面如果要加速，就會往這個方向走：

常數和質數放 constant
high / low 一起 precompute
7 個質數的結果不要全部落地
直接在 MMA accumulator 上拼起來

但這是優化。

理解上，今天先抓住獨立流程：

準備質數
必要時寫小 kernel 驗證 CRT
求餘數
GEMM
重建

Day 7、Day 8、Day 9 會再一段一段把它往實作靠近。

今天先停在這裡

今天重點先掌握三句：

1. 先定義極限與範圍。 質數決定 M，M 決定能重建多大的整數，也決定餘數要怎麼拼回來。這也是質數在 int8_t 能容許的範圍內越大越好的原因。
2. 餘數矩陣裝的是餘數，不是原始數值。 INT8 GEMM 看到的是 x mod p，不是原本那個 high / low 整數薄片。
3. 重建出來的東西還不是最後答案。 它只是某個 high / low 組合的整數乘積，必須經過 CRT 和 scaling，才會回到我們要的 double 尺度。

明天先專心做第一段：把 7 個質數、M、逆元、f_i 這張 CRT 表算清楚，再用一個最小例子確定真的拼得回來。

註解1
先補一個誠實聲明。

寫這種東西要為自己的產出負責。尤其是「FP64」、「double」、「高精度」這幾個詞，如果沒有講清楚，很容易讓人以為我們現在已經做出可以替代 cublasDgemm 的東西。

不是。

你從下面的數學其實也看得出來：Phase 22 的有效精度只有大約 34-bit。

程式介面吃的是 double，也就是 IEEE 754 binary64，通常大家直接叫 FP64。它不是 64 個有效小數位。它大概長這樣：

1 bit   sign
11 bits exponent
52 bits fraction
+ 1 個隱含 leading bit
= 53 bits significand precision

所以 full FP64 的有效精度大約是 53 bits。

Phase 22 不是這個意思。

Phase 22 是：

input  是 double
output 是 double
中間用 17-bit high + 17-bit low 的 Ozaki/CRT fast path

所以它比較像：

FP64 I/O
約 34-bit mantissa 的近似計算

不是一個正式的 fp34 型別，也不是「支援完整 FP64 精度」。

我們目前主線 Phase 22 不是 FP64 precision support。它吃 double、吐 double，但 fast path 大約只有 34-bit mantissa 的有效精度；在 2048³、U(-1,1) 測試上，max error 約 1e-9。

這比 FP32 好，但不是 cublasDgemm 的 ulp，也不是 1e-15。比較準確的說法是：這是一個相對簡化的 Ozaki/CRT 版本，用來證明「把高精度問題搬到 INT8 Tensor Core」這個方向是可行的，但它還不是完整 FP64 GEMM。

那原理能不能延伸到更高精度？

可以，原理是相通的。

當然不是「多換幾個數字就結束」。你可以切更多 Ozaki slices，也可以用更多 CRT 質數，或改成 Scheme I 那種 8-bit slices。可是代價也會一起上來：

Ozaki slice 變多 → GEMM pass 以 s² 成長
CRT 質數變多    → 每個 pass 要做更多模數下的 GEMM / 重建
M 變大          → uint64 / overflow / register pressure 都要重新檢查

數學上越高精度靠近，工程上要付出更多 pass、更多 memory traffic、更多暫存器壓力。Phase 22 是這次實作中拿出來講的一個部分 也很有噱頭跟篇幅的考量吧 但這幾天的例子 就麻煩大家將就了 (汗顏)
我會在第八或第九天的時候講到 轉換成更高位數的案例跟情況 (應該)