3、显存复用策略（上）：对象池模式、预分配与懒分配、Tensor复用技术、显存生命周期管理

显存复用，说白了就是「省着用」。

做吞吐压测的时候，显存就是你的命根子。我见过太多团队，模型推理性能上不去，一查发现显存被反复申请释放，碎片化严重，GPU利用率低得可怜。今天咱们就聊聊怎么把这部分优化做到极致。

3.1 对象池模式（Object Pool）

对象池这个概念，其实不新鲜。但在GPU编程里，它的价值被严重低估了。

你想想看，每次cudaMalloc和cudaFree，都是跟驱动打交道。这玩意儿开销不小。我在项目中遇到过，一个简单的推理服务，每秒处理几百次请求，每次都要分配临时缓冲区。结果呢？光分配释放就占了10%以上的时间。

对象池的实现，我一般这么搞：

class CudaMemoryPool {
private:
    std::vector<void*> free_blocks;
    std::mutex mtx;
    size_t block_size;
    size_t pool_size;
    
public:
    CudaMemoryPool(size_t size, size_t count) 
        : block_size(size), pool_size(count) {
        for (size_t i = 0; i < count; ++i) {
            void* ptr;
            cudaMalloc(&ptr, size);
            free_blocks.push_back(ptr);
        }
    }
    
    void* acquire() {
        std::lock_guard<std::mutex> lock(mtx);
        if (free_blocks.empty()) {
            // 池子空了，可以动态扩容
            void* ptr;
            cudaMalloc(&ptr, block_size);
            return ptr;
        }
        void* ptr = free_blocks.back();
        free_blocks.pop_back();
        return ptr;
    }
    
    void release(void* ptr) {
        std::lock_guard<std::mutex> lock(mtx);
        free_blocks.push_back(ptr);
    }
    
    ~CudaMemoryPool() {
        for (auto ptr : free_blocks) {
            cudaFree(ptr);
        }
    }
};

3.2 预分配与懒分配

这两个策略，其实是两种哲学。

预分配：启动时就把所有可能用到的显存都申请好。好处是运行时零延迟分配，坏处是——你得多占着茅坑不拉屎。

懒分配：用到了才分配。好处是省显存，坏处是第一次用的时候会卡一下。

我个人习惯是：核心路径用预分配，边缘路径用懒分配。

举个例子，做吞吐压测时，模型权重、输入输出缓冲区这些核心资源，必须预分配。而一些调试用的临时变量、日志缓冲区，懒分配就行。

3.3 Tensor复用技术（in-place操作）

Tensor复用，说白了就是「用完的显存别急着扔，看看能不能给下一个用」。

in-place操作是最直接的复用方式。比如ReLU、Dropout这些操作，输出可以直接覆盖输入。PyTorch里加个下划线就行：

# 普通操作
x = torch.relu(x)  # 分配了新显存

# in-place操作
torch.relu_(x)     # 直接覆盖，零额外显存

但要注意，不是所有操作都能in-place。比如BatchNorm，它需要保存输入的均值和方差用于反向传播。你要是直接覆盖了，反向传播就废了。

我做推理优化时，经常用这个技巧。推理不需要反向传播，所以几乎所有激活函数都可以in-place。显存占用直接砍半。

3.4 显存生命周期管理

这是最容易被忽视的部分。很多人只关注「怎么分配」，不关注「什么时候释放」。

显存生命周期管理，核心就一句话：谁分配，谁释放；谁借用，谁归还。

我画了一张图，帮你理解显存的生命周期流转：

实际项目中，我见过最典型的生命周期问题：

• 悬空指针：一个模块释放了显存，另一个模块还在用
• 内存泄漏：申请了显存，但忘记释放，或者异常路径没处理
• 生命周期不匹配：两个tensor生命周期重叠，但被错误地复用了同一块显存

举个例子，C++里我经常这么写：

class CudaTensor {
    float* data_;
    size_t size_;
    bool owned_;
    
public:
    // 构造函数：分配显存
    CudaTensor(size_t size) : size_(size), owned_(true) {
        cudaMalloc(&data_, size * sizeof(float));
    }
    
    // 构造函数：借用显存（不拥有所有权）
    CudaTensor(float* data, size_t size) : data_(data), size_(size), owned_(false) {}
    
    // 析构函数：只有owned_为true时才释放
    ~CudaTensor() {
        if (owned_ && data_) {
            cudaFree(data_);
        }
    }
    
    // 禁止拷贝，允许移动
    CudaTensor(const CudaTensor&) = delete;
    CudaTensor& operator=(const CudaTensor&) = delete;
    
    CudaTensor(CudaTensor&& other) noexcept 
        : data_(other.data_), size_(other.size_), owned_(other.owned_) {
        other.data_ = nullptr;
        other.owned_ = false;
    }
};

好了，显存复用策略的上半部分就聊到这儿。对象池、预分配与懒分配、Tensor复用、生命周期管理，这四个点你吃透了，吞吐压测的显存瓶颈基本就能解决大半。

想要更多资料，联系可可