CANN runtime 内存池——高效显存管理策略

前言

runtime 的内存池是昇腾 NPU 显存管理的核心。分配策略、碎片处理、生命周期管理，这些细节决定了多模型推理时的显存利用率。这篇文章把 runtime 内存池的设计思路掰开讲，帮助你在模型部署时把显存吃满、用透。

内存池架构：统一管理 vs 分块管理

昇腾 CANN runtime 的内存池设计，核心解决一个矛盾：NPU 显存昂贵且容量有限，而模型推理时内存分配频繁、生命周期短。如果每次推理都向系统申请显存，性能损耗会非常可观。

runtime 采用统一管理 + 分块分配的架构。内存池在初始化时从系统申请一大块显存（称为 arena），后续所有分配都从 arena 中切分。这种设计的好处是：避免了频繁的系统调用，同时让分配策略有更大的优化空间。

具体来说，内存池维护三层数据结构：

• Arena：向系统申请的大块显存，通常为几百 MB 到几 GB

• Block：Arena 内的逻辑分块，是分配的基本单位

• Bucket：按大小分组的 Block 集合，用于快速匹配请求

  ┌─────────────────────────────┐
  │ Arena (1GB) │
  ├──────────┬──────────┬────────┤
  │ Block 0 │ Block 1 │ Block 2│ (逻辑分块)
  │ 128KB │ 256KB │ 512KB │
  └──────────┴──────────┴────────┘
  ↓ ↓ ↓
  Bucket0 Bucket1 Bucket2 (大小分组)

这种架构的关键优势在于预分配 + 复用。模型加载时，runtime 会根据模型结构预估显存需求，一次性申请足够的 Arena。推理过程中，中间结果的显存申请/释放都在 Arena 内完成，不触碰系统层。

分配策略：Best-fit vs First-fit

内存池的分配策略直接影响显存利用率和分配速度。runtime 提供两种策略，可在初始化时配置：

First-fit：速度优先

First-fit 从 Bucket 的第一个满足大小的 Block 开始分配。优点是分配速度快，缺点是容易产生外部碎片。

// First-fit 伪代码逻辑
Block* first_fit_alloc(size_t size) {
    int bucket_idx = size_to_bucket(size);
    for (auto& block : buckets[bucket_idx]) {
        if (block.size >= size && block.is_free) {
            return split_and_use(block, size);  // 可能拆分
        }
    }
    return nullptr;  // 当前 Arena 不够，需扩展
}

Best-fit：利用率优先

Best-fit 在所有满足条件的 Block 中选择最小的那个。优点是碎片更少，缺点是需要遍历所有候选 Block，分配稍慢。

// Best-fit 伪代码逻辑
Block* best_fit_alloc(size_t size) {
    Block* best = nullptr;
    for (int i = bucket_idx; i < MAX_BUCKETS; ++i) {
        for (auto& block : buckets[i]) {
            if (block.size >= size && block.is_free) {
                if (!best || block.size < best->size) {
                    best = &block;
                }
            }
        }
    }
    return best ? split_and_use(best, size) : nullptr;
}

如何选择？

场景	推荐策略	理由
单模型推理	First-fit	分配频率低，速度优先
多模型并发	Best-fit	显存紧张时利用率更关键
动态 Shape 模型	Best-fit	变长请求下碎片更可控

实际部署时，可以通过环境变量切换策略：

export ASCEND_GLOBAL_LOG_LEVEL=3  # 开启调试日志
export ASCEND_MEMPOOL_POLICY=BEST_FIT  # 或 FIRST_FIT

碎片管理：主动碎片整理触发条件

长时间运行的推理服务，内存池难免出现碎片。runtime 提供两种碎片管理机制：

被动整理：分配失败时触发

当分配请求找不到合适的 Block，但 Arena 的空闲总量足够时，runtime 会触发被动碎片整理（compaction）。整理逻辑：

1. 遍历 Arena，找到所有空闲 Block
2. 按地址排序，尝试合并相邻 Block
3. 更新 Bucket 索引
4. 重新尝试分配

被动整理的问题是：整理期间会阻塞所有分配请求，影响推理延迟。

主动整理：定时触发

runtime 支持配置定时碎片整理，在服务负载较低时提前执行。触发条件通过参数控制：

// 碎片整理配置（概念代码，非实际 API）
struct MemPoolConfig {
    int compact_interval_ms = 5000;      // 每5秒检查一次
    float fragment_threshold = 0.3f;     // 碎片率超过30%触发
    int max_block_count = 10000;         // Block数过多时触发
};

// 碎片率计算逻辑
float calc_fragment_rate() {
    size_t free_total = 0;
    size_t max_free_block = 0;
    for (auto& block : free_blocks) {
        free_total += block.size;
        max_free_block = std::max(max_free_block, block.size);
    }
    return 1.0f - (float)max_free_block / free_total;
}

实际调优建议：

