gpu_input_batch.py 42.3 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, field
6
from typing import Optional, cast
7
8
9
10

import numpy as np
import torch

11
from vllm import envs
12
from vllm.config import get_current_vllm_config
13
from vllm.lora.request import LoRARequest
14
from vllm.multimodal.inputs import MultiModalKwargs, 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
21
22
from vllm.v1.sample.logits_processor import (BatchUpdateBuilder,
                                             MoveDirectionality,
                                             init_builtin_logitsprocs)
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
34
    prompt_token_ids: list[int]
    mm_inputs: list[MultiModalKwargs]
35
    mm_positions: list[PlaceholderRange]
36
37
    sampling_params: Optional[SamplingParams]
    pooling_params: Optional[PoolingParams]
38
39
    generator: Optional[torch.Generator]

40
    block_ids: tuple[list[int], ...]
41
    num_computed_tokens: int
42
    output_token_ids: list[int]
lizhigong's avatar
lizhigong committed
43
    spec_token_ids: list[int] = None
44

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

48
    lora_request: Optional[LoRARequest] = None
49
50
51
52
53
    # Lazily populated when `VLLM_V1_FAST_TOKEN_ID_COPY` is enabled.
    _prompt_token_ids_np: Optional[np.ndarray] = field(default=None,
                                                       init=False,
                                                       repr=False,
                                                       compare=False)
54

55
56
57
    def __post_init__(self):
        self.num_prompt_tokens = len(self.prompt_token_ids)

58
59
    @property
    def num_tokens(self) -> int:
60
61
62
63
64
65
66
        return self.num_prompt_tokens + len(self.output_token_ids)

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


class InputBatch:

    def __init__(
72
73
74
75
76
77
78
79
        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
80
        is_spec_decode: bool = False,
81
        logits_processing_needs_token_ids: bool = False,
82
    ):
83
84
85
86
        ori_max_num_reqs = max_num_reqs
        if is_spec_decode and envs.VLLM_REJECT_SAMPLE_OPT:
            vllm_config = get_current_vllm_config()
            max_num_reqs = max_num_reqs * (1 + vllm_config.speculative_config.num_speculative_tokens)
87
        self.is_spec_decode = is_spec_decode
88
89
        self.max_num_reqs = max_num_reqs
        self.max_model_len = max_model_len
90
        self.max_num_batched_tokens = max_num_batched_tokens
91
92
        self.device = device
        self.pin_memory = pin_memory
93
        self.vocab_size = vocab_size
94
95
        self.logits_processing_needs_token_ids = (
            logits_processing_needs_token_ids)
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
        self.token_ids_cpu_tensor = torch.zeros(
105
            (ori_max_num_reqs, max_model_len),
106
107
            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
132
133
134
135
136
137
138
139
140
141
        )

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

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

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

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

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

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

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

206
207
208
209
        # Track whether sampling metadata is currently expanded to
        # per-token shape (spec decode reject path).
        self._sampling_metadata_is_expanded = False

210
211
212
        # lora related
        self.request_lora_mapping = np.zeros((self.max_num_reqs, ),
                                             dtype=np.int32)
213
214
        self.lora_id_to_request_ids: dict[int, set[str]] = {}
        self.lora_id_to_lora_request: dict[int, LoRARequest] = {}
215

216
        # req_index -> generator
217
218
        # NOTE(woosuk): The indices of the requests that do not have their own
        # generator should not be included in the dictionary.
219
        self.generators: dict[int, torch.Generator] = {}
220

221
        self.num_logprobs: dict[str, int] = {}
222
223
        # NOTE(rob): num_prompt_logprobs only includes reqs
        # that are currently in the prefill phase.
224
        self.num_prompt_logprobs: dict[str, int] = {}
225

226
227
228
        # To accumulate prompt logprobs tensor chunks across prefill steps.
        self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}

