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

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

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


@dataclass
class CachedRequestState:

    req_id: str
32
    prompt_token_ids: Optional[list[int]]
33
    mm_features: list[MultiModalFeatureSpec]
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]
zhuwenwen's avatar
zhuwenwen committed
41
    spec_token_ids: list[int] = None
42

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

46
    lora_request: Optional[LoRARequest] = None
47
    prompt_embeds: Optional[torch.Tensor] = None
48

49
    def __post_init__(self):
50
51
        self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
            self.prompt_token_ids, self.prompt_embeds)
52

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
    def mm_inputs(self) -> list[MultiModalKwargsItems]:
        return [
63
64
            MultiModalKwargsItems.from_seq([f.data]) for f in self.mm_features
            if f.data is not None
65
        ]
66

67
68
    def get_token_id(self, idx: int) -> int:
        if idx < self.num_prompt_tokens:
69
70
71
72
            if self.prompt_token_ids is None:
                raise ValueError(
                    f"Tried to access token index {idx}, but that token was "
                    "provided via prompt_embeds, and its ID is unknown.")
73
            return self.prompt_token_ids[idx]
74
75
76
77
        elif idx - self.num_prompt_tokens < len(self.output_token_ids):
            return self.output_token_ids[idx - self.num_prompt_tokens]
        else:
            return -1
78
79
80
81
82


class InputBatch:

    def __init__(
83
84
85
86
87
88
89
90
        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
91
        logitsprocs: Optional[LogitsProcessors] = None,
92
        is_spec_decode: bool = False,
93
        is_pooling_model: bool = False,
94
        num_speculative_tokens: int = 0,
95
    ):
96
        self.is_pooling_model = is_pooling_model
97
        self.is_spec_decode = is_spec_decode
98
99
        self.max_num_reqs = max_num_reqs
        self.max_model_len = max_model_len
100
        self.max_num_batched_tokens = max_num_batched_tokens
101
102
        self.device = device
        self.pin_memory = pin_memory
103
        self.vocab_size = vocab_size
104

105
106
        self._req_ids: list[Optional[str]] = []
        self.req_id_to_index: dict[str, int] = {}
107

108
109
        # TODO(woosuk): This buffer could be too large if max_model_len is big.
        # Find a way to reduce the CPU memory usage.
110
111
        # This buffer is not directly transferred to the GPU, so it does not
        # need to be pinned.
112
113
114
115
        self.token_ids_cpu_tensor = torch.zeros(
            (max_num_reqs, max_model_len),
            device="cpu",
            dtype=torch.int32,
116
            pin_memory=False,
117
118
        )
        self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()
119
120
121
122
123
124
125
126
        self.is_token_ids = torch.zeros((max_num_reqs, max_model_len),
                                        device="cpu",
                                        dtype=bool,
                                        pin_memory=False)
        # Store prompt embeddings per request to avoid OOM from large upfront
        # allocation if max_model_len is big.
        # Maps req_index -> tensor of shape (num_prompt_tokens, hidden_size)
        self.req_prompt_embeds: dict[int, torch.Tensor] = {}
127
        self.num_tokens = np.zeros(max_num_reqs, dtype=np.int32)
128
        self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
129
        self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
130
131
132
133
134
135
136
137
        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()
138

139
        # Block table.
140
        self.block_table = MultiGroupBlockTable(
141
            max_num_reqs=max_num_reqs,
142
            max_model_len=max_model_len,
143
            max_num_batched_tokens=max_num_batched_tokens,
144
            pin_memory=pin_memory,
145
            device=device,
146
            block_sizes=block_sizes,
147
            num_speculative_tokens=num_speculative_tokens,
148
149
150
151
152
153
154
155
156
157
158
        )

        # 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()
159
160
        self.greedy_reqs: set[str] = set()
        self.random_reqs: set[str] = set()
161
162
163
164
165
166
167
168
169

        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()
170
        self.top_p_reqs: set[str] = set()
171
172
173
174
175
176
177
178
179

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

182
183
        # IDs of requests which do not support spec decoding
        self.spec_decode_unsupported_reqs: set[str] = set()
184

185
186
187
188
189
190
191
192
193
194
        # 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 = \
195
            self.frequency_penalties_cpu_tensor.numpy()
196
        self.frequency_penalties_reqs: set[str] = set()
197
198
199
200
201
202
203
204
205

        # 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)
206
207
        self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy(
        )
208
        self.presence_penalties_reqs: set[str] = set()
209
210
211
212
213
214
215
216
217
218
219

        # 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 = \
220
            self.repetition_penalties_cpu_tensor.numpy()
221
        self.repetition_penalties_reqs: set[str] = set()
222

223
224
225
226
227
228
229
230
        # 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()

231
232
233
        # lora related
        self.request_lora_mapping = np.zeros((self.max_num_reqs, ),
                                             dtype=np.int32)
