4. 内存分配机制
核心原理：为什么需要内存池与分配器
大语言模型推理时，张量的大小往往要到运行时才能确定。例如，KV Cache 的显存占用随请求数量、序列长度动态变化；中间激活值随 batch size 和输入长度波动。如果每次需要显存都直接调用 cudaMalloc，会带来两个严重问题：一是 cudaMalloc 是同步调用，会阻塞计算流，造成数百微秒甚至毫秒级的延迟；二是频繁分配与释放会导致显存碎片化，使可用显存逐渐减少，最终触发 OOM。
TensorRT-LLM 借鉴了传统内存管理思路，在显存之上建立了一层 内存池（Memory Pool）。推理引擎启动时，根据配置一次性向 CUDA 申请一大块连续显存，之后所有内部张量都从这块池中按需切分。池的管理者称为 分配器（Allocator），它负责将池中的空闲块分配给请求方，并在使用完毕后回收。分配器支持两种模式：直接模式（Direct） 与 池化模式（Pool）。直接模式每次仍调用 cudaMalloc，适合调试；池化模式则完全在用户态管理显存，是生产环境的默认选择。
关键点：内存池的粒度不是单个字节，而是“块（Block）”。分配器维护一个空闲块链表，分配时采用最佳适配（Best Fit）或首次适配（First Fit）策略，释放时尝试与相邻空闲块合并，从而抑制碎片化。
内存复用策略：从块分配到生命周期管理
内存池的核心价值在于 复用。TensorRT-LLM 的复用策略分为三个层次：
- 块级复用：同一大小的张量释放后，其块被放回空闲链表，后续相同大小的请求直接命中。例如，不同层的激活值大小相同，可以共享同一块显存。
- 流级复用：不同 CUDA 流上的张量，如果生命周期不重叠，也可以复用同一块物理显存。分配器通过记录每个块的“使用中”标志来保证安全。
- KV Cache 分页复用：KV Cache 按页（Block）管理，每页大小固定（如 64 的 KV）。请求结束后，其占用的页归还到全局页池，供下一个请求使用。这与操作系统虚拟内存的分页机制类似。
分配器内部维护一张 内存块表，每条记录包含：起始地址、块大小、是否空闲、所属流、最后使用时间。当收到分配请求时，分配器按以下流程操作：
- 根据请求大小，在空闲链表中查找满足条件的块（优先选择大小最接近的块，减少内部碎片）。
- 如果找到的块远大于请求大小，则进行 分裂：将块拆分为“已分配部分”和“新的空闲部分”。
- 如果找不到合适的空闲块，则从池末尾的“预留区”扩展新块；若池已满，则触发 池扩展（再次调用
cudaMalloc追加显存）。 - 释放时，将块标记为空闲，并检查其前后相邻块是否也为空闲；若是，则合并为一个更大的块。
工程实践：在构建引擎时，可通过 --mem_pool_size 参数预留 KV Cache 池的大小。预留过小会导致运行时频繁扩展池，产生额外延迟；预留过大会浪费显存。建议根据最大并发请求数和最大序列长度估算。
案例：多请求并发下的内存复用
假设显存池大小为 8GB，其中权重占 4GB，KV Cache 预留 3GB，激活区预留 1GB。两个请求同时到达：
| 时间点 | 请求 A（序列长 128） | 请求 B（序列长 256） | 分配器动作 |
|---|---|---|---|
| T1 | 申请 KV 页：2 页 | — | 从 KV 池分配 2 页（每页 64 ） |
| T2 | — | 申请 KV 页：4 页 | 从 KV 池分配 4 页，与 A 的页不重叠 |
| T3 | 请求结束，释放 2 页 | 继续生成 | 2 页标记为空闲，加入页池 |
| T4 | — | 序列增长，需追加 2 页 | 直接复用 T3 释放的 2 页，无需 cudaMalloc |
这个案例展示了 KV Cache 分页复用的核心价值：请求 A 释放的显存可以立即被请求 B 使用，整个过程不涉及 CUDA 显存分配，延迟仅为链表操作（微秒级）。相比之下，若每次请求都调用 cudaMalloc，T4 时刻将产生约 50~200 微秒的阻塞延迟。
误区与复盘：内存池使用中的常见问题
误区一：认为内存池越大越好。池大小受限于 GPU 总显存。如果池设置过大，留给 CUDA 上下文和其他库（如 cuBLAS）的显存不足，会导致引擎初始化失败。正确做法是：先测量模型权重大小，再根据最大并发数估算 KV Cache 需求，最后留出 10%~20% 余量。
误区二：忽略碎片化问题。即使有内存池，如果请求大小频繁变化（如 100MB、150MB、80MB 交替），仍会产生外部碎片。TensorRT-LLM 提供了 --fragment_size 参数来设置最小分配粒度（默认 256MB），小于该粒度的请求会向上取整，从而减少碎片。但粒度过大会增加内部碎片，需要权衡。
复盘：某团队在生产环境部署时，发现推理延迟偶尔飙升。排查后发现，每次请求的 KV Cache 大小不同，导致池中产生大量无法合并的小空闲块。最终通过设置 --kv_cache_page_size=128（增大页粒度）和 --fragment_size=512MB，将碎片率从 23% 降至 4%，P99 延迟降低 31%。
另一个常见问题是 内存池与 CUDA Graph 的交互。CUDA Graph 捕获时要求显存地址固定，如果内存池在 Graph 捕获期间发生扩展，会导致捕获失败。TensorRT-LLM 的解决方案是：在 Graph 捕获前预分配足够的池空间，并将分配器设置为“禁止扩展”模式。开发者应确保预热阶段完成所有池扩展，再进入 Graph 捕获流程。
最后需要强调的是，内存池并非万能。对于极小的张量（如几十字节的标量），池化管理的开销可能超过直接分配的代价。TensorRT-LLM 内部对小于 1KB 的请求走“快速路径”，直接从线程本地缓存分配，避免锁竞争。理解这些细节，有助于在自定义插件或扩展中正确使用内存分配 API，避免引入不必要的性能损耗。