229
230
231
232
233
234
235
236
237
238
239
240
241
242
        # Internal representation of per-step batch state changes, used for
        # reordering persistent batch and generating logitsprocs batch state
        # updates. Should reset each step.
        self.batch_update_builder = BatchUpdateBuilder()

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

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

249
250
251
        # req_index -> bad_words_token_ids
        self.bad_words_token_ids: dict[int, list[list[int]]] = {}

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

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

257
258
        self.pooling_params: dict[str, PoolingParams] = {}

259
    @property
260
    def req_ids(self) -> list[str]:
261
262
        # None elements should only be present transiently
        # while performing state updates to the batch.
263
        return cast(list[str], self._req_ids)
264

265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
    def _get_next_add_index(self) -> int:
        if (req_index := self.batch_update_builder.pop_removed()) is not None:
            # Fill the empty index.
            return req_index
        # Append to end
        return self.num_reqs

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

282
283
284
    def add_request(
        self,
        request: "CachedRequestState",
285
286
    ) -> int:
        req_index = self._register_add_request(request)
287
288

        req_id = request.req_id
289
290
291
292
293
294
295
        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

296
297
298
        self.req_id_to_index[req_id] = req_index

        # Copy the prompt token ids and output token ids.
299
300
301
        # OPTIMIZATION: Use np.copyto with pre-converted NumPy arrays
        # instead of slice assignment. This avoids the ~550 µs overhead
        # of converting Python list to NumPy array each time.
302
        num_prompt_tokens = len(request.prompt_token_ids)
303
        self.num_prompt_tokens[req_index] = num_prompt_tokens
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
        if not envs.VLLM_V1_FAST_TOKEN_ID_COPY:
            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
        else:
            prompt_token_ids_np = getattr(request, "_prompt_token_ids_np", None)
            rebuild_prompt_cache = True
            if prompt_token_ids_np is not None:
                try:
                    rebuild_prompt_cache = (prompt_token_ids_np.dtype != np.int32
                                            or prompt_token_ids_np.size !=
                                            num_prompt_tokens)
                except Exception:
                    rebuild_prompt_cache = True
            if rebuild_prompt_cache:
                prompt_token_ids_np = np.asarray(request.prompt_token_ids,
                                                dtype=np.int32)
                try:
                    request._prompt_token_ids_np = prompt_token_ids_np
                except Exception:
                    pass
            np.copyto(
                self.token_ids_cpu[req_index, :num_prompt_tokens],
                prompt_token_ids_np,
                casting='no',
            )
            start_idx = num_prompt_tokens
            output_token_ids_np = np.asarray(request.output_token_ids,
                                            dtype=np.int32)
            end_idx = start_idx + output_token_ids_np.size
            np.copyto(
                self.token_ids_cpu[req_index, start_idx:end_idx],
                output_token_ids_np,
                casting='no',
            )

lizhigong's avatar
lizhigong committed
343
344
345
346
347
348
        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

349
350
        # Number of token ids in token_ids_cpu.
        # NOTE(woosuk): This may include spec decode tokens.
lizhigong's avatar
lizhigong committed
351
        self.num_tokens[req_index] = request.num_tokens + num_spec_tokens
352
353
        # Number of tokens without spec decode tokens.
        self.num_tokens_no_spec[req_index] = request.num_tokens
354
355

        self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
356
        self.block_table.add_row(request.block_ids, req_index)
357

358
        if sampling_params := request.sampling_params:
359
360
361
            if (self.is_spec_decode
                    and is_spec_decode_unsupported(sampling_params)):
                self.spec_decode_unsupported_reqs.add(req_id)
362
363
364
365
366
367
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
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
            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:
                self.num_logprobs[req_id] = sampling_params.logprobs
            if sampling_params.prompt_logprobs is not None:
                self.num_prompt_logprobs[
                    req_id] = sampling_params.prompt_logprobs

            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
419
                # False means we don't fill with -inf.
