☰
LLVM内存管理对大模型推理KV Cache优化的启示
2026/10/14 11:04:49 网站建设 项目流程

1. 从一个编译基础设施的视角看大模型推理

LLVM 的 ADT(Abstract Data Types,抽象数据类型)和内存管理机制,乍看之下跟大模型推理八竿子打不着——一个是编译器后端的基础设施,一个是当下最热的 AI 推理场景。但如果你真正翻过 LLVM 的源码,再去看一遍大模型推理框架的内存分配逻辑,会发现两者在核心思路上有大量可以互相借鉴的地方。这篇随笔就是把我自己在啃 LLVM ADT 和内存管理时的一些理解,跟大模型推理中的内存管理问题做一个对照梳理,既是对知识的复盘,也希望能给正在做推理优化或者编译器相关工作的朋友一些参考。

这篇文章适合谁看?如果你正在做推理框架的性能优化,尤其是被 KV Cache 的内存碎片、显存分配效率这些问题困扰;或者你在写编译器相关的东西,想理解 LLVM 为什么要把 ADT 和内存管理设计成现在这个样子;再或者你只是对底层系统感兴趣,想找一个具体的切入点把编译器和 AI 推理串起来看——那这篇内容应该能给你一些有用的东西。我会尽量从实际代码和实际场景出发,把每个设计决策背后的“为什么”讲清楚,而不是只罗列结论。

整篇内容会围绕两条主线展开:一条是 LLVM 的 ADT 体系与内存管理机制,包括SmallVector、DenseMap、BumpPtrAllocator、SpecificBumpPtrAllocator这些核心组件;另一条是大模型推理中的内存管理,包括 KV Cache 的分配策略、PagedAttention 的思路、显存池化管理等。最后会把两条线合在一起,看看 LLVM 的设计哲学能给推理框架带来什么启发。

2. LLVM ADT 体系的设计哲学与核心组件

2.1 为什么 LLVM 不用 STL 而自己造一套 ADT

这是很多人第一次看 LLVM 源码时都会问的问题。标准模板库已经提供了std::vector、std::map、std::unordered_map,为什么 LLVM 还要费劲自己实现一套SmallVector、DenseMap、StringMap?

核心原因有三个。第一是性能可预测性。STL 容器的性能特征在不同标准库实现之间差异很大,比如std::unordered_map的哈希策略、桶的增长方式,在不同平台上行为不一致。编译器这种对性能极度敏感的场景,需要容器行为完全可控。第二是内存布局的优化空间。LLVM 的很多容器针对编译器场景做了特殊优化,比如SmallVector的小对象内联存储,在元素数量少的时候完全避免堆分配,这在编译器中极其常见——大量基本块、指令列表都很短。第三是避免异常和 RTTI 依赖。LLVM 在早期就禁用了 C++ 异常和 RTTI,STL 的很多容器在异常安全上有额外开销,而 LLVM 需要一套不依赖异常的容器体系。

我个人的理解是,LLVM 的 ADT 本质上是一套“为编译器工作量身定制的容器库”。它不是要替代 STL,而是在特定场景下提供更优解。你在写 LLVM Pass 的时候用SmallVector而不是std::vector,不是因为std::vector不好,而是因为你知道这个场景下元素数量通常很少,SmallVector的内联存储能帮你省掉大量小堆分配。

2.2 SmallVector 的内联存储机制拆解

SmallVector是 LLVM ADT 里最常用的容器之一,它的核心设计是“小对象内联存储”。模板参数里有一个N,表示内联存储的元素个数。当元素数量不超过N时,所有元素直接存在栈上的内联缓冲区里,不发生任何堆分配;超过N之后,才切换到堆分配,并把已有元素搬过去。

这个设计的精妙之处在于,它把“大多数情况下元素很少”这个编译器领域的统计规律直接编码进了类型系统。比如一个基本块的指令列表,绝大多数情况下不超过 8 条指令,那用SmallVector<Instruction*, 8>就非常合适。前 8 个元素零堆分配,超过之后才走堆。

但这里有几个容易踩的坑。第一,N的选择需要基于实际数据分布,选大了浪费栈空间,选小了频繁触发堆分配。第二,SmallVector在增长时的扩容策略跟std::vector类似,通常是翻倍增长,但它的内联缓冲区在切换到堆之后并不会被释放,而是继续占用栈空间直到对象析构。第三,SmallVector的迭代器在发生扩容后会失效,这一点跟std::vector一样,但在内联切换到堆的瞬间,所有指针和引用都会失效,这个时机比std::vector更隐蔽。

