↓ 跳过正文
111qqz的小窝
111qqz的小窝
  1. Posts/

从线上模型切换问题理解 PyTorch CUDA Caching Allocator：为什么显存管理需要理解生命周期

·6437 字·13 分钟
Deep-Learning CUDA PyTorch Gpu-Memory Memory-Allocator Inference Ml-Infra
目录

起因
#

我们的线上推理系统中存在大量模型——MLP、GNN、MoE，以及它们的各种组合。这个 workload 有几个重要特点：

  • 单张 GPU 上可能同时存在多个模型；
  • 很多模型的 architecture 相同或者高度相似；
  • 模型会随着业务时间轴不断上线、下线和切换；
  • 模型切换过程对 latency 比较敏感。

随着模型数量增加以及动态上下线越来越频繁，我们逐渐观察到两个问题：GPU memory fragmentation，以及 model loading / switching latency。

尤其是模型加载过程中，大量 weight tensor 都需要建立对应的 CUDA storage。直觉上可能首先怀疑：是不是频繁 cudaMalloc / cudaFree 导致的？

但很快会发现——PyTorch 本身早就有 CUDA Caching Allocator。也就是说，tensor allocation 并不等于每次 cudaMalloc。PyTorch 已经在帮助我们缓存和复用 GPU memory。

于是问题变成了：

既然 PyTorch 已经使用 Caching Allocator，为什么在频繁模型切换场景下，我们仍然会遇到 fragmentation 和 allocation latency？

为什么需要 Caching Allocator
#

先建立一个基本事实。

cudaMalloc 和 cudaFree 不是轻量操作。根据 NVIDIA 官方文档：

Managing memory allocations using cudaMalloc and cudaFree causes the GPU to synchronize across all executing CUDA streams.

也就是说，cudaMalloc / cudaFree 涉及 CUDA Runtime / Driver 层面的交互，可能导致 device-wide synchronization，带来较高的 latency 和 latency jitter。

如果深度学习框架真的按照这样的方式运行：

1Tensor 创建 → cudaMalloc
2Tensor 销毁 → cudaFree

那么训练和推理中大量临时 tensor 会产生很高的 allocator overhead。

因此 PyTorch 在 tensor 与 CUDA Runtime 之间增加了一层 CUDACachingAllocator。核心思想是：

从 CUDA Driver 获取较大块显存，然后由 PyTorch 自己管理、切分、缓存和复用。

用一个简单类比：PyTorch 从 CUDA Driver 批发显存，然后自己负责零售。

PyTorch CUDA Caching Allocator
#

以下重点讨论 allocator 长期稳定的设计思想，而不是依赖某一个版本的内部常数。

Cache 的基本思想
#

第一次 allocation：

1Tensor 申请 memory
2    ↓
3CUDACachingAllocator
4    ↓
5没有合适的 cached block
6    ↓
7向 CUDA 获取 segment

Tensor 生命周期结束后：

1Tensor 释放
2    ↓
3block 返回 allocator
4    ↓
5标记为 cached / free

下一次有新的 tensor 申请 memory：

1new Tensor 申请 memory
2    ↓
3在 cache 中找到合适的 free block
4    ↓
5直接复用

而不是重新执行 cudaMalloc。

整个过程用流程图表示：

caching-allocator-flow

Caching Allocator 的核心收益是：把很多 CUDA-level allocation 转化为 allocator 内部的 block bookkeeping 和 reuse。

Segment 与 Block
#

Allocator 的基本抽象可以分成两层。

Segment 是 allocator 从 CUDA 获取的一大块连续 memory。可以通过 cudaMalloc，也可以通过 CUDA virtual memory API（cuMemAddressReserve + cuMemCreate + cuMemMap，即 expandable segments 机制）获取。

Block 是 segment 内部的子区域，每个 block 对应一次 tensor allocation。

1Segment
2┌───────────────────────────────────────────┐
3│ Block A (allocated) │ Block B (allocated) │ Free Block │
4└───────────────────────────────────────────┘

当一个 tensor 申请 memory 时，allocator 在 free block 中找到合适的。如果 free block 比请求大，就 split——前半部分分配出去，后半部分变成新的 free block：

1申请 80MB，现有 200MB free block：
2┌───────────┬──────────────────┐
3│ 80 MB     │ 120 MB           │
4│ allocated │ cached / free    │
5└───────────┴──────────────────┘

当相邻的 block 都变成 free 时，allocator 会尝试 merge / coalesce——把它们合并回一个更大的 free block。

PyTorch 还按 size 把 block 分成两个 pool：small_blocks（<= 1 MiB）和 large_blocks（> 1 MiB）。两个 pool 完全独立，各自有各自的 segment 和 free list。

