CuPy 内存管理完全指南：内存池、显存限制与流有序分配-育师

CuPy 内存管理完全指南：内存池、显存限制与流有序分配

【免费下载链接】cupyNumPy & SciPy for GPU项目地址: https://gitcode.com/GitHub_Trending/cu/cupy

CuPy 默认采用**内存池（memory pool）**机制进行 GPU 显存与固定内存（pinned memory）的分配与复用，以显著降低内存分配开销和 CPU/GPU 同步带来的性能损耗。本文以官方用户指南 docs/source/user_guide/memory.rst 为骨架，结合仓库源码（cupy/cuda/memory.pyx、cupy/init.py、cupyx/_pinned_array.py）深入剖析两大内存池的工作原理、统计接口、显存硬上限、自定义分配器、流有序内存分配器（Stream Ordered Memory Allocator）以及实验性的统一内存编程（UMP）支持。读完本文，你将掌握如何监控、限制、替换乃至彻底禁用 CuPy 的内存池，并能针对多 GPU、多进程共存等场景做出正确的内存策略选型。

一、为什么需要内存池：默认分配策略与性能收益

CuPy 默认使用内存池来完成所有内存分配。内存池之所以能显著提升性能，核心在于两点：

减少分配开销：cudaMalloc这类底层分配调用本身成本高昂，内存池通过缓存已释放的内存块，使后续同尺寸分配直接命中复用，绕过底层驱动调用。
避免 CPU/GPU 同步：从 MemoryPool 类定义的源码注释可以看到，当分配命中池中预分配的空闲块时，不会调用cudaMalloc，因此不会触发 CPU 与 GPU 之间的同步。这使得「内存分配 + 内核调用」的交错执行变得非常快。

两大内存池

CuPy 中存在两个相互独立的内存池，分别服务不同的内存类型：

内存池	管理的内存	用途
设备内存池（Device Memory Pool）	GPU 显存	所有 GPU 内存分配，如`cupy.ndarray`的数据存储
固定内存池（Pinned Memory Pool）	不可交换的 CPU 内存（页锁定内存）	CPU → GPU 数据传输过程中的临时缓冲

固定内存（pinned/pagelocked memory）不会被操作系统换出到磁盘，DMA 引擎可以直接访问，因此是主机端到设备端数据传输的关键加速手段，但其分配成本也更高，所以同样值得池化复用。

监控内存时的「内存不释放」是预期行为

在使用nvidia-smi（监控 GPU 显存）或ps（监控 CPU 内存）观察内存占用时，你可能会发现即使数组实例已经超出作用域，内存依然没有被释放。这是完全正常的预期行为——默认内存池会把分配过的内存块「缓存」下来以便复用。池中的空闲块会在后续分配中再次被利用，而不是立即归还给操作系统或 CUDA 驱动。这一点也意味着内存池会尽可能长时间地持有设备内存，从 MemoryPool 源码注释可以看出，这可能导致与其他并行运行的 CUDA 程序产生显存竞争。

二、内存池操作：统计查询与块释放

通过cupy.get_default_memory_pool与cupy.get_default_pinned_memory_pool可以拿到默认的两个内存池实例。这两个函数在 cupy/init.py 中实现，而默认池本身在模块加载时就已创建并注册为 CuPy 的默认分配器：

# cupy/__init__.py（简化） _default_memory_pool = cuda.MemoryPool() _default_pinned_memory_pool = cuda.PinnedMemoryPool() cuda.set_allocator(_default_memory_pool.malloc) cuda.set_pinned_memory_allocator(_default_pinned_memory_pool.malloc)

MemoryPool和PinnedMemoryPool暴露了若干统计与操作接口，核心方法定义在 cupy/cuda/memory.pyx：

used_bytes()：当前被使用（未释放）的字节数；
free_bytes()：已从底层分配但当前空闲的字节数；
total_bytes()：池累计从底层分配的总字节数（used_bytes + free_bytes）；
n_free_blocks()：池中空闲内存块的数量；
free_all_blocks()：释放池中所有空闲块，归还给底层分配器；
set_limit(size=..., fraction=...)/get_limit()：设置/查询当前设备的显存分配上限。

