gpu_input_batch.py 32.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
# Datastructures defining a GPU input batch
4
5

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

import numpy as np
import torch

11
from vllm.lora.request import LoRARequest
12
from vllm.multimodal.inputs import MultiModalKwargs, PlaceholderRange
13
from vllm.pooling_params import PoolingParams
14
from vllm.sampling_params import SamplingParams, SamplingType
15
from vllm.utils import swap_dict_values
16
from vllm.v1.outputs import LogprobsTensors
17
from vllm.v1.pool.metadata import PoolingMetadata
18
19
20
from vllm.v1.sample.logits_processor import (BatchUpdateBuilder,
                                             MoveDirectionality,
                                             init_builtin_logitsprocs)
21
from vllm.v1.sample.metadata import SamplingMetadata
22
from vllm.v1.spec_decode.utils import is_spec_decode_unsupported
23
from vllm.v1.utils import copy_slice
24
from vllm.v1.worker.block_table import MultiGroupBlockTable
25
26
27
28
29
30


@dataclass
class CachedRequestState:

    req_id: str
31
32
    prompt_token_ids: list[int]
    mm_inputs: list[MultiModalKwargs]
33
    mm_positions: list[PlaceholderRange]
34
35
    sampling_params: Optional[SamplingParams]
    pooling_params: Optional[PoolingParams]
36
37
    generator: Optional[torch.Generator]

38
    block_ids: tuple[list[int], ...]
39
    num_computed_tokens: int
40
    output_token_ids: list[int]
41

42
43
44
    mrope_positions: Optional[torch.Tensor] = None
    mrope_position_delta: Optional[int] = None

45
46
    lora_request: Optional[LoRARequest] = None

47
48
49
    def __post_init__(self):
        self.num_prompt_tokens = len(self.prompt_token_ids)

50
51
    @property
    def num_tokens(self) -> int:
52
53
54
55
56
57
58
        return self.num_prompt_tokens + len(self.output_token_ids)

    def get_token_id(self, idx: int) -> int:
        if idx < self.num_prompt_tokens:
            return self.prompt_token_ids[idx]
        else:
            return self.output_token_ids[idx - self.num_prompt_tokens]
59
60
61
62
63


class InputBatch:

    def __init__(
64
65
66
67
68
69
70
71
        self,
        max_num_reqs: int,
        max_model_len: int,
        max_num_batched_tokens: int,
        device: torch.device,
        pin_memory: bool,
        vocab_size: int,
        block_sizes: list[int],  # The block_size of each kv cache group
72
        is_spec_decode: bool = False,
73
    ):
74
        self.is_spec_decode = is_spec_decode
75
76
        self.max_num_reqs = max_num_reqs
        self.max_model_len = max_model_len
77
        self.max_num_batched_tokens = max_num_batched_tokens
78
79
        self.device = device
        self.pin_memory = pin_memory
80
        self.vocab_size = vocab_size
81

82
83
        self._req_ids: list[Optional[str]] = []
        self.req_id_to_index: dict[str, int] = {}
84

85
86
        # TODO(woosuk): This buffer could be too large if max_model_len is big.
        # Find a way to reduce the CPU memory usage.
87
88
        # This buffer is not directly transferred to the GPU, so it does not
        # need to be pinned.
89
90
91
92
        self.token_ids_cpu_tensor = torch.zeros(
            (max_num_reqs, max_model_len),
            device="cpu",
            dtype=torch.int32,
93
            pin_memory=False,
94
95
        )
        self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()
96
        self.num_tokens = np.zeros(max_num_reqs, dtype=np.int32)
97
        self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
98
        self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
99
100
101
102
103
104
105
106
        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()
107

108
        # Block table.
109
        self.block_table = MultiGroupBlockTable(
110
            max_num_reqs=max_num_reqs,
111
            max_model_len=max_model_len,
112
            max_num_batched_tokens=max_num_batched_tokens,
113
            pin_memory=pin_memory,
114
            device=device,
115
            block_sizes=block_sizes,
116
117
118
119
120
121
122
123
124
125
126
        )

        # 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()
