Torch 显存管理与生命周期

del 掉了一个大张量，nvidia-smi 上的显存却一点没降——是泄漏了，还是根本不该降？
调用 torch.cuda.empty_cache() 后显存“降下来了”，为什么下一步 decode 反而更慢，OOM 也没解决？
活跃张量只有十几 GB，memory_reserved() 却涨到几百 GB，最后报 OOM——这些显存被谁占着？

推理服务的显存不只是权重和 KV cache。中间输出、kernel workspace、allocator 缓存、CUDA Graph 的私有内存池，都可能影响峰值。上面三个问题的答案都不一样：有的是对象仍被引用，有的只是分配器缓存，有的是缓存块因跨 stream 使用而迟迟不能复用。本篇先建立一个分层模型，把“谁持有数据”“谁保留了内存”“峰值在哪里”拆开回答，再看 SGLang 如何在这个模型上规划 KV cache 容量。张量布局与共享 storage 的基础见 Tensor Memory。

抽象：显存是三层账本

一块显存从“被张量使用”到“回到 GPU”，要经过三层，每层各记一本账：

引用层（Python 对象 → storage）：张量只是 storage 的透镜；只要还有任何张量、view、列表、闭包或缓存引用这块 storage，它就是活的。del x 只是删掉一个名字。
缓存分配器层（PyTorch caching allocator）：storage 的内存来自分配器管理的 block。storage 不再被引用时，block 通常回到分配器的缓存，等下一次同 stream、尺寸合适的分配复用，而不是立刻还给驱动——因为向驱动申请和释放显存既慢，又可能引入同步。
驱动/设备层：驱动看到的是进程实际向 GPU 申请到的内存，其中除了分配器保留的部分，还有 CUDA context、通信库缓冲区以及绕开 PyTorch allocator 的直接分配。

关键的一点：每个观察指标只看到其中一层，不能拿一层的数字去回答另一层的问题。

观察项	含义	不能直接推出什么
`x.numel() * x.element_size()`	张量逻辑元素的数据字节数	不代表其独占的 storage 大小
`torch.cuda.memory_allocated()`	当前张量等活跃分配占用的 CUDA 内存	不包含所有进程级 GPU 占用
`torch.cuda.memory_reserved()`	缓存分配器管理的内存	不代表这些内存都被活跃张量使用
`nvidia-smi`	驱动视角的进程显存占用	不能单独判断 Python 张量泄漏

逻辑字节数与 storage 也对不上：一个 view 可能只暴露几个元素，却让整块大 storage 持续存活；expand() 的逻辑元素数又可能远大于实际存储。共享 storage 的多个张量也不能简单相加来估算独立占用。

指标定义见 memory_allocated 和 memory_reserved。CUDA context、部分通信库或绕开 PyTorch allocator 的直接分配，会造成驱动统计与 PyTorch 统计之间的差异。

引用层：del 删除的是引用，不是强制释放显存

import torch

big = torch.empty(1024, 1024, device="cuda")
small = big[:1]  # view 仍引用整块 storage
del big         # storage 仍需为 small 保留
small_copy = small.clone()
del small       # 不再通过这个 view 持有原 storage

同一张量放进列表、对象属性、闭包或缓存后，删除局部变量不会消除其他引用。detach() 只断开 autograd 关系，仍共享 storage；需要独立数据时才考虑 clone()，并计入复制成本。见 Tensor.detach。

当最后一个相关引用消失且异步使用安全结束后，内存通常先回到 PyTorch 缓存分配器，供下一次分配复用，而不是立即还给 CUDA 驱动。这就是第一个问题的答案：memory_allocated() 下降了，memory_reserved() 和 nvidia-smi 不降是正常现象。

分配器层：empty_cache 不是清理活跃张量

torch.cuda.empty_cache() 只作用于第二层：把分配器缓存中当前未使用的 block 还给驱动。它不能释放仍被张量持有的内存，也不能修复引用泄漏。每个 decode step 调用它，会损失缓存复用的价值并增加分配开销；不应作为常规优化手段。

它可能在特定碎片化场景或让其他进程使用空闲显存时有帮助，但不能用“调用后显存下降”证明活跃数据变少。见 empty_cache。

在 SGLang 中，它出现在“生命周期边界”而不是热路径上：例如 Scheduler.flush_cache(empty_cache=True)（managers/scheduler.py）先重置 radix cache、req_to_token_pool 和 KV allocator 等业务缓存，再调用 current_platform.empty_cache() 归还分配器缓存；模型加载和权重导出路径也类似。utils/common.py 的 empty_device_cache() 注释写得很直白：它不清理 SGLang 的 KV/radix/请求缓存，也不释放活跃张量，只转发到后端 allocator。

峰值：每步拼接 vs 预分配

显存问题经常不是“稳态占用太高”，而是“某一瞬间的峰值”。考虑每步拼接 KV：

# 示意：old_k 为 [batch, head, sequence, dim]
k = torch.cat([old_k, new_k], dim=2)

这个实现容易理解，但每步都构造新结果并复制旧内容；在 cat 执行期间，旧、新存储同时存活，瞬时占用接近两份 KV。若每次增加一个 token，累计复制规模随序列长度 \(T\) 近似增长为 \(O(T^2)\)。这与 attention 本身的计算复杂度是不同的成本。尺寸不断变化的分配还会让分配器缓存里留下大量大小不一的 block，增加碎片。

预分配容量后，只写入新位置，可以减少反复分配和历史数据搬运：

# 示意：capacity、position、new_k 的 shape 均由调用方保证正确
k_buffer[:, :, position:position + 1, :].copy_(new_k)

预分配并不消除内存代价：它用较稳定的常驻容量换取更少的运行时分配。多请求服务还需要容量规划、slot 管理和释放策略，见 Memory Pool。

