MIT 6.5940 Fall 2026 Lab 1 補充：用 Roofline、Profiler 與 FlashAttention 看懂 GPU 瓶頸

🌏 English version

這篇是 Fall 2026 的材料。 MIT 6.5940 導讀系列以 Fall 2024 為主幹，這是第 17 篇，也是唯一一篇以 Fall 2026 材料為主的補充。

為什麼收錄：Fall 2024 沒有任何 GPU profiling 的 lab，第 13 講卻一再用「decode 被頻寬卡住」「FlashAttention 少搬資料」當論據。Fall 2026 的 Lab 1 剛好讓你親手量出這些事，放在 Lab 4＋Lab 5 後面，當作 LLM 推論這一段的收尾練習。

系列位置：上一篇 Lab 4＋Lab 5：AWQ 量化與筆電上的 LLaMA2-7B｜下一篇 L14 LLM 後訓練｜系列總覽

官方材料：Fall 2026 課程頁第 4 講那一列的 lab1_gpu_basics.zip，壓縮檔內的檔案日期是 2026-09-29。課程頁排程上它在 9 月 22 日發布，10 月 1 日（第 7 講）截止。以下題號與配分照壓縮檔裡的 README 與 notebook，2026-09-30 核對。

存取等級 A2（學期進行中）：壓縮檔公開可下載，但繳交走 MIT 的 Canvas，沒有公開解答，課程頁也寫明這學期不收 cross-registration。notebook 裡有幾個 public testcase 可以自我檢查。這篇不寫解答。

壓縮檔裡有什麼

README 的標題是「Lab1: Efficient AI Fundamentals」，致謝寫的是 Zhijian Liu 的 Z Lab 提供這份 lab。

• lab1(colab).ipynb：給 Colab 用，圖片內嵌。要把整個 lab1_gpu_basics 資料夾上傳到 Google Drive，從那裡打開，先跑 Setup。
• lab1.ipynb：同樣內容，給有 NVIDIA GPU 的本機用。
• utils/：benchmark、GPU 規格表（config.py 裡有 T4、L4、A100-40GB、A6000、H100 等）、畫 roofline 的函式。
• 兩個 notebook 只要交一個。

版本要鎖：README 寫明 transformers==4.57.6 等版本是必要的，較新的 transformers 改了 utils/config.py 依賴的 API。Colab 預裝的 transformers 太新，不跑 Setup 會在 import 時失敗。

GPU 需求：Part 1–4 在 Colab T4 上可以做。Part 3.3 的 prefill／decode 延遲量測要 A100 或 A6000，但這段沒有題目，用 T4 可以跳過。Part 5 在 Colab 上要把 runtime 版本設成 2025.10、用 A100，才能裝預編譯的 flash-attn，否則要從原始碼編譯，可能花上好幾個小時。

五個 Part 與配分

主題 80 分：

Bonus 20 分：整個 attention module 的 MAC 與 I/O（各 2.5）、decode 階段的 decoder layer profiling（2.5）、CUDA Graph profiling（2.5）、FlashAttention vs SDPA（5）、FlashAttention 分析（5）。

Part 1–2：先學會用三個數字描述一個運算

場景。Part 1 要你用純 Python 寫三層迴圈的矩陣乘法，不准用 NumPy 或 PyTorch 的運算。notebook 說明這是刻意的：迴圈在 CPU 上跑，之後要拿來跟 torch.matmul 比。

三個指標。Latency 量時間；MAC 是一次乘加，1 MAC 等於 2 FLOPs；I/O 是要搬多少 bytes。1.2 要你連續量幾次、不做 warmup，描述看到什麼。這一題在訓練「量測本身就有雜訊」的直覺。

Roofline。Part 2 定義運算強度 = FLOPs ÷ bytes，ridge point = 峰值 FLOPs ÷ 峰值頻寬。左邊的 memory-bound 區，效能 = 頻寬 × 運算強度；右邊的 compute-bound 區，效能 = 峰值 FLOPs。

2.1 用 N = 1024 到 8192 的 FP32 GEMM 畫 roofline，並推導在 A6000 上 N 要多大才會 compute-bound。2.2 換成 GEMV 再畫一次。兩張圖放在一起，你就會看到第 13 講說的「decode 是 GEMV、慢在搬權重」長什麼樣子。

記得把 notebook 預設的 gpu_name = "A6000" 改成你實際用的卡，註解裡有 Colab 的選項。

Part 3：拿一個真實的 decoder layer 來算

案例模型是 gemma-3-270m-it。notebook 先列出一層 decoder 的所有元件：input LayerNorm、self-attention（Q/K norm、QKV projection、rotary embedding、output projection）、兩個 LayerNorm、MLP（gate、up、activation、down）、post-MLP LayerNorm。

它先示範 MLP 的 MAC 與 I/O 怎麼算，再要你照同樣慣例算 attention（3.1、3.2）。notebook 寫明：多做的假設只要清楚寫在註解裡，合理就給滿分。

接著在 roofline 上畫 MLP 的 prefill 與 decode。notebook 自己給出結論：prefill 是 compute-bound，decode 是 memory-bound，而且這不只是 MLP 的特性，是現代 LLM 推論的普遍現象。3.4 的三題就在追問原因：

