跳到内容

双批次重叠

动机

vLLM 中 DBO 系统的核心动机是使 MoE 层中的稀疏全对全通信与周围的计算重叠。该系统目前仅针对 DP+EP 部署。

简介

双批次重叠系统通过在模型运行器中分割批次，创建两个工作线程，然后在每个工作线程上运行模型来工作。当启用 DBO 时，FusedMoEModularKernel 中的屈服点允许两个 CPU 工作线程（也称为 UBatch 线程）相互之间进行乒乓操作，以便当一个线程正在进行计算时，另一个线程正在等待通信。在代码中，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：解码仅批次中启用该批次的 DBO 所需的最小 token 数
  • --dbo-prefill-token-threshold：包含至少一个预填充的批次中启用该批次的 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 仅支持与 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 函数初始化两个 UBatchContexts，每个 UBatch 线程一个。它接受两个 CUDA 流、现有的 ForwardContexts 和一个 CPU 线程屏障。此函数应仅用于实例化 UBatchContexts。它将处理所有事件初始化。

dbo_register_recv_hook 方法注册一个回调，该回调可以由另一个 UBatch 线程的 UBatchContext 中的 FusedMoEPrepareAndFinalize 类返回。当另一个线程调用 dbo_maybe_run_recv_hook 时，将运行该回调。这通常用于等待全对全内核。

dbo_maybe_run_recv_hook 方法运行由 dbo_register_recv_hook 函数设置的回调，如果该回调存在的话。

dbo_yield 方法使当前线程休眠并唤醒另一个 UBatch 线程。