CUDA 内存管理实践

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 内存管理实践”