最近在调一个 GPU kernel 的性能，火焰图拍出来吓我一跳，热点根本不在计算逻辑上，而是大量时间耗在反复 malloc/free 临时数组上。项目里的老人给了一句提点：去看看 LLVM 怎么管内存的。我从SmallVector一路看到BumpPtrAllocator，又顺着这个思路去理大模型推理框架里的 KV cache 管理，突然发现这两件事的内核惊人地一致。这篇文章就是这段时间的笔记，核心围绕 LLVM ADT 与内存管理机制，以及我基于 nano-vllm 学习大模型推理关键路径的梳理。内容有点杂，但都是实际调试和读源码的真实体会，对做底层优化、推理框架二次开发的朋友应该有点参考价值。
1. 为什么 LLVM 要自建一套 ADT 而不是直接用 STL
1.1 一个真实的小场景
先看一个我踩过的实际例子。写算子融合逻辑的时候，需要给每个中间结果维护一个不固定长度的索引列表。最自然的写法是std::vector<int>，每次往里面push_back一个元素就可能在堆上开一块内存。如果这段代码在循环里跑几百万次，堆分配的开销立刻变成大头。你可能会说std::vector会预留容量，但关键是你不确定到底该预留多少，预留多了浪费，预留少了照样频繁扩容。
换成llvm::SmallVector<int, 4>之后，前 4 个元素直接存在栈上，只有超过 4 个才切换到堆分配。对于大多数场景里"元素个数通常没几个，偶尔会多到十几个"的情况，这个设计几乎把堆分配次数降到了零。实测下来，某个 pass 的总耗时直接砍掉了百分之十几，而代码改动只有一行。这就是 LLVM 造这套 ADT 的出发点：不要用通用方案去适配所有情况，而是针对编译器特有的数据访问规律做特化。
1.2 SmallVector 与 StringRef 的使用要领
SmallVector只是 LLVM ADT 家族里的冰山一角，但它最值得研究。模板签名是SmallVector<T, N>，其中N是栈上预留的元素个数。这里有个很容易被忽略的点：N直接决定这个对象占用多少栈空间。N=4时栈上就多出4 * sizeof(T)的字节，如果你开一个函数局部变量，问题不大；但如果一个结构体里塞了好几个SmallVector，而且N设置得很大，每个实例的栈占用会累积到一个危险的水平。我见过有人写SmallVector<char, 8192>，然后把它作为函数参数按值传递，每次拷贝都是 8KB 的栈复制，性能反而比std::string还差。
StringRef是另一个高频使用的 ADT，本质上就是const char* + size_t，只持有视图，不拥有数据。传参时能避免字符串拷贝，但生命周期问题也随之而来：如果返回一个局部std::string对应的StringRef，函数返回后字符串销毁，这个引用就悬垂了。我在实际项目里的习惯是：拿StringRef做参数和中间检索没问题，但只要涉及"把名字存下来以备后用"，一定立刻转成std::string。
1.3 LLVM 拒绝 STL 的几个硬核理由
很多人奇怪 LLVM 为什么不用现成的 STL，非要自己搞一套。从源码里能看出几个清晰的原因：
第一是性能控制。STL 容器是通用实现，要考虑各种分配场景，所以std::string即使有 SSO 优化，行为也不好精确控制。LLVM 的代码跑起来经常是几小时起步的长时间任务，每一处多余的堆分配都会累积成肉眼可见的耗时。它需要的是"针对编译器 pass 临时集合场景特化"的方案，而非通用容器。
第二是调试与断言。LLVM ADT 内建了大量断言和迭代器检查，在-DLLVM_ENABLE_ASSERTIONS=ON的构建里，越界访问会直接崩溃提示，定位问题快得多。而且这些断言在 release 构建里可以彻底去掉，不留下任何运行时开销。
第三是编译速度和依赖控制。STL 头文件又重又长，LLVM 这类大体量项目如果到处 include，编译时间会爆炸。自建一套轻量 ADT 头文件，很多核心模块只需引入很小的依赖，整个工程的编译层次会清爽很多。
这套 ADT 设计的思路，其实和推理框架里的 tensor 管理有不少相通之处。框架里最常见的就是形状、索引、名字这类小对象，如果每个都用std::vector包一层堆分配，调度器的延迟会非常难看。你完全可以照搬SmallVector的思路，用固定小数组加切换机制来管理这些临时元数据。
2. BumpPtrAllocator 与 LLVM 的内存管理哲学
2.1 碰撞指针分配到底做了什么
BumpPtrAllocator（也叫 arena allocator）是 LLVM 内存管理的一个基石。它的原理简单到不可思议：维护一个指针，指向当前内存块的"碰撞位置"。每次Allocate(size, align)调用时，先把指针按对齐要求向上取整，然后把这块内存返回给调用者，再把指针推进 size 字节。如果当前块剩余空间不够，就向系统申请一个新块，指针重新从块头开始。
来看一段简化示意：
class BumpPtrAllocator { char *CurPtr = nullptr; char *End = nullptr; public: void *Allocate(size_t Size, size_t Align) { uintptr_t Aligned = (uintptr_t(CurPtr) + Align - 1) & ~(Align - 1); if (Aligned + Size > uintptr_t(End)) { // 当前块不够，分配一个更大的新块 NewBlock(); Aligned = align(CurPtr, Align); } CurPtr = reinterpret_cast<char*>(Aligned + Size); return reinterpret_cast<void*>(Aligned); } };所有从这里拿出去的内存，没有单独的free。释放动作只有一个：当整个编译阶段结束时，把所有块的头部信息遍历一遍，整体归还给系统。单次分配的耗时只有几次指针运算和一次比较，代价是"内存一旦分配你就不能单独回收，只能等整个 arena 一起销毁"。
2.2 LLVM 里谁在用 BumpPtrAllocator
编译器场景天然适配这种"一代一代分配、一次性清空"的模式。LLVM 源码里最典型的就是LLVMContext和各个 pass 管理Instruction、BasicBlock的方式：一个函数被解析后，会创建成百上千个小对象，这些对象内部互相引用，生命周期几乎与这个函数在 pipeline 里的存活时间一致。
如果逐个new再逐个delete，不仅调用次数爆炸，还会因为对象分配顺序的碎片化显著降低缓存命中率。BumpPtrAllocator 把这一批对象放在连续内存里，遍历指令时 CPU 的 cache line 能连续命中，效果立竿见影。另外，这些对象之间含有大量指针引用，用引用计数或智能指针维护ownership 会陷入循环引用的泥潭。LLVM 的选择是：分析阶段对象本身就是"不可变的、被容器整体拥有"的，裸指针只做观测和关联，内存统一归 arena 管。
这是很重要的一个工程哲学：所有权模型一定要单一且清晰。如果你在框架里发现某个对象的析构时机难以确定，往往不是漏了 delete，而是 ownership 设计得太混乱。
2.3 与 KV Cache 显存管理的深层对照
我在读 nano-vllm 的 KV cache 实现时，脑子里一直浮现 BumpPtrAllocator 的影子，因为两者的核心策略实在太像了：先向系统要一大块连续资源，再在内部做精细切的切分和复用。
以 vLLM 系的 PagedAttention 为例，显存里会预分配一个大的 KV cache 池，切成固定大小（比如 16 个 token 一个 block）的物理块。请求进来时，KV cache 管理器按需分配 block，记录下来形成一张 block table，每个逻辑位置对应一个物理块编号。和 BumpPtrAllocator 一样，它根本不做随机的单独 free，而是把"释放"设计成了"把 block 归还到自由块队列"，等下一个请求复用。
更妙的是调度器也依赖这个模型：如果自由块不足，新请求就等下一轮调度；如果请求结束，它持有的 block 一次性全部归还。整个过程没有零散的显存分配，没有碎片化，性能也稳定。可以说，理解了 LLVM 的 arena 思路，再看 KV cache 的 block table 设计，几乎不需要额外费劲。
3. 基于 nano-vllm 梳理大模型推理的关键路径
3.1 从请求到 token：整个推理流程的骨架
nano-vllm 的价值在于它把大模型推理框架里那些工程上的边角料全部剥掉，只留核心路径，非常适合用来做全景式的功能梳理。把一个请求从进入系统到吐出最后一个 token，拆开来看就是一条清晰的流水线：
请求进来后先被封装成一个 Sequence 对象，包含 prompt 文本、请求 ID、最大生成长度等元数据。随后进入调度器，调度器决定它能不能参加下一轮 forward。注意这里的粒度不是"等上一个请求完全生成完再处理下一个"，而是每一轮都动态地组合一批可以执行的请求，这就是连续批处理（continuous batching）的精髓。
一旦被选中，框架要给这个请求分配 KV cache block。如果是这个请求第一次参与计算，就要执行 prefill：把整个 prompt 一次性前向计算，把过程中产生的 K 矩阵和 V 矩阵写入刚分配的 block 中。之后进入 decode 循环，每轮 forward 的输入只有一个新 token，模型读取之前存好的 KV cache，计算出当前步的输出 logits，再经过采样选出下一个 token。如此反复，直到生成结束条件出现。
3.2 调度器与连续批处理
调度器在 nano-vllm 里占据的位置，相当于编译器里的 pass manager：它决定哪些指令（请求）在哪个时刻（迭代轮次）被执行。核心数据结构是就绪队列和等待队列。等待队列里是还没拿到足够 KV cache block 的请求，就绪队列里的可以参与下一轮 forward。
调度器每轮要做的事可以归纳成三步：
- 统计当前自由物理 block 的数量。
- 从等待队列中尽量多地取出请求，估算它们需要的 block 数，能放下就转成就绪。
- 把就绪队列里的所有请求构造成一个 batch，交给模型执行层。
连续批处理的关键点是 prefill 和 decode 可以共存于同一个 batch。有些请求刚进来，在做漫长的 prefill；另一些已经在逐 token 生成。传统方案会强迫它们分开跑，GPU 在 prefill 阶段大量并行、decode 阶段又只能吃满极低利用率，计算资源被白白浪费。而这个混合 batch 能在时间轴上更均匀地铺满 GPU，整体吞吐高很多。
这和我们做系统设计时的直觉很像：不要把所有请求都当成同一种负载，要识别它们各自所处阶段，然后让不同阶段的任务在空间上共存。
3.3 KV cache 与 block table 的实现思路
KV cache 管理是推理框架里最容易被低估的模块。没有它，模型的 decode 步只能从头重新计算整个 prompt，那计算量是灾难级的。有了它，每个步只需要读取历史 K/V，计算复杂度只和已生成的序列长度有关。
nano-vllm 里 block table 的实现思路可以用一张逻辑到物理的映射表来理解：一个请求逻辑上拥有连续的 token 编号 0、1、2…… 但它们在显存里对应的 block 编号可能是 7、3、19，完全打散的。这样做的收益有两个：
第一，物理块可以按需分配和回收，避免 KV cache 整体搬迁。第二，一个请求的 block 数不再要求物理连续，显存碎片被限制在 block 粒度内，而不是字节粒度。
struct BlockTable { // logical_block_id -> physical_block_id std::vector<int> mapping; void FreeRequest(int logicalStart, int logicalEnd) { for (int i = logicalStart; i < logicalEnd; ++i) { int phys = mapping[i]; freeList.push_back(phys); // 归还物理块 mapping[i] = -1; } } };这里freeList相当于 BumpPtrAllocator 里的自由块链表。需要强调一点：这种设计不允许单个 block 内部的部分 token 被单独释放。释放粒度永远和 block 对齐，才能保证分配器简单高效。实际分配 block 时，我给框架加过一层按代际编号的 block 池，每次空闲检测只需要比较代际号，能比直接遍历 block 的引用计数快一个量级。这种"代际号 + 空闲队列"的思路，也是从 LLVM 的 free list allocator 里迁移过来的。
4. 算子自发现：LLVM 模式匹配与推理 kernel 调度的统一思想
4.1 LLVM 是怎么"找"到算子的
"算子自发现"这个说法，在 LLVM 生态里其实是指模式匹配与指令选择机制。编译器里有成百上千种指令，向量化 pass 怎么知道某段 IR 匹配哪个目标指令？靠的是 TableGen 自动生成的匹配器：先用一套.td描述文件声明指令模式，再用工具生成 C++ 查询代码，运行时以模式树的方式从操作数向上匹配。整个过程是由编译器“自动发现”适用的指令，而不是手动写一堆 if-else。
另一个更基础的模式匹配工具是isa、dyn_cast这一组接口。它们基于运行时类型信息，帮你快速判断一条指令是不是某个子类。比如：
if (auto *Call = dyn_cast<CallInst>(&I)) { // 只有 I 真的是 CallInst 时才会进入这里 }看起来很朴素，但它把"类型分支"变成了"声明式匹配"。我在写推理框架的算子分发时，把这种思路平移了过来：与其在每个调用点写满if (op == "gemm" || op == "matmul")这类硬编码，不如让每个算子自己声明它适用的 shape 和 dtype 约束，然后由一个统一的匹配器在注册表里找最优实现。
4.2 推理框架的 kernel dispatch 与注册表模式
GPU kernel 的选择远比表面看起来复杂。同一个矩阵乘法，输入是 fp16 还是 fp32、shape 是 1x4096 还是 64x4096，对应 kernel 的启动参数甚至 kernel 本身都可能不同。早期实现里，这些选择逻辑散落在各个调用点，耦合度极高，加一个新 kernel 要翻遍所有相关代码。而"算子自发现"的注册表模式能根治这个问题。
做法是：算子用一个全局的注册表保存"名字 -> 工厂函数"的映射。每个 kernel 通过静态变量在程序启动阶段自动把自己的工厂函数塞进注册表，调度器拿到一个算子请求时，只需查表并调用工厂函数即可。新增 kernel 只需要新增一个注册语句，完全不用改动调度主逻辑。
using KernelFactory = void (*)(const Tensor&, const Tensor&, Tensor&); struct KernelRegistry { static KernelRegistry& Get() { static KernelRegistry R; return R; } void Register(const char* name, KernelFactory f) { factories[std::string(name)] = f; } KernelFactory Find(const std::string& name) { auto it = factories.find(name); if (it == factories.end()) return nullptr; return it->second; } private: std::map<std::string, KernelFactory> factories; }; struct AutoRegister { AutoRegister(const char* name, KernelFactory f) { KernelRegistry::Get().Register(name, f); } }; // 在某个 kernel 文件的全局作用域里注册 static void GemmFp16Kernel(const Tensor&, const Tensor&, Tensor&) { /* ... */ } static AutoRegister reg_gemm_fp16("gemm_fp16", GemmFp16Kernel);这套模式的好处是内核之间完全解耦。你甚至可以把它推广成自发现的有向无环图调度：算子描述自己的输入输出和依赖，调度器在注册表里根据依赖关系自动连边、自动排序，这就是深度学习框架里 graph lowering 的雏形。相比手写连接代码，这种声明式方案在高版本 vLLM 的 op 自动查找机制里已成为默认策略。
4.3 实测：一个算子注册与自发现的小 demo
我在 nano-vllm 的代码基础上加过一个很小的自发现扩展，验证这个思路的效果。场景是这样的：框架里有多个计算算子，我新增了一个基于自定义切分策略的 GEMM 实现，如果按照传统写法，需要向所有 dispatch 分支里插入调用。用注册表模式后，只需要在扩展文件里加一个static AutoRegister，程序一启动，新算子自动被"发现"，调度主逻辑一行没动。
配合 shape 和 dtype 的判定条件，我做了三个版本的 kernel：fp16 大矩阵、fp16 小矩阵、fp32 通用回退。实测下来，一个大 shape 的 batch 在 fp16 专用 kernel 上获得了约 25% 的加速，而小 shape 的请求也不会因为误选大 kernel 而被拖慢。这个"按规模自动选择最优 kernel"的雏形，本质上就是在复刻 LLVM 指令选择器里"按模式选择最优指令"的思想。
5. 实战踩坑记录与排查心得
5.1 内存生命周期相关的坑
使用StringRef和ArrayRef最大的坑就是悬垂引用。我印象最深的是自己写过这么一段代码：
llvm::StringRef getOpName() { std::string name = "temp_name"; return llvm::StringRef(name); // name 销毁后，StringRef 悬垂 }调用方拿到的data()指向一块已经被释放的内存，一访问就是随机崩溃。排查方法很简单：拿到StringRef的引用者，生命周期必须严格短于 string 本身。如果拿不准，直接在函数边界转成std::string，等价于把所有权明确化。
在推理框架里也有对应的坑。比如从注册表里Find返回的 kernel 工厂函数指针，如果在框架卸载动态库之后还存在缓存引用，一旦调用就直接 segment fault。解决方法是注册时留一个生命周期守卫，动态库卸载前统一清理注册表条目。这个坑在开发插件式算子时极其容易碰到。
5.2 对齐、栈溢出与碎片问题
对齐问题：
BumpPtrAllocator::Allocate(size, align)如果你传的align小于对象的实际对齐需求，可能拿到未对齐的指针。我在自定结构体上遇到过alignas(16)的float4被塞进 8 字节对齐的内存里，跑起来经常性随机出错，排查到崩溃点才知道是对齐问题。解决办法是这段内存你在Allocate前要给足对齐参数，或者在分配时统一向上取一个保守的对齐值（如 16 或 32 字节），尤其是给 SIMD 相关数据分配时。栈溢出问题：
SmallVector的模板参数不宜设置过大。模板参数就是栈上预留的容量，设成 8192 就白占 8192 个元素的空间，作为局部变量一口气开几个就爆栈了。我的经验法则是，把 N 设成实际场景 90% 以上情况覆盖的值，超过就让它走堆分支，别贪心。KV cache 碎片问题：即使 block 粒度统一，请求交错分配和释放后，自由 block 在物理上也会呈现"空洞分布"，极端情况下分配器空闲块足够，但分散导致新请求无法凑成连续的合理集合。我在这类问题上的实际解法是引入一个小的
freeList压实机制：空闲块太少时把待释放请求的 block 数量统一化（对齐到某个较大粒度），或者在分配时优先取跨度大的自由块，减少碎片累积。fragment 问题永远不会完全消失，但压到低水位以下是完全可以做到的。
5.3 排查问题的一个私藏思路
调这类底层问题，我最大的体会是：先通过监控手段确定问题类型，不要急着改代码。比如怀疑内存分配性能差，就先挂上火焰图或统计 malloc 调用次数；怀疑是碎片问题，就先打印自由块数量和分布，而不是猜。有一次我为了查推理延迟抖动，写了个小工具，每隔一段时间记录 KV cache 空闲物理块的连续性，一下子就看出了每分钟周期性碎片化的规律。
用同样的思路，排查 LLVM pass 里的内存问题时，可以打开-debug-only=allocator这类调试输出，观察块分配与释放的模式。数据先到位，定位方向基本就不会错。
最后分享一个组合技巧：我在自己的调试工具链里会对齐打印“释放块总数”和“当前推理迭代耗时”，两者一起画曲线。实测下来，每次耗时尖峰都对应着释放块数量急剧增多的时间点，进一步定位发现是极端长度请求退出时一次性释放上百个 block 拖慢了显存回收逻辑。把"释放动作"延后到下一轮调度的空闲帧再做，曲线立刻平滑下来。这类"延迟回收"的思路，本质上也和编译器 pass 里把删除操作攒到统一清理阶段再做是一致的。
如果你在做推理框架调优或者底层 kernel 优化，建议抽半天时间把SmallVector和BumpPtrAllocator的源码读一遍。这两个类加起来可能不到两千行，但里面的设计决策能直接影响你在显存分配、请求调度、kernel 自发现方面的思路。我后来改 nano-vllm 的不少决定，回头看都是这些 LLVM 机制在推理场景的平移应用。
