March 26, 2025 最近对之前写的 TinyTorch 的内存管理(Allocator)进行了升级，主要是 Copy 了 PyTorch 的实现 CUDACachingAllocator.cpp，去除了其中 stream、cuda graph、expandable segments 以及数据统计相关的逻辑，仅保留最基础的 DeviceCachingAllocator 的核心实现，对第一次接触这块的同学来说会更容易看懂，完整代码详见：Allocator.cpp 基本思路 首先建议读下这篇文章： A guide to PyTorch’s CUDA Caching Allocator 该文章对 PyTorch CUDA 内存管理做了宏观介绍，并给出了伪代码，这里总结一下： 对要分配的 size 进行向上取整，比如 512 的倍数，从而提高 cache 复用概率进而减少碎片； 实际分配的时候，分配更大尺寸的内存块，这一块剩下的部分可以供之后的分配使用，从而减少调用系统 cudaMalloc 的次数； 区别对待小内存分配和大内存分配，使用不同的内存池管理，减少碎片产生； 当然还有其它的一些设计，比如每个 stream 使用独立内存池、通过 Events 跟踪对象生命周期、OOM 处理等，在我的项目里目前都不涉及，先忽略。 数据结构 首先是 Block，字段如下： struct Block { size_t size;Continue reading “CUDA 内存管理实践”