实操心得:在 LLVM Pass 里用SmallVector时,如果你不确定元素数量范围,可以先用一个保守的N值,然后通过 profiling 看堆分配次数。如果堆分配次数很高,说明N选小了;如果栈使用量报警,说明N选大了。我一般会从 4 或 8 开始试。

2.3 DenseMap 与 DenseSet 的哈希策略

DenseMap是 LLVM 里用得最多的哈希表,它的设计目标是“缓存友好”。跟std::unordered_map的链式哈希不同,DenseMap用的是开放寻址法,所有键值对存在一个连续的数组里,冲突时线性探测。这样做的好处是内存局部性好,查找时缓存命中率高;坏处是删除操作复杂,需要 tombstone 标记,而且负载因子不能太高。

DenseMap的默认负载因子是 3/4,超过之后会扩容。扩容时会重新哈希所有元素,这个操作是 O(n) 的。所以在性能敏感的代码里,如果你能预估元素数量,最好在构造时预留足够的空间,避免多次扩容。

另一个关键点是DenseMap的键类型要求。它要求键类型有一个DenseMapInfo特化,提供空键、 tombstone 键和哈希函数。对于指针类型,LLVM 已经内置了特化;对于自定义类型,你需要自己写。这里有个常见的坑:如果你的键类型有 padding 字节,直接按字节哈希会出问题,因为 padding 字节的内容是不确定的。LLVM 的做法是对每个字段分别哈希,而不是对整个对象做内存哈希。

2.4 StringMap 与 StringRef 的配合使用

StringMap是专门为字符串键优化的哈希表,它的特点是键的存储方式——它不复制字符串,而是把字符串数据直接存在哈希表的条目里。配合StringRef使用,可以避免大量字符串拷贝。

StringRef本身是一个轻量的字符串引用,只包含一个指针和一个长度,不拥有内存。它的设计哲学是“字符串的视图”,跟 C++17 的std::string_view类似,但 LLVM 在 C++17 之前就有了。StringRef不保证以 null 结尾,所以不能直接传给 C 风格字符串函数,需要先转成std::string或者用str()方法。

在 LLVM Pass 里处理字符串时,我一般会尽量用StringRef传递参数,只有在需要拥有字符串时才用std::string。这样可以避免很多不必要的拷贝。但要注意StringRef的生命周期问题——它不拥有内存,如果它引用的字符串被释放了,StringRef就成了悬空引用。

3. LLVM 内存管理机制的核心设计

3.1 BumpPtrAllocator 的分配策略与适用场景

BumpPtrAllocator是 LLVM 里最核心的内存分配器之一,它的思路极其简单:维护一个大的内存块,每次分配时只需要把指针往后移动,不需要维护空闲链表,也不需要处理碎片。释放的时候,要么整体释放,要么什么都不做。

这种分配器的优势是分配速度极快,基本上就是一次指针加法;缺点是只能整体释放,不能单独释放某个对象。所以它适合那种“一批对象一起创建、一起销毁”的场景。在 LLVM 里,AST 节点、IR 指令、类型对象这些生命周期跟编译单元一致的对象,都适合用BumpPtrAllocator。

BumpPtrAllocator的内部实现是维护一个 slab 列表,每个 slab 是一块较大的内存(默认 4KB 或更大)。当当前 slab 剩余空间不够时,就分配一个新的 slab。分配时如果对象大小超过 slab 剩余空间,会直接在新 slab 上分配,旧 slab 的剩余空间就浪费了。所以 slab 大小的选择会影响内存利用率。

注意事项:BumpPtrAllocator不调用对象的析构函数。如果你分配的对象有非平凡的析构逻辑(比如持有std::vector),需要自己手动调用析构,或者用SpecificBumpPtrAllocator来管理。

3.2 SpecificBumpPtrAllocator 与对象生命周期管理

SpecificBumpPtrAllocator是BumpPtrAllocator的一个特化版本,它针对单一类型做分配,并且在析构时会调用所有已分配对象的析构函数。它的内部实现是为每个类型维护一个独立的BumpPtrAllocator,分配时按类型大小对齐,析构时遍历所有对象调用析构函数。

