"tests/vscode:/vscode.git/clone" did not exist on "8e192ff967b44b186ea02d30e49fddf656fdfe50"
gpu_input_batch.py 28.5 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 TYPE_CHECKING, Optional, cast
6
7
8
9

import numpy as np
import torch

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

18
19
_SAMPLING_EPS = 1e-5

20
21
22
23
24
25
26
27
if TYPE_CHECKING:
    from vllm.multimodal.inputs import PlaceholderRange


@dataclass
class CachedRequestState:

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

35
    block_ids: list[int]
36
    num_computed_tokens: int
37
    output_token_ids: list[int]
38

39
40
41
    mrope_positions: Optional[torch.Tensor] = None
    mrope_position_delta: Optional[int] = None

42
43
    lora_request: Optional[LoRARequest] = None

44
45
46
47
48
49
50
51
52
53
54
55
56
57
    @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,
58
        vocab_size: int,
59
60
61
62
63
64
    ):
        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
65
        self.vocab_size = vocab_size
66

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

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

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

        # 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()
110
111
        self.greedy_reqs: set[str] = set()
        self.random_reqs: set[str] = set()
112
113
114
115
116
117
118
119
120

        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()
121
        self.top_p_reqs: set[str] = set()
122
123
124
125
126
127
128
129
130

        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()
131
        self.top_k_reqs: set[str] = set()
132

133
134
135
136
137
138
139
140
        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()
141
        self.min_p_reqs: set[str] = set()
142

143
144
145
146
147
148
149
150
151
152
        # 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 = \
153
            self.frequency_penalties_cpu_tensor.numpy()
154
        self.frequency_penalties_reqs: set[str] = set()
155
156
157
158
159
160
161
162
163

        # 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)
164
165
        self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy(
        )
166
        self.presence_penalties_reqs: set[str] = set()
167
168
169
170
171
172
173
174
175
176
177

        # 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 = \
178
            self.repetition_penalties_cpu_tensor.numpy()
179
        self.repetition_penalties_reqs: set[str] = set()
180

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

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

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

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

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

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

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

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

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

222
223
224
225
226
227
228
229
230
231
    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
232
233
234
235
236
237
238
        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

239
240
241
242
        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)
243
        self.num_prompt_tokens[req_index] = num_prompt_tokens
244
245
246
247
248
249
        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
250
251
        # Number of token ids in token_ids_cpu.
        # NOTE(woosuk): This may include spec decode tokens.
252
        self.num_tokens[req_index] = request.num_tokens
253
254
        # Number of tokens without spec decode tokens.
        self.num_tokens_no_spec[req_index] = request.num_tokens
255
256

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

        sampling_params = request.sampling_params
        if sampling_params.sampling_type == SamplingType.GREEDY:
261
262
            # Avoid later division by zero.
            self.temperature_cpu[req_index] = -1.0
263
264
            self.greedy_reqs.add(req_id)
        else:
265
            self.temperature_cpu[req_index] = sampling_params.temperature
266
267
268
269
270
            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)
271
272
        top_k = sampling_params.top_k
        if 0 < top_k < self.vocab_size:
273
            self.top_k_reqs.add(req_id)
274
275
276
        else:
            top_k = self.vocab_size
        self.top_k_cpu[req_index] = top_k
277
        self.min_p_cpu[req_index] = sampling_params.min_p
278
279
        self.frequency_penalties_cpu[
            req_index] = sampling_params.frequency_penalty
280
281
        if sampling_params.min_p > _SAMPLING_EPS:
            self.min_p_reqs.add(req_id)
282
283
        if sampling_params.frequency_penalty != 0.0:
            self.frequency_penalties_reqs.add(req_id)
284
285
        self.presence_penalties_cpu[
            req_index] = sampling_params.presence_penalty
286
287
        if sampling_params.presence_penalty != 0.0:
            self.presence_penalties_reqs.add(req_id)
288
289
        self.repetition_penalties_cpu[
            req_index] = sampling_params.repetition_penalty
290
291
        if sampling_params.repetition_penalty != 1.0:
            self.repetition_penalties_reqs.add(req_id)
292
293
294
        if sampling_params.min_tokens:
            self.min_tokens[req_index] = (sampling_params.min_tokens,
                                          sampling_params.all_stop_token_ids)
295

296
297
298
299
        # 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
300

301
302
303
304
        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
305
306
        if sampling_params.logit_bias is not None:
            self.logit_bias[req_index] = sampling_params.logit_bias
307

308
309
310
311
        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.
312
                # False means we don't fill with -inf.
313
314
315
316
317
318
319
320
321
                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")
322
323
            self.allowed_token_ids_mask_cpu_tensor[req_index] = True
            # False means we don't fill with -inf.