完整示例：观察内存池的完整生命周期

以下示例完整展示了内存池的统计行为（摘自官方指南并补充注释）：

import cupy import numpy mempool = cupy.get_default_memory_pool() pinned_mempool = cupy.get_default_pinned_memory_pool() # 在 CPU 上创建一个数组。 # NumPy 在 CPU 上分配 400 字节（不归 CuPy 内存池管理）。 a_cpu = numpy.ndarray(100, dtype=numpy.float32) print(a_cpu.nbytes) # 400 # 此时两个池都还没有任何分配。 print(mempool.used_bytes()) # 0 print(mempool.total_bytes()) # 0 print(pinned_mempool.n_free_blocks()) # 0 # 将数组从 CPU 传输到 GPU。 # 这会从设备内存池分配 400 字节，并从固定内存池分配另外 400 字节。 # 传输完成后，固定内存会被立即释放回固定内存池。 # 注意：出于性能考虑，实际分配大小会被向上取整（round up）。 a = cupy.array(a_cpu) print(a.nbytes) # 400 print(mempool.used_bytes()) # 512 print(mempool.total_bytes()) # 512 print(pinned_mempool.n_free_blocks()) # 1 # 当数组超出作用域时，设备内存被释放并保留在池中供后续复用。 a = None # 或使用 `del a` print(mempool.used_bytes()) # 0 print(mempool.total_bytes()) # 512 print(pinned_mempool.n_free_blocks()) # 1 # 调用 free_all_blocks 清空内存池。 mempool.free_all_blocks() pinned_mempool.free_all_blocks() print(mempool.used_bytes()) # 0 print(mempool.total_bytes()) # 0 print(pinned_mempool.n_free_blocks()) # 0

这个示例中值得注意的细节：

分配大小取整：请求 400 字节，设备池实际分配了 512 字节，这是 CuPy 底层内存分配对齐策略导致的向上取整行为；
固定内存的释放时机：cupy.array(a_cpu)触发的 CPU→GPU 传输会临时使用固定内存作为中转，传输完成后该固定内存块即被释放回池中（体现为n_free_blocks() == 1）；
池缓存 vs 真正释放：a = None后used_bytes()归零，但total_bytes()仍为 512，说明内存只是归还给池子，并未真正还给系统；只有free_all_blocks()才会把内存归还给底层分配器。

关于free_all_blocks还有一个细节：从源码注释可以看出，内存池为了空间利用率可能会拆分空闲块，被拆分的块即使调用free_all_blocks也不会立即释放，必须等到所有拆分部分重新合并为一个完整块后才会归还。此外free_all_blocks(stream=None)支持传入流参数，以仅释放指定流 arena 中的块。

三、固定内存高层 API：cupyx.empty_pinned 系列

为了更方便地使用固定内存，CuPy 在cupyx命名空间提供了四个高层工具函数，它们在 cupyx/_pinned_array.py 中实现，并在 cupyx/init.py 中导出：

cupyx.empty_pinned(shape, dtype=float, order='C')
cupyx.empty_like_pinned(a, dtype=None, order='K', subok=None, shape=None)
cupyx.zeros_pinned(shape, dtype=float, order='C')
cupyx.zeros_like_pinned(a, dtype=None, order='K', subok=None, shape=None)

这些函数返回的是由固定内存（页锁定内存）支撑的 NumPy 数组。从 empty_pinned 的实现可以看到其本质：先通过cuda.alloc_pinned_memory(nbytes)从 CuPy 的固定内存分配器（默认即固定内存池）取一块内存，再用numpy.ndarray(..., buffer=mem)把这块内存包装成 NumPy 数组。因此，只要 CuPy 的固定内存池处于启用状态，这些 API 分配的固定内存就来自该池，从而享受池化复用的性能收益。

典型用法示例：

