"tests/vscode:/vscode.git/clone" did not exist on "d00d65299835da159c161b29ee3b02c12440e777"
gpu_input_batch.py 28.8 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.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
26
27
28
if TYPE_CHECKING:
    from vllm.multimodal.inputs import PlaceholderRange


@dataclass
class CachedRequestState:

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

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

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

43
44
    lora_request: Optional[LoRARequest] = None

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

201
202
203
        # To accumulate prompt logprobs tensor chunks across prefill steps.
        self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}

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

212
213
214
        # req_index -> bad_words_token_ids
        self.bad_words_token_ids: dict[int, list[list[int]]] = {}

215
        self.req_output_token_ids: list[Optional[list[int]]] = []
216
217
218
219
220

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

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

226
227
228
229
230
231
232
233
234
235
    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
236
237
238
239
240
241
242
        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

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

        self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
261
        self.block_table.add_row(request.block_ids, req_index)
262
263
264

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

300
301
302
303
        # 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
304

305
306
307
308
        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
309
310
        if sampling_params.logit_bias is not None:
            self.logit_bias[req_index] = sampling_params.logit_bias
311

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

331
332
333
        if sampling_params.bad_words_token_ids:
            self.bad_words_token_ids[
                req_index] = sampling_params.bad_words_token_ids
334

335
336
337
338
339
340
341
342
343
344
345
346
347
        # 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

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

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

        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)
361
        self.min_p_reqs.discard(req_id)
362
        self.min_tokens.pop(req_index, None)
363
364
365
        self.frequency_penalties_reqs.discard(req_id)
        self.presence_penalties_reqs.discard(req_id)
        self.repetition_penalties_reqs.discard(req_id)
366
367
        self.generators.pop(req_index, None)
        self.num_logprobs.pop(req_id, None)
368
        self.num_prompt_logprobs.pop(req_id, None)
369
        self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
370
371
372
373
374
375
376
377
378
379

        # 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

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

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
418
419
420
    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]

421
422
423
424
425
426
427
428
429
        # 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

430
431
432
        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)
433
434
435
436
437

        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]
438
439
440
441
442
443

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

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

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

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

503
504
505
506
            min_token = self.min_tokens.pop(last_req_index, None)
            if min_token is not None:
                self.min_tokens[empty_index] = min_token

507
508
509
            self.request_lora_mapping[empty_index] = self.request_lora_mapping[
                last_req_index]

510
511
            self.logit_bias[empty_index] = self.logit_bias[last_req_index]

512
513
514
515
516
            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]

517
518
519
520
            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
521
522
523
            # Decrement last_req_index since it is now empty.
            last_req_index -= 1

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

563
564
565
566
567
568
569
        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]

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

591
592
593
594
595
596
    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,
597
598
            pin_memory=self.pin_memory,
        )
599
        prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()
600
601
        prompt_token_ids[:] = self.token_ids_cpu[:self.
                                                 num_reqs, :max_prompt_len]
602
603
604
605
606
607
608
        # 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)

609
610
    def make_lora_inputs(
        self, num_scheduled_tokens: np.ndarray
611
    ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
        """
        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))
627
        active_lora_requests: set[LoRARequest] = set(
628
629
630
631
            self.lora_id_to_lora_request.values())

        return prompt_lora_mapping, token_lora_mapping, active_lora_requests

632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
    @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

652
653
654
655
    @property
    def no_min_p(self) -> bool:
        return len(self.min_p_reqs) == 0

656
657
658
659
660
661
    @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)

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

    @property
    def no_prompt_logprob(self) -> bool:
668
        return not self.num_prompt_logprobs
669
670
671
672

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