此外，torch.empty() 创建的缓冲区内容未初始化；任何会被读取的位置都必须先写入。padding、graph bucket 和部分有效的 KV buffer 尤其需要确认有效区域。见 torch.empty。

跨 stream 生命周期：依赖与复用分别处理

分配器的 block 是按 stream 归属管理的：一个张量在 stream A 上分配，释放后，分配器只保证stream A 上之后排队的工作不会和旧数据冲突，所以可以立刻把这块 block 交给 stream A 上的下一次分配。它并不知道 stream B 还在读这块内存。

因此在侧 stream 使用一个 CUDA 张量，要解决两个问题：消费者何时可以读取，以及 allocator 何时可以复用存储。

wait_stream() 或 Event 处理 producer/consumer 的执行顺序。
record_stream() 通知 allocator：该张量的存储还在指定 stream 上使用，不能仅根据原 stream 的进度就复用。

record_stream() 不替代数据依赖；显式维护正确的 stream 同步也可以管理生命周期，但要求调用方更仔细地安排释放。见 Tensor.record_stream。

record_stream() 的代价是延迟释放：block 要等到被记录 stream 上已提交的工作完成后才能回到可复用状态。SGLang 的 DP attention two-batch-overlap 路径恰好踩到过这一点：dp_gather_partial_async()（layers/dp_attention.py）在通信 stream 上做 gather，并对输入输出调用 record_stream(comm)。紧随其后的 _TBO_PERSIST_BUF 注释记录了一次真实故障——每层、每个 micro-batch 都新 torch.empty 一块尺寸各异的 gather/combine buffer，又被通信 stream 的 record_stream 拖住，结果活跃工作集只有约 10GB，reserved 却涨到约 270GB，在 ROCm 上触发资源耗尽。修复方式是改用按 key 复用、只增不减的持久 buffer（get_tbo_persistent_buffer()）。layers/moe/token_dispatcher/moriep.py 的 dispatch 路径则出于同样原因刻意不调用 record_stream，改由同一调用内的 wait_event 保证安全。这正是第三个问题的典型答案：allocated 小、reserved 大，问题出在分配器层，而不在引用层。

结合 SGLang：KV cache 容量从哪里来

推理引擎会主动利用上面的分层模型来规划显存。SGLang 的思路是：先加载权重，再测“驱动层还剩多少”，扣掉给运行时临时分配预留的余量，剩下的一次性分配给 KV cache。

get_available_gpu_memory()（utils/common.py）在 CUDA 上先调用 empty_device_cache()，再用 torch.cuda.mem_get_info() 读驱动视角的空闲显存。先清缓存，是为了让分配器缓存里空闲的 block 也被算作“可用”；多卡时取所有 rank 的最小值。
KVCacheConfigurator._profile_available_bytes()（mem_cache/kv_cache_configurator.py）用这个空闲值减去 pre_model_load_memory * (1 - mem_fraction_static) 作为运行时余量（激活、workspace、CUDA Graph 等），剩余部分作为 KV pool 预算；若已加载的权重超出静态预算，会报错并给出建议的 --mem-fraction-static。
预算确定后，MHATokenToKVPool._create_buffers_normal()（mem_cache/memory_pool.py）按层一次性 torch.zeros 出 k_buffer / v_buffer；之后每步 forward 只经 set_kv_buffer() 按槽位 loc 写入，不再为 KV 增长做分配。

所以 --mem-fraction-static 调的并不是“PyTorch 能用多少显存”，而是常驻 KV 容量与运行时临时峰值之间的分界：设得过高，prefill 的临时 attention score、通信 buffer 或 graph capture 可能在余量里放不下而 OOM；设得过低，KV 槽位少，并发和 prefix 复用受限。槽位如何分配、回收与共享见 Memory Pool，CUDA Graph 的私有内存池见 Cuda Graph。

定位 OOM 要观察运行阶段与峰值

torch.cuda.synchronize()
torch.cuda.reset_peak_memory_stats()
# 执行需要测量的一段推理工作
torch.cuda.synchronize()
allocated = torch.cuda.memory_allocated()
reserved = torch.cuda.memory_reserved()
peak = torch.cuda.max_memory_allocated()

这里 peak 是重置后活跃分配的最高值，包含测量开始时已存在的分配，并不是本次 forward 的增量。

分别测量模型加载、预热、graph capture、prefill 和 decode。只看 forward 完成后的占用，可能错过临时 attention score 或 workspace 的峰值；只看 reserved 又可能误判泄漏。按三层账本排查：

allocated 持续增长 → 引用层：查谁还持有张量（列表、缓存、闭包、跨请求复用的对象）。
allocated 稳定但 reserved 远大于它 → 分配器层：查尺寸多变的临时分配、碎片和 record_stream 造成的延迟释放。
reserved 稳定但 nvidia-smi 更高 → 驱动层：查 CUDA context、通信库和绕开 allocator 的分配。

更深入的问题可使用 CUDA memory snapshot，但它只能看到 PyTorch allocator 能追踪的分配。

小结

把显存看成三层账本——引用层决定数据是否还活着，分配器层决定内存是否被保留以便复用，驱动层决定 GPU 上实际被占了多少——就能把“谁仍持有数据”“分配器保留了多少”“执行期间峰值在哪里”分开回答。del 只动第一层，empty_cache() 只动第二层，record_stream() 在两层之间加了一道延迟。SGLang 这类推理引擎的做法是在启动时测出驱动层余量、一次性预分配 KV、热路径只写不分配，从而把峰值限定在可预测的运行时余量里；稳定的 KV 容量和受控的临时分配，通常比反复清缓存更有助于推理性能和容量预测。