127
128
        self.greedy_reqs: set[str] = set()
        self.random_reqs: set[str] = set()
129
130
131
132
133
134
135
136
137

        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()
138
        self.top_p_reqs: set[str] = set()
139
140
141
142
143
144
145
146
147

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

150
151
        # IDs of requests which do not support spec decoding
        self.spec_decode_unsupported_reqs: set[str] = set()
152

153
154
155
156
157
158
159
160
161
162
        # 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 = \
163
            self.frequency_penalties_cpu_tensor.numpy()
164
        self.frequency_penalties_reqs: set[str] = set()
165
166
167
168
169
170
171
172
173

        # 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)
174
175
        self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy(
        )
176
        self.presence_penalties_reqs: set[str] = set()
177
178
179
180
181
182
183
184
185
186
187

        # 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 = \
188
            self.repetition_penalties_cpu_tensor.numpy()
189
        self.repetition_penalties_reqs: set[str] = set()
190

191
192
193
        # lora related
        self.request_lora_mapping = np.zeros((self.max_num_reqs, ),
                                             dtype=np.int32)
194
195
        self.lora_id_to_request_ids: dict[int, set[str]] = {}
        self.lora_id_to_lora_request: dict[int, LoRARequest] = {}
196

197
        # req_index -> generator
198
199
        # NOTE(woosuk): The indices of the requests that do not have their own
        # generator should not be included in the dictionary.
200
        self.generators: dict[int, torch.Generator] = {}
201

202
        self.num_logprobs: dict[str, int] = {}
203
204
        # NOTE(rob): num_prompt_logprobs only includes reqs
        # that are currently in the prefill phase.
205
        self.num_prompt_logprobs: dict[str, int] = {}
206

207
208
209
        # To accumulate prompt logprobs tensor chunks across prefill steps.
        self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}

210
211
212
213
214
215
216
217
218
219
220
221
222
223
        # Internal representation of per-step batch state changes, used for
        # reordering persistent batch and generating logitsprocs batch state
        # updates. Should reset each step.
        self.batch_update_builder = BatchUpdateBuilder()

        # Define logits processors.
        # TODO(andy): logits processor list should be extensible via engine
        # constructor argument; for now the list is fixed.
        self.logitsprocs = init_builtin_logitsprocs(
            pin_memory_available=pin_memory,
            max_num_reqs=max_num_reqs + 1,
            device=device)

        # TODO convert this to LogitsProcessor
224
        self.has_allowed_token_ids: set[str] = set()
225
226
        # NOTE(lufang): In the mask tensor, if the corresponding token allowed,
        # the value is False. Since we use masked_fill_ to set -inf.
227
228
        self.allowed_token_ids_mask: Optional[torch.Tensor] = None
        self.allowed_token_ids_mask_cpu_tensor: Optional[torch.Tensor] = None
229

230
231
232
        # req_index -> bad_words_token_ids
        self.bad_words_token_ids: dict[int, list[list[int]]] = {}

233
234
235
        self.logits_processing_needs_token_ids = np.zeros(max_num_reqs,
                                                          dtype=bool)

236
        self.req_output_token_ids: list[Optional[list[int]]] = []
237
238
239
240

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

241
242
        self.pooling_params: dict[str, PoolingParams] = {}

243
    @property
244
    def req_ids(self) -> list[str]:
245
246
        # None elements should only be present transiently
        # while performing state updates to the batch.
247
        return cast(list[str], self._req_ids)
248

249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
    def _get_next_add_index(self) -> int:
        if (req_index := self.batch_update_builder.pop_removed()) is not None:
            # Fill the empty index.
            return req_index
        # Append to end
        return self.num_reqs

    def _register_add_request(self, request: "CachedRequestState") -> int:
        """Track add-request operations"""
        req_index = self._get_next_add_index()
        assert req_index < self.max_num_reqs
        params = (request.sampling_params
                  if request.sampling_params else request.pooling_params)
        self.batch_update_builder.added.append(
            (req_index, params, request.output_token_ids))
        return req_index

