Skip to content

vllm.v1.worker.gpu.sample.output

Classes:

  • SamplingMaskTensors –

    Device-side masks pending async D2H: compact ids, plus the bitmask as

SamplingMaskTensors

Bases: NamedTuple

Device-side masks pending async D2H: compact ids, plus the bitmask as the exact fallback for rows wider than max_num_kept.

Methods:

  • cat –

    Join request-contiguous chunks that share one output layout.

  • from_logits –

    Capture committed target supports in fixed request-major slots.

  • tolists –

    Convert fixed output rows to the scheduler's CSR representation.

Source code in vllm/v1/worker/gpu/sample/output.py 
 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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
class SamplingMaskTensors(NamedTuple):
    """Device-side masks pending async D2H: compact ids, plus the bitmask as
    the exact fallback for rows wider than ``max_num_kept``."""

    # [num_requests * rows_per_request, max_num_kept]
    token_ids: torch.Tensor
    # [num_requests * rows_per_request, ceil(vocab_size / 8)]
    packed_mask: torch.Tensor
    # [num_requests * rows_per_request]
    counts: torch.Tensor
    vocab_size: int
    # Fixed request-major output slots. Ordinary sampling has one slot.
    rows_per_request: int = 1

    @classmethod
    def from_logits(
        cls,
        logits: torch.Tensor,
        cu_num_logits: torch.Tensor,
        num_sampled_tokens: torch.Tensor,
        max_num_kept: int,
        rows_per_request: int = 1,
    ) -> SamplingMaskTensors:
        """Capture committed target supports in fixed request-major slots."""
        num_reqs = num_sampled_tokens.shape[0]
        vocab_size = logits.shape[1]
        max_num_kept = min(max_num_kept, vocab_size, MAX_COMPACT_SUPPORT)
        num_output_rows = num_reqs * rows_per_request
        device = logits.device

        token_ids = torch.empty(
            (num_output_rows, max_num_kept), dtype=torch.int32, device=device
        )
        packed_mask = torch.empty(
            (num_output_rows, (vocab_size + 7) // 8),
            dtype=torch.uint8,
            device=device,
        )
        counts = torch.empty(num_output_rows, dtype=torch.int32, device=device)
        _compact_sampling_mask_kernel[(num_output_rows,)](
            logits,
            logits.stride(0),
            logits.stride(1),
            cu_num_logits,
            num_sampled_tokens,
            token_ids,
            token_ids.stride(0),
            packed_mask,
            packed_mask.stride(0),
            counts,
            vocab_size,
            max_num_kept,
            ROWS_PER_REQUEST=rows_per_request,
            BLOCK_SIZE=8192,
        )
        return cls(token_ids, packed_mask, counts, vocab_size, rows_per_request)

    @classmethod
    def cat(cls, chunks: Sequence[SamplingMaskTensors]) -> SamplingMaskTensors:
        """Join request-contiguous chunks that share one output layout."""
        first = chunks[0]
        assert len({(c.vocab_size, c.rows_per_request) for c in chunks}) == 1
        return cls(
            torch.cat([chunk.token_ids for chunk in chunks]),
            torch.cat([chunk.packed_mask for chunk in chunks]),
            torch.cat([chunk.counts for chunk in chunks]),
            first.vocab_size,
            first.rows_per_request,
        )

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

    def tolists(
        self, num_sampled_tokens: np.ndarray | None = None
    ) -> SamplingMaskLists:
        """Convert fixed output rows to the scheduler's CSR representation."""
        counts = self.counts.cpu().numpy()
        token_ids = self.token_ids.cpu().numpy()
        packed_mask = self.packed_mask.cpu().numpy()
        width = token_ids.shape[1]

        cu_num_generated_tokens = None
        if self.rows_per_request == 1:
            sampled_rows = np.arange(len(counts))
        else:
            assert num_sampled_tokens is not None
            active_slots = (
                np.arange(self.rows_per_request)[None, :] < num_sampled_tokens[:, None]
            )
            sampled_rows = np.flatnonzero(active_slots)
            cu_num_generated_tokens = np.cumsum(
                np.concatenate(([0], num_sampled_tokens))
            ).tolist()

        def support(row: int) -> np.ndarray:
            if counts[row] <= width:
                return token_ids[row, : counts[row]]
            # Wider than the compact row (ties or a huge top_k): use the bitmask.
            bits = np.unpackbits(
                packed_mask[row], count=self.vocab_size, bitorder="little"
            )
            return np.flatnonzero(bits).astype(np.int32, copy=False)

        supports = [support(row) for row in sampled_rows]
        offsets = np.zeros(len(supports) + 1, dtype=np.int64)
        np.cumsum([len(s) for s in supports], out=offsets[1:])
        flat = np.concatenate([np.empty(0, dtype=np.int32), *supports])
        return SamplingMaskLists(flat, offsets, cu_num_generated_tokens)

cat(chunks) classmethod

Join request-contiguous chunks that share one output layout.

Source code in vllm/v1/worker/gpu/sample/output.py 
143
144
145
146
147
148
149
150
151
152
153
154
@classmethod
def cat(cls, chunks: Sequence[SamplingMaskTensors]) -> SamplingMaskTensors:
    """Join request-contiguous chunks that share one output layout."""
    first = chunks[0]
    assert len({(c.vocab_size, c.rows_per_request) for c in chunks}) == 1
    return cls(
        torch.cat([chunk.token_ids for chunk in chunks]),
        torch.cat([chunk.packed_mask for chunk in chunks]),
        torch.cat([chunk.counts for chunk in chunks]),
        first.vocab_size,
        first.rows_per_request,
    )

from_logits(logits, cu_num_logits, num_sampled_tokens, max_num_kept, rows_per_request=1) classmethod

Capture committed target supports in fixed request-major slots.

Source code in vllm/v1/worker/gpu/sample/output.py 
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
132
133
134
135
136
137
138
139
140
141
@classmethod
def from_logits(
    cls,
    logits: torch.Tensor,
    cu_num_logits: torch.Tensor,
    num_sampled_tokens: torch.Tensor,
    max_num_kept: int,
    rows_per_request: int = 1,
) -> SamplingMaskTensors:
    """Capture committed target supports in fixed request-major slots."""
    num_reqs = num_sampled_tokens.shape[0]
    vocab_size = logits.shape[1]
    max_num_kept = min(max_num_kept, vocab_size, MAX_COMPACT_SUPPORT)
    num_output_rows = num_reqs * rows_per_request
    device = logits.device

    token_ids = torch.empty(
        (num_output_rows, max_num_kept), dtype=torch.int32, device=device
    )
    packed_mask = torch.empty(
        (num_output_rows, (vocab_size + 7) // 8),
        dtype=torch.uint8,
        device=device,
    )
    counts = torch.empty(num_output_rows, dtype=torch.int32, device=device)
    _compact_sampling_mask_kernel[(num_output_rows,)](
        logits,
        logits.stride(0),
        logits.stride(1),
        cu_num_logits,
        num_sampled_tokens,
        token_ids,
        token_ids.stride(0),
        packed_mask,
        packed_mask.stride(0),
        counts,
        vocab_size,
        max_num_kept,
        ROWS_PER_REQUEST=rows_per_request,
        BLOCK_SIZE=8192,
    )
    return cls(token_ids, packed_mask, counts, vocab_size, rows_per_request)

tolists(num_sampled_tokens=None)

Convert fixed output rows to the scheduler's CSR representation.

Source code in vllm/v1/worker/gpu/sample/output.py 
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def tolists(
    self, num_sampled_tokens: np.ndarray | None = None
) -> SamplingMaskLists:
    """Convert fixed output rows to the scheduler's CSR representation."""
    counts = self.counts.cpu().numpy()
    token_ids = self.token_ids.cpu().numpy()
    packed_mask = self.packed_mask.cpu().numpy()
    width = token_ids.shape[1]

    cu_num_generated_tokens = None
    if self.rows_per_request == 1:
        sampled_rows = np.arange(len(counts))
    else:
        assert num_sampled_tokens is not None
        active_slots = (
            np.arange(self.rows_per_request)[None, :] < num_sampled_tokens[:, None]
        )
        sampled_rows = np.flatnonzero(active_slots)
        cu_num_generated_tokens = np.cumsum(
            np.concatenate(([0], num_sampled_tokens))
        ).tolist()

    def support(row: int) -> np.ndarray:
        if counts[row] <= width:
            return token_ids[row, : counts[row]]
        # Wider than the compact row (ties or a huge top_k): use the bitmask.
        bits = np.unpackbits(
            packed_mask[row], count=self.vocab_size, bitorder="little"
        )
        return np.flatnonzero(bits).astype(np.int32, copy=False)

    supports = [support(row) for row in sampled_rows]
    offsets = np.zeros(len(supports) + 1, dtype=np.int64)
    np.cumsum([len(s) for s in supports], out=offsets[1:])
    flat = np.concatenate([np.empty(0, dtype=np.int32), *supports])
    return SamplingMaskLists(flat, offsets, cu_num_generated_tokens)

_compact_sampling_mask_kernel(logits_ptr, logits_row_stride, logits_col_stride, cu_num_logits_ptr, num_sampled_tokens_ptr, token_ids_ptr, token_ids_row_stride, packed_mask_ptr, packed_mask_row_stride, counts_ptr, vocab_size, max_num_kept, ROWS_PER_REQUEST, BLOCK_SIZE)

Per row: first max_num_kept finite-logit ids, the full count, the bitmask.

Source code in vllm/v1/worker/gpu/sample/output.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
@triton.jit
def _compact_sampling_mask_kernel(
    logits_ptr,
    logits_row_stride,
    logits_col_stride,
    cu_num_logits_ptr,
    num_sampled_tokens_ptr,
    token_ids_ptr,
    token_ids_row_stride,
    packed_mask_ptr,
    packed_mask_row_stride,
    counts_ptr,
    vocab_size,
    max_num_kept,
    ROWS_PER_REQUEST: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    """Per row: first ``max_num_kept`` finite-logit ids, the full count, the bitmask."""
    output_row = tl.program_id(0)
    req_idx = output_row // ROWS_PER_REQUEST
    slot_idx = output_row % ROWS_PER_REQUEST
    if slot_idx >= tl.load(num_sampled_tokens_ptr + req_idx):
        tl.store(counts_ptr + output_row, 0)
        return
    source_row = tl.load(cu_num_logits_ptr + req_idx) + slot_idx

    count = tl.zeros((), dtype=tl.int32)

    for start_idx in range(0, vocab_size, BLOCK_SIZE):
        offsets = start_idx + tl.arange(0, BLOCK_SIZE)
        logits = tl.load(
            logits_ptr + source_row * logits_row_stride + offsets * logits_col_stride,
            mask=offsets < vocab_size,
            other=-float("inf"),
        )
        keep = (logits > -float("inf")) & (logits < float("inf"))
        keep_i32 = keep.to(tl.int32)
        pos = count + tl.cumsum(keep_i32, axis=0) - keep_i32
        tl.store(
            token_ids_ptr + output_row * token_ids_row_stride + pos,
            offsets.to(tl.int32),
            mask=keep & (pos < max_num_kept),
        )
        count += tl.sum(keep_i32)

        bits = tl.reshape(keep_i32, (BLOCK_SIZE // 8, 8)) << tl.arange(0, 8)[None, :]
        byte_offsets = start_idx // 8 + tl.arange(0, BLOCK_SIZE // 8)
        tl.store(
            packed_mask_ptr + output_row * packed_mask_row_stride + byte_offsets,
            tl.sum(bits, axis=1).to(tl.uint8),
            mask=byte_offsets < tl.cdiv(vocab_size, 8),
        )

    tl.store(counts_ptr + output_row, count)