vllm.v1.worker.gpu.ubatch_utils ¶

def _slice_input_batch(
    input_batch: InputBatch,
    ubatch_slice: UBatchSlice,
    query_start_loc_buf: torch.Tensor,
    seq_lens_buf: torch.Tensor,
    dcp_local_seq_lens_buf: torch.Tensor | None = None,
    *,
    dcp_size: int = 1,
    dcp_rank: int = 0,
    cp_interleave: int = 1,
) -> InputBatch:
    """Build the sub-`InputBatch` a single microbatch runs on.

    A request straddling the boundary appears in both microbatches, its query
    truncated to the tokens each one owns.

    Query offsets and sequence lengths (including DCP-local lengths) are
    written to the caller's persistent buffers after truncation. Other GPU
    inputs are views of the full batch's buffers.

    The sub-batch describes the forward pass only. Sampling runs once over the
    merged batch, so `logits_indices`, `cu_num_logits` and the draft-token
    fields are carried over unchanged and must not be read off a sub-batch.
    """
    assert not ubatch_slice.is_empty(), f"Ubatch slice {ubatch_slice} is empty"

    req_start = ubatch_slice.request_slice.start
    req_stop = ubatch_slice.request_slice.stop
    tok_start = ubatch_slice.token_slice.start
    tok_stop = ubatch_slice.token_slice.stop

    num_reqs_padded = req_stop - req_start
    num_tokens_padded = tok_stop - tok_start
    # The trailing microbatch may extend past the real batch into CUDA graph
    # padding; the unpadded counts stop at the last real request/token.
    num_reqs = max(0, min(req_stop, input_batch.num_reqs) - req_start)
    num_tokens = max(0, min(tok_stop, input_batch.num_tokens) - tok_start)

    query_start_loc = _slice_query_start_loc(
        input_batch.query_start_loc,
        query_start_loc_buf[: num_reqs_padded + 1],
        req_start,
        req_stop,
        tok_start,
        num_tokens_padded,
    )
    query_start_loc_np = np.clip(
        input_batch.query_start_loc_np[req_start : req_stop + 1] - tok_start,
        0,
        num_tokens_padded,
    ).astype(np.int32)

    seq_lens = _slice_seq_lens(
        input_batch.seq_lens,
        seq_lens_buf[:num_reqs_padded],
        input_batch.query_start_loc,
        req_start,
        req_stop,
        tok_stop,
    )
    # Same truncation as `_slice_seq_lens`, on the host-side upper bound.
    seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound[
        req_start:req_stop
    ].clone()
    tokens_truncated = max(0, int(input_batch.query_start_loc_np[req_stop]) - tok_stop)
    if tokens_truncated:
        seq_lens_cpu_upper_bound[-1] -= tokens_truncated

    # Query lengths of the truncated requests, so consumers that derive
    # max_query_len from this array see the microbatch's own lengths.
    num_scheduled_tokens = np.diff(query_start_loc_np)[:num_reqs]

    dcp_local_seq_lens = None
    if dcp_size > 1:
        assert dcp_local_seq_lens_buf is not None
        dcp_local_seq_lens = maybe_prepare_dcp_local_seq_lens(
            dcp_local_seq_lens_buf,
            seq_lens,
            num_reqs,
            dcp_size,
            dcp_rank,
            cp_interleave,
            num_reqs_padded=num_reqs_padded,
        )

    return replace(
        input_batch,
        req_ids=input_batch.req_ids[req_start : min(req_stop, input_batch.num_reqs)],
        num_reqs=num_reqs,
        num_reqs_after_padding=num_reqs_padded,
        idx_mapping=input_batch.idx_mapping[req_start:req_stop],
        idx_mapping_np=input_batch.idx_mapping_np[req_start:req_stop],
        num_scheduled_tokens=num_scheduled_tokens,
        num_tokens=num_tokens,
        num_tokens_after_padding=num_tokens_padded,
        query_start_loc=query_start_loc,
        query_start_loc_np=query_start_loc_np,
        seq_lens=seq_lens,
        seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound,
        dcp_local_seq_lens=dcp_local_seq_lens,
        num_computed_tokens_np=input_batch.num_computed_tokens_np[req_start:req_stop],
        prefill_len_np=input_batch.prefill_len_np[req_start:req_stop],
        num_computed_prefill_tokens_np=input_batch.num_computed_prefill_tokens_np[
            req_start:req_stop
        ],
        is_prefilling_np=input_batch.is_prefilling_np[req_start:req_stop],
        input_ids=input_batch.input_ids[tok_start:tok_stop],
        positions=input_batch.positions[tok_start:tok_stop],
        is_padding=input_batch.is_padding[tok_start:tok_stop],
        prompt_lens=(
            None
            if input_batch.prompt_lens is None
            else input_batch.prompt_lens[req_start:req_stop]
        ),
        fast_prefill=None,
    )

