vLLM/vllm/v1/worker/gpu/sample/batch_shard.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

from dataclasses import dataclass, replace

import numpy as np
import torch
import torch.distributed

from vllm.distributed.communication_op import tensor_model_parallel_all_gather
from vllm.distributed.parallel_state import get_tp_group
from vllm.triton_utils import tl, triton
from vllm.v1.core.sched.output import GrammarOutput
from vllm.v1.outputs import LogprobsTensors
from vllm.v1.worker.gpu.input_batch import InputBatch
from vllm.v1.worker.gpu.sample.output import SamplerOutput


@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


@triton.jit(
    do_not_specialize=["num_reqs", "local_logits_start", "max_num_reqs_per_rank"]
)
def _build_shard_plan_kernel(
    idx_mapping_ptr,
    cu_num_logits_ptr,
    query_start_loc_ptr,
    seq_lens_ptr,
    sorted_logits_indices_ptr,
    gathered_src_indices_ptr,
    local_idx_mapping_ptr,
    local_cu_num_logits_ptr,
    local_logits_indices_ptr,
    local_expanded_idx_mapping_ptr,
    local_expanded_local_pos_ptr,
    local_seq_lens_ptr,
    num_reqs,
    local_logits_start,
    max_num_reqs_per_rank,
    TP_SIZE: tl.constexpr,
    TP_RANK: tl.constexpr,
    PADDED_NUM_REQS: tl.constexpr,
    PADDED_NUM_LOGITS_PER_REQ: tl.constexpr,
):
    req_idx = tl.program_id(0)
    req_state_idx = tl.load(idx_mapping_ptr + req_idx)
    owner = req_state_idx % TP_SIZE
    num_logits = tl.load(cu_num_logits_ptr + req_idx + 1) - tl.load(
        cu_num_logits_ptr + req_idx
    )

    req_block = tl.arange(0, PADDED_NUM_REQS)
    req_mask = req_block < num_reqs
    owners = tl.load(idx_mapping_ptr + req_block, mask=req_mask, other=0) % TP_SIZE
    num_logits_per_req = tl.load(
        cu_num_logits_ptr + req_block + 1, mask=req_mask, other=0
    ) - tl.load(cu_num_logits_ptr + req_block, mask=req_mask, other=0)

    # Owner-sorted layout is rank 0's requests in batch order, then rank 1's,
    # etc, so a request's starting offset can be derived from counting.
    lower = (owners < owner) & req_mask
    earlier = (owners == owner) & (req_block < req_idx) & req_mask

    # Compute the logits row offset for the current request in the sorted
    # batch.
    req_logits_start = tl.sum(tl.where(lower, num_logits_per_req, 0)) + tl.sum(
        tl.where(earlier, num_logits_per_req, 0)
    )

    # Copy the current request logit indices into the sorted batch.
    query_start = tl.load(query_start_loc_ptr + req_idx + 1) - num_logits
    logit_block = tl.arange(0, PADDED_NUM_LOGITS_PER_REQ)
    logit_mask = logit_block < num_logits
    logits_indices = (query_start + logit_block).to(tl.int64)
    tl.store(
        sorted_logits_indices_ptr + req_logits_start + logit_block,
        logits_indices,
        mask=logit_mask,
    )

    # Write the mapping of output tensor request indices => gathered tensor
    # request indices.
    # Batch request i's result lives in the padded, gathered results at
    # rank_of(i) * max_num_reqs_per_rank + offset_within_rank(i).
    local_req_idx = tl.sum(earlier.to(tl.int32))
    tl.store(
        gathered_src_indices_ptr + req_idx,
        (owner * max_num_reqs_per_rank + local_req_idx).to(tl.int64),
    )

    if req_idx == 0:
        tl.store(local_cu_num_logits_ptr, 0)

    if owner == TP_RANK:
        local_start = req_logits_start - local_logits_start
        tl.store(local_idx_mapping_ptr + local_req_idx, req_state_idx)
        tl.store(local_cu_num_logits_ptr + local_req_idx + 1, local_start + num_logits)
        tl.store(
            local_logits_indices_ptr + local_start + logit_block,
            logits_indices,
            mask=logit_mask,
        )
        tl.store(
            local_expanded_idx_mapping_ptr + local_start + logit_block,
            req_state_idx,
            mask=logit_mask,
        )
        tl.store(
            local_expanded_local_pos_ptr + local_start + logit_block,
            logit_block,
            mask=logit_mask,
        )
        tl.store(local_seq_lens_ptr + local_req_idx, tl.load(seq_lens_ptr + req_idx))


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


