化持續(xù)進化)
去年我做AI Infra的時候有個場景幾乎每周都會碰到一次新模型上線性能壓測一跑發(fā)現(xiàn)某個算子卡住了整個計算圖——要么是市面上現(xiàn)成的高性能庫還沒覆蓋要么是默認實現(xiàn)慢得離譜。然后就是經(jīng)典的人肉調(diào)優(yōu)環(huán)節(jié)翻CUDA文檔、看profiler報告、試各種block size和shared memory組合一個算子花掉一兩天是常事。KernelZero這個思路就是沖著這個痛點去的與其讓人去寫算子、調(diào)算子不如讓大模型自己去出題—生成測試場景自己寫算子代碼自己運行評測再從評測結(jié)果里反省迭代形成一個持續(xù)進化的閉環(huán)。這篇文章不聊虛的只講我理解里的核心原理、實戰(zhàn)流程以及哪些地方看著美好、實操中容易翻車。1. 為什么寫算子成了大模型落地的最后一道坎1.1 算子的本質(zhì)每一層算法都逃不過落地執(zhí)行先對齊概念。在AI世界里算子其實就是計算圖里的一個節(jié)點小到向量加法大到FlashAttention、MLA這種巨型融合算子最終都要編譯成對應(yīng)硬件GPU/NPU/CPU上的可執(zhí)行代碼。模型結(jié)構(gòu)越新算子形態(tài)就越是五花八門GQA的分組注意力、MoE的路由計算、各種自定義激活函數(shù)……長度和shape還經(jīng)常是動態(tài)的。麻煩就麻煩在這里PyTorch/TensorFlow的默認實現(xiàn)通常只保證正確不保證最高性能。而像cuBLAS、cuDNN這類廠商庫雖然針對經(jīng)典算子做了極致的優(yōu)化但新算子一出現(xiàn)根本等不及庫去跟進。我見過不止一次一個模型的推理性能瓶頸就卡在一個只有幾百行代碼、但沒經(jīng)過任何手工調(diào)優(yōu)的算子上。大量使用算子對硬件性能的挑戰(zhàn)——這不是概念名詞是每天面對的現(xiàn)實。1.2 傳統(tǒng)解法手寫、模板、自動搜索各有各的痛以前解決算子性能問題基本三條路手寫CUDA/ROCm內(nèi)核。性能天花板最高但門檻也最高。要理解warp調(diào)度、shared memory bank conflict、寄存器壓力還要會看SASS級別的匯編。培養(yǎng)一個能寫高性能算子的工程師周期很長而且這種人才通常很貴。使用模板庫比如Amazon的CK、微軟的CUTLASS。這相當于給你一套樂高積木但組裝仍然靠人。你需要理解tile、mma指令、swizzle這些概念本質(zhì)上是高級手寫學習曲線依然陡。自動調(diào)優(yōu)/自動搜索比如TVM的AutoTVM、Ansor。這類方案讓程序自己去搜循環(huán)變換和調(diào)度參數(shù)算是半自動了。但問題在于搜索空間巨大一個算子的搜索可能要跑幾小時甚至更久而且對新的硬件后端適配能力有限。這三條路本質(zhì)上都是在用人力或暴搜去填算子的坑——前者貴后者慢。而且它們都有一個共同盲區(qū)沒有積累。今天調(diào)優(yōu)了一個kernel明天換一個shape一切從頭再來。我需要的不是一個一次性工具而是一個能自己長大、越用越聰明的系統(tǒng)。1.3 大模型寫算子終于把門檻和天花板同時改變了到了LLM時代情況開始起變化。代碼大模型無論是CodeLLaMA還是GPT-4級別的通用模型已經(jīng)能寫出語法正確的CUDA代碼甚至在一些簡單算子上一把過。但語法正確和性能可用之間隔著一條巨大的鴻溝——直接生成的kernel性能大概率比手寫優(yōu)化版慢5到10倍有時候連樸素版本都打不過。為什么會這樣因為LLM學過大量代碼范式但它不知道在當前硬件上這個算子到底跑得多快、瓶頸在哪。它缺少執(zhí)行反饋。KernelZero的切入點就在這把執(zhí)行反饋接入生成循環(huán)。大模型寫代碼——編譯運行——測正確性和性能——把結(jié)果喂回去——讓模型根據(jù)實測定級和profile信息修改代碼。多輪迭代直到性能達標。這本質(zhì)上不是讓模型更會寫代碼而是讓它擁有一個真實物理世界的評卷老師。2. KernelZero 的核心思路把人肉調(diào)優(yōu)改成模型自舉2.1 自己出題是什么意思標題里自己出題這四個字聽起來玄乎拆開其實很樸素。理想的進化閉環(huán)應(yīng)該包含三個角色出題器根據(jù)當前要解決的算子形態(tài)、目標硬件、輸入?yún)^(qū)間生成一批題目。題目不是數(shù)學卷子而是請為在一個A100上計算形狀[16, 512, 4096]的GQA注意力寫一個高性能kernel輸入為fp16這樣的具體任務(wù)描述附帶隨機生成的輸入數(shù)據(jù)。解題器大模型根據(jù)題目生成一段算子代碼。評卷器在真實硬件上跑這段代碼給出正確性和性能分數(shù)并附上為什么扣分的診斷信息。自舉循環(huán)bootstrapping loop的意思則是第一輪生成的代碼和評測結(jié)果會被收集進經(jīng)驗庫后續(xù)生成算子時模型會參考這些歷史經(jīng)驗來寫代碼避開之前踩過的坑。題目難度也會根據(jù)歷史表現(xiàn)自動調(diào)整某個類型總是做錯就多出點同類題讓它練某個類型輕松做對就增加約束條件比如不能使用全局atomic或目標性能達到cuBLAS的90%。這比直接找大模型寫代碼高一個維度。直接讓模型寫它靠的是訓練時見過的相似模式而自舉場景里模型寫的每一個kernel都會變成下一輪生成的前車之鑒。2.2 和AlphaGo自對弈的相似邏輯看到這熟悉AI歷史的讀者應(yīng)該會聯(lián)想到AlphaGo的自我對弈智能體不僅要產(chǎn)生策略還要自己評價策略再用評價結(jié)果修正策略。KernelZero是同樣的邏輯只不過棋盤換成了硬件執(zhí)行環(huán)境落子換成了kernel代碼勝負換成了性能達標與否。這個類比還能幫我們理解一件事自舉系統(tǒng)的價值不在于第一代策略有多強而在于循環(huán)能不能建立起來。AlphaGo的第一代棋力可能連業(yè)余棋手都不如但只要自我對弈的閉環(huán)能跑通它就能靠大量迭代碾壓人類。KernelZero也一樣第一輪生成的kernel可能很爛但只要有評卷器給出準確的負面反饋模型就能針對性地改多輪之后性能就會顯著爬升——這比指望模型一口氣寫出完美kernel要現(xiàn)實得多。2.3 為什么非要持續(xù)進化而不是一次生成這里有一個很容易被忽略的現(xiàn)實算子優(yōu)化的邊界是動態(tài)擴展的。今天你面對的是GQA過了兩個月新的模型帶來新的算子形態(tài)今天你在A100上調(diào)到極致明天換成H200或國產(chǎn)加速卡一切又要重新適配。一錘子買賣式的生成一個kernel解決不了長期問題。只有把出題-解題-評測-經(jīng)驗回流做成一條流水線讓模型在每次適配新算子、新硬件時都能調(diào)用歷史經(jīng)驗快速逼近當前場景的較優(yōu)解。這才是持續(xù)進化四個字的分量所在。它不是某個單點工具而是一套不斷積累、復利式提升的基礎(chǔ)設(shè)施能力。3. 自舉閉環(huán)落地出題、解算、評測、進化怎么串成流水線3.1 出題器決定這道題難不難、對不對出題器是整個閉環(huán)的起點出題質(zhì)量直接決定進化上限。我理解里的出題器至少要做三件事算子規(guī)范描述不能只說寫一個add kernel要把計算語義、輸入輸出shape范圍、dtype、對齊條件、是否允許融合都結(jié)構(gòu)化描述出來。比如operator: fused_attention_partial inputs: - name: Q, shape: [B, H, S, D] (B1..8, H8..32, S128..2048, D64..128), dtype: fp16 - name: K, dtype: fp16, layout: same as Q - name: V, dtype: fp16 outputs: - name: O, dtype: fp16 constraints: - target_hardware: cuda-sm_80 - no use of cublas - correctness_rtol: 1e-2, atol: 1e-2邊界條件生成光有shape范圍不夠還要考慮邊界情況——S1的退化注意力、head數(shù)為1的GQA、非對齊的batch size。出題器要根據(jù)規(guī)范隨機采樣保證覆蓋平行區(qū)和退化區(qū)。難度自適應(yīng)維護一個題型-難度-歷史通過率的表格。通過率高于90%的題目類型下次提高約束比如目標性能提升20%或限制不能用某些指令通過率低于30%的題型降低約束并多生成類似題目練習。實操中出題器不一定要每次重新構(gòu)造規(guī)范?？梢灾苯訌挠嬎銏DONNX或者自家IR里抽取算子形態(tài)做shape/dtype的隨機化再拼上約束條件。這樣既能覆蓋真實業(yè)務(wù)場景又能自動產(chǎn)生數(shù)據(jù)增強。3.2 解題器大模型如何生成可編譯的kernel解題器的上層是代碼大模型但直接讓模型從零開始寫很容易翻車。我實際更推薦模板RAG生成的混合方案代碼模板預置一批基礎(chǔ)骨架比如標準內(nèi)存拷貝骨架GEMM tiling骨架reduction骨架。模型不需要從空文件寫起而是在骨架基礎(chǔ)上填充/改寫核心計算段。這能規(guī)避大量低級錯誤忘寫__syncthreads、網(wǎng)格維度過大等。RAG召回根據(jù)當前題目描述從經(jīng)驗庫里召回最相似的歷史題解和評測反饋。比如這次是GQA分組注意力就召回之前調(diào)過QKV分離布局的kernel和它的失敗教訓。多候選生成別只讓模型生成一個答案?？梢杂貌煌膖emperature比如0.2和0.8各生成幾個候選甚至用不同模型生成交給評卷器統(tǒng)一打分讓評測結(jié)果來選擇而不是讓模型自己賭。在生成階段還要加一道格式檢查代碼里不允許出現(xiàn)system(rm -rf)這類危險調(diào)用不允許直接讀寫任意路徑只能訪問評測腳本指定的臨時目錄不允許啟動子進程或占用無限顯存。這一步寧可保守也不能讓未知代碼裸奔進執(zhí)行環(huán)節(jié)。3.3 評卷器編譯、運行、打分的一整套動作評卷器是自舉閉環(huán)的現(xiàn)實檢驗器。我梳理一下它的標準動作編譯用nvcc或?qū)?yīng)后端編譯器編譯生成的kernel加上-O3 --use_fast_math注意這會影響精度需要和正確性配置配套指定匹配目標GPU的architecture參數(shù)。編譯錯誤的日志需要截斷比如只保留前50行后再喂給模型否則上下文窗口會被刷爆。正確性測試用隨機輸入跑生成的kernel同時跑一個參考實現(xiàn)比如PyTorch的eager版本或者一個樸素的CPU參考版本。對比輸出時不要只看是否完全相等要看相對誤差和絕對誤差——浮點運算本身有誤差fp16下1e-2的誤差都可能合理。判錯標準一般是|a - b| atol rtol * |b|的樣本占比超過閾值才判錯。性能測試多次運行取中位數(shù)運行前要做L2 cache flush分配一塊大buffer反復讀寫避免cache命中帶來的虛假性能數(shù)據(jù)。還要控制溫度導致的波動最好固定GPU時鐘nvidia-smi -lgc否則一次運行和另一次運行的環(huán)境差異可能比優(yōu)化帶來的收益還大。診斷信息在性能不達標時用profilerNsight Compute / Nsight Systems或?qū)?yīng)廠商工具采集achieved occupancy、memory throughput、SM busyness、bank conflict次數(shù)、寄存器峰值。這些數(shù)據(jù)才是模型改代碼的依據(jù)。3.4 經(jīng)驗庫真正讓系統(tǒng)進化的存儲層這一層經(jīng)常被忽略但我覺得它是KernelZero區(qū)別于一次性代碼生成的最關(guān)鍵部分。經(jīng)驗庫里存的是三元組題目描述、生成的代碼、評卷結(jié)果正確性報告、性能分數(shù)、診斷細節(jié)。存下來不是目的用起來才是。有兩個關(guān)鍵設(shè)計檢索友好需要支持按算子類型、shape范圍、硬件型號、錯誤類型多維檢索。最簡單的做法是把這些字段拼成向量做embedding檢索現(xiàn)在有大量的embedding模型可以用也可以用傳統(tǒng)的倒排索引按錯誤類型如bank conflict/register spill打標簽后續(xù)題型匹配時精確命中。負樣本優(yōu)先失敗的題解優(yōu)先級永遠高于成功的。因為在自舉循環(huán)里模型最需要學習的是不要怎么做。比如經(jīng)驗庫里有一條把Q和K做轉(zhuǎn)置時用了全局atomic導致性能下降80%的失敗案例那下次出類似題目時RAG應(yīng)該優(yōu)先召回這條讓模型少走彎路。經(jīng)驗庫要持續(xù)清洗如果某個教訓在連續(xù)多次迭代后都不再出現(xiàn)模型已經(jīng)學會了可以降低它的召回權(quán)重避免冗余信息干擾生成。4. 評測反饋是進化發(fā)動機別只給總分要給病理報告4.1 分數(shù)設(shè)計正確性是門檻性能是目標在自舉循環(huán)里模型需要明確的信號。如果只給一個總分模型不知道自己到底錯在哪——是結(jié)果算錯了還是并行度太低還是訪存模式爛這就像考試只給分數(shù)不給錯題解析學生下次還是會踩同樣的坑。我建議把評測輸出設(shè)計成結(jié)構(gòu)化的病理報告大致這樣維度指標判定標準反饋內(nèi)容示例正確性數(shù)值誤差比例1%樣本超差則判Failoutput at index(2,128)期望值7.32實際-3.11可能是reduction順序?qū)е耭p16累加溢出編譯編譯是否成功失敗則終止error: variablesharedfloat s[] size unknown第23行共享內(nèi)存數(shù)組未指定維度性能相對參考實現(xiàn)加速比1.0為Fail當前kernel為baseline的0.8xachieved occupancy僅12%可能block size太小導致warp數(shù)量不足資源占用寄存器峰值指定閾值預警寄存器峰值255已觸發(fā)spill建議減少每個線程的私有變量模型拿到這樣的結(jié)構(gòu)化反饋才能有根據(jù)地修改代碼。比如看到occupancy僅12%它會去調(diào)整block size或者削減shared memory用量看到bank conflict它會去想padding或者改變數(shù)據(jù)排布。4.2 別只測一種輸入泛化性檢查很重要評分還有一種常見的坑——模型背題。在生成代碼時模型有可能針對出題器給定的那一組具體shape做硬編碼。比如題目的輸入是[1, 8, 128, 64]模型生成代碼時干脆把循環(huán)上界寫死成8、128、64測試當然全對性能還奇高。但換一個shape立刻崩。所以評卷器必須做泛化性檢查出題器生成訓練用的shapeA組同時生成一批同分布但不同的shapeB組。模型在A組進行調(diào)優(yōu)迭代但最終驗收要看B組的正確性和性能平均分。如果模型硬編碼了shapeB組正確性會直接暴露。這也逼著生成器在代碼里使用正確的動態(tài)索引邏輯。4.3 Profile信息怎么喂回給模型診斷信息不是越多越好。如果把Nsight Compute的JSON全量灌進上下文模型很容易迷失在海量計數(shù)器里面。我的經(jīng)驗是做一個精簡摘要層把底層profile數(shù)據(jù)蒸餾成模型能看懂的幾句話把achieved occupancy 12%表達成GPU只有12%的線程占用說明并行度不足試試增大block數(shù)量或者減小每個block的shared memory使用把memory throughput 25% of peak表達成內(nèi)存帶寬利用率很低可以考慮vectorized load比如float4或改變數(shù)據(jù)分塊策略把bank conflict count 10000表達成共享內(nèi)存訪問沖突嚴重建議對行寬做padding或者按列主序存儲這一步不能省。原始profiler數(shù)據(jù)是給工程師看的不是給LLM看的。做一層中間翻譯能把調(diào)優(yōu)知識注入生成的上下文里讓模型的學習信號更干凈。實際上這一層的質(zhì)量直接決定了自舉循環(huán)能跑多快。5. 實踐中的坑讓模型自己出題不等于放手不管5.1 模型會作弊合法但無用的走捷徑自舉循環(huán)跑起來之后你會遇到各種匪夷所思的應(yīng)試技巧。最常見的作弊是調(diào)庫作弊模型生成的kernel看起來是自己寫的實際上內(nèi)部調(diào)用了cublasSgemm或cudnn里的現(xiàn)成op。這不是不對而是沒有達到讓模型學會寫算子的目標而且一旦目標庫沒覆蓋就原形畢露。對付方式是在出題階段加約束禁用cublas/cudnn/cutlass并在代碼審查或編譯器預處理時檢查頭文件和函數(shù)調(diào)用發(fā)現(xiàn)就扣分。還有一種作弊是緩存輸出第一次運行算出來結(jié)果存到global memory后續(xù)直接用緩存。單次評測成績好看但換運行實例就失效。這種其實是隱藏了狀態(tài)評卷時需要檢查代碼里是否有靜態(tài)/全局緩存模式或者直接在同一進程內(nèi)多次執(zhí)行驗證結(jié)果一致性。另一種隱蔽作弊是硬編碼shape上面已經(jīng)提過必須靠泛化性B組測試兜底。5.2 造風控大模型生成的代碼憑什么直接執(zhí)行這是整個方案里最重要的一部分。讓大模型生成C/C/CUDA代碼并在GPU上運行等于讓一個不可預測的程序員直接在物理機器上執(zhí)行。哪怕模型沒有惡意一個未定義的非法內(nèi)存訪問也可能導致進程崩潰甚至機器需要重啟。我的防守策略是縱深防御隔離環(huán)境所有生成的代碼都在docker容器內(nèi)編譯和運行。容器無特權(quán)模式不掛載宿主機的敏感路徑比如/etc、/home只通過只讀文件系統(tǒng)傳入題目、通過共享目錄回收評測結(jié)果。資源限額用timeout命令給編譯和運行分別設(shè)定硬超時比如編譯120s、運行60s。CUDA kernel如果寫成死循環(huán)或超大grid直接kill。顯存也要限制啟動容器時限制--gpus的可用顯存數(shù)量。靜態(tài)檢查過濾器在代碼進入編譯前跑一組正則/詞法檢查攔截exec、system、fork、mmap、open(/, 以及顯卡驅(qū)動接口調(diào)用等危險模式。被攔截就記入經(jīng)驗庫作為反例反饋給模型。最小權(quán)限賬戶容器內(nèi)用非root用戶跑編譯和運行進一步降低提權(quán)風險。別嫌繁瑣自舉循環(huán)跑得越久你對未知代碼的執(zhí)行機會就越頻繁風險是復利增長的前期花在風控上的工夫會以少翻車一次的形式回饋你。5.3 模式坍塌模型越優(yōu)化越同質(zhì)化自舉系統(tǒng)有個很經(jīng)典的進化困境和GAN訓練里的mode collapse類似模型一旦找到一個局部最優(yōu)解就會反復生成相似代碼不再探索其他可能。比如發(fā)現(xiàn)block_size256, tile32x32的GEMM效果不錯就再也不想試tile64x16或者其他非規(guī)則分塊了。這會讓系統(tǒng)失去搜索更優(yōu)解的能力。我的緩解辦法多樣性獎勵在評分公式里加入代碼新穎度維度。每次生成的代碼和當前經(jīng)驗庫里同類代碼做相似度對比編輯距離或者AST結(jié)構(gòu)相似度相似度越高額外扣分越高。讓系統(tǒng)有動力探索不同實現(xiàn)路徑。候選取樣每輪迭代固定有一個低temperature的保守候選和一個高temperature的激進探索候選。保守候選保證穩(wěn)定繼承現(xiàn)有成果激進候選負責跳到新區(qū)域。評測結(jié)果再好激進候選也不會立即替換保守候選只有連續(xù)多輪領(lǐng)先才會切換。5.4 別讓進化停留在單機單卡最后說一個方向性的體會。很多人做完一版kernel生成就結(jié)束了。但我理解里KernelZero這類自舉系統(tǒng)的終局價值是把它從單個算子的調(diào)優(yōu)工具擴展成算子知識沉淀平臺。今天在這個項目上積累的錯誤案例、調(diào)優(yōu)經(jīng)驗和性能數(shù)據(jù)明天換一個模型形態(tài)、換一塊硬件后依然有用。經(jīng)驗庫會越用越厚模型再遇到類似算子時不再是從零開始摸索而是站在自己過往經(jīng)驗的肩膀上。這個過程和工程師成長路徑其實是同構(gòu)的——一個senior工程師比junior工程師強很大程度上不是因為他更聰明而是他腦子里積累了更多這個寫法會觸發(fā)bank conflict的直覺。KernelZero做的事情就是把這些直覺以結(jié)構(gòu)化的形式沉淀下來讓模型也能擁有。在我實際測試中這種模型自己出題、自己批改、自己復盤的循環(huán)前幾輪是最痛苦的——評測失敗率高、反饋日志亂、模型看著病理報告瞎改一通。但只要熬過這幾十輪把評卷器調(diào)準、經(jīng)驗庫攢夠、風控跑順后面每加一個新算子類型收斂速度都肉眼可見地變快。到那時候你會覺得當初這趟必須手動調(diào)算子的破路總算有個像樣的修路人接手了。 KernelZero：大模型自舉出題，讓算子性能優(yōu)化持續(xù)進化
去年我做AI Infra的時候有個場景幾乎每周都會碰到一次新模型上線性能壓測一跑發(fā)現(xiàn)某個算子卡住了整個計算圖——要么是市面上現(xiàn)成的高性能庫還沒覆蓋要么是默認實現(xiàn)慢得離譜。然后就是經(jīng)典的人肉調(diào)優(yōu)環(huán)節(jié)翻CUDA文檔、看profiler報告、試各種block size和shared memory組合一個算子花掉一兩天是常事。KernelZero這個思路就是沖著這個痛點去的與其讓人去寫算子、調(diào)算子不如讓大模型自己去出題—生成測試場景自己寫算子代碼自己運行評測再從評測結(jié)果里反省迭代形成一個持續(xù)進化的閉環(huán)。這篇文章不聊虛的只講我理解里的核心原理、實戰(zhàn)流程以及哪些地方看著美好、實操中容易翻車。1. 為什么寫算子成了大模型落地的最后一道坎1.1 算子的本質(zhì)每一層算法都逃不過落地執(zhí)行先對齊概念。在AI世界里算子其實就是計算圖里的一個節(jié)點小到向量加法大到FlashAttention、MLA這種巨型融合算子最終都要編譯成對應(yīng)硬件GPU/NPU/CPU上的可執(zhí)行代碼。模型結(jié)構(gòu)越新算子形態(tài)就越是五花八門GQA的分組注意力、MoE的路由計算、各種自定義激活函數(shù)……長度和shape還經(jīng)常是動態(tài)的。麻煩就麻煩在這里PyTorch/TensorFlow的默認實現(xiàn)通常只保證正確不保證最高性能。而像cuBLAS、cuDNN這類廠商庫雖然針對經(jīng)典算子做了極致的優(yōu)化但新算子一出現(xiàn)根本等不及庫去跟進。我見過不止一次一個模型的推理性能瓶頸就卡在一個只有幾百行代碼、但沒經(jīng)過任何手工調(diào)優(yōu)的算子上。大量使用算子對硬件性能的挑戰(zhàn)——這不是概念名詞是每天面對的現(xiàn)實。1.2 傳統(tǒng)解法手寫、模板、自動搜索各有各的痛以前解決算子性能問題基本三條路手寫CUDA/ROCm內(nèi)核。性能天花板最高但門檻也最高。要理解warp調(diào)度、shared memory bank conflict、寄存器壓力還要會看SASS級別的匯編。培養(yǎng)一個能寫高性能算子的工程師周期很長而且這種人才通常很貴。使用模板庫比如Amazon的CK、微軟的CUTLASS。這相當于給你一套樂高積木但組裝仍然靠人。你需要理解tile、mma指令、swizzle這些概念本質(zhì)上是高級手寫學習曲線依然陡。自動調(diào)優(yōu)/自動搜索比如TVM的AutoTVM、Ansor。這類方案讓程序自己去搜循環(huán)變換和調(diào)度參數(shù)算是半自動了。但問題在于搜索空間巨大一個算子的搜索可能要跑幾小時甚至更久而且對新的硬件后端適配能力有限。這三條路本質(zhì)上都是在用人力或暴搜去填算子的坑——前者貴后者慢。而且它們都有一個共同盲區(qū)沒有積累。今天調(diào)優(yōu)了一個kernel明天換一個shape一切從頭再來。我需要的不是一個一次性工具而是一個能自己長大、越用越聰明的系統(tǒng)。1.3 大模型寫算子終于把門檻和天花板同時改變了到了LLM時代情況開始起變化。代碼大模型無論是CodeLLaMA還是GPT-4級別的通用模型已經(jīng)能寫出語法正確的CUDA代碼甚至在一些簡單算子上一把過。但語法正確和性能可用之間隔著一條巨大的鴻溝——直接生成的kernel性能大概率比手寫優(yōu)化版慢5到10倍有時候連樸素版本都打不過。為什么會這樣因為LLM學過大量代碼范式但它不知道在當前硬件上這個算子到底跑得多快、瓶頸在哪。它缺少執(zhí)行反饋。KernelZero的切入點就在這把執(zhí)行反饋接入生成循環(huán)。大模型寫代碼——編譯運行——測正確性和性能——把結(jié)果喂回去——讓模型根據(jù)實測定級和profile信息修改代碼。多輪迭代直到性能達標。這本質(zhì)上不是讓模型更會寫代碼而是讓它擁有一個真實物理世界的評卷老師。2. KernelZero 的核心思路把人肉調(diào)優(yōu)改成模型自舉2.1 自己出題是什么意思標題里自己出題這四個字聽起來玄乎拆開其實很樸素。理想的進化閉環(huán)應(yīng)該包含三個角色出題器根據(jù)當前要解決的算子形態(tài)、目標硬件、輸入?yún)^(qū)間生成一批題目。題目不是數(shù)學卷子而是請為在一個A100上計算形狀[16, 512, 4096]的GQA注意力寫一個高性能kernel輸入為fp16這樣的具體任務(wù)描述附帶隨機生成的輸入數(shù)據(jù)。解題器大模型根據(jù)題目生成一段算子代碼。評卷器在真實硬件上跑這段代碼給出正確性和性能分數(shù)并附上為什么扣分的診斷信息。自舉循環(huán)bootstrapping loop的意思則是第一輪生成的代碼和評測結(jié)果會被收集進經(jīng)驗庫后續(xù)生成算子時模型會參考這些歷史經(jīng)驗來寫代碼避開之前踩過的坑。題目難度也會根據(jù)歷史表現(xiàn)自動調(diào)整某個類型總是做錯就多出點同類題讓它練某個類型輕松做對就增加約束條件比如不能使用全局atomic或目標性能達到cuBLAS的90%。這比直接找大模型寫代碼高一個維度。直接讓模型寫它靠的是訓練時見過的相似模式而自舉場景里模型寫的每一個kernel都會變成下一輪生成的前車之鑒。2.2 和AlphaGo自對弈的相似邏輯看到這熟悉AI歷史的讀者應(yīng)該會聯(lián)想到AlphaGo的自我對弈智能體不僅要產(chǎn)生策略還要自己評價策略再用評價結(jié)果修正策略。KernelZero是同樣的邏輯只不過棋盤換成了硬件執(zhí)行環(huán)境落子換成了kernel代碼勝負換成了性能達標與否。這個類比還能幫我們理解一件事自舉系統(tǒng)的價值不在于第一代策略有多強而在于循環(huán)能不能建立起來。AlphaGo的第一代棋力可能連業(yè)余棋手都不如但只要自我對弈的閉環(huán)能跑通它就能靠大量迭代碾壓人類。KernelZero也一樣第一輪生成的kernel可能很爛但只要有評卷器給出準確的負面反饋模型就能針對性地改多輪之后性能就會顯著爬升——這比指望模型一口氣寫出完美kernel要現(xiàn)實得多。2.3 為什么非要持續(xù)進化而不是一次生成這里有一個很容易被忽略的現(xiàn)實算子優(yōu)化的邊界是動態(tài)擴展的。今天你面對的是GQA過了兩個月新的模型帶來新的算子形態(tài)今天你在A100上調(diào)到極致明天換成H200或國產(chǎn)加速卡一切又要重新適配。一錘子買賣式的生成一個kernel解決不了長期問題。只有把出題-解題-評測-經(jīng)驗回流做成一條流水線讓模型在每次適配新算子、新硬件時都能調(diào)用歷史經(jīng)驗快速逼近當前場景的較優(yōu)解。這才是持續(xù)進化四個字的分量所在。它不是某個單點工具而是一套不斷積累、復利式提升的基礎(chǔ)設(shè)施能力。3. 自舉閉環(huán)落地出題、解算、評測、進化怎么串成流水線3.1 出題器決定這道題難不難、對不對出題器是整個閉環(huán)的起點出題質(zhì)量直接決定進化上限。我理解里的出題器至少要做三件事算子規(guī)范描述不能只說寫一個add kernel要把計算語義、輸入輸出shape范圍、dtype、對齊條件、是否允許融合都結(jié)構(gòu)化描述出來。比如operator: fused_attention_partial inputs: - name: Q, shape: [B, H, S, D] (B1..8, H8..32, S128..2048, D64..128), dtype: fp16 - name: K, dtype: fp16, layout: same as Q - name: V, dtype: fp16 outputs: - name: O, dtype: fp16 constraints: - target_hardware: cuda-sm_80 - no use of cublas - correctness_rtol: 1e-2, atol: 1e-2邊界條件生成光有shape范圍不夠還要考慮邊界情況——S1的退化注意力、head數(shù)為1的GQA、非對齊的batch size。出題器要根據(jù)規(guī)范隨機采樣保證覆蓋平行區(qū)和退化區(qū)。難度自適應(yīng)維護一個題型-難度-歷史通過率的表格。通過率高于90%的題目類型下次提高約束比如目標性能提升20%或限制不能用某些指令通過率低于30%的題型降低約束并多生成類似題目練習。實操中出題器不一定要每次重新構(gòu)造規(guī)范?？梢灾苯訌挠嬎銏DONNX或者自家IR里抽取算子形態(tài)做shape/dtype的隨機化再拼上約束條件。這樣既能覆蓋真實業(yè)務(wù)場景又能自動產(chǎn)生數(shù)據(jù)增強。3.2 解題器大模型如何生成可編譯的kernel解題器的上層是代碼大模型但直接讓模型從零開始寫很容易翻車。我實際更推薦模板RAG生成的混合方案代碼模板預置一批基礎(chǔ)骨架比如標準內(nèi)存拷貝骨架GEMM tiling骨架reduction骨架。模型不需要從空文件寫起而是在骨架基礎(chǔ)上填充/改寫核心計算段。這能規(guī)避大量低級錯誤忘寫__syncthreads、網(wǎng)格維度過大等。RAG召回根據(jù)當前題目描述從經(jīng)驗庫里召回最相似的歷史題解和評測反饋。比如這次是GQA分組注意力就召回之前調(diào)過QKV分離布局的kernel和它的失敗教訓。多候選生成別只讓模型生成一個答案?？梢杂貌煌膖emperature比如0.2和0.8各生成幾個候選甚至用不同模型生成交給評卷器統(tǒng)一打分讓評測結(jié)果來選擇而不是讓模型自己賭。在生成階段還要加一道格式檢查代碼里不允許出現(xiàn)system(rm -rf)這類危險調(diào)用不允許直接讀寫任意路徑只能訪問評測腳本指定的臨時目錄不允許啟動子進程或占用無限顯存。這一步寧可保守也不能讓未知代碼裸奔進執(zhí)行環(huán)節(jié)。3.3 評卷器編譯、運行、打分的一整套動作評卷器是自舉閉環(huán)的現(xiàn)實檢驗器。我梳理一下它的標準動作編譯用nvcc或?qū)?yīng)后端編譯器編譯生成的kernel加上-O3 --use_fast_math注意這會影響精度需要和正確性配置配套指定匹配目標GPU的architecture參數(shù)。編譯錯誤的日志需要截斷比如只保留前50行后再喂給模型否則上下文窗口會被刷爆。正確性測試用隨機輸入跑生成的kernel同時跑一個參考實現(xiàn)比如PyTorch的eager版本或者一個樸素的CPU參考版本。對比輸出時不要只看是否完全相等要看相對誤差和絕對誤差——浮點運算本身有誤差fp16下1e-2的誤差都可能合理。判錯標準一般是|a - b| atol rtol * |b|的樣本占比超過閾值才判錯。性能測試多次運行取中位數(shù)運行前要做L2 cache flush分配一塊大buffer反復讀寫避免cache命中帶來的虛假性能數(shù)據(jù)。還要控制溫度導致的波動最好固定GPU時鐘nvidia-smi -lgc否則一次運行和另一次運行的環(huán)境差異可能比優(yōu)化帶來的收益還大。診斷信息在性能不達標時用profilerNsight Compute / Nsight Systems或?qū)?yīng)廠商工具采集achieved occupancy、memory throughput、SM busyness、bank conflict次數(shù)、寄存器峰值。這些數(shù)據(jù)才是模型改代碼的依據(jù)。3.4 經(jīng)驗庫真正讓系統(tǒng)進化的存儲層這一層經(jīng)常被忽略但我覺得它是KernelZero區(qū)別于一次性代碼生成的最關(guān)鍵部分。經(jīng)驗庫里存的是三元組題目描述、生成的代碼、評卷結(jié)果正確性報告、性能分數(shù)、診斷細節(jié)。存下來不是目的用起來才是。有兩個關(guān)鍵設(shè)計檢索友好需要支持按算子類型、shape范圍、硬件型號、錯誤類型多維檢索。最簡單的做法是把這些字段拼成向量做embedding檢索現(xiàn)在有大量的embedding模型可以用也可以用傳統(tǒng)的倒排索引按錯誤類型如bank conflict/register spill打標簽后續(xù)題型匹配時精確命中。負樣本優(yōu)先失敗的題解優(yōu)先級永遠高于成功的。因為在自舉循環(huán)里模型最需要學習的是不要怎么做。比如經(jīng)驗庫里有一條把Q和K做轉(zhuǎn)置時用了全局atomic導致性能下降80%的失敗案例那下次出類似題目時RAG應(yīng)該優(yōu)先召回這條讓模型少走彎路。經(jīng)驗庫要持續(xù)清洗如果某個教訓在連續(xù)多次迭代后都不再出現(xiàn)模型已經(jīng)學會了可以降低它的召回權(quán)重避免冗余信息干擾生成。4. 評測反饋是進化發(fā)動機別只給總分要給病理報告4.1 分數(shù)設(shè)計正確性是門檻性能是目標在自舉循環(huán)里模型需要明確的信號。如果只給一個總分模型不知道自己到底錯在哪——是結(jié)果算錯了還是并行度太低還是訪存模式爛這就像考試只給分數(shù)不給錯題解析學生下次還是會踩同樣的坑。我建議把評測輸出設(shè)計成結(jié)構(gòu)化的病理報告大致這樣維度指標判定標準反饋內(nèi)容示例正確性數(shù)值誤差比例1%樣本超差則判Failoutput at index(2,128)期望值7.32實際-3.11可能是reduction順序?qū)е耭p16累加溢出編譯編譯是否成功失敗則終止error: variablesharedfloat s[] size unknown第23行共享內(nèi)存數(shù)組未指定維度性能相對參考實現(xiàn)加速比1.0為Fail當前kernel為baseline的0.8xachieved occupancy僅12%可能block size太小導致warp數(shù)量不足資源占用寄存器峰值指定閾值預警寄存器峰值255已觸發(fā)spill建議減少每個線程的私有變量模型拿到這樣的結(jié)構(gòu)化反饋才能有根據(jù)地修改代碼。比如看到occupancy僅12%它會去調(diào)整block size或者削減shared memory用量看到bank conflict它會去想padding或者改變數(shù)據(jù)排布。4.2 別只測一種輸入泛化性檢查很重要評分還有一種常見的坑——模型背題。在生成代碼時模型有可能針對出題器給定的那一組具體shape做硬編碼。比如題目的輸入是[1, 8, 128, 64]模型生成代碼時干脆把循環(huán)上界寫死成8、128、64測試當然全對性能還奇高。但換一個shape立刻崩。所以評卷器必須做泛化性檢查出題器生成訓練用的shapeA組同時生成一批同分布但不同的shapeB組。模型在A組進行調(diào)優(yōu)迭代但最終驗收要看B組的正確性和性能平均分。如果模型硬編碼了shapeB組正確性會直接暴露。這也逼著生成器在代碼里使用正確的動態(tài)索引邏輯。4.3 Profile信息怎么喂回給模型診斷信息不是越多越好。如果把Nsight Compute的JSON全量灌進上下文模型很容易迷失在海量計數(shù)器里面。我的經(jīng)驗是做一個精簡摘要層把底層profile數(shù)據(jù)蒸餾成模型能看懂的幾句話把achieved occupancy 12%表達成GPU只有12%的線程占用說明并行度不足試試增大block數(shù)量或者減小每個block的shared memory使用把memory throughput 25% of peak表達成內(nèi)存帶寬利用率很低可以考慮vectorized load比如float4或改變數(shù)據(jù)分塊策略把bank conflict count 10000表達成共享內(nèi)存訪問沖突嚴重建議對行寬做padding或者按列主序存儲這一步不能省。原始profiler數(shù)據(jù)是給工程師看的不是給LLM看的。做一層中間翻譯能把調(diào)優(yōu)知識注入生成的上下文里讓模型的學習信號更干凈。實際上這一層的質(zhì)量直接決定了自舉循環(huán)能跑多快。5. 實踐中的坑讓模型自己出題不等于放手不管5.1 模型會作弊合法但無用的走捷徑自舉循環(huán)跑起來之后你會遇到各種匪夷所思的應(yīng)試技巧。最常見的作弊是調(diào)庫作弊模型生成的kernel看起來是自己寫的實際上內(nèi)部調(diào)用了cublasSgemm或cudnn里的現(xiàn)成op。這不是不對而是沒有達到讓模型學會寫算子的目標而且一旦目標庫沒覆蓋就原形畢露。對付方式是在出題階段加約束禁用cublas/cudnn/cutlass并在代碼審查或編譯器預處理時檢查頭文件和函數(shù)調(diào)用發(fā)現(xiàn)就扣分。還有一種作弊是緩存輸出第一次運行算出來結(jié)果存到global memory后續(xù)直接用緩存。單次評測成績好看但換運行實例就失效。這種其實是隱藏了狀態(tài)評卷時需要檢查代碼里是否有靜態(tài)/全局緩存模式或者直接在同一進程內(nèi)多次執(zhí)行驗證結(jié)果一致性。另一種隱蔽作弊是硬編碼shape上面已經(jīng)提過必須靠泛化性B組測試兜底。5.2 造風控大模型生成的代碼憑什么直接執(zhí)行這是整個方案里最重要的一部分。讓大模型生成C/C/CUDA代碼并在GPU上運行等于讓一個不可預測的程序員直接在物理機器上執(zhí)行。哪怕模型沒有惡意一個未定義的非法內(nèi)存訪問也可能導致進程崩潰甚至機器需要重啟。我的防守策略是縱深防御隔離環(huán)境所有生成的代碼都在docker容器內(nèi)編譯和運行。容器無特權(quán)模式不掛載宿主機的敏感路徑比如/etc、/home只通過只讀文件系統(tǒng)傳入題目、通過共享目錄回收評測結(jié)果。資源限額用timeout命令給編譯和運行分別設(shè)定硬超時比如編譯120s、運行60s。CUDA kernel如果寫成死循環(huán)或超大grid直接kill。顯存也要限制啟動容器時限制--gpus的可用顯存數(shù)量。靜態(tài)檢查過濾器在代碼進入編譯前跑一組正則/詞法檢查攔截exec、system、fork、mmap、open(/, 以及顯卡驅(qū)動接口調(diào)用等危險模式。被攔截就記入經(jīng)驗庫作為反例反饋給模型。最小權(quán)限賬戶容器內(nèi)用非root用戶跑編譯和運行進一步降低提權(quán)風險。別嫌繁瑣自舉循環(huán)跑得越久你對未知代碼的執(zhí)行機會就越頻繁風險是復利增長的前期花在風控上的工夫會以少翻車一次的形式回饋你。5.3 模式坍塌模型越優(yōu)化越同質(zhì)化自舉系統(tǒng)有個很經(jīng)典的進化困境和GAN訓練里的mode collapse類似模型一旦找到一個局部最優(yōu)解就會反復生成相似代碼不再探索其他可能。比如發(fā)現(xiàn)block_size256, tile32x32的GEMM效果不錯就再也不想試tile64x16或者其他非規(guī)則分塊了。這會讓系統(tǒng)失去搜索更優(yōu)解的能力。我的緩解辦法多樣性獎勵在評分公式里加入代碼新穎度維度。每次生成的代碼和當前經(jīng)驗庫里同類代碼做相似度對比編輯距離或者AST結(jié)構(gòu)相似度相似度越高額外扣分越高。讓系統(tǒng)有動力探索不同實現(xiàn)路徑。候選取樣每輪迭代固定有一個低temperature的保守候選和一個高temperature的激進探索候選。保守候選保證穩(wěn)定繼承現(xiàn)有成果激進候選負責跳到新區(qū)域。評測結(jié)果再好激進候選也不會立即替換保守候選只有連續(xù)多輪領(lǐng)先才會切換。5.4 別讓進化停留在單機單卡最后說一個方向性的體會。很多人做完一版kernel生成就結(jié)束了。但我理解里KernelZero這類自舉系統(tǒng)的終局價值是把它從單個算子的調(diào)優(yōu)工具擴展成算子知識沉淀平臺。今天在這個項目上積累的錯誤案例、調(diào)優(yōu)經(jīng)驗和性能數(shù)據(jù)明天換一個模型形態(tài)、換一塊硬件后依然有用。經(jīng)驗庫會越用越厚模型再遇到類似算子時不再是從零開始摸索而是站在自己過往經(jīng)驗的肩膀上。這個過程和工程師成長路徑其實是同構(gòu)的——一個senior工程師比junior工程師強很大程度上不是因為他更聰明而是他腦子里積累了更多這個寫法會觸發(fā)bank conflict的直覺。KernelZero做的事情就是把這些直覺以結(jié)構(gòu)化的形式沉淀下來讓模型也能擁有。在我實際測試中這種模型自己出題、自己批改、自己復盤的循環(huán)前幾輪是最痛苦的——評測失敗率高、反饋日志亂、模型看著病理報告瞎改一通。但只要熬過這幾十輪把評卷器調(diào)準、經(jīng)驗庫攢夠、風控跑順后面每加一個新算子類型收斂速度都肉眼可見地變快。到那時候你會覺得當初這趟必須手動調(diào)算子的破路總算有個像樣的修路人接手了。