gpu_input_batch.py 28.7 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
3
4
# Datastructures defining an input batch

from dataclasses import dataclass
5
from typing import Optional, cast
6
7
8
9

import numpy as np
import torch

10
from vllm.lora.request import LoRARequest
11
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
12
from vllm.sampling_params import SamplingParams, SamplingType
13
from vllm.utils import swap_dict_values
14
from vllm.v1.outputs import LogprobsTensors
15
from vllm.v1.sample.metadata import SamplingMetadata
16
from vllm.v1.utils import copy_slice
17
from vllm.v1.worker.block_table import BlockTable
18

19
20
_SAMPLING_EPS = 1e-5

21
22
23
24
25

@dataclass
class CachedRequestState:

    req_id: str
26
    prompt_token_ids: list[int]
27
    prompt: Optional[str]
28
    mm_inputs: list[MultiModalKwargs]
29
    mm_positions: list[PlaceholderRange]
30
31
32
    sampling_params: SamplingParams
    generator: Optional[torch.Generator]

33
    block_ids: list[int]
34
    num_computed_tokens: int
35
    output_token_ids: list[int]
36

37
38
39
    mrope_positions: Optional[torch.Tensor] = None
    mrope_position_delta: Optional[int] = None

40
41
    lora_request: Optional[LoRARequest] = None

42
43
44
45
46
47
48
49
50
51
52
53
54
55
    @property
    def num_tokens(self) -> int:
        return len(self.prompt_token_ids) + len(self.output_token_ids)


class InputBatch:

    def __init__(
        self,
        max_num_reqs: int,
        max_model_len: int,
        max_num_blocks_per_req: int,
        device: torch.device,
        pin_memory: bool,
56
        vocab_size: int,
57
58
59
60
61
62
    ):
        self.max_num_reqs = max_num_reqs
        self.max_model_len = max_model_len
        self.max_num_blocks_per_req = max_num_blocks_per_req
        self.device = device
        self.pin_memory = pin_memory
63
        self.vocab_size = vocab_size
64

65
66
        self._req_ids: list[Optional[str]] = []
        self.req_id_to_index: dict[str, int] = {}
67

68
69
        # TODO(woosuk): This buffer could be too large if max_model_len is big.
        # Find a way to reduce the CPU memory usage.
70
71
        # This buffer is not directly transferred to the GPU, so it does not
        # need to be pinned.
72
73
74
75
        self.token_ids_cpu_tensor = torch.zeros(
            (max_num_reqs, max_model_len),
            device="cpu",
            dtype=torch.int32,
76
            pin_memory=False,
77
78
        )
        self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()
79
        self.num_tokens = np.zeros(max_num_reqs, dtype=np.int32)
80
        self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
81
        self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
82
83
84
85
86
87
88
89
        self.num_computed_tokens_cpu_tensor = torch.zeros(
            (max_num_reqs, ),
            device="cpu",
            dtype=torch.int32,
            pin_memory=pin_memory,
        )
        self.num_computed_tokens_cpu = \
            self.num_computed_tokens_cpu_tensor.numpy()
90

91
92
93
94
        # Block table.
        self.block_table = BlockTable(
            max_num_reqs=max_num_reqs,
            max_num_blocks_per_req=max_num_blocks_per_req,
95
            pin_memory=pin_memory,
96
            device=device,
97
98
99
100
101
102
103
104
105
106
107
        )

        # Sampling-related.
        self.temperature = torch.empty((max_num_reqs, ),
                                       dtype=torch.float32,
                                       device=device)
        self.temperature_cpu_tensor = torch.empty((max_num_reqs, ),
                                                  dtype=torch.float32,
                                                  device="cpu",
                                                  pin_memory=pin_memory)
        self.temperature_cpu = self.temperature_cpu_tensor.numpy()
108
109
        self.greedy_reqs: set[str] = set()
        self.random_reqs: set[str] = set()
110
111
112
113
114
115
116
117
118

        self.top_p = torch.empty((max_num_reqs, ),
                                 dtype=torch.float32,
                                 device=device)
        self.top_p_cpu_tensor = torch.empty((max_num_reqs, ),
                                            dtype=torch.float32,
                                            device="cpu",
                                            pin_memory=pin_memory)
        self.top_p_cpu = self.top_p_cpu_tensor.numpy()