def all_to_all_logits(
    logits_shard: torch.Tensor,
    metadata: BatchShardMetadata,
) -> torch.Tensor:
    recv = logits_shard.new_empty(
        metadata.tp_size * metadata.num_local_logits, logits_shard.shape[-1]
    )
    torch.distributed.all_to_all_single(
        recv,
        logits_shard.contiguous(),
        output_split_sizes=[metadata.num_local_logits] * metadata.tp_size,
        input_split_sizes=metadata.num_logits_per_rank,
        group=get_tp_group().device_group,
    )
    shard_width = recv.shape[-1]
    return (
        recv.view(metadata.tp_size, metadata.num_local_logits, shard_width)
        .permute(1, 0, 2)
        .reshape(metadata.num_local_logits, metadata.tp_size * shard_width)
    )


def _shard_grammar_output(
    grammar_output: GrammarOutput,
    input_batch: InputBatch,
    local_req_ids: list[str],
) -> GrammarOutput | None:
    owned = set(local_req_ids)
    req_id_to_idx = {req_id: i for i, req_id in enumerate(input_batch.req_ids)}
    cu_num_logits_np = input_batch.cu_num_logits_np

    local_ids: list[str] = []
    keep_indices: list[int] = []
    cursor = 0
    for req_id in grammar_output.structured_output_request_ids:
        req_idx = req_id_to_idx[req_id]
        num_req_logits = int(cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx])
        if req_id in owned:
            local_ids.append(req_id)
            keep_indices.extend(range(cursor, cursor + num_req_logits))
        cursor += num_req_logits
    if not local_ids:
        return None
    return GrammarOutput(
        structured_output_request_ids=local_ids,
        grammar_bitmask=grammar_output.grammar_bitmask[keep_indices],
    )


@triton.jit(
    do_not_specialize=["max_num_logits_per_req", "num_src_cols", "packed_stride"]
)
def _pack_sampler_output_kernel(
    packed_ptr,
    packed_stride,
    sampled_token_ids_ptr,
    sampled_token_ids_stride,
    num_sampled_ptr,
    num_rejected_ptr,
    num_nans_ptr,
    local_cu_num_logits_ptr,
    max_num_logits_per_req,
    num_src_cols,
    BLOCK_SIZE: tl.constexpr,
):
    req_idx = tl.program_id(0)
    cols = tl.arange(0, BLOCK_SIZE)
    row_ptr = packed_ptr + req_idx * packed_stride

    token_ids = tl.load(
        sampled_token_ids_ptr + req_idx * sampled_token_ids_stride + cols,
        mask=cols < num_src_cols,
        other=0,
    )
    tl.store(row_ptr + cols, token_ids, mask=cols < max_num_logits_per_req)
    num_sampled = tl.load(num_sampled_ptr + req_idx)
    tl.store(row_ptr + max_num_logits_per_req, num_sampled.to(tl.int64))
    num_rejected = tl.load(num_rejected_ptr + req_idx)
    tl.store(row_ptr + max_num_logits_per_req + 1, num_rejected.to(tl.int64))
    if num_nans_ptr is not None:
        # num_nans is per logits row; reduce it per request. The replicated
        # samplers report per-row counts that downstream zips per request,
        # so a per-request sum is the well-defined equivalent.
        start = tl.load(local_cu_num_logits_ptr + req_idx)
        num_req_logits = tl.load(local_cu_num_logits_ptr + req_idx + 1) - start
        nans = tl.load(num_nans_ptr + start + cols, mask=cols < num_req_logits, other=0)
        tl.store(
            row_ptr + max_num_logits_per_req + 2, tl.sum(nans.to(tl.int64), axis=0)
        )


