深度解读：CUDA 缓存分配器中的内存碎片化问题

背景

在理想的抽象模型中，PyTorch 用户通常将 CUDA 内存管理视为一个简单的池：GPU 拥有固定数量的显存，每次分配时可用内存减少，释放时可用内存增加。然而，CUDA 缓存分配器（CUDA Caching Allocator）的内部实现机制更为复杂，特定的分配模式会导致内存碎片化（Fragmentation）。

这种现象表现为：虽然“技术上”存在足够的空闲空间来存储请求的分配，但 CUDA 缓存分配器无法实际服务该请求。这在现代 AI 应用场景中尤为常见，特别是大语言模型（LLM）推理服务。在这些场景中，用户希望尽可能利用 GPU 显存，同时避免内存溢出（OOM）。然而，由于分配顺序的影响，PyTorch 可能会保留比用户预期更多的内存，导致资源浪费甚至服务失败。

核心内容

1. 内存管理的基本结构：Segment 与 Block

CUDA 缓存分配器在两个层级上组织 GPU 显存：

• Segment（段）：从 CUDA 获取的连续内存区域（通过 cudaMalloc 或虚拟内存映射获得）。
• Block（块）：Segment 内部的子区域，用于跟踪单个分配。

分配与释放逻辑：

• 分配：当请求到来时，分配器寻找足够大的空闲 Block。如果 Block 比需求大，它会进行拆分（Split）：前部分用于服务分配，后部分成为新的空闲 Block。
• 释放：当 Block 被释放时，分配器尝试将其与直接相邻的 Block 合并（Merge），但前提是邻居也必须是空闲的。如果两个空闲 Block 之间隔着一个已分配的 Block，它们无法合并。

2. Expandable Segments 开启前：碎片化的根源

在未启用 expandable_segments（即 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False）的情况下，每个 cudaMalloc 调用都会创建一个独立的 Segment。

• 大小规则：
  • 对于 1 MiB 到 10 MiB 的分配，分配器请求一个 20 MiB 的 Segment。
  • 对于 $\ge$ 10 MiB 的分配，Segment 大小向上舍入到最接近的 2 MiB 倍数。
• 关键约束：不同 Segment 中的 Block 永远无法合并。每个 cudaMalloc 从 CUDA 的角度来看都是独立的分配。

场景演示：分配顺序的影响

为了说明问题，我们对比两种分配顺序（假设显存总量充足，仅关注碎片化逻辑）：

场景 A：小分配在前，大分配在后（坏顺序）

1. 分配 8 个 16 MiB：
   • 每个 16 MiB 请求触发一次独立的 cudaMalloc。
   • 由于 16 MiB $\ge$ 10 MiB，每个 Segment 被舍入为 16 MiB。
   • 结果：8 个独立的 16 MiB Segment，每个包含一个已分配 Block。保留内存：128 MiB。
2. 释放所有：
   • 每个 Segment 现在有一个 16 MiB 的空闲 Block。
   • 但由于它们是独立的 cudaMalloc 分配，彼此无法合并。池子中仍保留 128 MiB 内存，分布在 8 个独立 Segment 中。
3. 分配 4 个 32 MiB：
   • 分配器寻找 $\ge$ 32 MiB 的空闲 Block。
   • 现有的空闲 Block 最大仅为 16 MiB，且不能跨 Segment 合并。
   • 现有 Segment 均无法满足请求。
   • 分配器必须调用 cudaMalloc 4 次，每次分配 32 MiB。
   • 结果：保留内存激增至 256 MiB（8 个闲置的 16 MiB Segment + 4 个新的 32 MiB Segment）。

场景 B：大分配在前，小分配在后（好顺序）

1. 分配 4 个 32 MiB：
   • 4 次 cudaMalloc，4 个 32 MiB Segment。保留内存：128 MiB。
2. 释放所有：
   • 每个 Segment 有一个 32 MiB 的空闲 Block。保留内存：128 MiB。
3. 分配 8 个 16 MiB：
   • 第一个 16 MiB 请求找到一个 32 MiB 的空闲 Block。
   • 分配器将其拆分：16 MiB 用于分配，剩余 16 MiB 成为新的空闲 Block。
   • 第二个 16 MiB 请求直接使用该剩余 Block。
   • 两个分配共用一个 Segment，无需新的 cudaMalloc。
   • 此过程重复四次。
   • 结果：保留内存仍为 128 MiB，每个 Segment 被拆分为两个 16 MiB Block。

结论：同样的总工作量，场景 B 使用的内存仅为场景 A 的一半。经典的变通方法是始终按递减的批次大小顺序记录 CUDA 图（Large allocations establish segments, smaller ones split within them）。但这是一种“泄漏的抽象”，要求用户深入了解底层实现。

3. Expandable Segments 开启后：解决方案

当启用 expandable_segments（即 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True）时，行为发生根本性变化。

• 机制：利用 CUDA 的虚拟内存管理 API（cuMemMap），分配器不再为每个 Segment 调用独立的 cudaMalloc。
• 优势：这使得分配器可以在更大的地址空间内灵活管理内存，解决了不同 Segment 间无法合并的问题。虽然原文截断，但根据上下文，这意味着分配顺序对碎片化的影响大幅降低，系统能更有效地利用显存，特别是在 LLM 推理中需要为不同批次大小共享 CUDA 图缓冲区时。

关键要点

• 碎片化本质：即使总空闲显存足够，由于内存块不连续或无法跨 Segment 合并，分配器仍可能拒绝请求。
• Segment 隔离：在未启用 Expandable Segments 时，每个 cudaMalloc 创建的 Segment 是孤立的，块之间无法合并。
• 分配顺序至关重要：
  • 先小后大：导致大量小 Segment 闲置，大分配被迫创建新 Segment，造成内存保留量翻倍。
  • 先大后小：大 Segment 可被拆分服务小分配，内存利用率高。
• LLM 服务的挑战：LLM 推理中，KV Cache 和 CUDA Graph 缓冲区对显存极度敏感。不同批次大小的 CUDA 图需要共享内存池，而记录图时的分配顺序若不当，会导致严重的内存浪费。
• 解决方案：启用 expandable_segments:True 是解决此类碎片化问题的关键配置，它通过虚拟内存映射打破了传统 cudaMalloc 的隔离限制。

意义与影响

1. 性能优化：对于 LLM 推理服务，显存效率直接决定了吞吐量（Throughput）和并发能力。理解并避免碎片化可以显著减少显存浪费，允许部署更大的模型或支持更高的并发请求。
2. 开发实践：开发者在编写自定义内存管理逻辑或记录 CUDA 图时，应意识到分配顺序的影响。虽然启用 expandable_segments 是更通用的解决方案，但在旧版本或特定配置下，遵循“先大后小”的分配策略仍是有效的优化手段。
3. 抽象层的泄漏：PyTorch 试图隐藏底层内存管理的复杂性，但 CUDA 缓存分配器的行为揭示了这种抽象的局限性。用户需要了解底层机制才能在资源受限的环境中实现最优性能。
4. **技术