119
        self.top_p_reqs: set[str] = set()
120
121
122
123
124
125
126
127
128

        self.top_k = torch.empty((max_num_reqs, ),
                                 dtype=torch.int32,
                                 device=device)
        self.top_k_cpu_tensor = torch.empty((max_num_reqs, ),
                                            dtype=torch.int32,
                                            device="cpu",
                                            pin_memory=pin_memory)
        self.top_k_cpu = self.top_k_cpu_tensor.numpy()
129
        self.top_k_reqs: set[str] = set()
130

131
132
133
134
135
136
137
138
        self.min_p = torch.empty((max_num_reqs, ),
                                 dtype=torch.float32,
                                 device=device)
        self.min_p_cpu_tensor = torch.empty((max_num_reqs, ),
                                            dtype=torch.float32,
                                            device="cpu",
                                            pin_memory=pin_memory)
        self.min_p_cpu = self.min_p_cpu_tensor.numpy()
139
        self.min_p_reqs: set[str] = set()
140

141
142
143
144
145
146
147
148
149
150
        # Frequency penalty related data structures
        self.frequency_penalties = torch.empty((max_num_reqs, ),
                                               dtype=torch.float,
                                               device=device)
        self.frequency_penalties_cpu_tensor = torch.empty(
            (max_num_reqs, ),
            dtype=torch.float,
            device="cpu",
            pin_memory=pin_memory)
        self.frequency_penalties_cpu = \
151
            self.frequency_penalties_cpu_tensor.numpy()
152
        self.frequency_penalties_reqs: set[str] = set()
153
154
155
156
157
158
159
160
161

        # Presence penalty related data structures
        self.presence_penalties = torch.empty((max_num_reqs, ),
                                              dtype=torch.float,
                                              device=device)
        self.presence_penalties_cpu_tensor = torch.empty((max_num_reqs, ),
                                                         dtype=torch.float,
                                                         device="cpu",
                                                         pin_memory=pin_memory)
162
163
        self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy(
        )
164
        self.presence_penalties_reqs: set[str] = set()
165
166
167
168
169
170
171
172
173
174
175

        # Repetition penalty related data structures
        self.repetition_penalties = torch.empty((max_num_reqs, ),
                                                dtype=torch.float,
                                                device=device)
        self.repetition_penalties_cpu_tensor = torch.empty(
            (max_num_reqs, ),
            dtype=torch.float,
            device="cpu",
            pin_memory=pin_memory)
        self.repetition_penalties_cpu = \
176
            self.repetition_penalties_cpu_tensor.numpy()
177
        self.repetition_penalties_reqs: set[str] = set()
178

179
        # req_index -> (min_tokens, stop_token_ids)
180
        self.min_tokens: dict[int, tuple[int, set[int]]] = {}
181

182
183
184
        # lora related
        self.request_lora_mapping = np.zeros((self.max_num_reqs, ),
                                             dtype=np.int32)
185
186
        self.lora_id_to_request_ids: dict[int, set[str]] = {}
        self.lora_id_to_lora_request: dict[int, LoRARequest] = {}
187

188
        # req_index -> generator
189
190
        # NOTE(woosuk): The indices of the requests that do not have their own
        # generator should not be included in the dictionary.
191
        self.generators: dict[int, torch.Generator] = {}
192

193
        self.num_logprobs: dict[str, int] = {}
194
195
        # NOTE(rob): num_prompt_logprobs only includes reqs
        # that are currently in the prefill phase.
196
        self.num_prompt_logprobs: dict[str, int] = {}
197

198
199
200
        # To accumulate prompt logprobs tensor chunks across prefill steps.
        self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}

201
        self.logit_bias: list[Optional[dict[int,
202
                                            float]]] = [None] * max_num_reqs
203
        self.has_allowed_token_ids: set[str] = set()
204
205
        # NOTE(lufang): In the mask tensor, if the corresponding token allowed,
        # the value is False. Since we use masked_fill_ to set -inf.
206
207
        self.allowed_token_ids_mask: Optional[torch.Tensor] = None
        self.allowed_token_ids_mask_cpu_tensor: Optional[torch.Tensor] = None
