最近這兩周晚上的畫風有點分裂一邊在翻LLVM的ADT頭文件和BumpPtrAllocator的實現(xiàn)一邊跟著nano-vllm的調(diào)度循環(huán)調(diào)大模型推理的顯存分配。一個是編譯器后端的地基工程一個是AI推理服務(wù)的性能主戰(zhàn)場按理說八竿子打不著。但把兩份代碼同時攤開之后我越看越覺得它們其實在回答同一組問題數(shù)據(jù)怎么表達才不浪費、資源怎么分配才不碎片、依賴怎么梳理才不亂、生命周期到底該由誰負責。這篇隨筆不打算做教程式輸出而是把我這段交叉閱讀的收獲強制整理一遍。內(nèi)容大致會掃過LLVM ADT里幾個高頻數(shù)據(jù)結(jié)構(gòu)、LLVM的內(nèi)存管理哲學再順著顯存治理和算子自發(fā)現(xiàn)一路聊到大模型推理框架。標題既然掛了個隨筆01目標就定得實在一點每個主題不求講透但畫準畫像讓后面幾篇能在這個地基上繼續(xù)長。1. 編譯器藍本與推理引擎之間的同構(gòu)感在哪里1.1 三個繞不開的底層問題把LLVM和大模型推理引擎放到一起看不是標題工程而是它們都要回答同樣的一組底層問題。第一個問題是數(shù)據(jù)結(jié)構(gòu)的承載能力。LLVM的IR對象在pass之間來回穿梭需要大量只讀視圖不拷貝切片的操作推理引擎的tensor在模型前向過程里也是反復(fù)引用、變形、transpose理想情況同樣是零拷貝?？涩F(xiàn)實里沒人愿意為所有場景各寫一套容器所以兩邊都催生了高度定制的數(shù)據(jù)結(jié)構(gòu)。LLVM有自己的一套ADTvLLM則用自己的block管理器和tensor元信息描述來減少反復(fù)搬運。第二個問題是資源分配的粒度。編譯器在優(yōu)化階段反復(fù)創(chuàng)建和銷毀指令節(jié)點如果每次創(chuàng)建都走malloc百萬級IR的構(gòu)建成本會直接爆炸。推理引擎則要在幾百個并發(fā)請求之間分配KV Cache顯存的申請釋放頻率直接決定吞吐和碎片率。兩者都不約而同地選擇了批量分配、按池回收的路子。第三個問題是編排與依賴。編譯器里pass之間有AnalysisManager來維護依賴關(guān)系同一個分析結(jié)果可以被多個pass共享緩存推理引擎的scheduler要根據(jù)tensor依賴、顯存水位、請求優(yōu)先級來決定這一步該讓哪些序列前向、哪些序列先搶占。依賴關(guān)系處理得好不好直接決定整個執(zhí)行管線順不順。理解了這三點再回頭去看LLVM ADT和內(nèi)存管理機制很多設(shè)計就不會只看表面而是能明白它為什么這么設(shè)計。1.2 我這段時間的實際學習環(huán)境先說LLVM這邊。我本地用的是LLVM 19的源碼構(gòu)建時強烈建議只保留host targetcmake -G Ninja ../llvm \ -DCMAKE_BUILD_TYPERelease \ -DLLVM_ENABLE_ASSERTIONSON \ -DLLVM_TARGETS_TO_BUILDhost \ -DLLVM_BUILD_TOOLSONLLVM_ENABLE_ASSERTIONS不要關(guān)很多ADT和內(nèi)存管理的不變量比如迭代器失效、數(shù)組越界的assert靠它兜底。LLVM_TARGETS_TO_BUILDhost能省掉一大半編譯時間我們只是讀代碼不需要交叉后端。推理那邊我在看nano-vllm倉庫不大Python 3.10、PyTorch 2.x的環(huán)境下clone下來就能跑。用CUDA平臺跑一個小規(guī)模的Llama模型做推理然后專門在scheduler和block_manager上打斷點。這么做的原因很簡單同一套問題的兩套答題方式放在一起看記憶才深。只讀LLVM的arena分配器容易覺得那是編譯器特殊需求只有當你在nano-vllm里看到顯存池按塊分配、按請求歸還才會意識到這類方案是高性能系統(tǒng)的通用答案。2. LLVM ADT到底解決了什么問題如果你平時寫C業(yè)務(wù)代碼看LLVM ADT的第一反應(yīng)是這不就是加強版std::vector和std::string_view嗎。有這種感覺正常。但你對照GCC或MSVC庫里那套面向通用場景的實現(xiàn)再來看LLVM ADT會發(fā)現(xiàn)它完全是按編譯器的使用方式從零設(shè)計的犧牲通用性換取極致的分配次數(shù)和緩存局部性。2.1 StringRef/ArrayRef不擁有數(shù)據(jù)只描述視野StringRef是先于std::string_view出現(xiàn)的東西本質(zhì)就是一個(const char*, size_t)的二元組引用了別人的字符緩沖區(qū)自己不持有、不管理、不拷貝。它的好用之處在于幾乎所有字符串操作都可以原地完成StringRef full getFileContent(); StringRef trimmed full.trim(); StringRef firstLine trimmed.split(\n).first; fx: drop_front(2).take_front(16);注意split返回一對StringRef沒有產(chǎn)生任何字符串拷貝。這在lexer和parser里可以省掉成千上萬次臨時string構(gòu)造。ArrayRef則是把同樣的思路泛化成任意類型它描述一段連續(xù)內(nèi)存的只讀視圖。函數(shù)簽名用ArrayRef而不是vector好處是調(diào)用方可以傳SmallVector、std::vector、甚至std::initializer_list不用為了傳參構(gòu)造臨時vector。比如bool canFold(ArrayRefValue* ops); canFold({a, b, c}); // 不用臨時 vector我實際踩過的坑是永遠不要把ArrayRef存成成員變量除非你能嚴格保證底層數(shù)組的生命周期覆蓋它。StringRef/ArrayRef的哲學是用的時候再取不要保存它適合做參數(shù)和局部變量。2.2 SmallVector棧上預(yù)分配的小型倉庫SmallVectorT, N給容器預(yù)分配了N個元素的?？臻g當元素數(shù)量不超過N的時候完全不會觸發(fā)堆分配。LLVM里大量場景是一條指令三五個操作數(shù)一個基本塊十幾條指令用std::vector的話每次構(gòu)造都要malloc分配器開銷比實際元素拷貝還貴。這種設(shè)計的本質(zhì)是用小概率的浪費棧上預(yù)留N個元素空間換取大概率情況下的零堆分配。N的選擇很有講究選太大占用棧空間選太小命中率下降。LLVM內(nèi)部常用的N是4、8、16具體量級看容器在高頻路徑上的典型大小。使用SmallVector時有個和std::vector不太一樣的習慣LLVM風格里經(jīng)常直接構(gòu)造后調(diào)用push_back很少reserve因為當N夠大的時候reserve也只是移動cur指針。當超過N時SmallVector會轉(zhuǎn)到堆上分配擴容策略和張量分配器有相似之處。2.3 DenseMap開放尋址帶來的性能與紀律DenseMap是LLVM為高頻查找場景自研的哈希表。它采用開放尋址open addressing所有元素放在一塊連續(xù)內(nèi)存上而不是像std::unordered_map那樣每個節(jié)點單獨分配再串鏈表。連續(xù)內(nèi)存帶來的收益是緩存局部性一個cache line能載入多個bucket查找效率在數(shù)據(jù)量大時差距非常明顯。代價是使用紀律更嚴格。開放尋址的刪除使用墓碑標記插入觸達閾值會整體rehash。最影響使用體驗的是DenseMap的迭代器在插入導致rehash后全部失效。這一點和std::unordered_map節(jié)點式迭代器相對穩(wěn)定很不一樣。所以遍歷DenseMap時如果要插入新元素要么提前把key收集到SmallVector里要么先把迭代器推進完再做插入。類似這種為了快犧牲一點直覺的設(shè)計在LLVM ADT里到處都是。APInt表示任意位寬整數(shù)Twine延遲字符串拼接iterator_range把指針對封裝成可迭代對象。這套容器庫最值錢的地方不是單點性能而是它把整個編譯器的數(shù)據(jù)結(jié)構(gòu)底座統(tǒng)一了每個子系統(tǒng)都知道對方的數(shù)據(jù)布局是什么樣盡量減少適配層。3. BumpPtrAllocator與一次性釋放的內(nèi)存哲學聊LLVM就繞不開BumpPtrAllocator。我在讀llvm/Support/Allocator.h時反復(fù)看了好幾遍它其實簡單得驚人但哲學足夠深刻。3.1 原理內(nèi)存就是一條持續(xù)推進的水位線BumpPtrAllocator維護了一組大塊內(nèi)存。每次allocate時如果當前塊剩余空間夠用就做兩件事記錄當前指針、把指針往后推進Size字節(jié)然后返回原來的地址。不夠用就新開一塊更大的塊。整個過程沒有任何free也沒有任何空閑列表。void* allocate(size_t Size, size_t Alignment) { if (Size Threshold) { // 超過閾值單獨分配一塊大內(nèi)存 return malloc(Size); } // 否則從當前 slab 的水位線上切一塊 uintptr_t AlignedPtr alignTo(CurPtr, Alignment); ... CurPtr AlignedPtr Size; return reinterpret_castvoid*(AlignedPtr); }生活化的類比是裝修工地按項目統(tǒng)一訂一車磚泥瓦匠用多少就從磚堆里拿多少用完的磚頭渣子不單獨清理項目整體結(jié)束后連同工地垃圾一起清走。如果每個磚頭都要單獨記賬、單獨回收成本比磚本身還高。由此引出的特性是BumpPtrAllocator不回收到單個對象粒度只支持整塊釋放。整個arena析構(gòu)時所有從它分配出去的內(nèi)存一起失效。3.2 為什么編譯器敢不調(diào)用析構(gòu)函數(shù)這一點是我覺得最反直覺的地方。BumpPtrAllocator不管對象析構(gòu)直接整塊扔掉。這在普通業(yè)務(wù)代碼里是不可接受的但在LLVM IR構(gòu)建場景下完全成立原因有三層。第一層IR節(jié)點的成員幾乎都是POD或由同一arena分配的容器。比如整數(shù)常量就是APInt加一個opcode析構(gòu)不涉及外部資源。第二層如果某個節(jié)點需要字符串內(nèi)容LLVM不會單獨給字符串分配一塊獨立內(nèi)存而是用StringMap把字符串字符數(shù)據(jù)也放進同一個arena字符串的釋放成本和arena生命周期一致。第三層IR是在pass之間構(gòu)建、轉(zhuǎn)換、銷毀的典型的生命周期邊界就是某個pass的run函數(shù)或整個Module的編譯單元。所以這里有個樸素但關(guān)鍵的判斷當整個內(nèi)存區(qū)的生命周期明確且統(tǒng)一時逐個調(diào)用析構(gòu)是沒有意義的。把析構(gòu)工作推遲到整體回收時一起處理省掉的不只是析構(gòu)調(diào)用本身還有編譯器需要維護的哪個對象還活著的記賬信息。3.3 所有權(quán)與Use列表對象之間靠引用不靠計數(shù)再往深一層LLVM對象之間的所有權(quán)關(guān)系也很特別。IR里的Value、User、Instruction相互引用構(gòu)成一個巨大的def-use圖。對象之間不是通過shared_ptr計數(shù)的而是通過Use對象顯式連接User內(nèi)部持有一組Use每個Use指向被使用的ValueValue反過來維護一個use鏈表記錄誰在用我。這種設(shè)計的核心是容器所有權(quán)?；緣K里的指令由BasicBlock的指令鏈表ilist擁有函數(shù)的基本塊由Function擁有模塊又擁有函數(shù)。刪除一條指令是調(diào)用Instruction::eraseFromParent()它從容器的鏈表上解鏈同時把def-use邊全部清理掉。它不需要像shared_ptr那樣計算引用次數(shù)因為所有權(quán)鏈條是樹狀的、清晰的。把BumpPtrAllocator和ilist放在一起看才真正理解LLVM的內(nèi)存哲學它選擇在編譯器的生命周期邊界上做整體資源管理而不是在每個對象的生命周期上做細碎管理?？臻g換時間結(jié)構(gòu)換確定性。4. 帶著編譯器內(nèi)存哲學看大模型推理的顯存治理讀LLVM內(nèi)存管理讀到有點上頭的那個周末我正好在看nano-vllm的顯存分配代碼。兩個體系對照著看產(chǎn)生了一種這不就是同一道題嗎的既視感。4.1 KV Cache的預(yù)分配恐懼癥與顯存碎片大模型推理的性能大頭在KV Cache。生成每個token時注意力層需要把歷史token的Key和Value緩存下來否則每一步都要從頭重算計算量直接平方級爆炸。代價是顯存占用隨序列長度線性增長長上下文場景下KV Cache能占到70%以上顯存。樸素實現(xiàn)是每個請求進來就給它的KV Cache預(yù)分配最大長度max_seq_len的連續(xù)顯存塊。這個方案有兩個問題。一是預(yù)分配恐懼癥多開一個請求就多預(yù)分配一大塊可實際生成幾十個token就結(jié)束的請求占不滿顯存利用率極低二是碎片問題不同請求先后結(jié)束釋放的顯存大小不一新請求分配不上即便總顯存還夠用。這種場景和編譯器里每次new/delete指令的碎片化問題本質(zhì)是同一個問題只是對象從指令換成了顯存塊。4.2 PagedAttention把KV Cache做成操作系統(tǒng)的分頁vLLM的做法是把KV Cache切成固定大小的塊block比如按block_size16個token為一塊。邏輯上屬于同一個序列的KV物理上可以分散在顯存的任意塊位置靠一張block table記錄邏輯塊到物理塊的映射。這其實就是操作系統(tǒng)的分頁思想。LLVM的arena分配器是一次性分配一大塊然后線性切片vLLM則更進一步把切片粒度固定下來允許分散映射。兩者共同點是都不依賴逐個釋放的malloc/free語義而是用池塊的方式來管理高頻分配。不同點在于推理場景需要在請求結(jié)束時會回收部分塊給其他請求所以vLLM的block pool是有塊粒度的雙向借還而LLVM的arena可以更狠直接整區(qū)回收。這種設(shè)計帶來的實際收益非常直觀。在我壓測的LLaMA-7B場景里用固定預(yù)分配方案batch一超過4就頻繁O(jiān)OM換成分塊方案后同樣的顯存能跑的并發(fā)請求數(shù)差不多翻倍。原因很簡單每個請求只在最后一塊尾部浪費幾個slot大部分塊被填滿。4.3 continuous batching像arena一樣按批收放顯存治理之外scheduler的調(diào)度策略也反過來影響顯存分配。continuous batching的核心是不讓整個batch同步綁定到同一個生命周期。一個請求生成完了立刻在它空出的位置上插入新請求不用等一批全部結(jié)束。這里就很像編譯器里pass pipeline對IR的批處理一批IR整體進入優(yōu)化管線pass逐個處理處理完一個模塊整體回收。不同的是推理引擎的batch是動態(tài)的請求在任意時刻結(jié)束、任意時刻加入。所以scheduler需要精確知道每個序列當前占用幾個塊、哪些塊可以騰出來、新請求應(yīng)該優(yōu)先從哪個池里拿塊。這套邏輯在LLVM里沒有完全對應(yīng)的東西但理解分配器按塊管理之后再看scheduler代碼會順暢很多。5. 算子自發(fā)現(xiàn)與Pass依賴分析兩張同源的編排網(wǎng)標題里還有個熱搜詞是llvm算子自發(fā)現(xiàn)。這個說法本身比較模糊但順著LLVM的pass機制和推理引擎的graph優(yōu)化去理解能看出一個很清晰的同源結(jié)構(gòu)。5.1 LLVM里Pass是怎么被發(fā)現(xiàn)的傳統(tǒng)LLVM pass通過靜態(tài)注冊表暴露自己pass定義一個llvm::RegisterPass某個Pass的全局變量命令行工具按名字查找并創(chuàng)建。新PassManagernew PM在此基礎(chǔ)上做了升級pass之間通過AnalysisManager自動發(fā)現(xiàn)并緩存分析結(jié)果。比如一個優(yōu)化pass在run函數(shù)里調(diào)用getAnalysisDominatorTree()AnalysisManager發(fā)現(xiàn)之前已經(jīng)有人算過DominatorTree直接把緩存結(jié)果給你沒算過就先跑一次分析然后緩存。這就是一種按需發(fā)現(xiàn)依賴的機制pass不需要自己維護全局狀態(tài)。依賴信息不是白拿的拿完之后還要保證分析結(jié)果在pass修改IR后仍然有效。所以pass有preserve邏輯明確聲明自己改了什么、保留了哪些分析。這和推理引擎里liveness和alias分析的使用方式很像編譯器改IR之前要知道哪些指針還活著、哪些可以重新分配推理引擎在算子替換之前也要知道tensor被誰依賴。5.2 推理引擎里的算子自發(fā)現(xiàn)與pattern匹配大模型推理側(cè)的算子自發(fā)現(xiàn)更接近這種意思從一個kernel注冊表中自動嗅探當前計算圖中可被替換/融合的子模式。PyTorch的fx graph上做subgraph rewriting是典型的實現(xiàn)方式# 偽代碼示意 pattern matcher patterns [fused_attention_pattern, fused_mlp_pattern, rms_norm_pattern] for node in fx_graph.nodes: for pattern in patterns: if pattern.match(node): replace_subgraph(node, fused_kernel)匹配到的子圖替換為融合算子比如attention融合、SwiGLU融合這個過程和LLVM里instcombine識別add(x, 1)加add(y, -1)并做常量折疊思路幾乎一模一樣。區(qū)別只在于LLVM有統(tǒng)一的IR和明確定義的pass依賴推理引擎要面對的是TensorRT/ONNX Runtime/Triton各自不同的注冊體系和graph表示。Triton里更極端一點叫autotune同一段計算邏輯生成多個不同block_size、不同num_warps的kernel變體啟動前先跑一遍benchmark自動選最優(yōu)。這也可以算一種自發(fā)現(xiàn)——從候選池里發(fā)現(xiàn)當前輸入shape下最合適的實現(xiàn)。5.3 兩張編排網(wǎng)的對照把LLVM的AnalysisManager和推理引擎的scheduler/executor放在一張表里對照脈絡(luò)會清楚很多維度LLVM新Pass管線推理引擎依賴來源AnalysisManager按需計算緩存計算圖的tensor依賴、stream依賴資源約束內(nèi)存、pass修改IR后的有效性顯存水位、KV Cache塊、stream槽位執(zhí)行單位Optimization passkernel / op / step生命周期邊界Module/Function編譯單元request/sequence/scheduling cycle自發(fā)現(xiàn)體現(xiàn)PassRegistry按名字實例化、Analysis按需生成算子注冊表、pattern matcher、autotune我讀LLVM的pass pipeline代碼時經(jīng)常腦子里還在轉(zhuǎn)著nano-vllm里scheduler對sequence的調(diào)度。兩者一個偏靜態(tài)編排一個偏動態(tài)編排但編排的核心都是搞清楚誰依賴誰、誰現(xiàn)在占用什么資源、誰能釋放什么資源。6. 跟著nano-vllm把大模型推理關(guān)鍵功能串一遍前面都是概念層面的同構(gòu)真正落地還是要回到代碼。nano-vllm是個很適合串起一整條推理鏈路的小項目。6.1 nano-vllm的目錄級功能地圖我clone下來后按這個順序讀的scheduler.pySequence和SequenceGroup的管理決定每個step哪些序列可以前向block_manager.pyKV Cache block的分配、映射、釋放attention.pyPagedAttention的入口傳入block table完成attention計算model.pyLlama結(jié)構(gòu)的forward和采樣server.py把推理循環(huán)包成HTTP服務(wù)。這個順序符合推理主循環(huán)的推進邏輯請求進來 → scheduler分塊 → 模型前向 → attention讀KV → 采樣 → 更新block狀態(tài)。每一條都對應(yīng)vLLM里的大模塊但省掉了分布式、多進程、量化、動態(tài)bucketing等工程細節(jié)適合斷點跟蹤。6.2 scheduler和block_manager整個系統(tǒng)的CPU側(cè)大腦我在讀nano-vllm時覺得最值錢的是scheduler里誰先跑和block_manager里塊從哪來這兩個問題的交互。scheduler要維護若干SequenceGroup每個Group有多個Sequencebeam search場景下會分裂。每個step決定把哪些SequenceGroup調(diào)度到GPU上時要算清楚這些序列總共需要多少個block當前空閑block夠不夠。不夠的話要么等待要么搶占低優(yōu)先級序列并釋放它們的block。這套邏輯幾乎就是操作系統(tǒng)的內(nèi)存置換只是頁面換成了KV Cache塊。關(guān)鍵代碼是每個sequence維護自己的block_tableclass Sequence: def __init__(self, seq_id, prompt_token_ids): self.seq_id seq_id self.blocks [] # 邏輯 block 索引 self.num_generated_tokens 0 self.status Status.WAITING新token產(chǎn)生后如果當前最后一個邏輯塊還有空slot直接延續(xù)滿了就到block_manager申請新塊在block_table里追加一條。這個按需追加的動作就是KV Cache不會一次性打滿顯存的直接原因。6.3 更好的學習路徑動手拆光讀代碼容易漏細節(jié)我建議三個動手操作。第一個操作是調(diào)小block池容量觀察scheduler出現(xiàn)搶占preemption的時機。把能從十幾個block壓到四五個block讓兩個并發(fā)請求搶資源看scheduler怎么把一個請求的block交換出去之后又如何恢復(fù)。這個現(xiàn)象比任何文章都直觀。第二個操作是打印每個step的block_table變化。新生成一個token打印一次邏輯塊追加、物理塊分配、slot占用率。幾輪下來你對KV Cache為什么是顯存大頭、為什么要盡量復(fù)用塊會有肌肉記憶。第三個操作是拿nano-vllm的scheduler和vLLM的Scheduler源碼做對照。nano-vllm很多分支被砍掉了vLLM里有running、swapped、waiting三個隊列的完整狀態(tài)機。對照著看能發(fā)現(xiàn)簡化版砍掉的正是工程復(fù)雜度的核心vLLM里因為要支撐超高并發(fā)序列狀態(tài)轉(zhuǎn)換需要非常嚴格的不變性保證。這套學下來再回頭去想LLVM的BumpPtrAllocator和pass管理你會發(fā)現(xiàn)它們不是兩套知識而是一套底層能力在不同領(lǐng)域的投影理解資源邊界、理解生命周期、理解依賴關(guān)系。編譯器和推理引擎的高性能代碼最終都建立在把這三件事想得足夠清楚之上。最后再分享一個我自己的習慣讀這類底層基礎(chǔ)設(shè)施代碼時不要只盯著算法本身把每個數(shù)據(jù)結(jié)構(gòu)的生命周期圖畫出來——誰創(chuàng)建、誰持有、誰釋放、釋放后誰還能引用。LLVM這么畫能看懂IR對象的死活推理引擎這么畫能看懂顯存塊的流轉(zhuǎn)。畫過幾張之后再看任何高性能系統(tǒng)的內(nèi)存管理基本都不會迷路。