优化 CUDA Graph 开销

CUDA Graph 捕获 OOM 不一定来自模型参数或有效计算量增长。逐层创建辅助 CUDA Stream，可能限制临时张量的内存复用，使图内存池保留远超实际同时存活张量所需的容量。本文以 SGLang 的共享 Stream 优化为例，解释分配器行为、跨流依赖和捕获内存，并提供独立的 PyTorch 对照实验。

1. 问题背景

部分 MoE 层包含 Shared Experts 与 Routed Experts：前者处理所有 token，后者处理路由选中的 token。两条分支读取同一层输入，可以在不同 Stream 上执行，最后合并结果。

一种直接实现是在每层构造时创建辅助 Stream：

class MoELayer:
    def __init__(self):
        self.alt_stream = torch.cuda.Stream()

对于 32 层模型，这会产生 32 个不同的辅助 Stream。虽然层内需要并行，但后一层依赖前一层输出，通常不需要为每层保留不同的 Stream 身份。

层 0：主 Stream 执行 Routed Experts，辅助 Stream 0 执行 Shared Experts
      两路汇合，得到层 1 输入
层 1：主 Stream 执行 Routed Experts，辅助 Stream 1 执行 Shared Experts
      两路汇合，得到层 2 输入
……

这里的主要开销不是 Stream 对象本身，而是各 Stream 上分配的中间张量，以及分配器为这些张量保留的内存块。

2. Stream 与内存复用

2.1 逻辑生命周期

假设每层辅助分支需要一个大小为 $W$ 的临时工作区，使用后立即释放，下一层必须等当前层完成后才能开始。从计算依赖看，所有层可以先后使用同一段存储。

时间 →
层 0 工作区：[使用 A][释放]
层 1 工作区：              [复用 A][释放]
层 2 工作区：                            [复用 A][释放]

这要求临时张量没有被输出列表、闭包或其他对象继续引用，也没有尚未完成的跨 Stream 使用。Python 中执行 del 只减少对象引用；GPU 异步操作是否完成，还需要由执行顺序和内存生命周期管理保证。

2.2 分配器复用条件

PyTorch 的原生 CUDA 缓存分配器把分配记录与 Stream 关联。其缓存查找不仅考虑容量，也考虑分配所属的 Stream；不能把另一个 Stream 的空闲缓存块直接当成本 Stream 的普通可复用块。本文讨论的是 backend:native，不将这一行为推广到所有 CUDA 分配器。分配器源码

因此，“上一层逻辑上已用完”与“下一层分配时能找到可复用块”是两件事：

组织方式	相同大小临时工作区的分配路径
每层独立辅助 Stream	层 0 的空闲块属于 Stream 0；层 1 在 Stream 1 上继续查找或申请
跨层共享辅助 Stream	各层在同一 Stream 上申请，更容易复用上一层释放的块

同一 Stream 内的操作有顺序保证。即使 CPU 已经释放了张量对象，而 GPU 还未执行完旧操作，后续使用同一存储的操作只要排在旧操作之后，也可以安全执行。跨 Stream 则不能仅凭 CPU 提交顺序作出这个判断。

在只考虑这一个临时工作区、没有缓存回收且同流充分复用的理想模型中，保留容量可能呈现：

$$M_{\text{独立流}}\approx LW,M_{\text{共享流}}\approx W$$

其中 $L$ 为层数，$W$ 为单层临时工作区大小。这个公式用于解释增长来源，不是实际模型的显存预测；权重、持久输出、同时存活的中间结果、分配粒度和跨流事件都会改变结果。

2.3 显存统计口径

诊断时至少区分以下指标：

指标	含义
memory_allocated()	当前仍分配给张量等用途的分配器内存
memory_reserved()	PyTorch 分配器管理的总容量，包括缓存中的空闲块
max_memory_allocated()	自统计重置以来的分配峰值
驱动层显存占用	还包含 CUDA 上下文、图执行实例、库工作区等分配器外资源

临时张量已经释放时，allocated 可以很低，但 reserved 仍然较高。reserved - allocated 大也不必然代表泄漏或全是碎片：其中可能包含正常缓存、图池保留容量和暂时不能复用的块。

3. CUDA Graph 的放大效应

3.1 固定地址与图内复用

