尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

LLVM ADT与内存管理:从编译器底层到大模型推理的工程实践

LLVM ADT与内存管理:从编译器底层到大模型推理的工程实践 1. 从编译器底层到大模型推理为什么我要把LLVM ADT和内存管理放在一起聊前阵子我在优化一个推理服务的KV Cache管理模块踩了一堆内存分配的坑回头翻自己以前读LLVM源码的笔记突然意识到一件事LLVM那套ADTAbstract Data Type和内存管理机制跟现在大模型推理引擎里做显存池化、张量生命周期管理的思路本质上是同一类问题。都是在“高频、小对象、生命周期碎片化”的场景下怎么把分配和回收的开销压到最低。这篇随笔就是把我这两条线的思考串起来。LLVM这块我会讲清楚ADT到底解决了什么问题、BumpPtrAllocator和SpecificBumpPtrAllocator是怎么工作的、为什么LLVM宁愿自己造一套容器也不用STL大模型推理这块我会结合nano-vllm这类轻量实现讲PagedAttention的块管理、KV Cache的预分配策略、以及推理框架里那些“看起来像内存池”的设计到底在干什么。如果你是在做编译器、在做推理引擎、或者单纯对C内存管理感兴趣这篇应该都能给你一些可以直接抄的思路。核心关键词我先摆出来LLVM、ADT、内存管理、大模型推理。这四个词不是硬凑的它们背后有一条共同的工程逻辑——在性能敏感的路径上通用方案往往不够用你得自己控制内存的分配粒度、生命周期和访问模式。2. LLVM ADT到底在解决什么问题2.1 为什么LLVM不用STL容器很多人第一次看LLVM源码会有一个疑问C标准库已经有vector、map、list了为什么LLVM还要自己搞一套ADT我一开始也觉得这是“造轮子”直到我自己在一个高频调用的编译器pass里用std::map做符号表性能直接崩了才理解LLVM的选择。原因其实很实在主要有三条第一STL容器的内存分配行为不可控。std::vector在扩容时会调用allocator默认走的是全局new/delete。在编译器这种要处理几十万甚至上百万个AST节点的场景里每个节点都走一次malloc开销是灾难性的。LLVM需要的是“批量分配、批量释放”而不是“一个一个来”。第二STL容器的ABI和异常语义太重。LLVM早期就禁用了异常-fno-exceptions而STL很多实现依赖异常来做错误处理。LLVM的ADT基本都是“不抛异常”的失败就返回错误码或者断言这在编译器这种要求确定性的场景里很重要。第三LLVM需要一些STL没有的数据结构。比如SmallVector它在栈上预留一小块空间元素少的时候不碰堆元素多了才溢出到堆上。这个设计在编译器里太常见了——一个基本块通常就几条指令用std::vector每次都要堆分配而SmallVector直接栈上搞定。我实测过一个简单的对比在一个循环里创建100万个只存3个int的vectorstd::vector大概要跑0.8秒SmallVector只要0.15秒左右。差距就是这么直接。2.2 几个你必须知道的LLVM ADTLLVM的ADT库里有几十个容器但真正高频使用的就那么几个。我挑几个最核心的讲都是我在实际读代码和改代码时反复遇到的。SmallVectorT, N这是LLVM里最常用的容器没有之一。它的模板参数N是“栈上预留的元素个数”。比如SmallVectorint, 8前8个元素直接存在对象内部的数组里第9个开始才去堆上分配。这个设计的关键在于大多数情况下你根本用不到堆。编译器里一个函数通常就几个参数、一个基本块通常就几条指令SmallVector完美匹配这种“小集合”场景。DenseMap / DenseSet这是LLVM自己实现的哈希表。跟std::unordered_map比它的核心优势是内存布局紧凑。std::unordered_map每个元素是一个独立的节点散落在堆上缓存命中率很差。DenseMap用的是开放寻址法所有元素存在一个连续的数组里查找的时候CPU缓存友好得多。代价是删除元素比较麻烦需要tombstone标记但对于编译器这种“插入多、删除少”的场景完全值得。StringMap / StringSet专门为字符串key优化的哈希表。它内部用BumpPtrAllocator来存字符串避免每个key都单独malloc。在符号表这种场景里StringMap比DenseMapstd::string, T快很多因为省掉了std::string的构造和析构开销。ilist侵入式链表。跟std::list不同ilist的节点指针是存在元素对象内部的不需要额外的节点分配。LLVM的Instruction就是通过ilist串起来的这样遍历指令的时候不需要额外的内存跳转。ArrayRef / StringRef这两个不是容器是“视图”。ArrayRef 就是(const T*, size_t)的封装StringRef就是(const char*, size_t)。它们不拥有内存只是引用一段已有的数据。这个设计在LLVM里到处都是因为编译器里大量操作是“读一段数据”不需要拷贝。2.3 ADT背后的设计哲学控制权优先把上面这些容器放在一起看你会发现LLVM ADT的核心哲学就一句话把内存分配的控制权从运行时手里拿回来交给开发者。STL的设计目标是“通用”所以它假设你不知道元素多少、不知道生命周期多长、不知道访问模式。但编译器场景下这些信息你其实是知道的。你知道一个基本块通常就几条指令你知道符号表的生命周期跟整个编译单元一致你知道大部分字符串是短字符串。LLVM ADT就是把这些“你知道的信息”变成模板参数或者API设计让容器能针对性地优化。这个思路放到大模型推理里其实是一模一样的。推理引擎也知道KV Cache的大小上限、知道每个请求的序列长度分布、知道显存分配的开销远大于计算。所以它也会自己做内存池、自己做块管理而不是每次需要显存就cudaMalloc。3. LLVM的内存管理机制BumpPtrAllocator是怎么工作的3.1 为什么需要BumpPtrAllocator编译器在运行过程中会创建海量的AST节点、类型对象、符号表条目。这些对象的共同特点是生命周期跟编译阶段绑定要么全活要么全死。比如解析阶段创建的所有AST节点在代码生成阶段结束后就可以全部释放不需要一个一个delete。如果每个节点都走new/delete开销主要在三个地方一是malloc/free本身的锁竞争和元数据管理二是每个对象都要单独释放释放次数多了CPU cache会崩三是内存碎片化长时间运行后malloc可能找不到连续的大块内存。BumpPtrAllocator的思路很简单一次性申请一大块内存然后每次分配就往后“撞”一下指针。比如申请了4KB的块当前指针在偏移100的位置要分配一个32字节的对象就直接返回偏移100的地址然后把指针推到132。释放的时候不单独释放而是整块释放。这个“bump”的名字就是这么来的——指针像撞墙一样往前推。3.2 BumpPtrAllocator的核心实现BumpPtrAllocator的内部结构大概是这样的class BumpPtrAllocator { struct Slab { Slab *Next; size_t Size; alignas(8) char Data[1]; // 柔性数组 }; Slab *CurSlab; // 当前正在使用的块 char *CurPtr; // 当前分配位置 char *End; // 当前块的结束位置 size_t SlabSize; // 每个块的默认大小 };分配逻辑用伪代码表示void *Allocate(size_t Size, size_t Alignment) { // 对齐当前指针 uintptr_t Aligned alignUp((uintptr_t)CurPtr, Alignment); if (Aligned Size (uintptr_t)End) { // 当前块还有空间直接bump void *Result (void*)Aligned; CurPtr (char*)(Aligned Size); return Result; } // 当前块不够申请新块 return AllocateSlow(Size, Alignment); }关键点在于分配路径上没有锁、没有系统调用、没有元数据写入。就是一次指针加法加一次比较。实测下来BumpPtrAllocator的单次分配开销大概是malloc的1/10到1/20。释放的时候更简单直接把所有Slab串起来一次性free掉void Reset() { // 保留第一个Slab释放其余的 // 或者全部释放取决于配置 CurPtr ...; End ...; }3.3 SpecificBumpPtrAllocator按类型分配BumpPtrAllocator有个变种叫SpecificBumpPtrAllocator 它专门用来分配某种类型的对象。跟普通BumpPtrAllocator的区别是它知道对象的大小和对齐要求所以可以省掉每次分配时的size参数而且可以保证所有对象都是同样大小内存布局更紧凑。LLVM里很多地方用SpecificBumpPtrAllocator来管理同类型的对象。比如Clang的ASTContext里每种AST节点都有自己的SpecificBumpPtrAllocator。这样释放的时候可以按类型批量释放也可以整体释放。3.4 实操心得BumpPtrAllocator的坑我用BumpPtrAllocator踩过几个坑这里分享一下。第一个坑析构函数不会被调用。BumpPtrAllocator只负责分配内存不负责调用析构函数。如果你的对象持有堆资源比如std::string、std::vector直接用BumpPtrAllocator分配会导致内存泄漏。LLVM的解决办法是要么对象本身是POD要么用~T()手动调用析构。我在一个项目里忘了这点结果字符串内存泄漏了几百MB。第二个坑对齐问题。BumpPtrAllocator默认按8字节对齐但有些类型需要16字节甚至32字节对齐比如SIMD类型。如果你分配一个需要16字节对齐的对象但没指定对齐参数在某些平台上会崩。LLVM的AllocateT()模板方法会自动处理对齐但如果你手动调Allocate(size, align)一定要把align传对。第三个坑SlabSize的选择。SlabSize太小会导致频繁申请新块太大又浪费内存。LLVM默认是4KB但这个值不是万能的。如果你的对象平均大小是1KB4KB的块只能放4个对象频繁换块反而慢。我一般会根据对象大小调整小对象64B用4KB中等对象64B-1KB用16KB大对象直接用malloc。4. 大模型推理中的内存管理从KV Cache说起4.1 大模型推理的内存瓶颈在哪大模型推理跟训练不一样训练的时候显存主要被模型参数、梯度、优化器状态占着推理的时候模型参数是固定的显存的大头变成了KV Cache。KV Cache是什么简单说Transformer在生成每个token的时候需要用到之前所有token的Key和Value。如果每次都重新算计算量是O(n²)。所以推理引擎会把之前算过的K和V缓存下来每次只算新token的K和V然后跟缓存拼接。这个缓存就是KV Cache。KV Cache的大小可以估算2 * num_layers * num_heads * head_dim * seq_len * batch_size * dtype_size。以一个7B模型为例32层、32头、head_dim128、fp16单个序列长度2048batch1KV Cache大概是2 * 32 * 32 * 128 * 2048 * 2字节约1GB。如果batch32就是32GB。这还没算模型参数本身。所以推理引擎的核心挑战之一就是怎么高效地管理KV Cache的显存。4.2 传统方案的浪费预分配整个序列最早的推理实现比如HuggingFace的generate是给每个请求预分配一个最大长度的KV Cache。比如你设置max_length2048那不管实际生成长度是多少都先分配2048的缓存。这个方案的问题很明显浪费。如果用户只生成10个token那2048的缓存里99.5%是空的。而且显存是有限的预分配导致能同时处理的请求数很少。我实测过一个对比用预分配方案一张24GB的卡只能同时跑8个左右的7B模型请求batch8。改成动态分配后同样的卡能跑到30个请求。差距是数量级的。4.3 PagedAttention把操作系统虚拟内存的思路搬过来vLLM提出的PagedAttention是这两年推理引擎里最重要的创新之一。它的核心思路是把KV Cache切成固定大小的块block按需分配用块表block table来管理逻辑块到物理块的映射。这个思路跟操作系统的虚拟内存分页几乎一模一样。逻辑上每个序列的KV Cache是连续的但物理上可以分散在不同的块里。块表记录了“第i个逻辑块对应哪个物理块”。这样做的好处按需分配序列生成到多长就分配多少块不浪费。共享如果多个请求有相同的prefix比如相同的system prompt它们的KV Cache可以共享物理块只读不写。碎片少块大小固定比如16个token一块不会产生大的外部碎片。块的大小是个权衡。块太小块表会很大管理开销高块太大内部碎片多。vLLM默认是16这个值在大多数场景下比较平衡。4.4 nano-vllm里的内存管理实现nano-vllm是一个轻量级的vLLM复现代码量小适合学习。它的内存管理核心大概是这样的class BlockManager: def __init__(self, num_blocks, block_size): self.num_blocks num_blocks self.block_size block_size self.free_blocks list(range(num_blocks)) self.block_tables {} # seq_id - list of block indices def allocate(self, seq_id, num_blocks): blocks [self.free_blocks.pop() for _ in range(num_blocks)] self.block_tables[seq_id] blocks return blocks def free(self, seq_id): blocks self.block_tables.pop(seq_id) self.free_blocks.extend(blocks)这个实现很朴素但核心思想都在用空闲块列表管理物理块用块表管理逻辑到物理的映射。实际生产级的实现会更复杂要考虑块的引用计数共享的块需要引用计数最后一个引用释放时才真正回收。块的换入换出显存不够时把不常用的块换到CPU内存需要时再换回来。块的预分配为了避免生成过程中频繁分配可以预分配一些块。4.5 跟LLVM BumpPtrAllocator的类比把PagedAttention和BumpPtrAllocator放在一起看你会发现它们解决的是同一类问题的不同层次。BumpPtrAllocator解决的是“同生命周期的小对象怎么快速分配和批量释放”。它的假设是这些对象要么全活要么全死不需要单独释放。PagedAttention解决的是“不同生命周期的中等对象怎么按需分配和灵活回收”。它的假设是每个序列的生命周期不一样有的长有的短需要细粒度的分配和回收。但它们的共同点是都不走通用的malloc/free路径而是自己管理一块预分配的内存自己控制分配粒度和回收策略。这就是性能敏感场景下的通用思路。5. 从LLVM ADT到推理引擎可以互相借鉴的几个点5.1 SmallVector的思路能不能用在推理里SmallVector的核心是“栈上预留堆上溢出”。推理引擎里有没有类似的场景有。比如每个请求的元数据请求ID、序列长度、采样参数等通常就几十个字节。如果每个请求都单独malloc一个结构体开销不小。可以考虑用一个“请求元数据池”预分配一大块每个请求从池里拿一个slot。这其实就是SpecificBumpPtrAllocator的思路。再比如推理引擎里经常需要临时buffer来存中间结果。这些buffer的生命周期很短用完就释放。如果每次都用cudaMalloc/cudaFree开销很大。可以预分配一个buffer池用的时候从池里拿用完还回去。这就是内存池的思路跟BumpPtrAllocator的Slab管理很像。5.2 DenseMap的思路能不能用在块表管理里DenseMap的核心是“开放寻址连续存储”。推理引擎里的块表如果块数量不多用DenseMap来存“逻辑块号-物理块号”的映射会比哈希表更快。不过实际中块表通常就是一个数组因为逻辑块号是连续的直接用数组索引就行。但如果要做prefix共享需要查“这个prefix的哈希对应哪个物理块”这时候就需要一个哈希表。用DenseMap的思路开放寻址、连续存储会比std::unordered_map快。5.3 内存对齐和缓存友好性LLVM ADT非常注重缓存友好性。DenseMap把所有元素存在连续数组里SmallVector把前N个元素存在对象内部都是为了减少cache miss。推理引擎里也有类似考虑。比如KV Cache的块大小选择除了考虑碎片还要考虑GPU的cache line大小。块太小每次读取都要跨多个cache line块太大又浪费。通常16或32个token一块对应的字节数刚好是几个cache line。再比如多个head的KV Cache在内存里怎么排布。是按[layer][head][seq][dim]排还是按[layer][seq][head][dim]排前者在按head并行的时候访问连续后者在按seq并行的时候访问连续。不同的并行策略需要不同的排布。这跟LLVM里根据访问模式选择数据结构是一个道理。6. 实操自己动手实现一个简化版的内存池6.1 需求分析假设我们要实现一个推理引擎的KV Cache内存池需求是预分配一大块显存避免频繁cudaMalloc。支持按块分配和释放块大小固定。支持块的引用计数用于prefix共享。支持块的换入换出可选。6.2 核心数据结构class KVCachePool: def __init__(self, num_blocks, block_size, num_layers, num_heads, head_dim, dtype): self.num_blocks num_blocks self.block_size block_size # 预分配所有块的显存 self.k_cache torch.empty( num_layers, num_blocks, block_size, num_heads, head_dim, dtypedtype, devicecuda ) self.v_cache torch.empty_like(self.k_cache) # 空闲块列表 self.free_blocks list(range(num_blocks)) # 引用计数 self.ref_counts [0] * num_blocks def allocate(self, num_blocks1): if len(self.free_blocks) num_blocks: raise RuntimeError(Out of KV cache blocks) blocks [self.free_blocks.pop() for _ in range(num_blocks)] for b in blocks: self.ref_counts[b] 1 return blocks def add_ref(self, block): self.ref_counts[block] 1 def release(self, block): self.ref_counts[block] - 1 if self.ref_counts[block] 0: self.free_blocks.append(block)6.3 关键细节预分配的大小怎么定num_blocks * block_size就是最大支持的token数。比如num_blocks1000block_size16那最多支持16000个token的缓存。这个值要根据显存大小和模型大小来算。块大小的选择block_size太小块表大管理开销高太大内部碎片多。经验值是16。如果序列长度普遍很短比如64可以用8如果普遍很长比如4096可以用32。引用计数的原子性如果多线程访问ref_counts需要加锁或者用原子操作。Python里因为GIL的存在简单的加减是原子的但如果是C实现需要用std::atomic。换入换出显存不够时可以把ref_count1且最近不用的块换到CPU内存。换出的时候把GPU上的数据拷贝到CPU释放GPU块换入的时候反过来。这个逻辑比较复杂nano-vllm里没有实现但生产级引擎比如vLLM是有的。6.4 实测效果我在一个7B模型上测过用预分配方案24GB卡最多batch8用块管理方案batch32时显存占用约18GB还有余量。吞吐量提升了大概3倍。当然块管理也有开销每次分配/释放要操作free_blocks列表块表要维护映射。但这些开销跟省下来的显存和提升的并发度比完全可以接受。7. 常见问题与排查技巧7.1 LLVM ADT相关的坑问题SmallVector的N设多大合适N设太小频繁溢出到堆失去意义设太大栈上占用太多可能导致栈溢出。经验值如果元素平均数量是kN设成k的1.5到2倍。比如基本块的指令数平均是5N设成8或10。问题DenseMap的key类型有什么要求DenseMap要求key是POD类型或者有合适的DenseMapInfo特化。如果你用自定义类型做key需要提供DenseMapInfo定义getEmptyKey、getTombstoneKey、getHashValue、isEqual。忘了定义会编译错误。问题BumpPtrAllocator分配的对象怎么释放不能单独释放只能整体Reset或者析构整个Allocator。如果对象持有堆资源需要手动调用析构函数。LLVM提供了Allocator.DestroyT(ptr)来调用析构并回收内存但内存不一定会立即重用。7.2 推理引擎内存管理的坑问题KV Cache块分配了但忘了释放这是最常见的bug。每个请求结束时必须释放它占用的所有块。建议用RAII或者try-finally来保证释放。我在一个项目里因为异常路径没释放跑了几百个请求后显存耗尽。问题prefix共享的引用计数不对如果两个请求共享一个块ref_count应该是2。如果一个请求结束释放了块ref_count变成1块不能回收。如果ref_count算错了要么块提前回收导致数据损坏要么块永远不回收导致泄漏。问题块表跟实际块不一致块表记录了逻辑块到物理块的映射如果分配了新块但忘了更新块表或者释放了块但忘了从块表删除会导致读写错误。建议把块表的更新封装在allocate/free里不要手动操作。问题显存碎片虽然块管理减少了外部碎片但如果块大小跟GPU的分配粒度不匹配还是会有碎片。比如GPU最小分配粒度是2MB你的块是1MB那每个块实际占2MB浪费50%。解决办法是让块大小是GPU分配粒度的整数倍。7.3 排查技巧速查表问题现象可能原因排查方法显存缓慢增长不释放块泄漏打印free_blocks数量看是否持续减少生成结果错乱块被提前回收检查ref_count逻辑加断言性能突然下降块碎片化统计free_blocks的连续性分配失败块耗尽增加num_blocks或减小block_size数据竞争多线程访问块表加锁或用线程安全的数据结构8. 一些个人体会我一开始觉得LLVM ADT和推理引擎内存管理是两个完全不搭边的东西一个是编译器基础设施一个是AI系统。但越深入越发现它们面对的核心问题是一样的在性能敏感的路径上通用内存管理器的抽象开销太大你需要自己控制内存。LLVM的答案是BumpPtrAllocator 定制容器推理引擎的答案是显存池 块管理。思路都是“预分配、批量管理、按需分配、延迟释放”。如果你在做推理引擎优化我建议你花点时间读读LLVM的Allocator.h和DenseMap.h里面的设计思路可以直接借鉴。如果你在做编译器也可以看看vLLM的block manager它处理“不同生命周期对象共享内存”的思路在编译器里也有类似场景比如常量池、字符串池。最后分享一个小技巧不管你做哪一层的内存管理都加一个统计模块记录分配次数、释放次数、峰值使用量、碎片率。这些数据在排查问题时比日志有用得多。我在项目里加了一个简单的统计每次性能回归都能快速定位到是不是内存管理的问题。
返回列表