gpu_input_batch.py 35 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

import numpy as np
import torch
10
from typing_extensions import deprecated
11

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


@dataclass
class CachedRequestState:

    req_id: str
33
    prompt_token_ids: list[int]
34
    mm_kwargs: list[MultiModalKwargsItem]
35
    mm_positions: list[PlaceholderRange]
36
    mm_hashes: list[str]
37
38
    sampling_params: Optional[SamplingParams]
    pooling_params: Optional[PoolingParams]
39
40
    generator: Optional[torch.Generator]

41
    block_ids: tuple[list[int], ...]
42
    num_computed_tokens: int
43
    output_token_ids: list[int]
44

45
46
47
    mrope_positions: Optional[torch.Tensor] = None
    mrope_position_delta: Optional[int] = None

48
49
    lora_request: Optional[LoRARequest] = None

50
51
52
    def __post_init__(self):
        self.num_prompt_tokens = len(self.prompt_token_ids)

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

57
58
59
60
    # Temporary back-compatibility for plugins that define model runner
    @property
    @deprecated("`mm_inputs` is superseded by `mm_kwargs` and will be "
                "removed in v0.13. Please use `mm_kwargs` instead.")
61
62
63
64
    def mm_inputs(self) -> list[MultiModalKwargsItems]:
        return [
            MultiModalKwargsItems.from_seq([item]) for item in self.mm_kwargs
        ]
65

66
67
68
    def get_token_id(self, idx: int) -> int:
        if idx < self.num_prompt_tokens:
            return self.prompt_token_ids[idx]
69
        return self.output_token_ids[idx - self.num_prompt_tokens]
70
71
72
73
74


class InputBatch:

    def __init__(
75
76
77
78
79
80
81
82
        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
83
        logitsprocs: Optional[LogitsProcessors] = None,
84
        is_spec_decode: bool = False,
85
        is_pooling_model: bool = False,
86
        num_speculative_tokens: int = 0,
87
    ):
88
        self.is_pooling_model = is_pooling_model
89
        self.is_spec_decode = is_spec_decode
90
91
        self.max_num_reqs = max_num_reqs
        self.max_model_len = max_model_len
92
        self.max_num_batched_tokens = max_num_batched_tokens
93
94
        self.device = device
        self.pin_memory = pin_memory
95
        self.vocab_size = vocab_size
96

97
98
        self._req_ids: list[Optional[str]] = []
        self.req_id_to_index: dict[str, int] = {}
99

100
101
        # TODO(woosuk): This buffer could be too large if max_model_len is big.
        # Find a way to reduce the CPU memory usage.
102
103
        # This buffer is not directly transferred to the GPU, so it does not
        # need to be pinned.
104
105
106
107
        self.token_ids_cpu_tensor = torch.zeros(
            (max_num_reqs, max_model_len),
            device="cpu",
            dtype=torch.int32,
108
            pin_memory=False,
109
110
        )
        self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()
111
        self.num_tokens = np.zeros(max_num_reqs, dtype=np.int32)
112
        self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
113
        self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
114
115
116
117
118
119
120
121
        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()
122

123
        # Block table.
124
        self.block_table = MultiGroupBlockTable(
125
            max_num_reqs=max_num_reqs,
126
            max_model_len=max_model_len,
127
            max_num_batched_tokens=max_num_batched_tokens,
128
            pin_memory=pin_memory,
129
            device=device,
130
            block_sizes=block_sizes,
131
            num_speculative_tokens=num_speculative_tokens,
132
133
134
135
136
137
138
139
140
141
142
        )

        # 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()
143
144
        self.greedy_reqs: set[str] = set()
        self.random_reqs: set[str] = set()
145
146
147
148
149
150
151
152
153

        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()
154
        self.top_p_reqs: set[str] = set()
155
156
157
158
159
160
161
162
163

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

166
167
        # IDs of requests which do not support spec decoding
        self.spec_decode_unsupported_reqs: set[str] = set()