CUDA Graph 记录 Kernel 的执行关系和使用的地址。重放时不会重新执行 Python 分配逻辑，因此捕获的地址必须继续有效。

PyTorch 在捕获期间使用图私有内存池，图存活期间保留其需要的存储。但这不等于“每个中间张量都必须独占一块内存”：只要生命周期和依赖允许，捕获时仍然可以释放临时分配，并将其存储分配给后续操作。重放会遵循已经捕获的顺序。图内存池说明

Kernel A 写临时地址 P
Kernel B 读取 P，完成当前层计算
Kernel C 复用 P，写下一层临时结果

共享 Stream 改善的正是这种复用机会。多个 Stream 使缓存块难以互用，而图池又需要保留捕获所需容量，二者共同抬高持久占用。

3.2 BCG 的多桶与多图段

BCG 为不同 Token 桶捕获图段，并在动态操作处切开执行。相比单个固定形状的简单图，它有更多中间分配、图段和生命周期约束，更容易暴露 Stream 组织不当带来的内存问题。

不能直接用“桶数 × 层数 × 工作区”推算实际显存。图段和桶可能共享内存池，持久输出也可能被专门复用；最终结果取决于捕获顺序、地址约束和分配器实际决策。

empty_cache() 也不能替代结构优化。它可以归还可释放的空闲缓存，但不能随意归还存活图重放仍需要的存储。只有相关图、输出及共享池引用的生命周期结束后，相应容量才可能被回收。

4. 共享 Stream 实现

4.1 按用途分配

在一个模型执行上下文中创建一次辅助 Stream，再传给所有层，即可表达这项优化：

alt = torch.cuda.Stream()
layers = [MoELayer(alt_stream=alt) for _ in range(32)]

SGLang 使用 get_stream("alt") 按用途取得共享 Stream。它保留层内并行，同时避免每层独立创建 Stream。对应提交

主 Stream：  Routed 0 → 合并 0 → Routed 1 → 合并 1
辅助 Stream：Shared 0 ──┘        Shared 1 ──┘

适用前提是这些层没有需要保留的跨层并发。如果执行器还要同时运行多个批次或模型，共享范围应与执行上下文匹配，不能无条件把所有任务塞进一个全局 Stream。

4.2 分叉与汇合

下面是数据流示意，省略模型定义和额外的张量生命周期管理：

main = torch.cuda.current_stream()
alt.wait_stream(main)          # 辅助分支等待输入就绪

with torch.cuda.stream(alt):
    shared = shared_experts(x)

routed = routed_experts(x)

main.wait_stream(alt)          # 合并操作等待辅助分支完成
out = routed + shared

wait_stream() 给 GPU 建立执行依赖，不是 CPU 阻塞式的设备同步。捕获期间，这些依赖进入图，保证重放时先完成输入准备，再执行两条分支，最后合并结果。

执行依赖和内存生命周期需要同时正确。如果一个张量在其他 Stream 上仍被使用，而原分配 Stream 可能回收其存储，就需要 record_stream() 或等价的显式生命周期安排。record_stream() 用来延迟不安全的内存复用，不负责让消费者等待生产者完成，也不会把缓存块转移到另一个 Stream。

共享 Stream 不会自动修复原地写入冲突。例如一条分支覆写输入，另一条分支仍在读取相同输入，即使数学上独立也会产生竞争。

5. PyTorch 对照实验

5.1 实验结构

下面的程序只使用 PyTorch，不依赖 SGLang 或模型权重。它重复执行若干个“主分支 + 辅助分支”的阶段，每个阶段结束后汇合，再进入下一阶段。

辅助分支创建一个大型临时张量 tmp，在本 Stream 上归约成预分配的标量，然后释放 tmp。主分支同时执行另一个归约。大型临时张量不跨 Stream 传递，从而突出分配器对 Stream 身份的处理，而不是其他跨流生命周期问题。

输入 x 形状为 [N]，输出形状为 [layers]，每个输出元素均为：

$$y_i=\sum _{j=1}^N{x}_j+\sum _{j=1}^N{\sin }^2(x_j)$$

两种模式只有 Stream 分配方式不同：per-layer 每阶段一个辅助 Stream，shared 所有阶段共享一个。程序使用单个 CUDA Graph 演示底层机制，没有模拟完整的 MoE 或 BCG 多桶系统。

📄 完整 Demo：stream_graph_memory.py