@triton.jit(do_not_specialize=["max_num_logits_per_req", "gathered_stride"])
def _unpack_gathered_output_kernel(
    gathered_ptr,
    gathered_stride,
    gathered_src_indices_ptr,
    sampled_token_ids_ptr,
    num_sampled_ptr,
    num_rejected_ptr,
    num_nans_ptr,
    max_num_logits_per_req,
    BLOCK_SIZE: tl.constexpr,
):
    req_idx = tl.program_id(0)
    src = tl.load(gathered_src_indices_ptr + req_idx)
    row_ptr = gathered_ptr + src * gathered_stride
    cols = tl.arange(0, BLOCK_SIZE)
    mask = cols < max_num_logits_per_req

    token_ids = tl.load(row_ptr + cols, mask=mask)
    tl.store(
        sampled_token_ids_ptr + req_idx * max_num_logits_per_req + cols,
        token_ids,
        mask=mask,
    )
    num_sampled = tl.load(row_ptr + max_num_logits_per_req)
    tl.store(num_sampled_ptr + req_idx, num_sampled.to(tl.int32))
    num_rejected = tl.load(row_ptr + max_num_logits_per_req + 1)
    tl.store(num_rejected_ptr + req_idx, num_rejected.to(tl.int32))
    if num_nans_ptr is not None:
        num_nans = tl.load(row_ptr + max_num_logits_per_req + 2)
        tl.store(num_nans_ptr + req_idx, num_nans.to(tl.int32))


@triton.jit(
    do_not_specialize=[
        "max_num_logits_per_req",
        "num_logprob_cols",
        "num_src_cols",
        "send_ids_stride",
        "send_logprobs_stride",
        "src_ids_stride",
        "src_logprobs_stride",
    ]
)
def _pack_logprobs_kernel(
    send_ids_ptr,
    send_ids_stride,
    send_logprobs_ptr,
    send_logprobs_stride,
    src_ids_ptr,
    src_ids_stride,
    src_logprobs_ptr,
    src_logprobs_stride,
    selected_token_ranks_ptr,
    local_cu_num_logits_ptr,
    max_num_logits_per_req,
    num_logprob_cols,
    num_src_cols,
    BLOCK_SIZE: tl.constexpr,
):
    req_idx = tl.program_id(0)
    pos = tl.program_id(1)
    start = tl.load(local_cu_num_logits_ptr + req_idx)
    if pos < tl.load(local_cu_num_logits_ptr + req_idx + 1) - start:
        src = start + pos
        dst = req_idx * max_num_logits_per_req + pos
        cols = tl.arange(0, BLOCK_SIZE)
        # A rank whose sub-batch has no draft tokens runs the regular sampler
        # while others run the rejection sampler, and their column counts can
        # differ (e.g. logprob_token_ids widening); the two masks pad or
        # truncate to the agreed width.
        src_mask = cols < num_src_cols
        dst_mask = cols < num_logprob_cols
        ids = tl.load(src_ids_ptr + src * src_ids_stride + cols, mask=src_mask, other=0)
        logprobs = tl.load(
            src_logprobs_ptr + src * src_logprobs_stride + cols,
            mask=src_mask,
            other=float("-inf"),
        )
        ids_row = send_ids_ptr + dst * send_ids_stride
        tl.store(ids_row + cols, ids.to(tl.int64), mask=dst_mask)
        rank = tl.load(selected_token_ranks_ptr + src)
        tl.store(ids_row + num_logprob_cols, rank.to(tl.int64))
        tl.store(
            send_logprobs_ptr + dst * send_logprobs_stride + cols,
            logprobs.to(tl.float32),
            mask=dst_mask,
        )


