跳到内容

视觉编码器 (ViT) CUDA Graphs

vLLM 中的 CUDA Graphs 基础设施主要针对解码器(语言模型)的前向传播。vLLM 同时支持将编码器(视觉 Transformer)的前向传播捕获为 CUDA Graphs,并与解码器相互独立。这基于 Pull Request #35963

注意

编码器 CUDA Graphs 与解码器 CUDA Graphs 正交——两者可以同时启用。编码器图捕获视觉编码器执行(例如 Qwen3-VL 中的 ViT),而解码器图则捕获语言模型执行,如 CUDA Graphs 设计文档中所述。

动机

视觉编码器推理会在主机侧产生 CUDA 内核启动开销。当批量大小 (batch size) 或图像尺寸较小时,这种开销尤为显著。

编码器 CUDA Graphs 通过在模型初始化期间预先捕获多个令牌预算级别下的完整编码器前向传播,并在运行时重放对应的图,从而消除了这种开销。

设计

编码器 CUDA Graph 系统使用基于预算的捕获/重放策略,由 EncoderCudaGraphManager 管理。该系统包含以下核心组件:

  • EncoderCudaGraphManager:统筹编码器 CUDA Graphs 的捕获、重放、贪心装箱和数据并行执行。
  • SupportsEncoderCudaGraph:一种运行时可检查的协议,模型通过实现该协议来选择加入编码器 CUDA Graphs。
  • BudgetGraphMetadata:保存单个令牌预算级别下捕获的 CUDA Graph 及其关联的 I/O 缓冲区。

基于预算的图捕获

系统会在不同的令牌预算级别(例如 [2048, 4096, 8192, 13824])预先捕获多个 CUDA Graphs。每个预算定义了一个固定的令牌容量,所有预算共享相同的最大批量大小(图像数量)。每个级别的 BudgetGraphMetadata 存储了该图以及预分配的输入、元数据和输出缓冲区。

@dataclass
class BudgetGraphMetadata:
    token_budget: int
    max_batch_size: int
    graph: torch.cuda.CUDAGraph
    input_buffer: torch.Tensor       # e.g. pixel_values
    metadata_buffers: dict[str, torch.Tensor]  # e.g. embeddings, seq metadata
    output_buffer: torch.Tensor      # encoder hidden states

预算通过 get_encoder_cudagraph_budget_range() 从模型提供的范围自动生成为 2 的幂次级别,即使最大预算不落在 2 的幂次边界上,也会始终包含该最大预算。用户也可以通过 CompilationConfig 中的 encoder_cudagraph_token_budgets 明确指定预算。

运行时贪心装箱

当一批图像到达时,管理器会根据输出令牌数量对图像进行排序(从小到大),并贪心地将尽可能多的图像装入每个子批次中,同时保持在最大令牌预算和最大批量大小之内。一旦某个子批次确定(下一个图像会超出任一限制),管理器会找到适合该子批次总令牌数的最小预算,并重放相应的 CUDA Graph。此过程不断重复,直到该批次处理完毕。超过所有预算的图像将回退到即时 (eager) 执行。

对于每次图重放:

  1. 清零预分配的 input_buffer,然后将输入张量(如 pixel_values)复制到其中。
  2. 清零 metadata_buffers,然后切片复制预计算值(如旋转嵌入、序列元数据)。
  3. 重放 CUDA Graph。
  4. 克隆 output_buffer 中的输出(由于缓冲区在多次重放间复用,因此必须进行克隆)。

数据并行支持

mm_encoder_tp_mode="data" 时,管理器通过 get_load_balance_assignment 使用负载均衡分配方式将图像分发到各个 TP rank,在每个 rank 上本地执行,然后通过 tensor_model_parallel_all_gather 按原始顺序收集结果。

通过 SupportsEncoderCudaGraph 进行模型集成

模型通过实现 SupportsEncoderCudaGraph 协议来选择加入编码器 CUDA Graphs。该协议封装了所有与模型相关的逻辑,使得管理器保持模型无关性。协议定义了以下方法:

  • get_encoder_cudagraph_config() — 返回静态配置(支持的模态、输入键、缓冲区键、输出隐藏层大小)。
  • get_encoder_cudagraph_budget_range(vllm_config) — 返回 (min_budget, max_budget),用于自动推断令牌预算。
  • get_encoder_cudagraph_num_items(mm_kwargs) — 返回批次中的项目数量(例如图像数量)。
  • get_encoder_cudagraph_per_item_output_tokens(mm_kwargs) — 返回每个项目的输出令牌计数,用于贪心装箱。
  • get_encoder_cudagraph_per_item_input_sizes(mm_kwargs) — 返回每个项目的输入大小(例如补丁数量),用于数据并行负载均衡。
  • select_encoder_cudagraph_items(mm_kwargs, indices) — 按索引提取子批次项目,用于贪心装箱和数据并行分片。
  • prepare_encoder_cudagraph_capture_inputs(...) — 创建用于图捕获的虚拟输入。
  • prepare_encoder_cudagraph_replay_buffers(...) — 在重放前从实际批次输入计算新的缓冲区值。
  • encoder_cudagraph_forward(...) — 使用预计算缓冲区进行前向传播(在捕获和重放期间调用)。
  • encoder_eager_forward(...) — 当没有合适的图匹配时回退到的即时前向传播。