324
            self.allowed_token_ids_mask_cpu_tensor[req_index][
325
                sampling_params.allowed_token_ids] = False
326

327
328
329
        self.bad_words_token_ids[
            req_index] = sampling_params.bad_words_token_ids

330
331
332
333
334
335
336
337
338
339
340
341
342
        # 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

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

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

        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)
356
        self.min_p_reqs.discard(req_id)
357
        self.min_tokens.pop(req_index, None)
358
359
360
        self.frequency_penalties_reqs.discard(req_id)
        self.presence_penalties_reqs.discard(req_id)
        self.repetition_penalties_reqs.discard(req_id)
361
362
        self.generators.pop(req_index, None)
        self.num_logprobs.pop(req_id, None)
363
        self.num_prompt_logprobs.pop(req_id, None)
364
365
366
367
368
369
370
371
372
373

        # 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

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

382
383
384
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
    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]

415
416
417
418
419
420
421
422
423
        # 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

424
425
426
        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)
427
428
429
430
431

        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]
432
433
434
435
436
437

        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]
438
439
        self.block_table.swap_row(i1, i2)

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

        # NOTE(woosuk): This function assumes that the empty_req_indices
        # is sorted in descending order.
450
        last_req_index = num_reqs + len(empty_req_indices) - 1
451
452
453
454
455
456
457
458
459
460
461
        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.
462
463
            req_id = self._req_ids[last_req_index]
            output_token_ids = self.req_output_token_ids[last_req_index]
464
            assert req_id is not None
465
466
467
468
            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
469
470
            self.req_id_to_index[req_id] = empty_index

471
472
473
474
            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
475
476
            self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
                last_req_index]
477
478
            self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[
                last_req_index]
479
480
            self.num_computed_tokens_cpu[
                empty_index] = self.num_computed_tokens_cpu[last_req_index]
481
            self.block_table.move_row(last_req_index, empty_index)
482
483
484
485
            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]
486
487
488
489
490
491
            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]
492
            self.min_p_cpu[empty_index] = self.min_p_cpu[last_req_index]
493
494
495
496
            generator = self.generators.pop(last_req_index, None)
            if generator is not None:
                self.generators[empty_index] = generator

497
498
499
500
            min_token = self.min_tokens.pop(last_req_index, None)
            if min_token is not None:
                self.min_tokens[empty_index] = min_token

501
502
503
            self.request_lora_mapping[empty_index] = self.request_lora_mapping[
                last_req_index]

504
505
            self.logit_bias[empty_index] = self.logit_bias[last_req_index]

506
507
508
509
510
            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]

511
512
513
514
            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
515
516
517
            # Decrement last_req_index since it is now empty.
            last_req_index -= 1

518
519
520
521
522
523
524
525
526
        # 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
527
528
529
530
531
        if not self.all_greedy:
            temperature = copy_slice(self.temperature_cpu_tensor,
                                     self.temperature, num_reqs)
        else:
            temperature = None
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
        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
556

557
558
559
560
561
562
563
        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]

564
        return SamplingMetadata(
565
            temperature=temperature,
566
567
            all_greedy=self.all_greedy,
            all_random=self.all_random,
568
569
570
            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],
571
572
            generators=self.generators,
            max_num_logprobs=self.max_num_logprobs,
573
574
575
576
            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],
577
            output_token_ids=cast(list[list[int]], self.req_output_token_ids),
578
            min_tokens=self.min_tokens,
579
            no_penalties=self.no_penalties,
580
            logit_bias=self.logit_bias[:num_reqs],
581
            allowed_token_ids_mask=allowed_token_ids_mask,
582
            bad_words_token_ids=self.bad_words_token_ids,
583
584
        )

585
586
587
588
589
590
    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,
591
592
            pin_memory=self.pin_memory,
        )
593
        prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()
594
595
        prompt_token_ids[:] = self.token_ids_cpu[:self.
                                                 num_reqs, :max_prompt_len]
596
597
598
599
600
601
602
        # 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)

603
604
    def make_lora_inputs(
        self, num_scheduled_tokens: np.ndarray
605
    ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
        """
        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))
621
        active_lora_requests: set[LoRARequest] = set(
622
623
624
625
            self.lora_id_to_lora_request.values())

        return prompt_lora_mapping, token_lora_mapping, active_lora_requests

626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
    @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

646
647
648
649
    @property
    def no_min_p(self) -> bool:
        return len(self.min_p_reqs) == 0

650
651
652
653
654
655
    @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)

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

    @property
    def no_prompt_logprob(self) -> bool:
662
        return not self.num_prompt_logprobs
663
664
665
666

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