266
267
268
    def add_request(
        self,
        request: "CachedRequestState",
269
270
    ) -> int:
        req_index = self._register_add_request(request)
271
272

        req_id = request.req_id
273
274
275
276
277
278
279
        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

280
281
282
283
        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)
284
        self.num_prompt_tokens[req_index] = num_prompt_tokens
285
286
287
288
289
290
        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
291
292
        # Number of token ids in token_ids_cpu.
        # NOTE(woosuk): This may include spec decode tokens.
293
        self.num_tokens[req_index] = request.num_tokens
294
295
        # Number of tokens without spec decode tokens.
        self.num_tokens_no_spec[req_index] = request.num_tokens
296
297

        self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
298
        self.block_table.add_row(request.block_ids, req_index)
299

300
        if sampling_params := request.sampling_params:
301
302
303
            if (self.is_spec_decode
                    and is_spec_decode_unsupported(sampling_params)):
                self.spec_decode_unsupported_reqs.add(req_id)
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
            if sampling_params.sampling_type == SamplingType.GREEDY:
                # Avoid later division by zero.
                self.temperature_cpu[req_index] = -1.0
                self.greedy_reqs.add(req_id)
            else:
                self.temperature_cpu[req_index] = sampling_params.temperature
                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)
            top_k = sampling_params.top_k
            if 0 < top_k < self.vocab_size:
                self.top_k_reqs.add(req_id)
            else:
                top_k = self.vocab_size
            self.top_k_cpu[req_index] = top_k
            self.frequency_penalties_cpu[
                req_index] = sampling_params.frequency_penalty
            if sampling_params.frequency_penalty != 0.0:
                self.frequency_penalties_reqs.add(req_id)
            self.presence_penalties_cpu[
                req_index] = sampling_params.presence_penalty
            if sampling_params.presence_penalty != 0.0:
                self.presence_penalties_reqs.add(req_id)
            self.repetition_penalties_cpu[
                req_index] = sampling_params.repetition_penalty
            if sampling_params.repetition_penalty != 1.0:
                self.repetition_penalties_reqs.add(req_id)

            # 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

            if sampling_params.logprobs is not None:
340
341
342
                self.num_logprobs[req_id] = (self.vocab_size
                                             if sampling_params.logprobs == -1
                                             else sampling_params.logprobs)
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
            if sampling_params.prompt_logprobs is not None:
                self.num_prompt_logprobs[
                    req_id] = sampling_params.prompt_logprobs

            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.
                    # False means we don't fill with -inf.
                    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")
                self.allowed_token_ids_mask_cpu_tensor[req_index] = True
363
                # False means we don't fill with -inf.
364
365
                self.allowed_token_ids_mask_cpu_tensor[req_index][
                    sampling_params.allowed_token_ids] = False
366

367
368
369
            if sampling_params.bad_words_token_ids:
                self.bad_words_token_ids[
                    req_index] = sampling_params.bad_words_token_ids
370
371
372
373
        elif pooling_params := request.pooling_params:
            self.pooling_params[req_id] = pooling_params
            self.logits_processing_needs_token_ids[req_index] = (
                pooling_params.requires_token_ids)
374
        else:
375
            raise NotImplementedError(request)
376

377
378
379
380
381
382
383
384
385
386
387
388
389
        # 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

390
391
        return req_index

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