import cupy import cupyx # 直接创建固定内存 NumPy 数组 pinned_buf = cupyx.empty_pinned((1024, 1024), dtype=cupy.float32) pinned_zeros = cupyx.zeros_pinned(4096) # 按现有数组的形状/布局创建 a_gpu = cupy.random.random((64, 64)) pinned_like = cupyx.empty_like_pinned(a_gpu) # 固定内存是 CPU→GPU 高效传输的天然载体 a_cpu = cupyx.zeros_pinned(100) # 固定内存缓冲 gpu_arr = cupy.asarray(a_cpu) # 由固定内存出发的传输更快

empty_like_pinned和zeros_like_pinned支持order='K'（尽可能保持与a相同的内存布局），其中subok参数目前尚未支持，传入非None值会抛出TypeError（见源码第 75-76 行）。

四、限制 GPU 内存使用：CUPY_GPU_MEMORY_LIMIT 与 set_limit

在某些场景下（如显存较小的 GPU、需要与其他进程共享 GPU），你可能希望给 CuPy 设置一个显存分配的硬上限。

方式一：环境变量 CUPY_GPU_MEMORY_LIMIT

设置CUPY_GPU_MEMORY_LIMIT环境变量即可为每个 GPU 设备设定可分配显存的硬上限（默认值为0，即不限制）。该变量的完整说明见 docs/source/reference/environment.rst。它支持两种取值格式：

# 方式 A：绝对字节数。设置硬上限为 1 GiB export CUPY_GPU_MEMORY_LIMIT="1073741824" # 方式 B：占 GPU 总显存的百分比。若 GPU 显存为 2 GiB， # 则下面的配置与上面的绝对字节数配置等价。 export CUPY_GPU_MEMORY_LIMIT="50%"

在 Python 中可以通过get_limit()验证限制是否生效：

import cupy print(cupy.get_default_memory_pool().get_limit()) # 1073741824

方式二：MemoryPool.set_limit（按设备差异化限制）

环境变量为所有设备设置统一的默认值，而cupy.cuda.MemoryPool.set_limit则可以为当前设备单独设置（或覆盖环境变量指定的）限制，从而实现对每个 GPU 使用不同上限：

import cupy mempool = cupy.get_default_memory_pool() with cupy.cuda.Device(0): mempool.set_limit(size=1024**3) # 1 GiB with cupy.cuda.Device(1): mempool.set_limit(size=2*1024**3) # 2 GiB

从 set_limit 的源码实现可以归纳出以下关键语义：

size与fraction不能同时指定；
两者都未指定或指定为0时，表示取消限制；
fraction取值应在[0, 1]区间，表示占「当前设备可用显存」的比例；
通过方法设置的限制优先级高于环境变量；
该方法只影响当前设备，而环境变量为所有设备设置默认限制；
限制的修改不是线程安全的——在分配正在进行时修改限制，其他线程可能读取到过期的值，生产代码中需要注意。

注意事项：限制不覆盖池外内存

官方指南明确提醒：CUDA 还会在内存池之外分配一部分 GPU 内存（如 CUDA context、cuBLAS/cuDNN 等库的句柄与工作区）。这部分内存根据使用情况可能占用几十到几百 MiB，且不计入上述限制。因此实际显存占用峰值会略高于你所设置的限制值。

另外，CUPY_GPU_MEMORY_LIMIT也会被 MemoryAsyncPool 在初始化时读取（_parse_limit_string()），即流有序内存池同样支持通过该环境变量设限。

五、更换内存池与自定义分配器

除了默认的内存池，CuPy 允许你传入自定义的内存分配函数来替换默认分配策略，相关接口为：

cupy.cuda.set_allocator(func)：设置设备内存分配器；
cupy.cuda.set_pinned_memory_allocator(func)：设置固定内存分配器。

分配器函数需要满足统一签名：接收一个参数（请求的字节数），返回cupy.cuda.MemoryPointer/cupy.cuda.PinnedMemoryPointer。CuPy 为此提供了几种开箱即用的分配器。