420
421
                self.allowed_token_ids_mask_cpu_tensor[req_index][
                    sampling_params.allowed_token_ids] = False
422

423
424
425
426
427
428
            if sampling_params.bad_words_token_ids:
                self.bad_words_token_ids[
                    req_index] = sampling_params.bad_words_token_ids
        else:
            assert request.pooling_params is not None
            self.pooling_params[req_id] = request.pooling_params
429

430
431
432
433
434
435
436
437
438
439
440
441
442
        # 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

443
444
        return req_index

445
    def remove_request(self, req_id: str) -> Optional[int]:
446
447
448
449
450
451
452
453
        """This method must always be followed by a call to condense().
        
        Args:
          req_id: request to remove

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

455
456
457
        req_index = self.req_id_to_index.pop(req_id, None)
        if req_index is None:
            return None
458
        self.batch_update_builder.removed_append(req_index)
459
460
        self._req_ids[req_index] = None
        self.req_output_token_ids[req_index] = None
461
462
463
464
465

        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)
466
        self.spec_decode_unsupported_reqs.discard(req_id)
467
468
469
        self.frequency_penalties_reqs.discard(req_id)
        self.presence_penalties_reqs.discard(req_id)
        self.repetition_penalties_reqs.discard(req_id)
470
471
        self.generators.pop(req_index, None)
        self.num_logprobs.pop(req_id, None)
472
        self.num_prompt_logprobs.pop(req_id, None)
473
        self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
474
475
476
477
478
479
480
481
482
483

        # LoRA
        lora_id = self.request_lora_mapping[req_index]
        if lora_id != 0:
            self.lora_id_to_request_ids[lora_id].discard(req_id)
            if len(self.lora_id_to_request_ids[lora_id]) == 0:
                self.lora_id_to_request_ids.pop(lora_id)
                self.lora_id_to_lora_request.pop(lora_id)
            self.request_lora_mapping[req_index] = 0

484
485
        self.has_allowed_token_ids.discard(req_id)
        if self.allowed_token_ids_mask_cpu_tensor is not None:
486
            # False means we don't fill with -inf.
487
            self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
488
        self.bad_words_token_ids.pop(req_index, None)
489
        self.pooling_params.pop(req_id, None)
490
491
        return req_index

492
    def swap_states(self, i1: int, i2: int) -> None:
493
494
        self.batch_update_builder.moved.append(
            (i1, i2, MoveDirectionality.SWAP))
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
        old_id_i1 = self._req_ids[i1]
        old_id_i2 = self._req_ids[i2]
        self._req_ids[i1], self._req_ids[i2] =\
            self._req_ids[i2], self._req_ids[i1] # noqa
        self.req_output_token_ids[i1], self.req_output_token_ids[i2] =\
            self.req_output_token_ids[i2], self.req_output_token_ids[i1]
        assert old_id_i1 is not None and old_id_i2 is not None
        self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] =\
            self.req_id_to_index[old_id_i2], self.req_id_to_index[old_id_i1]
        self.num_tokens[i1], self.num_tokens[i2] =\
            self.num_tokens[i2], self.num_tokens[i1]
        self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] =\
            self.num_tokens_no_spec[i2], self.num_tokens_no_spec[i1]
        self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] =\
            self.num_prompt_tokens[i2], self.num_prompt_tokens[i1]
        self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] =\
            self.num_computed_tokens_cpu[i2], self.num_computed_tokens_cpu[i1]
        self.temperature_cpu[i1], self.temperature_cpu[i2] =\
            self.temperature_cpu[i2], self.temperature_cpu[i1]
        self.top_p_cpu[i1], self.top_p_cpu[i2] =\
            self.top_p_cpu[i2], self.top_p_cpu[i1]
        self.top_k_cpu[i1], self.top_k_cpu[i2] =\
            self.top_k_cpu[i2], self.top_k_cpu[i1]
        self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] =\
            self.frequency_penalties_cpu[i2], self.frequency_penalties_cpu[i1]
        self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] =\
            self.presence_penalties_cpu[i2], self.presence_penalties_cpu[i1]
        self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] =\
            self.repetition_penalties_cpu[i2], self.repetition_penalties_cpu[i1]

525
526
527
528
529
530
531
532
533
        # 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

534
535
        swap_dict_values(self.generators, i1, i2)
        swap_dict_values(self.bad_words_token_ids, i1, i2)
536
537
538

        self.request_lora_mapping[i1], self.request_lora_mapping[i2] =\
            self.request_lora_mapping[i2], self.request_lora_mapping[i1]
539
540
541
542
543
544

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

547
548
549
550
551
552
    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.

553
        Args:
554
555
556
557
558
          empty_req_indices: empty indices which may be filled.

        Returns:
          swaps: list of (from,to) swap tuples for moved requests
          empty_req_indices: indices not filled by condensation
559
        """