import argparse
import gc
import math

import torch


@torch.no_grad()
def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["per-layer", "shared"], required=True)
    parser.add_argument("--execution", choices=["eager", "graph"], default="graph")
    parser.add_argument("--layers", type=int, default=32)
    parser.add_argument("--temp-mib", type=int, default=8)
    args = parser.parse_args()
    if args.layers < 1 or args.temp_mib < 1:
        parser.error("layers 和 temp-mib 必须为正数")
    if not torch.cuda.is_available():
        raise RuntimeError("需要 CUDA GPU 和支持 CUDA 的 PyTorch")
    if torch.cuda.memory.get_allocator_backend() != "native":
        raise RuntimeError("本实验需要 backend:native")

    torch.cuda.set_device(0)
    n = args.temp_mib * 1024 * 1024 // 4  # FP32：每元素 4 字节
    x = torch.full((n,), 0.5, device="cuda", dtype=torch.float32)
    aux_scores = torch.empty(args.layers, device="cuda")
    main_scores = torch.empty_like(aux_scores)
    result = torch.empty_like(aux_scores)
    # 预先建立标量视图；整个实验期间保持输入、输出及视图存活。
    aux_rows = list(aux_scores.unbind())
    main_rows = list(main_scores.unbind())
    result_rows = list(result.unbind())

    count = 1 if args.mode == "shared" else args.layers
    streams = [torch.cuda.Stream() for _ in range(count)]
    capture_stream = torch.cuda.Stream()
    capture_stream.wait_stream(torch.cuda.current_stream())
    addresses = []

    def forward():
        main_stream = torch.cuda.current_stream()
        for i in range(args.layers):
            alt = streams[0] if args.mode == "shared" else streams[i]
            alt.wait_stream(main_stream)
            with torch.cuda.stream(alt):
                tmp = torch.sin(x)
                addresses.append(tmp.data_ptr())  # 只保存整数，不保留张量
                tmp.square_()
                torch.sum(tmp, dim=(0,), out=aux_rows[i])
                del tmp
            torch.sum(x, dim=(0,), out=main_rows[i])
            main_stream.wait_stream(alt)
            torch.add(main_rows[i], aux_rows[i], out=result_rows[i])

    # 在独立流上预热；清理普通缓存，避免其掩盖捕获新增容量。
    with torch.cuda.stream(capture_stream):
        for _ in range(2):
            forward()
    torch.cuda.synchronize()
    addresses.clear()
    gc.collect()
    torch.cuda.empty_cache()
    before = torch.cuda.memory_reserved()
    torch.cuda.reset_peak_memory_stats()

    graph = None
    if args.execution == "graph":
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph, stream=capture_stream):
            forward()
        graph.replay()
    else:
        with torch.cuda.stream(capture_stream):
            forward()
    torch.cuda.synchronize()

    mib = 1024 ** 2
    print(f"torch={torch.__version__}, cuda={torch.version.cuda}")
    print(f"gpu={torch.cuda.get_device_name(0)}")
    print(f"mode={args.mode}, execution={args.execution}, streams={count}")
    print(f"unique_tmp_addresses={len(set(addresses))}")
    print(f"allocated_MiB={torch.cuda.memory_allocated() / mib:.2f}")
    print(f"reserved_MiB={torch.cuda.memory_reserved() / mib:.2f}")
    print(f"reserved_delta_MiB={(torch.cuda.memory_reserved() - before) / mib:.2f}")
    print(f"peak_allocated_MiB={torch.cuda.max_memory_allocated() / mib:.2f}")

    # 改变输入验证重放，防止只检查到捕获/预热留下的结果。
    for value in (0.25, 0.75):
        x.fill_(value)
        if graph is not None:
            graph.replay()
        else:
            capture_stream.wait_stream(torch.cuda.current_stream())
            with torch.cuda.stream(capture_stream):
                forward()
        torch.cuda.synchronize()
        expected = n * (value + math.sin(value) ** 2)
        torch.testing.assert_close(
            result, torch.full_like(result, expected), rtol=1e-4, atol=1e-3
        )
    print("correctness=passed")


if __name__ == "__main__":
    main()

5.2 运行方式

将代码保存为 stream_graph_memory.py，在已有 CUDA PyTorch 环境中分别启动新进程。显式指定原生分配器，避免进程间缓存残留和分配器后端差异干扰对照：