5.1 托管内存（Managed Memory）：malloc_managed

托管内存（Unified Memory）允许 GPU 显存超量订阅（oversubscription），适用于具备cudaDevAttrConcurrentManagedAccess属性（通常为 Pascal 及更新的架构）的 GPU。对应分配器为cupy.cuda.malloc_managed，定义于 cupy/cuda/memory.pyx。

推荐的做法是用它构造一个内存池，而不是直接传入：

import cupy # 使用托管内存支撑的内存池 cupy.cuda.set_allocator(cupy.cuda.MemoryPool(cupy.cuda.malloc_managed).malloc)

需要注意：如果不构造MemoryPool，直接把malloc_managed传给set_allocator，那么内存释放时会立即归还给系统，失去池化复用带来的性能优势——这种直接传递的方式在某些场景下可能是有意为之，但要明确其代价。

5.2 流有序内存分配器（Stream Ordered Memory Allocator）：MemoryAsyncPool

CUDA 11.2 起引入了流有序内存分配器（Stream Ordered Memory Allocator），CuPy 通过cupy.cuda.MemoryAsyncPool提供了实验性接口。与 CuPy 自研内存池类似，它也是以「流有序」的方式异步地从一个内存池中分配/释放内存；关键区别在于，它由NVIDIA CUDA 驱动内置实现，因此同一进程内的其他 CUDA 应用程序也可以轻松地从同一个池中分配内存，天然支持跨库共享。

启用方式：

import cupy # 使用异步流有序内存 cupy.cuda.set_allocator(cupy.cuda.MemoryAsyncPool().malloc) # 创建自定义流 s = cupy.cuda.Stream() # 在流 s 上异步分配内存 with s: a = cupy.empty((100,), dtype=cupy.float64)

MemoryAsyncPool的构造参数pool_handles与MemoryPool不同，用于指定使用哪个 CUDA mempool：'default'表示设备默认 mempool，'current'表示当前 mempool（默认值），也可以传入外部创建的cudaMemPool_t整数句柄；还可以传入与可见设备数量等长的列表，为每个设备独立指定（详见 MemoryAsyncPool 定义）。如果直接把malloc_async传给set_allocator而不构造MemoryAsyncPool，则会使用设备当前的 mempool。

使用流有序内存时有几点必须注意：

流语义由用户自己维护：CuPy 不会自作聪明地替你管理流。你需要自行使用cupy.cuda.Stream与cupy.cuda.Event保证正确的流同步（相关 API 见 docs/source/reference/cuda.rst 的流与事件章节）。释放内存时，会优先在分配它的流上异步释放，其次在当前 CuPy 流上释放；分配内存的流在其上所有内存被释放前被销毁也是允许的。
与 cudaMalloc 的互操作问题：内部使用cudaMalloc（CUDA 默认的同步分配器）的应用与流有序分配器可能产生意外交互——归还到池中的内存对cudaMalloc而言可能不是立即可见的，从而引发意外的 OOM（内存不足）错误。此时可以调用MemoryAsyncPool.free_all_blocks()，或手动做一次（事件/流/设备）同步后重试。
实验性状态与驱动要求：MemoryAsyncPool的 API 与MemoryPool大体一致，但部分方法需要足够新的驱动（以及受支持的硬件、CUDA 版本与平台）。从源码可以看到，统计接口的可用性取决于驱动版本（driverGetVersion() >= 11030）。另外该池目前无法与内存钩子（memory hooks）配合使用，且底层池初始大小为 0，首次分配会按 32 MiB 的整数倍扩容（该内部行为可能随驱动版本变化，API 本身不依赖这些内部细节）。

5.3 彻底禁用内存池

如果出于调试或与外部内存管理机制协同的目的，你可以完全禁用默认内存池。必须在任何其他 CuPy 操作之前执行：

import cupy # 禁用设备内存（GPU）的内存池 cupy.cuda.set_allocator(None) # 禁用固定内存（CPU）的内存池 cupy.cuda.set_pinned_memory_allocator(None)