560
561
562
563
        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
564
565
        num_reqs = self.num_reqs
        if num_reqs == 0:
566
            # The batched states are empty.
567
568
            self._req_ids.clear()
            self.req_output_token_ids.clear()
569
570
571
572
            return

        # NOTE(woosuk): This function assumes that the empty_req_indices
        # is sorted in descending order.
573
        last_req_index = num_reqs + len(empty_req_indices) - 1
574
575
576
577
578
579
        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.
580
581
            empty_index = self.batch_update_builder.peek_removed()
            assert empty_index is not None
582
583
584
            if empty_index >= last_req_index:
                break

585
586
587
588
589
590
            # Move active request down into empty request
            # index.
            self.batch_update_builder.pop_removed()
            self.batch_update_builder.moved.append(
                (last_req_index, empty_index,
                 MoveDirectionality.UNIDIRECTIONAL))
591
592
            req_id = self._req_ids[last_req_index]
            output_token_ids = self.req_output_token_ids[last_req_index]
593
            assert req_id is not None
594
595
596
597
            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
598
599
            self.req_id_to_index[req_id] = empty_index

600
601
602
603
            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
604
605
            self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
                last_req_index]
606
607
            self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[
                last_req_index]
608
609
            self.num_computed_tokens_cpu[
                empty_index] = self.num_computed_tokens_cpu[last_req_index]
610
            self.block_table.move_row(last_req_index, empty_index)
611
612
613
614
            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]
615
616
617
618
619
620
            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]
621
622
623
624
            generator = self.generators.pop(last_req_index, None)
            if generator is not None:
                self.generators[empty_index] = generator

625
626
627
            self.request_lora_mapping[empty_index] = self.request_lora_mapping[
                last_req_index]

628
            # TODO convert these to LogitsProcessors
629
630
631
632
633
            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]

634
635
636
637
            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
638

639
640
641
            # Decrement last_req_index since it is now empty.
            last_req_index -= 1

642
643
644
645
        # Trim lists to the batch size.
        del self._req_ids[self.num_reqs:]
        del self.req_output_token_ids[self.num_reqs:]

646
    def refresh_metadata(self, repeat_counts: Optional[torch.Tensor] = None):
647
648
649
650
651
652
653
654
        """Apply batch updates, reset input batch at end of step
        
        * Apply batch add/remove/permute to logits procs' states
        * If batch state is modified, update sampling metadata
        """
        batch_update = self.batch_update_builder.get_and_reset(self.num_reqs)
        for logit_proc in self.logitsprocs.all:
            logit_proc.update_state(batch_update)
655
656
657
658
        needs_rebuild = (batch_update
                         or repeat_counts is not None
                         or self._sampling_metadata_is_expanded)
        if needs_rebuild:
659
660
661
662
663
664
            # if repeat_counts is None:
            #     self.sampling_metadata = self._make_sampling_metadata()
            # else:
            #     self.sampling_metadata = self._make_sampling_metadata_expanded(
            #         repeat_counts)
            self.sampling_metadata = self._make_sampling_metadata(repeat_counts)
665
666
            self._sampling_metadata_is_expanded = repeat_counts is not None
        # Expanded metadata is built on demand; do not cache a copy here.
667

668
    def _make_sampling_metadata(self, repeat_counts: Optional[torch.Tensor] = None) -> SamplingMetadata:
669
        num_reqs = self.num_reqs
670
671
        if not self.all_greedy:
            temperature = copy_slice(self.temperature_cpu_tensor,
672
673
                                     self.temperature, num_reqs,
                                     repeat_counts)
674
675
        else:
            temperature = None
676
        if not self.no_top_p:
677
            top_p = copy_slice(self.top_p_cpu_tensor, self.top_p, num_reqs, repeat_counts)
678
        if not self.no_top_k:
679
            top_k = copy_slice(self.top_k_cpu_tensor, self.top_k, num_reqs, repeat_counts)
680

681
682
683
        frequency_penalties = None
        presence_penalties = None
        repetition_penalties = None
684
685
686
687
        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.
688
689
690
691
692
693
694
695
696
            frequency_penalties = copy_slice(self.frequency_penalties_cpu_tensor,
                       self.frequency_penalties, num_reqs,
                       repeat_counts)
            presence_penalties = copy_slice(self.presence_penalties_cpu_tensor,
                       self.presence_penalties, num_reqs,
                       repeat_counts)
            repetition_penalties = copy_slice(self.repetition_penalties_cpu_tensor,
                       self.repetition_penalties, num_reqs,
                       repeat_counts)
697

698
699
700
701
702
703
704
705
        needs_prompt_token_ids = (not self.no_penalties or
                                  (self.num_reqs > 0
                                   and self.logits_processing_needs_token_ids))
        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.
706
707
708
            prompt_token_ids = self._make_prompt_token_ids_tensor()
        else:
            prompt_token_ids = None
709

710
711
712
        allowed_token_ids_mask: Optional[torch.Tensor] = None
        if not self.no_allowed_token_ids:
            assert self.allowed_token_ids_mask is not None
713
714
715
            allowed_token_ids_mask = copy_slice(self.allowed_token_ids_mask_cpu_tensor,
                       self.allowed_token_ids_mask, num_reqs,
                       repeat_counts)
716

717
718
719
720
721
722
723
724
725
        # Host-side summaries to avoid device synchronization in sampling
        # fast paths (e.g. reduced top-k/top-p sampling).
        max_top_k = None
        has_any_no_top_k = False
        if not self.no_top_k and num_reqs > 0:
            top_k_cpu = self.top_k_cpu[:num_reqs]
            max_top_k = int(top_k_cpu.max())
            has_any_no_top_k = bool((top_k_cpu == self.vocab_size).any())

726
        return SamplingMetadata(
727
            temperature=temperature,
728
729
            all_greedy=self.all_greedy,
            all_random=self.all_random,
730
731
            top_p=None if self.no_top_p else top_p,
            top_k=None if self.no_top_k else top_k,
732
733
            generators=self.generators,
            max_num_logprobs=self.max_num_logprobs,
734
            prompt_token_ids=prompt_token_ids,
735
736
737
            frequency_penalties=None if self.no_penalties else frequency_penalties,
            presence_penalties=None if self.no_penalties else presence_penalties,
            repetition_penalties=None if self.no_penalties else repetition_penalties,
738
            output_token_ids=cast(list[list[int]], self.req_output_token_ids),
739
            no_penalties=self.no_penalties,
740
            allowed_token_ids_mask=allowed_token_ids_mask,
741
            bad_words_token_ids=self.bad_words_token_ids,
742
            logitsprocs=self.logitsprocs,
743
744
            max_top_k=max_top_k,
            has_any_no_top_k=has_any_no_top_k,
745
746
        )