234
235
        self.lora_id_to_request_ids: dict[int, set[str]] = {}
        self.lora_id_to_lora_request: dict[int, LoRARequest] = {}
236

237
        # req_index -> generator
238
239
        # NOTE(woosuk): The indices of the requests that do not have their own
        # generator should not be included in the dictionary.
240
        self.generators: dict[int, torch.Generator] = {}
241

242
        self.num_logprobs: dict[str, int] = {}
243
244
        # NOTE(rob): num_prompt_logprobs only includes reqs
        # that are currently in the prefill phase.
245
        self.num_prompt_logprobs: dict[str, int] = {}
246

247
248
249
        # To accumulate prompt logprobs tensor chunks across prefill steps.
        self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}

250
251
252
253
254
255
        # 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
256
        self.has_allowed_token_ids: set[str] = set()
257
258
        # NOTE(lufang): In the mask tensor, if the corresponding token allowed,
        # the value is False. Since we use masked_fill_ to set -inf.
259
260
        self.allowed_token_ids_mask: Optional[torch.Tensor] = None
        self.allowed_token_ids_mask_cpu_tensor: Optional[torch.Tensor] = None
261

262
263
264
        # req_index -> bad_words_token_ids
        self.bad_words_token_ids: dict[int, list[list[int]]] = {}

265
266
267
        self.logits_processing_needs_token_ids = np.zeros(max_num_reqs,
                                                          dtype=bool)

268
        self.req_output_token_ids: list[Optional[list[int]]] = []
269

270
271
272
273
        # Store provided logitsprocs. If none are provided, initialize empty
        # data structure
        self.logitsprocs = logitsprocs or LogitsProcessors()

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

277
278
        self.pooling_params: dict[str, PoolingParams] = {}

279
280
281
282
283
        # 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

284
    @property
285
    def req_ids(self) -> list[str]:
286
287
        # None elements should only be present transiently
        # while performing state updates to the batch.
288
        return cast(list[str], self._req_ids)
289

290
    def _register_add_request(self, request: "CachedRequestState") -> int:
291
292
293
294
295
296
297
298
299
300
        """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
301
302
303
304
305
306
307
308
        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))

309
        return new_req_index
310

311
312
313
    def add_request(
        self,
        request: "CachedRequestState",
314
    ) -> int:
315
        req_index = self._register_add_request(request)
316
317

        req_id = request.req_id
318
319
320
321
322
323
324
        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

325
326
327
        self.req_id_to_index[req_id] = req_index

        # Copy the prompt token ids and output token ids.
328
329
        num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
            request.prompt_token_ids, request.prompt_embeds)
330
        self.num_prompt_tokens[req_index] = num_prompt_tokens
331
332
        start_idx = num_prompt_tokens
        end_idx = start_idx + len(request.output_token_ids)
333
334
335
336
337
338
339
340
        if request.prompt_token_ids is not None:
            self.token_ids_cpu[
                req_index, :num_prompt_tokens] = request.prompt_token_ids
            self.is_token_ids[req_index, :num_prompt_tokens] = True
        else:
            self.is_token_ids[req_index, :num_prompt_tokens] = False
        if request.prompt_embeds is not None:
            self.req_prompt_embeds[req_index] = request.prompt_embeds
341
342
        self.token_ids_cpu[req_index,
                           start_idx:end_idx] = request.output_token_ids
zhuwenwen's avatar
zhuwenwen committed
343
344
345
346
347
348
349
        
        num_spec_tokens = 0
        if request.spec_token_ids != None:
            num_spec_tokens = len(request.spec_token_ids)
            self.token_ids_cpu[req_index,
                            end_idx:end_idx + num_spec_tokens] = request.spec_token_ids

350

351
352
        self.is_token_ids[req_index, start_idx:end_idx] = True
        # Number of token ids in prompt (token_ids_cpu or prompt_embeds).
353
        # NOTE(woosuk): This may include spec decode tokens.
zhuwenwen's avatar
zhuwenwen committed
354
        self.num_tokens[req_index] = request.num_tokens + num_spec_tokens
355
356
        # Number of tokens without spec decode tokens.
        self.num_tokens_no_spec[req_index] = request.num_tokens
357
358

        self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
359
        self.block_table.add_row(request.block_ids, req_index)
360

361
        if sampling_params := request.sampling_params:
362
363
364
            if (self.is_spec_decode
                    and is_spec_decode_unsupported(sampling_params)):
                self.spec_decode_unsupported_reqs.add(req_id)
365
            if sampling_params.sampling_type == SamplingType.GREEDY:
366
367
                # Should avoid division by zero later when apply_temperature.
                self.temperature_cpu[req_index] = 0.0
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
                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:
401
402
403
                self.num_logprobs[req_id] = (self.vocab_size
                                             if sampling_params.logprobs == -1
                                             else sampling_params.logprobs)
404
            if sampling_params.prompt_logprobs is not None:
405
406
407
                self.num_prompt_logprobs[req_id] = (
                    self.vocab_size if sampling_params.prompt_logprobs == -1
                    else sampling_params.prompt_logprobs)
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424

            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
425
                # False means we don't fill with -inf.
426
427
                self.allowed_token_ids_mask_cpu_tensor[req_index][
                    sampling_params.allowed_token_ids] = False
428

429
430
431
            if sampling_params.bad_words_token_ids:
                self.bad_words_token_ids[
                    req_index] = sampling_params.bad_words_token_ids
432
433
434
435
        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)
436
        else:
437
            raise NotImplementedError("Unrecognized request type")
438

439
440
441
        # Speculative decoding: by default 1 token is generated.
        self.num_accepted_tokens_cpu[req_index] = 1

442
443
444
445
446
447
448
449
450
451
452
453
454
        # 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

455
456
        return req_index

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

460
461
462
463
464
465
        Args:
          req_id: request to remove

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

467
468
469
        req_index = self.req_id_to_index.pop(req_id, None)
        if req_index is None:
            return None
470
471

        self.batch_update_builder.removed_append(req_index)
472
473
        self._req_ids[req_index] = None
        self.req_output_token_ids[req_index] = None
474

475
476
477
478
479
480
481
482
483
484
485
486
487
488
        # 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

489
490
491
492
        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)
493
        self.spec_decode_unsupported_reqs.discard(req_id)
494
495
496
        self.frequency_penalties_reqs.discard(req_id)
        self.presence_penalties_reqs.discard(req_id)
        self.repetition_penalties_reqs.discard(req_id)
497
498
        self.generators.pop(req_index, None)
        self.num_logprobs.pop(req_id, None)
499
        self.num_prompt_logprobs.pop(req_id, None)
500
        self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
501

502
503
        self.has_allowed_token_ids.discard(req_id)
        if self.allowed_token_ids_mask_cpu_tensor is not None:
504
            # False means we don't fill with -inf.
505
            self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
506
        self.bad_words_token_ids.pop(req_index, None)
507
508
        return req_index

509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
    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]

528
529
530
531
532
533
534
535
536
        # 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

537
538
539
540
541
542
543
544
545
546
547
548
549
550
        self.is_token_ids[[i1, i2], ...] = self.is_token_ids[[i2, i1], ...]

        # Swap prompt embeddings if they exist
        embeds_i1 = self.req_prompt_embeds.get(i1)
        embeds_i2 = self.req_prompt_embeds.get(i2)
        if embeds_i1 is not None:
            self.req_prompt_embeds[i2] = embeds_i1
        else:
            self.req_prompt_embeds.pop(i2, None)
        if embeds_i2 is not None:
            self.req_prompt_embeds[i1] = embeds_i2
        else:
            self.req_prompt_embeds.pop(i1, None)

551
        self.block_table.swap_row(i1, i2)
552

553
        self.request_lora_mapping[i1], self.request_lora_mapping[i2] = \
554
            self.request_lora_mapping[i2], self.request_lora_mapping[i1]
555

556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
        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]
577
578
        self.num_accepted_tokens_cpu[i1], self.num_accepted_tokens_cpu[i2] =\
            self.num_accepted_tokens_cpu[i2], self.num_accepted_tokens_cpu[i1]
579
580
581
582

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

583
584
585
586
587
        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]
588

589
590
591
592
593
594
595
596
597
    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
598
        """
599
600
        num_reqs = self.num_reqs

601
602
603
604
        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
605
        if num_reqs == 0:
606
            # The batched states are empty.
607
608
            self._req_ids.clear()
            self.req_output_token_ids.clear()
609
610
611
612
            return

        # NOTE(woosuk): This function assumes that the empty_req_indices
        # is sorted in descending order.
613
        last_req_index = num_reqs + len(empty_req_indices) - 1
614
615
616
617
618
619
        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.
620
621
            empty_index = self.batch_update_builder.peek_removed()
            assert empty_index is not None
622
623
624
            if empty_index >= last_req_index:
                break

625
626
627
            # Move active request down into empty request
            # index.
            self.batch_update_builder.pop_removed()
628
629
            req_id = self._req_ids[last_req_index]
            output_token_ids = self.req_output_token_ids[last_req_index]
630
            assert req_id is not None
631
632
633
634
            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
635
636
            self.req_id_to_index[req_id] = empty_index

637
638
639
            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]
640
641
642
643
644
            self.is_token_ids[empty_index, :num_tokens] = self.is_token_ids[
                last_req_index, :num_tokens]
            if last_req_index in self.req_prompt_embeds:
                self.req_prompt_embeds[
                    empty_index] = self.req_prompt_embeds.pop(last_req_index)