当前支持:Qwen3-VL(见 vllm/model_executor/models/qwen3_vl.py)。

注意

SupportsEncoderCudaGraph 协议设计为模型无关。新的视觉编码器模型可以通过实现协议方法加入,而无需修改管理器。

注意

编码器 CUDA Graphs 目前已在 Blackwell GPU 上通过 --mm-encoder-attn-backend=FLASH_ATTN--mm-encoder-attn-backend=FLASHINFER 进行了测试。

配置

CompilationConfig 中的三个字段控制编码器 CUDA Graphs:

  • cudagraph_mm_encoder (bool, 默认 False) — 启用多模态编码器的 CUDA Graph 捕获。启用后,将为每个令牌预算级别捕获完整的编码器前向传播作为 CUDA Graph。
  • encoder_cudagraph_token_budgets (list[int], 默认 []) — 用于捕获的令牌预算级别。如果为空(默认),则从模型架构中按 2 的幂次自动推断。用户提供的值将覆盖自动推断。
  • encoder_cudagraph_max_images_per_batch (int, 默认 0) — 捕获期间每批次的最大图像数量。如果为 0(默认),则自动推断为 max_budget // min_budget

使用指南

通过 compilation_config 启用编码器 CUDA Graphs:

vllm serve Qwen/Qwen3-VL-32B \
  --compilation-config '{"cudagraph_mm_encoder": true}'

使用明确的预算:

vllm serve Qwen/Qwen3-VL-32B \
  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824], "encoder_cudagraph_max_images_per_batch": 8}'

Python 示例:

import vllm

compilation_config = {
    "cudagraph_mm_encoder": True,
    # Optional: override auto-inferred budgets
    # "encoder_cudagraph_token_budgets": [2048, 4096, 8192, 13824],
    # "encoder_cudagraph_max_images_per_batch": 8,
}

model = vllm.LLM(
    model="Qwen/Qwen3-VL-32B",
    compilation_config=compilation_config,
)

管理器会跟踪命中/未命中统计信息并定期记录。“命中”意味着图像通过 CUDA Graph 重放处理;“未命中”意味着回退到即时执行(图像超过了所有预算)。

关于性能

以下基准测试是在 Blackwell GPU (GB200) 上使用 vllm bench mm-processor 运行的。有关详细信息,请参阅 #35963

单 GPU (1x GB200)

模型:Qwen/Qwen3-VL-30B-A3B-Instruct,数据集:lmarena-ai/VisionArena-Chat (3000 条提示,300 条预热),max_model_len=32768

Backend 平均延迟提升 P99 延迟提升
FLASH_ATTN +11.8% (5.13→4.52ms) +31.6% (9.16→6.26ms)
FLASHINFER +19.6% (5.42→4.36ms) +40.3% (10.87→6.49ms)

重现方法:

vllm bench mm-processor \
  --model Qwen/Qwen3-VL-30B-A3B-Instruct \
  --dataset-name hf --dataset-path lmarena-ai/VisionArena-Chat \
  --num-prompts 3000 --num-warmups 300 \
  --max-model-len 32768 --seed 42 \
  --mm-encoder-attn-backend FLASH_ATTN \
  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'

多 GPU (4x GB200, TP=4, DP=4)

模型:Qwen/Qwen3-VL-32B-Instruct,数据集:random-mm (1000 条提示,200 条预热,每请求 20 张图像,分辨率 336x336),max_model_len=8192

Backend 平均延迟提升 P99 延迟提升
FLASH_ATTN +18.4% (28.39→23.16ms) +14.0% (238.78→205.28ms)
FLASHINFER +44.4% (23.24→12.91ms) +84.9% (172.41→26.05ms)

重现方法:

vllm bench mm-processor \
  --model Qwen/Qwen3-VL-32B-Instruct \
  --dataset-name random-mm \
  --random-mm-base-items-per-request 20 \
  --random-mm-num-mm-items-range-ratio 0.0 \
  --random-mm-bucket-config '{"(336,336,1)": 1.0}' \
  --num-prompts 1000 --num-warmups 200 \
  --max-model-len 8192 --seed 42 \
  --mm-encoder-attn-backend FLASHINFER \
  --tensor-parallel-size 4 --mm-encoder-tp-mode data \
  --compilation-config '{"cudagraph_mm_encoder": true, "encoder_cudagraph_token_budgets": [512, 1024, 1536, 2048, 2560, 3072, 3584, 4096, 4864], "encoder_cudagraph_max_images_per_batch": 8}'