168

169
170
171
172
173
174
175
176
177
178
        # 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 = \
179
            self.frequency_penalties_cpu_tensor.numpy()
180
        self.frequency_penalties_reqs: set[str] = set()
181
182
183
184
185
186
187
188
189

        # 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)
190
191
        self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy(
        )
192
        self.presence_penalties_reqs: set[str] = set()
193
194
195
196
197
198
199
200
201
202
203

        # 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 = \
204
            self.repetition_penalties_cpu_tensor.numpy()
205
        self.repetition_penalties_reqs: set[str] = set()
206

207
208
209
210
211
212
213
214
        # Speculative decoding
        self.num_accepted_tokens_cpu_tensor = torch.ones((max_num_reqs, ),
                                                         dtype=torch.int64,
                                                         device="cpu",
                                                         pin_memory=pin_memory)
        self.num_accepted_tokens_cpu = \
            self.num_accepted_tokens_cpu_tensor.numpy()

215
216
217
        # lora related
        self.request_lora_mapping = np.zeros((self.max_num_reqs, ),
                                             dtype=np.int32)
218
219
        self.lora_id_to_request_ids: dict[int, set[str]] = {}
        self.lora_id_to_lora_request: dict[int, LoRARequest] = {}
220

221
        # req_index -> generator
222
223
        # NOTE(woosuk): The indices of the requests that do not have their own
        # generator should not be included in the dictionary.
224
        self.generators: dict[int, torch.Generator] = {}
225

226
        self.num_logprobs: dict[str, int] = {}
227
228
        # NOTE(rob): num_prompt_logprobs only includes reqs
        # that are currently in the prefill phase.
229
        self.num_prompt_logprobs: dict[str, int] = {}
230

231
232
233
        # To accumulate prompt logprobs tensor chunks across prefill steps.
        self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}

234
235
236
237
238
239
        # 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()

        # TODO convert this to LogitsProcessor
240
        self.has_allowed_token_ids: set[str] = set()
241
242
        # NOTE(lufang): In the mask tensor, if the corresponding token allowed,
        # the value is False. Since we use masked_fill_ to set -inf.
243
244
        self.allowed_token_ids_mask: Optional[torch.Tensor] = None
        self.allowed_token_ids_mask_cpu_tensor: Optional[torch.Tensor] = None
245

246
247
248
        # req_index -> bad_words_token_ids
        self.bad_words_token_ids: dict[int, list[list[int]]] = {}

249
250
251
        self.logits_processing_needs_token_ids = np.zeros(max_num_reqs,
                                                          dtype=bool)

252
        self.req_output_token_ids: list[Optional[list[int]]] = []
253

254
255
256
257
        # Store provided logitsprocs. If none are provided, initialize empty
        # data structure
        self.logitsprocs = logitsprocs or LogitsProcessors()

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

261
262
        self.pooling_params: dict[str, PoolingParams] = {}

263
264
265
266
267
        # Cached reference to the GPU tensor of previously sampled tokens
        self.prev_sampled_token_ids: Optional[torch.Tensor] = None
        self.prev_sampled_token_ids_invalid_indices: Optional[set[int]] = None
        self.prev_req_id_to_index: Optional[dict[str, int]] = None

268
    @property
269
    def req_ids(self) -> list[str]:
270
271
        # None elements should only be present transiently
        # while performing state updates to the batch.
272
        return cast(list[str], self._req_ids)
273

274
    def _register_add_request(self, request: "CachedRequestState") -> int:
275
276
277
278
279
280
281
282
283
284
        """Track add-request operations for logits processors.
        Not applicable to pooling models.
        """

        # Fill the next empty index if there is one.
        if (new_req_index := self.batch_update_builder.pop_removed()) is None:
            # Append to end otherwise.
            new_req_index = self.num_reqs

        assert new_req_index < self.max_num_reqs