747
748
749
750
    def _make_sampling_metadata_expanded(
        self, repeat_counts: torch.Tensor
    ) -> SamplingMetadata:
        num_reqs = self.num_reqs
751
752
753
754
755
756
757
758
        # `repeat_counts` is expected to be a CPU torch tensor, but some
        # call sites may pass a NumPy array (or other array-likes). Normalize
        # to a CPU tensor to keep downstream ops (e.g. repeat_interleave)
        # consistent and avoid hard crashes.
        if isinstance(repeat_counts, torch.Tensor):
            repeat_counts_cpu = repeat_counts.to(device="cpu")
        else:
            repeat_counts_cpu = torch.as_tensor(repeat_counts, device="cpu")
759
760
761
762
763
764
        all_greedy = self.all_greedy
        all_random = self.all_random
        # For reject-sampling optimization, force greedy sampling to keep
        # rejection sampler assumptions (per-request shapes) intact.

        def _expand_cpu_to_gpu(
765
            t: Optional[object],
766
767
768
769
770
            *,
            dtype: Optional[torch.dtype] = None,
        ) -> Optional[torch.Tensor]:
            if t is None:
                return None
771
772
773
774
775
776
777
778
779
            # `t` should be a CPU torch tensor, but can be a NumPy array view
            # (e.g. created via `tensor.numpy()`). Convert if needed.
            if isinstance(t, torch.Tensor):
                base = t[:num_reqs]
            elif isinstance(t, np.ndarray):
                base = torch.from_numpy(t[:num_reqs])
            else:
                base = torch.as_tensor(t, device="cpu")[:num_reqs]
            base = base.repeat_interleave(repeat_counts_cpu, dim=0)
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
            return base.to(device=self.device,
                           dtype=dtype if dtype is not None else None,
                           non_blocking=True)

        needs_prompt_token_ids = (not self.no_penalties or
                                  (self.num_reqs > 0
                                   and self.logits_processing_needs_token_ids))
        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.
            prompt_token_ids = self._make_prompt_token_ids_tensor(
                repeat_counts_cpu)
        else:
            prompt_token_ids = None

        allowed_token_ids_mask: Optional[torch.Tensor] = None
        if not self.no_allowed_token_ids:
            assert self.allowed_token_ids_mask is not None
            allowed_token_ids_mask = self.allowed_token_ids_mask_cpu_tensor
        # Expand per-request metadata to per-token shape when repeat_counts
        # is provided (spec decode reject-sampling path).
        top_p_cpu = None if self.no_top_p else self.top_p_cpu_tensor
        top_k_cpu = None if self.no_top_k else self.top_k_cpu_tensor

        repeat_list = repeat_counts_cpu.tolist()
        row_offsets: list[int] = []
        total_rows = 0
        for repeat in repeat_list:
            row_offsets.append(total_rows)
            total_rows += int(repeat)
        expanded_output_token_ids: list[list[int]] = []
        expanded_bad_words_token_ids: dict[int, list[list[int]]] = {}
        expanded_generators: dict[int, torch.Generator] = {}
        row_idx = 0
        for req_idx in range(num_reqs):
            repeat = int(repeat_list[req_idx])
            if repeat <= 0:
                continue
            output_tokens = self.req_output_token_ids[req_idx]
            assert output_tokens is not None
            bad_words = self.bad_words_token_ids.get(req_idx)
            generator = self.generators.get(req_idx)
            for _ in range(repeat):
                expanded_output_token_ids.append(output_tokens)
                if bad_words is not None:
                    expanded_bad_words_token_ids[row_idx] = bad_words
                if generator is not None:
                    expanded_generators[row_idx] = generator
                row_idx += 1

832
833
834
835
836
837
838
        max_top_k = None
        has_any_no_top_k = False
        if not self.no_top_k and num_reqs > 0:
            top_k_cpu = self.top_k_cpu[:num_reqs]
            max_top_k = int(top_k_cpu.max())
            has_any_no_top_k = bool((top_k_cpu == self.vocab_size).any())

