汤不热吧在移动端进行AI推理时，显存（通常是共享内存DRAM或专用的VRAM）往往是瓶颈。对于参数量较大的模型（如轻量级LLM或大型CV模型），其激活值和中间计算结果可能会瞬间占用数百兆甚至超过1GB的内存。本文将聚焦于推理引擎中最关键的优化技术之一：显存复用（Tensor Reuse），并讲解如何设计内存池来有效控制峰值内存占用。
1. 为什么需要显存复用？
传统的模型推理过程，如果不进行优化，可能会为每个中间张量（Tensor）都独立分配内存。例如，一个包含100个计算层的模型，即使每层只消耗10MB的内存，累计下来也需要1GB的峰值内存。但实际上，某个张量 $T_A$ 完成计算后，如果后续没有任何操作依赖它，那么它所占用的底层显存块就可以立即释放或复用给下一个张量 $T_B$。
显存复用的核心思想是：分析模型的计算图，找出张量的生命周期，确保在任何时间点，内存池中仅保留当前和未来必要的操作所需的张量内存。
2. 内存池的设计原理
高效的内存池需要解决两个核心问题：内存分配速度和内存碎片化。
2.1 内存分配策略
端侧推理框架如 MNN、NCNN 和 TFLite 通常使用定制化的内存池，而不是依赖操作系统低效的 malloc/free（特别是在GPU或DSP上）。
- 预分配大块内存： 启动时，框架向系统申请一大块连续的内存作为内存池的底层存储。
- Best-Fit/First-Fit 查找： 当需要分配大小为 $S$ 的张量时，内存池首先查找已释放的内存块列表中是否存在大小 $\ge S$ 的块。使用 First-Fit (找到第一个合适的就分配) 或 Best-Fit (找到最接近 $S$ 的块) 策略，以减少外部碎片。
- 对齐： 确保分配的内存块地址满足硬件（如GPU纹理、CPU缓存线）的对齐要求，通常是4字节、16字节或32字节对齐。
2.2 核心：基于图分析的生命周期管理
在模型加载阶段，推理引擎会遍历计算图，并进行张量生命周期分析（Liveness Analysis）：
- 定义： 对于图中的每个操作 $O_i$，其输入张量 $T_{in}$ 和输出张量 $T_{out}$。
- 生命周期结束点： 张量 $T_{in}$ 的生命周期在其最后一个依赖 $O_j$ 使用完成后即结束。此时，其底层Buffer可以归还给内存池。
- Buffer ID 映射： 最终，计算图中的每个张量不再直接拥有内存，而是持有一个指向内存池中 Buffer ID 的引用。
3. 代码实操：模拟内存池的张量复用
我们使用 Python 模拟一个简化的内存池，展示如何通过复用机制来限制峰值内存。
假设我们有一个模型计算图，它需要的总内存为400MB，但我们设备只有300MB内存限制。
模型依赖图: $T_A \to OpB \to T_B \to OpC \to T_C \to OpD \to T_D$
| 张量 | 大小 (MB) | 生命周期结束点 |
|---|---|---|
| $T_A$ | 100 MB | OpB 完成后 |
| $T_B$ | 150 MB | OpC 完成后 |
| $T_C$ | 50 MB | OpD 完成后 |
| $T_D$ | 100 MB | 整个推理结束 |
预期峰值内存: (T_A + T_B) = 250 MB。 $T_C$ 在 $T_A$ 释放后分配。
import sys
# 辅助函数：MB转换为Byte
MB_TO_BYTE = 1024 * 1024
class SimpleInferenceMemoryPool:
def __init__(self, capacity_mb):
self.capacity = capacity_mb * MB_TO_BYTE
# 存储已释放的块：{size_bytes: [buffer_id, ...]}
self.free_buffers = {}
# 存储当前正在使用的块：{buffer_id: size_bytes}
self.allocated_buffers = {}
self.next_id = 1
self.current_reserved_usage = 0
def allocate(self, size_mb, tensor_name):
size_bytes = int(size_mb * MB_TO_BYTE)
# 1. 尝试复用精确大小的块
if size_bytes in self.free_buffers and self.free_buffers[size_bytes]:
buffer_id = self.free_buffers[size_bytes].pop(0)
self.allocated_buffers[buffer_id] = size_bytes
print(f"[Pool] R-Use: {tensor_name} (ID {buffer_id}), Size: {size_mb:.2f} MB")
return buffer_id
# 2. 如果无法复用，检查容量并分配新块
if self.current_reserved_usage + size_bytes > self.capacity:
raise MemoryError(
f"[OOM] Out of Pool Capacity! Required: {size_mb:.2f} MB, Current Peak: {self.get_peak_usage():.2f} MB, Cap: {self.capacity / MB_TO_BYTE:.2f} MB"
)
buffer_id = self.next_id
self.next_id += 1
self.allocated_buffers[buffer_id] = size_bytes
self.current_reserved_usage += size_bytes
print(f"[Pool] New: {tensor_name} (ID {buffer_id}), Size: {size_mb:.2f} MB. Peak Usage: {self.get_peak_usage():.2f} MB")
return buffer_id
def release(self, buffer_id, tensor_name):
if buffer_id not in self.allocated_buffers:
return
size = self.allocated_buffers.pop(buffer_id)
if size not in self.free_buffers:
self.free_buffers[size] = []
self.free_buffers[size].append(buffer_id)
# 注意：内存块仍然被池子保留，以便复用，因此 current_reserved_usage 不变
print(f"[Pool] Free: {tensor_name} (ID {buffer_id}) back to pool.")
def get_peak_usage(self):
return self.current_reserved_usage / MB_TO_BYTE
# --- 模拟推理过程 ---
POOL_CAPACITY_MB = 300 # 限制在 300MB
memory_manager = SimpleInferenceMemoryPool(POOL_CAPACITY_MB)
print(f"=== 开始推理 (限制容量: {POOL_CAPACITY_MB} MB) ===")
# 步骤 1: OpA -> T_A (100MB), T_A 仍在使用
t_a_id = memory_manager.allocate(100, "T_A")
# 步骤 2: OpB (输入 T_A) -> T_B (150MB)
t_b_id = memory_manager.allocate(150, "T_B")
# OpB 计算完成。T_A的生命周期结束。
memory_manager.release(t_a_id, "T_A")
# 步骤 3: OpC (输入 T_B) -> T_C (50MB)
t_c_id = memory_manager.allocate(50, "T_C") # 此时总占用 100+150+50 = 300MB
# OpC 计算完成。T_B的生命周期结束。
memory_manager.release(t_b_id, "T_B")
# 步骤 4: OpD (输入 T_C) -> T_D (100MB)
# T_D 需要 100MB。内存池中 T_A (100MB) 的块已被释放，可以复用！
t_d_id = memory_manager.allocate(100, "T_D")
# 如果没有复用，分配 T_D 会导致总峰值达到 100+150+50+100 = 400MB，从而 OOM。
# 由于复用了 T_A 的内存，峰值维持在 300MB。
print("\n=== 推理完成 ===")
print(f"最终峰值内存占用: {memory_manager.get_peak_usage():.2f} MB")
4. 实践中的考虑
- 动态复用 vs 静态复用： 许多端侧框架在模型加载时就静态计算好所有张量的复用方案，生成一个内存分配表。这避免了运行时分配的开销，是最优化的方式。
- 异构计算的挑战： 如果模型使用了 CPU、GPU、DSP 混合计算，每个设备可能拥有独立的内存池，需要处理不同内存类型（如 CPU 侧的 malloc 和 GPU 侧的 OpenCL/Vulkan 内存）之间的同步和零拷贝问题。
- 碎片化处理： 简单的精确匹配可能导致大量外部碎片。更复杂的内存池（如 Buddy System 或 Slab Allocator）在实际的端侧框架中被用于进一步优化内存使用和分配速度。