CUDA Graphs¶

本文档介绍了 vLLM v1 中除之前的 torch.compile 集成外的新 CUDA 图模式。总结如下：

增加了灵活的 cudagraph_mode 配置
使完整的 CUDA 图支持与编译过程正交化
引入了 CUDA 图调度器作为中央控制器，可根据批次自动选择所需的运行时模式和 CUDA 图

在本文档中，我们将讨论：

动机
CUDA 图模式
详细设计
不同 CUDA 图模式的使用示例
视觉编码器 (ViT) CUDA 图

注意

在本文档中，我们将纯解码（max_query_len=1）或投机解码（max_query_len = 1 + num_spec_tokens）称为均匀解码（uniform decode）批次，反之则为非均匀（non-uniform）批次（例如预填充或混合预填充-解码批次）。

注意

以下内容主要基于 Pull Request #20059 的最后一次提交。

动机¶

最初的分段编译旨在允许进行分段 CUDA 图捕获，排除不支持 CUDA 图的操作（主要是注意力机制）。这在保持与所有注意力后端兼容的同时，通过 CUDA 图实现了一定的加速。后来，我们通过放弃分段编译添加了对“完整 CUDA 图”的支持，以便在注意力机制支持 CUDA 图的情况下进一步降低延迟。然而，这种编译与 CUDA 图捕获之间的紧密耦合导致了“要么全有，要么全无”的体验，缺乏灵活性。此外，许多注意力后端尚未准备好统一的“完整” CUDA 图捕获（例如，目前仅 FlashAttention 3 支持），或者仅支持纯解码批次的 CUDA 图（例如 Flashinfer、FlashMLA、Mamba 等）。这导致了令人困惑的性能/兼容性权衡、不一致的 CUDA 图支持以及日益复杂的代码结构。

这促使我们寻求一种更细粒度的 CUDA 图解决方案，具备以下特性：

显式感知预填充/混合批次或（均匀）解码批次的 CUDA 图，并分别进行捕获。
将 CUDA 图捕获逻辑与编译逻辑（尽可能）分离，以实现特性的正交性，这意味着：
- 使用相同的编译图捕获分段和完整 CUDA 图；
- 无需编译即可进行完整 CUDA 图捕获。
在运行时根据批次组成在完整和分段 CUDA 图之间进行调度。
集中控制 CUDA 图行为，以降低代码复杂性并提高可扩展性。

这些特性为各种启动/性能权衡和特性支持提供了最大的灵活性。

`CudagraphModes`¶

CUDAGraphMode 是你在 CompilationConfig.cudagraph_mode 中调整的唯一旋钮：

NONE — 关闭 CUDA 图。适合调试。
PIECEWISE — 一种单模式策略（也是过去的默认设置）。它是最灵活的：注意力机制或其他不兼容 CUDA 图的操作保持 eager 执行，其他所有内容进入 CUDA 图。需要分段编译。
FULL — 一种单模式策略，仅为非均匀批次捕获完整 CUDA 图，然后均匀解码批次复用相同 batch_size 的非均匀批次 CUDA 图（因为它们兼容）；这对于小型模型或小提示词的工作负载非常有效。
FULL_DECODE_ONLY — 仅为均匀解码提供完整 CUDA 图，预填充/混合等不使用 CUDA 图；适用于 P/D 设置中的解码实例，其中预填充不太重要，这样可以节省 PIECEWISE CUDA 图所需的内存。
FULL_AND_PIECEWISE — （默认模式）为均匀解码提供完整 CUDA 图，为其他批次提供分段 CUDA 图；通常是性能最优的设置，特别是在小型模型或 MoE 的低延迟场景下，但也需要最多的内存，且捕获耗时最长。

默认值：如果您在 v1 上使用分段编译，我们默认为 FULL_AND_PIECEWISE 以获得更好的性能（对于池化模型，仍为 PIECEWISE）。否则，例如如果无法进行分段编译，我们默认为 NONE。

虽然 NONE、PIECEWISE 和 FULL 是单模式配置，分别等同于之前的 eager 执行、分段 CUDA 图和完整 CUDA 图实现，但 FULL_DECODE_ONLY 和 FULL_AND_PIECEWISE 是新增加的双模式配置，需要调度器根据运行时批次动态切换具体的运行时模式。

注意