208

209
210
211
        # req_index -> bad_words_token_ids
        self.bad_words_token_ids: dict[int, list[list[int]]] = {}

212
        self.req_output_token_ids: list[Optional[list[int]]] = []
213
214
215
216
217

        # This is updated each time the batch constituents change.
        self.sampling_metadata = self._make_sampling_metadata()

    @property
218
    def req_ids(self) -> list[str]:
219
220
        # None elements should only be present transiently
        # while performing state updates to the batch.
221
        return cast(list[str], self._req_ids)
222

223
224
225
226
227
228
229
230
231
232
    def add_request(
        self,
        request: "CachedRequestState",
        req_index: Optional[int] = None,
    ) -> None:
        if req_index is None:
            req_index = self.num_reqs
        assert req_index < self.max_num_reqs

        req_id = request.req_id
233
234
235
236
237
238
239
        if req_index == len(self._req_ids):
            self._req_ids.append(req_id)
            self.req_output_token_ids.append(request.output_token_ids)
        else:
            self._req_ids[req_index] = req_id
            self.req_output_token_ids[req_index] = request.output_token_ids

240
241
242
243
        self.req_id_to_index[req_id] = req_index

        # Copy the prompt token ids and output token ids.
        num_prompt_tokens = len(request.prompt_token_ids)
244
        self.num_prompt_tokens[req_index] = num_prompt_tokens
245
246
247
248
249
250
        self.token_ids_cpu[
            req_index, :num_prompt_tokens] = request.prompt_token_ids
        start_idx = num_prompt_tokens
        end_idx = start_idx + len(request.output_token_ids)
        self.token_ids_cpu[req_index,
                           start_idx:end_idx] = request.output_token_ids
251
252
        # Number of token ids in token_ids_cpu.
        # NOTE(woosuk): This may include spec decode tokens.
253
        self.num_tokens[req_index] = request.num_tokens
254
255
        # Number of tokens without spec decode tokens.
        self.num_tokens_no_spec[req_index] = request.num_tokens
256
257

        self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
258
        self.block_table.add_row(request.block_ids, req_index)
259
260
261

        sampling_params = request.sampling_params
        if sampling_params.sampling_type == SamplingType.GREEDY:
262
263
            # Avoid later division by zero.
            self.temperature_cpu[req_index] = -1.0
264
265
            self.greedy_reqs.add(req_id)
        else:
266
            self.temperature_cpu[req_index] = sampling_params.temperature
267
268
269
270
271
            self.random_reqs.add(req_id)

        self.top_p_cpu[req_index] = sampling_params.top_p
        if sampling_params.top_p < 1:
            self.top_p_reqs.add(req_id)
272
273
        top_k = sampling_params.top_k
        if 0 < top_k < self.vocab_size:
274
            self.top_k_reqs.add(req_id)
275
276
277
        else:
            top_k = self.vocab_size
        self.top_k_cpu[req_index] = top_k
278
        self.min_p_cpu[req_index] = sampling_params.min_p
279
280
        self.frequency_penalties_cpu[
            req_index] = sampling_params.frequency_penalty
281
282
        if sampling_params.min_p > _SAMPLING_EPS:
            self.min_p_reqs.add(req_id)
283
284
        if sampling_params.frequency_penalty != 0.0:
            self.frequency_penalties_reqs.add(req_id)
285
286
        self.presence_penalties_cpu[
            req_index] = sampling_params.presence_penalty
287
288
        if sampling_params.presence_penalty != 0.0:
            self.presence_penalties_reqs.add(req_id)
289
290
        self.repetition_penalties_cpu[
            req_index] = sampling_params.repetition_penalty
291
292
        if sampling_params.repetition_penalty != 1.0:
            self.repetition_penalties_reqs.add(req_id)
293
294
295
        if sampling_params.min_tokens:
            self.min_tokens[req_index] = (sampling_params.min_tokens,
                                          sampling_params.all_stop_token_ids)
296

297
298
299
300
        # NOTE(woosuk): self.generators should not include the requests that
        # do not have their own generator.
        if request.generator is not None:
            self.generators[req_index] = request.generator
301

