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/worker/model_runner.py — CUDAGraphRunner 负责 capture 和 replay，管理多个 batch size 的 graph。
  • vllm/compiler/ — 编译器辅助静态化模型以支持 graph capture。
  • vllm/config.py — CompilationConfig 和 enforce_eager 开关控制是否启用 CUDA Graph。
  • vllm/worker/worker.py — Worker 初始化时调用 warmup_model 触发 graph capture。

相关概念 ​