ops-memory：CANN Runtime 的 Tensor 内存管理

推理过程中显存的管理比 CPU 内存复杂得多。几十 GB 的模型参数、动态增长的 KV Cache、频繁创建销毁的中间激活值------全都在 NPU 的几十 GB 显存里竞争空间。

CANN 的 ops-memory 仓库负责 Tensor Buffer 的分配、回收和复用。它不是直接调 aclrtMalloc------它在上面封装了一整套内存池和生命周期管理机制。

为什么 Tensor Memory 很关键

一个 Decoder Block 的前向计算需要分配的 Tensor：

输入激活：[B, n, d]------8MB（B=1, n=4096, d=4096）
Q 投影输出：[B, n, d]------8MB
K 投影输出：[B, n, d]------8MB
V 投影输出：[B, n, d]------8MB
Attention Score：[B, H, n, n]------32MB
Context：[B, n, d]------8MB
FFN 中间激活：[B, n, 4d]------32MB
残差输出：[B, n, d]------8MB

总计约 112MB------一个 Block。40 个 Block 就是 4.5GB。这只是前向计算，训练的反向还要多一倍。

如果每次推理都重新分配这些 Tensor，分配+回收的开销可能占推理延迟的 5-10%。ops-memory 的优化是让 Tensor 复用------不同 Block 的同一层 Tensor 可以共享同一块显存。

Memory Pool 机制

ops-memory 的 Memory Pool 是一套多级分配器：

复制代码

进程级全局池（aclrtMalloc 的底层）
  └── 模型级池（当前加载的 OM 模型独占）
       └── Stream 级本地池（每个 Stream 独立加速）

全局池 在 CANN Runtime 初始化时从 NPU 显存中划走一大块。后续的 aclrtMalloc 都从池里分配------不需要跟内核交互。全局池的默认大小是 NPU 显存的 80%，剩余留给驱动和其他系统用途。

模型级池 在模型加载时创建。aclmdlLoadFromFile 加载 OM 模型时，GE 已经算出了模型需要的最大显存------模型级池的大小就是 GE 算出的最大值。模型级池从全局池中划分。

Stream 级本地池 在 Stream 创建时分配。aclrtCreateStream 创建 Stream 时可以指定 Stream 本地池的大小。推理时频繁创建销毁的临时 Tensor 从 Stream 本地池分配------不需要加锁（同一个 Stream 内的访问是单线程的）。

DMA 与 Buffer 对齐

DMA 的地址对齐要求比 CPU 严格------Buffer 的起始地址需要按 64 字节或 256 字节对齐。ops-memory 在分配 Buffer 时自动对齐。对齐方式由上层的 Tensor 格式决定：ND 格式用 64 字节对齐，NZ 格式用 256 字节对齐。

对齐浪费的显存空间通常很小（平均每个 Buffer 浪费 0.1-0.5%），但不符合对齐的 Buffer 会让 DMA 搬运退化到逐字节模式------带宽从 200GB/s 降到 10GB/s。

大模型中的显存优化

LLaMA-13B 推理中 ops-memory 的显存优化：

KV Cache 的增量分配。 Prefill 阶段结束后 KV Cache 已经满负荷。解码阶段每次追加一个 Token 的 K/V，ops-memory 不做整块重新分配------Cache 在 Prefill 时已经预分配了 max_length 的空间，解码阶段只更新偏移指针。

中间激活的复用。 LayerNorm 的 mean 和 var 在算完后立即释放。Attention Score 矩阵在第二阶段的矩阵乘结束后立即释放。ops-memory 的引用计数跟踪每个 Tensor 的使用次数------引用归零时返回内存池。

Memory Pool 的碎片管理

显存碎片是大模型推理中的隐蔽问题。KV Cache 在不同 Block 之间频繁分配和释放后，显存中会出现大量空闲但碎片化的 Block 间隙------总量够用但无法分配一个连续的 Buffer。

ops-memory 的碎片整理策略：当某个 Stream 的本地池碎片率超过 30% 时，触发碎片整理------把池中的空闲 Block 合并成更大的连续区域。整理过程需要暂停当前 Stream 的推理------约 50-100μs 的暂停。碎片整理的频率控制很关键：太频繁浪费延迟，太晚碎片化导致 OOM。

Runtime 的显存监控

ops-memory 通过 Runtime 的监控接口暴露显存使用指标：

total_memory：NPU 总显存
used_memory：当前分配量
pool_utilization：池利用率的百分比
fragmentation：碎片率

推理服务可以通过这些指标判断是否需要触发碎片整理或调整 Batch 大小。CANN 的 Runtime API 中 aclrtGetMemInfo 返回当前显存使用情况------推理框架可以定时调用这个接口做动态显存管理。

参考仓库

ops-memory 内存管理库

CANN Runtime