CUDA 算子工程：手写 FlashAttention v2 之路

附录 A · CUDA Graph 与 Stream

作者 杨艺韬 · 1,880 字 · 发布于 · 更新于

decoding 阶段真正的瓶颈常常不在 GPU，而在 CPU 一次一次地把 kernel 递过去。 这个附录讲 Stream、Event 与 CUDA Graph——把上千次 launch 压成一次 replay。

A.1 Stream 与 Event：异步执行的基础

CUDA 的 host/device 异步模型基于两个原语：

Stream（流）：一个 FIFO 队列，里面装的是 GPU 任务（kernel launch、memcpy）。同一 stream 内的任务严格按顺序执行；不同 stream 之间默认并行（除非用 event 同步）。

Event（事件）：一个时间戳/同步点。可以记录到某个 stream 上，也可以让某个 stream 等待某个 event。

cudaStream_t s1, s2;
cudaStreamCreate(&s1);
cudaStreamCreate(&s2);

cudaEvent_t event;
cudaEventCreate(&event);

// stream 1: 跑 kernel A
kernel_a<<<..., 0, s1>>>(...);
cudaEventRecord(event, s1);  // 在 s1 上记录 event

// stream 2: 等 event, 然后跑 kernel B
cudaStreamWaitEvent(s2, event, 0);
kernel_b<<<..., 0, s2>>>(...);  // 等 kernel A 完成才会跑

这套机制让程序员可以构造任意复杂的 DAG（有向无环图）执行计划。

A.1.1 默认 Stream 的坑

CUDA 有一个"默认 stream"（也叫 NULL stream），所有不指定 stream 的调用都跑在它上。默认（legacy）stream 与所有阻塞型用户 stream 隐式同步——默认 stream 上的任务要等这些 stream 里此前的任务做完才开始，这些 stream 里此后的任务也要等它做完。

// 反例: 误用默认 stream
kernel_a<<<...>>>(...);                        // 默认 stream
cudaMemcpyAsync(..., user_stream);              // 想跟 kernel_a 并行?
// 实际上: cudaMemcpyAsync 必须等默认 stream 上的 kernel_a 完成

修复：所有调用都指定 stream，避免默认 stream。或者用 cudaStreamCreateWithFlags(&s, cudaStreamNonBlocking) 创建"non-blocking"stream，它不会被默认 stream 阻塞（用 cudaStreamCreate 建的是阻塞型）。

A.2 为什么需要 CUDA Graph

LLM 推理的一次 forward 涉及几百次 kernel launch（attention、GEMM、LayerNorm 等）。每次 launch 都有一笔微秒量级的固定开销，包括：

  • CPU 调度
  • driver 级 dispatch
  • GPU 接收 launch 指令
  • 寄存器初始化

判断它值不值得治，只需要一条判据：在 nsys timeline 上，GPU kernel 行之间的空隙占了多大比例。

  • 训练 / prefill：单个 kernel 动辄几百微秒到毫秒，launch 开销被彻底摊掉，timeline 上几乎看不到空隙——Graph 收益很小。
  • Decoding：batch 小、每个 kernel 只跑几十微秒，几百次 launch 的开销和有效工作同一个量级——timeline 上会看到密密麻麻的空隙，这才是 Graph 的主场。

本专栏没有条件给出具体的毫秒分解（那强依赖模型、batch、卡型和 driver 版本），读者自己抓一次 trace，把 kernel 行的空隙加起来，就知道自己这条链路上限是多少。

CUDA Graph 的做法：把一系列 kernel launch 录制成一个 Graph，之后每次只提交这一个 Graph，省去逐个 kernel 重复走 driver 的开销。图里的依赖关系在实例化时就固定下来，GPU 端可以直接按拓扑顺序推进。

A.3 CUDA Graph 用法

A.3.1 录制方式

最常用的方式是 stream capture：

// 1. 开始捕获
cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal);

// 2. 跑一遍 forward (所有 kernel 都在这个 stream 上)
forward(input, output, stream);
// 内部:
//   layer_norm<<<..., 0, stream>>>(...);
//   gemm<<<..., 0, stream>>>(...);
//   attention<<<..., 0, stream>>>(...);
//   ...

// 3. 结束捕获, 得到 graph
cudaGraph_t graph;
cudaStreamEndCapture(stream, &graph);

// 4. 实例化 graph (生成可执行版本)
cudaGraphExec_t graphExec;
cudaGraphInstantiate(&graphExec, graph, /*flags=*/0);  // CUDA 12 起是 3 参数版

// 后续推理: 只 launch 这个 graph
for (int step = 0; step < N; ++step) {
    cudaGraphLaunch(graphExec, stream);
}

CUDA Graph 把整个 forward 当作一个整体提交。

A.3.2 Graph 的限制