• 单卡部署小模型（显存充裕）：关闭主动整理，减少开销
• 多卡部署大模型（显存紧张）：开启主动整理，碎片率阈值设为 20%
• 吞吐敏感服务：整理间隔设为推理周期的整数倍，避免打乱调度

显存泄漏排查：工具和方法

推理服务长期运行后显存持续增长，通常是显存泄漏的信号。runtime 提供几种排查工具：

1. 内存池状态快照

通过 AscendCL API 获取内存池当前状态：

#include "acl/acl.h"

void print_mempool_status() {
    size_t total_size, used_size, free_size;
    aclError ret = aclrtGetMemPoolInfo(&total_size, &used_size, &free_size);
    if (ret == ACL_SUCCESS) {
        printf("Arena总大小: %zu MB\n", total_size / 1024 / 1024);
        printf("已使用: %zu MB\n", used_size / 1024 / 1024);
        printf("空闲: %zu MB\n", free_size / 1024 / 1024);
    }
}

在推理前后分别调用，对比显存变化，可以判断是否泄漏。

2. 内存分配跟踪

设置环境变量开启分配日志：

export ASCEND_MEMPOOL_TRACE=ON
export ASCEND_MEMPOOL_TRACE_FILE=/tmp/mempool_trace.log

日志会记录每次分配/释放的调用栈、大小、时间戳。通过分析日志，可以定位泄漏点：

# 查找分配但未释放的内存
grep "alloc" /tmp/mempool_trace.log | awk '{print $4}' | sort | uniq -c | sort -rn | head -20

3. 常见泄漏场景

场景	原因	解决方案
模型卸载不完整	aclmdlUnload 未释放所有资源	先执行 aclmdlFinalize
Stream 未销毁	aclrtDestroyStream 遗漏	模型卸载前销毁所有 Stream
Event 泄漏	aclrtDestroyEvent 未调用	在推理循环外统一管理 Event

代码实操：多模型推理的内存池配置

下面是一个完整的示例：在 Ascend NPU 上同时部署三个模型，通过内存池配置实现显存隔离和最大化利用。

#include "acl/acl.h"
#include <vector>
#include <string>

class MultiModelInfer {
public:
    MultiModelInfer() : device_id_(0), stream_(nullptr) {}
    
    ~MultiModelInfer() {
        // 释放资源时必须按顺序：卸载模型 → 销毁 Stream → 反初始化 ACL
        for (auto& model : models_) {
            if (model.model_id != 0) {
                aclmdlUnload(model.model_id);
            }
        }
        if (stream_) aclrtDestroyStream(stream_);
        aclFinalize();
    }
    
    int Init(int device_id, const std::vector<std::string>& model_paths) {
        device_id_ = device_id;
        
        // 初始化 ACL
        aclError ret = aclInit(nullptr);
        if (ret != ACL_SUCCESS) {
            printf("aclInit failed: %d\n", ret);
            return -1;
        }
        
        ret = aclrtSetDevice(device_id_);
        if (ret != ACL_SUCCESS) {
            printf("aclrtSetDevice failed: %d\n", ret);
            return -1;
        }
        
        // 创建 Stream，所有模型共享一个 Stream 以减少资源占用
        ret = aclrtCreateStream(&stream_);
        if (ret != ACL_SUCCESS) {
            printf("aclrtCreateStream failed: %d\n", ret);
            return -1;
        }
        
        // 加载模型，内存池会自动分配
        for (const auto& path : model_paths) {
            ModelInfo info;
            info.path = path;
            ret = aclmdlLoadFromFile(path.c_str(), &info.model_id);
            if (ret != ACL_SUCCESS) {
                printf("Load model %s failed: %d\n", path.c_str(), ret);
                continue;
            }
            
            // 获取模型输入输出描述，用于后续推理
            info.desc = aclmdlCreateDesc();
            aclmdlGetDesc(info.desc, info.model_id);
            
            models_.push_back(info);
            printf("Loaded model: %s, ID: %u\n", path.c_str(), info.model_id);
        }
        
        // 打印内存池状态
        PrintMemPoolStatus();
        return 0;
    }
    
    void Infer(int model_idx, void* input_data, size_t input_size) {
        if (model_idx >= models_.size()) return;
        
        auto& model = models_[model_idx];
        
        // 创建输入 Dataset
        aclmdlDataset* input_dataset = aclmdlCreateDataset();
        aclDataBuffer* input_buffer = aclCreateDataBuffer(input_data, input_size);
        aclmdlAddDatasetBuffer(input_dataset, input_buffer);
        
        // 创建输出 Dataset（假设输出大小已知）
        size_t output_size = 1024 * 1024;  // 1MB 输出缓冲区
        void* output_data = nullptr;
        aclrtMalloc(&output_data, output_size, ACL_MEM_MALLOC_NORMAL_ONLY);
        aclDataBuffer* output_buffer = aclCreateDataBuffer(output_data, output_size);
        aclmdlDataset* output_dataset = aclmdlCreateDataset();
        aclmdlAddDatasetBuffer(output_dataset, output_buffer);
        
        // 执行推理，显存由内存池自动管理
        aclError ret = aclmdlExecute(model.model_id, input_dataset, output_dataset);
        if (ret == ACL_SUCCESS) {
            printf("Model %d inference success\n", model_idx);
        }
        
        // 释放 Dataset 和 Buffer，显存归还内存池
        aclDestroyDataBuffer(input_buffer);
        aclDestroyDataBuffer(output_buffer);
        aclmdlDestroyDataset(input_dataset);
        aclmdlDestroyDataset(output_dataset);
        aclrtFree(output_data);  // 显式释放推理输出显存
    }
    
