gpu_input_batch.py 41.1 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 cast
7
8
9
10

import numpy as np
import torch

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


@dataclass
class CachedRequestState:
    req_id: str
33
    prompt_token_ids: list[int] | None
34
    mm_features: list[MultiModalFeatureSpec]
35
36
    sampling_params: SamplingParams | None
    generator: torch.Generator | None
37

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

42
43
    mrope_positions: torch.Tensor | None = None
    mrope_position_delta: int | None = None
44

45
46
    xdrope_positions: torch.Tensor | None = None

47
48
    lora_request: LoRARequest | None = None
    prompt_embeds: torch.Tensor | None = None
49

50
51
52
    # Used when both async_scheduling and spec_decode are enabled.
    prev_num_draft_len: int = 0

53
54
55
56
    # for pooling models
    pooling_params: PoolingParams | None = None
    pooling_states: PoolingStates | None = None

57
    def __post_init__(self):
58
        self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
59
60
            self.prompt_token_ids, self.prompt_embeds
        )
61

62
63
64
        if self.pooling_params is not None:
            self.pooling_states = PoolingStates()

65
66
    @property
    def num_tokens(self) -> int:
67
68
69
70
        return self.num_prompt_tokens + len(self.output_token_ids)

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


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

109
        self._req_ids: list[str | None] = []
110
        self.req_id_to_index: dict[str, int] = {}
111

112
113
        # TODO(woosuk): This buffer could be too large if max_model_len is big.
        # Find a way to reduce the CPU memory usage.
114
115
        # This buffer is not directly transferred to the GPU, so it does not
        # need to be pinned.
116
117
118
119
        self.token_ids_cpu_tensor = torch.zeros(
            (max_num_reqs, max_model_len),
            device="cpu",
            dtype=torch.int32,
120
            pin_memory=False,
121
122
        )
        self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()
123
        self.is_token_ids_tensor = torch.zeros(
124
125
            (max_num_reqs, max_model_len), device="cpu", dtype=bool, pin_memory=False
        )
126
        self.is_token_ids = self.is_token_ids_tensor.numpy()
127
128
129
130
        # 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] = {}
131
        self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
132
        self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
133
        self.num_computed_tokens_cpu_tensor = torch.zeros(
134
            (max_num_reqs,),
135
136
137
138
            device="cpu",
            dtype=torch.int32,
            pin_memory=pin_memory,
        )
139
        self.num_computed_tokens_cpu = self.num_computed_tokens_cpu_tensor.numpy()
140

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

        # Sampling-related.
155
156
157
158
159
160
        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
        )
161
        self.temperature_cpu = self.temperature_cpu_tensor.numpy()
162
163
        self.greedy_reqs: set[str] = set()
        self.random_reqs: set[str] = set()
164

165
166
167
168
        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
        )
169
        self.top_p_cpu = self.top_p_cpu_tensor.numpy()
170
        self.top_p_reqs: set[str] = set()
171

172
173
174
175
        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
        )
176
        self.top_k_cpu = self.top_k_cpu_tensor.numpy()
177
        self.top_k_reqs: set[str] = set()
178

179
180
        # IDs of requests which do not support spec decoding
        self.spec_decode_unsupported_reqs: set[str] = set()
181

182
        # Frequency penalty related data structures
183
184
185
        self.frequency_penalties = torch.empty(
            (max_num_reqs,), dtype=torch.float, device=device
        )