Graph 不是万能的。它有几个关键限制：

  1. 形状固定：录制时所有 kernel 的 launch 参数（grid_size、block_size）固定。如果 batch_size 或 seq_len 变化，原图不会跟着变，只能换图或更新节点参数。
  2. 指针固定：录制时的输入/输出指针被烧到 graph 里。每次 launch 必须用同样的指针（或用 cudaGraphExecKernelNodeSetParams 更新）。
  3. 捕获期间不能有 host 同步：对正在捕获的 stream 调 cudaStreamSynchronize、或调 cudaDeviceSynchronize，都是非法操作，会让这次捕获失败。cudaStreamCaptureModeGlobal 模式下，捕获期间本线程（以及其他线程，只要还有 Global 模式的捕获在进行）调用 cudaMalloc 这类"可能不安全"的 API 也会报错；ThreadLocal 只限制本线程，Relaxed 不做这层限制。
  4. 首次 launch 开销大：实例化和首次 launch 比单独 launch 慢，但后续 launch 极快。

A.3.3 Graph 与 LLM 推理

LLM 推理对 graph 来说是完美场景——decoding 阶段每个 token 的 forward 输入形状完全一样（都是 batch_size×1×hidden_size）。所以可以一次录制，反复 launch。例外是 attention：它要读的 KV 长度每步都在涨，要么把 seq_len、block table 放进显存由 kernel 自己读（launch 参数不变），要么干脆让 attention 不进图（vLLM V1 的做法，见 A.3.4 末尾）。

vLLM 默认启用 graph（TensorRT-LLM、SGLang 也都支持 CUDA Graph）：

# vLLM 配置
llm = LLM(model="...", enforce_eager=False)  # enforce_eager 不设时即按 False 处理，启用 graph

但 prefill 阶段每次 prompt 长度可能不同——整图录制的做法（vLLM V0 就是这样）只录纯 decode 批，prefill 与混合批走 eager。

实际推理中常用 graph pool 策略：预先录制几个不同 batch_size 的 graph，运行时按 input 选最合适的 graph。（这里的 graph pool 指这一组图；PyTorch 的 torch.cuda.graph_pool_handle() 和 vLLM 源码里的 graph_pool 指的是多张图共享的显存池，读源码时别混。）

A.3.4 更新 Graph 输入

如果不希望每次都重新录制，可以用 cudaGraphExecKernelNodeSetParams 在 graph 实例上修改某个 kernel 的参数：

// 取得 graph 中的某个 kernel node
cudaGraphNode_t kernel_node = ...;

// 准备新参数
cudaKernelNodeParams params = {...};
params.func = (void*)my_kernel;
params.gridDim = new_grid;
params.blockDim = new_block;
params.kernelParams = ...;

// 更新
cudaGraphExecKernelNodeSetParams(graphExec, kernel_node, &params);

这种 in-place 更新比重新实例化快得多。

不过要注意：vLLM 应对动态 batch size 用的不是这条路，而是 A.3.3 说的 graph pool——按 cudagraph_capture_sizes 里的一组档位分别捕获多张图（v0.8.5 里这个字段在 vllm-0.8.5/vllm/config.py:3429，不设时由 vllm-0.8.5/vllm/config.py:3972 的 _set_cudagraph_sizes() 给出默认档位：默认的 V1 引擎取 [1, 2, 4] + [8, 16, ..., 512]，再按 max_num_batched_tokens 截断；V0 才是按 max_num_seqs 截取），运行时把实际规模 padding 到不小于它的最小档位再选图。V1 的档位按本拍 token 总数计，而且是分段（piecewise）捕获：图按 attention 算子切开，attention 本身走 eager、不进图，所以 prefill 与混合批只要 token 总数不超过最大档位也能用上 graph（vllm-0.8.5/vllm/v1/worker/gpu_model_runner.py:1022-1026；详见《vLLM 推理内核深度解析》第 8 章 §8.4）。用 padding 换掉图参数更新的复杂度，在 LLM 这种"形状只有有限几种"的场景里更划算。

A.4 Stream 与多 GPU

多 GPU 场景下，stream 是基本协调单位：

cudaSetDevice(0);
cudaStream_t s_gpu0;
cudaStreamCreate(&s_gpu0);

cudaSetDevice(1);
cudaStream_t s_gpu1;
cudaStreamCreate(&s_gpu1);

// GPU 0 跑 attention  (第 3 个尖括号参数是动态 SMEM, 第 4 个才是 stream)
cudaSetDevice(0);
attention<<<grid, block, 0, s_gpu0>>>(...);

// GPU 1 跑 FFN
cudaSetDevice(1);
ffn<<<grid, block, 0, s_gpu1>>>(...);

// 跨 GPU 同步用 event（event 必须建在与记录它的 stream 同一设备上）
cudaSetDevice(0);
cudaEvent_t e0;
cudaEventCreate(&e0);
cudaEventRecord(e0, s_gpu0);
cudaSetDevice(1);
cudaStreamWaitEvent(s_gpu1, e0, 0);

NCCL 的 collective 操作（AllReduce 等）也都在 stream 上。

A.5 这个附录的小结

CUDA 异步执行模型是 LLM 推理优化的隐藏维度：

  1. Stream + Event 让 host/device 异步并行。
  2. CUDA Graph 把逐个 kernel 的 launch 开销压成一次提交。
  3. LLM 推理 decoding 阶段 是 graph 的完美场景：kernel 又多又小，launch 开销与有效工作同量级。收益大小看 nsys timeline 上 kernel 之间的空隙占比。
  4. vLLM 默认启用 graph——如果你部署 LLM 推理服务，这是必须开的（vLLM 里对应 enforce_eager=False）。