285
286
287
288
289
290
291
292
        self.batch_update_builder.batch_changed = True
        if request.sampling_params:
            # Detailed added request metadata is only required for non-pooling
            # models, to support logitsprocs.
            self.batch_update_builder.added.append(
                (new_req_index, request.sampling_params,
                 request.prompt_token_ids, request.output_token_ids))

293
        return new_req_index
294

295
296
297
    def add_request(
        self,
        request: "CachedRequestState",
298
    ) -> int:
299
        req_index = self._register_add_request(request)
300
301

        req_id = request.req_id
302
303
304
305
306
307
308
        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

309
310
311
312
        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)
313
        self.num_prompt_tokens[req_index] = num_prompt_tokens
314
315
316
317
318
319
        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
320
321
        # Number of token ids in token_ids_cpu.
        # NOTE(woosuk): This may include spec decode tokens.
322
        self.num_tokens[req_index] = request.num_tokens
323
324
        # Number of tokens without spec decode tokens.
        self.num_tokens_no_spec[req_index] = request.num_tokens
325
326

        self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
327
        self.block_table.add_row(request.block_ids, req_index)
328

329
        if sampling_params := request.sampling_params:
330
331
332
            if (self.is_spec_decode
                    and is_spec_decode_unsupported(sampling_params)):
                self.spec_decode_unsupported_reqs.add(req_id)
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
            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:
369
370
371
                self.num_logprobs[req_id] = (self.vocab_size
                                             if sampling_params.logprobs == -1
                                             else sampling_params.logprobs)
372
            if sampling_params.prompt_logprobs is not None:
373
374
375
                self.num_prompt_logprobs[req_id] = (
                    self.vocab_size if sampling_params.prompt_logprobs == -1
                    else sampling_params.prompt_logprobs)
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392

            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
393
                # False means we don't fill with -inf.
394
395
                self.allowed_token_ids_mask_cpu_tensor[req_index][
                    sampling_params.allowed_token_ids] = False
396

397
398
399
            if sampling_params.bad_words_token_ids:
                self.bad_words_token_ids[
                    req_index] = sampling_params.bad_words_token_ids
400
401
402
403
        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)
404
        else:
405
            raise NotImplementedError("Unrecognized request type")
406

407
408
409
        # Speculative decoding: by default 1 token is generated.
        self.num_accepted_tokens_cpu[req_index] = 1

410
411
412
413
414
415
416
417
418
419
420
421
422
        # 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

423
424
        return req_index

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

428
429
430
431
432
433
        Args:
          req_id: request to remove

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

435
436
437
        req_index = self.req_id_to_index.pop(req_id, None)
        if req_index is None:
            return None
438
439

        self.batch_update_builder.removed_append(req_index)
440
441
        self._req_ids[req_index] = None
        self.req_output_token_ids[req_index] = None
442

443
444
445
446
447
448
449
450
451
452
453
454
455
456
        # LoRA
        lora_id = self.request_lora_mapping[req_index]
        if lora_id != 0:
            lora_req_ids = self.lora_id_to_request_ids[lora_id]
            lora_req_ids.discard(req_id)
            if not lora_req_ids:
                del self.lora_id_to_request_ids[lora_id]
                del self.lora_id_to_lora_request[lora_id]
            self.request_lora_mapping[req_index] = 0

        if self.is_pooling_model:
            self.pooling_params.pop(req_id, None)
            return req_index

457
458
459
460
        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)
461
        self.spec_decode_unsupported_reqs.discard(req_id)
462
463
464
        self.frequency_penalties_reqs.discard(req_id)
        self.presence_penalties_reqs.discard(req_id)
        self.repetition_penalties_reqs.discard(req_id)
465
466
        self.generators.pop(req_index, None)
        self.num_logprobs.pop(req_id, None)
467
        self.num_prompt_logprobs.pop(req_id, None)
468
        self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
469

