CUDA 图¶

本文档介绍了 vLLM v1 中新增的 CUDA 图模式，这是对之前 torch.compile 集成的改进。总结来说，我们

添加了灵活的 cudagraph_mode 配置
使完整的 CUDA 图支持与编译解耦
引入了 CUDA 图调度器作为中央控制器，自动为每个批次选择所需的运行时模式和 CUDA 图

本文档将讨论

动机
CUDA 图模式
详细设计
不同 CUDA 图模式的示例用法

注意

本文档中，我们将纯粹的解码（max_query_len=1）或投机解码（max_query_len =1+num_spec_tokens）称为统一解码批次，反之则为非统一批次（例如，预填充或预填充-解码混合批次）。

注意

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

动机¶

最初的分段编译是为了允许分段 CUDA 图捕获，排除了不支持 CUDA 图的操作（主要是 Attention）。这使得 CUDA 图能够提供一些加速，同时保持与所有 Attention 后端的兼容性。后来我们通过不进行分段编译来增加了对“完整 CUDA 图”的支持，以便在 Attention 支持 CUDA 图的情况下进一步降低延迟。然而，编译和 CUDA 图捕获之间的紧密耦合导致了一种“全有或全无”的体验，灵活性很低。许多 Attention 后端也还没有准备好统一的“完整” CUDA 图捕获（例如，目前只有 FlashAttention 3 支持），或者只支持纯解码批次的 CUDA 图（例如 Flashinfer, FlashMLA, Mamba 等）。这导致了令人困惑的性能/兼容性权衡、不一致的 CUDA 图支持以及日益复杂的代码结构。

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

明确区分预填充/混合批次或（统一）解码批次，并分别捕获。
将 CUDA 图捕获逻辑与编译分离（尽可能可行），以实现功能正交性，这表明
- 使用相同的编译图捕获分段和完整 CUDA 图，并且
- 在没有编译的情况下进行完整 CUDA 图捕获。
在运行时根据批次组成在完整和分段 CUDA 图之间进行调度。
集中控制 CUDA 图行为，以减少代码复杂性并提供更大的扩展性。

这些特性为各种启动/性能权衡和功能支持提供了最大的 CUDA 图捕获和编译灵活性。

`CudagraphModes`¶

CUDAGraphMode 是您在 CompilationConfig.cudagraph_mode 中调整的唯一参数。

NONE — 关闭 CUDA 图。适用于调试。
PIECEWISE — 单一模式策略（也是过去的默认设置）。这是最灵活的：Attention 或其他不支持 CUDA 图的操作保持即时执行，所有其他操作都进入 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 是单模式配置，并且分别相当于过去即时执行、分段 CUDA 图和完整 CUDA 图的实现，但 FULL_DECODE_ONLY 和 FULL_AND_PIECEWISE 是新添加的双模式配置，它们需要通过调度根据运行时批次动态切换到具体的运行时模式。

注意

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

虽然级联 Attention 不兼容 CUDA 图，但它现在与所有可能的 CUDA 图模式配置兼容。如果批次使用级联 Attention，它总是会被调度到 PIECEWISE 模式（如果可用）（否则是 NONE）。

注意

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

详细设计¶

概述¶

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

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

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

之前

之后

`BatchDescriptor`¶

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

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

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

此结构的目标是使用最少的项唯一标识一个（填充后的）批次，这些项对应于一个 CUDA 图项。

注意

未来，BatchDescriptor 的原型可能会扩展以适应更通用的情况，例如，包含更多项，如 uniform_query_len 来支持多种不同的统一解码长度设置（ Pull Request #23679），或支持 CUDA 图的非 token 长度感知输入模型所需的其他修改（例如，某些多模态输入）。

`CudagraphDispatcher`¶

CudagraphDispatcher 负责维护两组有效的调度键，一组用于 FULL 运行时模式，另一组用于 PIECEWISE 运行时模式，并在执行模型的前向传播之前调度正确的运行时模式和调度键。它将接收初始键（一个粗略的批次描述符，用于填充后的输入），并返回选定的运行时模式和最终的批次描述符，然后通过前向上下文将此决策告知 CUDAGraphWrapper 实例。请注意，CudagraphDispatcher 是可用 CUDA 图键的唯一事实来源，CUDAGraphWrapper 实例可以毫无顾虑地信任前向上下文关于要调度到哪个 CUDA 图的信息。这使我们能够简化 Wrapper 代码并将逻辑集中在调度器中。