    void PrintMemPoolStatus() {
        size_t total, used, free;
        aclrtGetMemPoolInfo(&total, &used, &free);
        printf("=== 内存池状态 ===\n");
        printf("Arena 总大小: %zu MB\n", total / 1024 / 1024);
        printf("已使用: %zu MB (%.1f%%)\n", used / 1024 / 1024, 
               100.0 * used / total);
        printf("空闲: %zu MB\n", free / 1024 / 1024);
    }

private:
    struct ModelInfo {
        std::string path;
        uint32_t model_id = 0;
        aclmdlDesc* desc = nullptr;
    };
    
    int device_id_;
    aclrtStream stream_;
    std::vector<ModelInfo> models_;
};

int main() {
    MultiModelInfer infer;
    
    // 加载三个模型，内存池按需扩展
    std::vector<std::string> models = {
        "/models/resnet50.om",
        "/models/bert_base.om",
        "/models/yolov5.om"
    };
    
    if (infer.Init(0, models) == 0) {
        // 模拟推理
        float input_data[224 * 224 * 3];
        infer.Infer(0, input_data, sizeof(input_data));
        
        // 再次打印内存池状态，观察显存变化
        infer.PrintMemPoolStatus();
    }
    
    return 0;
}

代码要点说明：

1. 共享 Stream：三个模型使用同一个 Stream，减少 ACL 资源占用，同时让内存池统一管理所有模型的显存
2. Dataset 生命周期：推理完成后立即销毁 Dataset 和 DataBuffer，让显存归还内存池复用
3. 内存池状态监控：推理前后打印内存池状态，可以观察 Arena 扩展和碎片情况
4. 资源释放顺序：必须先卸载模型、再销毁 Stream、最后 aclFinalize，否则会导致显存泄漏

踩坑实录：几个常见问题

问题 1：Arena 扩展失败

现象：加载大模型时报错 ACL_ERROR_MEM_ALLOC_FAIL，但 npu-smi info 显示显存充足。

原因：runtime 默认 Arena 大小有限制，单次申请超过阈值会失败。

解决：通过环境变量调整 Arena 上限：

export ASCEND_MEMORY_POOL_MAX_SIZE=8589934592  # 8GB，单位字节

问题 2：多模型间显存抢占

现象：模型 A 推理时，模型 B 的显存被覆盖，导致结果错误。

原因：默认情况下，所有模型共享同一个 Arena，没有显存隔离。

解决：为每个模型创建独立的 Context（不同 Context 使用独立内存池）：

aclrtContext ctx1, ctx2;
aclrtCreateContext(&ctx1, device_id);
aclrtCreateContext(&ctx2, device_id);

// 模型 A 在 ctx1 中加载
aclrtSetCurrentContext(ctx1);
aclmdlLoadFromFile(path_a, &model_a);

// 模型 B 在 ctx2 中加载
aclrtSetCurrentContext(ctx2);
aclmdlLoadFromFile(path_b, &model_b);

问题 3：碎片整理导致延迟抖动

现象：推理服务周期性出现延迟尖刺，监控发现与碎片整理时间点吻合。

原因：被动碎片整理阻塞了所有分配请求。

解决：

• 降低碎片整理触发阈值，在碎片率较低时提前整理
• 使用定时主动整理，避开业务高峰期

export ASCEND_MEMPOOL_COMPACT_THRESHOLD=0.2  # 碎片率20%触发
export ASCEND_MEMPOOL_COMPACT_INTERVAL=10000  # 每10秒检查一次

小结

runtime 内存池的设计，本质是在分配速度 和利用率之间找平衡。First-fit 偏向速度，Best-fit 偏向利用率。碎片管理这块，被动整理简单但有延迟风险，主动整理需要配合业务节奏。

实际部署时，建议先跑一轮基准测试：单模型用 First-fit 配默认 Arena 即可；多模型并发或显存紧张场景，切到 Best-fit，调整碎片整理参数，把显存吃满。

如果你正在做多模型推理部署，可以先从文中代码示例跑起，观察内存池状态变化。遇到问题用 ASCEND_MEMPOOL_TRACE 开日志，定位泄漏点。

runtime 仓库地址：https://atomgit.com/cann/runtime