470
471
        self.has_allowed_token_ids.discard(req_id)
        if self.allowed_token_ids_mask_cpu_tensor is not None:
472
            # False means we don't fill with -inf.
473
            self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
474
        self.bad_words_token_ids.pop(req_index, None)
475
476
        return req_index

477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
    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]

496
497
498
499
500
501
502
503
504
        # 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

505
        self.block_table.swap_row(i1, i2)
506

507
        self.request_lora_mapping[i1], self.request_lora_mapping[i2] = \
508
            self.request_lora_mapping[i2], self.request_lora_mapping[i1]
509

510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
        if self.is_pooling_model:
            # Sampling and logits parameters don't apply to pooling models.
            return

        # For autoregressive models, track detailed request reordering info
        # to support logitsprocs.
        self.batch_update_builder.moved.append(
            (i1, i2, MoveDirectionality.SWAP))

        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]
531
532
        self.num_accepted_tokens_cpu[i1], self.num_accepted_tokens_cpu[i2] =\
            self.num_accepted_tokens_cpu[i2], self.num_accepted_tokens_cpu[i1]
533
534
535
536

        swap_dict_values(self.generators, i1, i2)
        swap_dict_values(self.bad_words_token_ids, i1, i2)

537
538
539
540
541
        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]
542

543
544
545
546
547
548
549
550
551
    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.

        Returns:
          swaps: list of (from,to) swap tuples for moved requests
          empty_req_indices: indices not filled by condensation
