大模型架构

Appearance

CUDA Graph ​

CUDA Graph 将一组 GPU 操作及其依赖预先组织成图，后续通过一次图启动提交整组工作，减少 CPU 重复启动 Kernel 的开销。它主要优化工作提交方式，图中的 Kernel 仍按数据依赖执行。

本文先介绍 CUDA Graph 的生命周期与重放约束，再给出独立 PyTorch 示例，以及 LLM 推理中的形状分桶和捕获策略。

1. 提交开销 ​

一次模型前向包含投影、归一化、激活函数、Attention 和通信等操作。Eager 执行中，CPU 反复经过框架与驱动提交这些工作；当单个 Kernel 较短时，GPU 可能在等待下一次提交。

text
Eager
  CPU：提交 A → 提交 B → 提交 C → ……
  GPU：执行 A → 执行 B → 执行 C → ……

CUDA Graph
  初始化：记录 A → B → C 的操作与依赖
  CPU：启动图
  GPU：执行 A → 执行 B → 执行 C → ……

图示省略了 CPU 与 GPU 的异步重叠。CUDA Graph 降低的是重复提交成本，并不直接减少矩阵乘法的 FLOPs，也不会仅因为使用了图，就把多个 Kernel 融合成一个。

2. 生命周期 ​

2.1 图定义、实例化与启动 ​

CUDA Graph 是有向无环图。节点可以表示 Kernel、内存复制等工作，边表示执行依赖；依赖允许时，部分节点可以并行。

阶段作用典型 CUDA Runtime API
定义显式添加节点，或通过 Stream Capture 记录工作cudaGraphCreate、cudaStreamBeginCapture
捕获结束从被捕获的操作生成图cudaStreamEndCapture
实例化验证图并创建可执行实例cudaGraphInstantiate
启动把可执行图提交到 StreamcudaGraphLaunch

显式构图与 Stream Capture 是两种定义方式。PyTorch 的常用入口是 Stream Capture：进入 torch.cuda.graph() 后执行需要记录的张量操作，再使用 CUDAGraph.replay() 重放。底层阶段与节点类型见 CUDA Programming Guide。

2.2 捕获与普通执行 ​

捕获期间，Python 前向代码执行一次，用于发出被记录的 CUDA 操作；重放时不再重新运行这些 Python 语句。CPU 分支、计数器或打印不会因图重放而重新执行。

捕获阶段也不能当成一次普通 GPU 推理来消费结果。需要得到本轮输出时，应提交图重放，并在 CPU 实际读取前满足完成依赖。

PyTorch 默认的 CUDAGraph() 会在结束捕获时完成实例化；keep_graph=True 等高级用法另有生命周期要求。本文示例使用默认行为。CUDAGraph API

2.3 静态地址与动态数值 ​

普通 PyTorch 图重放读取捕获时记录的设备地址。输入与输出缓冲区需要保持存活，下一轮通过原地写入更新数据：

python
static_input.copy_(new_input)  # 更新图将读取的内容
# static_input = new_input    # 只改 Python 引用，图仍读取原地址

重放区域通常保持固定的形状、布局和操作结构。CUDA 提供图更新等更高级能力，但不会自动让上述 Python 重新赋值生效；推理框架常用静态缓冲区和形状分桶管理变化。

3. PyTorch 示例 ​

3.1 完整捕获与结果校验 ​

下面是无需模型权重的独立推理示例。输入和输出均为 [64, 128] 的 float32 CUDA 张量，计算逐元素函数 sin(x) + 0.25 * x²。运行需要 NVIDIA GPU 与支持 CUDA 的 PyTorch。

python
import torch


def forward(x):
    return torch.sin(x) + 0.25 * x.square()


@torch.no_grad()
def main():
    if not torch.cuda.is_available():
        raise RuntimeError("This example requires CUDA-enabled PyTorch and a GPU.")

    static_x = torch.zeros(64, 128, device="cuda")
    stream = torch.cuda.Stream()
    stream.wait_stream(torch.cuda.current_stream())

    # 在独立 Stream 上预热，提前完成首次调用的初始化。
    with torch.cuda.stream(stream):
        for _ in range(3):
            warm_output = forward(static_x)
    stream.synchronize()
    del warm_output

    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(graph, stream=stream):
        static_y = forward(static_x)

    saved_outputs = []
    for offset in (0.0, 0.5, -1.0):
        x = torch.full_like(static_x, offset)
        expected = forward(x)

        static_x.copy_(x)
        graph.replay()

        # clone 在重放之后排入当前 Stream，保存独立结果。
        actual = static_y.clone()
        torch.testing.assert_close(actual, expected)
        saved_outputs.append(actual)

    torch.cuda.synchronize()
    print("Three replays matched eager execution.")


