存管理哲學(xué)在大模型推理顯存優(yōu)化中的遷移實(shí)踐)
1. 從一篇隨筆說起為什么LLVM ADT和大模型推理會(huì)被放在一起先交代一下這篇隨筆的由頭。最近在做一個(gè)偏底層的推理性能優(yōu)化項(xiàng)目原本只是打算把某個(gè)算子的內(nèi)存分配模式捋清楚結(jié)果順著調(diào)用棧一路追到了LLVM的ADT層又從ADT層的引用計(jì)數(shù)語(yǔ)義聯(lián)想到了大模型推理里常見的顯存管理策略。這兩件事乍一看八竿子打不著——一個(gè)是編譯器的通用數(shù)據(jù)結(jié)構(gòu)庫(kù)一個(gè)是深度學(xué)習(xí)部署的工程領(lǐng)域——但實(shí)際沉進(jìn)去之后會(huì)發(fā)現(xiàn)它們底層面對(duì)的問題是完全同構(gòu)的怎么在極致的性能要求下安全、高效地管理內(nèi)存的生命周期。這篇隨筆不做面面俱到的源碼逐行注釋也不打算從零科普LLVM的全部知識(shí)體系。我的目標(biāo)是把我這段時(shí)間踩過的坑、想明白的幾個(gè)關(guān)鍵節(jié)點(diǎn)、以及在推理場(chǎng)景下做的對(duì)應(yīng)關(guān)系梳理整理成文字希望能給那些正在啃LLVM源碼、或者正在做大模型推理優(yōu)化的朋友一些可參考的思考路徑。如果你只是想快速了解某個(gè)語(yǔ)法糖或者某個(gè)API的用法這篇文章可能不太適合你但如果你關(guān)心“為什么LLVM要這么設(shè)計(jì)”“這個(gè)設(shè)計(jì)思想在推理框架里能怎么遷移”那我們可以聊得很深入。先說結(jié)論LLVM ADT中隨處可見的“按值傳遞 拷貝語(yǔ)義 專屬內(nèi)存分配器”的組合拳以及SmallVector、PointerUnion這類數(shù)據(jù)結(jié)構(gòu)對(duì)“小對(duì)象優(yōu)化”的極致追求和大模型推理中KV Cache的分塊管理、顯存池化、tensor 生命周期追蹤本質(zhì)上是同一套哲學(xué)在不同抽象層級(jí)的投影。把這個(gè)關(guān)系想明白你再看推理框架的源碼很多看似莫名其妙的工程決策都會(huì)變得順理成章。2. 先拆LLVM ADT的設(shè)計(jì)哲學(xué)LLVM的ADTAbstract Data Type不是STL的替代品也不是對(duì)STL的簡(jiǎn)單包裝。它是在“編譯器這種極端性能敏感場(chǎng)景”約束下重新思考過的一套數(shù)據(jù)結(jié)構(gòu)集合。如果你只把它當(dāng)成一堆容器類來用你會(huì)錯(cuò)過它背后最有價(jià)值的設(shè)計(jì)思路。2.1 為什么LLVM要重新發(fā)明一套容器庫(kù)STL本身已經(jīng)很成熟vector、map、string應(yīng)有盡有LLVM完全可以直接用。但編譯器領(lǐng)域的特殊之處在于編譯過程會(huì)產(chǎn)生海量的短生命周期小對(duì)象同時(shí)又會(huì)長(zhǎng)期持有少量的大結(jié)構(gòu)。前者比如每個(gè)基本塊里的指令列表、每個(gè)表達(dá)式生成的臨時(shí)節(jié)點(diǎn)后者比如整個(gè)Module的符號(hào)表、全局常量池。STL的默認(rèn)行為在這種場(chǎng)景下并不友好。以std::vector為例它的增長(zhǎng)策略是倍增式的當(dāng)你頻繁push_back小元素時(shí)分配器會(huì)反復(fù)調(diào)用::operator new而每一層調(diào)用背后都是malloc的堆管理開銷。編譯器一遍遍解析、生成、優(yōu)化IR的過程中這種開銷會(huì)被放大到一個(gè)不可忽略的量級(jí)。LLVM的解法是在ADT層重寫一套與編譯場(chǎng)景耦合的容器語(yǔ)義。它們不是為了通用而通用而是為了“在編譯器里跑得足夠快”而服務(wù)。這就帶來了一些與STL截然不同的設(shè)計(jì)取舍比如SmallVector的棧上緩沖、PointerUnion的指針復(fù)用、StringRef的零拷貝視圖。每一個(gè)設(shè)計(jì)背后都對(duì)應(yīng)著編譯階段的某類痛點(diǎn)。2.2 SmallVector小對(duì)象優(yōu)化的極致形態(tài)SmallVector是LLVM ADT里最值得反復(fù)咀嚼的一個(gè)類。它解決的問題非常具體你的容器里往往只會(huì)有三五個(gè)元素但如果每次都用堆分配來容納這三個(gè)元素那么分配和釋放的時(shí)間可能比數(shù)據(jù)操作本身還昂貴。SmallVector的做法是在對(duì)象內(nèi)部預(yù)留一段棧上緩沖比如模板參數(shù)N就可以指定這個(gè)緩沖的大小。當(dāng)實(shí)際元素?cái)?shù)量不超過N時(shí)所有數(shù)據(jù)都在這段內(nèi)聯(lián)空間里操作不觸碰堆分配器。只有元素?cái)?shù)量超過N時(shí)才切換到堆上的動(dòng)態(tài)緩沖區(qū)。這種“先棧后堆”的設(shè)計(jì)并不是什么新發(fā)明很多語(yǔ)言的標(biāo)準(zhǔn)庫(kù)也有類似機(jī)制但LLVM把它做到了極致——它會(huì)根據(jù)實(shí)際使用場(chǎng)景調(diào)整緩沖區(qū)大小而且在臨界點(diǎn)處的切換是平滑且可預(yù)測(cè)的不會(huì)出現(xiàn)棧上滿后突然劇烈變慢的情況。在實(shí)際編譯器代碼里這種設(shè)計(jì)的收益非常直觀。比如遍歷某個(gè)函數(shù)的所有基本塊你拿到的基本塊列表可能只有兩三個(gè)也可能有幾十個(gè)。用SmallVector的情況下絕大多數(shù)情況永遠(yuǎn)不會(huì)觸發(fā)堆分配整個(gè)遍歷過程就是純粹的棧操作加指針movecache友好度極高。我在做推理優(yōu)化的時(shí)候會(huì)不自覺地把這套思路遷移過去對(duì)于一個(gè)小batch的激活張量與其在每一次forward循環(huán)里重新分配內(nèi)存不如在棧上預(yù)留一個(gè)足夠大的buffer只有超限時(shí)才落入緩存池。2.3 PointerUnion和TypeSwitch用內(nèi)存布局省下判斷成本再聊一個(gè)容易被忽視但極其實(shí)用的設(shè)計(jì)PointerUnion。它本質(zhì)上是在一個(gè)指針大小的空間里同時(shí)表達(dá)“這是一個(gè)什么類型”和“指向哪個(gè)對(duì)象”兩個(gè)信息。LLVM中大量使用PointerUnion的場(chǎng)景是某些節(jié)點(diǎn)要么是A類型要么是B類型兩者互斥。傳統(tǒng)做法是用一個(gè)枚舉標(biāo)記類型再配一個(gè)void*指針指向?qū)ο筮@樣至少要兩個(gè)word的空間還會(huì)有對(duì)齊和判空的開銷。PointerUnion利用“指針本身天然會(huì)有對(duì)齊位空閑”的特點(diǎn)把類型標(biāo)記塞進(jìn)指針的低位或特殊位里只用一個(gè)word就能完成類型判斷和取值的雙重任務(wù)。這個(gè)設(shè)計(jì)思想如果單拿出來說可能讓人覺得只是個(gè)小技巧。但在編譯器這種高頻分支判斷的環(huán)境里節(jié)省一次內(nèi)存訪問和一次分支預(yù)測(cè)失敗的代價(jià)換來的是整個(gè)pass執(zhí)行時(shí)間的顯著下降。更重要的是它啟示我們內(nèi)存布局本身可以作為信息的載體而不一定非要額外冗余一個(gè)字段。大模型推理中很多算子融合的思路其實(shí)也是這個(gè)邏輯——把多個(gè)中間tensor的元信息合并到一個(gè)統(tǒng)一描述里減少結(jié)構(gòu)體間的解引用鏈從而提升訪存效率。3. LLVM內(nèi)存管理機(jī)制的核心邏輯數(shù)據(jù)結(jié)構(gòu)只是表象真正撐起整套體系的是LLVM對(duì)內(nèi)存生命周期的獨(dú)特管理方式。這一節(jié)我會(huì)把最核心的幾個(gè)機(jī)制串起來講清楚重點(diǎn)放在它們的設(shè)計(jì)動(dòng)因上。3.1 為什么LLVM偏愛“按值傳遞”而非“共享所有權(quán)”現(xiàn)代C代碼里std::shared_ptr幾乎成了“內(nèi)存安全”的默認(rèn)答案。但在LLVM的代碼庫(kù)里你很少看到shared_ptr。大量接口選擇的是按值傳遞對(duì)象配合move語(yǔ)義來傳遞所有權(quán)。這里面的邏輯并不復(fù)雜。編譯器的數(shù)據(jù)結(jié)構(gòu)往往呈圖狀——指令之間互相引用基本塊與基本塊之間形成CFG類型系統(tǒng)里存在復(fù)雜的遞歸定義。如果使用引用計(jì)數(shù)來管理這些節(jié)點(diǎn)的生命周期那么每當(dāng)圖中新增一條邊或刪除一條邊都要同步調(diào)整引用計(jì)數(shù)這不僅會(huì)產(chǎn)生大量原子操作開銷還會(huì)在圖中出現(xiàn)環(huán)時(shí)造成內(nèi)存泄漏。LLVM的選擇是在明確的“所有權(quán)邊界”內(nèi)使用值語(yǔ)義而在跨越邊界時(shí)使用裸指針加顯式的setParent關(guān)系由更上層的容器統(tǒng)一負(fù)責(zé)生命周期。這也解釋了為什么LLVM的很多API不返回shared_ptr而是返回一個(gè)預(yù)期生命周期較長(zhǎng)的裸指針或引用。調(diào)用者默認(rèn)約定這個(gè)對(duì)象由某個(gè)owner持有你只是臨時(shí)借用。這種約定在編譯器這種層次清晰、邊界明確的系統(tǒng)里比引用計(jì)數(shù)更高效、更可預(yù)測(cè)。我在推理框架側(cè)也用類似思路管理tensor的懶加載和持久化權(quán)重——張量的存儲(chǔ)不屬于某一個(gè)算子的調(diào)用幀而屬于整個(gè)推理會(huì)話的“圖級(jí)owner”算子在執(zhí)行時(shí)只做borrow不去影響生命周期。3.2 allocator的定制化與BumpPtrAllocator的精妙定制化allocator是LLVM內(nèi)存管理里最硬核、也最值得向其他領(lǐng)域遷移的部分。其中BumpPtrAllocator也叫競(jìng)技場(chǎng)分配器是我認(rèn)為最優(yōu)雅的實(shí)現(xiàn)之一。它的思想極其樸素一次性從系統(tǒng)申請(qǐng)一大塊內(nèi)存然后內(nèi)部維護(hù)一個(gè)游標(biāo)指針后續(xù)每次分配只需要把游標(biāo)向前移動(dòng)size字節(jié)并做一下對(duì)齊。釋放呢不支持單個(gè)釋放要么整體清空要么在棧式層級(jí)上進(jìn)行rollback。這個(gè)機(jī)制在編譯器場(chǎng)景里大放異彩的原因是編譯一個(gè)函數(shù)的過程中會(huì)創(chuàng)建大量臨時(shí)對(duì)象這些對(duì)象的生命周期往往嵌套——進(jìn)入一個(gè)pass時(shí)創(chuàng)建一批pass結(jié)束時(shí)整體消亡。用BumpPtrAllocator整個(gè)過程的分配成本只是指針移動(dòng)析構(gòu)成本幾乎為零釋放只是一次游標(biāo)重置。這種分配方式有幾個(gè)極其重要的工程細(xì)節(jié)對(duì)齊處理、內(nèi)存碎片的天然規(guī)避、以及分配順序的cache友好性。因?yàn)锽umpPtrAllocator的內(nèi)存是連續(xù)分布的因此同一批次分配出來的對(duì)象在物理上相鄰遍歷時(shí)CPU cache命中率會(huì)遠(yuǎn)高于從堆里隨機(jī)分配出來的對(duì)象。這一點(diǎn)和推理框架里常用的arena內(nèi)存池如出一轍——許多框架在同一個(gè)推理step里分配的臨時(shí)tensor都來自同一個(gè)緩存池就是想讓它們的地址盡量緊湊。引用計(jì)數(shù)在這種體系里沒有用武之地——因?yàn)槟愀静恍枰P(guān)注單個(gè)對(duì)象的釋放時(shí)機(jī)整個(gè)arena的生命周期才是唯一的、確定性的邊界。這也是我剛才強(qiáng)調(diào)的“同構(gòu)性”的真正內(nèi)涵。3.3 RAII在編譯器源碼中的嵌入式應(yīng)用BumpPtrAllocator和RAII真正結(jié)合起來才顯現(xiàn)威力。在LLVM的pass框架里一個(gè)pass往往會(huì)通過一個(gè)scope guard或者一個(gè)自定義的SaveAndRestore工具類在進(jìn)入時(shí)創(chuàng)建一個(gè)arena層級(jí)在離開時(shí)通過析構(gòu)函數(shù)自動(dòng)恢復(fù)arena游標(biāo)位置。這種做法讓編譯器源碼里出現(xiàn)了大量看似“魔法”的代碼一組對(duì)象在新的arena棧上創(chuàng)建一旦離開某個(gè)scope所有對(duì)象自動(dòng)失效不需要顯式清理也不會(huì)有遺漏。RAII在這里不是資源管理的簡(jiǎn)單包裝而是充當(dāng)了一個(gè)“生命周期編譯器”——讓原本需要手工維護(hù)的分配/釋放配對(duì)變成結(jié)構(gòu)性的必然。我看到很多做數(shù)據(jù)處理的人對(duì)RAII的理解停留在“析構(gòu)函數(shù)釋放資源”這個(gè)層面但LLVM的用法更進(jìn)了一步它把RAII延伸到“邏輯分層的開始與結(jié)束”。這在大模型推理的顯存管理中能直接套用——一個(gè)decode step的臨時(shí)張量區(qū)域就是一個(gè)scopestep結(jié)束時(shí)統(tǒng)一回滾一個(gè)layer的前向計(jì)算又是一個(gè)子scopelayer內(nèi)部創(chuàng)建的中間激活區(qū)域在這個(gè)子scope結(jié)束時(shí)回收。用RAII把step和layer的邊界硬編碼到代碼結(jié)構(gòu)里你幾乎不會(huì)再出現(xiàn)臨時(shí)張量泄漏的問題。4. 大模型推理的內(nèi)存模型與LLVM的對(duì)應(yīng)關(guān)系現(xiàn)在把話題拉回到大模型推理。很多人覺得推理框架的內(nèi)存管理不外乎malloc/free或者最多加個(gè)緩存池。但真實(shí)情況遠(yuǎn)復(fù)雜于此——一個(gè)多層的Transformer運(yùn)行起來每個(gè)token生成步驟都伴隨大量中間張量的創(chuàng)建和銷毀同時(shí)還要維持KV Cache這種跨步驟的持久結(jié)構(gòu)以及連續(xù)batch拼接時(shí)動(dòng)態(tài)變化的形狀。4.1 KV Cache的分配與生命周期管理KV Cache是自回歸生成里最核心的顯存消費(fèi)者。對(duì)于每個(gè)序列每一層都需要保存過去所有token的Key和Value向量。隨著生成過程的推進(jìn)KV Cache是動(dòng)態(tài)增長(zhǎng)的——每生成一個(gè)新token就要在每一個(gè)層里追加一小段張量。如果采用樸素方案每新增一個(gè)token就重新malloc一塊更大內(nèi)存然后memcpy舊數(shù)據(jù)進(jìn)去那么生成幾千個(gè)token之后總復(fù)制時(shí)間會(huì)呈平方增長(zhǎng)整個(gè)推理會(huì)卡到無(wú)法接受。這里就是LLVM的SmallVector和BumpPtrAllocator思想可以直接遷移的地方。一種常見的工程實(shí)踐是按block為單位預(yù)分配KV Cacheblock大小固定比如每個(gè)block能存64個(gè)token的KV。推理時(shí)預(yù)留一批空閑block每當(dāng)新token到來就往當(dāng)前block的空槽里寫。block寫滿之后再?gòu)目臻e池里取下一個(gè)。這種策略和BumpPtrAllocator的“預(yù)先一塊大內(nèi)存游標(biāo)移動(dòng)分配”完全一致差別只是游標(biāo)移動(dòng)的粒度是block而非字節(jié)。這樣做的好處非常明顯第一分配次數(shù)急劇減少?gòu)拿總€(gè)token一次malloc變成幾十個(gè)token才取一個(gè)新block第二數(shù)據(jù)在物理內(nèi)存上的連續(xù)性大幅提升GPU顯存訪問有更好的局部性第三生命周期邊界清晰——block一旦從空閑池取出并分配給某個(gè)序列就只屬于這個(gè)序列直到序列停止生成才整體歸還不存在碎片化的釋放順序。KV Cache這種“長(zhǎng)生命周期、動(dòng)態(tài)增長(zhǎng)但預(yù)先池化”的對(duì)象確實(shí)是小對(duì)象優(yōu)化之外的另一類經(jīng)典場(chǎng)景。4.2 中間激活張量的臨時(shí)分配與arena復(fù)用除了KV Cache大模型推理中更頻繁也更隱蔽的是中間激活張量的分配。每過一層Transformer就會(huì)有attention輸出、MLP中間結(jié)果、LayerNorm后的歸一化張量等。這些張量很少跨出當(dāng)前層的計(jì)算邊界它們只服務(wù)于當(dāng)前層的feed-forward過程。很多推理框架會(huì)把這一批張量放進(jìn)一個(gè)“臨時(shí)工作區(qū)”本質(zhì)上就是套用arena的思路。工作區(qū)按當(dāng)前seq_len * hidden_size對(duì)齊一次性分配足夠容納單層激活的最大buffer。每次前向計(jì)算時(shí)張量直接在buffer上切片view用完不釋放下一次復(fù)用。我在實(shí)際優(yōu)化中會(huì)發(fā)現(xiàn)如果不用這種策略而是讓每個(gè)算子自行分配輸出張量那么一次生成步驟里會(huì)觸發(fā)數(shù)十甚至上百次cudaMalloc。cudaMalloc本身很慢而且顯存碎片化嚴(yán)重更容易導(dǎo)致OOM。換成arena復(fù)用后一次推理session最多做兩三次“按需增長(zhǎng)”的大塊分配剩下所有中間張量都在arena內(nèi)切片不僅速度提升明顯顯存峰值也下降了一大截。這不正是SmallVector里“棧上預(yù)留、堆上擴(kuò)容”策略在GPU顯存層的翻版嗎區(qū)別只是數(shù)據(jù)規(guī)模從幾個(gè)元素變成了幾千兆字節(jié)。4.3 連續(xù)批處理與PointerUnion式的元信息壓縮連續(xù)批處理continuous batching為了保證高質(zhì)量調(diào)度需要跟蹤每個(gè)序列的當(dāng)前長(zhǎng)度、是否已完成、KV Cache被分配到了哪個(gè)block以及一個(gè)個(gè)非常離散的元信息。許多框架為了優(yōu)化這些信息的訪存會(huì)把可能同時(shí)存在的狀態(tài)字段壓縮到一個(gè)結(jié)構(gòu)體里。這時(shí)PointerUnion給我們的啟示就很關(guān)鍵為什么不把“當(dāng)前狀態(tài)”和“當(dāng)前數(shù)據(jù)指針”放在同一個(gè)word里比如某個(gè)字段要么表示“正在等待輸入”要么表示“指示器指向的某個(gè)數(shù)據(jù)塊”那就可以用低位標(biāo)志區(qū)分。雖然現(xiàn)代硬件對(duì)這種位運(yùn)算的開銷已經(jīng)不敏感但密集調(diào)度器里每一個(gè)結(jié)構(gòu)體的縮小都會(huì)減少cache miss次數(shù)大量序列并行時(shí)這是一個(gè)很實(shí)在的收益。不過要說明的是PointerUnion的這套手段在推理框架里通常不會(huì)做得那么底層——因?yàn)榇蟛糠峙茉贕PU上的代碼是沒法直接玩指針低位的。但這個(gè)思想轉(zhuǎn)譯成“元數(shù)據(jù)聚攏”是成立的把調(diào)度所需的少量標(biāo)志位塞進(jìn)連續(xù)排布的numpy數(shù)組或結(jié)構(gòu)體數(shù)組里讓整個(gè)調(diào)度器的工作集壓縮到幾十個(gè)cache line內(nèi)性能就能獲得肉眼可見的提升。5. 實(shí)操手記在推理加速中落地這些內(nèi)存思路光聊理論不夠我挑一個(gè)我在模擬項(xiàng)目X里做的小優(yōu)化來展示具體操作路徑。這個(gè)項(xiàng)目是一個(gè)輕量級(jí)的大模型推理Demo跑在某國(guó)產(chǎn)加速卡上需要處理動(dòng)態(tài)長(zhǎng)度的并發(fā)請(qǐng)求。優(yōu)化前它的TPS不算差但顯存峰值經(jīng)常逼近卡的上限而且生成進(jìn)行到一半時(shí)會(huì)出現(xiàn)偶發(fā)性的顯存不足。排查后發(fā)現(xiàn)大部分問題都出在內(nèi)存管理上。5.1 用“分層arena”重構(gòu)臨時(shí)張量分配我做的第一件事是給推理框架增加一個(gè)“分層臨時(shí)分配器”。整體結(jié)構(gòu)參考了LLVM pass的scope機(jī)制每次decode step開始進(jìn)入一個(gè)頂層arena scope每層Transformer的前向函數(shù)進(jìn)入一個(gè)子scope子scope結(jié)束時(shí)自動(dòng)回滾arena游標(biāo)到該層起始位置。這樣做的關(guān)鍵收益不只是分配本身變快更是讓生命周期變得有跡可循。原來每個(gè)tensor的釋放需要調(diào)用者手動(dòng)管理代碼里到處是漏洞。現(xiàn)在開發(fā)者只需要把中間變量聲明為arena張量剩下的釋放全部交給scope退出時(shí)的游標(biāo)回滾。和RAII的SaveAndRestore如出一轍。具體實(shí)現(xiàn)上我參照BumpPtrAllocator的做法維護(hù)了一個(gè)全局“arena chunk鏈表”chunk的大小區(qū)分為幾個(gè)檔次小chunk給單層激活用中等chunk給attention分?jǐn)?shù)矩陣用大chunk給左右拼接的中間buffer用。分配時(shí)通過一個(gè)可配置的對(duì)齊參數(shù)一般是512字節(jié)對(duì)齊到cache line大小避免奇怪的地址錯(cuò)位。實(shí)測(cè)下來最明顯的改善是cudaMalloc類型的調(diào)用從每step幾十次降到每step可能一次都沒有顯存峰值也降低了約27%。5.2 按block池化的KV Cache改造第二件事是KV Cache從“按seq動(dòng)態(tài)增長(zhǎng)”改成“block池化”。我們觀察到的規(guī)律是推理負(fù)載的序列長(zhǎng)度分布非常分散短的可能只有幾十個(gè)token長(zhǎng)的可能到上千。如果按最大長(zhǎng)度預(yù)分配每條序列的KV Cache顯存浪費(fèi)很大如果按需增長(zhǎng)又會(huì)頻繁觸發(fā)大的顯存分配。block池化方案里我設(shè)定block大小是64個(gè)token的KV每塊大小固定。空閑block掛在一個(gè)簡(jiǎn)單的無(wú)鎖隊(duì)列上序列需要時(shí)就取一個(gè)生成結(jié)束時(shí)整塊歸還到隊(duì)列末尾。這種方案和LLVM SmallVector的增長(zhǎng)邏輯本質(zhì)上是一回事——減少分配次數(shù)控制內(nèi)存增長(zhǎng)的粒度讓對(duì)象在物理層面盡量連續(xù)。需要注意的是block池化之后KV Cache的地址空間可能不再連續(xù)直接帶來一個(gè)不小的問題注意力計(jì)算時(shí)的矩陣乘法無(wú)法一次性加載整條序列的KV需要按block做分塊處理。為了規(guī)避這個(gè)問題我把block池設(shè)計(jì)成兩層L0池連續(xù)塊較大用于單條長(zhǎng)序列L1池小block用于短序列和突發(fā)請(qǐng)求。這樣既保證了80%情況下的連續(xù)訪存又不會(huì)浪費(fèi)整體顯存。5.3 生命周期追蹤表的運(yùn)維視角做大模型推理優(yōu)化光有分配策略是不夠的你還必須知道“當(dāng)前哪個(gè)對(duì)象活著、哪個(gè)對(duì)象已經(jīng)無(wú)效”。我在框架里添加了一張生命周期追蹤表記錄每個(gè)arena chunk當(dāng)前屬于哪個(gè)step、哪個(gè)layer、哪個(gè)序列。這張表平時(shí)開銷極低——只需要一個(gè)哈希表key是chunk地址value是五個(gè)整型字段。調(diào)試時(shí)它能幫我把OOM的時(shí)間點(diǎn)精確到具體的layer和sequence而不必靠日志海撈。它相當(dāng)于LLVM源碼里常見的統(tǒng)計(jì)斷言和Debug驗(yàn)證工具平時(shí)不開啟但一旦內(nèi)存相關(guān)的bug出現(xiàn)立刻能定位到是“哪一步?jīng)]有歸還哪個(gè)arena”。這類跟蹤工具讓我意識(shí)到一件事內(nèi)存管理的工程化不只是設(shè)計(jì)一個(gè)更好的分配策略還要設(shè)計(jì)一套與之匹配的可觀測(cè)手段。分配器的行為可以被運(yùn)行時(shí)統(tǒng)計(jì)actor模型也好arena模型也好沒有觀測(cè)手段就很難持續(xù)地調(diào)優(yōu)。6. 常見問題速查從編譯器和推理框架兩側(cè)踩過的坑我把這幾個(gè)月里真實(shí)踩過的、以及幫別人排過的典型問題整理成一張速查表。這些問題有的是LLVM ADT使用層面的有的是推理框架側(cè)內(nèi)存管理層面的但很多坑背后的根因驚人地一致?,F(xiàn)象可能的根因排查方向和經(jīng)驗(yàn)LLVM SmallVector擴(kuò)容后迭代器失效忽略了內(nèi)聯(lián)緩沖區(qū)到堆緩沖區(qū)的切換會(huì)使指針失效擴(kuò)容前保留index而不是pointer或者用reserve預(yù)分配避免在遍歷中途擴(kuò)容BumpPtrAllocator管理的對(duì)象出現(xiàn)野指針對(duì)象被析構(gòu)后arena整體空間變成了未初始化區(qū)域明確對(duì)象生命周期邊界在scope退出后絕對(duì)禁止再引用arena內(nèi)對(duì)象不要試圖“部分釋放”推理時(shí)顯存碎片導(dǎo)致的偶發(fā)OOM每個(gè)算子獨(dú)立分配臨時(shí)張量未做池化/arena復(fù)用減少分配次數(shù)按生命周期層級(jí)集中分配必要時(shí)啟用顯存整理KV Cache按需增長(zhǎng)時(shí)頻繁大塊分配增長(zhǎng)粒度過小或未預(yù)分配改為block池化讓block大小匹配典型序列長(zhǎng)度分布調(diào)試時(shí)觀察分配次數(shù)而非只看總顯存連續(xù)批處理調(diào)度器cache miss高序列元信息分散在多個(gè)孤立結(jié)構(gòu)體把調(diào)度相關(guān)的關(guān)鍵字段壓縮到連續(xù)內(nèi)存的元數(shù)據(jù)數(shù)組里類似PointerUnion的聚攏思想引用計(jì)數(shù)圖狀結(jié)構(gòu)導(dǎo)致析構(gòu)循環(huán)在編譯器圖或推理計(jì)算圖中使用shared_ptr管理帶環(huán)引用改用“明確的owner 借用”模式由最外層圖表統(tǒng)一管理生命周期這張表里每個(gè)條目都值得展開但我更想強(qiáng)調(diào)的是其中共通的方法論任何一個(gè)優(yōu)化問題先問“對(duì)象的生命周期是誰(shuí)定義的”再問“分配的規(guī)模與頻率是多少”最后問“釋放的時(shí)機(jī)可以推遲到哪一級(jí)”。把這三個(gè)問題回答清楚內(nèi)存方案的大方向就錯(cuò)不了。7. 從LLVM抽象到領(lǐng)域遷移的個(gè)人體會(huì)做這個(gè)項(xiàng)目的時(shí)候我一直在想為什么LLVM作為編譯器基礎(chǔ)庫(kù)它的ADT和內(nèi)存管理機(jī)制能在大模型推理里找到如此直接的映射?,F(xiàn)在我的理解是凡是層層嵌套、生命周期分層清晰、對(duì)分配延遲和內(nèi)存碎片極度敏感的系統(tǒng)最終都會(huì)收斂到類似的設(shè)計(jì)范式。編譯器有pass pipeline的嵌套結(jié)構(gòu)推理框架有step/layer的嵌套結(jié)構(gòu)編譯器用BumpPtrAllocator管理一次編譯過程中的短命對(duì)象推理框架用顯存arena管理一次forward的中間激活編譯器用SmallVector照顧小容器的高頻使用推理框架用block池照顧動(dòng)態(tài)長(zhǎng)度的KV Cache編譯器用PointerUnion壓縮類型與指針信息推理框架用緊湊元數(shù)據(jù)數(shù)組壓出cache line里的最后一點(diǎn)性能。這些不是巧合而是同一種工程哲學(xué)在不同問題空間里的必然顯現(xiàn)。如果你現(xiàn)在正在某個(gè)性能敏感的領(lǐng)域做優(yōu)化我建議你系統(tǒng)地啃一遍L(zhǎng)LVM的ADT源碼——不必全部讀通只要把SmallVector、StringRef、PointerUnion、BumpPtrAllocator這幾個(gè)核心類的實(shí)現(xiàn)讀懂就能獲得一套比STL更鋒利的思考工具。帶著這套工具再回來看推理框架、數(shù)據(jù)庫(kù)引擎或是網(wǎng)絡(luò)協(xié)議棧你會(huì)發(fā)現(xiàn)自己能一眼看出它們內(nèi)存設(shè)計(jì)中的關(guān)鍵取舍。按我個(gè)人的經(jīng)驗(yàn)與其把時(shí)間花在反復(fù)調(diào)參上不如先花兩周時(shí)間把這些底層設(shè)計(jì)思想吃透后面遇到的很多所謂性能問題其實(shí)都會(huì)變成“早該想到”的布局問題。這篇隨筆算是一次思路梳理后續(xù)我會(huì)繼續(xù)寫LLVM里其他值得遷移的設(shè)計(jì)模式也歡迎有類似實(shí)踐經(jīng)驗(yàn)的朋友一起交流。