跳到內容

雙批次重疊

動機

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 執行緒。