这里，单模式 NONE、PIECEWISE 和 FULL 被视为 CUDA 图调度的运行时模式。如果使用双模式，调度器将始终根据批次组成调度到其成员模式之一（如果没有合适的 CUDA 图，则可能包含 NONE）。

虽然级联注意力（cascade attention）与 CUDA 图不兼容，但它现在与所有可能的 CUDA 图模式配置兼容。如果批次使用了级联注意力，它始终被调度到 PIECEWISE 模式（如果可用，否则为 NONE）。

注意

并非所有 CUDA 图模式都与每个注意力后端兼容。我们会自动将模式“降级”为最接近的受支持模式。例如，如果某个后端仅支持纯解码/均匀批次的 CUDA 图，我们会将 FULL 转换为 FULL_AND_PIECEWISE（如果启用了分段编译），否则转换为 FULL_DECODE_ONLY。

Detailed Design (详细设计)¶

概述¶

新的 CUDA 图逻辑构建在分段编译之上，并支持双 CUDA 图运行时模式切换。系统包含以下核心组件：

CUDAGraphWrapper：处理包装的可调用对象的 CUDA 图捕获与重放的包装器。
CudagraphDispatcher：中央控制器，包含关于 CUDA 图的单一事实来源，并处理它们之间的调度。
CUDAGraphMode：描述支持模式和运行时模式的枚举（如上所述）。
BatchDescriptor：作为运行时批次的唯一表示，用于调度。

参见下图，了解 vLLM 以前和当前在 Inductor 编译下 CUDA 图设计模式的快速比较。我们可以看到，以前 CUDA 图逻辑和编译逻辑紧密耦合在 vllm PiecewiseBackend 中，CUDA 图由 batch_size 隐式调度。现在 CUDA 图逻辑被分离到 CUDAGraphWrapper 类中，负责完整和分段 CUDA 图功能，调度是通过运行时模式加上 BatchDescriptor 作为调度键，通过 CudagraphDispatcher 显式完成的。

以前

现在

`BatchDescriptor`¶

BatchDescriptor 是 ForwardContext 中的一个组件，与 CUDA 图运行时模式一起，作为运行时调度的核心结构。其原型为：

class BatchDescriptor(NamedTuple):
    num_tokens: int
    num_reqs: int
    uniform: bool = False
    has_lora: bool = False

其中 num_tokens 可以是填充后的 token 长度，uniform 表示所有请求是否具有相同的查询长度。许多注意力后端仅在批次均匀时才支持完整 CUDA 图；纯解码批次是均匀的，但查询长度可能不为 1（即 num_tokens == num_reqs），这发生在投机解码的验证过程中，此时“解码”批次的查询长度为 1+num_spec_tokens。

该结构的目标是使用最少且对应于 CUDA 图条目的项来唯一标识一个（填充后的）批次。

注意

BatchDescriptor 的原型在未来可能会扩展以适应更通用的情况，例如包含更多条目（如 uniform_query_len 以支持多个不同的均匀解码长度设置，参考 Pull Request #23679），或者为了支持输入不一定感知 token 长度的模型（例如某些多模态输入）所需的其他修改。

`CudagraphDispatcher`¶

CudagraphDispatcher 负责维护两组有效的调度键，一组用于 FULL 运行时模式，一组用于 PIECEWISE 运行时模式，并在执行模型前向传递之前调度正确的运行时模式和调度键。它接收初始键（填充后输入的简略 batch_descriptor），并返回选定的运行时模式和最终的 batch_descriptor，然后通过前向上下文告诉 CUDAGraphWrapper 实例该决策。注意 CudagraphDispatcher 是可用 CUDA 图键的唯一事实来源，CUDAGraphWrapper 实例可以盲目信任前向上下文中应该调度到哪个 CUDA 图。这使我们可以简化包装器代码并将逻辑集中在调度器中。

调度键通过调度器的 initialize_cudagraph_keys 方法初始化，该方法在所有可能的注意力后端初始化后由 gpu_model_runner 调用。这是我们未来可以做更多花样的地方，例如“准备”各种 CUDA 图组合。目前，我们只是根据编译配置中 cudagraph_mode 的 decode_mode/mixed_mode 以及 cudagraph_capture_sizes 的有效组合来添加可用键。

调度代码如下：

batch_descriptor=BatchDescriptor(num_tokens=num_input_tokens, uniform_decode=...)
runtime_mode, batch_descriptor = cudagraphdispatcher.dispatch(batch_descriptor)
# execution
with set_forward_context(
    ..., 
    cudagraph_runtime_mode=runtime_mode, 
    batch_descriptor=batch_descriptor,
):
     output = self.model(...)

