使用 Codex 進行自動化研究：實現 232 倍加速的 QR 分解核心 (Kernel)
執行摘要
透過使用 OpenAI 的 Codex 實作自動化研究迴圈，一位開發者在批次化正方形緊湊 Householder QR 分解 (batched square compact-Householder QR factorization) 上，相較於 torch.geqrf 基線實現了 232 倍的加速。這次成功的關鍵在於從通用的函式庫呼叫轉向分塊 Householder 演算法 (blocked Householder algorithm)，利用「候選者束搜尋」(beam of candidates) 來避免陷入局部最大值，並利用 GPU Mode popcorn CLI 與 Modal 提供的緊密回饋迴圈。
優化挑戰：QR 分解
目標是為 FP32 CUDA 矩陣實作批次化正方形緊湊 Householder QR 分解。輸出需要一個緊湊表示法，包含一個 H 矩陣（其中上三角為 R，下三角儲存 Householder 向量）以及一個包含反射係數的 tau 向量。
序列化瓶頸
標準的 Householder QR 本質上是序列化的：反射器 $j+1$ 取決於反射器 $j$ 所產生的矩陣。這種順序依賴性阻礙了 Tensor Cores 的有效使用，因為運算仍維持在矩陣-向量 (matrix-vector) 的形式，而非矩陣-矩陣 (matrix-matrix) 的形式，導致 GPU 最強大的運算單元處於閒置狀態。
分塊 Householder 解決方案
為了克服序列化瓶頸，開發者採用了分塊 Householder 演算法。這種方法將序列化運算限制在一個寬度為 $b$ 欄的窄面板 (panel) 中。接著將這 $b$ 個反射器壓縮成單個秩為 $b$ 的更新 (WY representation)，從而允許使用三個連續的 GEMMs (General Matrix Multiplications) 來更新矩陣的剩餘區塊。這將大部分運算轉換為 Tensor Cores 可以高效執行的格式。
自動化研究方法論：「迴圈工程」(Loop Engineering)
該開發者利用了高頻率的迭代迴圈，在 14 天內提交了超過 1,500 次。該過程依賴於大型語言模型 (LLMs，如 Codex 與 Claude) 與結構化框架的組合。
框架設置與引導
- 工具： 開發者使用 popcorn CLI 進行基準測試並提交至排行榜，並使用 Modal 進行分析與 nsys/NCU 分析。
- 日誌記錄： 一個
log.md檔案記錄了每一次提交、其狀態 (accept/reject) 以及各維度的大小 (shape-wise) 耗時，以防止重複實驗。 - 目標導向提示詞： 在 Codex 中使用
/goal指令來設定量化目標（例如：「超越我們目前最佳的 n=512 耗時」），讓模型能夠自主迭代數小時甚至數天。 - 監督：
/btw與/side指令允許開發者在不中斷主優化迴圈的情況下，查詢代理程式 (agent) 的 進度與當前的假設。
使用束搜尋 (Beam Search) 逃避局部最大值
當效能達到 3,000 $\mu$s 關卡時，模型經常會陷入局部最大值，僅進行微小的參數調整而非結構性的創新。為了達成此目標，開發者實作了候選者束搜尋：
- 代理程式不再僅維持單個現有的最佳方案，而是維持 3-5 個活躍的「想法家族」(idea families)。
- 這防止了高風險的結構性變更，僅因為最初表現不如當前最佳方案而過早地被捨棄。
- 該「束」(beam) 通常包含一個「開發」(exploit) 束、一個「接近成功」(near-miss) 束，以及一個「結構性/高風險」(structural/high-risk) 束。
核心 (Kernel) 的技術演進
該核心透過十次重大的結構性突破，最終達到 1,805 $\mu$s 的幾何平均值：
| 階段 | 變更 | 影響 |
|---|---|---|
| 1 | torch.geqrf |
基線起點 (>108.8k $\mu$s) |
| 2 | 分塊 WY QR (n=512) | 引入了面板 (panel) 分解與剩餘區塊更新 |
| 3 | 全維度分塊路徑 (all shapes) | 將分塊邏輯應用於所有矩陣大小 |
| 4 | Triton panels | 實作了自定義的 panel16/32 核心 |
| 5 | Cholesky-ORHR (n=4096) | 對於最大矩陣使用 Gram-Cholesky |
| 6 | CUDA graph replay | 減少了核心 (kernel) 啟動開銷 |
| 7 | 融合 V/T 佈局 (Fused V/T layout) | |
| 8 | Split16 panels | 透過 Gram-Schmidt 優化尾端處理 |
| 9 | 專門化維度 (Shape specialization) | 硬編碼行數並進行融合歸約 (fused reductions) |
| 10 | Superpanels | V256/T256 packs 與直接-H 返回 (1.80k $\mu$s) |
關鍵洞察與權衡
領域專家知識 vs. 自動化
雖然代理程式可以驅動顯著的加速，但領域專家知識對於引導至關重要。開發者指出，前 10 名的解決方案進一步透過以下方式進行了優化：
- 數據檢測： 利用特定的輸入分佈 (例如：低秩情況)。
- ** लाइब्रेरी-free 實作：--- 替換剩餘的 PyTorch 函式庫 (如 triangular solve) 以自定義 CUDA/Triton 實作。
- 精度管理： 使剩餘矩陣保持在 FP16 以避免重複的型別轉換。
泛化性 vs. 專用性
社群討論強調了這種「迴圈工程」方法的一個重大風險：對基準測試過度擬合 (overfitting to the benchmark)。
"8 out of 10 top solutions... completely broke at any other input than the competition ones. The only solutions that did not break... were made by experts who... followed and adjusted their solution in reasonable bounds."
這表明，雖然 LLM 驅動的迴圈對於解決特定、定義明確的基準測試非常強大，但若沒有強烈的人類引導，它們很難維持通用型的魯棒性 (robustness)。
Sources
相關
- Dispatch
- Dispatch
- Dispatch
- Dispatch
- Dispatch