这个设计解决了一个关键问题:如何在批量分配的场景下,仍然保证对象的析构函数被正确调用。在 LLVM 里,很多 IR 对象需要在销毁时做一些清理工作,比如从某些全局数据结构中注销自己。SpecificBumpPtrAllocator让这些清理工作可以自动完成。

但它的限制是只能管理单一类型。如果你需要管理多种类型,要么用多个SpecificBumpPtrAllocator,要么用普通的BumpPtrAllocator加手动析构。我个人的经验是,如果一个模块里有三四种核心对象类型,用多个SpecificBumpPtrAllocator分别管理是最清晰的方案。

3.3 内存分配器的组合使用与性能对比

在实际的 LLVM 模块里,内存分配器往往是组合使用的。比如 Clang 的 ASTContext 里,既有BumpPtrAllocator用于分配 AST 节点,也有SpecificBumpPtrAllocator用于管理特定类型的对象,还有标准的malloc用于那些需要单独释放的对象。

这种组合使用的思路是:根据对象的生命周期和释放模式,选择最合适的分配器。生命周期跟编译单元一致的,用BumpPtrAllocator;需要单独释放的,用标准分配器;需要调用析构的批量对象,用SpecificBumpPtrAllocator。

从性能角度看,BumpPtrAllocator的分配速度比malloc快一个数量级左右,因为它不需要维护空闲链表,也不需要处理锁竞争(单线程场景下)。但它的内存利用率可能不如malloc,因为 slab 内部的碎片无法回收。所以在内存敏感的场景下,需要权衡分配速度和内存利用率。

4. 大模型推理中的内存管理问题梳理

4.1 KV Cache 的内存占用分析与瓶颈定位

大模型推理跟训练最大的不同在于,推理是自回归的,每生成一个 token,都需要用到之前所有 token 的 Key 和 Value 向量。这些向量如果每次都重新计算,计算量会大到无法接受。所以推理框架会把之前计算过的 Key 和 Value 缓存起来,这就是 KV Cache。

KV Cache 的内存占用可以用一个公式估算:2 * num_layers * num_heads * head_dim * seq_len * batch_size * dtype_size。以一个 7B 模型为例,32 层、32 个注意力头、head_dim 为 128、FP16 精度,单个 token 的 KV Cache 大小是2 * 32 * 32 * 128 * 2字节,也就是 512KB。如果序列长度是 2048,batch size 是 8,那 KV Cache 的总大小就是512KB * 2048 * 8,约 8GB。这还只是一个 7B 模型,如果是 70B 模型,KV Cache 的占用会大到单卡放不下。

所以 KV Cache 的内存管理是推理框架的核心问题之一。传统的做法是预分配一个最大长度的连续缓冲区,但这样会造成大量浪费,因为实际序列长度往往远小于最大长度。而且不同请求的序列长度不同,连续分配会导致严重的内存碎片。

4.2 PagedAttention 的分页管理思路

PagedAttention 的核心思路借鉴了操作系统的虚拟内存分页机制。它把 KV Cache 分成固定大小的块(block),每个块存储固定数量 token 的 Key 和 Value。不同请求的 KV Cache 可以共享块,也可以按需分配新块,不需要连续的内存空间。

这样做的好处是:第一,内存利用率大幅提升,因为块可以按需分配,不会预分配大量用不到的空间;第二,支持块级别的共享,比如多个请求有相同的 prompt 前缀时,可以共享前缀部分的 KV Cache 块;第三,内存碎片问题得到缓解,因为块的大小固定,分配和回收都很简单。

但 PagedAttention 也引入了新的复杂性。第一,注意力计算需要按块 gather,增加了 kernel 的实现难度;第二,块表的管理需要额外的内存和计算开销;第三,块的大小选择需要权衡,块太小管理开销大,块太大内存利用率下降。

4.3 显存池化与分配器设计

推理框架的显存分配器需要解决几个问题:分配速度快、碎片少、支持多卡和多流。常见的做法是实现一个显存池,预分配一大块显存,然后在这个池子里做子分配。

显存池的设计可以参考BumpPtrAllocator的思路:维护一个大的显存块,分配时移动指针,释放时标记空闲。但跟BumpPtrAllocator不同的是,显存池需要支持单独释放,因为不同请求的生命周期不同。所以通常会用空闲链表来管理已释放的块,分配时从空闲链表里找合适大小的块。

