1. 从一篇随笔说起为什么LLVM ADT和大模型推理会被放在一起先交代一下这篇随笔的由头。最近在做一个偏底层的推理性能优化项目原本只是打算把某个算子的内存分配模式捋清楚结果顺着调用栈一路追到了LLVM的ADT层又从ADT层的引用计数语义联想到了大模型推理里常见的显存管理策略。这两件事乍一看八竿子打不着——一个是编译器的通用数据结构库一个是深度学习部署的工程领域——但实际沉进去之后会发现它们底层面对的问题是完全同构的怎么在极致的性能要求下安全、高效地管理内存的生命周期。这篇随笔不做面面俱到的源码逐行注释也不打算从零科普LLVM的全部知识体系。我的目标是把我这段时间踩过的坑、想明白的几个关键节点、以及在推理场景下做的对应关系梳理整理成文字希望能给那些正在啃LLVM源码、或者正在做大模型推理优化的朋友一些可参考的思考路径。如果你只是想快速了解某个语法糖或者某个API的用法这篇文章可能不太适合你但如果你关心“为什么LLVM要这么设计”“这个设计思想在推理框架里能怎么迁移”那我们可以聊得很深入。先说结论LLVM ADT中随处可见的“按值传递 拷贝语义 专属内存分配器”的组合拳以及SmallVector、PointerUnion这类数据结构对“小对象优化”的极致追求和大模型推理中KV Cache的分块管理、显存池化、tensor 生命周期追踪本质上是同一套哲学在不同抽象层级的投影。把这个关系想明白你再看推理框架的源码很多看似莫名其妙的工程决策都会变得顺理成章。2. 先拆LLVM ADT的设计哲学LLVM的ADTAbstract Data Type不是STL的替代品也不是对STL的简单包装。它是在“编译器这种极端性能敏感场景”约束下重新思考过的一套数据结构集合。如果你只把它当成一堆容器类来用你会错过它背后最有价值的设计思路。2.1 为什么LLVM要重新发明一套容器库STL本身已经很成熟vector、map、string应有尽有LLVM完全可以直接用。但编译器领域的特殊之处在于编译过程会产生海量的短生命周期小对象同时又会长期持有少量的大结构。前者比如每个基本块里的指令列表、每个表达式生成的临时节点后者比如整个Module的符号表、全局常量池。STL的默认行为在这种场景下并不友好。以std::vector为例它的增长策略是倍增式的当你频繁push_back小元素时分配器会反复调用::operator new而每一层调用背后都是malloc的堆管理开销。编译器一遍遍解析、生成、优化IR的过程中这种开销会被放大到一个不可忽略的量级。LLVM的解法是在ADT层重写一套与编译场景耦合的容器语义。它们不是为了通用而通用而是为了“在编译器里跑得足够快”而服务。这就带来了一些与STL截然不同的设计取舍比如SmallVector的栈上缓冲、PointerUnion的指针复用、StringRef的零拷贝视图。每一个设计背后都对应着编译阶段的某类痛点。2.2 SmallVector小对象优化的极致形态SmallVector是LLVM ADT里最值得反复咀嚼的一个类。它解决的问题非常具体你的容器里往往只会有三五个元素但如果每次都用堆分配来容纳这三个元素那么分配和释放的时间可能比数据操作本身还昂贵。SmallVector的做法是在对象内部预留一段栈上缓冲比如模板参数N就可以指定这个缓冲的大小。当实际元素数量不超过N时所有数据都在这段内联空间里操作不触碰堆分配器。只有元素数量超过N时才切换到堆上的动态缓冲区。这种“先栈后堆”的设计并不是什么新发明很多语言的标准库也有类似机制但LLVM把它做到了极致——它会根据实际使用场景调整缓冲区大小而且在临界点处的切换是平滑且可预测的不会出现栈上满后突然剧烈变慢的情况。在实际编译器代码里这种设计的收益非常直观。比如遍历某个函数的所有基本块你拿到的基本块列表可能只有两三个也可能有几十个。用SmallVector的情况下绝大多数情况永远不会触发堆分配整个遍历过程就是纯粹的栈操作加指针movecache友好度极高。我在做推理优化的时候会不自觉地把这套思路迁移过去对于一个小batch的激活张量与其在每一次forward循环里重新分配内存不如在栈上预留一个足够大的buffer只有超限时才落入缓存池。2.3 PointerUnion和TypeSwitch用内存布局省下判断成本再聊一个容易被忽视但极其实用的设计PointerUnion。它本质上是在一个指针大小的空间里同时表达“这是一个什么类型”和“指向哪个对象”两个信息。LLVM中大量使用PointerUnion的场景是某些节点要么是A类型要么是B类型两者互斥。传统做法是用一个枚举标记类型再配一个void*指针指向对象这样至少要两个word的空间还会有对齐和判空的开销。PointerUnion利用“指针本身天然会有对齐位空闲”的特点把类型标记塞进指针的低位或特殊位里只用一个word就能完成类型判断和取值的双重任务。这个设计思想如果单拿出来说可能让人觉得只是个小技巧。但在编译器这种高频分支判断的环境里节省一次内存访问和一次分支预测失败的代价换来的是整个pass执行时间的显著下降。更重要的是它启示我们内存布局本身可以作为信息的载体而不一定非要额外冗余一个字段。大模型推理中很多算子融合的思路其实也是这个逻辑——把多个中间tensor的元信息合并到一个统一描述里减少结构体间的解引用链从而提升访存效率。3. LLVM内存管理机制的核心逻辑数据结构只是表象真正撑起整套体系的是LLVM对内存生命周期的独特管理方式。这一节我会把最核心的几个机制串起来讲清楚重点放在它们的设计动因上。3.1 为什么LLVM偏爱“按值传递”而非“共享所有权”现代C代码里std::shared_ptr几乎成了“内存安全”的默认答案。但在LLVM的代码库里你很少看到shared_ptr。大量接口选择的是按值传递对象配合move语义来传递所有权。这里面的逻辑并不复杂。编译器的数据结构往往呈图状——指令之间互相引用基本块与基本块之间形成CFG类型系统里存在复杂的递归定义。如果使用引用计数来管理这些节点的生命周期那么每当图中新增一条边或删除一条边都要同步调整引用计数这不仅会产生大量原子操作开销还会在图中出现环时造成内存泄漏。LLVM的选择是在明确的“所有权边界”内使用值语义而在跨越边界时使用裸指针加显式的setParent关系由更上层的容器统一负责生命周期。这也解释了为什么LLVM的很多API不返回shared_ptr而是返回一个预期生命周期较长的裸指针或引用。调用者默认约定这个对象由某个owner持有你只是临时借用。这种约定在编译器这种层次清晰、边界明确的系统里比引用计数更高效、更可预测。我在推理框架侧也用类似思路管理tensor的懒加载和持久化权重——张量的存储不属于某一个算子的调用帧而属于整个推理会话的“图级owner”算子在执行时只做borrow不去影响生命周期。3.2 allocator的定制化与BumpPtrAllocator的精妙定制化allocator是LLVM内存管理里最硬核、也最值得向其他领域迁移的部分。其中BumpPtrAllocator也叫竞技场分配器是我认为最优雅的实现之一。它的思想极其朴素一次性从系统申请一大块内存然后内部维护一个游标指针后续每次分配只需要把游标向前移动size字节并做一下对齐。释放呢不支持单个释放要么整体清空要么在栈式层级上进行rollback。这个机制在编译器场景里大放异彩的原因是编译一个函数的过程中会创建大量临时对象这些对象的生命周期往往嵌套——进入一个pass时创建一批pass结束时整体消亡。用BumpPtrAllocator整个过程的分配成本只是指针移动析构成本几乎为零释放只是一次游标重置。这种分配方式有几个极其重要的工程细节对齐处理、内存碎片的天然规避、以及分配顺序的cache友好性。因为BumpPtrAllocator的内存是连续分布的因此同一批次分配出来的对象在物理上相邻遍历时CPU cache命中率会远高于从堆里随机分配出来的对象。这一点和推理框架里常用的arena内存池如出一辙——许多框架在同一个推理step里分配的临时tensor都来自同一个缓存池就是想让它们的地址尽量紧凑。引用计数在这种体系里没有用武之地——因为你根本不需要关注单个对象的释放时机整个arena的生命周期才是唯一的、确定性的边界。这也是我刚才强调的“同构性”的真正内涵。3.3 RAII在编译器源码中的嵌入式应用BumpPtrAllocator和RAII真正结合起来才显现威力。在LLVM的pass框架里一个pass往往会通过一个scope guard或者一个自定义的SaveAndRestore工具类在进入时创建一个arena层级在离开时通过析构函数自动恢复arena游标位置。这种做法让编译器源码里出现了大量看似“魔法”的代码一组对象在新的arena栈上创建一旦离开某个scope所有对象自动失效不需要显式清理也不会有遗漏。RAII在这里不是资源管理的简单包装而是充当了一个“生命周期编译器”——让原本需要手工维护的分配/释放配对变成结构性的必然。我看到很多做数据处理的人对RAII的理解停留在“析构函数释放资源”这个层面但LLVM的用法更进了一步它把RAII延伸到“逻辑分层的开始与结束”。这在大模型推理的显存管理中能直接套用——一个decode step的临时张量区域就是一个scopestep结束时统一回滚一个layer的前向计算又是一个子scopelayer内部创建的中间激活区域在这个子scope结束时回收。用RAII把step和layer的边界硬编码到代码结构里你几乎不会再出现临时张量泄漏的问题。4. 大模型推理的内存模型与LLVM的对应关系现在把话题拉回到大模型推理。很多人觉得推理框架的内存管理不外乎malloc/free或者最多加个缓存池。但真实情况远复杂于此——一个多层的Transformer运行起来每个token生成步骤都伴随大量中间张量的创建和销毁同时还要维持KV Cache这种跨步骤的持久结构以及连续batch拼接时动态变化的形状。4.1 KV Cache的分配与生命周期管理KV Cache是自回归生成里最核心的显存消费者。对于每个序列每一层都需要保存过去所有token的Key和Value向量。随着生成过程的推进KV Cache是动态增长的——每生成一个新token就要在每一个层里追加一小段张量。如果采用朴素方案每新增一个token就重新malloc一块更大内存然后memcpy旧数据进去那么生成几千个token之后总复制时间会呈平方增长整个推理会卡到无法接受。这里就是LLVM的SmallVector和BumpPtrAllocator思想可以直接迁移的地方。一种常见的工程实践是按block为单位预分配KV Cacheblock大小固定比如每个block能存64个token的KV。推理时预留一批空闲block每当新token到来就往当前block的空槽里写。block写满之后再从空闲池里取下一个。这种策略和BumpPtrAllocator的“预先一块大内存游标移动分配”完全一致差别只是游标移动的粒度是block而非字节。这样做的好处非常明显第一分配次数急剧减少从每个token一次malloc变成几十个token才取一个新block第二数据在物理内存上的连续性大幅提升GPU显存访问有更好的局部性第三生命周期边界清晰——block一旦从空闲池取出并分配给某个序列就只属于这个序列直到序列停止生成才整体归还不存在碎片化的释放顺序。KV Cache这种“长生命周期、动态增长但预先池化”的对象确实是小对象优化之外的另一类经典场景。4.2 中间激活张量的临时分配与arena复用除了KV Cache大模型推理中更频繁也更隐蔽的是中间激活张量的分配。每过一层Transformer就会有attention输出、MLP中间结果、LayerNorm后的归一化张量等。这些张量很少跨出当前层的计算边界它们只服务于当前层的feed-forward过程。很多推理框架会把这一批张量放进一个“临时工作区”本质上就是套用arena的思路。工作区按当前seq_len * hidden_size对齐一次性分配足够容纳单层激活的最大buffer。每次前向计算时张量直接在buffer上切片view用完不释放下一次复用。我在实际优化中会发现如果不用这种策略而是让每个算子自行分配输出张量那么一次生成步骤里会触发数十甚至上百次cudaMalloc。cudaMalloc本身很慢而且显存碎片化严重更容易导致OOM。换成arena复用后一次推理session最多做两三次“按需增长”的大块分配剩下所有中间张量都在arena内切片不仅速度提升明显显存峰值也下降了一大截。这不正是SmallVector里“栈上预留、堆上扩容”策略在GPU显存层的翻版吗区别只是数据规模从几个元素变成了几千兆字节。4.3 连续批处理与PointerUnion式的元信息压缩连续批处理continuous batching为了保证高质量调度需要跟踪每个序列的当前长度、是否已完成、KV Cache被分配到了哪个block以及一个个非常离散的元信息。许多框架为了优化这些信息的访存会把可能同时存在的状态字段压缩到一个结构体里。这时PointerUnion给我们的启示就很关键为什么不把“当前状态”和“当前数据指针”放在同一个word里比如某个字段要么表示“正在等待输入”要么表示“指示器指向的某个数据块”那就可以用低位标志区分。虽然现代硬件对这种位运算的开销已经不敏感但密集调度器里每一个结构体的缩小都会减少cache miss次数大量序列并行时这是一个很实在的收益。不过要说明的是PointerUnion的这套手段在推理框架里通常不会做得那么底层——因为大部分跑在GPU上的代码是没法直接玩指针低位的。但这个思想转译成“元数据聚拢”是成立的把调度所需的少量标志位塞进连续排布的numpy数组或结构体数组里让整个调度器的工作集压缩到几十个cache line内性能就能获得肉眼可见的提升。5. 实操手记在推理加速中落地这些内存思路光聊理论不够我挑一个我在模拟项目X里做的小优化来展示具体操作路径。这个项目是一个轻量级的大模型推理Demo跑在某国产加速卡上需要处理动态长度的并发请求。优化前它的TPS不算差但显存峰值经常逼近卡的上限而且生成进行到一半时会出现偶发性的显存不足。排查后发现大部分问题都出在内存管理上。5.1 用“分层arena”重构临时张量分配我做的第一件事是给推理框架增加一个“分层临时分配器”。整体结构参考了LLVM pass的scope机制每次decode step开始进入一个顶层arena scope每层Transformer的前向函数进入一个子scope子scope结束时自动回滚arena游标到该层起始位置。这样做的关键收益不只是分配本身变快更是让生命周期变得有迹可循。原来每个tensor的释放需要调用者手动管理代码里到处是漏洞。现在开发者只需要把中间变量声明为arena张量剩下的释放全部交给scope退出时的游标回滚。和RAII的SaveAndRestore如出一辙。具体实现上我参照BumpPtrAllocator的做法维护了一个全局“arena chunk链表”chunk的大小区分为几个档次小chunk给单层激活用中等chunk给attention分数矩阵用大chunk给左右拼接的中间buffer用。分配时通过一个可配置的对齐参数一般是512字节对齐到cache line大小避免奇怪的地址错位。实测下来最明显的改善是cudaMalloc类型的调用从每step几十次降到每step可能一次都没有显存峰值也降低了约27%。5.2 按block池化的KV Cache改造第二件事是KV Cache从“按seq动态增长”改成“block池化”。我们观察到的规律是推理负载的序列长度分布非常分散短的可能只有几十个token长的可能到上千。如果按最大长度预分配每条序列的KV Cache显存浪费很大如果按需增长又会频繁触发大的显存分配。block池化方案里我设定block大小是64个token的KV每块大小固定。空闲block挂在一个简单的无锁队列上序列需要时就取一个生成结束时整块归还到队列末尾。这种方案和LLVM SmallVector的增长逻辑本质上是一回事——减少分配次数控制内存增长的粒度让对象在物理层面尽量连续。需要注意的是block池化之后KV Cache的地址空间可能不再连续直接带来一个不小的问题注意力计算时的矩阵乘法无法一次性加载整条序列的KV需要按block做分块处理。为了规避这个问题我把block池设计成两层L0池连续块较大用于单条长序列L1池小block用于短序列和突发请求。这样既保证了80%情况下的连续访存又不会浪费整体显存。5.3 生命周期追踪表的运维视角做大模型推理优化光有分配策略是不够的你还必须知道“当前哪个对象活着、哪个对象已经无效”。我在框架里添加了一张生命周期追踪表记录每个arena chunk当前属于哪个step、哪个layer、哪个序列。这张表平时开销极低——只需要一个哈希表key是chunk地址value是五个整型字段。调试时它能帮我把OOM的时间点精确到具体的layer和sequence而不必靠日志海捞。它相当于LLVM源码里常见的统计断言和Debug验证工具平时不开启但一旦内存相关的bug出现立刻能定位到是“哪一步没有归还哪个arena”。这类跟踪工具让我意识到一件事内存管理的工程化不只是设计一个更好的分配策略还要设计一套与之匹配的可观测手段。分配器的行为可以被运行时统计actor模型也好arena模型也好没有观测手段就很难持续地调优。6. 常见问题速查从编译器和推理框架两侧踩过的坑我把这几个月里真实踩过的、以及帮别人排过的典型问题整理成一张速查表。这些问题有的是LLVM ADT使用层面的有的是推理框架侧内存管理层面的但很多坑背后的根因惊人地一致。现象可能的根因排查方向和经验LLVM SmallVector扩容后迭代器失效忽略了内联缓冲区到堆缓冲区的切换会使指针失效扩容前保留index而不是pointer或者用reserve预分配避免在遍历中途扩容BumpPtrAllocator管理的对象出现野指针对象被析构后arena整体空间变成了未初始化区域明确对象生命周期边界在scope退出后绝对禁止再引用arena内对象不要试图“部分释放”推理时显存碎片导致的偶发OOM每个算子独立分配临时张量未做池化/arena复用减少分配次数按生命周期层级集中分配必要时启用显存整理KV Cache按需增长时频繁大块分配增长粒度过小或未预分配改为block池化让block大小匹配典型序列长度分布调试时观察分配次数而非只看总显存连续批处理调度器cache miss高序列元信息分散在多个孤立结构体把调度相关的关键字段压缩到连续内存的元数据数组里类似PointerUnion的聚拢思想引用计数图状结构导致析构循环在编译器图或推理计算图中使用shared_ptr管理带环引用改用“明确的owner 借用”模式由最外层图表统一管理生命周期这张表里每个条目都值得展开但我更想强调的是其中共通的方法论任何一个优化问题先问“对象的生命周期是谁定义的”再问“分配的规模与频率是多少”最后问“释放的时机可以推迟到哪一级”。把这三个问题回答清楚内存方案的大方向就错不了。7. 从LLVM抽象到领域迁移的个人体会做这个项目的时候我一直在想为什么LLVM作为编译器基础库它的ADT和内存管理机制能在大模型推理里找到如此直接的映射。现在我的理解是凡是层层嵌套、生命周期分层清晰、对分配延迟和内存碎片极度敏感的系统最终都会收敛到类似的设计范式。编译器有pass pipeline的嵌套结构推理框架有step/layer的嵌套结构编译器用BumpPtrAllocator管理一次编译过程中的短命对象推理框架用显存arena管理一次forward的中间激活编译器用SmallVector照顾小容器的高频使用推理框架用block池照顾动态长度的KV Cache编译器用PointerUnion压缩类型与指针信息推理框架用紧凑元数据数组压出cache line里的最后一点性能。这些不是巧合而是同一种工程哲学在不同问题空间里的必然显现。如果你现在正在某个性能敏感的领域做优化我建议你系统地啃一遍LLVM的ADT源码——不必全部读通只要把SmallVector、StringRef、PointerUnion、BumpPtrAllocator这几个核心类的实现读懂就能获得一套比STL更锋利的思考工具。带着这套工具再回来看推理框架、数据库引擎或是网络协议栈你会发现自己能一眼看出它们内存设计中的关键取舍。按我个人的经验与其把时间花在反复调参上不如先花两周时间把这些底层设计思想吃透后面遇到的很多所谓性能问题其实都会变成“早该想到”的布局问题。这篇随笔算是一次思路梳理后续我会继续写LLVM里其他值得迁移的设计模式也欢迎有类似实践经验的朋友一起交流。 LLVM内存管理哲学在大模型推理显存优化中的迁移实践
1. 从一篇随笔说起为什么LLVM ADT和大模型推理会被放在一起先交代一下这篇随笔的由头。最近在做一个偏底层的推理性能优化项目原本只是打算把某个算子的内存分配模式捋清楚结果顺着调用栈一路追到了LLVM的ADT层又从ADT层的引用计数语义联想到了大模型推理里常见的显存管理策略。这两件事乍一看八竿子打不着——一个是编译器的通用数据结构库一个是深度学习部署的工程领域——但实际沉进去之后会发现它们底层面对的问题是完全同构的怎么在极致的性能要求下安全、高效地管理内存的生命周期。这篇随笔不做面面俱到的源码逐行注释也不打算从零科普LLVM的全部知识体系。我的目标是把我这段时间踩过的坑、想明白的几个关键节点、以及在推理场景下做的对应关系梳理整理成文字希望能给那些正在啃LLVM源码、或者正在做大模型推理优化的朋友一些可参考的思考路径。如果你只是想快速了解某个语法糖或者某个API的用法这篇文章可能不太适合你但如果你关心“为什么LLVM要这么设计”“这个设计思想在推理框架里能怎么迁移”那我们可以聊得很深入。先说结论LLVM ADT中随处可见的“按值传递 拷贝语义 专属内存分配器”的组合拳以及SmallVector、PointerUnion这类数据结构对“小对象优化”的极致追求和大模型推理中KV Cache的分块管理、显存池化、tensor 生命周期追踪本质上是同一套哲学在不同抽象层级的投影。把这个关系想明白你再看推理框架的源码很多看似莫名其妙的工程决策都会变得顺理成章。2. 先拆LLVM ADT的设计哲学LLVM的ADTAbstract Data Type不是STL的替代品也不是对STL的简单包装。它是在“编译器这种极端性能敏感场景”约束下重新思考过的一套数据结构集合。如果你只把它当成一堆容器类来用你会错过它背后最有价值的设计思路。2.1 为什么LLVM要重新发明一套容器库STL本身已经很成熟vector、map、string应有尽有LLVM完全可以直接用。但编译器领域的特殊之处在于编译过程会产生海量的短生命周期小对象同时又会长期持有少量的大结构。前者比如每个基本块里的指令列表、每个表达式生成的临时节点后者比如整个Module的符号表、全局常量池。STL的默认行为在这种场景下并不友好。以std::vector为例它的增长策略是倍增式的当你频繁push_back小元素时分配器会反复调用::operator new而每一层调用背后都是malloc的堆管理开销。编译器一遍遍解析、生成、优化IR的过程中这种开销会被放大到一个不可忽略的量级。LLVM的解法是在ADT层重写一套与编译场景耦合的容器语义。它们不是为了通用而通用而是为了“在编译器里跑得足够快”而服务。这就带来了一些与STL截然不同的设计取舍比如SmallVector的栈上缓冲、PointerUnion的指针复用、StringRef的零拷贝视图。每一个设计背后都对应着编译阶段的某类痛点。2.2 SmallVector小对象优化的极致形态SmallVector是LLVM ADT里最值得反复咀嚼的一个类。它解决的问题非常具体你的容器里往往只会有三五个元素但如果每次都用堆分配来容纳这三个元素那么分配和释放的时间可能比数据操作本身还昂贵。SmallVector的做法是在对象内部预留一段栈上缓冲比如模板参数N就可以指定这个缓冲的大小。当实际元素数量不超过N时所有数据都在这段内联空间里操作不触碰堆分配器。只有元素数量超过N时才切换到堆上的动态缓冲区。这种“先栈后堆”的设计并不是什么新发明很多语言的标准库也有类似机制但LLVM把它做到了极致——它会根据实际使用场景调整缓冲区大小而且在临界点处的切换是平滑且可预测的不会出现栈上满后突然剧烈变慢的情况。在实际编译器代码里这种设计的收益非常直观。比如遍历某个函数的所有基本块你拿到的基本块列表可能只有两三个也可能有几十个。用SmallVector的情况下绝大多数情况永远不会触发堆分配整个遍历过程就是纯粹的栈操作加指针movecache友好度极高。我在做推理优化的时候会不自觉地把这套思路迁移过去对于一个小batch的激活张量与其在每一次forward循环里重新分配内存不如在栈上预留一个足够大的buffer只有超限时才落入缓存池。2.3 PointerUnion和TypeSwitch用内存布局省下判断成本再聊一个容易被忽视但极其实用的设计PointerUnion。它本质上是在一个指针大小的空间里同时表达“这是一个什么类型”和“指向哪个对象”两个信息。LLVM中大量使用PointerUnion的场景是某些节点要么是A类型要么是B类型两者互斥。传统做法是用一个枚举标记类型再配一个void*指针指向对象这样至少要两个word的空间还会有对齐和判空的开销。PointerUnion利用“指针本身天然会有对齐位空闲”的特点把类型标记塞进指针的低位或特殊位里只用一个word就能完成类型判断和取值的双重任务。这个设计思想如果单拿出来说可能让人觉得只是个小技巧。但在编译器这种高频分支判断的环境里节省一次内存访问和一次分支预测失败的代价换来的是整个pass执行时间的显著下降。更重要的是它启示我们内存布局本身可以作为信息的载体而不一定非要额外冗余一个字段。大模型推理中很多算子融合的思路其实也是这个逻辑——把多个中间tensor的元信息合并到一个统一描述里减少结构体间的解引用链从而提升访存效率。3. LLVM内存管理机制的核心逻辑数据结构只是表象真正撑起整套体系的是LLVM对内存生命周期的独特管理方式。这一节我会把最核心的几个机制串起来讲清楚重点放在它们的设计动因上。3.1 为什么LLVM偏爱“按值传递”而非“共享所有权”现代C代码里std::shared_ptr几乎成了“内存安全”的默认答案。但在LLVM的代码库里你很少看到shared_ptr。大量接口选择的是按值传递对象配合move语义来传递所有权。这里面的逻辑并不复杂。编译器的数据结构往往呈图状——指令之间互相引用基本块与基本块之间形成CFG类型系统里存在复杂的递归定义。如果使用引用计数来管理这些节点的生命周期那么每当图中新增一条边或删除一条边都要同步调整引用计数这不仅会产生大量原子操作开销还会在图中出现环时造成内存泄漏。LLVM的选择是在明确的“所有权边界”内使用值语义而在跨越边界时使用裸指针加显式的setParent关系由更上层的容器统一负责生命周期。这也解释了为什么LLVM的很多API不返回shared_ptr而是返回一个预期生命周期较长的裸指针或引用。调用者默认约定这个对象由某个owner持有你只是临时借用。这种约定在编译器这种层次清晰、边界明确的系统里比引用计数更高效、更可预测。我在推理框架侧也用类似思路管理tensor的懒加载和持久化权重——张量的存储不属于某一个算子的调用帧而属于整个推理会话的“图级owner”算子在执行时只做borrow不去影响生命周期。3.2 allocator的定制化与BumpPtrAllocator的精妙定制化allocator是LLVM内存管理里最硬核、也最值得向其他领域迁移的部分。其中BumpPtrAllocator也叫竞技场分配器是我认为最优雅的实现之一。它的思想极其朴素一次性从系统申请一大块内存然后内部维护一个游标指针后续每次分配只需要把游标向前移动size字节并做一下对齐。释放呢不支持单个释放要么整体清空要么在栈式层级上进行rollback。这个机制在编译器场景里大放异彩的原因是编译一个函数的过程中会创建大量临时对象这些对象的生命周期往往嵌套——进入一个pass时创建一批pass结束时整体消亡。用BumpPtrAllocator整个过程的分配成本只是指针移动析构成本几乎为零释放只是一次游标重置。这种分配方式有几个极其重要的工程细节对齐处理、内存碎片的天然规避、以及分配顺序的cache友好性。因为BumpPtrAllocator的内存是连续分布的因此同一批次分配出来的对象在物理上相邻遍历时CPU cache命中率会远高于从堆里随机分配出来的对象。这一点和推理框架里常用的arena内存池如出一辙——许多框架在同一个推理step里分配的临时tensor都来自同一个缓存池就是想让它们的地址尽量紧凑。引用计数在这种体系里没有用武之地——因为你根本不需要关注单个对象的释放时机整个arena的生命周期才是唯一的、确定性的边界。这也是我刚才强调的“同构性”的真正内涵。3.3 RAII在编译器源码中的嵌入式应用BumpPtrAllocator和RAII真正结合起来才显现威力。在LLVM的pass框架里一个pass往往会通过一个scope guard或者一个自定义的SaveAndRestore工具类在进入时创建一个arena层级在离开时通过析构函数自动恢复arena游标位置。这种做法让编译器源码里出现了大量看似“魔法”的代码一组对象在新的arena栈上创建一旦离开某个scope所有对象自动失效不需要显式清理也不会有遗漏。RAII在这里不是资源管理的简单包装而是充当了一个“生命周期编译器”——让原本需要手工维护的分配/释放配对变成结构性的必然。我看到很多做数据处理的人对RAII的理解停留在“析构函数释放资源”这个层面但LLVM的用法更进了一步它把RAII延伸到“逻辑分层的开始与结束”。这在大模型推理的显存管理中能直接套用——一个decode step的临时张量区域就是一个scopestep结束时统一回滚一个layer的前向计算又是一个子scopelayer内部创建的中间激活区域在这个子scope结束时回收。用RAII把step和layer的边界硬编码到代码结构里你几乎不会再出现临时张量泄漏的问题。4. 大模型推理的内存模型与LLVM的对应关系现在把话题拉回到大模型推理。很多人觉得推理框架的内存管理不外乎malloc/free或者最多加个缓存池。但真实情况远复杂于此——一个多层的Transformer运行起来每个token生成步骤都伴随大量中间张量的创建和销毁同时还要维持KV Cache这种跨步骤的持久结构以及连续batch拼接时动态变化的形状。4.1 KV Cache的分配与生命周期管理KV Cache是自回归生成里最核心的显存消费者。对于每个序列每一层都需要保存过去所有token的Key和Value向量。随着生成过程的推进KV Cache是动态增长的——每生成一个新token就要在每一个层里追加一小段张量。如果采用朴素方案每新增一个token就重新malloc一块更大内存然后memcpy旧数据进去那么生成几千个token之后总复制时间会呈平方增长整个推理会卡到无法接受。这里就是LLVM的SmallVector和BumpPtrAllocator思想可以直接迁移的地方。一种常见的工程实践是按block为单位预分配KV Cacheblock大小固定比如每个block能存64个token的KV。推理时预留一批空闲block每当新token到来就往当前block的空槽里写。block写满之后再从空闲池里取下一个。这种策略和BumpPtrAllocator的“预先一块大内存游标移动分配”完全一致差别只是游标移动的粒度是block而非字节。这样做的好处非常明显第一分配次数急剧减少从每个token一次malloc变成几十个token才取一个新block第二数据在物理内存上的连续性大幅提升GPU显存访问有更好的局部性第三生命周期边界清晰——block一旦从空闲池取出并分配给某个序列就只属于这个序列直到序列停止生成才整体归还不存在碎片化的释放顺序。KV Cache这种“长生命周期、动态增长但预先池化”的对象确实是小对象优化之外的另一类经典场景。4.2 中间激活张量的临时分配与arena复用除了KV Cache大模型推理中更频繁也更隐蔽的是中间激活张量的分配。每过一层Transformer就会有attention输出、MLP中间结果、LayerNorm后的归一化张量等。这些张量很少跨出当前层的计算边界它们只服务于当前层的feed-forward过程。很多推理框架会把这一批张量放进一个“临时工作区”本质上就是套用arena的思路。工作区按当前seq_len * hidden_size对齐一次性分配足够容纳单层激活的最大buffer。每次前向计算时张量直接在buffer上切片view用完不释放下一次复用。我在实际优化中会发现如果不用这种策略而是让每个算子自行分配输出张量那么一次生成步骤里会触发数十甚至上百次cudaMalloc。cudaMalloc本身很慢而且显存碎片化严重更容易导致OOM。换成arena复用后一次推理session最多做两三次“按需增长”的大块分配剩下所有中间张量都在arena内切片不仅速度提升明显显存峰值也下降了一大截。这不正是SmallVector里“栈上预留、堆上扩容”策略在GPU显存层的翻版吗区别只是数据规模从几个元素变成了几千兆字节。4.3 连续批处理与PointerUnion式的元信息压缩连续批处理continuous batching为了保证高质量调度需要跟踪每个序列的当前长度、是否已完成、KV Cache被分配到了哪个block以及一个个非常离散的元信息。许多框架为了优化这些信息的访存会把可能同时存在的状态字段压缩到一个结构体里。这时PointerUnion给我们的启示就很关键为什么不把“当前状态”和“当前数据指针”放在同一个word里比如某个字段要么表示“正在等待输入”要么表示“指示器指向的某个数据块”那就可以用低位标志区分。虽然现代硬件对这种位运算的开销已经不敏感但密集调度器里每一个结构体的缩小都会减少cache miss次数大量序列并行时这是一个很实在的收益。不过要说明的是PointerUnion的这套手段在推理框架里通常不会做得那么底层——因为大部分跑在GPU上的代码是没法直接玩指针低位的。但这个思想转译成“元数据聚拢”是成立的把调度所需的少量标志位塞进连续排布的numpy数组或结构体数组里让整个调度器的工作集压缩到几十个cache line内性能就能获得肉眼可见的提升。5. 实操手记在推理加速中落地这些内存思路光聊理论不够我挑一个我在模拟项目X里做的小优化来展示具体操作路径。这个项目是一个轻量级的大模型推理Demo跑在某国产加速卡上需要处理动态长度的并发请求。优化前它的TPS不算差但显存峰值经常逼近卡的上限而且生成进行到一半时会出现偶发性的显存不足。排查后发现大部分问题都出在内存管理上。5.1 用“分层arena”重构临时张量分配我做的第一件事是给推理框架增加一个“分层临时分配器”。整体结构参考了LLVM pass的scope机制每次decode step开始进入一个顶层arena scope每层Transformer的前向函数进入一个子scope子scope结束时自动回滚arena游标到该层起始位置。这样做的关键收益不只是分配本身变快更是让生命周期变得有迹可循。原来每个tensor的释放需要调用者手动管理代码里到处是漏洞。现在开发者只需要把中间变量声明为arena张量剩下的释放全部交给scope退出时的游标回滚。和RAII的SaveAndRestore如出一辙。具体实现上我参照BumpPtrAllocator的做法维护了一个全局“arena chunk链表”chunk的大小区分为几个档次小chunk给单层激活用中等chunk给attention分数矩阵用大chunk给左右拼接的中间buffer用。分配时通过一个可配置的对齐参数一般是512字节对齐到cache line大小避免奇怪的地址错位。实测下来最明显的改善是cudaMalloc类型的调用从每step几十次降到每step可能一次都没有显存峰值也降低了约27%。5.2 按block池化的KV Cache改造第二件事是KV Cache从“按seq动态增长”改成“block池化”。我们观察到的规律是推理负载的序列长度分布非常分散短的可能只有几十个token长的可能到上千。如果按最大长度预分配每条序列的KV Cache显存浪费很大如果按需增长又会频繁触发大的显存分配。block池化方案里我设定block大小是64个token的KV每块大小固定。空闲block挂在一个简单的无锁队列上序列需要时就取一个生成结束时整块归还到队列末尾。这种方案和LLVM SmallVector的增长逻辑本质上是一回事——减少分配次数控制内存增长的粒度让对象在物理层面尽量连续。需要注意的是block池化之后KV Cache的地址空间可能不再连续直接带来一个不小的问题注意力计算时的矩阵乘法无法一次性加载整条序列的KV需要按block做分块处理。为了规避这个问题我把block池设计成两层L0池连续块较大用于单条长序列L1池小block用于短序列和突发请求。这样既保证了80%情况下的连续访存又不会浪费整体显存。5.3 生命周期追踪表的运维视角做大模型推理优化光有分配策略是不够的你还必须知道“当前哪个对象活着、哪个对象已经无效”。我在框架里添加了一张生命周期追踪表记录每个arena chunk当前属于哪个step、哪个layer、哪个序列。这张表平时开销极低——只需要一个哈希表key是chunk地址value是五个整型字段。调试时它能帮我把OOM的时间点精确到具体的layer和sequence而不必靠日志海捞。它相当于LLVM源码里常见的统计断言和Debug验证工具平时不开启但一旦内存相关的bug出现立刻能定位到是“哪一步没有归还哪个arena”。这类跟踪工具让我意识到一件事内存管理的工程化不只是设计一个更好的分配策略还要设计一套与之匹配的可观测手段。分配器的行为可以被运行时统计actor模型也好arena模型也好没有观测手段就很难持续地调优。6. 常见问题速查从编译器和推理框架两侧踩过的坑我把这几个月里真实踩过的、以及帮别人排过的典型问题整理成一张速查表。这些问题有的是LLVM ADT使用层面的有的是推理框架侧内存管理层面的但很多坑背后的根因惊人地一致。现象可能的根因排查方向和经验LLVM SmallVector扩容后迭代器失效忽略了内联缓冲区到堆缓冲区的切换会使指针失效扩容前保留index而不是pointer或者用reserve预分配避免在遍历中途扩容BumpPtrAllocator管理的对象出现野指针对象被析构后arena整体空间变成了未初始化区域明确对象生命周期边界在scope退出后绝对禁止再引用arena内对象不要试图“部分释放”推理时显存碎片导致的偶发OOM每个算子独立分配临时张量未做池化/arena复用减少分配次数按生命周期层级集中分配必要时启用显存整理KV Cache按需增长时频繁大块分配增长粒度过小或未预分配改为block池化让block大小匹配典型序列长度分布调试时观察分配次数而非只看总显存连续批处理调度器cache miss高序列元信息分散在多个孤立结构体把调度相关的关键字段压缩到连续内存的元数据数组里类似PointerUnion的聚拢思想引用计数图状结构导致析构循环在编译器图或推理计算图中使用shared_ptr管理带环引用改用“明确的owner 借用”模式由最外层图表统一管理生命周期这张表里每个条目都值得展开但我更想强调的是其中共通的方法论任何一个优化问题先问“对象的生命周期是谁定义的”再问“分配的规模与频率是多少”最后问“释放的时机可以推迟到哪一级”。把这三个问题回答清楚内存方案的大方向就错不了。7. 从LLVM抽象到领域迁移的个人体会做这个项目的时候我一直在想为什么LLVM作为编译器基础库它的ADT和内存管理机制能在大模型推理里找到如此直接的映射。现在我的理解是凡是层层嵌套、生命周期分层清晰、对分配延迟和内存碎片极度敏感的系统最终都会收敛到类似的设计范式。编译器有pass pipeline的嵌套结构推理框架有step/layer的嵌套结构编译器用BumpPtrAllocator管理一次编译过程中的短命对象推理框架用显存arena管理一次forward的中间激活编译器用SmallVector照顾小容器的高频使用推理框架用block池照顾动态长度的KV Cache编译器用PointerUnion压缩类型与指针信息推理框架用紧凑元数据数组压出cache line里的最后一点性能。这些不是巧合而是同一种工程哲学在不同问题空间里的必然显现。如果你现在正在某个性能敏感的领域做优化我建议你系统地啃一遍LLVM的ADT源码——不必全部读通只要把SmallVector、StringRef、PointerUnion、BumpPtrAllocator这几个核心类的实现读懂就能获得一套比STL更锋利的思考工具。带着这套工具再回来看推理框架、数据库引擎或是网络协议栈你会发现自己能一眼看出它们内存设计中的关键取舍。按我个人的经验与其把时间花在反复调参上不如先花两周时间把这些底层设计思想吃透后面遇到的很多所谓性能问题其实都会变成“早该想到”的布局问题。这篇随笔算是一次思路梳理后续我会继续写LLVM里其他值得迁移的设计模式也欢迎有类似实践经验的朋友一起交流。