552
        """
553
554
        num_reqs = self.num_reqs

555
556
557
558
        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
559
        if num_reqs == 0:
560
            # The batched states are empty.
561
562
            self._req_ids.clear()
            self.req_output_token_ids.clear()
563
564
565
566
            return

        # NOTE(woosuk): This function assumes that the empty_req_indices
        # is sorted in descending order.
567
        last_req_index = num_reqs + len(empty_req_indices) - 1
568
569
570
571
572
573
        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.
574
575
            empty_index = self.batch_update_builder.peek_removed()
            assert empty_index is not None
576
577
578
            if empty_index >= last_req_index:
                break

579
580
581
            # Move active request down into empty request
            # index.
            self.batch_update_builder.pop_removed()
582
583
            req_id = self._req_ids[last_req_index]
            output_token_ids = self.req_output_token_ids[last_req_index]
584
            assert req_id is not None
585
586
587
588
            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
589
590
            self.req_id_to_index[req_id] = empty_index

591
592
593
594
            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
595
596
            self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
                last_req_index]
597
598
            self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[
                last_req_index]
599
600
            self.num_computed_tokens_cpu[
                empty_index] = self.num_computed_tokens_cpu[last_req_index]
601
            self.block_table.move_row(last_req_index, empty_index)
602
603
604
605
606
607

            self.request_lora_mapping[empty_index] = self.request_lora_mapping[
                last_req_index]

            if self.is_pooling_model:
                last_req_index -= 1
co63oc's avatar
co63oc committed
608
                # Sampling state not used by pooling models.
609
610
611
612
613
614
615
616
                continue

            # Autoregressive models require detailed tracking of condense
            # operations to support logitsprocs
            self.batch_update_builder.moved.append(
                (last_req_index, empty_index,
                 MoveDirectionality.UNIDIRECTIONAL))

617
618
619
620
            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]
621
622
623
624
625
626
            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]
627
628
            self.num_accepted_tokens_cpu[
                empty_index] = self.num_accepted_tokens_cpu[last_req_index]
629
630
631
632
            generator = self.generators.pop(last_req_index, None)
            if generator is not None:
                self.generators[empty_index] = generator

633
            # TODO convert these to LogitsProcessors
634
635
636
637
638
            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]

639
640
641
642
            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
643

644
645
646
            # Decrement last_req_index since it is now empty.
            last_req_index -= 1

647
        # Trim lists to the batch size.
648
649
        del self._req_ids[num_reqs:]
        del self.req_output_token_ids[num_reqs:]
650

651
    def refresh_metadata(self):
652
        """Apply any batch updates to sampling metadata."""
653

654
        if self.is_pooling_model:
655
656
657
            batch_changed = self.batch_update_builder.reset()
            if batch_changed:
                self.sampling_metadata = self._make_sampling_metadata()
658
659
660
661
662
            return

        # For non-pooling models - generate and apply logitsprocs update;
        # reset batch update tracking.
        # Update sampling metadata if batch state is changed.
663
664
665
666
667
        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()
668
669
670

    def _make_sampling_metadata(self) -> SamplingMetadata:
        num_reqs = self.num_reqs
671
672
673
674
675
        if not self.all_greedy:
            temperature = copy_slice(self.temperature_cpu_tensor,
                                     self.temperature, num_reqs)
        else:
            temperature = None
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
        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)

692
693
694
        needs_prompt_token_ids = (
            not self.no_penalties
            or self.logits_processing_needs_token_ids[:num_reqs].any())
695
696
697
698
699
        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.
700
701
702
            prompt_token_ids = self._make_prompt_token_ids_tensor()
        else:
            prompt_token_ids = None
703

704
705
706
707
708
709
710
        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]

711
        return SamplingMetadata(
712
            temperature=temperature,
713
714
            all_greedy=self.all_greedy,
            all_random=self.all_random,
715
716
            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],
717
718
            generators=self.generators,
            max_num_logprobs=self.max_num_logprobs,
719
720
721
722
            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],
723
            output_token_ids=cast(list[list[int]], self.req_output_token_ids),
724
            no_penalties=self.no_penalties,
725
            allowed_token_ids_mask=allowed_token_ids_mask,
726
            bad_words_token_ids=self.bad_words_token_ids,
727
            logitsprocs=self.logitsprocs,
728
729
        )

730
731
732
733
734
735
    def get_pooling_params(self) -> list[PoolingParams]:
        assert len(self.req_ids) == len(self.pooling_params)
        return [self.pooling_params[req_id] for req_id in self.req_ids]

    def get_pooling_metadata(self) -> PoolingMetadata:
        pooling_params = self.get_pooling_params()
736
737
738

        return PoolingMetadata(
            prompt_lens=torch.from_numpy(
739
                self.num_prompt_tokens[:self.num_reqs]),
740
741
742
743
            prompt_token_ids=self.sampling_metadata.prompt_token_ids,
            pooling_params=pooling_params,
        )

744
    def _make_prompt_token_ids_tensor(self) -> torch.Tensor:
745
746
        num_reqs = self.num_reqs
        max_prompt_len = self.num_prompt_tokens[:num_reqs].max()
747
748
749
750
        prompt_token_ids_cpu_tensor = torch.empty(
            (self.num_reqs, max_prompt_len),
            device="cpu",
            dtype=torch.int64,
751
752
            pin_memory=self.pin_memory,
        )
753
        prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()
754
        prompt_token_ids[:] = self.token_ids_cpu[:num_reqs, :max_prompt_len]
755
756
        # Use the value of vocab_size as a pad since we don't have a
        # token_id of this value.
757
        for i in range(num_reqs):
758
759
760
761
            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)

762
763
    def make_lora_inputs(
        self, num_scheduled_tokens: np.ndarray
764
    ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
        """
        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))
780
        active_lora_requests: set[LoRARequest] = set(
781
782
783
784
            self.lora_id_to_lora_request.values())

        return prompt_lora_mapping, token_lora_mapping, active_lora_requests

785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
    @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

805
806
807
808
809
810
    @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)

811
    @property
812
813
    def max_num_logprobs(self) -> Optional[int]:
        return max(self.num_logprobs.values()) if self.num_logprobs else None
814
815
816

    @property
    def no_prompt_logprob(self) -> bool:
817
        return not self.num_prompt_logprobs
818
819
820
821

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