302
303
304
305
        if sampling_params.logprobs is not None:
            self.num_logprobs[req_id] = sampling_params.logprobs
        if sampling_params.prompt_logprobs is not None:
            self.num_prompt_logprobs[req_id] = sampling_params.prompt_logprobs
306
307
        if sampling_params.logit_bias is not None:
            self.logit_bias[req_index] = sampling_params.logit_bias
308

309
310
311
312
        if sampling_params.allowed_token_ids:
            self.has_allowed_token_ids.add(req_id)
            if self.allowed_token_ids_mask_cpu_tensor is None:
                # Lazy allocation for this tensor, which can be large.
313
                # False means we don't fill with -inf.
314
315
316
317
318
319
320
321
322
                self.allowed_token_ids_mask = torch.zeros(self.max_num_reqs,
                                                          self.vocab_size,
                                                          dtype=torch.bool,
                                                          device=self.device)
                self.allowed_token_ids_mask_cpu_tensor = torch.zeros(
                    self.max_num_reqs,
                    self.vocab_size,
                    dtype=torch.bool,
                    device="cpu")
323
324
            self.allowed_token_ids_mask_cpu_tensor[req_index] = True
            # False means we don't fill with -inf.
325
            self.allowed_token_ids_mask_cpu_tensor[req_index][
326
                sampling_params.allowed_token_ids] = False
327

328
329
330
        if sampling_params.bad_words_token_ids:
            self.bad_words_token_ids[
                req_index] = sampling_params.bad_words_token_ids
331

332
333
334
335
336
337
338
339
340
341
342
343
344
        # Add request lora ID
        if request.lora_request:
            lora_id = request.lora_request.lora_int_id
            if lora_id not in self.lora_id_to_request_ids:
                self.lora_id_to_request_ids[lora_id] = set()

            self.request_lora_mapping[req_index] = lora_id
            self.lora_id_to_request_ids[lora_id].add(request.req_id)
            self.lora_id_to_lora_request[lora_id] = request.lora_request
        else:
            # No LoRA
            self.request_lora_mapping[req_index] = 0

345
    def remove_request(self, req_id: str) -> Optional[int]:
346
347
        """This method must always be followed by a call to condense()."""

348
349
350
        req_index = self.req_id_to_index.pop(req_id, None)
        if req_index is None:
            return None
351
352
        self._req_ids[req_index] = None
        self.req_output_token_ids[req_index] = None
353
354
355
356
357

        self.greedy_reqs.discard(req_id)
        self.random_reqs.discard(req_id)
        self.top_p_reqs.discard(req_id)
        self.top_k_reqs.discard(req_id)
358
        self.min_p_reqs.discard(req_id)
359
        self.min_tokens.pop(req_index, None)
360
361
362
        self.frequency_penalties_reqs.discard(req_id)
        self.presence_penalties_reqs.discard(req_id)
        self.repetition_penalties_reqs.discard(req_id)
363
364
        self.generators.pop(req_index, None)
        self.num_logprobs.pop(req_id, None)
365
        self.num_prompt_logprobs.pop(req_id, None)
366
        self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
367
368
369
370
371
372
373
374
375
376

        # LoRA
        lora_id = self.request_lora_mapping[req_index]
        if lora_id != 0:
            self.lora_id_to_request_ids[lora_id].discard(req_id)
            if len(self.lora_id_to_request_ids[lora_id]) == 0:
                self.lora_id_to_request_ids.pop(lora_id)
                self.lora_id_to_lora_request.pop(lora_id)
            self.request_lora_mapping[req_index] = 0

377
        self.logit_bias[req_index] = None
378
379
        self.has_allowed_token_ids.discard(req_id)
        if self.allowed_token_ids_mask_cpu_tensor is not None:
380
            # False means we don't fill with -inf.
381
            self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
382
        self.bad_words_token_ids.pop(req_index, None)
383
384
        return req_index

