## Optimizing Large Language Model Inference on RISC-V Platforms Using TVM Compiler 資工四 莊翔鈞 資工四 洪子傑 指導教授:陳鵬升 note: 我們的題目是在嵌入式設備上運行大語言模型。因為嵌入式設備的算力有限，直接把大語言模型丟上去跑會很慢，所以我們用TVM BYOC 技術來優化模型，最後總體運行速度提升9.6倍。 --- ## Apache TVM  note: 為了讓大語言模型能在資源受限的嵌入式設備上運行，我們重寫了Whisper-tiny模型的矩陣乘法運算子，並利用TVM compiler中的BYOC (Bring Your Own Codegen) 技術，將優化後的運算子整合至模型的計算圖中，大幅提升計算效率。優化過的模型與原版相比，運行速度提升9.6倍。 --- ## Goal  --- ## Banana pi f3  --- ## TVM 編譯流程  note: 那接下來我來介紹一下我們用到的技術。首先是TVM，他是一個深度學習的編譯框架，可以把多種前端編譯到一個特定的設備上，然後BYOC允許我們在這個流程中添加自己設計的運算子，來替換AI模型原有的運算子，以此來優化模型。 ---  note: 那我們在實驗的一開始先用Netro分析模型，發現矩陣乘法佔了95%的執行時間，所以我們打算改寫這個運算子。 --- ## 如何加速矩陣乘法？ 1. 使用 RVV 指令集 2. 矩陣分塊 ---- ## RISC-V RVV 指令 1. 一次處理多筆資料（SIMD） 2. 向量長度可變（Vector-Length Agnostic） 3. 提升諸如矩陣乘法、卷積、影像處理等資料密集型運算效能。 ----  ----  ---- ## 矩陣分塊 首先，假設我們有兩個邊長為n=4的矩陣 $X = \begin{bmatrix} a_1 & a_2 & b_1 & b_2 \\ a_3 & a_4 & b_3 & b_4 \\ c_1 & c_2 & d_1 & d_2 \\ c_3 & c_4 & d_3 & d_4 \end{bmatrix}$ $Y = \begin{bmatrix} e_1 & e_2 & f_1 & f_2 \\ e_3 & e_4 & f_3 & f_4 \\ g_1 & g_2 & h_1 & h_2 \\ g_3 & g_4 & h_3 & h_4 \end{bmatrix}.$ ---- 將scalar如 $X=\begin{bmatrix} a_1 & a_2 & | & b_1 & b_2 \\ a_3 & a_4 & | & b_3 & b_4 \\ -- & -- & + & -- & -- \\ c_1 & c_2 & | & d_1 & d_2 \\ c_3 & c_4 & | & d_3 & d_4 \end{bmatrix}$ 分塊切割成邊長為n/2的方形子矩陣可以得到 $X=\left[\begin{array}{ll} A & B \\ C & D \end{array}\right], Y=\left[\begin{array}{ll} E & F \\ G & H \end{array}\right],A = \begin{bmatrix} a_1 & a_2 \\ a_3 & a_4 \end{bmatrix}$ ---- $X=\left[\begin{array}{ll} A & B \\ C & D \end{array}\right], Y=\left[\begin{array}{ll} E & F \\ G & H \end{array}\right]$ 那麼矩陣乘法就可以變成 $X Y=Z=\left[\begin{array}{ll} A E+B G & A F+B H \\ C E+D G & C F+D H \end{array}\right]$ ---- ## Proof $\begin{gathered} Z_{1,1}=\sum_{k=1}^n X_{1, k} Y_{k, 1} \\ =\sum_{k=1}^{n / 2} X_{1, k} Y_{k, 1}+\sum_{k=n / 2+1}^n X_{1, k} Y_{k, 1}\\= \sum_{k=1}^{n / 2} A_{1, k} E_{k, 1}+\sum_{k=1}^{n / 2} B_{1, k} G_{k, 1}\\=(A E)_{1,1}+(B G)_{1,1} \end{gathered}$ ---- ## 為何使用切割矩陣乘法? ---- ## Ans: 減少 cache miss ---- ## 範例 ``` void classic_matmul(double A[N][N], double B[N][N], double C[N][N]) { for (int i = 0; i < N; i++) { for (int j = 0; j < N; j++) { for (int k = 0; k < N; k++) { C[i][j] += A[i][k] * B[k][j]; // A[i][k] → row-wise 存取，連續，快 // C[i][j] → row-wise 存取，內層 j 連續，快 // B[k][j] → column-wise 存取，內層 k 不連續 → cache miss } } } } ``` ---- $X=\begin{bmatrix} a_1 & a_2 & b_1 & b_2 \\ a_3 & a_4 & b_3 & b_4 \\ c_1 & c_2 & d_1 & d_2 \\ c_3 & c_4 & d_3 & d_4 \end{bmatrix},Y = \begin{bmatrix} e_1 & e_2 & f_1 & f_2 \\ e_3 & e_4 & f_3 & f_4 \\ g_1 & g_2 & h_1 & h_2 \\ g_3 & g_4 & h_3 & h_4 \end{bmatrix}. $ ---- ```cpp! #define N 2048 #define BLOCK_SIZE 64 void matmul_blockwise(double* A, double* B, double* C) { // 分塊矩陣乘法 for (int i = 0; i < N; i += BLOCK_SIZE) { for (int j = 0; j < N; j += BLOCK_SIZE) { for (int k = 0; k < N; k += BLOCK_SIZE) { // 處理一個塊 for (int ii = i; ii < i + BLOCK_SIZE && ii < N; ii++) { for (int jj = j; jj < j + BLOCK_SIZE && jj < N; jj++) { for (int kk = k; kk < k + BLOCK_SIZE && kk < N; kk++) { C[ii * N + jj] += A[ii * N + kk] * B[kk * N + jj]; } } } } } } } ``` ---- $X=\begin{bmatrix} a_1 & a_2 & | & b_1 & b_2 \\ a_3 & a_4 & | & b_3 & b_4 \\ -- & -- & + & -- & -- \\ c_1 & c_2 & | & d_1 & d_2 \\ c_3 & c_4 & | & d_3 & d_4 \end{bmatrix}$ $Y = \begin{bmatrix} e_1 & e_2 & | & f_1 & f_2 \\ e_3 & e_4 & | & f_3 & f_4 \\ -- & -- & + & -- & -- \\ g_1 & g_2 & | & h_1 & h_2 \\ g_3 & g_4 & | & h_3 & h_4 \end{bmatrix}.$ note: block 內 B 仍然 column-wise 存取 → 仍有些 cache miss 但因為每次處理的 block 很小，cache line 可以被重複利用 → miss 數量大幅降低 若 block 太大 → 無法完全放入 cache → miss 仍多 ---- ## 小規模測試  ---- ## 經典 vs 向量化 內積 ``` function dot_product_scalar(A[0..N-1], B[0..N-1]): dot = 0 for i = 0 to N-1: dot = dot + (A[i] * B[i]) return dot ``` ``` function dot_product_vector(A[0..N-1], B[0..N-1]): dot = 0 VL = vector_length # 硬體一次可處理的元素數 for i = 0 to N-1 step VL: # 每次都是 i+=VL chunkA = A[i : i+VL-1] # 取出一個 vector chunk chunkB = B[i : i+VL-1] # 取出對應元素 vtmp = vector_mul(chunkA, chunkB) # 每個元素對應相乘 (VL 個同時做) dot = dot + vector_reduce_sum(vtmp) # 將 vector 內元素加總到 scalar return dot ``` note: 一次 vector 乘法處理 VL 個元素 reduce_sum 將這 VL 個元素累加成單一標量 ---- | 項目 | 傳統 for loop | RVV 向量化 | | ----- | -------------------- | ----------------------------------- | | 單次運算量 | 1 個元素 | VL 個元素（如 8 或 16） | | 指令數 | N 乘法 + N 加法 | N/VL 次向量乘法 + N/VL 次向量累加（每次 VL 元素並行） | | 效率 | 1 個元素/時鐘週期 | VL 個元素/時鐘週期 → 潛在加速 8×、16× | note: Banana Pi BPI-F3 SpacemiT K1 RISC-V chip is of vector 256bit and float32 is 32 bit => 8 scalar in one insturction --- ## 數據分析  note: 由以上數據可見，無論單獨使用 RVV 版矩陣乘法或開啟 O3 最佳化，都能顯著提升計算速度；而兩者兼用時，效能提升更為突出。 在未開啟 O3 的情況下，經典矩陣乘法（O(n³)）甚至比原版模型更慢，顯示其效率不足；但啟用 O3 後，即便使用此演算法，執行時間仍可降至原本的 66，反映出編譯器在迴圈轉換與記憶體排程上的優化效果。 若採用 RVV 向量化矩陣乘法，即使未啟用 O3，執行時間也降至原版的 42%，顯示向量化能有效提升吞吐量。 當 RVV 與 O3 同時啟用時，效能提升最為明顯——Encoder 計算時間由 1390 秒降至 69 秒（約 20 倍），Decoder Prefill 由 62 秒降至 5.3 秒（約 11.7 倍），整體速度提升達 9.6 倍，充分展現向量化與編譯器優化的協同效益。 --- $Speedup_{overall}=9.6$ --- ## 延伸討論 1. 最佳化矩陣分塊大小？ 2. 為何不用strassen演算法 $O(n^{2.8})$？ ---- ## Future work 1. decoder_with_past.onnx (dynamic shape) 2. 回饋開源社區 3. 拓展到其他模型 ---- ## Future work  --- # Thanks  --- pic1 --- pic2