在 dispatch() 方法内部，调度器将搜索合适的 CUDA 图运行时模式和现有的调度键以返回。我们基本上按照优先级搜索现有键：FULL > PIECEWISE > None。如果调度键不存在，则默认返回 NONE 模式以进行 eager 执行。实现可以在这里找到。

以下是模型执行器中运行时工作流程的简化说明：

`CUDAGraphWrapper`¶

CUDAGraphWrapper 实例包装了一个可运行对象，并简单地通过增加 CUDA 图能力来模拟该可运行对象。每个包装器实例都绑定到一个特定的 runtime_mode（限制为 PIECEWISE 和 FULL 模式），并负责捕获/重放以及透传（直接调用）可运行对象。在运行时，每个包装器会：

从全局前向上下文中检查 runtime_mode 和 batch_descriptor（调度键）。
如果 runtime_mode 为 NONE 或 runtime_mode 与包装器的模式不匹配，则直接调用可运行对象。
否则，即如果 runtime_mode 与包装器的模式匹配，包装器将执行 CUDA 图捕获（如果键不存在，则创建新条目并缓存）或重放（如果键存在于缓存中）。

上述步骤基于这样的假设：CUDA 图包装器将直接信任前向上下文中的内容（由调度器控制）。这使我们可以简化和集中逻辑，降低复杂性以及包装器和调度器之间状态不匹配的风险。它还允许在 FULL 和 PIECEWISE 运行时模式中复用包装器类。参见实现这里。

Nested Wrapper design (嵌套包装器设计)¶

使完整 CUDA 图和分段 CUDA 图共存并兼容的核心机制是嵌套 CUDA 图包装器设计，该设计构建在具有单个分段 FX 图的分段编译之上。我们在整个模型外部包装了一个 FULL 模式包装器以实现完整 CUDA 图功能；同时，每个分段后端都在编译内部通过 PIECEWISE 模式包装器进行包装。

下面的流程图应清楚地描述其工作原理。

因此，对于 FULL 运行时模式，捕获/重放完整 CUDA 图是安全的，因为分段包装器未被激活。PIECEWISE 模式的情况类似，因为 FULL 模式包装器和 PIECEWISE 模式包装器之间没有冲突。对于 NONE 运行时模式，FULL 和 PIECEWISE 包装器都不会被激活，因此我们直接进行 eager 执行。

Full CUDA Graph capturing & warm-up (完整 CUDA 图捕获与预热)¶

CUDA 图捕获发生在 runner 首次使用非 NONE 运行时模式调用模型前向（使用 _dummy_run）时。对于完整 CUDA 图捕获，我们通过适当地设置注意力元数据来显式捕获不同情况（即预填充/混合批次或均匀解码批次），以确保底层的注意力后端启动所需的内核例程。为了区分预填充/混合批次或均匀解码批次，最重要的属性是 attn_metadata 中的 max_query_len（对于大多数注意力后端为真）。我们将其设置为均匀解码所需的 uniform_query_len，否则将其设置为非均匀解码批次的 num_tokens。

CUDA 图包装器不再管理预热逻辑。预热过程现在直接由 GPU 模型 runner 控制，其中分配 NONE 运行时模式以进行预热的 eager 执行。在为完整 CUDA 图预热时，在预热 dummy_run 调用期间显式运行注意力机制也很重要。

CUDA Graphs Compatibility of Attention Backends (注意力后端的 CUDA 图兼容性)¶

为了标识注意力后端的 CUDA 图兼容性，我们引入了一个新的枚举类型 AttentionCGSupport，它跟踪注意力后端支持 CUDA 图的能力。该值按能力顺序排序，即 ALWAYS > UNIFORM_BATCH > UNIFORM_SINGLE_TOKEN_DECODE > NEVER。

class AttentionCGSupport(enum.Enum):
    """ Constants for the CUDA Graphs support of the attention backend
    Here we do not consider the cascade attention, as currently
    it is never CUDA Graphs supported."""

    ALWAYS = 3
    """CUDA Graphs always supported; supports mixed-prefill-decode"""
    UNIFORM_BATCH = 2
    """CUDA Graphs supported for batches the only contain query lengths that are
    the same, this can be used for spec-decode 
        i.e. "decodes" are 1 + num_speculative_tokens"""
    UNIFORM_SINGLE_TOKEN_DECODE = 1
    """CUDA Graphs supported for batches the only contain query_len==1 decodes"""
    NEVER = 0
    """NO CUDA Graphs support"""