385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
    def swap_states(self, i1: int, i2: int) -> None:
        old_id_i1 = self._req_ids[i1]
        old_id_i2 = self._req_ids[i2]
        self._req_ids[i1], self._req_ids[i2] =\
            self._req_ids[i2], self._req_ids[i1] # noqa
        self.req_output_token_ids[i1], self.req_output_token_ids[i2] =\
            self.req_output_token_ids[i2], self.req_output_token_ids[i1]
        assert old_id_i1 is not None and old_id_i2 is not None
        self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] =\
            self.req_id_to_index[old_id_i2], self.req_id_to_index[old_id_i1]
        self.num_tokens[i1], self.num_tokens[i2] =\
            self.num_tokens[i2], self.num_tokens[i1]
        self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] =\
            self.num_tokens_no_spec[i2], self.num_tokens_no_spec[i1]
        self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] =\
            self.num_prompt_tokens[i2], self.num_prompt_tokens[i1]
        self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] =\
            self.num_computed_tokens_cpu[i2], self.num_computed_tokens_cpu[i1]
        self.temperature_cpu[i1], self.temperature_cpu[i2] =\
            self.temperature_cpu[i2], self.temperature_cpu[i1]
        self.top_p_cpu[i1], self.top_p_cpu[i2] =\
            self.top_p_cpu[i2], self.top_p_cpu[i1]
        self.top_k_cpu[i1], self.top_k_cpu[i2] =\
            self.top_k_cpu[i2], self.top_k_cpu[i1]
        self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] =\
            self.frequency_penalties_cpu[i2], self.frequency_penalties_cpu[i1]
        self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] =\
            self.presence_penalties_cpu[i2], self.presence_penalties_cpu[i1]
        self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] =\
            self.repetition_penalties_cpu[i2], self.repetition_penalties_cpu[i1]
        self.min_p_cpu[i1], self.min_p_cpu[i2] =\
            self.min_p_cpu[i2], self.min_p_cpu[i1]

418
419
420
421
422
423
424
425
426
        # NOTE: the following is unsafe
        # self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\
        #     self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...]
        # instead, we need to temporiarily copy the data for one of the indices
        # TODO(lucas): optimize this by only copying valid indices
        tmp = self.token_ids_cpu[i1, ...].copy()
        self.token_ids_cpu[i1, ...] = self.token_ids_cpu[i2, ...]
        self.token_ids_cpu[i2, ...] = tmp

427
428
429
        swap_dict_values(self.generators, i1, i2)
        swap_dict_values(self.min_tokens, i1, i2)
        swap_dict_values(self.bad_words_token_ids, i1, i2)
430
431
432
433
434

        self.request_lora_mapping[i1], self.request_lora_mapping[i2] =\
            self.request_lora_mapping[i2], self.request_lora_mapping[i1]
        self.logit_bias[i1], self.logit_bias[i2] =\
            self.logit_bias[i2], self.logit_bias[i1]
435
436
437
438
439
440

        if self.allowed_token_ids_mask_cpu_tensor is not None:
            self.allowed_token_ids_mask_cpu_tensor[i1], \
                self.allowed_token_ids_mask_cpu_tensor[i2] =\
                self.allowed_token_ids_mask_cpu_tensor[i2], \
                    self.allowed_token_ids_mask_cpu_tensor[i1]
441
442
        self.block_table.swap_row(i1, i2)

443
    def condense(self, empty_req_indices: list[int]) -> None:
444
445
        num_reqs = self.num_reqs
        if num_reqs == 0:
446
            # The batched states are empty.
447
448
            self._req_ids.clear()
            self.req_output_token_ids.clear()
449
450
451
452
            return

        # NOTE(woosuk): This function assumes that the empty_req_indices
        # is sorted in descending order.
453
        last_req_index = num_reqs + len(empty_req_indices) - 1
454
455
456
457
458
459
460
461
462
463
464
        while empty_req_indices:
            # Find the largest non-empty index.
            while last_req_index in empty_req_indices:
                last_req_index -= 1

            # Find the smallest empty index.
            empty_index = empty_req_indices.pop()
            if empty_index >= last_req_index:
                break

            # Swap the states.
465
466
            req_id = self._req_ids[last_req_index]
            output_token_ids = self.req_output_token_ids[last_req_index]
467
            assert req_id is not None
468
469
470
471
            self._req_ids[empty_index] = req_id
            self._req_ids[last_req_index] = None
            self.req_output_token_ids[empty_index] = output_token_ids
            self.req_output_token_ids[last_req_index] = None
472
473
            self.req_id_to_index[req_id] = empty_index