186
        self.frequency_penalties_cpu_tensor = torch.empty(
187
188
189
            (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory
        )
        self.frequency_penalties_cpu = self.frequency_penalties_cpu_tensor.numpy()
190
        self.frequency_penalties_reqs: set[str] = set()
191
192

        # Presence penalty related data structures
193
194
        self.presence_penalties = torch.empty(
            (max_num_reqs,), dtype=torch.float, device=device
195
        )
196
197
198
199
        self.presence_penalties_cpu_tensor = torch.empty(
            (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory
        )
        self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy()
200
        self.presence_penalties_reqs: set[str] = set()
201
202

        # Repetition penalty related data structures
203
204
205
        self.repetition_penalties = torch.empty(
            (max_num_reqs,), dtype=torch.float, device=device
        )
206
        self.repetition_penalties_cpu_tensor = torch.empty(
207
208
209
            (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory
        )
        self.repetition_penalties_cpu = self.repetition_penalties_cpu_tensor.numpy()
210
        self.repetition_penalties_reqs: set[str] = set()
211

212
        # Speculative decoding
213
214
215
216
        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()
217

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

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

228
        self.num_logprobs: dict[str, int] = {}
229

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

233
234
235
236
237
238
        # 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
239
        self.has_allowed_token_ids: set[str] = set()
240
241
        # NOTE(lufang): In the mask tensor, if the corresponding token allowed,
        # the value is False. Since we use masked_fill_ to set -inf.
242
243
        self.allowed_token_ids_mask: torch.Tensor | None = None
        self.allowed_token_ids_mask_cpu_tensor: torch.Tensor | None = None
244

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

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

250
        self.req_output_token_ids: list[list[int] | None] = []
251

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

257
        # Store last speculative tokens for sampler.
258
        self.spec_token_ids: list[list[int]] = [[] for _ in range(max_num_reqs)]
259

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

263
        # for pooling models
264
        self.pooling_params: dict[str, PoolingParams] = {}
265
        self.pooling_states: dict[str, PoolingStates] = {}
266

267
        # Cached reference to the GPU tensor of previously sampled tokens
268
269
        self.prev_sampled_token_ids: torch.Tensor | None = None
        self.prev_req_id_to_index: dict[str, int] | None = None
270
271
272
        # These are used to update output_token_ids with real sampled
        # ids from prior step, if required by current sampling params
        # (e.g. penalties).
273
        self.sampled_token_ids_cpu: torch.Tensor | None = None
274
        self.async_copy_ready_event: torch.Event | None = None
275

276
    @property
277
    def req_ids(self) -> list[str]:
278
279
        # None elements should only be present transiently
        # while performing state updates to the batch.
280
        return cast(list[str], self._req_ids)
281

282
    def _register_add_request(self, request: "CachedRequestState") -> int:
283
284
285
286
287
288
289
290
291
292
        """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
293
294
295
296
297
        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(
298
299
300
301
302
303
304
                (
                    new_req_index,
                    request.sampling_params,
                    request.prompt_token_ids,
                    request.output_token_ids,
                )
            )
305

306
        return new_req_index
307

308
309
310
    def add_request(
        self,
        request: "CachedRequestState",
311
    ) -> int:
312
        req_index = self._register_add_request(request)
313
314

        req_id = request.req_id
315
316
317
        if req_index == len(self._req_ids):
            self._req_ids.append(req_id)
            self.req_output_token_ids.append(request.output_token_ids)
318
            self.spec_token_ids.append([])
319
320
321
        else:
            self._req_ids[req_index] = req_id
            self.req_output_token_ids[req_index] = request.output_token_ids
322
            self.spec_token_ids[req_index].clear()
323

324
325
326
        self.req_id_to_index[req_id] = req_index

        # Copy the prompt token ids and output token ids.
327
        num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
328
329
            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
        if request.prompt_token_ids is not None:
334
            self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids
335
336
337
338
339
            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
340
        self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids
341
        self.is_token_ids[req_index, start_idx:end_idx] = True
342
343
        # Number of tokens without spec decode tokens.
        self.num_tokens_no_spec[req_index] = request.num_tokens
344
345

        self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
346
        self.block_table.add_row(request.block_ids, req_index)
347

348
        if sampling_params := request.sampling_params:
349
            if self.is_spec_decode and is_spec_decode_unsupported(sampling_params):
350
                self.spec_decode_unsupported_reqs.add(req_id)
351
            if sampling_params.sampling_type == SamplingType.GREEDY:
352
353
                # Should avoid division by zero later when apply_temperature.
                self.temperature_cpu[req_index] = 0.0
354
355
356
357
358
359
360
361
362
363
364
365
366
367
                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
368
            self.frequency_penalties_cpu[req_index] = sampling_params.frequency_penalty
369
370
            if sampling_params.frequency_penalty != 0.0:
                self.frequency_penalties_reqs.add(req_id)
371
            self.presence_penalties_cpu[req_index] = sampling_params.presence_penalty
372
373
            if sampling_params.presence_penalty != 0.0:
                self.presence_penalties_reqs.add(req_id)
374
375
376
            self.repetition_penalties_cpu[req_index] = (
                sampling_params.repetition_penalty
            )
377
378
379
380
381
382
383
384
385
            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:
386
387
388
389
390
                self.num_logprobs[req_id] = (
                    self.vocab_size
                    if sampling_params.logprobs == -1
                    else sampling_params.logprobs
                )
391
392
393
394
395
396
397
398
399
400

            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,
401
402
                        device=self.device,
                    )
403
404
405
406
                    self.allowed_token_ids_mask_cpu_tensor = torch.zeros(
                        self.max_num_reqs,
                        self.vocab_size,
                        dtype=torch.bool,
407
408
                        device="cpu",
                    )
409
                self.allowed_token_ids_mask_cpu_tensor[req_index] = True
410
                # False means we don't fill with -inf.
411
                self.allowed_token_ids_mask_cpu_tensor[req_index][
412
413
                    sampling_params.allowed_token_ids
                ] = False
414

415
            if sampling_params.bad_words_token_ids:
416
417
418
                self.bad_words_token_ids[req_index] = (
                    sampling_params.bad_words_token_ids
                )
419
        elif pooling_params := request.pooling_params:
420
421
422
            pooling_states = request.pooling_states
            assert pooling_states is not None

423
            self.pooling_params[req_id] = pooling_params
424
            self.pooling_states[req_id] = pooling_states
425
            self.logits_processing_needs_token_ids[req_index] = (
426
427
                pooling_params.requires_token_ids
            )
428
        else:
429
            raise NotImplementedError("Unrecognized request type")
430

431
432
433
        # Speculative decoding: by default 1 token is generated.
        self.num_accepted_tokens_cpu[req_index] = 1

434
435
436
437
438
439
440
441
442
443
444
445
446
        # 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

447
448
        return req_index

449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
    def update_req_spec_token_ids(
        self, request: CachedRequestState, scheduled_spec_tokens: dict[str, list[int]]
    ) -> None:
        req_id = request.req_id
        req_index = self.req_id_to_index[req_id]
        cur_spec_token_ids = self.spec_token_ids[req_index]
        # When speculative decoding is used with structured output,
        # the scheduler can drop draft tokens that do not
        # conform to the schema. This can result in
        # scheduler_output.scheduled_spec_decode_tokens being empty,
        # even when speculative decoding is enabled.
        cur_spec_token_ids.clear()
        spec_token_ids = scheduled_spec_tokens.get(req_id, ())
        num_spec_tokens = len(spec_token_ids)
        request.prev_num_draft_len = num_spec_tokens
        if not spec_token_ids:
            return

        # For async scheduling, token_ids_cpu assigned from
        # spec_token_ids are placeholders and will be overwritten in
        # _prepare_input_ids.
        start_index = self.num_tokens_no_spec[req_index]
        end_token_index = start_index + num_spec_tokens
        self.token_ids_cpu[req_index, start_index:end_token_index] = spec_token_ids
        cur_spec_token_ids.extend(spec_token_ids)

475
    def remove_request(self, req_id: str) -> int | None:
476
        """This method must always be followed by a call to condense().
477

478
479
480
481
482
483
        Args:
          req_id: request to remove

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

485
486
487
        req_index = self.req_id_to_index.pop(req_id, None)
        if req_index is None:
            return None
488
489

        self.batch_update_builder.removed_append(req_index)
490
491
        self._req_ids[req_index] = None
        self.req_output_token_ids[req_index] = None
492
        self.spec_token_ids[req_index].clear()
493

494
495
496
497
498
499
500
501
502
503
504
505
        # 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)
506
            self.pooling_states.pop(req_id, None)
507
508
            return req_index

509
510
511
512
        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)
513
        self.spec_decode_unsupported_reqs.discard(req_id)
514
515
516
        self.frequency_penalties_reqs.discard(req_id)
        self.presence_penalties_reqs.discard(req_id)
        self.repetition_penalties_reqs.discard(req_id)
517
518
        self.generators.pop(req_index, None)
        self.num_logprobs.pop(req_id, None)
519
        self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
520
521
        if self.prev_req_id_to_index is not None:
            self.prev_req_id_to_index.pop(req_id, None)
522

523
524
        self.has_allowed_token_ids.discard(req_id)
        if self.allowed_token_ids_mask_cpu_tensor is not None:
525
            # False means we don't fill with -inf.
526
            self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
527
        self.bad_words_token_ids.pop(req_index, None)
528
529
        return req_index

530
531
532
    def swap_states(self, i1: int, i2: int) -> None:
        old_id_i1 = self._req_ids[i1]
        old_id_i2 = self._req_ids[i2]
533
534
535
536
537
        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],
        )
