AI 显存分配器深度解析：从 cudaMalloc 系统开销、PyTorch Caching Allocator 块分裂合并到流序复用与 Expandable Segments 碎片治理

AI 显存分配器深度解析：从 cudaMalloc 系统开销、PyTorch Caching Allocator 块分裂合并到流序复用与 Expandable Segments 碎片治理核心痛点：nvidia-smi 明明显示还剩 20GB，PyTorch 却报 CUDA out of memory——显存到底被谁吃了？碎片卡在哪一层？分配器的每个设计决策如何决定一次大模型训练能不能跑起来？适配人群：使用 PyTorch 做训练/推理、遭遇过 OOM 与显存碎片、想读懂 CUDACachingAllocator 源码的算法工程师与 AI 系统工程师收获能力：掌握 cudaMalloc 系统开销的本质、Caching Allocator 的 Block/双池/分裂合并全生命周期、跨流复用的事件护栏与 record_stream 语义、OOM 重试链与碎片治理参数、Expandable Segments 的 VMM 三步曲、零拷贝与 IPC 机制，以及显存诊断工具链的完整实战用法技术背景与演进逻辑：为什么 AI 训练不能直接用 cudaMalloccudaMalloc 的三宗罪，决定了它撑不起训练负载：罪一（隐式同步）：cudaMalloc/cudaFree 在多数驱动路径上具有全设备同步语义 - 每次分配都要等所有流上的异步内核排空 - 训练循环里每步几十上百次分配，等于把异步流水线打回同步慢放罪二（驱动开销）：分配要穿越用户态运行时、驱动、内核态页表更新全链路，单次开销在几十微秒到毫秒级 - 一次迭代累计毫秒级纯开销，乘以每天百万次迭代就是小时级浪费罪三（无缓存语义）：驱动按需分配释放，既不做尺寸圆整，也不管理可复用块 - 训练的动态 shape（batch 波动、变长序列）让物理页反复进出，外部碎片迅速累积AI 训练负载的分配特征恰好放大三宗罪：每个迭代步都重复"分配激活 - 反向传播 - 释放"的循环（时间局部性极强）、shape 随输入波动（尺寸不可预测）、多流多卡并行（跨流复用引入正确性问题）- 缓存式分配器成为框架的必然选择演进时间线（text 树表达）：GPU 显存分配器演进时间线 ├── 2007 - CUDA 1.0 发布 cudaMalloc: 同步语义，适合图形负载不适合训练循环 ├── 2011-2016 - Torch7/早期 PyTorch 引入 THC 缓存分配器: 训练循环免同步分配的开端 ├── 2017 - c10/cuda/CUDACachingAllocator 重构定型: 双池 + Block 分裂合并成为现代形态 ├── 2020-11 - CUDA 11.2 引入 cudaMallocAsync 流序分配: 分配/释放进入流序，驱动级内存池 ├── 2023-10 - PyTorch 2.1 引入 expandable_segments（实验）: VMM 虚拟内存映射治碎片 ├── 2024-2025 - memory snapshot 可视化成熟 + expandable_segments 转为推荐配置 └── 2026-06 - PyTorch 官方 DevLog 系统分析 CUDA Graph 池碎片 + RFC 推进与 NCCL VMM 互操作[总结]：分配器的演进主线只有一条——把"昂贵的驱动级操作"从训练热路径上移走：先缓存（caching）避免频繁 cudaMalloc，再流序化（stream-ordered）避免同步，最后虚拟化（VMM）消灭碎片核心原理深度解析：Caching Allocator 的数据结构与全生命周期层级架构（text 树表达）：DeviceCachingAllocator 内部结构 ├── per-device Malloc 状态 │ ├── large_blocks: 大块池（= 1MB 请求，2MB 圆整） │ ├── small_blocks: 小块池（ 1MB 请求，512B 圆整） │ ├── active_blocks: 已分配块索引（按地址查块） │ ├── cudaMalloc 扩容链: 池未命中 - 驱动申请新 segment │ └── 统计器: allocated / reserved / active / retries / ooms └── Block（最小管理单元） ├── ptr / size: 设备地址与字节数 ├── stream: 分配时所处流的身份 ├── prev / next: 同一 segment 内分裂链的相邻块指针 ├── stream_uses: 跨流使用记录（事件护栏的依据） └── allocated: 缓存态（false=已释放待复用）与活跃态尺寸策略常量（text 树表达）：圆整与分界常量 ├── 小请求界: kSmallSize = 1MB，之下进 small pool ├── 小圆整: 512 字节粒度，小请求向上取整 ├── 小 segment: kSmallBuffer = 2MB（小块从 2MB 段上切出来） ├── 大圆整: 2MB 粒度，大请求向上取整 └── 大 segment: kLargeBuffer = 20MB（大块至少按 20MB 段向驱动申请）分配路径（text 树 + 箭头）：malloc(size) 全路径 ├── 圆整 - 小于 1MB 进 small pool / 否则进 large pool ├── best-fit 查找 - pool.blocks.lower_bound(size) │ ├── 命中且剩余量值得跟踪 - split: 大块裂成 [已用块 + 余量回池] │ ├── 命中且无需分裂 - 直接标 allocated 返回 │ └── 未命中 ↓ ├── 池扩充 - cudaMalloc( max(请求, 段下限) ) 新 segment │ ├── 成功 - 切块返回，reserved 增长 │ └── CUDA_ERROR_OUT_OF_MEMORY ↓ └── OOM 重试链 ├── 释放全部缓存块（release_cached_blocks）- cudaDeviceSynchronize - 重试 ├── 再失败 - 遍历释放 + 打印完整统计报告 - 抛出 RuntimeError └── 全程计数: num_alloc_retries / num_ooms 递增（诊断金指标）释放路径（text 树 + 箭头）：free(block) 全路径 ├── 标记 allocated=false - 插回所属池的有序集合 ├── 合并尝试 - prev 空闲则并块 / next 空闲则并块（Doubly-linked 分裂链上原地合并） ├── 结果: 块驻留缓存（地址不还给驱动）- 下次同尺寸请求零开销复用 └── 例外: 跨流使用过 - 先走事件护栏（见下一章）best-fit 的数据结构选择：每个池是一棵按 (size, ptr) 排序的红黑树（std::set）- 查找O ( l o g n ) O(log n)