474
475
476
477
            num_tokens = self.num_tokens[last_req_index]
            self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[
                last_req_index, :num_tokens]
            self.num_tokens[empty_index] = num_tokens
478
479
            self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
                last_req_index]
480
481
            self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[
                last_req_index]
482
483
            self.num_computed_tokens_cpu[
                empty_index] = self.num_computed_tokens_cpu[last_req_index]
484
            self.block_table.move_row(last_req_index, empty_index)
485
486
487
488
            self.temperature_cpu[empty_index] = self.temperature_cpu[
                last_req_index]
            self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index]
            self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index]
489
490
491
492
493
494
            self.frequency_penalties_cpu[
                empty_index] = self.frequency_penalties_cpu[last_req_index]
            self.presence_penalties_cpu[
                empty_index] = self.presence_penalties_cpu[last_req_index]
            self.repetition_penalties_cpu[
                empty_index] = self.repetition_penalties_cpu[last_req_index]
495
            self.min_p_cpu[empty_index] = self.min_p_cpu[last_req_index]
496
497
498
499
            generator = self.generators.pop(last_req_index, None)
            if generator is not None:
                self.generators[empty_index] = generator

500
501
502
503
            min_token = self.min_tokens.pop(last_req_index, None)
            if min_token is not None:
                self.min_tokens[empty_index] = min_token

504
505
506
            self.request_lora_mapping[empty_index] = self.request_lora_mapping[
                last_req_index]

507
508
            self.logit_bias[empty_index] = self.logit_bias[last_req_index]

509
510
511
512
513
            if self.allowed_token_ids_mask_cpu_tensor is not None:
                self.allowed_token_ids_mask_cpu_tensor[
                    empty_index] = self.allowed_token_ids_mask_cpu_tensor[
                        last_req_index]

514
515
516
517
            bad_words_token_ids = self.bad_words_token_ids.pop(
                last_req_index, None)
            if bad_words_token_ids is not None:
                self.bad_words_token_ids[empty_index] = bad_words_token_ids
518
519
520
            # Decrement last_req_index since it is now empty.
            last_req_index -= 1

521
522
523
524
525
526
527
528
529
        # Trim lists to the batch size.
        del self._req_ids[self.num_reqs:]
        del self.req_output_token_ids[self.num_reqs:]

    def refresh_sampling_metadata(self):
        self.sampling_metadata = self._make_sampling_metadata()

    def _make_sampling_metadata(self) -> SamplingMetadata:
        num_reqs = self.num_reqs
530
531
532
533
534
        if not self.all_greedy:
            temperature = copy_slice(self.temperature_cpu_tensor,
                                     self.temperature, num_reqs)
        else:
            temperature = None
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
        if not self.no_top_p:
            copy_slice(self.top_p_cpu_tensor, self.top_p, num_reqs)
        if not self.no_top_k:
            copy_slice(self.top_k_cpu_tensor, self.top_k, num_reqs)
        if not self.no_min_p:
            copy_slice(self.min_p_cpu_tensor, self.min_p, num_reqs)

        if not self.no_penalties:
            # Since syncing these tensors is expensive only copy them
            # if necessary i.e. if there are requests which require
            # penalties to be applied during sampling.
            copy_slice(self.frequency_penalties_cpu_tensor,
                       self.frequency_penalties, num_reqs)
            copy_slice(self.presence_penalties_cpu_tensor,
                       self.presence_penalties, num_reqs)
            copy_slice(self.repetition_penalties_cpu_tensor,
                       self.repetition_penalties, num_reqs)

            # The prompt tokens are used only for applying penalties during
            # the sampling process. Hence copy these tensors only when
            # there are requests which need penalties to be applied.
            prompt_token_ids = self._make_prompt_token_ids_tensor()
        else:
            prompt_token_ids = None
559

560
561
562
563
564
565
566
        allowed_token_ids_mask: Optional[torch.Tensor] = None
        if not self.no_allowed_token_ids:
            assert self.allowed_token_ids_mask is not None
            copy_slice(self.allowed_token_ids_mask_cpu_tensor,
                       self.allowed_token_ids_mask, num_reqs)
            allowed_token_ids_mask = self.allowed_token_ids_mask[:num_reqs]

