矩陣乘法要算幾次？AI晶片的運算量與運算強度

生活中的例子

AI晶片的規格表上，通常會列出2種數字。1秒能計算幾次，以及1秒能從記憶體搬運幾位元組。例如NVIDIA的H100（2022年發表，SXM版），16位元的矩陣計算每秒約1,000兆次。能搬運的量是每秒3.35TB。

那麼，手上AI的速度，是由哪個數字決定的呢？答案是「看計算的內容」。這篇文章會計算矩陣乘法的工作量和搬運的資料量，分辨是哪一邊決定了速度。

GPU用大量計算電路同時進行矩陣計算，已在「AI為什麼要用GPU？」中說明過。這篇文章則聚焦在更進一步的「計算數量」。

原理1：向量和矩陣是「數字的表」

把數字排成1排的東西，稱為向量。某個學生國文、數學、英文的分數（80, 60, 70），就是3個數字的向量。在文字AI裡，1個詞會用數千個數字的向量來表示。

把數字排成縱橫的表，就是矩陣。有m列（橫排）、k行（直排）的矩陣，寫成「m×k的矩陣」。AI的權重（表示要多看重哪個輸入的數字），也是以矩陣的形式存放。

矩陣乘上向量時，會對表的每一列做1次「和輸入相乘再相加」，得到和列數一樣多的答案。這就是矩陣×向量的乘積。從1個學生的分數，算出理組和文組2個分數的計算，就屬於這種。

把100人份的分數一起計算，輸入也會變成表（3×100），就成了矩陣×矩陣的乘積。特色是同一張權重表，可以重複用在100人份上。聊天AI一個詞一個詞產生答案的計算，只有1人份時接近前者，把很多人份合在一起時則接近後者。

原理2：計算次數的算法（2mnk）

m×k的矩陣A乘上k×n的矩陣B，答案C是m×n的矩陣。A的行數和B的列數，必須同樣是k才能相乘。答案的1格，是從A的1列和B的1行，把k組數字相乘再相加做出來的。

1次乘法加1次加法的組合，稱為乘積累加運算。也取英文字首寫成MAC。答案有m×n格，每1格需要k次乘積累加，所以總共是m×n×k次。

在計算機的世界，常把乘法和加法分別算成1次，這時就是2×m×n×k次。這種算法的1次，稱為FLOP（1次浮點數計算）。NVIDIA的效能指南，也用這個2mnk來計算矩陣乘法。

用小例子確認看看。2×3和3×2的矩陣相乘，答案有4格，每格3次乘積累加，所以是12次乘積累加，也就是24 FLOP。用Python寫3層的重複（迴圈）來數，乘法是12次，加法是8次。加法比較少，是因為每格第一個乘積，不需要加到任何東西上。

這個差距只有答案的格數（m×n）。在k有數千的實際計算中，差距不到1,000分之1，所以用2mnk來算成了標準。縱橫都是4,096的正方形矩陣互乘，是2×4,096×4,096×4,096，約1,370億FLOP。

1秒能做幾次FLOP，就是FLOPS，用來表示晶片的速度。加上代表兆的「tera」，寫成TFLOPS。整數的計算以「運算」來算，1秒幾兆次寫成TOPS。

原理3：搬運資料量的算法

計算電路要工作，必須從記憶體讀出A和B的數字，再把答案C寫回記憶體。讀寫的數字有m×k＋k×n＋m×n個。乘上1個數字的位元組數，就是搬運的資料量。

1個數字的大小，由數字的格式決定。AI常用的16位元浮點數（FP16）是2位元組，8位元整數（INT8）是1位元組。格式的差異，整理在「FP32、FP16、BF16、INT8的差異」中。

縱橫4,096的例子用FP16來算，3個矩陣約有5,000萬個數字，約1億位元組（約100MB）。這裡算的是每個數字只要搬一次的理想情況。後面會看到，實際搬運的量會因程式的寫法而變成好幾倍。

原理4：運算強度是「1位元組計算幾次」

把計算次數除以搬運的位元組數，得到的值稱為運算強度。它表示搬來的1位元組能計算幾次，單位是FLOP/位元組。

矩陣×向量的乘積，運算強度只有約1。因為FP16的1個權重（2位元組），只用在1次乘積累加（2 FLOP），之後就沒有用處了。NVIDIA的指南也說明，矩陣×向量的乘積運算強度小於1，速度總是由記憶體決定。

矩陣×矩陣的情況就不同了。在縱橫4,096的例子中，約1,370億FLOP除以約1億位元組，運算強度約1,365。因為搬來的每一個數字，都會用在答案的1列或1行上，被使用好幾千次。運算強度可以想成是「搬來的數字會重複使用幾次」的指標。

原理5：Roofline：2個上限中較低的那個

晶片的速度有2個上限。1個是計算電路1秒能做的計算次數。另1個是1秒能搬運的位元組數（記憶體頻寬）乘上運算強度。數字送不到就無法計算，所以實際的速度，不會超過2者中較低的那個。

以橫軸為運算強度、縱軸為計算速度，畫出這個關係的圖，稱為Roofline（屋頂線）。這是2009年Williams等人提出、用來估計效能的圖。左邊的斜線是搬運速度的上限，右邊的平線是計算的上限。

