Skip to content

vllm.v1.worker.gpu.sample.output

Classes:

SamplingMaskTensors

Bases: NamedTuple

Bit-packed device-side sampling mask data pending async D2H.

Methods:

  • from_logits –

    Pack the finite-logit support for requests that sampled tokens.

  • tolists –

    Convert the packed masks to the scheduler's CSR representation.

Source code in vllm/v1/worker/gpu/sample/output.py 
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
class SamplingMaskTensors(NamedTuple):
    """Bit-packed device-side sampling mask data pending async D2H."""

    # [num_requests, ceil(vocab_size / 8)]
    packed_mask: torch.Tensor
    # [num_requests]
    counts: torch.Tensor
    vocab_size: int

    @classmethod
    def from_logits(
        cls,
        logits: torch.Tensor,
        num_sampled_tokens: torch.Tensor,
    ) -> SamplingMaskTensors:
        """Pack the finite-logit support for requests that sampled tokens."""
        num_reqs, vocab_size = logits.shape
        packed_width = (vocab_size + 7) // 8

        packed_mask = torch.empty(
            (num_reqs, packed_width), dtype=torch.uint8, device=logits.device
        )
        counts = torch.empty(num_reqs, dtype=torch.int32, device=logits.device)
        _pack_sampling_mask_kernel[(num_reqs,)](
            logits,
            logits.stride(0),
            logits.stride(1),
            num_sampled_tokens,
            packed_mask,
            packed_mask.stride(0),
            counts,
            vocab_size,
            BLOCK_SIZE=8192,
        )

        return cls(packed_mask, counts, vocab_size)

    def to_cpu_nonblocking(self) -> SamplingMaskTensors:
        if self.packed_mask.device.type == "cpu":
            return self
        return SamplingMaskTensors(
            self.packed_mask.to("cpu", non_blocking=True),
            self.counts.to("cpu", non_blocking=True),
            self.vocab_size,
        )

    def tolists(self, num_sampled_tokens: np.ndarray) -> SamplingMaskLists:
        """Convert the packed masks to the scheduler's CSR representation."""
        sampled_rows = np.flatnonzero(num_sampled_tokens)
        counts = self.counts.cpu().numpy()[sampled_rows]
        offsets = np.empty(len(counts) + 1, dtype=np.int64)
        offsets[0] = 0
        np.cumsum(counts, dtype=np.int64, out=offsets[1:])
        unpacked = np.unpackbits(
            self.packed_mask.cpu().numpy()[sampled_rows],
            axis=1,
            count=self.vocab_size,
            bitorder="little",
        )
        token_ids = np.nonzero(unpacked)[1].astype(np.int32, copy=False)
        return SamplingMaskLists(
            token_ids=token_ids,
            offsets=offsets,
            cu_num_generated_tokens=np.cumsum(
                np.concatenate(([0], num_sampled_tokens))
            ).tolist(),
        )

from_logits(logits, num_sampled_tokens) classmethod

Pack the finite-logit support for requests that sampled tokens.

Source code in vllm/v1/worker/gpu/sample/output.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
@classmethod
def from_logits(
    cls,
    logits: torch.Tensor,
    num_sampled_tokens: torch.Tensor,
) -> SamplingMaskTensors:
    """Pack the finite-logit support for requests that sampled tokens."""
    num_reqs, vocab_size = logits.shape
    packed_width = (vocab_size + 7) // 8

    packed_mask = torch.empty(
        (num_reqs, packed_width), dtype=torch.uint8, device=logits.device
    )
    counts = torch.empty(num_reqs, dtype=torch.int32, device=logits.device)
    _pack_sampling_mask_kernel[(num_reqs,)](
        logits,
        logits.stride(0),
        logits.stride(1),
        num_sampled_tokens,
        packed_mask,
        packed_mask.stride(0),
        counts,
        vocab_size,
        BLOCK_SIZE=8192,
    )

    return cls(packed_mask, counts, vocab_size)

tolists(num_sampled_tokens)

Convert the packed masks to the scheduler's CSR representation.

Source code in vllm/v1/worker/gpu/sample/output.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def tolists(self, num_sampled_tokens: np.ndarray) -> SamplingMaskLists:
    """Convert the packed masks to the scheduler's CSR representation."""
    sampled_rows = np.flatnonzero(num_sampled_tokens)
    counts = self.counts.cpu().numpy()[sampled_rows]
    offsets = np.empty(len(counts) + 1, dtype=np.int64)
    offsets[0] = 0
    np.cumsum(counts, dtype=np.int64, out=offsets[1:])
    unpacked = np.unpackbits(
        self.packed_mask.cpu().numpy()[sampled_rows],
        axis=1,
        count=self.vocab_size,
        bitorder="little",
    )
    token_ids = np.nonzero(unpacked)[1].astype(np.int32, copy=False)
    return SamplingMaskLists(
        token_ids=token_ids,
        offsets=offsets,
        cu_num_generated_tokens=np.cumsum(
            np.concatenate(([0], num_sampled_tokens))
        ).tolist(),
    )