PyTorch pr(2) cudaGraph下多流显存复用

Mon, 05 Jan 2026 23:11:11 +0800

[CUDA] Reuse blocks with record_stream during CUDA Graph capture in the CUDACachingAllocator #158352

https://github.com/pytorch/pytorch/pull/158352

专栏另一篇文章解读了record_stream相关内容，但CUDACachingAllocator源码中存在一个分支，会根据CUDAAllocatorConfig::graph_capture_record_stream_reuse()判断是否free和malloc某个block，这里填个坑，找一下对应的pr看是怎么复用的。

背景：CUDA Graph 下显存管理限制

在 CUDA Graph 捕获过程中，CUDACachingAllocator对内存的管理存在一个关键限制：必须等到捕获结束后才能回收内存块。这一限制源于 CUDA 的底层约束 —— 捕获阶段不允许查询event状态（因为此时 CUDA 操作尚未执行），而CUDACachingAllocator依赖事件驱动逻辑判断内存块是否可回收。而这会导致本该被释放的block未能被正确释放，直到捕获完全结束，比如free中的代码片段：

    if (!block->stream_uses.empty()) {
      if (C10_UNLIKELY(!captures_underway.empty())) {
        // 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);

cudaGraph下安全的多流复用

为实现安全重用，PR 首先明确了两个核心术语，作为后续判断逻辑的基础：

Free marker：通过cudaGraphAddEmptyNode创建的 “捕获合法” 空节点，插入到每个使用过该内存块的stream中，且位于该块最后一次被使用的操作之后，用于标记 “内存块已空闲”。
Terminal Node：流或捕获图中 “最新操作” 的集合，新捕获的操作会附加在终端节点之后。对于正在捕获的流，可通过cudaStreamGetCaptureInfo的dependencies_out参数获取终端节点集合。

内存块可重用性判断规则

cudaGraph生成一个DAG，因此作者提出了两种平衡 “安全性” 与 “灵活性” 的判断规则：

PyTorch pr(1)之torch.accelerator

Sun, 07 Dec 2025 20:44:02 +0800

Add unified memory APIs for torch.accelerator #152932

https://github.com/pytorch/pytorch/pull/152932

为torch.accelerator增加了如下API：

empty_cache
max_memory_allocated
max_memory_reserved
memory_allocated
memory_reserved
memory_stats
reset_accumulated_memory_stats
reset_peak_memory_stats

这个pr进一步说明了PyTorch正在去CUDA化，从而转变为一系列加速器提供计算框架，即使他们的github主页仍然写着：

Tensors and Dynamic neural networks in Python with strong GPU acceleration

这些用法很长一段时间是为cuda提供的，调用方式主要为torch.cuda.empty_cache等，接口主要作用是查看device上的显存占用如memory_allocated和memory_reserved，以及清理和重置相关状态。一起来看下相关函数的调用流程，以cuda为例，从源码分析这个接口是怎么一步一步获取底层的状态的。首先所有memory相关的数据都是从memory_stats这个字典内提取对于的value

def memory_reserved(device_index: _device_t = None, /) -> int:
    r"""Return the current :ref:`accelerator<accelerators>` device memory managed by the caching allocator
    in bytes for a given device index.
    Args:
        device_index (:class:`torch.device`, str, int, optional): the index of the device to target.
            If not given, use :func:`torch.accelerator.current_device_index` by default.
            If a :class:`torch.device` or str is provided, its type must match the current
            :ref:`accelerator<accelerators>` device type.
    """
    return memory_stats(device_index).get("reserved_bytes.all.current", 0)

而memory_stas通过底层_accelerator_getDeviceStats接口获取相关device的状态