Dual Batch Overlap

Dual Batch Overlap (DBO)

동기 (Motivation)

vLLM의 DBO 시스템의 핵심 동기는 MoE 레이어의 희소한 all-to-all 통신을 주변 계산과 겹치는 것입니다. 이 시스템은 현재 DP+EP 배포만 대상으로 합니다.

출처: 문서

본문

소개 (Introduction)

Dual Batch Overlap 시스템은 모델 러너에서 배치를 분할해 두 개의 워커 스레드를 만들고 각 워커 스레드에서 모델을 실행하는 방식으로 동작합니다. DBO가 활성화되면 FusedMoEModularKernel 안의 yield 지점들이 두 CPU 워커 스레드(UBatch 스레드라고도 함)가 서로 핑퐁하도록 하여, 하나가 계산을 실행하는 동안 다른 하나는 커뮤니케이션을 기다리게 합니다. 코드 전반에서 ubatch는 microbatch의 짧은 형태로 쓰일 수 있는데, 이는 µ-batch(마이크로배치)의 ASCII 친화적인 표기입니다.

DBO 시스템은 GpuModelRunner와 ModularKernel의 수정을 포함하며, 두 유틸리티 클래스 UBatchWrapper와 UBatchContext를 정의합니다. UBatchWrapper는 모델의 스레드 생명주기와 CUDA graph 실행을 관리합니다. UBatchContext는 두 UBatch 스레드의 동기화를 조정하기 위해 ForwardContext를 감쌉니다.

다음은 vLLM에 현재 구현된 overlap 스케줄입니다.

# 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로 실행 (Running with DBO)

DBO 시스템을 활성화하려면 vllm serve 명령에 --enable-dbo 인자를 전달합니다. 반드시 --data-parallel-size N(N은 1보다 큼)과 --enable-expert-parallel과 함께 실행해야 합니다. 추가로 두 가지 구성 노브가 있습니다.

  • --dbo-decode-token-threshold — decode 전용 배치에 DBO를 활성화하는 데 필요한 최소 토큰 수
  • --dbo-prefill-token-threshold — prefill이 하나 이상 포함된 배치에 DBO를 활성화하는 데 필요한 최소 토큰 수

현재 DBO는 DeepEP에서만 지원되므로, DeepEP가 설치돼 있어야 하고 워크로드가 주로 decode 요청이면 --all2all-backend 인자가 deepep_low_latency로, 주로 prefill 요청이면 deepep_high_throughput으로 설정돼 있어야 합니다.

다음은 전문가 병렬과 DBO를 활성화한 2개 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 컴포넌트 (DBO Components)

  • GPUModelRunner
  • UBatchWrapper
  • UBatchContext

GPU Model Runner

배치는 GPUModelRunner 클래스에 의해 마이크로배치로 분할됩니다. 이는 두 단계로 수행됩니다. 먼저 모든 DP 랭크 간 조정을 통해 마이크로배칭을 적용할지 결정합니다. 마이크로배칭은 모든 DP 랭크에서 균일해야 합니다. 어떤 DP 랭크에서 마이크로배칭이 불가능하면 모든 랭크에서 비활성화됩니다. 모든 DP 랭크가 마이크로배칭한다면 총 토큰 수는 모든 랭크 중 최대 토큰 수까지 패딩됩니다. 패딩 적용 후 어떤 랭크가 빈 두 번째 마이크로배치로 끝나게 되면 마이크로배칭은 중단되고 어떤 랭크도 마이크로배칭하지 않습니다. 모든 랭크가 마이크로배칭을 시작하면 두 번째 단계가 수행됩니다. CommonAttentionMetadata는 GPUModelRunner에 의해 반으로 슬라이스돼 마이크로배치당 하나의 attention 메타데이터가 생깁니다.

UBatchWrapper

UBatchWrapper 클래스는 DBO의 모든 스레드, UBatchContext, CUDA graph 관리를 담당하는 모델 래퍼입니다. GPU 모델 러너에 비교적 투명하게 설계됐습니다.

구현은 마이크로배치 각각에 대해 모델을 두 번 실행합니다. 각 모델 호출은 UBatch 스레드 안에서 발생합니다. 이 스레드들은 병렬로 실행되며 UBatchContext로 동기화됩니다. 각 스레드는 배치의 자기 절반을 실행하는 데 사용되는 슬라이스된 attention 메타데이터를 제공받습니다.

DBO용 CUDA graph는 전적으로 UBatchWrapper가 관리합니다. 그 때문에 DBO는 Full CUDA graphs로만 실행을 지원합니다. 그러나 일단 DBO CUDA graph가 캡처되면 멀티스레딩이나 CPU 동기화 없이 재생(replay)할 수 있습니다.

인터페이스 (Interfaces)

__init__ 메서드는 모델, VllmConfig, CUDAGraphMode, device를 받습니다.

forward 메서드는 모델 인자만 받습니다. forward_context에 ubatch_slices 객체가 있는지에 따라 DBO로 실행할지 결정합니다. 그렇지 않으면 DBO 없이 모델을 실행합니다.

UBatchContext

UBatchContext 클래스는 UBatchWrapper가 두 UBatch 스레드를 동기화하는 데 사용하는 ForwardContext 래퍼 클래스입니다. 반드시 make_ubatch_contexts로만 인스턴스화해야 합니다.

UBatch 스레드 중 하나가 dbo_yield 호출에 도달하면 일시 정지하고, 다른 스레드를 시작해 같은 dbo_yield 호출에 도달할 때까지 실행합니다. 이런 "ping-pong" 동적이 계속되며 각 dbo_yield 호출에서 스레드가 교체되고, 모델 실행이 완료될 때까지 반복됩니다.

현재 구현에서는 모든 dbo_yield와 dbo_maybe_run_recv_hook 호출이 FusedMoEModularKernel.forward 메서드에 있습니다.

인터페이스 (Interfaces)

make_ubatch_context 함수는 UBatch 스레드 각각에 대해 하나씩 두 개의 UBatchContext를 초기화합니다. 두 CUDA 스트림, 기존 ForwardContexts, CPU 스레드 배리어를 받습니다. 이 함수를 통해서만 UBatchContext를 인스턴스화해야 하며, 모든 이벤트 초기화를 처리합니다.

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 스레드를 깨웁니다.

더 알아보기 (Learn more)