PyTorch 源码学习：GPU 内存管理之深入分析 CUDACachingAllocator

因引入 expandable_segments 机制，PyTorch 2.1.0 版本发生了较大变化。本文关注的是 PyTorch 原生的 GPU 内存管理机制，故研究的 PyTorch 版本为 2.0.0。代码地址：

• c10/cuda/CUDACachingAllocator.h

• c10/cuda/CUDACachingAllocator.cpp
  更多内容请参考：

• Ubuntu 22.04 LTS 源码编译安装 PyTorch

• 【翻译】pytorch/CONTRIBUTING.md

• 【翻译】Pytorch机制，源代码分析与内存管理调研

• 深度学习框架与静态/动态计算图【笔记】

• PyTorch 源码学习：阅读经验 & 代码结构

• PyTorch 源码学习：从 Tensor 到 Storage

• PyTorch 源码学习：Dispatch & Autograd & Operators

• PyTorch 源码学习：GPU 内存管理之它山之石------TensorFlow BFC 算法

文章目录

• [1 PyTorch 为什么要有 CUDACachingAllocator？](#1 PyTorch 为什么要有 CUDACachingAllocator？)
• [2 核心类](#2 核心类)
• • [2.1 class NativeCachingAllocator](#2.1 class NativeCachingAllocator)
  • [2.2 Allocator 和 CUDAAllocator](#2.2 Allocator 和 CUDAAllocator)
• [3 全局变量](#3 全局变量)
• [4 核心数据结构](#4 核心数据结构)
• • [4.1 struct Block 内存块](#4.1 struct Block 内存块)
  • • [4.1.1 struct SegmentInfo 内存段的信息](#4.1.1 struct SegmentInfo 内存段的信息)
    • [4.1.2 struct BlockInfo 分配块的信息](#4.1.2 struct BlockInfo 分配块的信息)
    • [4.1.3 进一步区分 Segment 和 Block](#4.1.3 进一步区分 Segment 和 Block)
  • [4.2 struct BlockPool 内存池](#4.2 struct BlockPool 内存池)
  • [4.3 struct PrivatePool 私有内存池](#4.3 struct PrivatePool 私有内存池)
  • [4.4 struct AllocParams 分配内存时的参数和状态](#4.4 struct AllocParams 分配内存时的参数和状态)
  • [4.5 class DeviceCachingAllocator 类成员变量](#4.5 class DeviceCachingAllocator 类成员变量)
• [5 PyTorch BFC 算法](#5 PyTorch BFC 算法)
• [6 核心函数](#6 核心函数)
• • [6.1 Block* malloc 分配函数](#6.1 Block* malloc 分配函数)
  • [6.2 void free 释放函数](#6.2 void free 释放函数)
• [7 功能函数](#7 功能函数)
• • [7.1 static bool BlockComparator 自定义比较函数](#7.1 static bool BlockComparator 自定义比较函数)
  • [7.2 static size_t round_size 对内存大小 size 进行调整](#7.2 static size_t round_size 对内存大小 size 进行调整)
  • [7.3 BlockPool& get_pool 返回对应的内存块池](#7.3 BlockPool& get_pool 返回对应的内存块池)
  • [7.4 static size_t get_allocation_size 确定实际要进行分配的内存大小](#7.4 static size_t get_allocation_size 确定实际要进行分配的内存大小)
  • [7.5 bool get_free_block 寻找满足分配需求的空闲内存块](#7.5 bool get_free_block 寻找满足分配需求的空闲内存块)
  • [7.6 void garbage_collect_cached_blocks 释放未使用的缓存块](#7.6 void garbage_collect_cached_blocks 释放未使用的缓存块)
  • [7.7 bool release_available_cached_blocks 优先从最大的块开始释放](#7.7 bool release_available_cached_blocks 优先从最大的块开始释放)
  • [7.8 bool release_cached_blocks 释放所有未使用的内存块](#7.8 bool release_cached_blocks 释放所有未使用的内存块)
  • [7.9 void release_blocks 释放所有非拆分块](#7.9 void release_blocks 释放所有非拆分块)
  • [7.10 void release_block 释放指定的内存块](#7.10 void release_block 释放指定的内存块)
  • [7.11 bool alloc_block 为指定的分配请求分配一个内存块](#7.11 bool alloc_block 为指定的分配请求分配一个内存块)
  • [7.12 cudaError_t cudaMallocMaybeCapturing 一个封装的内存分配工具](#7.12 cudaError_t cudaMallocMaybeCapturing 一个封装的内存分配工具)
  • [7.13 bool is_split 判断 block 是否被拆分过](#7.13 bool is_split 判断 block 是否被拆分过)
  • [7.14 bool should_split 判断给定的内存块 block 是否应该进行分割](#7.14 bool should_split 判断给定的内存块 block 是否应该进行分割)
  • [7.15 void free_block 将一个块移动到缓存的空闲块池中](#7.15 void free_block 将一个块移动到缓存的空闲块池中)
  • [7.16 size_t try_merge_blocks 合并先前拆分的块](#7.16 size_t try_merge_blocks 合并先前拆分的块)
• [8 问题与挑战](#8 问题与挑战)
• 参考资料

1 PyTorch 为什么要有 CUDACachingAllocator？

CUDACachingAllocator是 PyTorch 的 CUDA 缓存分配器。其目的是通过缓存和复用内存块来减少频繁的cudaMalloc和cudaFree操作，提升 CUDA 程序的性能。

GPU 底层机制分析：显存分配开销 这篇文章通过实验分析了 CUDA 原生显存分配接口的显存分配开销和显存分配碎片。结论就是，虽然 CUDA 原生显存分配接口逻辑比较简单，但显存分配开销无法接受，显存分配碎片也需要优化。

以下内容来自：Pytorch 显存管理机制与显存占用分析方法，这篇文章对显存管理机制的总结写得很好。

GPU 作为一种通用的数据处理设备，为了满足更广泛客户的需求且保证更小的维护成本，其 API 在设计的时候比较开放，++尽管 CUDA 生态中也有高阶 API，但并没有针对某个深度学习框架做设计优化，其中显存的精细管理留给上层的深度学习框架去完成++。

cudaMalloc（CUDA API）是从 GPU 申请显存最常用的方式，给定指针和数据大小即可进行 API 调用，其调用有着不小的时间开销，且是 stream 内的同步操作 。当深度学习框架使用的数据非常零碎且数量多时，需要反复调用cudaMalloc，该行为会直接影响程序的整体性能，因此++深度学习框架的显存管理机制在设计时要尽量降低cudaMalloc的调用频次++。

PyTorch 框架基于 CUDA API 实现了一套显存管理逻辑/机制，可更好地满足框架的日常使用需求，相比原生的 CUDA API 可做到管理细化、使用相对高效，其采用动态申请与二次分配的设计思路：

• 动态申请：在使用的时候根据用量实时地向 GPU 发出请求，最大优点是不会占用过量的显存，方便多人同时使用一个设备（与之相对的是 TensorFlow 早期版本在启动前就把 GPU 上的大部分显存都申请到，然后再去分配使用）。
• 二次分配 ：将显存的申请与使用进行分离，即显存申请后会进行二次分配。显存管理机制会先通过cudaMalloc向 GPU 申请一个显存块 Segment，然后从 Segment 分离出子块 Block，我们使用的是分离后的 Block 显存，而不直接使用 Segment。

---

---

2 核心类

Allocator -> CUDAAllocator -> NativeCachingAllocator -> DeviceCachingAllocator

struct C10_API Allocator

class CUDAAllocator : public Allocator

class NativeCachingAllocator : public CUDAAllocator {
 public:
  std::vector<std::unique_ptr<DeviceCachingAllocator>> device_allocator;

class DeviceCachingAllocator

以下内容来自：PyTorch显存管理介绍与源码解析（二）

显存管理代码中的类的大体关系如下图所示。cudaCachingAllocator类的结构相对来说比较清晰（代码不是特别清晰），DeviceCachingAllocator管理一个设备显存 ；THCCachingAllocator管理一个进程上所有的DeviceCachingAllocator；CudaCachingAllocator继承自Allocator，方便框架调用。

DeviceCachingAllocator类是显存管理机制真正实现的地方，它负责完成 device 上的显存管理，包括显存开辟、分块、合并、释放等。

说明 ：这张图分析的应该不是作者在文中提到的 PyTorch 1.13.0，因为在该版本的代码注释中提到THCCachingAllocator这种说法已经是 old-style 了。作者这里提到的THCCachingAllocator在最新的版本里对应的是NativeCachingAllocator，而CudaCachingAllocator对应的是CUDAAllocator。但上面的内容仍值得参考。

---

以下内容来自：Pytorch 1.6 显存管理CudaCachingAllocator剖析

Pytorch 通过 Allocator 实现显存管理，实现如下图：

• DeviceCachingAllocator，每个 GPU 设备卡都维护一个这样的结构，用于对该设备进行显存管理；
• THCCachingAllocator，维护一个 DeviceCachingAllocator 列表及一些全局的状态，核心逻辑委托给 DeviceCachingAllocator 实现；
• CudaCachingAllocator，这是 Pytorch 默认实现的分配器包装器，向上层暴露一些有用的接口，核心逻辑还是委托给 THCCachingAllocator 来实现，用户也可以自定义一个实现分配器替代 CudaCachingAllocator。

在 PyTorch 1.6.0 里确实是上面这样的类关系。

---

2.1 class NativeCachingAllocator

以下内容来自：PyTorch显存管理介绍与源码解析（二）

NativeCachingAllocator该类主要是管理 CPU 主进程调用的 GPU 的显存使用状况。通过创建一个数组，每个 GPU 有一个与之对应的DeviceCachingAllocator实例管理；同时记录下进程中创建的全部的 blocks。

---

NativeCachingAllocator的代码实现：

class NativeCachingAllocator : public CUDAAllocator {
 private:
 
  // 存储已分配的显存块
  // allocated blocks by device pointer
  ska::flat_hash_map<void*, Block*> allocated_blocks;
  
  // 将分配的显存块添加到 allocated_blocks 中
  void add_allocated_block(Block* block) {
    allocated_blocks[block->ptr] = block;
  }

 public:
  // 每个设备（GPU）都有一个 DeviceCachingAllocator 实例，用于管理该设备上的显存分配
  std::vector<std::unique_ptr<DeviceCachingAllocator>> device_allocator;
  
  // 根据设备指针获取对应的显存块
  Block* get_allocated_block(void* ptr, bool remove = false) {
    auto it = allocated_blocks.find(ptr);  // 查找 ptr 对应的 Block*
    if (it == allocated_blocks.end()) {
      return nullptr;
    }
    Block* block = it->second;
    if (remove) {
      allocated_blocks.erase(it);
    }
    return block;
  }
 
  // 初始化 device_allocator，为每个设备创建一个 DeviceCachingAllocator 实例
  void init(int device_count) override {
    const auto size = static_cast<int64_t>(device_allocator.size());
    if (size < device_count) {
      device_allocator.resize(device_count);
      for (const auto i : c10::irange(size, device_count)) {
        device_allocator[i] = std::make_unique<DeviceCachingAllocator>();
      }
    }
  }
  
  // 在指定设备上分配显存，并将分配的显存块添加到 allocated_blocks 中
  /** allocates a block which is safe to use from the provided stream */
  void malloc(void** devPtr, int device, size_t size, cudaStream_t stream) {
     // 1) 检查设备号是否有效
    TORCH_INTERNAL_ASSERT(
        0 <= device && static_cast<size_t>(device) < device_allocator.size(), 
        "Allocator not initialized for device ",
        device,
        ": did you call init?");
    // 2) 调用 6.1 Block* malloc 分配显存块
    Block* block = device_allocator[device]->malloc(device, size, stream);
    // 3) 将分配的显存块添加到 allocated_blocks 中
    add_allocated_block(block);
  }
 
  // 释放指定的显存块
  void free(void* ptr) {
    if (!ptr) {
      return;
    }
    // 1) 从 allocated_blocks 中获取对应的 Block*，并将其从 allocated_blocks 中移除
    Block* block = get_allocated_block(ptr, true /* remove */);
    if (!block) {
      TORCH_CHECK(false, "invalid device pointer: ", ptr);
    }
    // 2) 调用 6.2 void free 释放显存块
    device_allocator[block->device]->free(block);
  }
  
  // 分配显存并返回一个 DataPtr 对象
  DataPtr allocate(size_t size) const override {
    constexpr size_t one_exa_bytes = 1152921504606846976ULL;
    TORCH_CHECK_WITH(
        OutOfMemoryError,
        size < one_exa_bytes,
        "CUDA out of memory. Tried to allocate more than 1EB memory.");
    int device;
    C10_CUDA_CHECK(cudaGetDevice(&device));
    void* r = nullptr;
    
	// 关于 forceUncachedAllocator() 有下面这样的注释
	// 其实就是直接不使用 CUDACachingAllocator，直接调用 cudaMalloc 分配内存
	// Returns whether to force all allocations to bypass the caching allocator and
	// go straight to cudaMalloc.  This setting is useful when debugging GPU memory
	// errors, since the caching allocator foils cuda-memcheck.
    if (forceUncachedAllocator()) {
      // Deliberately don't use cudaMallocMaybeCapturing here, to force an error
      // if someone tries to use forceUncachedAllocator while capturing.
      C10_CUDA_CHECK(cudaMalloc(&r, size));
      return {r, r, &uncached_delete, Device(DeviceType::CUDA, device)};
    }
	
	// 使用 CUDACachingAllocator 分配内存
	// 调用了 void malloc(void** devPtr, int device, size_t size, cudaStream_t stream)
    if (size != 0) {
      // Allocator declars allocate const!?
      const_cast<NativeCachingAllocator*>(this)->malloc(
          &r, device, size, cuda::getCurrentCUDAStream(device));
    }
    return {r, r, &local_raw_delete, Device(DeviceType::CUDA, device)};
  }
  
  // 返回删除器函数指针
  DeleterFnPtr raw_deleter() const override {
    if (forceUncachedAllocator()) {
      return &uncached_delete;
    } else {
      return &local_raw_delete;
    }
  }
 
  // 分配指定大小的显存
  void* raw_alloc(size_t nbytes) override {
    if (nbytes == 0) {
      return nullptr;
    }
    int device;
    C10_CUDA_CHECK(cudaGetDevice(&device));
    void* r = nullptr;
    // 调用了 void malloc(void** devPtr, int device, size_t size, cudaStream_t stream)
    malloc(&r, device, nbytes, cuda::getCurrentCUDAStream(device));
    return r;
  }
  
  // 在指定的 CUDA 流上分配显存
  void* raw_alloc_with_stream(size_t nbytes, cudaStream_t stream) override {
    if (nbytes == 0) {
      return nullptr;
    }
    int device;
    C10_CUDA_CHECK(cudaGetDevice(&device));
    void* r = nullptr;
    // 调用了 void malloc(void** devPtr, int device, size_t size, cudaStream_t stream)
    malloc(&r, device, nbytes, stream);
    return r;
  }
  
  // 释放指定的显存
  void raw_delete(void* ptr) override {
    this->free(ptr);
  }

};

这里要区分下DataPtr allocate(size_t size)和void* raw_alloc(size_t nbytes)：

• allocate调用链：allocate -> malloc -> add_allocated_block。在分配显存后，会将显存块添加到allocated_blocks中，并返回一个 DataPtr 对象。
• raw_alloc调用链：raw_alloc -> malloc。仅分配显存并返回指针。

TODO: allocate和raw_alloc两者的使用场景还是有很大区别的，因本人精力有限，暂不深入研究了。

一些线索：

1）struct C10_API Allocator的代码实现：

struct C10_API Allocator {
  virtual ~Allocator() = default;

  virtual DataPtr allocate(size_t n) const = 0;

  // If this returns a non nullptr, it means that allocate()
  // is guaranteed to return a unique_ptr with this deleter attached;
  // it means the rawAllocate and rawDeallocate APIs are safe to use.
  // This function MUST always return the same BoundDeleter.
  virtual DeleterFnPtr raw_deleter() const {
    return nullptr;
  }
  void* raw_allocate(size_t n) {
    auto dptr = allocate(n);
    AT_ASSERT(dptr.get() == dptr.get_context());
    return dptr.release_context();
  }
  void raw_deallocate(void* ptr) {
    auto d = raw_deleter();
    AT_ASSERT(d);
    d(ptr);
  }
};

2）class CUDAAllocator : public Allocator的代码实现：

class CUDAAllocator : public Allocator {
 public:
  virtual void* raw_alloc(size_t nbytes) = 0;
  virtual void* raw_alloc_with_stream(size_t nbytes, cudaStream_t stream) = 0;
  virtual void raw_delete(void* ptr) = 0;
  // ......
};

2.2 Allocator 和 CUDAAllocator

不是本文的重点，就不深入分析了。前文的"一些线索"也展示了 Allocator 和 CUDAAllocator 两个类的部分代码，感兴趣可以直接学习源码。

本文剩余内容都实现于 DeviceCachingAllocator 类中，这也是与设备打交道的最底层的一个类。

3 全局变量

全局变量的代码实现：

constexpr size_t kMinBlockSize = 512; 
// all sizes are rounded to at least 512 bytes
// 为了优化内存管理，显存块的大小通常会对齐到特定的值。
// 表示显存分配的最小粒度为 512 字节。

constexpr size_t kSmallSize = 1048576; 
// largest "small" allocation is 1 MiB
// 定义了 "小分配" 的最大阈值，为 1 MiB（1024 * 1024 字节）。
// 小于或等于 kSmallSize 的分配被视为 "小分配"。

constexpr size_t kMinLargeAlloc = 10485760; 
// allocations between 1 and 10 MiB may use kLargeBuffer
// 定义了 "大分配" 的最小阈值，为 10 MiB（10 * 1024 * 1024 字节）。
// 大于或等于 kMinLargeAlloc 的分配被视为 "大分配"。

constexpr size_t kSmallBuffer = 2097152; 
// "small" allocations are packed in 2 MiB blocks
// 定义了 "小分配" 的默认打包大小，为 2 MiB（2 * 1024 * 1024 字节）。

constexpr size_t kLargeBuffer = 20971520; 
// "large" allocations may be packed in 20 MiB blocks
// 定义了 "大分配" 的默认打包大小，为 20 MiB（20 * 1024 * 1024 字节）。

constexpr size_t kRoundLarge = 2097152; 
// round up large allocations to 2 MiB
// 定义了 "大分配" 的对齐单位，为 2 MiB（2 * 1024 * 1024 字节）。
// 所有大分配的大小都会被向上取整到 2 MiB 的倍数，以简化分配管理并减少碎片。

constexpr size_t kRoundUpPowerOfTwoIntervals = 16;

4 核心数据结构

4.1 struct Block 内存块

Block表示内存池中的一个内存块，封装了内存分配的信息。

以下内容来自：一文读懂 PyTorch 显存管理机制

Block：

• 分配 / 管理内存块的基本单位，(stream_id, size, ptr) 三元组可以特异性定位一个 Block，即 Block 维护一个 ptr 指向大小为 size 的内存块，隶属于 stream_id 的 CUDA Stream。
• 所有地址连续的 Block（不论是否为空闲，只要是由Allocator::malloc得来的）都被组织在一个双向链表里，便于在释放某一个 Block 时快速检查前后是否存在相邻碎片，若存在可以直接将这三个 Block 合成为一个。
• Block 在 Allocator 内有两种组织方式 ，
  • 一种是显式地组织在 BlockPool（红黑树）中，按照大小排列；
  • 另一种是具有连续地址的 Block 隐式地组织在一个双向链表里（通过结构体内的 prev, next 指针），可以以 O(1) 时间查找前后 Block 是否空闲，便于在释放当前 Block 时合并碎片。

一段连续空间内部（由双向链表组织的 Block 们）如下图所示：

---

Block的代码实现：

struct Block {
  int device; // gpu 表示分配内存的 GPU 设备 ID。
  
  cudaStream_t stream; 
  // allocation stream 
  // 分配内存时关联的 CUDA 流。
  stream_set stream_uses; 
  // streams on which the block was used 
  // 记录该内存块被哪些流使用过。
  
  size_t size; // block size in bytes 块的大小（字节）。
  size_t requested_size; // memory originally requested 最初请求的内存大小。
  
  BlockPool* pool{nullptr}; // owning memory pool 该内存块所属的内存池。
  
  void* ptr{nullptr}; // memory address 指向实际分配的内存地址。
  
  bool allocated{false}; // in-use flag 标志是否正在使用。
  
  Block* prev{nullptr}; // prev block if split from a larger allocation
  Block* next{nullptr}; // next block if split from a larger allocation
  // 若该块是从更大的块分割而来，指向前后关联块。
  
  int event_count{0}; // number of outstanding CUDA events
  
  int gc_count{0}; 
  // counter for prioritizing older / less useful blocks for garbage collection
  // 用于垃圾回收计数，优先回收不活跃的块。
  
  // 构造函数1：用于创建实际的 Block 对象。
  Block(
      int device,
      cudaStream_t stream,
      size_t size,
      BlockPool* pool,
      void* ptr)
      : device(device),
        stream(stream),
        stream_uses(),
        size(size),
        requested_size(0),
        pool(pool),
        ptr(ptr) {}

  // constructor for search key
  // 构造函数2：用于搜索时的键。
  Block(int device, cudaStream_t stream, size_t size)
      : device(device),
        stream(stream),
        stream_uses(),
        size(size),
        requested_size(0) {}
        
  // 检查该块是否是通过分割其他内存块生成的。
  bool is_split() const {
    return (prev != nullptr) || (next != nullptr);
  }
};

注意：以下 4.1.1 和 4.1.2 中的 Segment 和 Block 其实都是逻辑上的概念，本质上其实都是 4.1 中的 Block。可以理解为，Segment 是cudaMalloc直接申请得到的一段内存，而 Block 是实际分配给 Tensor 的内存块，一个 Segment 可能包含多个 Block。
SegmentInfo、BlockInfo 记录 segment 和 block 的信息，用于 Snapshot 数据记录（可用于显存可视化）。

4.1.1 struct SegmentInfo 内存段的信息

Segment 是通过一次cudaMalloc调用分配的一块连续的显存区域。一个内存段 (Segment) 可以包含多个分配块 (Block) 。

SegmentInfo的代码实现：

// Struct containing info of a memory segment (i.e. one contiguous cudaMalloc).
struct SegmentInfo {
  int64_t device = 0; // 分配此段的 CUDA 设备。
  int64_t address = 0; // 内存段的起始地址。
  int64_t total_size = 0; // 整个内存段的总大小。
  int64_t requested_size = 0; // 用户请求的总大小。
  int64_t allocated_size = 0; // 段中实际被分配的内存大小。
  int64_t active_size = 0; // 段中当前正在使用的内存大小。
  cudaStream_t stream = 0; // 与此段相关联的 CUDA 流。
  bool is_large = false; // 标记此段是否是一个大段分配（large allocation）。
  std::vector<BlockInfo> blocks; // 此段内的分配块列表。
};

4.1.2 struct BlockInfo 分配块的信息

Block 是 CUDA 显存分配的基本单元。一个分配块 (Block) 可以是通过cudaMalloc获取的内存段 (Segment) 的一部分。

BlockInfo的代码实现：

// Struct containing info of an allocation block (i.e. a fractional part of a cudaMalloc)..
struct BlockInfo {
  int64_t size = 0; // 当前块的大小（以字节为单位）。
  int64_t requested_size = 0; // 用户实际请求的大小。分配器通常会将请求大小对齐到某个块大小。
  int32_t gc_counter = 0; // 垃圾回收计数器，记录该块的回收次数。
  bool allocated = false; // 标记该块是否已被分配给用户。
  bool active = false; // 标记该块是否正在被某个 CUDA 流使用（例如分配或计算）。
  std::vector<History> history;
};

4.1.3 进一步区分 Segment 和 Block

以下内容来自：PyTorch显存机制分析

直观点来说，PyTorch 的显存管理是一个层级结构。

然后他们又是包含与被包含的关系，即 PyTorch Allocated memory 使用的是 PyTorch Cached Memory 里的显存，PyTorch Cached Memory 则用的是 GPU 的显存。

这里 PyTorch Allocated memory 指的就是实际分配给上层应用 的 Block 占用的内存，而 PyTorch Cached Memory 指的就是直接通过cudaMalloc申请得到 并缓存在内存池的 Block（也就是逻辑上的 Segment）占用的内存。

---

以下内容来自：PyTorch显存管理介绍与源码解析（一）和 PyTorch显存管理介绍与源码解析（二）

在管理机制中，将显存的申请与使用过程进行了分离，即显存申请后会进行二次分配 ，其过程是：先通过cudaMalloc申请一个显存块 segment，然后从 segment 分离出来的子显存块 block，框架的上层应用使用的是分离后的 block 显存，上层应用不直接使用 segment。

进一步，通过池化的方式将 block 按照块放入不同的显存池中，进行分类管理。

PyTorch 内存管理机制里，对显存分了两级 （segment、blocks），显存池分了两类 （active pool、remaining pool）。在结构体的实现上面主要是靠块 (block) 结构体来完成承上启下的功能，显存池结构体（blockPool）主要是记录 block 的归属；segment 是一个从cudaMalloc创建出来的显存块（承载依然是 block 结构体），segment 再次切分就是能够提供给用户使用的显存。

---

以下内容来自：Pytorch 1.6 显存管理CudaCachingAllocator剖析

下图为多种进程使用显存的布局：

• allocated 是 Pytorch 分配给 Tensor 使用的块，即torch.cuda.memory_allocated()的值
• cached 是 Pytorch 空闲块
• reserved 是 Pytorch 进程管理的所有显存块，即torch.cuda.memory_reserved()的值
• reserved = allocated + cached

逻辑上来看，这里的 reserved 就是所有 Segment 占用的内存，allocated 就是所有分配出去的 Block 占用的内存，cached 就是所有未分配出去的 Block 占用的内存。

---

4.2 struct BlockPool 内存池

BlockPool用于管理 多个Block，可以分为小块池 和大块池。

以下内容来自：一文读懂 PyTorch 显存管理机制

BlockPool：

• 内存池，用std::set存储 Block 的指针，按照 (cuda_stream_id -> block size -> addr) 的优先级从小到大排序。所有保存在 BlockPool 中的 Block 都是空闲的。
• DeviceCachingAllocator中维护两种 BlockPool (large_blocks, small_blocks)，分别存放较小的块和较大的块（为了分别加速小需求和大需求），简单地将 <= 1MB 的 Block 归类为小块，> 1MB 的为大块。

直观理解 Block、BlockPool 见下图：

---

以下内容来自：PyTorch显存管理介绍与源码解析（一）和 PyTorch显存管理介绍与源码解析（二）

DeviceCachingAllocator类创建两个 BlockPool（Large_blocks & Small_blocks），将未使用的 blocks 存入其中。对于在使用的 block，通过哈希集合active_blocks存储指向 block 的指针。

目前未使用的块管理有两种类型的池子：large_blocks，small_blocks。分类阈值默认值设置为 1M，小于 1M 放入small_blocks，大于 1M 放入large_blocks。

关于active_blocks，见 4.5 class DeviceCachingAllocator 类成员变量 一节的内容。另外，在 PyTorch显存管理介绍与源码解析（二） 文中，作者提到 "显存池分了两类（active pool、remaining pool）"，这里的 active pool 指的就是active_blocks。

• active_blocks 是一个 ska::flat_hash_set<Block*>，基于哈希表 的集合。
  • 哈希表的特性是插入、删除和查找的平均时间复杂度为 O ( 1 ) O(1) O(1)，适合频繁的动态操作。
  • active_blocks目的是快速跟踪所有活跃的内存块 ，从代码来看，active_blocks主要用于 Snapshot。
• large_blocks和small_blocks是BlockPool类型的对象，内部使用std::set<Block*, Comparison>来存储和排序 Block，是一个基于红黑树 的集合。
  • 红黑树的特性是元素有序，插入、删除和查找的时间复杂度为 O ( l o g n ) O(log n) O(logn)，适合需要排序的场景。
  • large_blocks和small_blocks目的是高效管理空闲内存块。

---

以下内容来自：Pytorch 1.6 显存管理CudaCachingAllocator剖析

空闲块管理规则如下：

---

❗关于 Small Pool 和 Large Pool，笔者有一些和上述博客不同的看法，在这里根据笔者的理解做一些补充，如果有不对之处还请指出。

Small Pool 和 Large Pool 的作用是管理空闲 Block，他们本质上是红黑树集合。

PyTorch 根据请求的大小去这两个 Pool 里寻找空闲内存 Block：对于 ≤ 1 MB 的请求，首先去 Small Pool 寻找空闲 Block；对于＞1 MB 的请求，首先去 Large Pool 寻找空闲 Block。对应代码 L817-L820 和 L824-L829。

当在 Small Pool 和 Large Pool 找不到空闲 Block 时，PyTorch 会去申请新的 Block。对应代码 L839-L847。新申请的 Block 确实也有一个 Pool 的属性，这里还是对应代码 L817-L820，其中：

• ≤ 1MB 的请求对应 Small Pool，会去申请一个 2MB 的 Cache Block；
• ＞1MB 的请求对应 Large Pool，会按规则申请更大的 Cache Block。

但从alloc_block这个函数的具体实现来看，其实并没有将新创建的 Block 插入到 Small Pool 或 Large Pool。对应代码 L626-L646。

而把空闲 Block 插入 Small Pool 或 Large Pool 的只有下面两处：

1. split ：对应代码 L959，这里会把拆分后的 remaining Block 插入到与被拆分的 Block 一样的 Pool 中。并且，根据should_split函数，插入 Small Pool 的一定都是 512B ≤ size＜2MB 的空闲块，插入 Large Pool 的一定都是 size＞1MB 的空闲块。对应代码 L1546-L1554。
2. free ：对应代码 L1433，这里会把释放后的 Block 插入到对应的 Pool 中。因为可能涉及合并操作（也可能不合并），所以在这个过程中，插入 Small Pool 的一定都是 512B ≤ size ≤ 2MB 的空闲块，插入 Large Pool 的一定都是 size＞1MB 的空闲块。

这么分析来看，Small Pool 和 Large Pool 管理的空闲 Block 的大小其实是有一些交叉的。主要是 1MB＜ size ≤ 2MB 这个区间。

以下内容来自：Pytorch 显存管理机制与显存占用分析方法

这篇文章的说法比较妥当：

• 显存管理机制会依据未分配 Block 所在 Segment 的大小，将未分配的 Block 划入 large pool（Segment > 2MB）或 small pool（Segment ≤ 2MB）。
• 用户创建 tensor 申请显存时，会先从 tensor size 对应未分配显存的 pool 中查找是否有满足 size 要求的 Block，如果没有才会向 GPU 申请新的 Segment 显存块。

比如显存管理器当前有且仅有一个 2MB 的 Segment，已分配了 0.5MB，还剩 1.5MB，用户此时需要创建一个 1.1MB 的 tensor，那么显存管理器不会从这 1.5MB 的未分配 Block 中划分一部分空间给 tensor，而是额外申请一个 20MB 的 Segment 再进行分配。

---

BlockPool的代码实现：

struct BlockPool {

  std::set<Block*, Comparison> blocks;
  // 一个 std::set 容器，用于存储和排序 Block 指针。
  
  const bool is_small;
  // 标识该池是否用于小块管理。
  
  PrivatePool* owner_PrivatePool;
  // 指向关联的 PrivatePool（可选）。

  // 构造函数
  BlockPool(
      Comparison comparator,
      // 初始化时需要自定义比较函数
      bool small,
      PrivatePool* private_pool = nullptr)
      : blocks(comparator), is_small(small), owner_PrivatePool(private_pool) {}
};

关于 std::set 容器：【C++】set 容器最全解析（什么是 set? set容器的常用接口有那些？）

4.3 struct PrivatePool 私有内存池

PrivatePool是一个支持 CUDA 图（CUDA Graphs）的私有内存池，用于管理显存块。

PrivatePool的代码实现：

// CUDA graphs helper
struct PrivatePool {
  int use_count;
  // Number of live graphs using this pool
  // 记录使用该池的 CUDA 图数量。
  
  int cudaMalloc_count;
  // Number of unfreed cudaMallocs made for this pool.
  // 记录未释放的 cudaMalloc 调用计数。
  
  // When use_count and cudaMalloc_count drop to zero, we can delete this PrivatePool from graph_pools.
  // 当引用计数（use_count）和 cudaMalloc 计数都降为零时，我们可以从图内存池（graph_pools）中删除这个私有内存池（PrivatePool）。

  BlockPool large_blocks;
  BlockPool small_blocks;
  // BlockPool 类型的成员，分别管理大块和小块。
  
  // 构造函数
  PrivatePool()
      : use_count(1),
        cudaMalloc_count(0),
        // 初始化时传入比较函数。
        large_blocks(BlockComparator, /*is_small=*/false, this),
        small_blocks(BlockComparator, /*is_small=*/true, this) {}
  PrivatePool(const PrivatePool&) = delete;
  PrivatePool(PrivatePool&&) = delete;
  PrivatePool& operator=(const PrivatePool&) = delete;
};

4.4 struct AllocParams 分配内存时的参数和状态

AllocParams封装分配内存时的参数和状态，用于寻找或创建合适的内存块。

AllocParams的代码实现：

struct AllocParams {
  Block search_key;
  // 存储查找所需的信息
  
  BlockPool* pool;
  // 目标内存池。
  
  size_t alloc_size;
  // 分配大小（可能大于 size，因为需要对齐）。
  
  Block* block;
  // 指向找到或创建的内存块。
  
  // 构造函数
  AllocParams(
      int device,
      size_t size,
      cudaStream_t stream,
      BlockPool* pool,
      size_t alloc_size,
      DeviceStats& stats)
      : search_key(device, stream, size),
        pool(pool),
        alloc_size(alloc_size),
        block(nullptr),
        err(cudaSuccess) {}
  
  // 提供访问设备、流和大小的接口。
  int device() const {
    return search_key.device;
  }
  cudaStream_t stream() const {
    return search_key.stream;
  }
  size_t size() const {
    return search_key.size;
  }
};

4.5 class DeviceCachingAllocator 类成员变量

DeviceCachingAllocator的代码实现：

class DeviceCachingAllocator {
 private:
 
  BlockPool large_blocks;
  BlockPool small_blocks;
  // unallocated cached blocks larger than 1 MB
  // unallocated cached blocks 1 MB or smaller
  // 管理未分配的大内存块（> 1 MB）、小内存块（≤ 1 MB）。
  
  ska::flat_hash_set<Block*> active_blocks;
  // allocated or in use by a stream. Holds all active allocations,
  // whether they came from graph_pools or one of the BlockPools above.
  // 用于跟踪当前被分配（活跃）的内存块。
  // 无论内存块来自 large_blocks、small_blocks，还是 CUDA 图专用内存池，它都会被记录在这里。

  size_t total_allocated_memory = 0;
  // record used memory.
  // 记录当前分配的总显存大小。

  size_t allowed_memory_maximum = 0;
  // 最大允许的显存分配量。

  bool set_fraction = false;
  // 标记是否已经设置了允许的最大内存。
 
 public:
  // 默认构造函数
  DeviceCachingAllocator()
      : large_blocks(BlockComparator, /*is_small=*/false),
        // 初始化大块内存池
        small_blocks(BlockComparator, /*is_small=*/true),
        // 初始化小块内存池
        alloc_trace(new std::vector<TraceEntry>()) {
    stats.max_split_size = CachingAllocatorConfig::max_split_size();
    context_recorder_.store(nullptr);
  }
  
}

结构体和类的区别：Struct 和 Class 的区别以及使用场景_class和struct区别

5 PyTorch BFC 算法

以下内容来自：PyTorch显存管理介绍与源码解析（一）

整体的运行逻辑：

1. 查找：从显存池里面查找看是否有满足当前条件的 block，如果有直接跳入步骤3，如果没有跳入步骤2；
2. 创建 ：调用cudaMalloc创建新的 segment，如果满足 split 条件跳入步骤3，不满足直接跳入步骤5；
3. 切分：搜索/创建到满足条件的 block 根据需求进行切分，如果有剩余跳入步骤4，如果没有剩余跳入步骤5；
4. 保存：将切分后的 block 放入显存池中，执行完成后再执行步骤5；
5. 返回：将创建好的 block 数据指针返回给上层API；
6. 回收：当 block 释放时，会先去池子里面查找是否有可以与之合并的 block，如果有先进行合并，然后将合并后的 block 存入显存池中。
7. 释放 ：当 block 为一个 segment 时，可以触发cudaFree操作释放显存。

---

6 核心函数

6.1 Block* malloc 分配函数

以下内容来自：PyTorch CUDA backend，作者这种作图方式值得学习。

---

以下内容来自：Pytorch 显存管理机制与显存占用分析方法

---

malloc的代码实现：

  Block* malloc(int device, size_t orig_size, cudaStream_t stream) {
  
    // 对齐请求的内存大小。
    size_t size = round_size(orig_size);
    // 获取内存池。
    auto& pool = get_pool(size, stream);
    // 实际的分配大小。
    const size_t alloc_size = get_allocation_size(size);
    // 创建一个 AllocParams 对象，封装所有分配参数。
    AllocParams params(device, size, stream, &pool, alloc_size, stats);

    // First, try to get a block from the existing pool.
    // 1、从现有池中尝试分配
    bool block_found =
        // Search pool
        // 尝试直接从内存池中获取空闲块。
        get_free_block(params)
        // Trigger callbacks and retry search
        // 如果上一步失败，触发释放内存的回调，然后再次尝试获取空闲块。
        || (trigger_free_memory_callbacks(params) && get_free_block(params));

    // Can't reuse an existing block; try to get a new one.
    // 2、如果复用失败，尝试分配新的内存块
    if (!block_found) {
    
      // Do garbage collection if the flag is set.
      // 如果启用了垃圾回收
      if (C10_UNLIKELY(
              set_fraction &&
              CachingAllocatorConfig::garbage_collection_threshold() > 0.0)) {
        // 清理不再使用的缓存块。
        garbage_collect_cached_blocks();
      }
      
      // Attempt allocate
      // 尝试分配新内存块
      block_found = alloc_block(params, false)
          // Free enough available cached blocks to satisfy alloc and retry alloc.
          // 释放足够的可用缓存块以满足 alloc 并重试 alloc。
          || (release_available_cached_blocks(params) &&
              alloc_block(params, false))
          // Free all non-split cached blocks and retry alloc.
          // 释放所有未拆分的缓存块并重试 alloc。
          || (C10_LIKELY(captures_underway == 0) && release_cached_blocks() &&
              alloc_block(params, true));
    }
    
    // 如果多次尝试仍然失败，报告内存不足（OOM）错误
    if (!block_found) {
    // ......
    }

    // 检查 params 参数的有效性，确保：
    //     CUDA 操作成功（params.err == cudaSuccess）。
    //     内存块（params.block）和其指针（params.block->ptr）不为空。
    TORCH_INTERNAL_ASSERT(
        params.err == cudaSuccess && params.block != nullptr &&
        params.block->ptr != nullptr);
    
    // 当前要操作的内存块。
    Block* block = params.block;
    // 用于存储分割后的剩余内存块。
    Block* remaining = nullptr;
    
    // 检查 block 是否已经是分割块。
    const bool already_split = block->is_split();
    
    // 判断是否需要分割
    if (should_split(block, size)) {
    
      // 内存块分割处理
      
      // 1.保存 remaining 为原始块。
      remaining = block;
      
      // 2.创建新的 block，它使用原始块的前 size 内存地址。
      block = new Block(device, stream, size, &pool, block->ptr);
      
      // 3.block 链接到前面的块。
      block->prev = remaining->prev;
      if (block->prev) {
        block->prev->next = block;
      }
      
      // 4.block 链接到 remaining
      block->next = remaining;
      remaining->prev = block;
      
      // 5.调整 remaining 起始地址和大小。
      remaining->ptr = static_cast<char*>(remaining->ptr) + size;
      remaining->size -= size;
      
      // 6.将剩余块重新插入到内存池，通过断言确保插入成功。
      bool inserted = pool.blocks.insert(remaining).second;
      TORCH_INTERNAL_ASSERT_DEBUG_ONLY(inserted);

      // 7.更新分割块的统计信息
      // ......

    } else if (already_split) {
      // An already-split block is becoming active
      // 如果块不需要分割但已经是分割块，则直接更新统计信息，表示该块变为"活动"块。
      // ......
    }
    
    // 标记内存块为"已分配"。
    block->allocated = true;
    // 设置用户请求的大小。
    block->requested_size = orig_size;
    
    // 将分配完成的块插入到 active_blocks 集合中。
    bool inserted = active_blocks.insert(block).second;
    TORCH_INTERNAL_ASSERT_DEBUG_ONLY(inserted);
    
    // 更新统计数据
    // ......

    return block;
  }

malloc依赖的函数：

• 7.2 static size_t round_size(size_t size)
• 7.3 BlockPool& get_pool(size_t size, cudaStream_t stream)
• 7.4 static size_t get_allocation_size(size_t size)
• 7.5 bool get_free_block(AllocParams& p)
• 7.6 void garbage_collect_cached_blocks()
• 7.7 bool release_available_cached_blocks(const AllocParams& p)
• 7.8 bool release_cached_blocks()
• 7.11 bool alloc_block(AllocParams& p, bool isRetry)
• 7.13 bool is_split()
• 7.14 bool should_split(const Block* block, size_t size)

6.2 void free 释放函数

以下内容来自：PyTorch CUDA backend，作者这种作图方式值得学习。

---

free的代码实现：

  void free(Block* block) {
    
    // 标记内存块为未分配状态
    block->allocated = false;

    if (!block->stream_uses.empty()) {
      if (C10_UNLIKELY(captures_underway)) {
        // It's forbidden to cudaEventQuery an event recorded during CUDA graph
        // capture. We conservatively defer recording end-of-life events until
        // the next call to process_events() (which won't happen until no
        // captures are underway)
        needs_events_deferred_until_no_capture.push_back(block);
      } else {
        insert_events(block);
      }
    } else {
      // 如果没有活动事件，则释放内存块并返回到池中。
      free_block(block);
    }
    
    // 报告内存使用情况
    // ......
  }

free依赖的函数：7.15 void free_block(Block* block) → 调用 7.16 size_t try_merge_blocks(Block* dst, Block* src, BlockPool& pool)

7 功能函数

7.1 static bool BlockComparator 自定义比较函数

用于在BlockPool的std::set中对Block指针排序。

BlockComparator的代码实现：

static bool BlockComparator(const Block* a, const Block* b) {
  // 按 stream 指针的地址排序（升序）。
  if (a->stream != b->stream) {
    return (uintptr_t)a->stream < (uintptr_t)b->stream;
  }
  // 若流相同，按块大小排序（升序）。
  if (a->size != b->size) {
    return a->size < b->size;
  }
  // 若大小也相同，按内存地址排序（升序）。
  return (uintptr_t)a->ptr < (uintptr_t)b->ptr;
}

7.2 static size_t round_size 对内存大小 size 进行调整

对size调整到合适的大小，确保其满足块大小的最小值 或接近划分点。

roundup_power2_next_division将一个数值size调整为靠近的 2 的次幂划分点（多个划分点中的一个）。例子：

• 输入: size = 1200, divisions = 4
• 过程 :
  • power2_floor = 1024
  • 区间划分点: [1024, 1280, 1536, 1792]
  • size = 1200 在划分点 [1024, 1280) 内，向上调整到 1280。
• 输出: 1280

kMinBlockSize * ((size + kMinBlockSize - 1) / kMinBlockSize)。例子：

• 输入: size = 700, kMinBlockSize = 512
• 过程 :
  • size 大于 kMinBlockSize，但划分点不足。
  • 对齐为 512 的倍数，size = 1024。
• 输出: 1024

round_size的代码实现：

  static size_t round_size(size_t size) {
    // 如果 size 小于最小块大小 kMinBlockSize，直接返回 kMinBlockSize。512B
    if (size < kMinBlockSize) {
      return kMinBlockSize;
    } else {
      auto divisions = CachingAllocatorConfig::roundup_power2_divisions(size);
      // 调用配置函数 roundup_power2_divisions(size) 获取划分点的数量。
      if (divisions > 0 && size > (kMinBlockSize * divisions)) {
        return roundup_power2_next_division(size, divisions);
        // 如果划分点数大于 0 且 size 大于划分点范围，调用 roundup_power2_next_division 调整。
      } else {
        // 否则，将 size 向上舍入到最接近 kMinBlockSize（512字节）的倍数。
        // 用于确保分配的内存块大小是满足最小粒度要求的合适值。
        return kMinBlockSize * ((size + kMinBlockSize - 1) / kMinBlockSize);
      }
    }
  }

7.3 BlockPool& get_pool 返回对应的内存块池

get_pool的代码实现：

BlockPool& get_pool(size_t size, cudaStream_t stream) {
    // kSmallSize 1 MiB
    if (size <= kSmallSize) {
      return small_blocks;
    } else {
      return large_blocks;
    }
  }

7.4 static size_t get_allocation_size 确定实际要进行分配的内存大小

以下内容来自：PyTorch显存管理介绍与源码解析（一）

显存管理机制是根据申请 size 来决定从 GPU 创建多大的 segment，以及是否要进行切分（split）。基本执行逻辑如下图所示，

---

以下内容来自：Pytorch 1.6 显存管理CudaCachingAllocator剖析

Pytorch 向 cuda 申请显存时，并不是严格按实际所需大小申请的，而是按块大小申请的。块是一段地址连续 的显存，但块与块之间，地址不一定连续。

块计算规则：

Pytorch 按上述大小的块向 cuda 申请显存。如果直接将这些规格的块分给 Tensor 使用，将会产生大量的内部碎片 ，造成浪费。因此，Pytorch 会对这些大小的块先尝试进行切分（split）。

Pytorch 向 cuda 申请块的方式是惰性的。当只有从管理的块池子中找不到满足要求的空闲块时，才向 cuda 申请新的块。新申请的块会根据实际请求大小切分成更小的块。

文中还给了一些具体的例子，可以有助于理解。

---

get_allocation_size的代码实现：

  static size_t get_allocation_size(size_t size) {
    if (size <= kSmallSize) { // 1 MiB
      return kSmallBuffer; // 2 MiB
    } else if (size < kMinLargeAlloc) { // 10 MiB
      return kLargeBuffer; // 20 MiB
    } else {
      return kRoundLarge * ((size + kRoundLarge - 1) / kRoundLarge); // 2 MiB
    }
  }

7.5 bool get_free_block 寻找满足分配需求的空闲内存块

get_free_block尝试从指定内存池中BlockPool找到一个满足分配需求的空闲内存块Block，并将其分配给传入的参数对象AllocParams。

get_free_block的代码实现：

  bool get_free_block(AllocParams& p) { // 包含当前内存分配请求的参数
    // 获取当前操作的内存块池的引用，后续操作都基于这个内存池。
    BlockPool& pool = *p.pool;

    if (C10_UNLIKELY(
            set_fraction &&
            CachingAllocatorConfig::garbage_collection_threshold() > 0.0)) {
      // Track block reuse interval only when garbage collection is enabled.
      // 如果启用了垃圾回收（GC），每个空闲块的 gc_count 自增，用于跟踪未使用的时间间隔。
      for (auto& b : pool.blocks) {
        ++b->gc_count;
      }
    }
    
    // 从内存池 pool.blocks 中找到第一个大小不小于 p.search_key 的内存块。
    auto it = pool.blocks.lower_bound(&p.search_key);
    
    // 没有找到合适的块，或者找到的块的流 (stream) 不匹配当前请求流 p.stream()。
    if (it == pool.blocks.end() || (*it)->stream != p.stream())
      return false;
    
    // 当前请求的块大小 p.size() 小于允许的最大拆分大小，但找到的块太大（超出 CachingAllocatorConfig::max_split_size()）。
    // Do not return an oversized block for a large request
    // 目的：避免浪费内存，尽量分配更接近需求的块。
    if ((p.size() < CachingAllocatorConfig::max_split_size()) &&
        ((*it)->size >= CachingAllocatorConfig::max_split_size()))
      return false;
    
    // 如果当前请求的大小较大 (p.size() >= max_split_size())，但找到的块超出需求太多（大于 p.size() + kLargeBuffer）。
    // Allow oversized block size to be rounded up but within a limit
    // 允许超大块大小向上取整但在一定范围内。
    if ((p.size() >= CachingAllocatorConfig::max_split_size()) &&
        ((*it)->size >= p.size() + kLargeBuffer)) // 20 MiB
      return false;
    
    // 找到合适的块
    p.block = *it;
    (*it)->gc_count = 0; // Denote this block has been used
    pool.blocks.erase(it);
    return true;
  }

7.6 void garbage_collect_cached_blocks 释放未使用的缓存块

garbage_collect_cached_blocks的代码实现：

  void garbage_collect_cached_blocks() {
    // Unlike release_cached_blocks(), this does not enforce synchronization and
    // therefore should be of less overheads.
    // 避免了强制同步操作，以降低性能开销
    
    // 计算垃圾回收的阈值
    size_t gc_threshold = static_cast<size_t>(
        CachingAllocatorConfig::garbage_collection_threshold() *
        allowed_memory_maximum);

    // No need to trigger GC yet
    // 内存使用尚未超出阈值，无需回收。
    if (total_allocated_memory <= gc_threshold) {
      return;
    }
    
    // 计算需要释放的目标内存量
    const auto target_size = total_allocated_memory - gc_threshold;
    // 初始化已回收内存
    size_t gc_reclaimed = 0;

    // Calculate the total age of the free-able blocks. We'll use it later to get "avg age" threshold.
    // 计算可释放块的总年龄。稍后我们将使用它来获取"平均年龄"阈值。
    double total_age = 0.0;
    int freeable_block_count = 0;
    
    // 遍历当前的大块内存块列表，找出可以回收的块
    for (auto& b : large_blocks.blocks) {
      // 只有未拆分的大块才可回收。
      if (!b->is_split()) {
        // b->gc_count 表示该块未使用的累计时间，用作垃圾回收的优先级。
        // 将所有可回收块的 gc_count 累加，表示总的未使用时间。
        total_age += b->gc_count;    // 这里如果用真实时间呢？
        // 记录可以回收的块总数。
        ++freeable_block_count;
      }
    }
    
    // No free-able blocks?
    // 没有可回收块，直接返回
    if (freeable_block_count == 0) {
      return;
    }
    
    // 初始化为 true，表示可能会有块被释放。
    bool block_freed = true;
    
    // Repeat GC until we reach reclaim > target size.
    // 继续执行回收，直到以下条件之一成立：
    //     已回收的内存量（gc_reclaimed）达到或超过目标释放量（target_size）。
    //     当前循环中没有任何块被释放（block_freed == false）。
    //     没有剩余的可回收块（freeable_block_count == 0）。
    while (gc_reclaimed < target_size && block_freed == true &&
           freeable_block_count > 0) {
      // Free blocks exceeding this age threshold first.
      // 首先释放超过此年龄阈值的区块 -> 优先释放未使用时间最长的块。
      double age_threshold = total_age / freeable_block_count;
      
      // Stop iteration if we can no longer free a block.
      // 如果某轮循环中没有任何块被释放，退出循环，避免陷入无效操作。
      block_freed = false;

      // Free blocks of > avg age. Don't stop upon reaching the target_size,
      // we don't want this GC to be triggered frequently.
      // 即使达到了目标释放量（target_size），仍然尝试多释放一些内存，减少垃圾回收频率。
      
      // 遍历当前的内存块列表，检查每个块。
      auto it = large_blocks.blocks.begin();
      while (it != large_blocks.blocks.end()) {
        Block* block = *it;
        ++it;
        
        // 块必须未拆分，块的未使用时间需要大于等于年龄阈值。
        if (!block->is_split() && block->gc_count >= age_threshold) {
          block_freed = true;
          
          // 累计回收的内存量
          gc_reclaimed += block->size;
          
          // Decrement the age
          // 从总未使用时间中减去该块的 gc_count。
          total_age -= block->gc_count; 
          
          // One less block that can be freed
          // 减少可回收块的计数。
          freeable_block_count--; 
          
          // 释放内存块
          release_block(block);
        }
      }
    }
  }

7.7 bool release_available_cached_blocks 优先从最大的块开始释放

release_available_cached_blocks的代码实现：

  // Free one or more oversize blocks to the system allocator.  
  // But only enough to satisfy the target size.
  // 释放一个或多个超大块到系统分配器。但只足以满足目标大小。
  bool release_available_cached_blocks(const AllocParams& p) {
     
    // 检查是否允许释放超大块
    if (CachingAllocatorConfig::max_split_size() ==
        std::numeric_limits<size_t>::max())
      return false;

    // 提取目标内存池
    BlockPool& pool = *p.pool;

    // because of std::unique_ptr, block cannot be trivially copied
    // 由于std::unique_ptr，块不能简单地复制
    
    // 创建一个 Block 对象 key，用于搜索与当前需求匹配的块。
    Block key(
        p.search_key.device,
        p.search_key.stream,
        p.search_key.size,
        p.search_key.pool,
        p.search_key.ptr);
    
    // 调整目标大小 key.size，避免释放小于超大块的块。
    key.size = (key.size < CachingAllocatorConfig::max_split_size())
        ? CachingAllocatorConfig::max_split_size()
        : key.size;
    
    // 在内存池的块集合找到第一个大小大于或等于 key.size 的块。
    auto it = pool.blocks.lower_bound(&key);
    if (it == pool.blocks.end() || (*it)->stream != p.stream()) {
    
      // No single block is large enough; free multiple oversize blocks, starting with the largest
      // 没有一个块足够大；释放多个oversize块，从最大的开始
      
      if (it == pool.blocks.begin())
        return false;
      
      // 累计已释放的内存大小。
      size_t totalReleased = 0;
      
      // Back up one item.  Now on the largest block for the correct stream
      // 将迭代器退回到前一个块，确保从最大的块开始释放。
      --it; 
      
      // 在以下条件满足时继续释放：
      //     已释放的总大小 totalReleased 小于目标大小 key.size。
      //     当前块是超大块（(*it)->size >= max_split_size）。
      //     当前块的流与请求流匹配。      
      while ((totalReleased < key.size) &&
             ((*it)->size >= CachingAllocatorConfig::max_split_size()) &&
             ((*it)->stream == p.stream())) {
             
        auto cur = it;
        totalReleased += (*it)->size;
        
        // 如果当前块不是第一个块，迭代器退回到前一个块继续释放。
        if (it != pool.blocks.begin()) {
          --it;
          release_block(*cur);
        } else {
          // 如果已经到达集合起始位置，退出循环。
          release_block(*cur);
          break;
        }
      }
      
      // 如果释放的总内存仍不足以满足目标大小，返回 false。
      if (totalReleased < key.size)
        return false;
        
    } else {
      // 如果搜索结果是一个单个块且满足需求，直接调用 release_block 释放该块。
      release_block(*it);
    }
    
    // 如果成功释放了足够的内存（无论是单个块还是多个块），返回 true。
    return true;
  }

7.8 bool release_cached_blocks 释放所有未使用的内存块

release_cached_blocks的代码实现：

  bool release_cached_blocks() {
  
    // First ensure that all blocks that can't currently be allocated due to
    // outstanding events are returned to the pool.
    // 有些内存块可能被异步 CUDA 操作（如内核执行）锁定。这些块不能立即被释放，因为尚未完成的操作需要这些内存。调用 synchronize_and_free_events 确保所有与事件相关的操作完成后，这些块可以被安全释放并返回到内存池中。
    synchronize_and_free_events();

    // Free all non-split cached blocks to system allocator
    // 释放所有未拆分的缓存块
    release_blocks(large_blocks);
    release_blocks(small_blocks);

    // 遍历并释放 CUDA 图的可释放内存池
    for (auto it = graph_pools_freeable.begin();
         it != graph_pools_freeable.end();) {
      // See notifyCaptureDestroy for the strategy here.
      TORCH_INTERNAL_ASSERT(it->second->use_count == 0);
      release_blocks(it->second->small_blocks);
      release_blocks(it->second->large_blocks);
      if (it->second->cudaMalloc_count == 0) {
        auto erase_count = graph_pools.erase(it->first);
        TORCH_INTERNAL_ASSERT(erase_count == 1);
        it = graph_pools_freeable.erase(it);
      } else {
        ++it;
      }
    }

    return true;
  }

函数调用关系：release_cached_blocks → release_blocks → release_block

7.9 void release_blocks 释放所有非拆分块

release_blocks的代码实现：

  void release_blocks(BlockPool& pool) {
    // Frees all non-split blocks
    auto it = pool.blocks.begin();
    while (it != pool.blocks.end()) {
      Block* block = *it;
      ++it;
      // 非分割块，可以直接释放
      if (!block->prev && !block->next) {
        release_block(block);
      }
    }
  }

7.10 void release_block 释放指定的内存块

free函数把 block 标记为未使用 ，把 block 从active_blocks中清除，之后将 block 回收到Large_blocks /Small_blocks中。这个过程不会触发cudaFree，真正要释放掉一个 block 需要在release_block中完成。

release_block的代码实现：

  void release_block(Block* block) {
  
    // 释放该内存块的指针 block->ptr 所占用的 GPU 内存。
    C10_CUDA_CHECK(cudaFree((void*)block->ptr));
    // 减少已分配的总内存量。
    total_allocated_memory -= block->size;
    
    // 如果块属于私有内存池，释放后减少分配计数。
    auto* pool = block->pool;
    if (pool->owner_PrivatePool) {
      // The cudaFreed block belonged to a CUDA graph's PrivatePool.
      TORCH_INTERNAL_ASSERT(pool->owner_PrivatePool->cudaMalloc_count > 0);
      pool->owner_PrivatePool->cudaMalloc_count--;
    }
    
    // ...... 更新内存统计信息
    
    // 从块所属的内存池（pool->blocks）中移除该块。
    pool->blocks.erase(block);
    
    // 释放 Block 对象本身占用的内存。
    // 这是分配器的元数据，而非 GPU 内存。
    delete block;
  }

7.11 bool alloc_block 为指定的分配请求分配一个内存块

alloc_block的代码实现：

  bool alloc_block(AllocParams& p, bool isRetry) {
    // Defensively checks for preexisting CUDA error state.
    // 防御性地检查是否有未捕获的 CUDA 错误状态，确保代码运行前的环境干净。
    C10_CUDA_CHECK(cudaGetLastError());

    size_t size = p.alloc_size; // 需要分配的内存大小（以字节为单位）。
    void* ptr; // 分配后的内存指针，初始值为未定义。
    
    // 记录重试分配的次数，便于统计和调试。
    if (isRetry) {
      stats.num_alloc_retries += 1;
    }
    
    // 检查当前总分配的内存是否会超过允许的最大内存。
    if (set_fraction && // 启用内存限制机制
        total_allocated_memory + size > allowed_memory_maximum) {
      p.err = cudaErrorMemoryAllocation;
      return false;
    } else {
      // 分配 CUDA 设备内存。
      p.err = cudaMallocMaybeCapturing(&ptr, size);
      if (p.err != cudaSuccess) {
        if (p.err == cudaErrorMemoryAllocation) {
          // If this is the first attempt (!isRetry), we can forgive and clear
          // CUDA's internal error state.
          //
          // If this is the second attempt (isRetry), malloc's TORCH_CHECK_WITH
          // will take over to throw a helpful exception. The user can choose
          // to catch the exception, free some stuff in their script, and
          // attempt the allocation again. In this case, we can also forgive and
          // clear CUDA's internal error state.
          cudaGetLastError();
        } else {
          // If the error's unrelated to memory allocation, we should throw immediately.
          // 如果错误与内存分配无关，立即抛出异常。
          C10_CUDA_CHECK(p.err);
        }
        return false;
      }
    }
    
    // 如果分配的内存块属于某个 CUDA 图（CUDA Graph）的私有内存池。
    if (p.pool->owner_PrivatePool) {
      // The block is for a CUDA graph's PrivatePool.
      // 记录分配的调用次数。
      p.pool->owner_PrivatePool->cudaMalloc_count++;
    }
    
    // 记录总分配的内存量。
    total_allocated_memory += size;
    
    // 为分配的内存创建一个 Block 对象，并存储在 p.block 中。
    p.block = new Block(p.device(), p.stream(), size, p.pool, (char*)ptr);
    
    // 更新内存分配的统计信息。
    for_each_selected_stat_type(p.stat_types, [&](size_t stat_type) {
      // 分配的内存块数量。
      update_stat(stats.segment[stat_type], 1);
      // 分配的总字节数。
      update_stat(stats.reserved_bytes[stat_type], size);
    });
    // 记录过大的分配。
    if (size >= CachingAllocatorConfig::max_split_size())
      update_stat(stats.oversize_segments, 1);

    // p.block came from new, not cudaMalloc. It should not be nullptr here.
    // 验证内存块 p.block 和其指针 p.block->ptr 是否有效。
    TORCH_INTERNAL_ASSERT(p.block != nullptr && p.block->ptr != nullptr);
    return true;
  }

7.12 cudaError_t cudaMallocMaybeCapturing 一个封装的内存分配工具

cudaMallocMaybeCapturing的代码实现：

cudaError_t cudaMallocMaybeCapturing(void** p, size_t size) {
#if !defined(USE_ROCM) || ROCM_VERSION >= 50300
  if (at::cuda::currentStreamCaptureStatusMayInitCtx() ==
      at::cuda::CaptureStatus::None) {
#endif
    // 如果返回值为 None，表示未进入 CUDA 图捕获模式，执行普通内存分配。
    return C10_CUDA_ERROR_HANDLED(cudaMalloc(p, size));
#if !defined(USE_ROCM) || ROCM_VERSION >= 50300
  } else {
    // It's ok to capture cudaMallocs, as long as we never cudaFree those
    // addresses before replay.
    // Capturing cudaMalloc behaves nicely: it gives the graph new VA,
    // but is ignored (won't leakily allocate new memory) in replays.
    // 在捕获模式下，CUDA 图会为 cudaMalloc 提供新的虚拟地址（Virtual Address，VA），但在图重放时不会实际重新分配物理内存。
    at::cuda::CUDAStreamCaptureModeGuard g{cudaStreamCaptureModeRelaxed};
    return C10_CUDA_ERROR_HANDLED(cudaMalloc(p, size));
  }
#endif
}

7.13 bool is_split 判断 block 是否被拆分过

以下内容来自：一文读懂 PyTorch 显存管理机制

当 Block 被释放时，会检查其 prev、next 指针是否为空，及若非空是否正在被使用。若没有在被使用，则会使用try_merge_blocks合并相邻的 Block。++由于每次释放 Block 都会检查，因此不会出现两个相邻的空闲块，于是只须检查相邻的块是否空闲即可++。

---

is_split的代码实现：

  bool is_split() const {
    return (prev != nullptr) || (next != nullptr);
  }

7.14 bool should_split 判断给定的内存块 block 是否应该进行分割

以下内容来自：Pytorch 1.6 显存管理CudaCachingAllocator剖析

在决定了申请的块由哪个池子管理之后，Pytorch 通过should_split()来判断申请的块是否应该切分。规则如下：

• 对于 2M 的块，如果说切分后剩余的大小小于 512B，那么就不进行切分了，直接把整个块给用户好了；
• 对于大于 2M 的块，如果说切分后剩余的大小小于了 1M 了，那么就不进行切分了，直接把整个块给用户好了。

这样做的好处是可以保证 Small BlockPool 中所有空闲的块的大小至少是 512B；Large BlockPool 中所有空闲的块至少是 1M 大小。

---

should_split的代码实现：

  bool should_split(const Block* block, size_t size) {
    // 计算出如果进行分配操作后该内存块剩余的空间大小
    size_t remaining = block->size - size;
    if (block->pool->is_small) {
      return remaining >= kMinBlockSize; // 512B
    } else {
      // 要分配的内存大小 size 必须小于最大分割尺寸
      // 剩余空间大小要大于 kSmallSize
      return (size < CachingAllocatorConfig::max_split_size()) &&
          (remaining > kSmallSize); // 1 MiB
    }
  }

7.15 void free_block 将一个块移动到缓存的空闲块池中

free_block的代码实现：

  /** moves a block into a pool of cached free blocks */
  void free_block(Block* block) {
  
    // 确保 block 是可释放的。
    TORCH_INTERNAL_ASSERT(
        !block->allocated && block->event_count == 0 &&
        block->stream_uses.empty());
    
    // 获取其所属的内存池
    auto& pool = *block->pool;
    
    // 用于记录内存块合并等操作导致的非活跃拆分内存块数量和大小净变化
    int64_t net_change_inactive_split_blocks = 0;
    int64_t net_change_inactive_split_size = 0;
    
    // 数组 merge_candidates 包含了要释放的内存块 block 的前一个内存块和后一个内存块
    const std::array<Block*, 2> merge_candidates = {block->prev, block->next};
    for (Block* merge_candidate : merge_candidates) {
      // 尝试进行合并操作，并返回被合并的内存块大小
      const int64_t subsumed_size =
          try_merge_blocks(block, merge_candidate, pool);
      // 成功合并了一个内存块，更新前面定义的统计净变化变量。
      if (subsumed_size > 0) {
        net_change_inactive_split_blocks -= 1;
        net_change_inactive_split_size -= subsumed_size;
      }
    }
    
    // 从活跃块列表中移除
    active_blocks.erase(block);
    
    // 插入到空闲池中
    bool inserted = pool.blocks.insert(block).second;
    TORCH_INTERNAL_ASSERT(inserted);
    
    // 如果 block 是分裂块：增加分裂块的计数。增加分裂块的总大小。
    if (block->is_split()) {
      net_change_inactive_split_blocks += 1;
      net_change_inactive_split_size += block->size;
    }
    
    // 更新统计信息
    // ......
  }

7.16 size_t try_merge_blocks 合并先前拆分的块

以下内容来自：PyTorch显存管理介绍与源码解析（一）

当用户不需要使用某个显存块时，显存管理机制并不会直接从 GPU 设备上删除（free）该块，而是将其先回收到 BlockPool。

在 block 块释放后会触发一种块的融合机制 ，逻辑如下所示，当释放一个 Block_A 显存块时，去池子里面寻找是否有空闲的、地址与之连续的 Block，当匹配到符合条件的 Block_C 时触发了合并操作，这样 Block_A 和 Block_C 融合成了 Block_D。这种机制能够降低显存碎片问题。

---

try_merge_blocks的代码实现：

  // returns the size of the subsumed block, or 0 on failure.
  // 返回包含块的大小，失败时返回 0。
  size_t try_merge_blocks(Block* dst, Block* src, BlockPool& pool) {
  
    // 检查可合并条件，dst 是 block，src 是准备和 block 合并的块 
    if (!src || src->allocated || src->event_count > 0 ||
        !src->stream_uses.empty()) {
      return 0;
    }
    
    // 进一步确保 dst 和 src 两个内存块都是拆分块
    AT_ASSERT(dst->is_split() && src->is_split());

    if (dst->prev == src) { // [src dst] 合并 block 前的块
      // dst 内存块指向合并后内存块应该指向的内存起始位置，实现内存区域的合并效果。
      dst->ptr = src->ptr;
      
      // 构建合并后内存块在链表中的前向连接关系
      dst->prev = src->prev;
      if (dst->prev) {
        dst->prev->next = dst;
      }
    } else { // [dest src] 合并 block 后的块
      // 构建合并后内存块在链表中的后向连接关系
      dst->next = src->next;
      if (dst->next) {
        dst->next->prev = dst;
      }
    }
    
    const size_t subsumed_size = src->size;
    // 两个内存块合并后的实际内存占用情况
    dst->size += subsumed_size;
    
    // 从内存块所属的内存池中移除被合并的内存块
    auto erased = pool.blocks.erase(src);
    TORCH_INTERNAL_ASSERT_DEBUG_ONLY(erased == 1);
    
    // 释放被合并内存块 src 的资源
    delete src;
    
    // 返回被合并内存块的大小
    return subsumed_size;
  }

8 问题与挑战

以下内容来自：PyTorch显存管理介绍与源码解析（一）

当前机制下的一些问题：

• 机制的参数为静态参数，没有根据需求进行自动调整。 比如 Large_blocks 和 Small_blocks 分类的阈值、切分的阈值。
• 合并回收的触发方式单一，只有显存不足或者手工调用时才会触发回收，没有一个算法对回收时机进行分析。
• 当前机制下有不可避免的碎片问题；

---

还可以参考 深入剖析大模型训练的显存优化 提到的五种设备内存优化路径。

参考资料

说明：资料内容依据的 PyTorch 版本见标题前[]，最新版本的源码实现还请查看 PyTorch 仓库。

• unknown\] [PyTorch源码浅析(2)：THC \| NIUHE](https://www.52coding.com.cn/2019/05/05/PyTorch2/#thccachingallocator)（THCCachingAllocator 小节）

• 1.9.1\] [PyTorch显存机制分析](https://zhuanlan.zhihu.com/p/424512257)（提到 CUDA Context 开销） → 补充 \[1.10.0\] [PyTorch显存机制分析 - Angry_Panda](https://www.cnblogs.com/xyz/p/15578068.html)

• unknown\] [GPU 底层机制分析：显存分配开销](https://zhuanlan.zhihu.com/p/544879511)

• 1.13.0\] [PyTorch显存管理介绍与源码解析（一）](https://zhuanlan.zhihu.com/p/680769942) ⭐ * \[1.13.0\] [PyTorch显存管理介绍与源码解析（二）](https://zhuanlan.zhihu.com/p/681651660) ⭐

• unknown\] [Pytorch内存管理机制小记](https://mp.weixin.qq.com/s/qIgH6-UiIeb_JHt27XtcDA)（文中分析`max_spilt_size_mb`的部分写得较清晰）

• 2.3.0\] [Pytorch 显存管理机制与显存占用分析方法](https://zhuanlan.zhihu.com/p/699254132) ⭐