PYTORCH_CUDA_ALLOC_CONF=backend:native python stream_graph_memory.py --mode per-layer
PYTORCH_CUDA_ALLOC_CONF=backend:native python stream_graph_memory.py --mode shared

再比较不捕获图的执行，观察同样的 Stream 复用约束如何出现在普通缓存中：

PYTORCH_CUDA_ALLOC_CONF=backend:native python stream_graph_memory.py --mode per-layer --execution eager
PYTORCH_CUDA_ALLOC_CONF=backend:native python stream_graph_memory.py --mode shared --execution eager

默认每个临时张量为 8 MiB，32 阶段。可分别改变 --layers 和 --temp-mib 观察趋势；这里不需要主动制造 OOM。

5.3 结果解释

重点观察 unique_tmp_addresses 和 reserved_delta_MiB：共享 Stream 通常对应更少的临时地址、更小的保留容量增量，而 allocated 未必有同等幅度的差异，因为大型临时张量在阶段结束时已释放。

观测	解释
多个阶段记录相同临时地址	同一块存储被依次复用，不代表多个仍存活张量发生非法重叠
当前分配小、保留容量大	张量已释放，但分配器或图池仍保留容量
共享流没有达到单工作区容量	分配粒度、归约内部临时空间或版本策略仍带来额外占用
共享流未出现预期改善	先核对后端、独立进程、Stream 数量和临时地址，再检查实际分配行为

程序在 GPU 同步后取样，报告的是整个 PyTorch 分配器的统计，不是图池专属统计，也没有测量 CUDA 驱动的全部资源。本文未在 CUDA GPU 上实测该 Demo，不提供虚构的输出数值或固定降幅。

6. SGLang 案例

SGLang PR #36624 在 Cohere Command-A-Plus 上将逐层辅助 Stream 改为共享 Stream。作者在 4×GB300 环境报告以下 Prefill 图捕获数据：

配置	捕获显存：优化前 → 优化后	捕获时间：优化前 → 优化后
NVFP4	23.57 → 5.34 GB	22.33 → 19.23 s
BF16	21.36 → 5.21 GB	24.04 → 17.47 s

对应捕获显存降幅约为 77% 和 76%。这些沿用 PR 的统计口径，不能直接等同于 Demo 的 reserved_delta_MiB，也不是所有模型的预期收益。

该 PR 同时包含 Decode 的 Attention/MLP 并行和 MoE 后端选择。因此，其端到端吞吐结果属于组合优化，不能全部归因于 Stream 共享。本文关注的因果链是：减少不必要的 Stream 身份 → 改善临时分配复用 → 降低捕获保留容量 → 缓解启动 OOM 和 KV Cache 容量压力。

已经分配好的 KV Cache 不会因为其他内存被释放而自动扩容。要把节省的容量转为更多缓存，需要结合框架启动时的内存规划；也不能把所有节省空间都直接分给 KV Cache，还要保留运行期峰值所需空间。

7. 排查顺序与适用边界

对类似问题，可以按以下顺序定位：

1. Stream 来源：检查模型层构造函数是否重复调用 torch.cuda.Stream()，确认实例数量及用途。
2. 生命周期：确认工作区是否真正释放，是否被列表、闭包、图输出或其他对象持有。
3. 内存口径：区分当前分配、保留容量、峰值和驱动层占用，必要时用内存快照追踪分配 Stream。
4. 单变量对照：固定模型、桶、精度和其他后端，仅改变 Stream 组织方式，并用独立进程测量。
5. 依赖与精度：检查分叉/汇合、跨流张量生命周期和原地写入，在多次更新输入后验证重放结果。

共享 Stream 最适合“层间顺序执行、层内辅助分支结构相似、临时工作区生命周期不重叠”的情况。如果显存主要由权重、KV Cache 或持久输出占用，这项优化不会直接减少它们；如果不同层或批次本来需要并发，共享 Stream 还可能增加串行化。

8. 参考

1. SGLang PR #36624：Cohere Command-A-Plus Decode 与 BCG 优化
2. SGLang：对应模型实现，提交 e6f21cdadc
3. PyTorch v2.11.0：CUDA 原生缓存分配器
4. PyTorch：CUDA 语义与内存管理
5. CUDA Graph
6. Breakable CUDA Graph