@triton.jit(
    do_not_specialize=[
        "max_num_logits_per_req",
        "num_logprob_cols",
        "gathered_ids_stride",
        "gathered_logprobs_stride",
    ]
)
def _unpack_logprobs_kernel(
    gathered_ids_ptr,
    gathered_ids_stride,
    gathered_logprobs_ptr,
    gathered_logprobs_stride,
    gathered_src_indices_ptr,
    cu_num_logits_ptr,
    logprob_token_ids_ptr,
    logprobs_ptr,
    selected_token_ranks_ptr,
    max_num_logits_per_req,
    num_logprob_cols,
    BLOCK_SIZE: tl.constexpr,
):
    req_idx = tl.program_id(0)
    pos = tl.program_id(1)
    start = tl.load(cu_num_logits_ptr + req_idx)
    if pos < tl.load(cu_num_logits_ptr + req_idx + 1) - start:
        src = tl.load(gathered_src_indices_ptr + req_idx) * max_num_logits_per_req + pos
        dst = start + pos
        cols = tl.arange(0, BLOCK_SIZE)
        mask = cols < num_logprob_cols
        src_ids_row = gathered_ids_ptr + src * gathered_ids_stride
        ids = tl.load(src_ids_row + cols, mask=mask)
        tl.store(logprob_token_ids_ptr + dst * num_logprob_cols + cols, ids, mask=mask)
        rank = tl.load(src_ids_row + num_logprob_cols)
        tl.store(selected_token_ranks_ptr + dst, rank)
        logprobs = tl.load(
            gathered_logprobs_ptr + src * gathered_logprobs_stride + cols, mask=mask
        )
        tl.store(logprobs_ptr + dst * num_logprob_cols + cols, logprobs, mask=mask)


def _gather_logprobs_tensors(
    local_output: SamplerOutput | None,
    metadata: BatchShardMetadata,
    global_batch: InputBatch,
    local_batch: InputBatch,
    logprobs_dims: tuple[int, int],
    device: torch.device,
) -> LogprobsTensors:
    num_logprobs, max_token_ids = logprobs_dims
    num_logprob_cols = 1 + max(num_logprobs, max_token_ids)
    max_num_logits_per_req = metadata.max_num_logits_per_req
    num_send_rows = metadata.max_num_reqs_per_rank * max_num_logits_per_req
    block_size = triton.next_power_of_2(num_logprob_cols)

    # Padded slots are never read back: gathered_src_indices only addresses
    # rows that an owner rank wrote.
    send_ids_ranks = torch.empty(
        num_send_rows, num_logprob_cols + 1, dtype=torch.int64, device=device
    )
    send_logprobs = torch.empty(
        num_send_rows, num_logprob_cols, dtype=torch.float32, device=device
    )
    if local_output is not None and local_output.logprobs_tensors is not None:
        lp = local_output.logprobs_tensors
        _pack_logprobs_kernel[(metadata.num_local_reqs, max_num_logits_per_req)](
            send_ids_ranks,
            send_ids_ranks.stride(0),
            send_logprobs,
            send_logprobs.stride(0),
            lp.logprob_token_ids,
            lp.logprob_token_ids.stride(0),
            lp.logprobs,
            lp.logprobs.stride(0),
            lp.selected_token_ranks,
            local_batch.cu_num_logits,
            max_num_logits_per_req,
            num_logprob_cols,
            lp.logprob_token_ids.shape[1],
            BLOCK_SIZE=block_size,
        )

    gathered_ids_ranks = tensor_model_parallel_all_gather(send_ids_ranks, dim=0)
    gathered_logprobs = tensor_model_parallel_all_gather(send_logprobs, dim=0)

    num_reqs = global_batch.num_reqs
    num_logits = int(global_batch.cu_num_logits_np[-1])
    logprob_token_ids = torch.empty(
        num_logits, num_logprob_cols, dtype=torch.int64, device=device
    )
    logprobs = torch.empty(
        num_logits, num_logprob_cols, dtype=torch.float32, device=device
    )
    selected_token_ranks = torch.empty(num_logits, dtype=torch.int64, device=device)
    _unpack_logprobs_kernel[(num_reqs, max_num_logits_per_req)](
        gathered_ids_ranks,
        gathered_ids_ranks.stride(0),
        gathered_logprobs,
        gathered_logprobs.stride(0),
        metadata.gathered_src_indices,
        global_batch.cu_num_logits,
        logprob_token_ids,
        logprobs,
        selected_token_ranks,
        max_num_logits_per_req,
        num_logprob_cols,
        BLOCK_SIZE=block_size,
    )
    return LogprobsTensors(
        logprob_token_ids=logprob_token_ids,
        logprobs=logprobs,
        selected_token_ranks=selected_token_ranks,
        cu_num_generated_tokens=(
            global_batch.cu_num_logits_np.tolist() if num_logits != num_reqs else None
        ),
    )