def _slice_query_start_loc(
    query_start_loc: torch.Tensor,
    out: torch.Tensor,
    req_start: int,
    req_stop: int,
    tok_start: int,
    num_tokens_padded: int,
) -> torch.Tensor:
    """Rebase query_start_loc onto the microbatch's token range.

    The clamp truncates the requests straddling either boundary to the tokens
    this microbatch owns.
    """
    torch.sub(query_start_loc[req_start : req_stop + 1], tok_start, out=out)
    return out.clamp_(0, num_tokens_padded)

def _slice_seq_lens(
    seq_lens: torch.Tensor,
    out: torch.Tensor,
    query_start_loc: torch.Tensor,
    req_start: int,
    req_stop: int,
    tok_stop: int,
) -> torch.Tensor:
    """Slice seq_lens, shortening the request truncated at the boundary.

    Only tokens past the end of this microbatch are missing; tokens it owns in
    an *earlier* microbatch are computed before this one reads them, so they
    still count. Deriving the truncation from this microbatch's own query
    lengths would subtract them twice. It is computed from tensors, not Python
    ints, so this stays valid inside a CUDA graph.
    """
    out.copy_(seq_lens[req_start:req_stop])
    last = req_stop - req_start - 1
    out[last] -= (query_start_loc[req_stop] - tok_stop).clamp_(min=0)
    return out

def create_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices:
    """Split a DP-padded batch into the slices its microbatches run on.

    Splitting at the midpoint of the DP-padded token count leaves the trailing
    microbatch anywhere from full to entirely inside the padding. Clamping its
    request slice onto the last request leaves it holding that request with
    zero query tokens -- the shape a straddling request already takes -- so it
    stays well-formed with no work to do, like a dummy run.
    """
    _, ubatch_slices = maybe_create_ubatch_slices(
        True,
        input_batch.num_scheduled_tokens,
        input_batch.num_tokens_after_padding,
        input_batch.num_reqs_after_padding,
        num_ubatches,
    )
    assert ubatch_slices is not None
    return [
        UBatchSlice(
            slice(
                min(s.request_slice.start, input_batch.num_reqs - 1),
                min(s.request_slice.stop, input_batch.num_reqs_after_padding),
            ),
            s.token_slice,
        )
        for s in ubatch_slices
    ]

def maybe_build_ubatch_runner(
    vllm_config: VllmConfig,
    device: torch.device,
    model_state: ModelState,
    attn_groups: list[list[AttentionGroup]],
    kv_cache_config: KVCacheConfig,
    max_num_reqs: int,
) -> UBatchRunner | None:
    """Build the microbatch runner, or None when DBO is not in use.

    Microbatching needs the DP handshake to agree on it, so it is only
    available with more than one DP rank (as in the V1 runner).
    """
    parallel_config = vllm_config.parallel_config
    if not parallel_config.use_ubatching or parallel_config.data_parallel_size <= 1:
        return None

    logger.info_once(
        "Dual batch overlap is enabled. Microbatched steps run without "
        "CUDA graphs on the V2 model runner."
    )
    return UBatchRunner(
        vllm_config, device, model_state, attn_groups, kv_cache_config, max_num_reqs
    )

def merge_ubatch_outputs(outputs: list[Any]) -> Any:
    """Reassemble the full-batch model output from the per-microbatch ones.

    Models return hidden states, a tuple of those plus auxiliary hidden states
    (EAGLE3), or IntermediateTensors on non-final pipeline ranks.
    """
    first = outputs[0]
    if isinstance(first, IntermediateTensors):
        return IntermediateTensors(
            {
                key: torch.cat([out.tensors[key] for out in outputs], dim=0)
                for key in first.tensors
            }
        )
    if isinstance(first, tuple):
        return tuple(torch.cat(parts, dim=0) for parts in zip(*outputs))
    return torch.cat(outputs, dim=0)

def slice_model_inputs(
    model_inputs: dict[str, Any], token_slice: slice
) -> dict[str, Any]:
    """Narrow the model's per-token inputs to one microbatch."""
    sliced = dict(model_inputs)
    for key in ("input_ids", "inputs_embeds"):
        value = model_inputs.get(key)
        if value is not None:
            sliced[key] = value[token_slice]

    positions = model_inputs["positions"]
    # M-RoPE carries a leading section dim.
    sliced["positions"] = (
        positions[:, token_slice] if positions.ndim == 2 else positions[token_slice]
    )

    intermediate_tensors = model_inputs.get("intermediate_tensors")
    if intermediate_tensors is not None:
        sliced["intermediate_tensors"] = intermediate_tensors[token_slice]
    return sliced