显存分配机制：从cudaMalloc到虚拟内存管理
聊到GPU编程，显存分配是绕不开的第一道坎。我刚开始接触CUDA时，觉得不就是调个cudaMalloc嘛，有啥好研究的？结果第一次跑大模型训练，直接OOM，程序崩溃，显卡风扇狂转——嗯，从那以后我再也不敢小看显存分配了。
说白了，显存分配就是告诉GPU：「这块内存我要用，你给我留着」。但怎么留、留多少、什么时候释放，这里面门道不少。今天咱们就掰开揉碎，把CUDA运行时分配、驱动API分配、还有虚拟内存管理这三层机制讲清楚。
核心要点：显存分配不是简单的malloc，它涉及运行时库、驱动层、操作系统虚拟内存三个层次。理解这些，你才能写出不爆显存的代码。
一、cudaMalloc：最常用的显存分配方式
cudaMalloc是CUDA运行时API提供的函数，用法和C语言的malloc几乎一样。我个人习惯在项目初期先用它，因为简单直接，不容易出错。
float *d_data;
cudaMalloc((void**)&d_data, size * sizeof(float));
// 使用d_data...
cudaFree(d_data);这里有个坑：cudaMalloc返回的是设备指针，你不能在主机端直接解引用。我曾经见过新手把d_data当主机指针用，结果段错误，排查了半天。
小技巧：每次cudaMalloc后，建议检查返回值。用cudaError_t err = cudaMalloc(...); if (err != cudaSuccess)来捕获分配失败的情况。显存不够时，cudaMalloc会返回cudaErrorMemoryAllocation。
cudaMalloc底层做了什么？它实际上调用了驱动API的cuMemAlloc，然后做了一些封装。所以性能上两者差别不大，但运行时API更方便——它会自动管理上下文，你不需要手动指定哪个GPU。
二、cuMemAlloc：驱动层的直接控制
如果你需要更精细的控制，比如指定分配在哪个GPU上，或者想绕过运行时的一些开销，那就得用驱动API了。cuMemAlloc是驱动层的函数，用法稍微复杂一点。
CUdeviceptr d_ptr;
cuMemAlloc(&d_ptr, size);
// 使用d_ptr...
cuMemFree(d_ptr);注意这里的CUdeviceptr，它本质上是一个unsigned long long，代表设备内存地址。驱动API要求你手动管理CUDA上下文——你得先cuCtxCreate或cuCtxSetCurrent，然后才能分配。
我在做多GPU通信时，就吃过这个亏。当时想用驱动API直接分配显存给NVLink传输用，结果忘了设置当前上下文，cuMemAlloc一直返回CUDA_ERROR_INVALID_CONTEXT。排查了半小时才发现问题。
| 特性 | cudaMalloc | cuMemAlloc |
|---|---|---|
| API层级 | 运行时API | 驱动API |
| 上下文管理 | 自动 | 手动 |
| 多GPU支持 | 通过cudaSetDevice | 直接指定上下文 |
| 性能开销 | 略高（封装层） | 更低 |
| 使用场景 | 日常开发 | 底层优化、驱动开发 |
警告：不要混用运行时API和驱动API的分配/释放函数。比如用cudaMalloc分配的内存，不要用cuMemFree释放。这会导致未定义行为，严重时可能让驱动崩溃。
三、虚拟内存管理：现代GPU的显存分配新范式
从CUDA 10.2开始，NVIDIA引入了虚拟内存管理（VMM）功能。说白了，就是让GPU也能像CPU那样，用虚拟地址映射物理内存。这玩意儿在显存复用场景下特别有用。
为什么需要VMM？你想想看，传统cudaMalloc分配的是连续的物理显存。如果你的程序需要频繁分配释放小块内存，很容易产生碎片。VMM允许你预先保留一大块虚拟地址空间，然后按需映射物理页。
CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.location.id = 0; // GPU 0
size_t granularity;
cuMemGetAllocationGranularity(&granularity, &prop,
CU_MEM_ALLOC_GRANULARITY_MINIMUM);
CUmemGenericAllocationHandle handle;
cuMemCreate(&handle, size, &prop, 0);
CUdeviceptr ptr;
cuMemAddressReserve(&ptr, size, 0, 0, 0);
cuMemMap(ptr, size, 0, handle, 0);
// 使用ptr...
cuMemUnmap(ptr, size);
cuMemRelease(handle);这段代码看着复杂，其实核心就三步：
- 创建物理内存块：
cuMemCreate分配物理显存，返回一个句柄 - 预留虚拟地址：
cuMemAddressReserve在GPU的虚拟地址空间中划一块区域 - 建立映射：
cuMemMap把物理内存映射到虚拟地址
我在做显存池（memory pool）时，就用了这套机制。预先分配一大块虚拟地址空间，然后按需映射物理页。当某个tensor不再需要时，只解除映射，不释放物理内存。这样下次分配同样大小的tensor时，直接复用物理页，速度飞快。
关键点：VMM的粒度很重要。不同GPU的granularity不同，Tesla V100是64KB，A100是2MB。分配时一定要对齐到这个粒度，否则cuMemCreate会报错。
三种分配方式的对比与选择
说了这么多，到底该用哪种？我个人的经验是：
- 日常开发、原型验证：用
cudaMalloc，省心省力 - 需要多GPU协同、或追求极致性能：用
cuMemAlloc，控制更精细 - 做显存池、动态分配、或处理超大模型：用VMM，灵活且高效
记得有一次，我在优化一个推荐系统的推理服务。模型有几十个embedding表，每个表大小不一，而且请求来了才动态加载。如果用cudaMalloc，每次加载都要分配释放，碎片化严重，跑几个小时就OOM了。后来改用VMM，预先保留2GB虚拟地址空间，按需映射物理页，问题迎刃而解。
避坑指南：我曾经在VMM上踩过一个坑——忘记调用cuMemAddressFree释放虚拟地址。结果虚拟地址空间被耗尽，后续分配全部失败。记住：cuMemUnmap只解除映射，不释放虚拟地址；cuMemAddressFree才释放虚拟地址。
知识体系总览
下面这张图，把三种分配方式的关系和适用场景梳理清楚了：
从这张图能看出来，cudaMalloc和cuMemAlloc是「分配即映射」的模式——分配的同时就占用了物理显存。而VMM把「分配虚拟地址」和「映射物理内存」拆成了两步，给了你更大的灵活性。
嗯，关于显存分配机制，今天就聊到这儿。记住一句话：没有银弹。不同的场景选不同的方式，关键是要理解每层API背后的代价和限制。下次遇到OOM，别急着加显存，先想想是不是分配方式没选对。
一句话总结：cudaMalloc省心，cuMemAlloc精细，VMM灵活。从应用层到硬件层，每层都有它的设计哲学。
想要更多资料，联系可可