538
539
540
541
        self.spec_token_ids[i1], self.spec_token_ids[i2] = (
            self.spec_token_ids[i2],
            self.spec_token_ids[i1],
        )
542
        assert old_id_i1 is not None and old_id_i2 is not None
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
        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_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],
        )
559

560
561
562
        # 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, ...]
563
        # instead, we need to temporarily copy the data for one of the indices
564
565
566
567
568
        # 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

569
570
571
572
573
574
575
576
577
578
579
580
581
582
        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)

583
        self.block_table.swap_row(i1, i2)
584

585
586
587
588
        self.request_lora_mapping[i1], self.request_lora_mapping[i2] = (
            self.request_lora_mapping[i2],
            self.request_lora_mapping[i1],
        )
589

590
591
592
593
594
595
        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.
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
        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],
        )
        self.num_accepted_tokens_cpu[i1], self.num_accepted_tokens_cpu[i2] = (
            self.num_accepted_tokens_cpu[i2],
            self.num_accepted_tokens_cpu[i1],
        )
620
621
622
623

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

624
        if self.allowed_token_ids_mask_cpu_tensor is not None:
625
626
627
628
629
630
631
            (
                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],
            )
632

633
634
635
636
637
638
639
640
641
    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
642
        """
643
644
        num_reqs = self.num_reqs

645
646
647
648
        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
649
        if num_reqs == 0:
650
            # The batched states are empty.
651
652
            self._req_ids.clear()
            self.req_output_token_ids.clear()
653
            self.spec_token_ids.clear()
654
655
656
657
            return

        # NOTE(woosuk): This function assumes that the empty_req_indices
        # is sorted in descending order.
658
        last_req_index = num_reqs + len(empty_req_indices) - 1
659
660
661
662
663
664
        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.
665
666
            empty_index = self.batch_update_builder.peek_removed()
            assert empty_index is not None
667
668
669
            if empty_index >= last_req_index:
                break

670
671
672
            # Move active request down into empty request
            # index.
            self.batch_update_builder.pop_removed()
673
674
            req_id = self._req_ids[last_req_index]
            output_token_ids = self.req_output_token_ids[last_req_index]
675
            assert req_id is not None
676
677
678
679
            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
680
681
            self.req_id_to_index[req_id] = empty_index

682
683
684
685
686
687
688
689
690
            num_tokens = self.num_tokens_no_spec[last_req_index] + len(
                self.spec_token_ids[last_req_index]
            )

            (self.spec_token_ids[last_req_index], self.spec_token_ids[empty_index]) = (
                self.spec_token_ids[empty_index],
                self.spec_token_ids[last_req_index],
            )
            self.spec_token_ids[last_req_index].clear()
691

692
            self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[
693
694
                last_req_index, :num_tokens
            ]
695
            self.is_token_ids[empty_index, :num_tokens] = self.is_token_ids[
696
697
                last_req_index, :num_tokens
            ]
698
            if last_req_index in self.req_prompt_embeds:
699
700
701
                self.req_prompt_embeds[empty_index] = self.req_prompt_embeds.pop(
                    last_req_index
                )
702
            self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
703
704
705
706
707
708
                last_req_index
            ]
            self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[last_req_index]
            self.num_computed_tokens_cpu[empty_index] = self.num_computed_tokens_cpu[
                last_req_index
            ]
709
            self.block_table.move_row(last_req_index, empty_index)
710
711

            self.request_lora_mapping[empty_index] = self.request_lora_mapping[
712
713
                last_req_index
            ]
714
715
716

            if self.is_pooling_model:
                last_req_index -= 1
co63oc's avatar
co63oc committed
717
                # Sampling state not used by pooling models.
718
719
720
721
722
                continue

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

726
            self.temperature_cpu[empty_index] = self.temperature_cpu[last_req_index]
727
728
            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]
729
730
731
732
733
734
735
736
737
738
739
740
            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
            ]
            self.num_accepted_tokens_cpu[empty_index] = self.num_accepted_tokens_cpu[
                last_req_index
            ]
741
742
743
744
            generator = self.generators.pop(last_req_index, None)
            if generator is not None:
                self.generators[empty_index] = generator

745
            # TODO convert these to LogitsProcessors
746
            if self.allowed_token_ids_mask_cpu_tensor is not None:
747
748
749
                self.allowed_token_ids_mask_cpu_tensor[empty_index] = (
                    self.allowed_token_ids_mask_cpu_tensor[last_req_index]
                )
750

751
            bad_words_token_ids = self.bad_words_token_ids.pop(last_req_index, None)
752
753
            if bad_words_token_ids is not None:
                self.bad_words_token_ids[empty_index] = bad_words_token_ids
754

755
756
757
            # Decrement last_req_index since it is now empty.
            last_req_index -= 1

758
        # Trim lists to the batch size.
759
760
        del self._req_ids[num_reqs:]
        del self.req_output_token_ids[num_reqs:]
761
        del self.spec_token_ids[num_reqs:]
762

763
    def refresh_metadata(self):
764
        """Apply any batch updates to sampling metadata."""
765

766
        if self.is_pooling_model:
767
768
769
            batch_changed = self.batch_update_builder.reset()
            if batch_changed:
                self.sampling_metadata = self._make_sampling_metadata()
770
771
772
773
774
            return

        # For non-pooling models - generate and apply logitsprocs update;
        # reset batch update tracking.
        # Update sampling metadata if batch state is changed.
775
776
777
778
779
        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()
780
781
782

    def _make_sampling_metadata(self) -> SamplingMetadata:
        num_reqs = self.num_reqs
783
        if not self.all_greedy:
784
785
786
            temperature = copy_slice(
                self.temperature_cpu_tensor, self.temperature, num_reqs
            )
787
788
        else:
            temperature = None
789
790
791
792
793
794
795
796
797
        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.
798
799
800
801
802
803
804
805
806
807
808
            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,
            )
809

810
811
        needs_prompt_token_ids = (
            not self.no_penalties
812
813
            or self.logits_processing_needs_token_ids[:num_reqs].any()
        )
814
815
816
817
818
819
820
        # 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.
        prompt_token_ids = (
            self._make_prompt_token_ids_tensor() if needs_prompt_token_ids else None
        )
821

822
823
824
825
826
827
828
829
830
831
832
833
834
        # Only set output_token_ids if required by the current requests'
        # sampling parameters.
        needs_output_token_ids = (
            not self.no_penalties
            or bool(self.bad_words_token_ids)
            or self.logitsprocs_need_output_token_ids
        )
        output_token_ids = (
            cast(list[list[int]], self.req_output_token_ids)
            if needs_output_token_ids
            else []
        )

835
        allowed_token_ids_mask: torch.Tensor | None = None
836
837
        if not self.no_allowed_token_ids:
            assert self.allowed_token_ids_mask is not None
838
839
840
841
842
            copy_slice(
                self.allowed_token_ids_mask_cpu_tensor,
                self.allowed_token_ids_mask,
                num_reqs,
            )
843
844
            allowed_token_ids_mask = self.allowed_token_ids_mask[:num_reqs]

845
        return SamplingMetadata(
846
            temperature=temperature,
847
848
            all_greedy=self.all_greedy,
            all_random=self.all_random,
849
850
            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],
851
852
            generators=self.generators,
            max_num_logprobs=self.max_num_logprobs,
853
854
855
856
            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],
857
            output_token_ids=output_token_ids,
858
            spec_token_ids=self.spec_token_ids,
859
            no_penalties=self.no_penalties,
860
            allowed_token_ids_mask=allowed_token_ids_mask,
861
            bad_words_token_ids=self.bad_words_token_ids,
862
            logitsprocs=self.logitsprocs,
863
864
        )

865
866
867
868
    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]

869
870
871
872
    def get_pooling_states(self) -> list[PoolingStates]:
        assert len(self.req_ids) == len(self.pooling_states)
        return [self.pooling_states[req_id] for req_id in self.req_ids]

873
874
    def get_pooling_metadata(self) -> PoolingMetadata:
        pooling_params = self.get_pooling_params()
875
        pooling_states = self.get_pooling_states()
876
877

        return PoolingMetadata(
878
            prompt_lens=torch.from_numpy(self.num_prompt_tokens[: self.num_reqs]),
879
880
            prompt_token_ids=self.sampling_metadata.prompt_token_ids,
            pooling_params=pooling_params,
881
            pooling_states=pooling_states,
882
883
        )

884
    def _make_prompt_token_ids_tensor(self) -> torch.Tensor:
885
886
        num_reqs = self.num_reqs
        max_prompt_len = self.num_prompt_tokens[:num_reqs].max()
887
888
889
890
        prompt_token_ids_cpu_tensor = torch.empty(
            (self.num_reqs, max_prompt_len),
            device="cpu",
            dtype=torch.int64,
891
892
            pin_memory=self.pin_memory,
        )
893
        prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()
894
        prompt_token_ids[:] = self.token_ids_cpu[:num_reqs, :max_prompt_len]
895
896
        # Use the value of vocab_size as a pad since we don't have a
        # token_id of this value.
897
        for i in range(num_reqs):
898
899
            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)
900

901
    def make_lora_inputs(
902
        self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray
903
    ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
904
905
906
907
        """
        Given the num_scheduled_tokens for each request in the batch, return
        datastructures used to activate the current LoRAs.
        Returns:
908
909
910
            1. prompt_lora_mapping: A tuple of size np.sum(num_sampled_tokens)
               where, prompt_lora_mapping[i] is the LoRA id to use for the ith
               sampled token.
911
912
913
914
915
            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.
        """

916
        req_lora_mapping = self.request_lora_mapping[: self.num_reqs]
917
        prompt_lora_mapping = tuple(req_lora_mapping.repeat(num_sampled_tokens))
918
        token_lora_mapping = tuple(req_lora_mapping.repeat(num_scheduled_tokens))
919

920
        active_lora_requests: set[LoRARequest] = set(
921
922
            self.lora_id_to_lora_request.values()
        )
923
924
925

        return prompt_lora_mapping, token_lora_mapping, active_lora_requests

926
927
928
    def set_async_sampled_token_ids(
        self,
        sampled_token_ids_cpu: torch.Tensor,
929
        async_copy_ready_event: torch.Event,
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
    ) -> None:
        """
        In async scheduling case, store ref to sampled_token_ids_cpu
        tensor and corresponding copy-ready event. Used to repair
        output_token_ids prior to sampling, if needed by logits processors.
        """
        if self.sampling_metadata.output_token_ids:
            self.sampled_token_ids_cpu = sampled_token_ids_cpu
            self.async_copy_ready_event = async_copy_ready_event
        else:
            self.sampled_token_ids_cpu = None
            self.async_copy_ready_event = None

    def update_async_output_token_ids(self) -> None:
        """
        In async scheduling case, update output_token_ids in sampling metadata
        from prior steps sampled token ids once they've finished copying to CPU.
        This is called right before they are needed by the logits processors.
        """
        output_token_ids = self.sampling_metadata.output_token_ids
        if self.sampled_token_ids_cpu is None or not output_token_ids:
            # Output token ids not needed or not async scheduling.
            return

        assert self.prev_req_id_to_index is not None
        sampled_token_ids = None
        for index, req_id in enumerate(self.req_ids):
            prev_index = self.prev_req_id_to_index.get(req_id)
            if prev_index is None:
                continue
            req_output_token_ids = output_token_ids[index]
            if not req_output_token_ids or req_output_token_ids[-1] != -1:
                # Final output id is not a placeholder, some tokens must have
                # been discarded after a kv-load failure.
                continue
            if sampled_token_ids is None:
                assert self.async_copy_ready_event is not None
                self.async_copy_ready_event.synchronize()
                sampled_token_ids = self.sampled_token_ids_cpu.squeeze(-1).tolist()
            # Replace placeholder token id with actual sampled id.
            req_output_token_ids[-1] = sampled_token_ids[prev_index]

972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
    @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

992
993
    @property
    def no_penalties(self) -> bool:
994
995
996
997
998
        return (
            len(self.presence_penalties_reqs) == 0
            and len(self.frequency_penalties_reqs) == 0
            and len(self.repetition_penalties_reqs) == 0
        )
999

1000
    @property
1001
    def max_num_logprobs(self) -> int | None:
1002
        return max(self.num_logprobs.values()) if self.num_logprobs else None
1003

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