395
396
397
398
399
400
        Args:
          req_id: request to remove

        Returns:
          Removed request index, or `None` if `req_id` not recognized
        """
401

402
403
404
        req_index = self.req_id_to_index.pop(req_id, None)
        if req_index is None:
            return None
405
        self.batch_update_builder.removed_append(req_index)
406
407
        self._req_ids[req_index] = None
        self.req_output_token_ids[req_index] = None
408
409
410
411
412

        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)
413
        self.spec_decode_unsupported_reqs.discard(req_id)
414
415
416
        self.frequency_penalties_reqs.discard(req_id)
        self.presence_penalties_reqs.discard(req_id)
        self.repetition_penalties_reqs.discard(req_id)
417
418
        self.generators.pop(req_index, None)
        self.num_logprobs.pop(req_id, None)
419
        self.num_prompt_logprobs.pop(req_id, None)
420
        self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
421
422
423
424
425
426
427
428
429
430

        # 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

431
432
        self.has_allowed_token_ids.discard(req_id)
        if self.allowed_token_ids_mask_cpu_tensor is not None:
433
            # False means we don't fill with -inf.
434
            self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
435
        self.bad_words_token_ids.pop(req_index, None)
436
        self.pooling_params.pop(req_id, None)
437
438
        return req_index

439
    def swap_states(self, i1: int, i2: int) -> None:
440
441
        self.batch_update_builder.moved.append(
            (i1, i2, MoveDirectionality.SWAP))
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
        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]

472
473
474
475
476
477
478
479
480
        # 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

481
482
        swap_dict_values(self.generators, i1, i2)
        swap_dict_values(self.bad_words_token_ids, i1, i2)
483
484
485

        self.request_lora_mapping[i1], self.request_lora_mapping[i2] =\
            self.request_lora_mapping[i2], self.request_lora_mapping[i1]
486
487
488
489
490
491

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

494
495
496
497
498
499
    def condense(self) -> None:
        """Slide non-empty requests down into lower, empty indices.

        Any consecutive empty indices at the very end of the list are not
        filled.

500
        Args:
501
502
503
504
505
          empty_req_indices: empty indices which may be filled.

        Returns:
          swaps: list of (from,to) swap tuples for moved requests
          empty_req_indices: indices not filled by condensation
506
        """
507
508
509
510
        if not (empty_req_indices := self.batch_update_builder.removed):
            # All removed requests were replaced by added requests, or else no
            # requests were removed at all. No condense() needed
            return
511
512
        num_reqs = self.num_reqs
        if num_reqs == 0:
513
            # The batched states are empty.
514
515
            self._req_ids.clear()
            self.req_output_token_ids.clear()
516
517
518
519
            return

        # NOTE(woosuk): This function assumes that the empty_req_indices
        # is sorted in descending order.
520
        last_req_index = num_reqs + len(empty_req_indices) - 1
521
522
523
524
525
526
        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.
527
528
            empty_index = self.batch_update_builder.peek_removed()
            assert empty_index is not None
529
530
531
            if empty_index >= last_req_index:
                break

532
533
534
535
536
537
            # Move active request down into empty request
            # index.
            self.batch_update_builder.pop_removed()
            self.batch_update_builder.moved.append(
                (last_req_index, empty_index,
                 MoveDirectionality.UNIDIRECTIONAL))
538
539
            req_id = self._req_ids[last_req_index]
            output_token_ids = self.req_output_token_ids[last_req_index]
540
            assert req_id is not None
541
542
543
544
            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
545
546
            self.req_id_to_index[req_id] = empty_index

547
548
549
550
            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
551
552
            self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
                last_req_index]
553
554
            self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[
                last_req_index]
555
556
            self.num_computed_tokens_cpu[
                empty_index] = self.num_computed_tokens_cpu[last_req_index]
557
            self.block_table.move_row(last_req_index, empty_index)
558
559
560
561
            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]
562
563
564
565
566
567
            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]
568
569
570
571
            generator = self.generators.pop(last_req_index, None)
            if generator is not None:
                self.generators[empty_index] = generator

572
573
574
            self.request_lora_mapping[empty_index] = self.request_lora_mapping[
                last_req_index]

575
            # TODO convert these to LogitsProcessors
576
577
578
579
580
            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]

581
582
583
584
            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
585

586
587
588
            # Decrement last_req_index since it is now empty.
            last_req_index -= 1

589
590
591
592
        # Trim lists to the batch size.
        del self._req_ids[self.num_reqs:]
        del self.req_output_token_ids[self.num_reqs:]

593
594
    def refresh_metadata(self):
        """Apply batch updates, reset input batch at end of step
