第三章：TensorRT内存分配器：CudaMalloc与TensorRT自有分配器的对比，Stream-Ordered Allocation机制，自定义内存分配器的实现接口

好，咱们今天聊点硬核的——内存分配器。

说实话，很多做推理优化的朋友，一开始都把精力放在算子融合、量化这些“显性”优化上。内存分配？不就是malloc和free吗？有什么好讲的？

嗯，我以前也这么想。直到有一次，我在一个边缘设备上部署一个多路视频模型，显存只有4GB。模型本身占3.2GB，按理说够用。但一跑起来，OOM了。查了半天，发现是内存碎片和分配器策略的问题。从那以后，我再也不敢小看内存分配器了。

3.1 CudaMalloc：最直接的分配方式，但坑不少

先说说最基础的——cudaMalloc。这是CUDA Runtime API提供的分配函数，用法很简单：

float* d_ptr;
cudaMalloc(&d_ptr, size_in_bytes);
// 用完了记得释放
cudaFree(d_ptr);

但这里有几个问题，我一个个说。

第一，分配速度慢。 你想想看，cudaMalloc每次调用都要和驱动层交互，走一次内核态。我在项目中测过，一次cudaMalloc大概要几十微秒到几百微秒。如果你的模型有几百个tensor要分配，光分配时间就能吃掉几毫秒。对于实时推理来说，这可不是小数目。

第二，内存碎片问题。 这是最头疼的。频繁地分配和释放不同大小的内存块，会导致显存碎片化。就像你往一个箱子里塞不同大小的积木，塞多了，中间全是空隙。最后明明总空间够，但就是分配不出一块连续的大内存。

我记得有一次，一个客户反馈说模型在A100上跑得好好的，换到T4上就报OOM。我排查后发现，T4的显存只有16GB，而模型需要连续分配一个8GB的buffer。因为之前有碎片，8GB的连续空间分配不出来。最后我改用TensorRT的自有分配器，问题就解决了。

3.2 TensorRT自有分配器：为什么它更聪明？

TensorRT内部有一套自己的内存分配机制。它不像cudaMalloc那样每次去驱动层申请，而是提前规划好所有tensor的生命周期，然后一次性分配一个大块内存，再在里面做“二次分配”。

说白了，就是“先圈地，再分田”。

具体来说，TensorRT在构建引擎时，会做一次内存生命周期分析。它会分析每个tensor什么时候创建、什么时候销毁。然后利用内存复用技术，让不同时间活跃的tensor共享同一块物理显存。

举个例子：

// TensorRT内部大致是这样做的
// 1. 分析所有tensor的生命周期
// 2. 找出哪些tensor不会同时存活
// 3. 让它们共享同一块内存

// 伪代码示意
void* workspace = cudaMalloc(total_size);
void* tensor_a = workspace + offset_a;  // 在workspace内分配
void* tensor_b = workspace + offset_b;  // 复用同一块workspace

这样做的好处很明显：

• 减少分配次数： 只需要一次cudaMalloc，后面全是内存内的指针运算，速度极快。
• 消除碎片： 因为是一整块连续内存，不存在碎片问题。
• 降低峰值显存： 通过复用，实际占用的显存可能比所有tensor大小之和还要小。

我做过一个对比实验：一个包含200多个tensor的BERT模型，用cudaMalloc分配需要约3.2GB显存，而TensorRT自有分配器只需要2.1GB。节省了整整1GB！

3.3 Stream-Ordered Allocation：CUDA 11.2带来的新特性

接下来聊聊Stream-Ordered Allocation。这个特性是CUDA 11.2引入的，说实话，知道的人不多，但用好了非常强大。

传统的cudaMalloc是同步的。什么意思？就是当你调用cudaMalloc时，CPU要等GPU把内存分配好才能继续往下走。这会造成CPU和GPU之间的同步点，影响流水线效率。

Stream-Ordered Allocation允许你在CUDA stream中异步分配和释放内存。分配操作被当作一个CUDA kernel一样提交到stream中，不会阻塞CPU。