这里的一个关键问题是碎片管理。如果频繁分配和释放不同大小的块,会产生外部碎片,导致明明有足够的空闲显存,但无法分配出连续的大块。解决方案包括:按大小分类的空闲链表(segregated free list)、伙伴系统(buddy system)、或者定长块分配(类似 PagedAttention 的思路)。

5. LLVM 内存管理思路对大模型推理的启发

5.1 BumpPtrAllocator 思路在 KV Cache 管理中的应用

BumpPtrAllocator的核心思路是“批量分配、整体释放”,这个思路在 KV Cache 管理中有直接的应用场景。比如在一个 batch 的推理过程中,所有请求的 KV Cache 可以一起分配,等这个 batch 处理完之后一起释放。这样就不需要维护复杂的空闲链表,分配速度可以做到极快。

具体做法是:为每个 batch 预分配一块显存,按最大序列长度和 batch size 计算所需空间。每个请求的 KV Cache 在这块显存里按顺序排列,用一个偏移量表记录每个请求的起始位置。batch 处理完后,整块显存归还给显存池。这种方案适合离线推理或者 batch size 固定的在线推理场景。

但它的局限性也很明显:如果请求的序列长度差异很大,预分配最大长度会造成浪费;如果请求是动态到达的,batch 的组成会变化,整体释放的时机不好把握。所以这种方案更适合作为显存池的一个分配策略,而不是唯一的策略。

5.2 SmallVector 内联思路对推理请求调度的启示

SmallVector的内联存储思路可以借鉴到推理请求的调度中。比如每个推理请求的元数据(请求 ID、序列长度、采样参数等)通常不多,可以用一个内联的小数组来存储,避免为每个请求单独分配堆内存。当请求数量超过内联容量时,再切换到动态分配。

这个思路在请求调度器里特别有用,因为调度器需要频繁地遍历和修改请求列表。如果每个请求的元数据都是分散在堆上的,缓存局部性会很差;如果内联存储,遍历时缓存命中率会高很多。我实测过,在请求数量不多(比如几十个)的场景下,内联存储的调度器比堆分配的调度器快 20% 到 30%。

5.3 DenseMap 思路在块表管理中的借鉴

PagedAttention 的块表本质上是一个映射:逻辑块号到物理块号的映射。这个映射可以用DenseMap的思路来实现——用开放寻址法的哈希表,所有条目存在连续内存里,查找时缓存友好。

块表的特点是条目数量不大(通常几百到几千),但访问极其频繁(每次注意力计算都要查)。所以缓存局部性非常重要。用DenseMap风格的开放寻址哈希表,比用std::unordered_map的链式哈希,查找速度可以快 2 到 3 倍。而且块表的键是整数,哈希函数可以很简单,进一步降低开销。

实操心得:在实现块表时,如果块号是连续的整数,其实可以直接用数组而不是哈希表。数组的查找是 O(1) 且缓存最友好。只有当块号不连续或者需要支持稀疏映射时,才需要哈希表。我见过一些实现上来就用哈希表,其实用数组就够了,性能还能更好。

6. 实操中的常见问题与排查技巧

6.1 LLVM ADT 使用中的典型陷阱

第一个陷阱是SmallVector的引用失效。前面提到过,当SmallVector从内联存储切换到堆存储时,所有指针和引用都会失效。这个时机很隐蔽,因为从外部看,容器的大小可能只增加了一个元素,但底层存储已经换了。我踩过这个坑:在一个循环里持有SmallVector元素的引用,然后往容器里 push 新元素,结果引用失效导致数据错乱。

第二个陷阱是DenseMap的迭代器失效。DenseMap在扩容时,所有迭代器都会失效。而且DenseMap的迭代顺序是不确定的,不能依赖它做有序遍历。如果需要有序遍历,应该用MapVector或者先导出到std::vector再排序。

第三个陷阱是StringRef的生命周期。StringRef不拥有内存,如果它引用的字符串是临时对象,那StringRef很快就会变成悬空引用。常见的错误是把std::string的临时对象传给返回StringRef的函数,函数返回后临时对象销毁,StringRef就悬空了。

6.2 大模型推理内存问题的排查思路