645
            self.num_tokens[empty_index] = num_tokens
646
647
            self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
                last_req_index]
648
649
            self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[
                last_req_index]
650
651
            self.num_computed_tokens_cpu[
                empty_index] = self.num_computed_tokens_cpu[last_req_index]
652
            self.block_table.move_row(last_req_index, empty_index)
653
654
655
656
657
658

            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
659
                # Sampling state not used by pooling models.
660
661
662
663
664
665
666
667
                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))

668
669
670
671
            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]
672
673
674
675
676
677
            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]
678
679
            self.num_accepted_tokens_cpu[
                empty_index] = self.num_accepted_tokens_cpu[last_req_index]
680
681
682
683
            generator = self.generators.pop(last_req_index, None)
            if generator is not None:
                self.generators[empty_index] = generator

684
            # TODO convert these to LogitsProcessors
685
686
687
688
689
            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]

690
691
692
693
            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
694

695
696
697
            # Decrement last_req_index since it is now empty.
            last_req_index -= 1

698
        # Trim lists to the batch size.
699
700
        del self._req_ids[num_reqs:]
        del self.req_output_token_ids[num_reqs:]
701

702
    def refresh_metadata(self):
703
        """Apply any batch updates to sampling metadata."""
704

705
        if self.is_pooling_model:
706
707
708
            batch_changed = self.batch_update_builder.reset()
            if batch_changed:
                self.sampling_metadata = self._make_sampling_metadata()
709
710
711
712
713
            return

        # For non-pooling models - generate and apply logitsprocs update;
        # reset batch update tracking.
        # Update sampling metadata if batch state is changed.
714
715
716
717
718
        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()
719
720
721

    def _make_sampling_metadata(self) -> SamplingMetadata:
        num_reqs = self.num_reqs
722
723
724
725
726
        if not self.all_greedy:
            temperature = copy_slice(self.temperature_cpu_tensor,
                                     self.temperature, num_reqs)
        else:
            temperature = None
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
        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)

743
744
745
        needs_prompt_token_ids = (
            not self.no_penalties
            or self.logits_processing_needs_token_ids[:num_reqs].any())
746
747
748
749
750
        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.
751
752
753
            prompt_token_ids = self._make_prompt_token_ids_tensor()
        else:
            prompt_token_ids = None
754

755
756
757
758
759
760
761
        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]

762
        return SamplingMetadata(
763
            temperature=temperature,
764
765
            all_greedy=self.all_greedy,
            all_random=self.all_random,
766
767
            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],
768
769
            generators=self.generators,
            max_num_logprobs=self.max_num_logprobs,
770
771
772
773
            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],
774
            output_token_ids=cast(list[list[int]], self.req_output_token_ids),
775
            no_penalties=self.no_penalties,
776
            allowed_token_ids_mask=allowed_token_ids_mask,
777
            bad_words_token_ids=self.bad_words_token_ids,
778
            logitsprocs=self.logitsprocs,
779
780
        )

781
782
783
784
785
786
    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()
787
788
789

        return PoolingMetadata(
            prompt_lens=torch.from_numpy(
790
                self.num_prompt_tokens[:self.num_reqs]),
791
792
793
794
            prompt_token_ids=self.sampling_metadata.prompt_token_ids,
            pooling_params=pooling_params,
        )

795
    def _make_prompt_token_ids_tensor(self) -> torch.Tensor:
796
797
        num_reqs = self.num_reqs
        max_prompt_len = self.num_prompt_tokens[:num_reqs].max()
798
799
800
801
        prompt_token_ids_cpu_tensor = torch.empty(
            (self.num_reqs, max_prompt_len),
            device="cpu",
            dtype=torch.int64,
802
803
            pin_memory=self.pin_memory,
        )
804
        prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()
805
        prompt_token_ids[:] = self.token_ids_cpu[:num_reqs, :max_prompt_len]
806
807
        # Use the value of vocab_size as a pad since we don't have a
        # token_id of this value.
808
        for i in range(num_reqs):
809
810
811
812
            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)

813
814
    def make_lora_inputs(
        self, num_scheduled_tokens: np.ndarray
815
    ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
        """
        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))
831
        active_lora_requests: set[LoRARequest] = set(
832
833
834
835
            self.lora_id_to_lora_request.values())

        return prompt_lora_mapping, token_lora_mapping, active_lora_requests

836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
    @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

856
857
858
859
860
861
    @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)

862
    @property
863
864
    def max_num_logprobs(self) -> Optional[int]:
        return max(self.num_logprobs.values()) if self.num_logprobs else None
865
866
867

    @property
    def no_prompt_logprob(self) -> bool:
868
        return not self.num_prompt_logprobs
869
870
871
872

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