allocated、reserved 与 empty_cache
#

PyTorch 中两个很容易混淆的指标：

1torch.cuda.memory_allocated()   # 当前 live tensors 使用的 memory
2torch.cuda.memory_reserved()    # PyTorch 从 CUDA 获取的、目前仍管理的 memory

概念上 allocated 是 reserved 的子集：

1Reserved（PyTorch 从 CUDA 获取的总量）
2┌──────────────────────────────────┐
3│ Allocated（当前 live tensors）    │
4│                                  │
5│ Cached Free（已释放但未还给 CUDA）│
6└──────────────────────────────────┘

之前的文章讨论过这个区别。这里再强调一次：

Tensor 生命周期结束，不代表显存立即返回 CUDA Driver。

所以 del tensor 以后，nvidia-smi 显示的显存不一定下降。

torch.cuda.empty_cache() 释放的是 allocator 当前没有使用的 cached memory。它不会释放正在被 tensor 使用的 memory，也不应该在正常 inference / training loop 中频繁调用——那样反而会破坏 caching 带来的性能收益。

Caching 没有消灭 Fragmentation
#

CachingAllocator 解决了频繁 cudaMalloc / cudaFree 的问题。但它带来了另一个经典 allocator 问题：fragmentation。

这里需要区分两种 fragmentation。

Internal fragmentation：分配出去的 block 比实际请求大，多出的部分被浪费。PyTorch 通过 size rounding 和 split 来控制这个问题。

External fragmentation：free memory 总量足够，但被分散成很多不连续的小 block，无法满足一个较大的 contiguous allocation。

重点讨论 external fragmentation。

假设一个 segment 里有这样的 block 布局：

1┌────┬────┬────┬────┬────┬────┬────┐
2│ A  │ B  │ C  │ D  │ E  │ F  │ G  │
3└────┴────┴────┴────┴────┴────┴────┘

其中各个 block 的 lifetime 完全不同。一段时间后，B、D、F 释放了：

1┌────┬──────┬────┬──────┬────┬──────┬────┐
2│ A  │ FREE │ C  │ FREE │ E  │ FREE │ G  │
3└────┴──────┴────┴──────┴────┴──────┴────┘

虽然总 free memory 不少（free1 + free2 + free3），但如果需要一个更大的 contiguous block——大于任何一个单独的 free hole——仍然无法直接分配。

这就是 external fragmentation。

需要特别注意：fragmentation 不等于 memory leak。 Leaked memory 是应用逻辑上不再需要但没有被释放的 memory。Fragmentation 是 free memory 确实存在，但布局不够 contiguous。两者完全不同。

Allocation Lifetime
#

大多数关于 allocator 的讨论会集中在 block size 上——best fit、size rounding、split threshold。但我们的线上 workload 暴露出了一个更重要的问题：

不同生命周期的 allocation 被混合放进了同一个 memory arena。

在在线推理中，GPU memory 中的对象有非常不同的 lifetime。

Long-lived 对象：

1model weights
2persistent buffers

可能存活数分钟、数小时，甚至整个服务时段。

Short-lived 对象：

1activation
2temporary tensor
3operator workspace

可能只存活几十微秒、几毫秒、一次 inference 的时间。

如果把这些东西全部交给同一个 allocator，segment 内部的 block 分布可能变成：

1┌────────┬────────────┬────────┬──────┬────────┬───────────┬────────┐
2│ Weight │ Activation │ Weight │ Temp │ Weight │ Workspace │ Weight │
3└────────┴────────────┴────────┴──────┴────────┴───────────┴────────┘

随着 short-lived allocation 不断创建和销毁：

1┌────────┬────────────┬────────┬──────┬────────┬───────────┬────────┐
2│ Weight │    FREE    │ Weight │ FREE │ Weight │   FREE    │ Weight │
3└────────┴────────────┴────────┴──────┴────────┴───────────┴────────┘

Long-lived weight 就像一根根"柱子"，把整个 address space 分割成很多 hole。

Free memory 看起来不少，但每个 hole 都不大。下一次需要分配一个较大的 contiguous block 时，这些 hole 用不上。

PyTorch 2026 年 6 月的 CUDACachingAllocator DevLog 中有一句话恰好描述了这个问题：

But if your use case has long-lived allocations interleaved with short-lived ones in the same pool, expandable segments won’t save you from fragmentation within the segment.

也就是说，即使使用了 expandable segments——一种通过 virtual memory mapping 让同一 segment 内的 block 总是可以 merge 的机制——如果 long-lived 和 short-lived allocation 交错存在，fragmentation 仍然无法避免。

因为 merge 的前提是相邻 block 都是 free。只要中间有一个 long-lived block 还活着，两边的 free block 就无法合并。