推理内存问题的排查,我一般按这个顺序来:第一,看显存占用是否符合预期,用公式估算一下 KV Cache 的理论大小,跟实际占用对比;第二,看是否有内存泄漏,长时间运行后显存占用是否持续增长;第三,看是否有碎片问题,分配大块显存时是否失败,但总空闲显存明明够;第四,看是否有频繁的分配释放,用 profiling 工具统计分配次数和耗时。

排查碎片问题时,一个有用的技巧是打印显存池的空闲块分布。如果空闲块很多但都很小,说明碎片严重;如果空闲块很少但都很大,说明池子快用完了。根据这个分布,可以决定是调整块大小,还是引入更复杂的碎片整理机制。

6.3 性能调优的实操建议

对于 LLVM ADT 的性能调优,我的建议是:先用 profiling 找到热点,再针对性优化。不要一上来就把所有std::vector换成SmallVector,因为SmallVector的内联存储会增加对象大小,如果对象本身很大或者数量很多,反而会拖慢性能。

对于推理内存的性能调优,关键是减少分配次数和提升缓存局部性。减少分配次数的方法包括:预分配、池化、批量分配。提升缓存局部性的方法包括:把频繁访问的数据放在一起、用连续内存而不是链表、避免指针追逐。

下面这个表格总结了我遇到过的典型问题、排查方法和解决方案:

问题现象可能原因排查方法解决方案
显存占用远超理论值KV Cache 预分配过大对比理论公式和实际占用改用按需分配或分页管理
长时间运行显存持续增长内存泄漏定期打印显存池状态检查释放逻辑,用 RAII 管理
分配大块显存失败但总空闲够内存碎片打印空闲块分布引入碎片整理或定长块分配
推理速度波动大频繁分配释放profiling 分配耗时预分配或池化
注意力计算 kernel 慢块表查找开销大profiling kernel 耗时优化块表数据结构

7. 从编译器到推理框架的跨领域思考

7.1 内存管理抽象层次的对比

LLVM 的内存管理抽象层次比较清晰:最底层是malloc和free,往上是BumpPtrAllocator和SpecificBumpPtrAllocator,再往上是各种容器(SmallVector、DenseMap等),最上层是具体的 IR 对象和 AST 节点。每一层都有明确的职责和适用场景。

推理框架的内存管理抽象层次相对模糊一些。很多框架直接在最上层做显存分配,没有中间的分配器层。这导致的问题是:分配策略跟业务逻辑耦合,难以复用和替换。我觉得推理框架可以借鉴 LLVM 的分层思路,把显存分配抽象成独立的分配器层,上层业务只关心“我要一块显存”,不关心底层是怎么分配的。

7.2 生命周期管理的异同

LLVM 里对象的生命周期通常跟编译单元一致,所以可以用BumpPtrAllocator做批量管理。推理框架里对象的生命周期差异很大:模型权重跟进程一致,KV Cache 跟请求一致,中间激活值跟算子一致。所以推理框架需要更细粒度的生命周期管理。

但两者的共同点是:都需要根据生命周期选择合适的分配策略。LLVM 用SpecificBumpPtrAllocator管理需要析构的批量对象,推理框架可以用类似思路管理 KV Cache——按请求批量分配,请求结束时批量释放。

7.3 可借鉴的设计模式总结

从 LLVM 的内存管理里,我觉得推理框架可以借鉴这几个设计模式:第一,分层抽象,把分配策略跟业务逻辑解耦;第二,批量分配整体释放,减少分配次数和管理开销;第三,内联存储,对小对象避免堆分配;第四,开放寻址哈希表,提升缓存局部性;第五,按生命周期选择分配器,而不是一刀切。

这些模式不是万能的,需要根据具体场景权衡。比如内联存储会增加对象大小,如果对象数量很多,反而会浪费内存。开放寻址哈希表在删除频繁的场景下性能会下降。所以关键是理解每个模式的适用条件和代价,而不是盲目套用。

我在实际做推理优化的时候,最大的体会是:不要一开始就追求极致的优化,先把功能做对,再用 profiling 找瓶颈,最后针对性优化。很多性能问题不是靠复杂的技巧解决的,而是靠减少不必要的分配、提升缓存局部性这些基础工作解决的。LLVM 的内存管理之所以高效,很大程度上是因为它在设计之初就把这些基础工作做对了,而不是靠什么黑科技。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询