567
        return SamplingMetadata(
568
            temperature=temperature,
569
570
            all_greedy=self.all_greedy,
            all_random=self.all_random,
571
572
573
            top_p=None if self.no_top_p else self.top_p[:num_reqs],
            top_k=None if self.no_top_k else self.top_k[:num_reqs],
            min_p=None if self.no_min_p else self.min_p[:num_reqs],
574
575
            generators=self.generators,
            max_num_logprobs=self.max_num_logprobs,
576
577
578
579
            prompt_token_ids=prompt_token_ids,
            frequency_penalties=self.frequency_penalties[:num_reqs],
            presence_penalties=self.presence_penalties[:num_reqs],
            repetition_penalties=self.repetition_penalties[:num_reqs],
580
            output_token_ids=cast(list[list[int]], self.req_output_token_ids),
581
            min_tokens=self.min_tokens,
582
            no_penalties=self.no_penalties,
583
            logit_bias=self.logit_bias[:num_reqs],
584
            allowed_token_ids_mask=allowed_token_ids_mask,
585
            bad_words_token_ids=self.bad_words_token_ids,
586
587
        )

588
589
590
591
592
593
    def _make_prompt_token_ids_tensor(self) -> torch.Tensor:
        max_prompt_len = self.num_prompt_tokens[:self.num_reqs].max()
        prompt_token_ids_cpu_tensor = torch.empty(
            (self.num_reqs, max_prompt_len),
            device="cpu",
            dtype=torch.int64,
594
595
            pin_memory=self.pin_memory,
        )
596
        prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()
597
598
        prompt_token_ids[:] = self.token_ids_cpu[:self.
                                                 num_reqs, :max_prompt_len]
599
600
601
602
603
604
605
        # Use the value of vocab_size as a pad since we don't have a
        # token_id of this value.
        for i in range(self.num_reqs):
            prompt_token_ids[i, self.num_prompt_tokens[i]:] = self.vocab_size
        return prompt_token_ids_cpu_tensor.to(device=self.device,
                                              non_blocking=True)

606
607
    def make_lora_inputs(
        self, num_scheduled_tokens: np.ndarray
608
    ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
        """
        Given the num_scheduled_tokens for each request in the batch, return
        datastructures used to activate the current LoRAs.
        Returns:
            1. prompt_lora_mapping: A tuple of size self.num_reqs where,
               prompt_lora_mapping[i] is the LoRA id to use for the ith prompt.
            2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens)
               where, token_lora_mapping[i] is the LoRA id to use for ith token.
            3. lora_requests: Set of relevant LoRA requests.
        """

        req_lora_mapping = self.request_lora_mapping[:self.num_reqs]
        prompt_lora_mapping = tuple(req_lora_mapping)
        token_lora_mapping = tuple(
            req_lora_mapping.repeat(num_scheduled_tokens))
624
        active_lora_requests: set[LoRARequest] = set(
625
626
627
628
            self.lora_id_to_lora_request.values())

        return prompt_lora_mapping, token_lora_mapping, active_lora_requests

629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
    @property
    def num_reqs(self) -> int:
        return len(self.req_id_to_index)

    @property
    def all_greedy(self) -> bool:
        return len(self.random_reqs) == 0

    @property
    def all_random(self) -> bool:
        return len(self.greedy_reqs) == 0

    @property
    def no_top_p(self) -> bool:
        return len(self.top_p_reqs) == 0

    @property
    def no_top_k(self) -> bool:
        return len(self.top_k_reqs) == 0

649
650
651
652
    @property
    def no_min_p(self) -> bool:
        return len(self.min_p_reqs) == 0

653
654
655
656
657
658
    @property
    def no_penalties(self) -> bool:
        return (len(self.presence_penalties_reqs) == 0
                and len(self.frequency_penalties_reqs) == 0
                and len(self.repetition_penalties_reqs) == 0)

659
    @property
660
661
    def max_num_logprobs(self) -> Optional[int]:
        return max(self.num_logprobs.values()) if self.num_logprobs else None
662
663
664

    @property
    def no_prompt_logprob(self) -> bool:
665
        return not self.num_prompt_logprobs
666
667
668
669

    @property
    def no_allowed_token_ids(self) -> bool:
        return len(self.has_allowed_token_ids) == 0