839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
        return SamplingMetadata(
            temperature=_expand_cpu_to_gpu(
                None if all_greedy else self.temperature_cpu_tensor),
            all_greedy=all_greedy,
            all_random=all_random,
            top_p=_expand_cpu_to_gpu(top_p_cpu),
            top_k=_expand_cpu_to_gpu(top_k_cpu, dtype=torch.int32),
            generators=expanded_generators,
            max_num_logprobs=self.max_num_logprobs,
            prompt_token_ids=prompt_token_ids,
            frequency_penalties=(
                None if self.no_penalties else _expand_cpu_to_gpu(
                    self.frequency_penalties_cpu_tensor)),
            presence_penalties=(
                None if self.no_penalties else _expand_cpu_to_gpu(
                    self.presence_penalties_cpu_tensor)),
            repetition_penalties=(
                None if self.no_penalties else _expand_cpu_to_gpu(
                    self.repetition_penalties_cpu_tensor)),
            output_token_ids=expanded_output_token_ids,
            no_penalties=self.no_penalties,
            allowed_token_ids_mask=_expand_cpu_to_gpu(
                allowed_token_ids_mask, dtype=torch.bool),
            bad_words_token_ids=expanded_bad_words_token_ids,
863
            logitsprocs=self.logitsprocs,
864
865
            max_top_k=max_top_k,
            has_any_no_top_k=has_any_no_top_k,
866
867
        )

868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
    @property
    def pooling_metadata(self) -> PoolingMetadata:
        if len(self.pooling_params) == 0:
            pooling_params = []
        else:
            # Note, for now this assumes that all request in the batch
            # are either sampling or pooling requests
            assert len(self.req_ids) == len(self.pooling_params)
            pooling_params = [
                self.pooling_params[req_id] for req_id in self.req_ids
            ]

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

887
888
889
    def _make_prompt_token_ids_tensor(
        self, repeat_counts_cpu: Optional[torch.Tensor] = None
    ) -> torch.Tensor:
890
891
892
893
894
        max_prompt_len = self.num_prompt_tokens[:self.num_reqs].max()
        prompt_token_ids_cpu_tensor = torch.empty(
            (self.num_reqs, max_prompt_len),
            device="cpu",
            dtype=torch.int64,
895
896
            pin_memory=self.pin_memory,
        )
897
        prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()
898
899
        prompt_token_ids[:] = self.token_ids_cpu[:self.
                                                 num_reqs, :max_prompt_len]
900
901
902
903
        # Use the value of vocab_size as a pad since we don't have a
        # token_id of this value.
        for i in range(self.num_reqs):
            prompt_token_ids[i, self.num_prompt_tokens[i]:] = self.vocab_size
904
905
906
        if repeat_counts_cpu is not None:
            prompt_token_ids_cpu_tensor = prompt_token_ids_cpu_tensor \
                .repeat_interleave(repeat_counts_cpu, dim=0)
907
908
909
        return prompt_token_ids_cpu_tensor.to(device=self.device,
                                              non_blocking=True)

910
911
    def make_lora_inputs(
        self, num_scheduled_tokens: np.ndarray
912
    ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
        """
        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))
928
        active_lora_requests: set[LoRARequest] = set(
929
930
931
932
            self.lora_id_to_lora_request.values())

        return prompt_lora_mapping, token_lora_mapping, active_lora_requests

933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
    @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

953
954
955
956
957
958
    @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)

959
    @property
960
961
    def max_num_logprobs(self) -> Optional[int]:
        return max(self.num_logprobs.values()) if self.num_logprobs else None
962
963
964

    @property
    def no_prompt_logprob(self) -> bool:
965
        return not self.num_prompt_logprobs
966
967
968
969

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