跳到内容

双批次重叠 (Dual Batch Overlap)

动机

vLLM 中 DBO 系统的核心动机是在 MoE 层中将稀疏的全对全 (all-to-all) 通信与周围的计算进行重叠。该系统目前仅针对 DP+EP（数据并行 + 专家并行）部署。

简介

双批次重叠系统的工作原理是：在模型运行器中拆分批次，创建两个工作线程，然后在每个工作线程上运行模型。当启用 DBO 时，FusedMoEModularKernel 中的让出点 (yield points) 允许两个 CPU 工作线程（也称为 UBatch 线程）相互交替（ping-pong），以便当一个线程在执行计算时，另一个线程在等待通信。在整个代码中，ubatch 常被用作 microbatch（微批次）的缩写；这是 µ-batch 的 ASCII 友好版本。

DBO 系统包含对 GpuModelRunner 和 ModularKernel 的修改，并定义了两个工具类：UBatchWrapper 和 UBatchContext。UBatchWrapper 管理线程生命周期和模型的 CUDA 图执行。UBatchContext 封装了 ForwardContext，以协调两个 UBatch 线程之间的同步。

以下是 vLLM 目前实现的重叠调度。

# Schedule notation legend:
#    S = Shared expert
#    A0 = MLA qkv proj,
#    A1 = Core attn + out proj + MoE gate
#    D = Dispatch
#    C = Combine

# Comp: |-A0₀-A1₀-||-MLP₁-||-S₁-MLP₀-||-S₀-A0₁-A1₁-|
# Comm: |----D₁---||--D₀--||----C₁---||-----C₀-----|
# Order: D₁ send, A0₀, A1₀, D₁ recv, D₀ send, MLP₁, D₀ recv,
#        C₁ send, S₁, MLP₀, C₁ recv, C₀ send, S₀, A0₁, A1₁, C₀ recv.
# MLP_SHARED_OVERLAP = "mlp_shared_overlap"

使用 DBO 运行

要启用 DBO 系统，请在 vllm serve 命令中传入 --enable-dbo 参数。这必须与 --data-parallel-size N（其中 N 大于 1）和 --enable-expert-parallel 结合使用。此外，还有两个配置旋钮。

  • --dbo-decode-token-threshold：在仅解码 (decode-only) 批次中启用 DBO 所需的最小 token 数量。
  • --dbo-prefill-token-threshold：在包含至少一个预填充 (prefill) 的批次中启用 DBO 所需的最小 token 数量。

目前，DBO 仅支持 DeepEP，因此必须安装 DeepEP。如果您的工作负载主要是解码请求，则 --all2all-backend 参数必须设置为 deepep_low_latency；如果工作负载主要是预填充请求，则设置为 deepep_high_throughput。

以下命令将启动一个具有专家并行和 DBO 启用的双 DP 秩服务器。示例：vllm serve deepseek-ai/DeepSeek-V2-Lite --trust-remote-code --data-parallel-size 2 --enable-expert-parallel --enable-dbo --all2all-backend deepep_low_latency

请注意，CUDA_VISIBLE_DEVICES 中必须至少可见两个 GPU。

DBO 组件

  • GPUModelRunner
  • UBatchWrapper
  • UBatchContext

GPU 模型运行器

批次由 GPUModelRunner 类拆分为微批次。这通过两个步骤完成。首先，跨所有 DP 秩进行协调，以确定是否应用微批处理。微批处理必须在所有 DP 秩上保持一致。如果对于任何 DP 秩微批处理不可行，则对所有秩禁用它。如果所有 DP 秩都将进行微批处理，则 token 总数将被填充至所有秩中的最大 token 数。如果在应用填充后任何秩最终出现空的第二个微批次，则微批处理将被中止，所有秩都不会进行微批处理。一旦所有秩都启动了微批处理，就会执行第二步。CommonAttentionMetadata 会被 GPUModelRunner 切分，使得每个微批次对应一份注意力元数据。

UBatchWrapper

gpu_ubatch_wrapper

UBatchWrapper 类是一个模型包装器，负责 DBO 的所有线程、UBatchContext 和 CUDA 图管理。它的设计对 GPU 模型运行器相对透明。

该实现运行模型两次，每个微批次一次。每次模型调用都在一个 UBatch 线程中发生。这些线程并行启动，并使用 UBatchContext 进行同步。每个线程都被提供一份切分后的注意力元数据，用于运行其对应的半个批次。

DBO 的 CUDA 图完全由 UBatchWrapper 管理。因此，DBO 仅支持在完整 CUDA 图 (Full CUDA graphs) 下运行。但是，一旦捕获了 DBO CUDA 图，就可以在没有任何多线程或 CPU 同步的情况下重放它。

接口

__init__ 方法接收模型、VllmConfig、CUDAGraphMode 和设备。

forward 方法专门接收模型参数。它根据 forward_context 中是否存在 ubatch_slices 对象来决定是否使用 DBO 运行。否则，模型将不使用 DBO 运行。

UBatchContext

ubatch_context

UBatchContext 类是一个 ForwardContext 包装类，供 UBatchWrapper 类用来同步两个 UBatch 线程。它应该仅通过使用 make_ubatch_contexts 来实例化。

当其中一个 UBatch 线程到达 dbo_yield 调用时，它会暂停，并启动另一个线程，该线程将运行直到到达相同的 dbo_yield 调用。这种“乒乓”动态持续进行，线程在每次 dbo_yield 调用时进行交换，直到模型执行完毕。

当前的实现将所有 dbo_yield 和 dbo_maybe_run_recv_hook 调用放在 FusedMoEModularKernel.forward 方法中。

接口

make_ubatch_context 函数初始化两个 UBatchContext，每个 UBatch 线程一个。它接收两个 CUDA 流、预先存在的 ForwardContexts 和一个 CPU 线程屏障。此函数应专门用于实例化 UBatchContexts。它将处理所有事件初始化。

dbo_register_recv_hook 方法注册一个回调，该回调可以由另一个 UBatch 线程的 UBatchContext 中的 FusedMoEPrepareAndFinalizeModular 类返回。当另一个线程调用 dbo_maybe_run_recv_hook 时，该回调将被执行。这通常用于等待 all-to-all 内核。

dbo_maybe_run_recv_hook 方法运行一个由 dbo_register_recv_hook 函数设置的回调（如果该回调存在）。

dbo_yield 方法使当前线程进入睡眠状态并唤醒另一个 UBatch 线程。