if __name__ == "__main__":
    main()

这里保留 static_x、static_y 和 graph，使捕获地址在重放期间有效。每次重放都会覆盖 static_y；若直接把它多次加入列表，列表元素会指向同一块输出存储，因此示例使用 clone() 保存快照。

stream.synchronize() 用于初始化完成后的衔接；在线循环不要求每个 Kernel 后都同步。结果断言可能触发等待，适合正确性检查，不应放进正式性能计时区间。预热与地址约束依据 PyTorch CUDA Graph 说明，显式捕获 Stream 的行为见 torch.cuda.graph。

3.2 计时口径 ​

GPU 提交是异步的。用 CPU 时钟包住 replay() 而不等待，只能测到部分提交时间。测量连续调用完成的总耗时，可以在区间边界同步：

python
import time
import torch


def average_ms(run, repeats=100):
    for _ in range(5):
        run()
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(repeats):
        run()
    torch.cuda.synchronize()
    return (time.perf_counter() - start) * 1000 / repeats

这段函数测量一组连续调用的平均完成耗时，不等于在线请求的单次延迟。比较 Eager 与 Graph 时，应使用相同的数据准备和输出消费方式；若 Graph 分支包含输入 copy_()，就要明确报告这部分成本。CUDA Event 更适合测量设备时间线上的区间，端到端收益还需要包含 CPU 调度与数据传输。

4. LLM 推理中的应用 ​

4.1 形状分桶 ​

普通 Decode 每个请求通常贡献一个新 token，模型输入的主要变化是请求数；Prefill 的主要变化还包括本轮新增 token 总数与请求结构。框架可以预先捕获有限组尺寸，将实际输入补齐到能容纳它的桶。

例如，已捕获 8、16、32 行，实际有 13 行时使用 16 行图，其余 3 行作为 padding。静态形状有利于重放，但 padding 的稠密投影仍会产生计算成本；Attention 还需要正确的长度、索引或掩码，避免把无效行当成真实请求。

4.2 SGLang 捕获边界 ​

SGLang 将形状、静态缓冲区与执行准备放在 Runner，将图的捕获方式放在 Backend。其 FullCudaGraphBackend 对每个形状保存一张图及输出，执行时调用对应图的 replay()。

Full 表示对 Runner 交给 Backend 的区域整体捕获。本文核对的提交中，Prefill Runner 可将 Transformer 主体交给 Full Backend，而把 LM Head 与 Logits Processor 留在图外。因此，Full Graph 不意味着请求接收、Tokenization、采样或所有 CPU 调度都进入同一张图。SGLang Full Backend、Prefill Runner

4.3 分段捕获 ​

部分操作需要动态准备或不能直接捕获时，可以将稳定区域录成多张图：

text
Graph A → Eager 动态操作 → Graph B

Piecewise CUDA Graph 在 SGLang 的 TC 路径中通过编译器跟踪与切图形成这些区域；BCG 则在捕获过程中按显式断点分段。两者都使用底层 CUDA Graph，并承担多次图启动与图外调用的成本。

5. 性能与约束 ​

CUDA Graph 更适合操作结构重复、Kernel 较短且提交开销明显的工作负载。大型矩阵乘法已经占据主要时间时，减少提交开销对总耗时的影响可能较小。

约束工程处理
固定地址保留静态输入输出，并原地更新内容
形状变化分桶、重新捕获，或退回普通执行
CPU 取值与动态分支在图外处理，如依赖 GPU 结果的 .item() 分支
共享输出被覆盖在下一轮写入前消费，或复制保存结果
图与工作区常驻评估捕获尺寸数量、内存池和显存预算
初始化成本分开统计预热、捕获、首次重放与稳态执行

张量值变化与 Python 状态变化需要区别处理：设备 Kernel 可以读取更新后的输入值，但捕获时选中的普通 Python 分支不会在重放时重新判断。本文示例用于理解生命周期和结果一致性，未在当前环境执行 CUDA 测试，不提供加速倍数。

6. 参考 ​

  1. NVIDIA：CUDA Graphs
  2. PyTorch：CUDA Graph 使用说明
  3. PyTorch：torch.cuda.graph
  4. PyTorch：torch.cuda.CUDAGraph
  5. SGLang：FullCudaGraphBackend