vllm.v1.worker.gpu.sample.batch_shard ¶

@dataclass
class BatchShardMetadata:
    """Collective layout for one sampling step.

    Requests (and their logits) are owner-sorted: all requests owned by
    rank 0 first, then rank 1, etc. Within an owner, batch order is
    preserved (stable sort). Every field is a pure function of the
    replicated idx_mapping and cu_num_logits, so all ranks build identical
    plans without communication — which relies on request slots
    (idx_mapping) being assigned identically on every rank (see
    finish_requests in the model runner).
    """

    tp_size: int
    # Per-owner-rank logits counts (all-to-all send splits).
    num_logits_per_rank: list[int]
    num_local_logits: int
    num_local_reqs: int
    # The number of request entries contributed by each rank to the
    # gathered result (shorter shards are padded up to this value).
    max_num_reqs_per_rank: int
    # [num_reqs] For batch request i, its source index in the rank-major
    # gathered results. Indexing the gathered tensors with this restores
    # the original batch order.
    gathered_src_indices: torch.Tensor
    # The number of entries occupied by each request in the gathered
    # tensor. Derived from the global cu_num_logits, so each rank agrees
    # ont he gather shapes, including empty shards.
    max_num_logits_per_req: int

class BatchSharder:
    """Shards the sampler inputs across TP ranks along the batch dimension."""

    def __init__(
        self,
        max_num_reqs: int,
        max_num_logits_per_req: int,
        device: torch.device,
    ):
        tp_group = get_tp_group()
        self.tp_rank = tp_group.rank_in_group
        self.tp_size = tp_group.world_size
        self.device = device
        self._padded_num_reqs = triton.next_power_of_2(max_num_reqs)
        self._padded_num_logits_per_req = triton.next_power_of_2(max_num_logits_per_req)
        self._num_warps = max(1, min(8, self._padded_num_reqs // 128))

    def shard_sampler_inputs(
        self,
        input_batch: InputBatch,
        grammar_output: GrammarOutput | None,
    ) -> tuple[InputBatch, torch.Tensor, GrammarOutput | None, BatchShardMetadata]:
        """Owner-sort the batch and build this rank's local sampler inputs.

        Returns the local sub-batch, the owner-sorted logits_indices (gather
        the sampling hidden states with these so `compute_logits_local` emits
        logits in all-to-all send order), the local grammar output (None if
        this rank owns none of the structured-output requests), and the shard
        metadata (used for all-gathering the sampler outputs).
        """
        tp_rank = self.tp_rank
        tp_size = self.tp_size
        num_reqs = input_batch.idx_mapping_np.shape[0]
        num_logits = int(input_batch.cu_num_logits_np[-1])

        # Deterministically assign requests to ranks, round-robin over slot
        # indices so ownership stays balanced when the slot allocator fills
        # low slots first (partial occupancy).
        # NOTE: This mirrors the assignment used in _build_shard_plan_kernel.
        req_owner_np = input_batch.idx_mapping_np % tp_size
        local_req_indices_np = np.flatnonzero(req_owner_np == tp_rank)
        local_idx_mapping_np = input_batch.idx_mapping_np[local_req_indices_np]
        num_local_reqs = local_req_indices_np.shape[0]
        num_reqs_per_rank_np = np.bincount(req_owner_np, minlength=tp_size)
        max_num_reqs_per_rank = int(num_reqs_per_rank_np.max()) if num_reqs else 1
        # Derive the number of logits owned by each rank, as well as the
        # local rank.
        num_logits_per_req_np = np.diff(input_batch.cu_num_logits_np)
        num_logits_per_rank_np = np.bincount(
            req_owner_np, weights=num_logits_per_req_np, minlength=tp_size
        ).astype(np.int64)
        num_local_logits = int(num_logits_per_rank_np[tp_rank])
        local_logits_start = int(num_logits_per_rank_np[:tp_rank].sum())
        local_cu_num_logits_np = np.zeros(num_local_reqs + 1, dtype=np.int32)
        np.cumsum(
            num_logits_per_req_np[local_req_indices_np], out=local_cu_num_logits_np[1:]
        )
        max_num_logits_per_req = int(num_logits_per_req_np.max()) if num_reqs else 1

        # Shard the input batch GPU tensors.
        sorted_logits_indices = torch.empty(
            num_logits, dtype=torch.int64, device=self.device
        )
        gathered_src_indices = torch.empty(
            num_reqs, dtype=torch.int64, device=self.device
        )
        local_logits_indices = torch.empty(
            num_local_logits, dtype=torch.int64, device=self.device
        )
        local_idx_mapping = torch.empty(
            num_local_reqs, dtype=torch.int32, device=self.device
        )
        local_cu_num_logits = torch.empty(
            num_local_reqs + 1, dtype=torch.int32, device=self.device
        )
        local_expanded_idx_mapping = torch.empty(
            num_local_logits, dtype=torch.int32, device=self.device
        )
        local_expanded_local_pos = torch.empty(
            num_local_logits, dtype=torch.int32, device=self.device
        )
        local_seq_lens = torch.empty(
            num_local_reqs, dtype=torch.int32, device=self.device
        )
        if num_reqs > 0:
            _build_shard_plan_kernel[(num_reqs,)](
                input_batch.idx_mapping,
                input_batch.cu_num_logits,
                input_batch.query_start_loc,
                input_batch.seq_lens,
                sorted_logits_indices,
                gathered_src_indices,
                local_idx_mapping,
                local_cu_num_logits,
                local_logits_indices,
                local_expanded_idx_mapping,
                local_expanded_local_pos,
                local_seq_lens,
                num_reqs,
                local_logits_start,
                max_num_reqs_per_rank,
                TP_SIZE=tp_size,
                TP_RANK=tp_rank,
                PADDED_NUM_REQS=self._padded_num_reqs,
                PADDED_NUM_LOGITS_PER_REQ=self._padded_num_logits_per_req,
                num_warps=self._num_warps,
            )

        # Compute the local number of draft tokens.
        num_draft_tokens_per_req = None
        num_draft_tokens = 0
        if input_batch.num_draft_tokens_per_req is not None:
            num_draft_tokens_per_req = input_batch.num_draft_tokens_per_req[
                local_req_indices_np
            ]
            num_draft_tokens = int(num_draft_tokens_per_req.sum())

        local_req_ids = [input_batch.req_ids[i] for i in local_req_indices_np.tolist()]
        local_batch = replace(
            input_batch,
            req_ids=local_req_ids,
            num_reqs=num_local_reqs,
            idx_mapping=local_idx_mapping,
            idx_mapping_np=local_idx_mapping_np,
            expanded_idx_mapping=local_expanded_idx_mapping,
            expanded_local_pos=local_expanded_local_pos,
            seq_lens=local_seq_lens,
            logits_indices=local_logits_indices,
            cu_num_logits=local_cu_num_logits,
            cu_num_logits_np=local_cu_num_logits_np,
            num_draft_tokens=num_draft_tokens,
            num_draft_tokens_per_req=num_draft_tokens_per_req,
        )
        local_grammar_output = None
        if grammar_output is not None:
            local_grammar_output = _shard_grammar_output(
                grammar_output, input_batch, local_req_ids
            )
        metadata = BatchShardMetadata(
            tp_size=tp_size,
            num_logits_per_rank=num_logits_per_rank_np.tolist(),
            num_local_logits=num_local_logits,
            num_local_reqs=num_local_reqs,
            max_num_reqs_per_rank=max_num_reqs_per_rank,
            gathered_src_indices=gathered_src_indices,
            max_num_logits_per_req=max_num_logits_per_req,
        )
        return local_batch, sorted_logits_indices, local_grammar_output, metadata