allocation-lifetime-gantt

Weight 从头活到尾。Activation 和 temporary tensor 来了又走。如果它们混合分配在同一段 memory 里，short-lived 对象释放后留下的 hole 永远被 long-lived 对象隔开。

生命周期相近的对象应该放在一起
#

这是经典 allocator / arena design 里的一个基本原则：

Objects with similar lifetimes should be allocated together.

生命周期相近的对象应该尽量放在同一个 allocator domain / arena / pool 中。

目的不是单纯让 allocation 更快，而是：

1相近时间创建 + 相近时间释放
2        ↓
3memory region 更容易整体回收
4        ↓
5减少长期 fragmentation

当一次 inference 结束时，如果所有 short-lived 对象都在同一个 arena 里，大量 block 可以一起变成 free，甚至整个 arena 可以被整体回收。而 weight arena 的布局保持稳定，不会被 short-lived allocation 的创建和销毁搅乱。

arena-lifetime-separation

很多 allocator 优化会首先关注 size class、best fit、rounding、split、merge。这些都是空间维度上的优化。

但我们的线上 workload 暴露出了另一个更重要的维度：时间维度，也就是 allocation lifetime。

两个 tensor 即使 size 完全相同——都是 1MB——但如果一个的 lifetime 是 3 小时，另一个的 lifetime 是 200 微秒，从 allocator 角度，它们并不是同一种对象。

也就是说：

Memory allocation 不只有 size 这个维度，还有 lifetime。

甚至可以说：

Fragmentation 很大程度上来源于不同生命周期对象的混合，而不仅仅是 block size 不匹配。

Fragmentation 本质上不仅是空间布局问题，也是时间行为问题。一个 allocator 如果只按 size 做 best-fit，但完全不考虑 allocation 会存活多久，就没有利用到 workload 中非常强的时间结构信息。

回到线上模型：Lifetime Separation
#

把前面的原理映射回我们的 workload。

在线推理中，GPU memory 中的对象恰好具有非常清晰的 lifetime 分类：

1Model Weight = long-lived
2Activation / Workspace / Temporary = short-lived

所以第一步不是发明一个更复杂的 block search algorithm，而是：首先把完全不同生命周期的 allocation 分开。

 1┌───────────────────────────┐
 2│ Weight Memory Pool        │
 3│                           │
 4│ long-lived                │
 5│ stable layout             │
 6└───────────────────────────┘
 7
 8┌───────────────────────────┐
 9│ Runtime Memory Pool       │
10│                           │
11│ activation                │
12│ workspace                 │
13│ temporary tensor          │
14└───────────────────────────┘

这个 separation 本身就能减少 short-lived allocation 对 long-lived weight layout 的"污染"。

Weight Pool 内利用 Size Repetition
#

完成 lifetime separation 后，可以进一步观察 weight workload 的特点。

模型 architecture 决定了参数的 shape。如果大量模型属于相同或者相似结构：

1Model A:  weight1 = [256, 128]   weight2 = [128, 64]   weight3 = [64, 32]
2Model B:  weight1 = [256, 128]   weight2 = [128, 64]   weight3 = [64, 32]
3Model C:  weight1 = [256, 128]   weight2 = [128, 64]   weight3 = [64, 32]

那么 tensor allocation 的 size 并不是随机的，而是高度重复的。

weight-pool-size-repetition

在大量同构模型中，相同的 allocation size 不断重复出现。可以针对性地建立 size 到 reusable blocks 的映射，进一步提高复用率。

这里要强调：这是建立在 lifetime separation 之后的第二层优化。 顺序是先按 lifetime 分开，然后在 Weight Pool 内利用 size repetition——而不是反过来。

不是重新发明 Caching
#

这里容易混淆。

PyTorch CachingAllocator 本身当然会 reuse cached block，也会做 size rounding、best-fit matching、split、merge。所以我们的优化不是"PyTorch 不 cache，所以我们 cache"，也不是"PyTorch 不会复用相同 size，所以我们会"。

真正的区别是：PyTorch 面向 generic tensor allocation。而我们知道：

  • 这个 allocation 是 weight
  • 它通常是 long-lived
  • 这个 size 会频繁再次出现
  • 它属于某一类 model architecture

我们引入的是 workload semantics，而不是重新发明 caching。

Batch Suballocation 与 Arena
#

对于大量 weight tensor，还有一种进一步的优化方式：不要让所有 tensor 都成为完全独立的底层 allocation。

假设一个模型有 w1、w2、w3 …… wn 个 weight tensor，可以先规划 total size，然后一次获得一个 large segment：

1┌──────────────────────────────────────┐
2│ w1 │ w2 │ w3 │ w4 │ ... │ wn        │
3└──────────────────────────────────────┘