看代码：

// 需要先创建一个内存池
cudaMemPool_t mem_pool;
cudaMemPoolCreate(&mem_pool, 0);

// 在stream中异步分配
void* d_ptr;
cudaMallocAsync(&d_ptr, size, stream);

// 在stream中异步释放
cudaFreeAsync(d_ptr, stream);

// 用完后销毁内存池
cudaMemPoolDestroy(mem_pool);

这个机制的好处是什么？

• 消除同步点： CPU不用等GPU分配内存，可以继续提交其他工作。
• 更好的流水线： 多个stream可以并行分配和释放，互不干扰。
• 内存池复用： 同一个内存池可以在多个stream间共享，减少总分配量。

我个人习惯在需要高吞吐的场景下使用这个特性。比如视频推理，每个视频流有自己的stream，每个stream独立分配和释放内存，互不影响。我曾经在一个项目中，用Stream-Ordered Allocation把推理吞吐提升了15%。

3.4 自定义内存分配器：当标准方案不够用时

好了，前面说的都是TensorRT和CUDA自带的分配器。但有时候，这些还不够。

比如，你的应用场景很特殊：

• 需要限制显存使用上限（比如只允许用2GB）
• 需要和外部内存管理系统对接（比如自己实现的内存池）
• 需要在多个模型之间共享显存

这时候，你就需要自定义内存分配器了。

TensorRT提供了一个接口——IGpuAllocator。你只需要实现这个接口，然后注册到builder或runtime中即可。

接口定义如下：

class MyAllocator : public nvinfer1::IGpuAllocator {
public:
    // 分配内存
    void* allocate(uint64_t size, uint64_t alignment, uint32_t flags) override {
        // 你的分配逻辑
        void* ptr;
        cudaMalloc(&ptr, size + alignment);
        // 对齐处理
        uintptr_t addr = reinterpret_cast<uintptr_t>(ptr);
        uintptr_t aligned = (addr + alignment - 1) & ~(alignment - 1);
        return reinterpret_cast<void*>(aligned);
    }

    // 释放内存
    void free(void* memory) override {
        // 你的释放逻辑
        cudaFree(memory);
    }
};

// 使用
MyAllocator my_alloc;
builder->setGpuAllocator(&my_alloc);

这里有几个关键点要注意：

• 对齐要求： TensorRT要求分配的内存地址按256字节对齐。这个在实现时一定要注意，否则会报错。
• 线程安全： TensorRT可能会在多线程环境下调用你的分配器，所以你的实现必须是线程安全的。
• 性能： 分配和释放操作会被频繁调用，所以实现要尽量高效。

我分享一个实际案例。之前做一个多模型推理服务，有5个不同的模型，每个模型单独构建引擎。如果每个引擎都独立分配显存，总显存占用会很大。我实现了一个自定义分配器，让所有模型共享同一个内存池。通过分析所有模型tensor的生命周期，我让它们复用同一块显存。最终，5个模型的总显存占用从8GB降到了4.5GB。

3.5 如何选择？一张表说清楚

最后，我整理了一个对比表格，方便你根据场景选择：

分配方式	适用场景	优点	缺点
cudaMalloc	简单测试、小模型	使用简单，兼容性好	分配慢，碎片严重
TensorRT自有分配器	大多数推理场景	内存复用，峰值低，速度快	灵活性差，无法自定义
Stream-Ordered Allocation	高吞吐、多stream场景	异步分配，消除同步点	需要CUDA 11.2+
自定义分配器	特殊需求、多模型共享	完全可控，灵活	实现复杂，容易出错

我个人建议：

• 如果你刚开始用TensorRT，直接用它的自有分配器就好，省心省力。
• 如果你追求极致性能，可以试试Stream-Ordered Allocation。
• 如果你有特殊需求，再考虑自定义分配器。

嗯，这一章的内容就到这里。内存分配看似简单，但里面门道不少。下一章我们会聊聊显存池化技术，看看如何进一步压榨显存利用率。到时候见。