PyTorch源码：PyTorch是怎么管理显存的 (3)
源码分析：图2-15\n \nTL;DR：DeepSeek：\n1️⃣ 核心对象\nBlock：基础显存单元，由cudaMalloc或ExpandableSegment分配\nBlockPool：分大小缓存池（≤1MB进small_blocks，>1MB进large_blocks）\nPrivatePool：用户自定义池，支持细粒度策略（如ncclMemAlloc）\n2️⃣ 分配流程四部曲\n🔹 预处理\n加全局锁🔒\n处理完成事件释放Block（process_events）\n🔹 参数计算\n用户尺寸→512倍数对齐（e.g. 960B→1024B）\n按尺寸选池 + 物理显存策略（e.g. 1024B→实际分配2MB）\n🔹 检索/分配\n优先从对应BlockPool找空闲Block\n失败→触发回调清理→再次检索\n仍失败→阶梯式分配：\n垃圾回收（若有内存限额）🗑️\n新分配Block → 释放大块 → 清空缓存池（三次尝试❗）\n🔹 后处理\n成功：按需分裂Block（碎片优化）✂️\n失败：抛出经典 CUDA OOM 错误💥\n3️⃣ 关键策略\n碎片管理：超尺寸Block自动分裂缓存\n多级回收：事件驱动/回调/强制释放三级保障\n定制化：环境变量+PrivatePool支持特殊分配需求\n总结：PyTorch显存分配是 事件驱动+池化缓存+阶梯fallback 的精巧平衡⚖️！