def gather_sampler_output(
    local_output: SamplerOutput | None,
    metadata: BatchShardMetadata,
    device: torch.device,
    global_batch: InputBatch,
    local_batch: InputBatch,
    gather_num_nans: bool = False,
    logprobs_dims: tuple[int, int] | None = None,
) -> SamplerOutput:
    max_num_logits_per_req = metadata.max_num_logits_per_req
    num_packed_cols = max_num_logits_per_req + 2 + (1 if gather_num_nans else 0)
    block_size = triton.next_power_of_2(max_num_logits_per_req)

    # Pack the sampler output tensors (excluding logprobs) into a single
    # tensor for the all-gather.
    packed = torch.empty(
        metadata.max_num_reqs_per_rank,
        num_packed_cols,
        dtype=torch.int64,
        device=device,
    )
    if local_output is not None:
        assert local_output.num_sampled is not None
        assert local_output.num_rejected is not None
        assert not gather_num_nans or local_output.num_nans is not None
        num_src_cols = min(
            local_output.sampled_token_ids.shape[1], max_num_logits_per_req
        )
        _pack_sampler_output_kernel[(metadata.num_local_reqs,)](
            packed,
            packed.stride(0),
            local_output.sampled_token_ids,
            local_output.sampled_token_ids.stride(0),
            local_output.num_sampled,
            local_output.num_rejected,
            local_output.num_nans if gather_num_nans else None,
            local_batch.cu_num_logits if gather_num_nans else None,
            max_num_logits_per_req,
            num_src_cols,
            BLOCK_SIZE=block_size,
        )

    # All-gather the sampler output tensors (excluding logprobs).
    gathered = tensor_model_parallel_all_gather(packed, dim=0)
    # All-gather the logprobs tensors.
    logprobs_tensors = None
    if logprobs_dims is not None:
        logprobs_tensors = _gather_logprobs_tensors(
            local_output, metadata, global_batch, local_batch, logprobs_dims, device
        )

    # Unpack the gathered tensor into a sampler output object.
    num_reqs = metadata.gathered_src_indices.shape[0]
    sampled_token_ids = torch.empty(
        num_reqs, max_num_logits_per_req, dtype=torch.int64, device=device
    )
    num_sampled = torch.empty(num_reqs, dtype=torch.int32, device=device)
    num_rejected = torch.empty(num_reqs, dtype=torch.int32, device=device)
    num_nans = (
        torch.empty(num_reqs, dtype=torch.int32, device=device)
        if gather_num_nans
        else None
    )
    _unpack_gathered_output_kernel[(num_reqs,)](
        gathered,
        gathered.stride(0),
        metadata.gathered_src_indices,
        sampled_token_ids,
        num_sampled,
        num_rejected,
        num_nans,
        max_num_logits_per_req,
        BLOCK_SIZE=block_size,
    )
    return SamplerOutput(
        sampled_token_ids=sampled_token_ids,
        logprobs_tensors=logprobs_tensors,
        num_nans=num_nans,
        num_sampled=num_sampled,
        num_rejected=num_rejected,
    )