每个 tensor 使用 base_ptr + offset 形成 suballocation。

这就是 arena allocator 的思想。当对象属于同一个模型，并且生命周期相近时，把它们放在同一个 arena 里，本身就比把它们当成完全无关的 allocation 更符合 workload 的实际行为。

这里再次联系到 lifetime locality：同一个模型的 weight tensor 几乎总是同时创建、同时存在、同时销毁。

从 Tensor Reuse 到 Model Layout Reuse
#

完成 Weight Pool 后，还有一个更强的业务先验。

很多时候模型切换是这样的：

1Model A 下线
2Model B 上线
3
4architecture(A) == architecture(B)

也就是说：tensor shapes 完全相同，tensor sizes 完全相同，memory layout 可以完全相同。

如果只站在 allocator 角度，这是一个"free blocks → 重新 allocate blocks"的过程。但如果站在 model runtime 角度：

为什么还需要重新做一次 allocation planning？

可以进一步做到：Model A 的 memory layout 直接保留，Model B 直接复用。

model-layout-reuse

这里的核心不是单纯的 block reuse，而是 layout reuse。

Block caching 只知道"这块内存目前空闲"；model-level reuse 还知道"这套内存布局恰好适合下一个模型"。

这是从 Memory Allocator 走向 Memory Planner 的重要一步。

Lazy Reclaim 与 TTL
#

模型 offline 后，也不一定马上把 Weight Pool 中的 segment 真正释放。可以这样做：

1model offline
2    ↓
3mark reusable
4    ↓
5设置 TTL

如果很快出现 architecture compatible 的新模型，直接 reuse。否则 TTL 到期后再 reclaim。

这利用的是 temporal locality：最近刚刚使用过的 model memory layout，在不久以后可能再次有价值。类似 connection pool、object cache、CPU cache 的思想。

这些复用依赖什么信息
#

optimization-levels-flow

PyTorch 的通用 allocator 能缓存空闲 block，但不知道某个 weight 会跟着哪一个模型下线。线上 runtime 多知道了模型的生命周期、结构和 weight size 分布，才有机会按模型组织内存，并复用相同布局。

这也限制了这些方法的适用范围。模型结构变化大、生命周期不稳定时，整模型布局复用的机会就少；pool 留得太久，也会增加常驻显存。不能只看省了多少次 allocation，还要看模型切换延迟和高峰时的实际占用。

回到模型切换
#

最初的问题是碎片和加载延迟。Caching allocator 已经减少了底层分配次数，但上层仍然可以利用 weight 的长期存活和模型结构重复，减少布局重建。

这部分优化依赖 runtime 已经掌握的模型信息。 对另一套 workload，是否值得单独管理 weight pool，需要重新看生命周期分布和模型切换行为。

补充：Allocator 配置与静态规划
#

PyTorch allocator tuning 与 CUDA Graph、compiler memory planning

PyTorch Allocator Tuning
#

PyTorch 也提供了一些 allocator tuning 机制。通过环境变量 PYTORCH_CUDA_ALLOC_CONF（或其统一名称 PYTORCH_ALLOC_CONF）可以配置的参数包括：

  • max_split_size_mb：阻止 allocator split 超过指定大小的 block，可以在某些 workload 下减少 fragmentation
  • expandable_segments：使用 CUDA virtual memory API 让 segment 可以动态增长，同一 segment 内的 block 总是可以 merge
  • roundup_power2_divisions：控制 size rounding 策略
  • garbage_collection_threshold：控制何时触发 GC 回收

这些机制主要是在 block、segment、size、reuse policy 层面改善通用 allocator 的行为。

而 weight lifetime、model identity、architecture equivalence、model replacement 属于更高层的 runtime semantics。

调 allocator 参数和做 workload-specific memory management 并不矛盾。前者优化通用 allocator 的行为，后者引入 workload 层面的语义信息。

联系 CUDA Graph 和 Compiler Memory Planning
#

Tensor shape 和 lifetime 可预测时，也可以提前规划部分内存。

CUDA Graph 要求执行过程中的 memory address 保持稳定。TensorRT 在 build 阶段会做 memory planning，预先规划每一层的 tensor 使用哪一段 memory。XLA 和 AOTInductor 也有类似的 static memory planning pass。

这些技术细节不展开，但它们说明了一件事：我们在线上 inference 中做的 lifetime-aware memory management 不是孤立的技巧。它属于更大的趋势：

1dynamic allocator
2    ↓
3lifetime-aware allocator
4    ↓
5memory planner
6    ↓
7compile / runtime co-design

越能够在 allocation 发生之前就预测 tensor 的 shape、lifetime 和使用模式，就越有机会用 planning 替代 dynamic allocation。

参考资料
#

相关文章