vllm.v1.worker.gpu.sample.batch_shard

Classes:

• BatchShardMetadata –

  Collective layout for one sampling step.

• BatchSharder –

  Shards the sampler inputs across TP ranks along the batch dimension.

Functions:

• batch_sharded_sampling_enabled –

  Whether the target's sampling and verification are batch-sharded.

• shard_owner –

  TP rank that samples a request and holds its draft logits.

• shard_row –

  Row of a request in its owner's batch-sharded draft logits cache.

BatchShardMetadata dataclass

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).

Source code in vllm/v1/worker/gpu/sample/batch_shard.py

@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
    adaptive_verification_enabled: bool

BatchSharder

Shards the sampler inputs across TP ranks along the batch dimension.

Methods:

• shard_sampler_inputs –

  Owner-sort the batch and build this rank's local sampler inputs.

Source code in vllm/v1/worker/gpu/sample/batch_shard.py

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,
        adaptive_verification_enabled: bool,
    ):
        tp_group = get_tp_group()
        self.tp_rank = tp_group.rank_in_group
        self.tp_size = tp_group.world_size
        self.device = device
        self.adaptive_verification_enabled = adaptive_verification_enabled
        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]

        # 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
        # Get the number of logits per request on the host.
        num_logits_per_req_np = np.diff(input_batch.cu_num_logits_np)
        num_admitted_draft_tokens_np = None
        if input_batch.num_admitted_draft_tokens is not None:
            # Under adaptive verification, cu_num_logits_np holds the scheduled counts.
            # When the budget binds, the admitted counts are copied to the host so the
            # plan matches the device layout.
            assert input_batch.num_draft_tokens_per_req is not None
            num_admitted_draft_tokens_np = input_batch.num_admitted_draft_tokens.wait()
            num_logits_per_req_np = (
                num_logits_per_req_np
                - input_batch.num_draft_tokens_per_req
                + num_admitted_draft_tokens_np
            )
        num_logits = int(num_logits_per_req_np.sum())
        # Derive the number of logits owned by each rank, as well as the
        # local rank.
        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
        local_seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound[
            torch.from_numpy(local_req_indices_np)
        ]

        # 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 = input_batch.idx_mapping.new_empty(num_local_reqs)
        local_cu_num_logits = input_batch.cu_num_logits.new_empty(num_local_reqs + 1)
        local_expanded_idx_mapping = input_batch.expanded_idx_mapping.new_empty(
            num_local_logits
        )
        local_expanded_local_pos = input_batch.expanded_local_pos.new_empty(
            num_local_logits
        )
        local_seq_lens = input_batch.seq_lens.new_empty(num_local_reqs)
        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
            ]
            if num_admitted_draft_tokens_np is not None:
                # Under adaptive verification num_draft_tokens_per_req holds the
                # scheduled counts. Follow the admitted counts so every rank picks
                # the same sampler as the replicated path, including a budget of 0.
                num_draft_tokens = int(
                    num_admitted_draft_tokens_np[local_req_indices_np].sum()
                )
            elif input_batch.num_draft_tokens > 0:
                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,
            seq_lens_cpu_upper_bound=local_seq_lens_cpu_upper_bound,
            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, 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,
            adaptive_verification_enabled=self.adaptive_verification_enabled,
        )
        return local_batch, sorted_logits_indices, local_grammar_output, metadata

shard_sampler_inputs(input_batch, grammar_output)

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).

Source code in vllm/v1/worker/gpu/sample/batch_shard.py

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]

    # 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
    # Get the number of logits per request on the host.
    num_logits_per_req_np = np.diff(input_batch.cu_num_logits_np)
    num_admitted_draft_tokens_np = None
    if input_batch.num_admitted_draft_tokens is not None:
        # Under adaptive verification, cu_num_logits_np holds the scheduled counts.
        # When the budget binds, the admitted counts are copied to the host so the
        # plan matches the device layout.
        assert input_batch.num_draft_tokens_per_req is not None
        num_admitted_draft_tokens_np = input_batch.num_admitted_draft_tokens.wait()
        num_logits_per_req_np = (
            num_logits_per_req_np
            - input_batch.num_draft_tokens_per_req
            + num_admitted_draft_tokens_np
        )
    num_logits = int(num_logits_per_req_np.sum())
    # Derive the number of logits owned by each rank, as well as the
    # local rank.
    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
    local_seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound[
        torch.from_numpy(local_req_indices_np)
    ]

    # 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 = input_batch.idx_mapping.new_empty(num_local_reqs)
    local_cu_num_logits = input_batch.cu_num_logits.new_empty(num_local_reqs + 1)
    local_expanded_idx_mapping = input_batch.expanded_idx_mapping.new_empty(
        num_local_logits
    )
    local_expanded_local_pos = input_batch.expanded_local_pos.new_empty(
        num_local_logits
    )
    local_seq_lens = input_batch.seq_lens.new_empty(num_local_reqs)
    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
        ]
        if num_admitted_draft_tokens_np is not None:
            # Under adaptive verification num_draft_tokens_per_req holds the
            # scheduled counts. Follow the admitted counts so every rank picks
            # the same sampler as the replicated path, including a budget of 0.
            num_draft_tokens = int(
                num_admitted_draft_tokens_np[local_req_indices_np].sum()
            )
        elif input_batch.num_draft_tokens > 0:
            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,
        seq_lens_cpu_upper_bound=local_seq_lens_cpu_upper_bound,
        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, 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,
        adaptive_verification_enabled=self.adaptive_verification_enabled,
    )
    return local_batch, sorted_logits_indices, local_grammar_output, metadata

batch_sharded_sampling_enabled(vllm_config, target_model)

Whether the target's sampling and verification are batch-sharded.

Source code in vllm/v1/worker/gpu/sample/batch_shard.py

def batch_sharded_sampling_enabled(
    vllm_config: VllmConfig, target_model: nn.Module
) -> bool:
    """Whether the target's sampling and verification are batch-sharded."""
    return bool(
        vllm_config.parallel_config.enable_batch_sharded_sampling
        and hasattr(target_model, "compute_logits_local")
    )

shard_owner(req_state_idx, NUM_SHARDS)

TP rank that samples a request and holds its draft logits.

Mirrored on the host in BatchSharder.shard_sampler_inputs.

Source code in vllm/v1/worker/gpu/sample/batch_shard.py

@triton.jit
def shard_owner(req_state_idx, NUM_SHARDS: tl.constexpr):
    """TP rank that samples a request and holds its draft logits.

    Mirrored on the host in `BatchSharder.shard_sampler_inputs`.
    """
    return req_state_idx % NUM_SHARDS

shard_row(req_state_idx, NUM_SHARDS)

Row of a request in its owner's batch-sharded draft logits cache.

Source code in vllm/v1/worker/gpu/sample/batch_shard.py

@triton.jit
def shard_row(req_state_idx, NUM_SHARDS: tl.constexpr):
    """Row of a request in its owner's batch-sharded draft logits cache."""
    return req_state_idx // NUM_SHARDS