利用 LLM Agent 自動化開發 GPU Kernel

手寫 GPU kernel 需要耐心、努力，而且還會消耗你的理智。我一直在實驗用現代 LLM（Claude Opus 5、GPT 5.6 Sol）來直接幫我「印出」kernel。我想聊聊我從中學到的事情。

老實說，寫 GPU kernel 是一項容易驗證到令人不好意思的任務。

首先，你決定要為哪些操作撰寫 kernel。接著，你寫出第一個版本，並確保它能編譯且沒有明顯的錯誤。第三步，你用一個慢速的參考實作來檢查正確性。如果結果不符合，你就嘗試修正正確性的問題。一旦完成，你就對 kernel 的執行時間進行基準測試。後續的版本就建立在這之上，你持續最佳化，直到達到 roofline 指標或你滿意為止。

可驗證迴圈：GPU kernel 開發

上圖顯示它是一個可驗證迴圈：

編譯檢查（compile check）在你考慮正確性之前，就是一個緊密的本地迴圈（B <-> C）。
正確性檢查（D <-> E <-> F）是核心的可驗證獎勵迴圈。正是這部分讓 kernel 撰寫成為適合自動化驗證的好題目，因為你有一個 ground-truth 參考可以比對。
最佳化（G -> H -> 回到 D）會為每個新版本重用相同的正確性迴圈，因為一個快但錯誤的 kernel 是沒有價值的，每個版本都必須驗證正確性。
roofline/滿意度檢查（I）是決定要繼續最佳化還是停下來的外部迴圈。

第一版

上面的圖仍然是比較高層次的視角，魔鬼藏在細節裡。我們需要確保 LLM agent 擁有所需的全部 context，才能開始寫出一個好的第一版。

CUDA DSL 在這裡登場。Triton、CuTeDSL 和 Tilelang 是其中非常容易在 Python 中上手的選擇。它們的學習曲線比 CUDA C++ 平緩許多，但這些 DSL 的抽象概念反而可能讓 agent 更困惑。我們需要一種方法，把這些抽象的 context 傳遞給 agent。

現代 LLM 已經知道如何寫出「好的」 Triton。即使沒有任何 context，它們也能很好地處理 Triton 的抽象。不過，對於像 CuTeDSL 這種提供了比 Triton 更多控制的 DSL，我發現建立一個 context 目錄，讓 agent 可以在裡面查閱、理解 DSL 的抽象概念，會有很大的幫助。

舉例來說，將
NVIDIA cutlass
倉庫克隆到 context 目錄中，是讓 agent 在撰寫 CuTeDSL kernel 時查找相關抽象的好方法，例如
Layout Algebra、Copy/GEMM atoms、記憶體階層、範例 kernel
等等。

根據我的經驗，一個好的第一版 kernel 能無明顯錯誤地編譯，並通過我下面要講的正確性測試。

測試、基準測試與效能剖析

在給 agent 足夠的 context 之後，真正的瓶頸就轉移到驗證上。參考實作本身，以及與其對照驗證，變得越來越重要。我把這個階段稱為正確性測試（correctness testing），或簡稱測試。參考實作的速度遠不如它的設計意圖來得重要。你打算測量和驗證的目標，正是你的 agent 會去最佳化的方向。

通常，當計算並非以較低精度進行時，我會測量最大絕對/相對誤差（MAE）、均方誤差（MSE/RMSE）和PSNR（峰值信噪比）。當涉及較低精度時，我傾向於測量PSNR和餘弦相似度（cossim）。

要讓 kernel 版本真正在 GPU 上執行，取決於 GPU 是本機可用還是透過雲端取得。無論如何，我們的 agent 都應該要有某種方式能存取它的輸出。

我發現下面這個 rung 方法論是建立 N 個測試函式的好方法：

python

1def rung(name):
2    def deco(fn):
3        try:
4            out = fn()
5            results[name] = {"ok": True, **(out or {})}
6            print(f"[{name}] ok  " + "  ".join(f"{k}={v}" for k, v in (out or {}).items()))
7        except Exception as e:
8            results[name] = {"ok": False, "err": f"{type(e).__name__}: {e}"}
9            print(f"[{name}] FAILED  {type(e).__name__}: {e}")
10            traceback.print_exc()
11        return fn
12    return deco

你可以像這樣呼叫：

python

1out = {}
2
3@rung("pre-checks")
4def _():
5    run_pure_checks()
6    run_dsl_checks()
7
8@rung("run")
9def _():
10    out["o"] = custom_kernel(*inputs)
11    torch.cuda.synchronize()
12    return {"shape": tuple(out["o"].shape),
13            "finite": bool(torch.isfinite(out["o"]).all())}

在 benchmark rung 方面，你可以做很多事情：

端到端基準測試 kernel 的執行時間，了解整體花費的時間
使用 kernel 內部追蹤（intra-kernel tracing），基準測試 kernel 內的各區段，並將結果傾印到輸出中（使用自訂 tracer 或 CUPTI）
將產生的 IR、PTX、SASS 和 CUBIN 傾印到 dumps 目錄，讓 agent 去檢視

最後一點可以再多說一些。有的時候，DSL 在 lower 的過程中可能會產生次佳的 PTX（最終是 SASS），而你會發現有更好的指令或 shape 可以替代。我們的 agent 可以閱讀 PTX/SASS 文字檔，直接內聯較低層級的程式碼，而不是讓 DSL 處理那個次佳的部分。同樣地，把「可搜尋的」 PTX 文件當作 context 傳遞，在這裡非常有幫助。

最後一個把所有東西串起來的是效能剖析（Profiling）。如果你的 agent 能存取 NCU（Nsight Compute Systems）CLI，你可以請它剖析你的 kernel 並產生報告，作為上述可驗證回饋迴圈的一部分。

結語

「那麼 GPU kernel 開發已經死透了嗎？」

「嗯，是，但其實也不是。」

是啊，因為 layout、indexing、抽象和整體結構這些困難的部分，在給 agent 足夠的 context 之後，大多都能被解決。你可以輕鬆地把 2-3 週的工作縮短成 1-2 天。又不是，因為真正的瓶頸已經從 kernel 轉移到驗證上。現在沒有單一的做法：你的 context 和 harness（測試框架）越好，整個流程就越快。特殊案例會受益更多，你需要的只是一個好的 harness。

最後，與其把 agent 當成完全自主的個體，不如把它們當成你可以引導的、非常聰明的助手，這仍然是必要的。這正是你對 GPU 和 kernel 的基本理解派上用場的地方。這裡仍然需要人的部分——也就是你。

我知道，這種感覺苦樂參半 :)