2條線交點的運算強度，可以用計算的上限除以記憶體頻寬求得。H100是約989兆÷3.35兆，約300。運算強度比這個小的工作，計算電路等記憶體的時間會變長。這種狀態稱為記憶體受限，相反地，計算電路不停工作的狀態稱為運算受限。

一次1人份的矩陣×向量（運算強度約1），位在交點左邊很遠的地方，只能用到H100計算能力的約0.3%。大型矩陣乘法（約1,365）在右側，能以計算的上限運作。

原理6：批次：把同一份權重用在好幾人份上

提高運算強度最直接的方法，是把輸入合在一起。一起處理的輸入數量，稱為批次大小。權重表只要讀一次，就能用在合併起來的所有人份的計算上。

以縱橫4,096的權重表（FP16約3,400萬位元組）來說，批次設為B時，在B還小的範圍內，運算強度也大約是B。因為計算次數變成B倍，搬運的量卻幾乎還是權重那一份。要達到H100交點的約300，需要合併數百人份（見後面「從數字來看」的表格）。

原理7：分塊：在晶片內的小記憶體裡重複使用

到目前為止，都假設每個數字只搬一次。但是，計算電路旁邊的高速記憶體（晶片內記憶體）比較小。它比HBM等外部記憶體小得多，大的矩陣無法整個放進去。

如果不下任何功夫，每做出答案的1格，就要從外部記憶體重新讀一次A的1列和B的1行。這樣，搬運的量會增加到和乘積累加次數差不多，運算強度會降到和矩陣×向量一樣的1左右。

所以，要把矩陣切成小的正方形區塊。把A的區塊和B的區塊讀進晶片內，做完這一組能做的所有計算後，再換下一個區塊。這種方法稱為分塊（tiling）。區塊的1邊設為b個時，從外部記憶體搬運的量，會減少到不下功夫時的約b分之1。

根據NVIDIA的指南，GPU會把答案矩陣分成區塊來計算。張量核心這類專用電路，也會用1個指令，一次算完小區塊之間的乘積。

原理8：影像的卷積也能變成矩陣乘法

影像AI（CNN）常用的卷積，是一邊移動小窗口（例如3×3像素），一邊把窗口裡的數字乘上權重再相加的計算。把窗口內容排成1列，依窗口位置的數量縱向堆疊，就會得到一個大矩陣。再乘上權重矩陣，卷積就變成了矩陣乘法。這種重新排列稱為im2col。

不過，窗口會互相重疊，所以重新排列後，同一個數字會被複製好幾次。3×3的窗口每次移動1像素的話，1個數字最多出現9次，記憶體也用得比較多。所以，也有只在晶片內一點一點做出重排矩陣的方法（cuDNN 2014年的論文）。每一層的乘積累加次數，可以用CNN輸出尺寸計算工具試算。

原理9：量化：把數字變短，減少搬運量

把浮點數換成位元數較少的整數等，減少1個數字的大小，在AI領域稱為量化。FP16換成INT8，搬運的位元組數會減半；從32位元浮點數（FP32）換的話，則是4分之1。

在記憶體受限的工作中，搬運量減半，時間也幾乎減半。1人份的矩陣×向量，換成INT8後，運算強度會提高到約2。運算受限的工作也有好處。在H100的規格中，INT8的計算上限是16位元的2倍。不過，數字太粗的話答案會偏掉，所以要一邊確認精度一邊使用。

原理10：脈動陣列：把重複使用做成電路的形狀

把乘積累加的電路排成像棋盤的格子，讓數字依序傳給旁邊的電路來計算矩陣乘法，這種結構稱為脈動陣列。它是1970年代末，由H. T. Kung和C. E. Leiserson提出的。資料從記憶體出來，在電路之間規律地流動，再回到記憶體。Kung在1982年的解說中，把這種流動比喻成進出心臟的血液。

Google的TPU（初代）就是用這種方式計算矩陣乘法。權重一直放在各個電路裡，讓輸入的數字從左往右流，部分和從上往下流。從記憶體讀出的1個數字，會被橫向排列的256個電路使用，所以讀記憶體的次數大幅減少。

初代TPU以256×256＝65,536個8位元乘積累加電路，每秒最多能做92兆次運算（2017年的論文）。這是把提高運算強度的技巧，直接做進電路排列方式本身的設計。

常見的誤解

從數字來看：改變批次會發生什麼事

用H100規格的值，估算縱橫4,096的權重表（FP16）乘上批次數量輸入的計算。這是假設計算和搬運同時進行、所花時間取較長那一邊的理想值。「計算」「搬運」分別是各自所需的時間（µs是100萬分之1秒），「使用率」是計算電路在工作的時間比例。

從批次1到256，所花時間幾乎維持在10～11 µs，合併多少就賺多少。到了512，計算時間超過搬運時間，之後批次加倍，時間也會加倍。每1人的處理速度不再提高，只有1次的等待時間變長。

進一步學習

• 名詞的詳細說明：運算強度 / Roofline模型 / 脈動陣列 / 訓練後量化 / 混合精度
• 用表格看差異：訓練與推論的差異 / CPU、GPU、NPU的差異 / DDR、LPDDR、GDDR、HBM的差異
• 用數字試試看：記憶體頻寬計算 / CNN的輸出尺寸、乘積累加次數
• 用圖來看：圖解一覽（Roofline模型、GPU執行緒的階層）
• 從基礎讀起：「AI為什麼要用GPU？」 / 「CPU是什麼？」