• 3.4.1：用 query 形狀解釋差異（提示：想想 GEMM 與 GEMV）。
• 3.4.2：decode 時把 batch size 從 1 調到 64，點在 roofline 上怎麼移動。
• 3.4.3：prefill 時把長度從 64 調到 4096，點怎麼移動。

3.4.2 的答案，就是第 13 講說「W8A8 適合批次 serving、單人 decode 要 W4A16」的理由。

Part 4：看見 kernel，再把它們合起來

Profiler。4.1 用 torch.profiler 看一次 torch.matmul 實際呼叫了哪個 CUDA kernel。4.2 換成 MLP，kernel 一下子變很多，並示範匯出 Chrome trace，丟到 Perfetto 看時間軸。

Kernel fusion。notebook 引用 Horace He 的工廠與倉庫比喻：像 torch.cos 這種一元運算，資料從倉庫（記憶體）運到工廠（運算單元）只做一點點事又運回去，時間幾乎都花在運輸。一串這樣的運算，每一步都來回一趟，fusion 就是讓資料留在工廠裡做完。

4.3 讓你親手體會：用 PyTorch 基本運算寫 tanh 近似的 GeLU，跟 torch.nn.functional.gelu(x, approximate="tanh") 比延遲，再用 profiler 解釋差距從哪來。

torch.compile 與 CUDA Graph。4.4 先把你的 GeLU 丟給 torch.compile，看它被合成什麼樣子；再用 max-autotune-no-cudagraphs 編譯 MLP，比較前後的 profile。Bonus 4.4.2 改用 max-autotune（會啟用 CUDA Graph），要在 torch.no_grad() 下才會重播。notebook 也提醒：T4 這類 SM 少的卡，max-autotune 不會調 matmul，主要看得到的是 element-wise 運算被合併。

Part 5：FlashAttention（全部改成 bonus）

notebook 寫明這段改成 bonus，因為標準的 Colab T4 跑不動，不想讓學生在截止前卡在搶 GPU。

• 5.1：用序列長度 512 到 32768 量標準 SDPA 的峰值記憶體，回答記憶體複雜度、在你的 GPU 上多長會出問題、主要的記憶體開銷是什麼、哪種 LLM 情境會很麻煩。
• 5.2：比較 flash_attn_func 與 SDPA 的延遲和記憶體，再用 profiler 看 FlashAttention 怎麼把 attention 合成一個 kernel。

notebook 對 FlashAttention 的說明和第 13 講第 82–83 頁一致：不把 $N \times N$ 的 attention 矩陣寫進 HBM，把 Q、K、V 切塊放進 SRAM，用 online softmax 逐塊計算。它列的效果是記憶體從 $O(N^2)$ 降到 $O(N)$、減少 HBM 存取帶來 2–4 倍加速，而且是精確計算、不是近似。

這份 lab 和 Fall 2024 主幹怎麼對照

差別在於 Fall 2024 的 Lab 4／5 讓你在 CPU 上優化 kernel，這份 lab 讓你在 GPU 上量測和診斷。兩邊合起來，才是「先找到瓶頸，再動手優化」的完整流程。

另外要注意：Fall 2026 因為 Lab 1 換成這份，已經沒有剪枝 lab。想練剪枝只能用 Fall 2024 的 Lab 1。

自學怎麼做

1. 先決定跑在哪。只有 Colab 免費 T4 就做 Part 1–4，Part 3.3 與 Part 5 跳過；有 A100 或 A6000 再做完整版。
2. 跑 Setup 後一定要重啟 runtime，README 特別提醒過，不重啟會 import 失敗。
3. Part 1–2 先用紙筆算。GEMM 與 GEMV 的 MAC、I/O 先手推，再寫程式、對 public testcase。
4. Part 4 一定要匯出一次 trace 打開看。表格只給總和，時間軸才看得到 kernel 之間的空隙。

今晚可以做的一件事：下載壓縮檔，只做 2.1.2，用你手邊 GPU 的峰值算力與頻寬推出 ridge point。之後讀任何「這個運算 memory-bound」的說法，你都能自己驗算。

延伸閱讀

• 同系列：L13 LLM 部署、Lab 4＋Lab 5、L11 TinyEngine 與平行運算
• GPU 與 kernel：CS336 GPU 與 TPU、CS336 Kernel 與 Triton、CMU 11-868 GPU 程式設計
• FlashAttention：CMU 11-868 FlashAttention

參考資料

• MIT 6.5940 Fall 2026 課程頁 — Lab 1 發布與截止日、lab 清單、cross-registration 政策
• lab1_gpu_basics.zip（Fall 2026，Dropbox） — README、notebook、utils；本文所有題號、配分、硬體需求
• MIT 6.5940 Fall 2024 課程頁 — 主幹學期的 lab 清單（Lab 1 為 Pruning）
• Lec13-LLM-Deployment.pdf（Fall 2024） — decode 瓶頸、kernel fusion、FlashAttention 的對照頁
• google/gemma-3-270m-it（Hugging Face） — Part 3 的案例模型
• Horace He, Making Deep Learning Go Brrrr From First Principles — notebook 引用的 kernel fusion 比喻
• Dao-AILab/flash-attention（GitHub） — Part 5 使用的 flash-attn 套件
• Dao et al., FlashAttention（arXiv:2205.14135）
• PyTorch Profiler 文件