资讯动态

LLVM ADT与内存管理:从DenseMap陷阱到大模型推理内存池

发布时间:2026/10/21 17:06:40 来源:尧图企业网站定制
1. 从一次编译期崩溃说起为什么要啃LLVM的ADT和内存管理前阵子帮一个做静态分析工具的朋友排查问题他的工具基于LLVM做了一套自定义的IR遍历Pass跑小项目一直很稳结果换到一个稍大的代码库上编译到一半直接段错误堆栈打出来指向的是DenseMap的某个内部bucket操作。他第一反应是LLVM的容器有bug我让他先别急着甩锅把代码里几处DenseMap的key类型和生命周期打印出来一看问题很清楚他用了一个指向临时对象的指针当key对象析构之后指针还挂在map里后续rehash的时候访问了已经释放的内存。这个坑其实特别典型。很多人从STL转过来写LLVM相关代码脑子里默认std::unordered_map那一套语义觉得容器就是容器能存能取就行。但LLVM的ADTAbstract Data Types库从设计之初就不是为了通用而生的它是为了编译器的性能特征而生的——大量小对象、频繁的插入查找、对内存布局极度敏感、对分配次数锱铢必较。你如果不理解它背后的内存管理假设写出来的代码要么慢得离谱要么就像我朋友那样直接崩掉。这篇随笔想聊三件事LLVM ADT里几个核心容器到底是怎么管内存的、BumpPtrAllocator这套分配机制为什么长这样、以及顺着这条线往下大模型推理引擎在内存管理上其实和编译器面对的是同一类问题。前两部分是硬核的工程细节第三部分是我自己从编译器领域切到大模型推理时的一些对照思考不一定严谨但我觉得对理解内存管理这件事的整体脉络有帮助。适合谁看如果你正在写LLVM Pass、做基于LLVM的编译器或工具链开发或者你对高性能场景下怎么设计容器和分配器这个话题感兴趣那这篇应该能给你一些能直接上手的东西。如果你只是听说过LLVM但没写过建议先把官方那套Kaleidoscope教程过一遍再回来不然有些细节会有点跳。2. DenseMap和SmallVector两个最容易被误用的ADT2.1 DenseMap的bucket布局与空/墓碑状态机先从最常被误用的DenseMap说起。它的名字里dense这个词是关键——和std::unordered_map那种每个元素独立分配一个节点的做法不同DenseMap把所有键值对连续存在一块数组里通过开放寻址open addressing解决冲突。这个设计带来的直接好处是缓存友好遍历的时候内存是连续的CPU预取器能很好地工作坏处是删除操作不能真的把元素抹掉否则会破坏探测链。所以DenseMap用了三态标记每个bucket要么是空从未使用、要么是占用、要么是墓碑tombstone曾经占用但被删除了。查找的时候遇到空就停遇到墓碑继续往后探。这个状态机是理解它所有行为的基础。// 简化示意非LLVM真实源码 enum BucketState { Empty, Tombstone, Full }; template typename KeyT, typename ValueT struct Bucket { KeyT key; ValueT value; unsigned state; // 实际实现里用key/value的特殊值编码省空间 };真实实现更精巧它不额外存state字段而是用key和value的特定值来编码状态比如key等于某个哨兵值就代表空。这样每个bucket省下4个字节在大map上就是可观的内存节省。这个技巧你在自己设计紧凑数据结构时完全可以借鉴。那什么时候会触发rehash当占用墓碑的数量超过bucket总数的3/4时LLVM里这个阈值是写死的DenseMap的getMaxBucketCount相关逻辑就会重新分配一块更大的数组把所有有效元素重新插入。注意这里墓碑也算进负载因子因为墓碑多了探测链会变长查找性能会退化。这就是为什么我朋友那个bug会在稍大的代码库上才暴露——小项目map没满不触发rehash悬垂指针就一直没被访问到。2.2 用指针当key的陷阱生命周期必须由你兜底DenseMap对key的要求和STL不一样。STL的unordered_map要求key有hash和equalDenseMap则依赖DenseMapInfoKeyT这个trait来提供getEmptyKey()、getTombstoneKey()和isEqual()。对于指针类型LLVM内置的DenseMapInfoT*直接用指针地址做hash用地址相等做比较。问题就出在这它只比较地址不关心地址指向的对象是否还活着。你用一个栈上临时对象的地址当key对象出了作用域地址可能被复用也可能指向已释放内存但map完全不知道。rehash的时候它会去读这个key读到的就是垃圾。正确的做法有几种。如果key的生命周期确实比map短那就别用裸指针用StringRef配合一个稳定的字符串池或者干脆用值类型。如果非要用指针确保被指向的对象由某个更长寿的owner持有比如放在BumpPtrAllocator里分配的对象只要allocator不销毁对象就一直有效。// 危险写法 DenseMapconst std::string*, int BadMap; { std::string tmp hello; BadMap[tmp] 1; // tmp析构后这个key就是悬垂的 } // 相对安全的写法key指向长生命周期对象 BumpPtrAllocator Alloc; DenseMapconst std::string*, int GoodMap; auto* Stable new (Alloc) std::string(hello); GoodMap[Stable] 1; // 只要Alloc活着Stable就有效提示DenseMap的迭代顺序是不确定的而且rehash之后顺序会变。任何依赖遍历顺序的代码都是错的这一点和std::unordered_map一样但很多人写LLVM代码时会忘。2.3 SmallVector的小到底省在哪SmallVectorT, N是另一个高频误用点。它的设计目标很明确对于元素数量不超过N的场景完全不做堆分配元素直接存在对象内部的inline buffer里超过N了才去堆上分配并把已有元素搬过去。这个N的选择是有讲究的。N太大对象本身占的空间就大如果这个SmallVector是某个大结构体的成员会白白浪费栈空间N太小稍微多一点元素就触发堆分配省了个寂寞。LLVM自己的代码里N一般取4、8、16这种值取决于典型使用场景。// 典型用法函数参数列表大多数函数参数不超过8个 SmallVectorValue*, 8 Args; // 如果N设成64每个这种对象就占512字节64*8 // 一个递归下降的parser里几百层嵌套栈直接爆我见过有人为了避免堆分配把N设得特别大结果栈溢出。这里有个经验法则SmallVector作为局部变量时N乘以sizeof(T)最好不要超过128字节作为结构体成员时要算一下这个结构体被大量实例化时的总内存开销。还有一个坑是SmallVector的引用失效。它和std::vector一样扩容之后所有迭代器、引用、指针都失效。但因为它小的时候不分配很多人会误以为它像数组一样稳定。不是的一旦超过N触发扩容之前拿到的data()指针就废了。3. BumpPtrAllocator为什么编译器偏爱只分配不释放3.1 从malloc的代价说起要理解BumpPtrAllocator下面简称BPA得先想清楚通用分配器malloc/free贵在哪。一次malloc要维护空闲链表、要处理碎片、要加锁多线程场景、要做size class分类这些开销在分配-释放模式不规则的时候是必要的但在编译器这种场景下往往是纯浪费。编译器前端处理一个源文件时会创建海量的AST节点、类型对象、符号表条目。这些对象的生命周期高度一致要么在整个编译单元结束时统一销毁要么在某个Pass结束时统一销毁。几乎没有这个节点先创建、那个节点后创建然后中间某个时刻单独释放中间那个的需求。这种批量创建、批量销毁的模式用通用分配器就是杀鸡用牛刀。BPA的思路极其简单维护一块大内存slab每次分配就把指针往后挪size字节对齐一下返回旧指针。释放不存在的单个对象不释放。要清理就整块slab一起还回去。// BPA分配的核心逻辑简化版 void* Allocate(size_t Size, size_t Alignment) { uintptr_t Aligned alignUp(CurPtr, Alignment); if (Aligned Size EndPtr) { // 当前slab不够了申请新slab grow(); Aligned alignUp(CurPtr, Alignment); } void* Result reinterpret_castvoid*(Aligned); CurPtr Aligned Size; return Result; }就这么几行。没有空闲链表没有锁单线程场景没有size class。分配就是指针加法快到离谱。代价是你不能单独free一个对象析构函数也不会自动调用。3.2 slab大小与对齐那些影响性能的细节BPA默认的slab大小是4KBDefaultSlabSize但这个值可以调。slab太小大对象频繁触发新slab申请slab太大小规模编译会浪费内存。LLVM里有些场景会显式指定更大的slab比如处理超大源文件时。对齐这件事也有讲究。BPA的Allocate接受一个Alignment参数默认是alignof(std::max_align_t)也就是16字节64位平台。如果你分配的对象有更严格的对齐要求比如某些SIMD类型要32字节对齐必须显式传进去否则拿到的指针可能不对齐访问时直接崩。// 分配一个需要32字节对齐的对象 void* P Alloc.Allocate(sizeof(MySIMDType), 32);还有一个容易被忽略的点BPA分配的内存不会自动调用析构函数。如果你往里放的是有非平凡析构的对象比如持有std::string的结构体必须自己记录哪些对象需要析构在销毁allocator之前手动调用。LLVM提供了BumpPtrAllocator配合SpecificBumpPtrAllocatorT的用法后者会在销毁时对所有分配过的T调用析构。SpecificBumpPtrAllocatorMyClass Alloc; MyClass* Obj Alloc.Allocate(); // 分配但不构造 new (Obj) MyClass(...); // placement new 构造 // Alloc析构时会对所有分配过的MyClass调用~MyClass()注意SpecificBumpPtrAllocator只析构不释放单块内存——内存还是整块还回去的。这个语义要搞清楚不然会以为它能像普通容器那样精细管理。3.3 什么时候不该用BPABPA不是万能的。如果你的对象生命周期不规则比如一个长期存活的符号表里偶尔要删几个条目用BPA就会导致内存只增不减。这种情况下要么用DenseMap配合墓碑标记逻辑删除要么老老实实用通用分配器。另一个不适合的场景是对象大小差异极大。BPA的slab是固定大小的如果你一会儿分配8字节、一会儿分配1MB大对象会浪费大量slab空间。LLVM的BumpPtrAllocator对大对象有特殊处理超过slab一半的直接单独分配但如果你自己实现类似机制这点要考虑进去。注意BPA不是线程安全的。多线程环境下每个线程要有自己的BPA或者在外面加锁。LLVM的ThreadSafeBumpPtrAllocator如果版本里有或者自己包一层是常见做法。4. 从编译器到大模型推理内存管理的同构性4.1 KV Cache和AST节点不同场景同一类问题聊完LLVM这套东西我想说说为什么我觉得它和大模型推理有可比性。表面上看编译器处理的是符号和IR推理引擎处理的是张量和注意力八竿子打不着。但如果你把视角拉到内存管理这一层会发现它们面对的核心矛盾惊人地相似。大模型推理里最吃内存的是什么KV Cache。自回归生成时每生成一个token就要把当前层的Key和Value缓存下来供后续token的注意力计算使用。序列越长KV Cache越大而且是只增不减的——除非整个序列生成完否则中间不会释放任何一块。这和编译器前端创建AST节点是不是一模一样都是批量创建、生命周期一致、中间不单独释放。那推理引擎怎么管KV Cache主流做法是预分配一块大显存然后按需切分。这本质上就是BumpPtrAllocator的思路一次性要一大块内部用指针偏移来分配不做精细的free。vLLM那套PagedAttention更进一步把KV Cache切成固定大小的block用页表来映射逻辑位置到物理block这其实是在解决BPA的一个固有问题——内存碎片和浪费。4.2 内存池化推理引擎里的slab分配推理引擎里还有个东西叫内存池memory pool专门用来管理中间激活值activation的分配。一次前向传播会产生大量中间张量这些张量的生命周期很短用完就扔。如果每个都走CUDA的cudaMalloc/cudaFree光是分配开销就能把推理速度拖垮。所以引擎会预先分配一大块显存然后自己维护一个空闲块列表。分配时从列表里找合适大小的块释放时还回去。这不就是通用分配器的思路吗只不过它跑在GPU上管理的是显存而不是主机内存。有意思的是很多推理引擎的内存池实现里能看到BPA的影子它们也会把内存分成不同size class的slab小对象走小slab大对象走大slab避免碎片。区别在于推理引擎需要支持释放所以比BPA复杂但核心的slab指针偏移思想是一致的。4.3 一个具体的对照DenseMap的墓碑 vs PagedAttention的block回收再举个更细的对照。前面说DenseMap删除元素时不能真删要留墓碑否则破坏探测链。PagedAttention里也有类似的问题一个序列生成完了它占用的KV Cache block要回收但block在物理上是不连续的回收之后怎么让后续序列复用vLLM的做法是维护一个空闲block列表序列结束时把它的block还回列表。但这里有个细节block里可能还有旧数据新序列复用时需要覆盖。这和DenseMap的墓碑有点像——墓碑占着位置但不能直接用得等rehash时统一清理。PagedAttention的block回收更直接因为它有显式的空闲列表不需要靠探测链来找位置。这个对照让我意识到内存管理的本质问题在不同领域是相通的怎么在分配快和不浪费之间找平衡怎么处理生命周期不一致的对象怎么在碎片和性能之间取舍。LLVM的答案偏向分配快、不释放推理引擎的答案偏向池化显式回收但底层逻辑是一套。5. 几个我踩过的坑和对应的排查思路5.1 DenseMap的key类型选错导致编译期报错有一次我想用DenseMapstd::pairint, int, Value结果编译报了一堆模板错误核心是找不到DenseMapInfostd::pairint,int的特化。LLVM没有为std::pair提供默认的DenseMapInfo因为pair的空和墓碑状态没法用pair本身的值表示。解决办法是自己写一个特化template struct DenseMapInfostd::pairint, int { static std::pairint, int getEmptyKey() { return {-1, -1}; } static std::pairint, int getTombstoneKey() { return {-2, -2}; } static unsigned getHashValue(const std::pairint, int P) { return DenseMapInfoint::getHashValue(P.first) ^ DenseMapInfoint::getHashValue(P.second); } static bool isEqual(const std::pairint, int A, const std::pairint, int B) { return A B; } };这里的关键是空和墓碑的取值必须是你业务里不会用到的值。如果{-1,-1}在你的场景里是合法key那就得换别的哨兵值。这个坑我在两个不同项目里都见过都是因为哨兵值和业务值撞了。5.2 BumpPtrAllocator忘记对齐导致SIMD崩溃前面提过对齐的事我自己真踩过。当时在写一个基于LLVM的向量化Pass往BPA里分配一个包含__m256成员的结构体没传对齐参数结果在某些机器上跑得好好的换一台就崩。原因是BPA默认16字节对齐而__m256要32字节。这种bug特别难查因为崩溃位置和真正的原因隔了十万八千里。排查思路任何涉及SIMD、原子操作、cache line对齐的分配都要显式检查对齐参数。可以在分配后加个断言void* P Alloc.Allocate(Size, 32); assert(reinterpret_castuintptr_t(P) % 32 0 alignment violated);这个断言在debug构建里几乎零成本但能帮你省下几小时的调试时间。5.3 推理引擎里KV Cache预分配过大导致OOM这个坑是从编译器转过来做推理时踩的。我按编译器的思路觉得预分配大块内存总是好的就把KV Cache的预分配量设得很大。结果在小显存卡上直接OOM因为预分配是立即占用的不像BPA那样按需增长。后来改成按需增长分块预分配初始只分配一个较小的块序列变长时再追加。这和SmallVector的思路一模一样——小的时候用inline buffer大了再堆分配。看来不管哪个领域按需增长都是省内存的不二法门。提示推理引擎的KV Cache预分配量要根据目标硬件和典型序列长度来定。盲目设大值在开发机上可能没事部署到边缘设备就炸了。6. 写LLVM相关代码时我会遵守的几条规矩聊了这么多原理和坑最后分享几条我自己写LLVM相关代码时雷打不动的习惯都是被坑出来的。第一条能用值类型就别用指针当key。DenseMapStringRef, X比DenseMapconst std::string*, X安全得多前提是StringRef指向的字符串生命周期够长。如果字符串是动态生成的用一个BumpPtrAllocator或者StringSaver来持有它们。第二条SmallVector的N值要算过再定。别拍脑袋写个16看看这个vector典型存多少元素算一下对象大小再决定。我一般会写个注释说明为什么选这个N。第三条BPA分配的对象如果有非平凡析构用SpecificBumpPtrAllocator。别自己手动记列表容易漏。如果对象没有析构需求POD类型直接用BumpPtrAllocator最省事。第四条任何涉及内存布局的代码debug构建里加断言。对齐、大小、生命周期能断言的都断言上。这些断言在release里会被去掉不影响性能但在开发阶段能帮你抓住大部分低级错误。第五条跨领域对照着理解。编译器和大模型推理在内存管理上的相似性不是巧合而是因为底层约束是一样的。多看看不同领域怎么解决同类问题往往能给你自己的场景带来新思路。我那次把KV Cache改成按需增长灵感就是来自SmallVector的设计。这套东西没有银弹每个选择都有取舍。关键是理解你面对的场景到底需要什么——是分配速度、是内存利用率、还是代码简洁性。想清楚这个再去看LLVM的ADT和分配器设计很多为什么这么设计的问题就自然有答案了。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