禁用后，每次分配都会直接调用底层cudaMalloc/cudaHostAlloc，释放时立即归还给系统，这通常会导致明显的性能下降，仅建议在特殊场景下使用。

六、统一内存编程（UMP）支持（实验性）

对于启用了异构内存管理（HMM，Heterogeneous Memory Management）或地址转换服务（ATS，Address Translation Services）的系统——例如 NVIDIA Grace Hopper 超级芯片——CuPy 提供了一种实验性的**统一内存编程（Unified Memory Programming，UMP）**能力，让 NumPy 与 CuPy共享同一份系统内存，CPU 与 GPU 之间不再有显式拷贝。

启用步骤

启用 UMP 需要完成以下 4 个步骤：

第 1 步：安装numpy_allocator包（一个提供 NumPy 对齐分配器的第三方扩展包）。

第 2 步：设置环境变量CUPY_ENABLE_UMP=1。

第 3 步：让 CuPy 从系统内存（malloc_system）分配。cupy.cuda.memory.malloc_system定义于 cupy/cuda/memory.pyx，用于在 HMM/ATS 系统上分配可被 CPU 和 GPU 共同访问的系统内存：

import cupy as cp cp.cuda.set_allocator(cp.cuda.MemoryPool(cp.cuda.memory.malloc_system).malloc)

第 4 步：让 NumPy 切换到对齐分配器，使其也从系统内存分配（此处直接复用 CuPy 编译产物中导出的对齐分配符号）：

import cupy._core.numpy_allocator as ac import numpy_allocator import ctypes lib = ctypes.CDLL(ac.__file__) class my_allocator(metaclass=numpy_allocator.type): _calloc_ = ctypes.addressof(lib._calloc) _malloc_ = ctypes.addressof(lib._malloc) _realloc_ = ctypes.addressof(lib._realloc) _free_ = ctypes.addressof(lib._free) my_allocator.__enter__() # 全局切换分配器

启用后的效果

完成上述配置后，所有数据移动 API——如cupy.ndarray.get()、cupy.asnumpy、cupy.asarray——都变成no-op（零拷贝），因为 CPU 和 GPU 访问的是同一份物理内存。官方指南给出的加速示例：

# a, b, c 是 np.ndarray，d 是 cp.ndarray a = np.random.random(100) b = np.random.random(100) c = np.add(a, b) # CPU 上执行 d = cp.matmul(cp.asarray(a), cp.asarray(c)) # GPU 上执行，但无数据拷贝

本质上，在这种模式下CPU/GPU「内存空间」的区分消失了，np与cp退化为纯粹表征「执行空间」（代码跑在 CPU 还是 GPU 上）的标签。除了 NumPy/CuPy 的这组配置外，用户的业务代码完全不需要任何改动。

需要强调的是，UMP 属于实验性特性，依赖 HMM/ATS 硬件与系统支持（如 Grace Hopper 平台），在普通 x86 + 独立 GPU 的常见环境中并不适用。

七、补充：FFT 计划缓存与设备内存

官方指南还特别提醒（CuPy v8 及以上）：CuPy 提供了 FFT 计划缓存（FFT plan cache），当使用 FFT 及相关函数时，该缓存可能占用一部分设备内存。如果希望回收这部分内存，可以通过缩小或禁用缓存来实现。对于同时使用 FFT 与内存池监控的场景，排查显存占用时不要忽略这一来源。

八、总结与参考

CuPy 的内存管理设计围绕「池化复用」这一核心展开：设备内存池与固定内存池各司其职，通过get_default_memory_pool/get_default_pinned_memory_pool暴露统计接口，通过CUPY_GPU_MEMORY_LIMIT与set_limit控制显存上限，通过set_allocator/set_pinned_memory_allocator支持托管内存、流有序内存乃至 UMP 系统内存等多种分配后端。理解这些机制，可以帮助你在多进程共存、显存受限、跨库共享内存等真实场景下做出正确的配置决策。