vllm 学习笔记

Appearance

将多个 CUDA kernel 的启动和执行录制为一张静态图，通过一次性重放消除 kernel launch 开销。

为什么需要 CUDA Graph ​

LLM 推理的 decode 阶段每次 forward 计算量很小，但 CPU 向 GPU 提交 kernel 的 launch 开销（约 5-10us/kernel）相对显著。一个 Transformer 层涉及数十个 kernel，累积的 launch 延迟可能占总延迟的 30% 以上。CUDA Graph 将整张计算图录制下来，每次 decode 只需重放一次 graph，消除所有 kernel launch 开销。

核心原理 ​

  • Capture：在推理开始前，用 dummy 输入执行一次 forward，录制所有 kernel 为 CUDA Graph。
  • Replay：实际推理时将输入数据拷贝到预留 buffer 并 replay graph，无需 CPU 逐个 launch kernel。
  • 静态形状限制：Graph 中所有张量形状必须固定，因此通常只为 decode 阶段（固定 batch）捕获 graph。
  • 多 graph 管理：为不同 batch size 捕获不同 graph，运行时选择最接近的 graph 并 padding。

在源码中的实现 ​

  • vllm/v1/worker/gpu/cudagraph_utils.py — CUDAGraph 的 capture 与 replay，管理多个 batch size 的 graph。
  • vllm/config/compilation.py — CompilationConfig 控制 graph capture 的静态化。
  • vllm/config/model.py — ModelConfig.enforce_eager 开关关闭 CUDA Graph。
  • vllm/v1/worker/gpu/warmup.py — Worker 初始化时 warmup 触发 graph capture。
  • vllm/v1/cudagraph_dispatcher.py — 运行时按 batch size 分发到对应 graph。

相关概念 ​

反向链接