从LLVM ADT到大模型推理：内存管理与算子注册的底层共通逻辑

最近整理笔记的时候我把两件看起来八竿子打不着的事情放在了一起一边是 LLVM 内部的 ADT 与内存管理机制一边是大模型推理框架里的显存调度和算子组织。本来只是顺手翻翻结果越看越觉得它们其实是同一套底层逻辑的不同舞台。这篇随笔主要记录我自己的梳理路径也顺手把踩过的坑和验证过的心得写出来无论你是刚开始接触 LLVM 工具链的 C 开发者还是正在啃大模型推理框架的算法工程师应该都能从里面找到一点能立刻用上的东西。先说结论LLVM ADT 并不是“编译器专属数据结构”它是一套被高性能工程反复验证过的 C 基础设施而大模型推理里的 KV Cache、算子注册、显存池管理本质上也躲不开内存生命周期、零拷贝视图和编译器调用约定这几件事。把两边一起看能少走很多弯路。1. 为什么值得先啃 LLVM ADT 这口硬骨头我第一次看到llvm/ADT目录的时候心里想的是“这是编译器内部的数据结构我平时写业务代码又用不上”。后来跟别人讨论大模型推理框架的源码发现里面大量出现 small-vector、string-view、arena allocator 这类设计才意识到自己把宝山当成了普通山头。1.1 它不是编译器专属的工具库LLVM 项目的llvm/ADT下有一批完全独立的头文件比如SmallVector.h、StringRef.h、ArrayRef.h、StringMap.h、DenseMap.h它们不依赖 LLVM 后端的编译流程单独拿出来就能作为通用 C 工具集使用。StringRef和ArrayRef本质是“指针 长度”的轻量视图。SmallVectorT, N则是一种内嵌了 N 个元素缓冲区的动态数组元素数量少的时候直接使用栈上/对象内的一部分内存超过阈值才切到堆分配。StringMap和DenseMap是专门为高频查找优化过的哈希表内部实现比std::unordered_map更加可预测。这些类在 LLVM 里被用来存放指令列表、基本块引用、寄存器名字符串、编译期临时对象性能要求极高。但放在大模型推理场景里同样成立推理引擎的调度器要反复收集 token id 列表op 名要反复查找tensor shape 和 kernel 参数要反复传递这些东西用同一套思路去做收益非常直接。1.2 很多“高性能”来自对自己内存的掌控我学习 LLVM 内存管理观最大的收获是理解了“高性能不是靠某个魔法而是靠把分配和释放控制住”。LLVM 内部大量使用BumpPtrAllocator它一次性从底层分配一大块内存然后在内存池里按需切出小块用完以后整块释放。这么做的好处非常明显分配小块对象时只是移动一个指针开销远比malloc低同时因为大量对象集中在一整块连续内存里缓存命中率也更好。malloc本身有锁、有元数据、可能有内存碎片问题在高频编译路径上根本耗不起。大模型推理也一样。如果你在 CUDA 上频繁调用cudaMalloc给中间 tensor 分配显存开销会高到让人崩溃。所以几乎所有的推理框架都会维护显存池或者 block 缓存把一次分配、反复使用、批量释放这套逻辑做进去。从抽象层面看这就是 LLVM 里BumpPtrAllocator的显存版本。1.3 llvm gcc msvc 背后的工具链差异网上一搜 “llvm gcc msvc”大多数讨论都在对比编译器前端和工具链很多人觉得这跟推理沾不上边。实际上当你写自定义 GPU 算子或者给推理框架写 C 扩展时调用约定、结构体对齐、异常处理这些问题直接决定了程序能不能稳定跑起来。比如同一个结构体用 MSVC 和用 Clang/GCC 编译内存布局可能有区别函数的 calling convention 不同参数传递方式也不同。如果你在 PyTorch 里通过pybind11绑定一个自定义算子宿主侧和扩展侧用的编译器不一致很容易出现数据错位或者诡异的段错误。理解 LLVM 是怎么设计内存布局和调用边界的能帮你在排查这类问题时更快定位方向。2. 高频 ADT 与内存机制的使用要点这一节挑几个我在实际项目里真正用过的 LLVM ADT 和内存工具来讲每一个都会配上代码片段和“为什么这么设计”的解释。2.1 SmallVector小对象分配优先走栈llvm::SmallVectorT, N是 LLVM 里特别常用的容器它的核心思想是如果元素不超过 N 个就直接使用对象内部预留的缓冲区不走堆分配。只有当元素超过 N 时才切换到类似std::vector的动态扩容逻辑。#include llvm/ADT/SmallVector.h llvm::SmallVectorint, 8 collect_even(llvm::ArrayRefint input) { llvm::SmallVectorint, 8 out; for (int v : input) { if (v % 2 0) out.push_back(v); } return out; }当out中不超过 8 个元素时这次函数调用不会触发堆分配。这类局部小列表在编译器里到处都是在推理调度器里也很多比如收集当前批次需要 decode 的 sequence id、收集一组 kernel 参数。我踩过的一个坑是不要把这个优化误解为“完全不会导致内存分配”。当元素数量超过 N 后它仍然会走堆分配而且为了兼容移动语义还会涉及拷贝或移动。所以N的选择要有依据拍脑袋选一个 1024 并不会更好反而会把结构体撑得很大。最合理的做法是统计线上真实数据的分位数选择覆盖绝大多数情况的阈值。2.2 ArrayRef / StringRef零拷贝的视图但也最容易翻车ArrayRefT和StringRef是我最喜欢的两个 LLVM ADT因为它们把一个“数组/字符串”表示成“指针 长度”传参的时候不复制底层数据。写出来的代码更清爽运行时开销也更小。#include llvm/ADT/StringRef.h #include llvm/ADT/ArrayRef.h void print_names(llvm::ArrayRefllvm::StringRef names) { for (auto name : names) { llvm::outs() name.size() : name \n; } }StringRef还提供了split、find、drop_front等方法处理字符串时完全不需要创建一堆临时std::string。比如在推理框架里解析模型名和算子名的时候省掉这些临时对象能让高频路径的 CPU 开销降一个量级。但这里有一个非常关键的陷阱视图类不拥有底层数据。如果底层的std::string或者字符缓冲区被销毁StringRef就成了悬空引用。我在早期写过类似的代码llvm::StringRef get_name(int id) { std::string s tensor_ std::to_string(id); return s; // 返回后就悬空了 }这样做必炸。正确的做法是要么让底层数据的生命周期覆盖到视图使用结束要么显式把数据拷贝到拥有所有权的容器里。用 C 写高性能服务才能真正理解“所有权”这三个字而不是靠垃圾回收兜底。2.3 BumpPtrAllocator一次性分配、整体释放llvm::BumpPtrAllocator是我觉得最值得研究的一个内存分配器。它的逻辑是维护一个或几个大块内存Allocate的时候直接返回当前偏移处的地址然后偏移向前推进空间不够时再向系统申请新的大块。它不提供单个对象的free而是提供整块Reset。#include llvm/Support/Allocator.h struct OpMetadata { int type; const char* name; }; llvm::BumpPtrAllocator alloc; void create_metadata(const char* name) { auto *meta new (alloc.AllocateOpMetadata()) OpMetadata; meta-name name; // 使用完不逐个 delete } void clear_metadata() { alloc.Reset(); }Reset不会为已构造的对象调用析构函数它只是把当前偏移恢复到初始位置下一轮分配会覆盖之前的内存。所以这种分配器适合 trivially destructible 的临时对象。如果你需要管理持有智能指针或引用计数的对象就得手动清理否则会出现内存泄漏或者更隐蔽的悬空引用。这个特性很难不让人联想到显存池把一整块显存圈起来按 block 或者按线性 offset 分配最后整段释放或者整段回收。两者在设计哲学上是完全一致的。2.4 StringMap 与算子注册表的关系llvm::StringMapT是专门为字符串 key 设计的哈希表内部使用线性探测对短键和只读查找特别友好。我在看某些推理框架的算子注册表实现时发现思路跟它几乎一样算子名字符串作为 key 映射到 kernel 函数指针或者一个注册条目。#include llvm/ADT/StringMap.h using KernelFunc void (*)(const void* args); llvm::StringMapKernelFunc registry; void register_all_ops() { registry[rms_norm] rms_norm_kernel; registry[matmul] matmul_kernel; } KernelFunc find_op(const char* op_name) { auto it registry.find(op_name); return it registry.end() ? nullptr : it-second; }这种“注册中心 字符串查找”的模型在 LLVM 的 Pass 注册机制里也很常见在大模型推理框架中同样普遍。它的本质不是魔法而是编译期或加载期注册、运行期查表。理解了这一点你再听别人讨论“算子自发现”时就不会觉得里面有什么玄学核心就是一张表加一个查找流程。3. 拉通大模型推理的关键模块从 LLVM ADT 往大模型推理上迁移的时候首先要做的是把“编译器内部机制”和“推理引擎机制”之间的对应关系找出来。下面我用自己在 nano-vllm 这类最小化推理工程里学到的几个关键功能做一个梳理。3.1 从 nano-vllm 这类最小工程入手我比较推荐从名字里带 “nano” 的最小化推理工程入手而不是一上来就啃仓库几十万行代码的完整推理框架。这类工程通常会砍掉大量工程冗余只保留几个核心环节请求调度、KV Cache 管理、算子执行、采样输出。举个例子很多最小化推理工程会用一个调度器来判定当前这批请求应该做 prefill预填充还是 decode逐 token 生成然后按批把请求分组统一执行。这里每一步都会产生临时列表比如当前批次的 token id、sequence length、attention mask 索引。如果你用标准库容器频繁构造开销很大的临时对象性能就会直接受影响用SmallVectorArrayRef这类思路能明显降低 CPU 开销。研究这类小工程还有个额外好处你能看到真实推理框架里 KV Cache 是怎么被划分和复用的。等到再看完整实现时就不会被细节淹没能直接抓住主线。3.2 算子自发现是如何发生的“算子自发现”是我最近反复看到的一个热词也是很多人觉得神秘的地方。大型模型推理框架在加载模型时经常需要在众多 ops 里找到对应 runtime 的实现这个过程叫 “self-discovery” 也好叫 “operator registry” 也好落到代码层面都是三步第一步定义统一的算子接口例如输入统一结构体或统一的参数包第二步在算子实现位置通过静态变量注册把 “名字 - 函数指针/工厂函数” 插入全局注册表第三步运行时通过字符串 key 去注册表查找到实现然后调用。在 CUDA 场景里自定义算子常常会注册一个 kernel 启动函数在 MLIR 或者 LLVM 生态里Pass 和 Dialect 也是用类似方式做自发现。因此你会发现用StringMap当注册表简直再合适不过查找稳定、内存可控、没有std::unordered_map那种迭代器失效和哈希桶抖动问题。如果你自己写一个推理框架的插件系统也可以照着这个模式插件加载时调register_plugin_ops把 op 名写入一张静态StringMap引擎启动时统一构建。这样做的好处是动态库加载顺序不影响最终结果因为注册是在加载期发生的不是运行时扫描。3.3 KV Cache 生命周期与内存管理器思路KV Cache 是解码阶段存储历史 key 和 value 的缓存显存占用往往占大头。我刚开始学的时候以为它只是给每一条请求一次性分配一块大显存后来发现根本不是这样推理过程里序列长度不断增长如果按照最大长度预分配显存浪费严重如果每次增长都重新分配又会造成大量拷贝和碎片。很多框架会采用分块缓存的方式把 KV Cache 切成固定大小的 block按需分配。这跟BumpPtrAllocator的“一次分配一个大块逻辑上切成小节”相比多了一个“可以回收单块”的能力。KV Cache 分配器里通常维护两个列表空闲 block 和已分配 block。请求新 block 时从空闲列表取出释放时标记回空闲列表。为了快速判断某个 block 是否可以释放还会维护引用计数。这里用到的底层数据结构并不复杂但是“谁持有 block、谁负责释放”的生命周期管理非常严格稍不留神就会导致显存泄漏。从 LLVM ADT 的角度去思考这个问题我体会到的是分配器的设计不是孤立的必须配合数据的生命周期模型。谁来构造对象谁来负责析构谁来释放底层资源这三个问题必须在代码层面非常明确。3.4 一张表把两套知识对齐每次梳理这类知识我都会画一张对应表让 LLVM 的底层机制和大模型推理的上层概念对齐起来后面再复习也会很快。LLVM ADT / 内存机制大模型推理中的对应点共同关注点SmallVector请求 token id 列表、局部索引数组小对象避免频繁堆分配ArrayRef / StringReftensor 名字、op 名字的零拷贝传递与查找视图生命周期管理StringMap算子注册表、参数索引表字符串哈希查找效率BumpPtrAllocator显存池、block 缓存预分配一次性大块分配、整块释放llvm / gcc / msvc 调用约定自定义算子的宿主与设备调用边界结构体对齐、ABI 稳定性静态注册机制算子自发现、插件注册表注册期解析、运行期查表表格写出来之后你会发现两件事一是“高性能”没有灵丹妙药都是把内存分配和对象生命周期管明白二是不同领域的术语很不一样但底层的设计模式惊人一致。4. 常见问题与排查心得这篇文章里的知识我并不是第一次就完全弄懂的。相反好多问题是在实操里踩出来的简单整理成速查表希望你不用再踩一遍。4.1 视图类生命周期错乱症状StringRef或者ArrayRef在使用时出现随机数据错乱有时候打印出来是对的有时候是一堆乱码。 原因要么是视图超出底层缓冲区生命周期要么是底层容器发生了 reallocation之前的指针失效。 排查手法先用 AddressSanitizer 跑一遍小规模用例。如果是在 LLVM 头文件环境下写 C可以给编译加-fsanitizeaddress如果是在推理框架里可以先在 Python 侧用gc.disable()临时关闭回收看是否会稳定复现再往下找 C 侧的生命周期问题。 修复原则让视图类只是临时参数不要长期存储如果必须长期存储就换成具有所有权的容器例如std::string或者自定义的SharedString。4.2 显存碎片和频繁 cudaMalloc症状推理过程刚开始正常运行一段时间后偶发 CUDA out of memory但显存总量并没有全部占满。 原因高频cudaMalloc/cudaFree造成碎片化显存被切成很多不可用的小块。 排查手法打开推理框架的显存日志统计分配次数和释放次数重点看是否存在每次请求动态分配大块显存的情况。 对策启动时按照gpu_memory_utilization规划好显存池尽量把大块显存预分配出来然后在这个池子上做逻辑分配。在 LLVM 的BumpPtrAllocator思想里这就是“一次性大块、内部切分”。4.3 学习过程中的“假懂”状态我犯过的一个典型错误看文档看得懂看源码觉得眼熟但让自己动手写一个小的分配器或者算子注册表时卡了很久。这就是“假懂”状态。 判断方法很简单能不能不看源码自己把BumpPtrAllocator的核心逻辑写出来能不能自己设计一张算子注册表并保证并发安全如果写不出来就说明还没有真正消化。 改变学习的动作不要只看不写可以自己写一个 200 行以内的小例子把分配、注册、查找、清理这四个环节都走一遍。这个过程比读十篇博客都有效。4.4 一个值得跑的最小实验如果你想亲手试一下 LLVM ADT 的编译方式可以在装了 LLVM 开发包的机器上写一个小的 C 文件然后编译运行。#include llvm/ADT/SmallVector.h #include llvm/ADT/StringMap.h #include llvm/Support/raw_ostream.h int main() { llvm::SmallVectorint, 4 data {1, 2, 3, 4}; data.push_back(5); llvm::outs() size data.size() \n; llvm::StringMapint map; map[prefill] 1; map[decode] 2; llvm::outs() decode map[decode] \n; return 0; }编译时只要让编译器能找到 LLVM 头文件路径即可。在 Debian/Ubuntu 系统上安装好llvm-dev后通常可以通过llvm-config --includedir拿到头文件路径。这个实验本身不需要链接庞大的 libLLVM因为只用到了头文件模板实现。这个实验跑通以后你再回头去看大模型推理框架里的调度器代码会发现很多局部索引收集、op 名查找、临时元数据存储用的都是同一套模式。5. 如果重来一次我会这么切入如果时间倒流回到刚开始接触大模型推理的时候我不会一上来就去细扣每个算子的数学公式而是先花两周时间把 LLVM 的 ADT 和内存管理机制过一遍然后用一个最小的推理工程做对照。我会先写出自己的 “最小分配器” 和 “最小算子注册表”然后问自己三个问题它能不能处理生命周期它能不能在频繁分配时保持低开销它能不能在并发状态下保证稳定这三个问题一旦想清楚整条知识线就通了。我个人在实际操作中的体会是内存管理的代码看起来枯燥却是理解所有高性能系统的钥匙。LLVM ADT 里的SmallVector、StringRef、BumpPtrAllocator和大模型推理里的显存池、KV Cache、算子注册表本质上都是同一套思路在不同领域的实现。你不需要把两套体系当成孤立的知识去背把底层逻辑打通后面再看任何框架源码都能快很多。最后再分享一个小技巧记录这类笔记时不要只记结论要把当时的排查过程写下来比如“我第一次用 StringRef 返回局部字符串时发生了什么”“显存碎片出现前有哪些征兆”。这些过程才是真正有价值的经验之后回看会有很强的复现感。