假设我们有混合注意力后端（例如在 Mamba 混合模型中）。在这种情况下，我们寻求所有后端的最低能力来确定模型的最终能力，并且我们可能会通过将模式降级为最合适的模式来解决不兼容的 CUDA 图模式。例如，如果最低能力为 UNIFORM_BATCH，则将 FULL 模式降级为 FULL_AND_PIECEWISE 模式，或者如果对于 -O3 编译模式最低能力为 NEVER，则降级为 PIECEWISE 模式。关于完整的降级策略，请参见此处的代码。

下表列出了在撰写本文时支持完整 CUDA 图的后端。

注意力后端	cudagraph_support	评论
FlashAttention v2	`UNIFORM_BATCH`	实际上是 `ALWAYS`，但出于性能原因权宜之计回退到 `FULL_AND_PIECEWISE`
FlashAttention v3	`ALWAYS`	针对两种批次都有统一的例程，因此 `FULL` 模式很好
Triton Attention	`ALWAYS`	偏好 `FULL_AND_PIECEWISE`，因为它为预填充/混合和纯解码批次有不同的内核
AITER FlashAttention	`UNIFORM_BATCH`
FlashInfer	`UNIFORM_SINGLE_TOKEN_DECODE`	在 Blackwell 上使用 TRTLLM 注意力时将设置为 `UNIFORM_BATCH`
FlashMLA	`UNIFORM_BATCH`
FlashInferMLA	`UNIFORM_BATCH`
FlashInferMLASparse	`UNIFORM_BATCH`
AITER MLA	`UNIFORM_SINGLE_TOKEN_DECODE`
CUTLASS MLA	`UNIFORM_SINGLE_TOKEN_DECODE`
Mamba attention	`UNIFORM_SINGLE_TOKEN_DECODE`

未列出的后端均声明为 NEVER。

Usage guide (使用指南)¶

现在 CLI 直接在 compilation_config 中使用 cudagraph_mode 的大写字符串： --compilation-config '{"cudagraph_mode": "..."}'，其中 ... 应该是 NONE、PIECEWISE、FULL、FULL_DECODE_ONLY 和 FULL_AND_PIECEWISE 之一。请注意，所有 PIECEWISE 相关模式都需要分段编译，所有 FULL 相关模式都需要注意力后端的 CUDA 图支持。例如：

vllm serve --model meta-llama/Llama-3.1-8B-Instruct --compilation-config '{"cudagraph_mode": "FULL_AND_PIECEWISE"}'

Python examples (Python 示例)¶

import os
os.environ.setdefault("VLLM_LOGGING_LEVEL", "DEBUG")

import vllm
from vllm.config import CUDAGraphMode

compilation_config = {"mode": 3, "cudagraph_mode": "FULL_AND_PIECEWISE"}
model = vllm.LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    dtype="auto",
    compilation_config=compilation_config,
)
sampling_params = vllm.SamplingParams(
    temperature=0,  # greedy decoding
    max_tokens=1024,
)
outputs = model.generate(
    ["My name is John and"],
    sampling_params=sampling_params,
)

Piecewise compilation and full graph custom passes (分段编译与全图自定义优化，如注意力融合、序列并行)¶

遗憾的是，一些自定义编译优化必须查看整个图才能有效，因此与分段编译不兼容。这包括 AttnQuantFusionPass 和 SequenceParallelismPass。作为短期解决方案，我们在启用注意力融合时自动禁用分段编译（通过设置 splitting_ops=[]）。我们使用 CUDA 图模式 FULL 或 FULL_DECODE_ONLY（取决于后端支持）。然而，这导致了另一种优化不兼容和令人困惑的性能权衡。

从长远来看，我们增加了在 Inductor 中而不是在 Dynamo 之后对图进行分区的能力。可以通过 CompilationConfig.use_inductor_graph_partition=True 启用，但目前是实验性的，仅在 torch>=2.9 下可用。这也增加了编译时间，因为它必须编译整个图且无法复用分段编译产物。一旦 vLLM 支持 2.9，我们计划将其作为默认方法，因为它也将加速分段 CUDA 图捕获。

About the Performance (关于性能)¶

请参见以下链接获取示例：