调度键通过调度器的 initialize_cudagraph_keys 方法进行初始化，该方法在所有可能的 Attention 后端初始化完成后由 gpu_model_runner 调用。这是我们将来可以做得更花哨的地方，并“准备”所有可能的 CUDA 图组合。目前，我们只是根据编译配置中 decode_mode/mixed_mode 的 cudagraph_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 模式以进行即时执行。实现可以在此处找到。

下面是在模型执行器中运行时的工作流程的简化图示：

`CUDAGraphWrapper`¶

CUDAGraphWrapper 实例包装了一个可运行对象，并简单地模仿了可运行对象，附加了 CUDA 图功能。每个 Wrapper 实例都绑定到一个特定的 runtime_mode，该模式被限制为 PIECEWISE 和 FULL 模式，并负责捕获/重放和传递（直接调用）可运行对象。运行时，每个 Wrapper 会

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

以上步骤基于 CUDA 图 Wrapper 将直接信任前向上下文中的内容的假设（由调度器控制）。这使我们能够简化和集中逻辑，减少复杂性以及 Wrapper 和调度器之间状态不匹配的风险。它还允许 Wrapper 类同时用于 FULL 和 PIECEWISE 运行时模式。实现可以在此处找到。

嵌套 Wrapper 设计¶

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

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

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

完整 CUDA 图捕获 & 热身¶

CUDA 图捕获发生在 Runner 第一次使用非 NONE 运行时模式调用模型前向传播时（使用 _dummy_run）。对于完整的 CUDA 图捕获，我们通过正确设置 Attention 元数据来显式捕获不同的情况（例如，预填充/混合批次或统一解码批次），以确保底层的 Attention 后端启动所需的内核例程。区分预填充/混合批次或统一解码批次，最重要的属性是 attn_metadata 中的 max_query_len（对大多数 Attention 后端而言）。我们将其设置为所需的 uniform_query_len 以用于统一解码，否则将其设置为非统一解码批次的 num_tokens。

CUDA 图 Wrapper 不再管理热身逻辑。热身过程现在由 GPU 模型 Runner 直接控制，其中 NONE 运行时模式被分配用于热身的即时执行。在为完整 CUDA 图进行热身时，在热身 dummy_run 调用期间显式运行 Attention 也很重要。

Attention 后端对 CUDA 图的兼容性¶

为了表示 Attention 后端对 CUDA 图的兼容性，我们引入了一种新的枚举类型 AttentionCGSupport，它是一种枚举类型，用于跟踪 Attention 后端支持 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"""

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

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

Attention 后端	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 Attention 时将设置为 `UNIFORM_BATCH`
FlashMLA	`UNIFORM_BATCH`
FlashInferMLA	`UNIFORM_BATCH`
AITER MLA	`UNIFORM_SINGLE_TOKEN_DECODE`
CUTLASS MLA	`UNIFORM_SINGLE_TOKEN_DECODE`
Mamba attention	`UNIFORM_SINGLE_TOKEN_DECODE`

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

使用指南¶

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

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

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,
)

分段编译和完整图自定义通道（Attention 融合，序列并行）¶

不幸的是，一些自定义编译通道需要看到整个图才能生效，因此与分段编译不兼容。这包括 AttnFusionPass 和 SequenceParallelismPass。作为短期解决方案，当启用 Attention 融合时，我们自动禁用分段编译（通过设置 splitting_ops=[]）。我们使用 CUDA 图模式 FULL 或 FULL_DECODE_ONLY（取决于后端支持）。然而，这会导致另一个优化不兼容和令人困惑的性能权衡。

长期来看，我们增加了在 Inductor 中分区图的能力，而不是在 Dynamo 之后立即分区。可以通过 CompilationConfig.use_inductor_graph_partition=True 启用，但目前仍处于实验阶段，并且仅在 torch>=2.9 时可用。这也增加了编译时间，因为它需要编译整个图，并且无法重用分段编译的工件。一旦 vLLM 支持 2.9，我们计划将其作为默认方法，因为它也将加速分段 CUDA 图捕获。

关于性能¶

有关示例，请参阅以下链接