# HW4: FlashAttention

[HackMD link](https://hackmd.io/@elliehung/HJKdAbNrJl)

| CS542200 Parallel Programming | HW4: FlashAttention | 112065528 | 洪巧雲 |
| ----------------------------- | ---------------------------- | --------- | ------ |

## Implementation
- GPU memory hierarchy
- GPU SRAM --> shared memory
- HBM --> Global memory

(a) Describe how you implemented the FlashAttention forward pass using CUDA. Mention the algorithm's key steps, such as matrix blocking, SRAM usage, and how intermediate results like scaling factors (ℓ 和 𝑚) were calculated.
- 這份作業中，我實作的部分是將 Q、K、V 分塊，並在 GPU 的 shared memory (SRAM) 中進行計算的部分。
- Step 1: Read data
- Read `Q`、`K`、`V` 三個矩陣，並分配到 GPU device 上。
- Step 2: 分塊
- 將 `Q`、`K`、`V` 三個矩陣照 `Br = 16`、`Bc = 16` 的方式分塊如下 (以 `testcases/t01` 舉例之: `N = 128`、`d = 32`):


- 分塊後的矩陣 size 、計算過程如下:


- 每個 Cuda Block 會需要處理 `Br` 筆的 data，並且在 for loop 中需要遍歷所有 `Bc` 欄位的 blocks。
- Step 3:
- 將當前要處理的 Q 之資料載入至 shared memory，以減少反覆讀取 global memory。
- SRAM / Shared Memory 的宣告:


- Step 4:
- Compute : S_ij = Q_i · K_j^T
- shared memory 中進行 dot product 運算乘上 softmaxScale
- Step 5: (ℓ 和 𝑚) 的計算 (Intermediate Results)
- 先計算目前區塊上的最大值 localMax，再計算所有分數的 exponent 和 localSum
- 宣告與初始變數 `maxVal`、`sumExp`
- 計算 dot product 及局部最大值 (localMax)
- 先對 `Q_i`、`K_j` 做點積乘上 `softmaxScale = 1 / sqrt(dim)`。


- 計算 softmax → exponent & sumExp


- 合併新舊最大值 (maxVal) 與 sumExp


- 更新 partial output:


(b) Explain how matrices Q, K, and V are divided into blocks and processed in parallel.
- `Q` 的分法: block 大小為 `Br×1`，將整個 `Q` 分成大小為 `Br` 的塊，共有`totalRowTiles = (N + Br - 1) / Br` 個 row tiles。
- `K`、`V` 的分法: 由於每個 block 都會迭代所有的 `K`、`V` tiles，所以將 `K`、`V` 依照欄位（col-wise）切成大小 Bc 的 blocks，共有 `totalColTiles = (N + Bc - 1) / Bc` 個 col tiles。
- 這些 blocks/tiles 都能在 GPU 上並行執行，因為在啟動 kernel 時使用了 `gridDim(totalRowTiles)`、`blockDim(Br)`，每個 block 同時都會在自己的共享記憶體中處理一部分 `Q` 的 row 搭配所有的 `K`、`V` blocks。
(c) Describe how you chose the block sizes 𝐵𝑟 and 𝐵𝑐 and why.
- 在程式中我選擇的大小為:`Br = 16`、`Bc = 16`，原因是我選 32 的話會 t12 & t22 這兩筆大測資會 TLE。
(d) Specify the configurations for CUDA kernel launches, such as the number of threads per block, shared memory allocation, and grid dimensions.
- threads per block (blockDim) : `blockDim(Br)` → 也就是 `blockDim.x = 16`。每個 thread block 包含 16 條 thread，對應到要處理 `Q` 的 16 列 (`tile` 大小)。
- grid dimensions (gridDim): `gridDim(totalRowTiles)`，其中 `totalRowTiles = (N + Br - 1) / Br`。每個 block 負責一個 row tile（大小為 `Br` 列），所以總共多少 tile 就多少個 block。
- shared memory allocation:

size_t sharedMemorySize = (Br * (dim + 1) 
                         + Bc * (dim + 1) 
                         + Bc * (dim + 1) 
                         + Br * (dim + 1) 
                         + Br * Bc) 
                         * sizeof(float);

這塊記憶體包括：
`Br * (dim + 1)` 用來存 `Q` tile
`Bc * (dim + 1)` 用來存 `K` tile
`Bc * (dim + 1)` 用來存 `V` tile
`Br * (dim + 1)` 用來存 `O` (partial output)
`Br * Bc` 用來存 `S_ij`
(e) Justify your choices and how they relate to the blocking factors and the SRAM size.
- `16 × 16` 的分塊有不錯的記憶體/計算效能比，又能容納在共享記憶體中，同時能很好地搭配 CUDA warp 大小(通常 32 threads)。
- 將 `Q`、`K`、`V` 以及中間計算結果 (S_ij, partial output) 全都放在 shared memory 裡一次做計算，可以減少對 global memory 的頻繁讀寫。但若選擇太大的 Br, Bc，會直接導致 shared memory 不足或 block 無法排滿 SM。
- 下圖為把 在 HBM 中的 `Q` 載入 SRAM `shared_Q` 的例子:


### Flash Attention V1 vs V2
- Flash Attention V1
由下圖示可以發現其實在 v1 的演算法中，已經是把 Attention 演算法加速很多了，最主要的就是讓 `Q`、`K`、`V` 三個矩陣不需要一直做 HBM 與 SRAM 的讀寫，這是最耗時的，FlashAttention 的核心就是用時間換取空間，然後再用 CUDA 將可以平行化的計算平行化，達到即便是使用重複計算去取代一直需要從 HBM 中讀取東西到 SRAM，也可以透過 GPU 加速整體效能。


- Flash Attention V2
但是 v1 依舊不夠快，在實作程式時發現有些測資會一直 Time limit exceeded，所以才去研究了 v2 版本的演算法，發現 v1 跟 v2 版本的演算法其實只有一部分不一樣，就是 inner for loop 與 outer for loop 的值相反了，如下圖所示:


之所以有這裏不同就造成加速的原因是，v2 真正優化的部分是對於 `Q` 矩陣少了非常多次的讀取，因為 v2 是 Row-wise 的計算，而 v1 則是 column-wise 的計算，這導致 `Q` 矩陣在 v2 版本的演算法中可以在每次做完一整個 Row 得到 row_sum `l` 與 row_MAX `m` 後就不會再需要 load 這個 block 的值了，但是在 v1 版本是 column-wise 的計算順序的話，每換一個 column `Q` 都會需要再把同一個 block 的 value 從 HBM load 進 SRAM 中進行 S value 的計算，如上下我畫的圖所示。還有一點很重要的是，在做完一整個 row 的計算後，v2 才會需要去更新 `l`、`m` 這兩個 value，但是 v1 因為是 column-wise 所以說它的計算方向是每計算完一個 block 後就更新回 `l`、`m` 這兩個值回去 HBM，這也造成了 HBM SRAM 很多多餘的讀寫。最後一點的優化則是在開的 thread 數量也是不同的，v1 中，可以開的 thread 數量為 `Br` 條，v2 中則是可以開 `N` 條 threads，也因為它是一整個 row 去做計算 + 平行化。


下圖藍色螢光筆畫記的部分是指 v1 與 v2 可以平行化的地方:


## Experiment and Result
### Optimization

`Testcases: t21`
| B | N | d |
| --- | ---- | --- |
| 500 | 2048 | 32 |

這邊做了 3 個優化:
- Shared Memory
分配 `sharedQueries`, `sharedKeys`, `sharedValues`, `sharedOutput`, `sharedScores` 這五塊 Shared Memory:


這樣做之所以可以達到優化是因為先把每一個 tile (`Br×dim` 或 `Bc×dim`) 的資料載入到 shared memory，再重複使用，降低 global memory 的存取。
由圖表可見，Baseline 跑約 4.53s，而使用 Shared Memory 後降至 3.97s。
Speed-up 從 1.00 提升到 1.14，改善約 14%。

- Handle Bank Conflict
同樣使用 Shared Memory，但處理「Bank Conflict」
Bank Conflict 是指在 GPU 的 Shared Memory（共享記憶體）中，當多個 threads 同時存取位於同一個 memory bank 內的不同位址時，會導致存取衝突，需要序列化讀寫，進而降低效能，如下圖所示:

例如在圖示的 2-way Bank Conflict 裡，T-0、T-1 這兩條 thread 分別讀取了地址 0 與 1，但這兩個地址落在同一個 bank (bank-0) 裡，導致讀取衝突。
解決方法其實不難，只要在每一行多空一格，讓他們不要讀到同一個位址的 memory 就好。
如下所示，我的程式把 shared memory index 的部分 (dim + 1) 可以使每次的存取都分散到不同的 memory bank 上，這會加速的原因是可以減少 bank conflict 帶來的延遲。:



- Stream
在 `main` 中創建多個 `cudaStream_t stream[i]`，並用 `cudaMemcpyAsync`、`kernel<<<..., stream[...]>>>` 去將多個 batch 的 `Q`、`K`、`V` 讀取和計算重疊(overlap) 執行:



加了 Stream 後優化非常多，因為這樣做可以讓 GPU 在資料傳輸（Host ↔ Device）與計算同時進行。在圖表中可以看到時間進一步降到 1.56s，對比 handle bank conflict 時的 1.82s 再進一步縮短。



### Time Distribution

- NVPROF Profiling:

- Time Distribution 的實驗中，我使用了 7 個不同大小的 testcases 做實驗:

| Testcases | t02 | t04 | t06 | t08 | t10 | t26 | t27 |
| --------- | --- | --- | --- | ---- | ---- |:----:| ----- |
| B | 160 | 80 | 40 | 20 | 10 | 10 | 10 |
| N | 128 | 256 | 512 | 1024 | 2048 | 8192 | 16384 |
| d | 64 | 64 | 32 | 64 | 64 | 64 | 32 |

由實驗結果發現，我的程式 bottleneck 是 compute time，這其實蠻合理的，因為 Flash Attention 的做法就是利用重複計算的原則去大大減少資料在 HBM 與 SRAM 之間傳遞的耗時風險，這樣看似以時間換取空間的方法，卻可以利用 GPU 平行加速的特性不只省到空間還能加速計算。但在 `N` 變到非常大的 size，則會發現 compute time 增長的幅度非常快，畢竟 GPU 有 memory bound，所以導致後面 compute time 迅速增加。



### Blocking Factor
- NVPROF Profiling:


- Blocking Factor 的實驗中，我使用了 `Testcase: t18` 做實驗，block 大小測試了四個 4, 8, 16 跟 32。由實驗結果發現 block size = 16 竟然是跑最快的，這也代表了 block size 並不是開越大越快


- Integer GOPS: 當 block size = 16 時，GOPS 為 0.19，是最高的，代表同時間內完成了更多計算操作。`block size = 4` 只有 0.13 GOPS，`block size = 8` 為 0.16，`block size = 32` 下降至 0.15，都不及 16 來得高。


- Global Store Throughput: block size = 16 時可達 1.05 GB/s，也是測試中最高的 Store 吞吐量。表示在寫入結果（或中間資料）到 global memory 時，block size = 16 的配置最能有效利用裝置頻寬與 SM（Streaming Multiprocessor）資源。


- Global Load Throughput: 雖然從圖中可以看到 block size = 4 時 Load Throughput 非常高（356.25 GB/s），但這通常不代表計算效率高，有時反而表示程式頻繁讀取記憶體但計算不夠密集，會導致總體時間拖長。block size = 16 的 Load Throughput 為 135.00 GB/s，看似比 4 更低，但透過**更好的計算密度（compute intensity）** 與 blocking 配合，就能更有效率地利用已載入的資料，不需要過度頻繁地來回抓資料。


- sm efficiency: block size = 4 時，SM Efficiency 幾乎到 99.88%，8 時有 96.93%，16 時降到 94.09%，32 時剩 79.64%。SM Efficiency 可解釋為 SM 在執行 thread 時的「active cycles」占「總 cycles」的比例，或是 warp 執行效率的度量。這裡 block size=4 近乎 100% 效率，意即當 block size 很小，SM 基本上都在忙著跑某些 warp，但實際上計算量(或單次工作量)不足，可能並未帶來最終整體吞吐量的上升。當 block size 變大（32），SM Efficiency 下降到約 80%，表示每個 SM 可能在更多時間等待記憶體、或因為資源不足無法 fully schedule warp，導致 SM 不是那麼忙碌。block size=16 雖然 Efficiency 只有 94.09%，但執行時間最短，代表說「雖然沒達到最高使用率，卻在計算/記憶體之間取得了更好平衡」。


- Occupancy: 圖中顯示：block size = 4 時，有最高的 19.99% Occupancy，之後 block size 越大，Occupancy 越低，直到 32 時只有 2.56%。
- Occupancy 代表每個 SM（Streaming Multiprocessor）上可以同時排程、運行的 thread 比例。如果程式使用的暫存器與 shared memory 太多，會使得一次只能排很少的 block、warp，同時執行的 thread 數量就比較少。當 block size 小（例如 4），每個 block 占用資源很少，SM 可以一次排很多個 block，因此占用率（occupancy）高。反之，block size 大（例如 32），每個 block 需要更多資源（暫存器或 shared memory 等），SM 可能一次只能排少數幾個 block，甚至只有一個 block，同時執行的執行緒就更少，occupancy 因此下降。
- 雖然 block size=4 有快 20% 的 occupancy，但該設定在「execution time」上實際表現反而比較差，原因可能是單個 block 做的工作太小，需要更多輪迴調度（kernel launch 或 block 切換）才能處理完整個資料集，也可能導致對記憶體的取用太過頻繁，造成瓶頸。


- Shared Load Throughput: 當 block size=4 時，Shared Load Throughput 高達 2850.70 GB/s，之後隨 block size 增加而下降 (8 → 2019.20 GB/s, 16 → 1433.80 GB/s, 32 → 1184.70 GB/s)。block size 小（4）時，可能因為反覆進行讀取(許多 tile、kernel 週期)，在統計上呈現很高的 load 次數，但實際計算效率卻不一定好。block size 大（32）時，雖然 Load Throughput 看起來相對小，但可能每次讀取的效率也未必好，並且 block 數量也少。


- Shared Store Throughput: block size=4 時高達 556.95 GB/s，後續 8→208.54 GB/s、16→94.71 GB/s、32→78.28 GB/s 越來越低。


- Blocking factor = 64: 會超過 memory 48 KB，這是因為在 kernel 中我使用了 5 個 shared memory 的關係。

### Stream Factor
- 除了 block factor 以外，為了可以找到最好的 Stream Factor，我也做了 stream 數量的實驗，實驗結果是當 stream 數量為 4 的時候執行時間最短。
`Testcase: t18`
- NVPROF Profiling:

- Result









# Conclusion

透過本次作業的實作與實驗，我觀察到 FlashAttention 演算法在 GPU 上能透過分塊 (blocking) 與共享記憶體 (shared memory) 的使用，大幅減少對全域記憶體 (HBM) 的頻繁存取，進而提升效能。結果顯示，Block size 並非越大越好，像在多組測試中，16×16 的分塊配置最能兼顧記憶體頻寬與計算密度，取得最短執行時間。此外，配合處理 bank conflict 的技巧（於每行插入一格 padding），可進一步降低 shared memory 存取的序列化。最後，利用多個 CUDA Streams 重疊主機與裝置間的資料傳輸並行度，使 Kernel 計算與記憶體拷貝同時進行，效能再獲提升。這次的作業證明了在 GPU 上，需同時考量分塊大小、shared memory 排布、以及指令與資料傳輸的並行化，才能在不同維度的序列長度與向量維度下達到平衡，最終有效縮短運行時間並充分發揮 GPU 資源的潛能。
其實這份作業要優化的難度很高，因為相對於 hw3，這次作業涉及到的計算相對複雜，而且 shared memory 數量也要開很多，這在優化上有一定難度。FlashAttention 演算法是很新的論文，覺得在這次作業最難的地方是要先讀懂 flash attention 優化的步驟，接下來則是難在實作 v1 的話 testcases 還是沒辦法全過，此時需要去看更近期發表的 FlashAttention v2，v2 跟 v1 要理解其中的差異以及優化的地方也需要一段時間。但在這次作業學到很多，也發現更深入了解了 GPU 優化的關鍵。