3. 显存分配器机制剖析
要理解 GPU 推理服务为什么会在显存看似充足时仍然 OOM，必须把“显存分配”拆成两层来看：一层是 CUDA 驱动/运行时向操作系统申请的物理显存，另一层是框架在已申请到的显存池内部做的二次分配。PyTorch 的 Caching Allocator 正是第二层的核心，它把 cudaMalloc 得到的显存块缓存起来，避免每次张量创建都触发昂贵的驱动调用。而 CUDA VMM（Virtual Memory Management）则提供了更底层的虚拟地址到物理页的映射能力，是理解大块预留、按需提交和碎片治理的关键。
一、从 cudaMalloc 到 Caching Allocator：两级分配模型
当 PyTorch 第一次需要显存时，Caching Allocator 会向 CUDA 申请一块较大的显存，通常以 2MB 为基本粒度（小分配）或更大块（大分配）。这块显存被记录在分配器的空闲块列表中，之后所有张量申请都优先从这个池子里切分，而不是直接调用 cudaMalloc。释放张量时，显存不会立刻还给 CUDA，而是回到空闲块列表，等待下一次复用。这种“缓存不归还”的策略显著降低了分配延迟，但也意味着 torch.cuda.memory_allocated() 与 torch.cuda.memory_reserved() 会出现差异：前者是张量实际占用，后者是分配器向 CUDA 预留的总量。
分配器内部维护多个 size bucket，按请求大小分类管理。小请求走 slab 分配，大请求走独立块分配。slab 分配的核心思想是：把一块较大的显存切成多个固定大小的小块，用位图或空闲链表管理，从而减少元数据开销和外部碎片。例如，一个 20MB 的 slab 可以切成 10 个 2MB 块，当有 1.5MB 的请求时，直接取一个 2MB 块返回，剩余 0.5MB 内部碎片由该块承担。这种设计在大量小张量场景下非常高效，但也会带来内部碎片。
二、块合并：碎片治理的关键动作
当张量释放时，分配器会把对应块标记为空闲，并尝试与地址相邻的空闲块合并。合并的前提是两块在物理地址上连续，且都属于同一个大块预留。如果中间夹着一个仍被占用的块，合并就无法进行，于是空闲块被“切碎”成多个不连续的小块。随着推理请求不断变化，这种碎片会逐渐累积，最终出现“总空闲显存足够，但没有一块连续显存能满足新请求”的情况。
PyTorch 提供了 torch.cuda.memory_summary() 来观察碎片状况。其中 allocated_bytes.all.allocated 表示当前张量占用，reserved_bytes.all.allocated 表示分配器向 CUDA 预留的总量，两者差值就是缓存池中的空闲块。如果 reserved 远大于 allocated，且 num_alloc_retries 持续上升，说明碎片已经影响到新分配。此时可以调用 torch.cuda.empty_cache() 把未使用的缓存块归还给 CUDA，但这会带来后续分配延迟，不适合高频调用。
三、CUDA VMM：从物理连续到虚拟连续
传统 cudaMalloc 要求物理显存连续，这在大模型推理中很容易成为瓶颈。CUDA VMM 通过 cuMemCreate 创建物理页句柄，再用 cuMemMap 把物理页映射到一段连续的虚拟地址空间。这样，虚拟地址看起来是连续的，但底层物理页可以离散分布。对于推理框架而言，这意味着可以预留一大段虚拟地址，按需提交物理页，从而减少预留浪费和外部碎片。
在 PyTorch 2.x 中，可以通过 PYTORCH_CUDA_ALLOC_CONF 启用相关后端。例如设置 backend:cudaMallocAsync 会使用 CUDA 流有序内存分配器，它内部也依赖 VMM 机制。另一种常见配置是 expandable_segments:True，它允许分配器把多个物理页拼成一个可扩展的虚拟段，减少因块不连续导致的 OOM。下面是一个示例配置：
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True,max_split_size_mb:128,garbage_collection_threshold:0.8
其中 max_split_size_mb 控制大块切分阈值，避免把大块切得太碎；garbage_collection_threshold 触发空闲块回收。这些参数需要根据模型张量尺寸分布来调整，而不是照搬默认值。
四、碎片产生的典型过程与观察方法
假设推理服务依次加载三个模型分片，每个分片需要 3GB 连续显存。分配器先预留 10GB 大块，切出三个 3GB 块和一个 1GB 剩余块。当中间分片被卸载后，释放的 3GB 块与相邻的 1GB 空闲块合并成 4GB。但如果此时另一个 2GB 张量占用了 1GB 块旁边的位置，合并就会失败，空闲块被分割成 3GB 和 1GB 两块。后续若需要 4GB 连续块，就会分配失败，尽管总空闲显存有 4GB。
观察碎片最直接的方式是定期打印内存摘要，并关注 num_alloc_retries 和 num_ooms。在服务化场景中，建议在每次请求批次结束后记录 allocated 与 reserved 的比值。如果该比值持续低于 0.7，说明缓存池中空闲块占比过高，可能需要调整 max_split_size_mb 或启用 expandable_segments。同时，避免在推理循环中频繁创建和销毁不同尺寸的张量，尽量复用固定形状的缓冲区，可以从源头减少碎片。
memory_summary() 中 reserved 是否下降、num_alloc_retries 是否归零。参数优化不是一次性的，模型分片、批大小和请求长度变化后都需要重新评估。
总结来看，PyTorch Caching Allocator 通过缓存和 slab 分配提升了分配效率，但块合并失败会带来外部碎片；CUDA VMM 用虚拟连续、物理离散的方式缓解了连续显存压力。理解这两层机制，才能在推理服务中合理设置分配参数、观察碎片指标，并在 OOM 发生时快速判断是真实显存不足还是碎片导致的假象。