本文基于昇腾CANN和昇腾NPU，围绕 Memory Pool 显存池技术展开。推理过程中 Tensor 的生灭频繁得像呼吸——一次推理几十个中间 Tensor，每个都要申请和释放显存。Memory Pool 把这件事从每次推理的循环里拿出来。

NPU 显存申请不是 CPU 上的 malloc——调用 acl.rt.malloc 要经过驱动层、内存管理单元、页表映射，延迟在 0.5ms 到 2ms 之间。一次推理如果有 20 个中间 Tensor，malloc+free 的累积延迟就达到 10-40ms。而真正的矩阵运算可能只有 15ms。等于 40-70% 的时间花在管理显存上。

Pool 是如何工作的

Memory Pool 的逻辑很直白：Runtime 初始化时申请一大块显存（比如 4GB），动态分配 的逻辑全在 Pool 内部——后续所有中间 Tensor 的分配都在这一块内部做，不再走驱动层。Pool 内部维护一个空闲块链表，分配时找够大的空闲块切下来，释放时把块加回链表并合并相邻空闲块。

显存管理 的最终形态就是这个：CANN Runtime 的 Pool 还支持两个关键特性：碎片整理（compact）和按需扩容。碎片多到找不到大块时，compact 把已分配块"搬家"到连续区域，合并出大空闲块。Pool 用光时，再向驱动层申请一块新的扩大池子。

# CANN Runtime Memory Pool 的简化实现

class CANNMemoryPool:
    def __init__(self, pool_size=4 * 1024**3):
        self.ptr = acl.rt.malloc(pool_size,
                                 acl.ACL_MEM_MALLOC_HUGE_FIRST)
        self.total = pool_size
        self.used = 0
        self.free_blocks = [(0, pool_size)]
        self.allocated = {}

    def allocate(self, name, size):
        # 找够大的空闲块
        for i, (off, sz) in enumerate(self.free_blocks):
            if sz >= size:
                self.allocated[name] = (off, size)
                self.used += size
                if sz > size:
                    self.free_blocks[i] = (off + size, sz - size)
                else:
                    self.free_blocks.pop(i)
                return off
        # Pool 不够——尝试扩容
        new_pool = acl.rt.malloc(size)
        self.allocated[name] = ("external", size)
        return new_pool

    def deallocate(self, name):
        off, size = self.allocated.pop(name)
        if off != "external":
            self.used -= size
            self.free_blocks.append((off, size))
            self.free_blocks.sort()
            self._merge_adjacent()

    def _merge_adjacent(self):
        # 合并相邻空闲块——防止碎片化
        merged = [self.free_blocks[0]]
        for off, sz in self.free_blocks[1:]:
            prev_off, prev_sz = merged[-1]
            if prev_off + prev_sz == off:
                merged[-1] = (prev_off, prev_sz + sz)
            else:
                merged.append((off, sz))
        self.free_blocks = merged

大模型推理中的显存场景

LLaMA-7B 推理的显存分布：模型权重 14GB（固定），KV Cache 按需增长（128K 上下文时约 8GB），中间 Tensor Buffer 峰值约 3.5GB。如果不做 Pool，中间 Tensor 的 malloc/free 额外吃掉约 500MB 的碎片空间和 15ms 的延迟。

Pool 把 15ms 降到 1ms 以内，500MB 碎片降到零。省下的显存可以直接加 Batch Size。

推荐继续学习

Pool 解决了分配延迟，碎片问题要看 Memory Fragmentation。

参考仓库

Runtime 运行时

GE 图引擎

CANN Samples

CANN 学习中心