595

596
597
598
599
600
601
602
603
        * Apply batch add/remove/permute to logits procs' states
        * If batch state is modified, update sampling metadata
        """
        batch_update = self.batch_update_builder.get_and_reset(self.num_reqs)
        for logit_proc in self.logitsprocs.all:
            logit_proc.update_state(batch_update)
        if batch_update:
            self.sampling_metadata = self._make_sampling_metadata()
604
605
606

    def _make_sampling_metadata(self) -> SamplingMetadata:
        num_reqs = self.num_reqs
607
608
609
610
611
        if not self.all_greedy:
            temperature = copy_slice(self.temperature_cpu_tensor,
                                     self.temperature, num_reqs)
        else:
            temperature = None
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
        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_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)

628
629
630
        needs_prompt_token_ids = (
            not self.no_penalties
            or self.logits_processing_needs_token_ids[:num_reqs].any())
631
632
633
634
635
        if needs_prompt_token_ids:
            # The prompt tokens are used only for applying penalties or
            # step pooling during the sampling/pooling process.
            # Hence copy these tensors only when there are requests which
            # need penalties/step_pooler to be applied.
636
637
638
            prompt_token_ids = self._make_prompt_token_ids_tensor()
        else:
            prompt_token_ids = None
639

640
641
642
643
644
645
646
        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]

647
        return SamplingMetadata(
648
            temperature=temperature,
649
650
            all_greedy=self.all_greedy,
            all_random=self.all_random,
651
652
            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],
653
654
            generators=self.generators,
            max_num_logprobs=self.max_num_logprobs,
655
656
657
658
            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],
659
            output_token_ids=cast(list[list[int]], self.req_output_token_ids),
660
            no_penalties=self.no_penalties,
661
            allowed_token_ids_mask=allowed_token_ids_mask,
662
            bad_words_token_ids=self.bad_words_token_ids,
663
            logitsprocs=self.logitsprocs,
664
665
        )

666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
    @property
    def pooling_metadata(self) -> PoolingMetadata:
        if len(self.pooling_params) == 0:
            pooling_params = []
        else:
            # Note, for now this assumes that all request in the batch
            # are either sampling or pooling requests
            assert len(self.req_ids) == len(self.pooling_params)
            pooling_params = [
                self.pooling_params[req_id] for req_id in self.req_ids
            ]

        return PoolingMetadata(
            prompt_lens=torch.from_numpy(
                self.num_prompt_tokens[:self.num_reqs]).to(self.device),
            prompt_token_ids=self.sampling_metadata.prompt_token_ids,
            pooling_params=pooling_params,
        )

685
686
687
688
689
690
    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,
691
692
            pin_memory=self.pin_memory,
        )
693
        prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()
694
695
        prompt_token_ids[:] = self.token_ids_cpu[:self.
                                                 num_reqs, :max_prompt_len]
696
697
698
699
700
701
702
        # 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)

703
704
    def make_lora_inputs(
        self, num_scheduled_tokens: np.ndarray
705
    ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
        """
        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))
721
        active_lora_requests: set[LoRARequest] = set(
722
723
724
725
            self.lora_id_to_lora_request.values())

        return prompt_lora_mapping, token_lora_mapping, active_lora_requests

726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
    @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

746
747
748
749
750
751
    @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)

752
    @property
753
754
    def max_num_logprobs(self) -> Optional[int]:
        return max(self.num_logprobs.values()) if self.num_logprobs else None
755
756
757

    @property
    def no_prompt_logprob(self) -> bool:
758
        return not self.num_prompt_logprobs
759
760
761
762

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