正式驗證AI生成的GPU內核
AI智能體生成的高性能GPU內核雖然通過數值測試，但仍可能存在隱藏缺陷。本文介紹了Gimlet Labs開發的一款早期研究系統，利用形式化驗證補充傳統數值測試，確保AI生成及人工編寫的內核的正確性。通過一個缺失中間裁剪的注意力機制案例，展示了形式化驗證如何發現測試遺漏的等價性問題。
2026年7月8日，Gimlet Labs的研究團隊發表了一項關於形式化驗證AI生成GPU內核的早期工作。隨着AI智能體生成高性能內核的能力增強，如何信任其輸出成為新的瓶頸。傳統數值測試只能覆蓋有限輸入空間，而AI智能體可能針對測試集進行優化，導致獎勵攻擊或功能缺失。
以縮放點積注意力（SDPA）為例：參考實現包含中間值裁剪（clamp）至[-5,5]的操作，而AI生成的優化內核為了融合操作移除了該裁剪。隨機測試中兩者輸出相似，但實際生產環境可能因輸入邊界觸發差異。這類缺陷難以通過增加測試用例徹底解決，因為輸入空間本質上是無限的。
Gimlet的驗證系統將參考模型（PyTorch）和候選內核（Triton）分別降級為張量代數DAG，分解為基本數學運算，再標量化處理為單個輸出元素的公式。隨後調用Z3求解器比較兩個公式，返回等價證明（UNSAT）或反例（SAT）。該方法類似編譯器領域正式驗證的實踐，如LLVM的Alive2和CompCert。
當前系統仍是研究原型，但已能捕獲真實工作中通過數值測試的錯誤。團隊在ARRAY 2026會議上展示了完整工作。未來，形式化驗證有望成為AI生成內核生產部署的關鍵保障。
此外，Gimlet Labs還在探索將這一方法擴展到CUDA內核等更低層次的實現。由於CUDA不提供類似TTIR的可檢查中間表示，需要開發專門的提取管道來恢復計算圖。研究團隊計劃進一步優化驗證器的性能，使其能夠處理更大規模的內核，並最終集成到CI/CD流程中，為AI生成內核的部署提供自動化的正確性保證。