Dual-batch overlap (DBO): two stages per step

Code version: afd-plugin @ 7452d52 · 2026-10-27 · vLLM @ ced6857

The idea#

In AFD the attention side sends hidden states to the FFN side and waits for the result. If it waited with nothing else to do, the GPU would idle. Dual-batch overlap (DBO) cuts each step into two stages. The attention side sends stage 0, then works on stage 1 while the FFN side computes stage 0. The two roles overlap.

DBO is a native vLLM feature. The plugin reuses it and adds two things: a yield op that hands control to the other stage between send and receive, and scoped patches that let DBO run even with one data-parallel rank.

An everyday comparison. Think of a cook with two pans. They start pan 1, and while it simmers they start pan 2, instead of standing and watching pan 1. The comparison stops here: AFD always uses exactly two pans, never three, and a validator enforces that.

The picture#

Read the top panel left to right. The overlap is the point: stage 1 of the attention side runs at the same time as stage 0 on the FFN side. The yield is what hands control across. The visual below lets you step through the same events.

In the code#

Step 1: Exactly two stages#

DBO turns on with --enable-dbo. In vLLM that sets the number of microbatches to two (vllm:vllm/config/parallel.py:599-601). AFD requires exactly two. The GPU validator raises for any other count (afd_plugin/validation.py:126-131).

    if (
        parallel.enable_dbo or parallel.use_ubatching
    ) and parallel.num_ubatches != GPU_MRV2_NUM_UBATCHES:
        raise RuntimeError(
            "AFD ModelRunnerV2 DBO/ubatching requires exactly two microbatches",
        )

Step 2: Each stage carries a stage index on the wire#

Each stage gets a stage index, 0 or 1. The ubatch helper clones the parent AFD metadata per stage and sets the index and the stage count (afd_plugin/v1/worker/ubatch_wrapper.py:26-36).

    ubatch_slice = ubatch_slices[ubatch_idx]
    clone = afd_metadata.clone()
    clone.stage_idx = ubatch_idx
    clone.num_stages = len(ubatch_slices)
    clone.tokens_start_loc = [int(ubatch_slice.token_slice.start)]
    clone.requests_start_loc = [int(ubatch_slice.request_slice.start)]
    clone.tokens_lens = [int(ubatch_slice.num_tokens)]
    clone.tokens_unpadded_lens = [
        _resolve_ubatch_unpadded_tokens(afd_metadata, ubatch_slice, ubatch_idx),
    ]
    return clone

The RemoteFFNProxy reads the stage index from the forward context and records it in the transfer metadata (afd_plugin/model_executor/models/deepseek_v2.py:136-143). So the FFN side knows which stage each tensor belongs to.

Step 3: The proxy yields between send and receive#

Inside the proxy, the order is: send the hidden states, yield, then receive the FFN output (afd_plugin/model_executor/models/deepseek_v2.py:146-158). The yield is what lets the other stage run while this one waits.

        afd_metadata.connector.send_attn_output(
            hidden_states,
            context,
            **send_kwargs,
        )
        hidden_states = maybe_apply_dbo_yield(
            hidden_states,
            role="attention",
        )
        return afd_metadata.connector.recv_ffn_output(
            ref_tensor=hidden_states,
            ubatch_idx=stage_idx,
        )

The yield calls a custom op, manual_dbo_yield, registered once per process (afd_plugin/v1/worker/dbo.py:27-49). The op yields only when vLLM DBO is active; otherwise it does nothing (afd_plugin/v1/worker/dbo.py:52-54).

def maybe_apply_dbo_yield(
    tensor: torch.Tensor,
    *,
    role: str,
) -> torch.Tensor:
    """Yield to the peer ubatch thread when vLLM DBO is active."""
    try:
        register_dbo_yield_custom_op()
    except ImportError:
        return tensor

    torch.ops.vllm.manual_dbo_yield(tensor)
    return tensor

Step 4: The thresholds decide when a step splits#

A step splits into two stages only when it is large enough. The recipes set two thresholds: --dbo-decode-token-threshold and --dbo-prefill-token-threshold (recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_dbo_dp1tp2.sh:23-25).

    --enable-dbo \
    --dbo-decode-token-threshold 2 \
    --dbo-prefill-token-threshold 12 \

vLLM reads these in check_ubatch_thresholds: a decode batch splits when its token count reaches the decode threshold, a batch with a prefill when it reaches the prefill threshold (vllm:vllm/v1/worker/ubatch_utils.py:141-149).

Step 5: DP1 DBO through scoped patches#

Native vLLM builds no ubatch runner when the data-parallel size is one. AFD coordinates the two stages locally, so it re-enables DP1 with two scoped patches (afd_plugin/compat/patches/gpu_mrv2_ubatching.py:263-282).

@contextmanager
def use_gpu_mrv2_ubatch_factory() -> Iterator[None]:
    """Scope the DP1 factory to AFD's native KV initialization call."""
    original_factory = model_runner.maybe_build_ubatch_runner
    model_runner.maybe_build_ubatch_runner = maybe_build_ubatch_runner
    try:
        yield
    finally:
        model_runner.maybe_build_ubatch_runner = original_factory


@contextmanager
def use_gpu_mrv2_dbo_dispatch() -> Iterator[None]:
    """Scope DP1 eligibility to one AFD native execute_model call."""
    original_dispatch = model_runner.dispatch_cg_and_sync_dp
    model_runner.dispatch_cg_and_sync_dp = dispatch_cg_and_sync_dp
    try:
        yield
    finally:
        model_runner.dispatch_cg_and_sync_dp = original_dispatch

Deep dive#

The blocks below give the two scoped patches, the threshold logic with numbers, and the yield mechanism.

What can go wrong#

• The wrong number of ubatches. Any count other than two raises at config validation (afd_plugin/validation.py:126-131).
• Thresholds set below the microbatch count. vLLM requires each threshold to be at least the number of microbatches (vllm:vllm/config/parallel.py:1093-1099).
• A step that is too small. A step under the threshold does not split, so it runs with no overlap (vllm:vllm/v1/worker/ubatch_utils.py:144-149). This is expected, not an error.
• DBO with a graph but not two ubatches. The graph policy raises when ubatching is on with any count other than two (afd_plugin/v1/worker/cuda_graph.py:72-77).

Why it matters. DBO keeps both roles busy. The attention side does useful work on one stage while the FFN side works on the other, instead of waiting. The two-stage limit and the yield placement are what make the overlap correct.