A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 06:00 UTC
[V1] Decouple GPU and TPU `InputBatch` (#19778)
Signed-off-by: Andrew Feldman <afeldman@redhat.com>
afeldman-nm committed
Jun 18, 2025 at 02:38 UTC
19a53b27833d767632c9eaff6ddf5ef08ba3af2f
5 files changed
+597
-4
vllm/v1/sample/tpu/metadata.py
+1
-1
@@ -5,7 +5,7 @@ from typing import Optional
5
6
import torch
7
8
-from vllm.v1.worker.gpu_input_batch import InputBatch
8
+from vllm.v1.worker.tpu_input_batch import InputBatch
9
10
DEFAULT_SAMPLING_PARAMS = dict(
11
temperature=-1.0,
vllm/v1/worker/gpu_input_batch.py
+6
-1
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-# Datastructures defining an input batch
3
+# Datastructures defining a GPU input batch
4
5
from dataclasses import dataclass
6
from typing import Optional, cast
@@ -453,6 +453,11 @@ class InputBatch:
453
self.block_table.swap_row(i1, i2)
454
455
def condense(self, empty_req_indices: list[int]) -> None:
456
+ """Move non-empty requests down into lower, empty indices.
457
+
458
+ Args:
459
+ empty_req_indices: empty batch indices, sorted descending.
460
+ """
461
num_reqs = self.num_reqs
462
if num_reqs == 0:
463
# The batched states are empty.
vllm/v1/worker/lora_model_runner_mixin.py
+5
-1
@@ -5,6 +5,7 @@ Define LoRA functionality mixin for model runners.
5
"""
6
7
from contextlib import contextmanager
8
+from typing import Union
9
10
import numpy as np
11
import torch.nn as nn
@@ -15,7 +16,10 @@ from vllm.lora.layers import LoRAMapping
16
from vllm.lora.request import LoRARequest
17
from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
18
from vllm.model_executor.models import supports_lora, supports_multimodal
18
-from vllm.v1.worker.gpu_input_batch import InputBatch
19
+from vllm.v1.worker.gpu_input_batch import InputBatch as GPUInputBatch
20
+from vllm.v1.worker.tpu_input_batch import InputBatch as TPUInputBatch
21
+
22
+InputBatch = Union[TPUInputBatch, GPUInputBatch]
23
24
logger = init_logger(__name__)
25
vllm/v1/worker/tpu_input_batch.py
new
+584
@@ -0,0 +1,584 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# Datastructures defining a TPU input batch
4
+
5
+from typing import Optional, cast
6
+
7
+import numpy as np
8
+import torch
9
+
10
+from vllm.lora.request import LoRARequest
11
+from vllm.sampling_params import SamplingType
12
+from vllm.utils import swap_dict_values
13
+from vllm.v1.outputs import LogprobsTensors
14
+from vllm.v1.worker.block_table import MultiGroupBlockTable
15
+from vllm.v1.worker.gpu_input_batch import CachedRequestState
16
+
17
+_SAMPLING_EPS = 1e-5
18
+
19
+
20
+class InputBatch:
21
+
22
+ def __init__(
23
+ self,
24
+ max_num_reqs: int,
25
+ max_model_len: int,
26
+ max_num_batched_tokens: int,
27
+ device: torch.device,
28
+ pin_memory: bool,
29
+ vocab_size: int,
30
+ block_sizes: list[int], # The block_size of each kv cache group
31
+ ):
32
+ self.max_num_reqs = max_num_reqs
33
+ self.max_model_len = max_model_len
34
+ self.max_num_batched_tokens = max_num_batched_tokens
35
+ self.device = device
36
+ self.pin_memory = pin_memory
37
+ self.vocab_size = vocab_size
38
+
39
+ self._req_ids: list[Optional[str]] = []
40
+ self.req_id_to_index: dict[str, int] = {}
41
+
42
+ # TODO(woosuk): This buffer could be too large if max_model_len is big.
43
+ # Find a way to reduce the CPU memory usage.
44
+ # This buffer is not directly transferred to the GPU, so it does not
45
+ # need to be pinned.
46
+ self.token_ids_cpu_tensor = torch.zeros(
47
+ (max_num_reqs, max_model_len),
48
+ device="cpu",
49
+ dtype=torch.int32,
50
+ pin_memory=False,
51
+ )
52
+ self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()
53
+ self.num_tokens = np.zeros(max_num_reqs, dtype=np.int32)
54
+ self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
55
+ self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
56
+ self.num_computed_tokens_cpu_tensor = torch.zeros(
57
+ (max_num_reqs, ),
58
+ device="cpu",
59
+ dtype=torch.int32,
60
+ pin_memory=pin_memory,
61
+ )
62
+ self.num_computed_tokens_cpu = \
63
+ self.num_computed_tokens_cpu_tensor.numpy()
64
+
65
+ # Block table.
66
+ self.block_table = MultiGroupBlockTable(
67
+ max_num_reqs=max_num_reqs,
68
+ max_model_len=max_model_len,
69
+ max_num_batched_tokens=max_num_batched_tokens,
70
+ pin_memory=pin_memory,
71
+ device=device,
72
+ block_sizes=block_sizes,
73
+ )
74
+
75
+ # Sampling-related.
76
+ self.temperature = torch.empty((max_num_reqs, ),
77
+ dtype=torch.float32,
78
+ device=device)
79
+ self.temperature_cpu_tensor = torch.empty((max_num_reqs, ),
80
+ dtype=torch.float32,
81
+ device="cpu",
82
+ pin_memory=pin_memory)
83
+ self.temperature_cpu = self.temperature_cpu_tensor.numpy()
84
+ self.greedy_reqs: set[str] = set()
85
+ self.random_reqs: set[str] = set()
86
+
87
+ self.top_p = torch.empty((max_num_reqs, ),
88
+ dtype=torch.float32,
89
+ device=device)
90
+ self.top_p_cpu_tensor = torch.empty((max_num_reqs, ),
91
+ dtype=torch.float32,
92
+ device="cpu",
93
+ pin_memory=pin_memory)
94
+ self.top_p_cpu = self.top_p_cpu_tensor.numpy()
95
+ self.top_p_reqs: set[str] = set()
96
+
97
+ self.top_k = torch.empty((max_num_reqs, ),
98
+ dtype=torch.int32,
99
+ device=device)
100
+ self.top_k_cpu_tensor = torch.empty((max_num_reqs, ),
101
+ dtype=torch.int32,
102
+ device="cpu",
103
+ pin_memory=pin_memory)
104
+ self.top_k_cpu = self.top_k_cpu_tensor.numpy()
105
+ self.top_k_reqs: set[str] = set()
106
+
107
+ self.min_p = torch.empty((max_num_reqs, ),
108
+ dtype=torch.float32,
109
+ device=device)
110
+ self.min_p_cpu_tensor = torch.empty((max_num_reqs, ),
111
+ dtype=torch.float32,
112
+ device="cpu",
113
+ pin_memory=pin_memory)
114
+ self.min_p_cpu = self.min_p_cpu_tensor.numpy()
115
+ self.min_p_reqs: set[str] = set()
116
+
117
+ # Frequency penalty related data structures
118
+ self.frequency_penalties = torch.empty((max_num_reqs, ),
119
+ dtype=torch.float,
120
+ device=device)
121
+ self.frequency_penalties_cpu_tensor = torch.empty(
122
+ (max_num_reqs, ),
123
+ dtype=torch.float,
124
+ device="cpu",
125
+ pin_memory=pin_memory)
126
+ self.frequency_penalties_cpu = \
127
+ self.frequency_penalties_cpu_tensor.numpy()
128
+ self.frequency_penalties_reqs: set[str] = set()
129
+
130
+ # Presence penalty related data structures
131
+ self.presence_penalties = torch.empty((max_num_reqs, ),
132
+ dtype=torch.float,
133
+ device=device)
134
+ self.presence_penalties_cpu_tensor = torch.empty((max_num_reqs, ),
135
+ dtype=torch.float,
136
+ device="cpu",
137
+ pin_memory=pin_memory)
138
+ self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy(
139
+ )
140
+ self.presence_penalties_reqs: set[str] = set()
141
+
142
+ # Repetition penalty related data structures
143
+ self.repetition_penalties = torch.empty((max_num_reqs, ),
144
+ dtype=torch.float,
145
+ device=device)
146
+ self.repetition_penalties_cpu_tensor = torch.empty(
147
+ (max_num_reqs, ),
148
+ dtype=torch.float,
149
+ device="cpu",
150
+ pin_memory=pin_memory)
151
+ self.repetition_penalties_cpu = \
152
+ self.repetition_penalties_cpu_tensor.numpy()
153
+ self.repetition_penalties_reqs: set[str] = set()
154
+
155
+ # req_index -> (min_tokens, stop_token_ids)
156
+ self.min_tokens: dict[int, tuple[int, set[int]]] = {}
157
+
158
+ # lora related
159
+ self.request_lora_mapping = np.zeros((self.max_num_reqs, ),
160
+ dtype=np.int32)
161
+ self.lora_id_to_request_ids: dict[int, set[str]] = {}
162
+ self.lora_id_to_lora_request: dict[int, LoRARequest] = {}
163
+
164
+ # req_index -> generator
165
+ # NOTE(woosuk): The indices of the requests that do not have their own
166
+ # generator should not be included in the dictionary.
167
+ self.generators: dict[int, torch.Generator] = {}
168
+
169
+ self.num_logprobs: dict[str, int] = {}
170
+ # NOTE(rob): num_prompt_logprobs only includes reqs
171
+ # that are currently in the prefill phase.
172
+ self.num_prompt_logprobs: dict[str, int] = {}
173
+
174
+ # To accumulate prompt logprobs tensor chunks across prefill steps.
175
+ self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}
176
+
177
+ self.logit_bias: list[Optional[dict[int,
178
+ float]]] = [None] * max_num_reqs
179
+ self.has_allowed_token_ids: set[str] = set()
180
+ # NOTE(lufang): In the mask tensor, if the corresponding token allowed,
181
+ # the value is False. Since we use masked_fill_ to set -inf.
182
+ self.allowed_token_ids_mask: Optional[torch.Tensor] = None
183
+ self.allowed_token_ids_mask_cpu_tensor: Optional[torch.Tensor] = None
184
+
185
+ # req_index -> bad_words_token_ids
186
+ self.bad_words_token_ids: dict[int, list[list[int]]] = {}
187
+
188
+ self.req_output_token_ids: list[Optional[list[int]]] = []
189
+
190
+ @property
191
+ def req_ids(self) -> list[str]:
192
+ # None elements should only be present transiently
193
+ # while performing state updates to the batch.
194
+ return cast(list[str], self._req_ids)
195
+
196
+ def add_request(
197
+ self,
198
+ request: "CachedRequestState",
199
+ req_index: Optional[int] = None,
200
+ ) -> None:
201
+ if req_index is None:
202
+ req_index = self.num_reqs
203
+ assert req_index < self.max_num_reqs
204
+
205
+ req_id = request.req_id
206
+ if req_index == len(self._req_ids):
207
+ self._req_ids.append(req_id)
208
+ self.req_output_token_ids.append(request.output_token_ids)
209
+ else:
210
+ self._req_ids[req_index] = req_id
211
+ self.req_output_token_ids[req_index] = request.output_token_ids
212
+
213
+ self.req_id_to_index[req_id] = req_index
214
+
215
+ # Copy the prompt token ids and output token ids.
216
+ num_prompt_tokens = len(request.prompt_token_ids)
217
+ self.num_prompt_tokens[req_index] = num_prompt_tokens
218
+ self.token_ids_cpu[
219
+ req_index, :num_prompt_tokens] = request.prompt_token_ids
220
+ start_idx = num_prompt_tokens
221
+ end_idx = start_idx + len(request.output_token_ids)
222
+ self.token_ids_cpu[req_index,
223
+ start_idx:end_idx] = request.output_token_ids
224
+ # Number of token ids in token_ids_cpu.
225
+ # NOTE(woosuk): This may include spec decode tokens.
226
+ self.num_tokens[req_index] = request.num_tokens
227
+ # Number of tokens without spec decode tokens.
228
+ self.num_tokens_no_spec[req_index] = request.num_tokens
229
+
230
+ self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
231
+ self.block_table.add_row(request.block_ids, req_index)
232
+
233
+ sampling_params = request.sampling_params
234
+ if sampling_params.sampling_type == SamplingType.GREEDY:
235
+ # Avoid later division by zero.
236
+ self.temperature_cpu[req_index] = -1.0
237
+ self.greedy_reqs.add(req_id)
238
+ else:
239
+ self.temperature_cpu[req_index] = sampling_params.temperature
240
+ self.random_reqs.add(req_id)
241
+
242
+ self.top_p_cpu[req_index] = sampling_params.top_p
243
+ if sampling_params.top_p < 1:
244
+ self.top_p_reqs.add(req_id)
245
+ top_k = sampling_params.top_k
246
+ if 0 < top_k < self.vocab_size:
247
+ self.top_k_reqs.add(req_id)
248
+ else:
249
+ top_k = self.vocab_size
250
+ self.top_k_cpu[req_index] = top_k
251
+ self.min_p_cpu[req_index] = sampling_params.min_p
252
+ self.frequency_penalties_cpu[
253
+ req_index] = sampling_params.frequency_penalty
254
+ if sampling_params.min_p > _SAMPLING_EPS:
255
+ self.min_p_reqs.add(req_id)
256
+ if sampling_params.frequency_penalty != 0.0:
257
+ self.frequency_penalties_reqs.add(req_id)
258
+ self.presence_penalties_cpu[
259
+ req_index] = sampling_params.presence_penalty
260
+ if sampling_params.presence_penalty != 0.0:
261
+ self.presence_penalties_reqs.add(req_id)
262
+ self.repetition_penalties_cpu[
263
+ req_index] = sampling_params.repetition_penalty
264
+ if sampling_params.repetition_penalty != 1.0:
265
+ self.repetition_penalties_reqs.add(req_id)
266
+ if sampling_params.min_tokens:
267
+ self.min_tokens[req_index] = (sampling_params.min_tokens,
268
+ sampling_params.all_stop_token_ids)
269
+
270
+ # NOTE(woosuk): self.generators should not include the requests that
271
+ # do not have their own generator.
272
+ if request.generator is not None:
273
+ self.generators[req_index] = request.generator
274
+
275
+ if sampling_params.logprobs is not None:
276
+ self.num_logprobs[req_id] = sampling_params.logprobs
277
+ if sampling_params.prompt_logprobs is not None:
278
+ self.num_prompt_logprobs[req_id] = sampling_params.prompt_logprobs
279
+ if sampling_params.logit_bias is not None:
280
+ self.logit_bias[req_index] = sampling_params.logit_bias
281
+
282
+ if sampling_params.allowed_token_ids:
283
+ self.has_allowed_token_ids.add(req_id)
284
+ if self.allowed_token_ids_mask_cpu_tensor is None:
285
+ # Lazy allocation for this tensor, which can be large.
286
+ # False means we don't fill with -inf.
287
+ self.allowed_token_ids_mask = torch.zeros(self.max_num_reqs,
288
+ self.vocab_size,
289
+ dtype=torch.bool,
290
+ device=self.device)
291
+ self.allowed_token_ids_mask_cpu_tensor = torch.zeros(
292
+ self.max_num_reqs,
293
+ self.vocab_size,
294
+ dtype=torch.bool,
295
+ device="cpu")
296
+ self.allowed_token_ids_mask_cpu_tensor[req_index] = True
297
+ # False means we don't fill with -inf.
298
+ self.allowed_token_ids_mask_cpu_tensor[req_index][
299
+ sampling_params.allowed_token_ids] = False
300
+
301
+ if sampling_params.bad_words_token_ids:
302
+ self.bad_words_token_ids[
303
+ req_index] = sampling_params.bad_words_token_ids
304
+
305
+ # Add request lora ID
306
+ if request.lora_request:
307
+ lora_id = request.lora_request.lora_int_id
308
+ if lora_id not in self.lora_id_to_request_ids:
309
+ self.lora_id_to_request_ids[lora_id] = set()
310
+
311
+ self.request_lora_mapping[req_index] = lora_id
312
+ self.lora_id_to_request_ids[lora_id].add(request.req_id)
313
+ self.lora_id_to_lora_request[lora_id] = request.lora_request
314
+ else:
315
+ # No LoRA
316
+ self.request_lora_mapping[req_index] = 0
317
+
318
+ def remove_request(self, req_id: str) -> Optional[int]:
319
+ """This method must always be followed by a call to condense()."""
320
+
321
+ req_index = self.req_id_to_index.pop(req_id, None)
322
+ if req_index is None:
323
+ return None
324
+ self._req_ids[req_index] = None
325
+ self.req_output_token_ids[req_index] = None
326
+
327
+ self.greedy_reqs.discard(req_id)
328
+ self.random_reqs.discard(req_id)
329
+ self.top_p_reqs.discard(req_id)
330
+ self.top_k_reqs.discard(req_id)
331
+ self.min_p_reqs.discard(req_id)
332
+ self.min_tokens.pop(req_index, None)
333
+ self.frequency_penalties_reqs.discard(req_id)
334
+ self.presence_penalties_reqs.discard(req_id)
335
+ self.repetition_penalties_reqs.discard(req_id)
336
+ self.generators.pop(req_index, None)
337
+ self.num_logprobs.pop(req_id, None)
338
+ self.num_prompt_logprobs.pop(req_id, None)
339
+ self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
340
+
341
+ # LoRA
342
+ lora_id = self.request_lora_mapping[req_index]
343
+ if lora_id != 0:
344
+ self.lora_id_to_request_ids[lora_id].discard(req_id)
345
+ if len(self.lora_id_to_request_ids[lora_id]) == 0:
346
+ self.lora_id_to_request_ids.pop(lora_id)
347
+ self.lora_id_to_lora_request.pop(lora_id)
348
+ self.request_lora_mapping[req_index] = 0
349
+
350
+ self.logit_bias[req_index] = None
351
+ self.has_allowed_token_ids.discard(req_id)
352
+ if self.allowed_token_ids_mask_cpu_tensor is not None:
353
+ # False means we don't fill with -inf.
354
+ self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
355
+ self.bad_words_token_ids.pop(req_index, None)
356
+ return req_index
357
+
358
+ def swap_states(self, i1: int, i2: int) -> None:
359
+ old_id_i1 = self._req_ids[i1]
360
+ old_id_i2 = self._req_ids[i2]
361
+ self._req_ids[i1], self._req_ids[i2] =\
362
+ self._req_ids[i2], self._req_ids[i1] # noqa
363
+ self.req_output_token_ids[i1], self.req_output_token_ids[i2] =\
364
+ self.req_output_token_ids[i2], self.req_output_token_ids[i1]
365
+ assert old_id_i1 is not None and old_id_i2 is not None
366
+ self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] =\
367
+ self.req_id_to_index[old_id_i2], self.req_id_to_index[old_id_i1]
368
+ self.num_tokens[i1], self.num_tokens[i2] =\
369
+ self.num_tokens[i2], self.num_tokens[i1]
370
+ self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] =\
371
+ self.num_tokens_no_spec[i2], self.num_tokens_no_spec[i1]
372
+ self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] =\
373
+ self.num_prompt_tokens[i2], self.num_prompt_tokens[i1]
374
+ self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] =\
375
+ self.num_computed_tokens_cpu[i2], self.num_computed_tokens_cpu[i1]
376
+ self.temperature_cpu[i1], self.temperature_cpu[i2] =\
377
+ self.temperature_cpu[i2], self.temperature_cpu[i1]
378
+ self.top_p_cpu[i1], self.top_p_cpu[i2] =\
379
+ self.top_p_cpu[i2], self.top_p_cpu[i1]
380
+ self.top_k_cpu[i1], self.top_k_cpu[i2] =\
381
+ self.top_k_cpu[i2], self.top_k_cpu[i1]
382
+ self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] =\
383
+ self.frequency_penalties_cpu[i2], self.frequency_penalties_cpu[i1]
384
+ self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] =\
385
+ self.presence_penalties_cpu[i2], self.presence_penalties_cpu[i1]
386
+ self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] =\
387
+ self.repetition_penalties_cpu[i2], self.repetition_penalties_cpu[i1]
388
+ self.min_p_cpu[i1], self.min_p_cpu[i2] =\
389
+ self.min_p_cpu[i2], self.min_p_cpu[i1]
390
+
391
+ # NOTE: the following is unsafe
392
+ # self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\
393
+ # self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...]
394
+ # instead, we need to temporiarily copy the data for one of the indices
395
+ # TODO(lucas): optimize this by only copying valid indices
396
+ tmp = self.token_ids_cpu[i1, ...].copy()
397
+ self.token_ids_cpu[i1, ...] = self.token_ids_cpu[i2, ...]
398
+ self.token_ids_cpu[i2, ...] = tmp
399
+
400
+ swap_dict_values(self.generators, i1, i2)
401
+ swap_dict_values(self.min_tokens, i1, i2)
402
+ swap_dict_values(self.bad_words_token_ids, i1, i2)
403
+
404
+ self.request_lora_mapping[i1], self.request_lora_mapping[i2] =\
405
+ self.request_lora_mapping[i2], self.request_lora_mapping[i1]
406
+ self.logit_bias[i1], self.logit_bias[i2] =\
407
+ self.logit_bias[i2], self.logit_bias[i1]
408
+
409
+ if self.allowed_token_ids_mask_cpu_tensor is not None:
410
+ self.allowed_token_ids_mask_cpu_tensor[i1], \
411
+ self.allowed_token_ids_mask_cpu_tensor[i2] =\
412
+ self.allowed_token_ids_mask_cpu_tensor[i2], \
413
+ self.allowed_token_ids_mask_cpu_tensor[i1]
414
+ self.block_table.swap_row(i1, i2)
415
+
416
+ def condense(self, empty_req_indices: list[int]) -> None:
417
+ """Move non-empty requests down into lower, empty indices.
418
+
419
+ Args:
420
+ empty_req_indices: empty batch indices, sorted descending.
421
+ """
422
+ num_reqs = self.num_reqs
423
+ if num_reqs == 0:
424
+ # The batched states are empty.
425
+ self._req_ids.clear()
426
+ self.req_output_token_ids.clear()
427
+ return
428
+
429
+ # NOTE(woosuk): This function assumes that the empty_req_indices
430
+ # is sorted in descending order.
431
+ last_req_index = num_reqs + len(empty_req_indices) - 1
432
+ while empty_req_indices:
433
+ # Find the largest non-empty index.
434
+ while last_req_index in empty_req_indices:
435
+ last_req_index -= 1
436
+
437
+ # Find the smallest empty index.
438
+ empty_index = empty_req_indices.pop()
439
+ if empty_index >= last_req_index:
440
+ break
441
+
442
+ # Swap the states.
443
+ req_id = self._req_ids[last_req_index]
444
+ output_token_ids = self.req_output_token_ids[last_req_index]
445
+ assert req_id is not None
446
+ self._req_ids[empty_index] = req_id
447
+ self._req_ids[last_req_index] = None
448
+ self.req_output_token_ids[empty_index] = output_token_ids
449
+ self.req_output_token_ids[last_req_index] = None
450
+ self.req_id_to_index[req_id] = empty_index
451
+
452
+ num_tokens = self.num_tokens[last_req_index]
453
+ self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[
454
+ last_req_index, :num_tokens]
455
+ self.num_tokens[empty_index] = num_tokens
456
+ self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
457
+ last_req_index]
458
+ self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[
459
+ last_req_index]
460
+ self.num_computed_tokens_cpu[
461
+ empty_index] = self.num_computed_tokens_cpu[last_req_index]
462
+ self.block_table.move_row(last_req_index, empty_index)
463
+ self.temperature_cpu[empty_index] = self.temperature_cpu[
464
+ last_req_index]
465
+ self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index]
466
+ self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index]
467
+ self.frequency_penalties_cpu[
468
+ empty_index] = self.frequency_penalties_cpu[last_req_index]
469
+ self.presence_penalties_cpu[
470
+ empty_index] = self.presence_penalties_cpu[last_req_index]
471
+ self.repetition_penalties_cpu[
472
+ empty_index] = self.repetition_penalties_cpu[last_req_index]
473
+ self.min_p_cpu[empty_index] = self.min_p_cpu[last_req_index]
474
+ generator = self.generators.pop(last_req_index, None)
475
+ if generator is not None:
476
+ self.generators[empty_index] = generator
477
+
478
+ min_token = self.min_tokens.pop(last_req_index, None)
479
+ if min_token is not None:
480
+ self.min_tokens[empty_index] = min_token
481
+
482
+ self.request_lora_mapping[empty_index] = self.request_lora_mapping[
483
+ last_req_index]
484
+
485
+ self.logit_bias[empty_index] = self.logit_bias[last_req_index]
486
+
487
+ if self.allowed_token_ids_mask_cpu_tensor is not None:
488
+ self.allowed_token_ids_mask_cpu_tensor[
489
+ empty_index] = self.allowed_token_ids_mask_cpu_tensor[
490
+ last_req_index]
491
+
492
+ bad_words_token_ids = self.bad_words_token_ids.pop(
493
+ last_req_index, None)
494
+ if bad_words_token_ids is not None:
495
+ self.bad_words_token_ids[empty_index] = bad_words_token_ids
496
+ # Decrement last_req_index since it is now empty.
497
+ last_req_index -= 1
498
+
499
+ # Trim lists to the batch size.
500
+ del self._req_ids[self.num_reqs:]
501
+ del self.req_output_token_ids[self.num_reqs:]
502
+
503
+ def _make_prompt_token_ids_tensor(self) -> torch.Tensor:
504
+ max_prompt_len = self.num_prompt_tokens[:self.num_reqs].max()
505
+ prompt_token_ids_cpu_tensor = torch.empty(
506
+ (self.num_reqs, max_prompt_len),
507
+ device="cpu",
508
+ dtype=torch.int64,
509
+ pin_memory=self.pin_memory,
510
+ )
511
+ prompt_token_ids = prompt_token_ids_cpu_tensor.numpy()
512
+ prompt_token_ids[:] = self.token_ids_cpu[:self.
513
+ num_reqs, :max_prompt_len]
514
+ # Use the value of vocab_size as a pad since we don't have a
515
+ # token_id of this value.
516
+ for i in range(self.num_reqs):
517
+ prompt_token_ids[i, self.num_prompt_tokens[i]:] = self.vocab_size
518
+ return prompt_token_ids_cpu_tensor.to(device=self.device,
519
+ non_blocking=True)
520
+
521
+ def make_lora_inputs(
522
+ self, num_scheduled_tokens: np.ndarray
523
+ ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
524
+ """
525
+ Given the num_scheduled_tokens for each request in the batch, return
526
+ datastructures used to activate the current LoRAs.
527
+ Returns:
528
+ 1. prompt_lora_mapping: A tuple of size self.num_reqs where,
529
+ prompt_lora_mapping[i] is the LoRA id to use for the ith prompt.
530
+ 2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens)
531
+ where, token_lora_mapping[i] is the LoRA id to use for ith token.
532
+ 3. lora_requests: Set of relevant LoRA requests.
533
+ """
534
+
535
+ req_lora_mapping = self.request_lora_mapping[:self.num_reqs]
536
+ prompt_lora_mapping = tuple(req_lora_mapping)
537
+ token_lora_mapping = tuple(
538
+ req_lora_mapping.repeat(num_scheduled_tokens))
539
+ active_lora_requests: set[LoRARequest] = set(
540
+ self.lora_id_to_lora_request.values())
541
+
542
+ return prompt_lora_mapping, token_lora_mapping, active_lora_requests
543
+
544
+ @property
545
+ def num_reqs(self) -> int:
546
+ return len(self.req_id_to_index)
547
+
548
+ @property
549
+ def all_greedy(self) -> bool:
550
+ return len(self.random_reqs) == 0
551
+
552
+ @property
553
+ def all_random(self) -> bool:
554
+ return len(self.greedy_reqs) == 0
555
+
556
+ @property
557
+ def no_top_p(self) -> bool:
558
+ return len(self.top_p_reqs) == 0
559
+
560
+ @property
561
+ def no_top_k(self) -> bool:
562
+ return len(self.top_k_reqs) == 0
563
+
564
+ @property
565
+ def no_min_p(self) -> bool:
566
+ return len(self.min_p_reqs) == 0
567
+
568
+ @property
569
+ def no_penalties(self) -> bool:
570
+ return (len(self.presence_penalties_reqs) == 0
571
+ and len(self.frequency_penalties_reqs) == 0
572
+ and len(self.repetition_penalties_reqs) == 0)
573
+
574
+ @property
575
+ def max_num_logprobs(self) -> Optional[int]:
576
+ return max(self.num_logprobs.values()) if self.num_logprobs else None
577
+
578
+ @property
579
+ def no_prompt_logprob(self) -> bool:
580
+ return not self.num_prompt_logprobs
581
+
582
+ @property
583
+ def no_allowed_token_ids(self) -> bool:
584
+ return len(self.has_allowed_token_ids) == 0
vllm/v1/worker/tpu_model_runner.py
+1
-1
@@ -42,8 +42,8 @@ from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,
42
from vllm.v1.sample.tpu.metadata import TPUSupportedSamplingMetadata
43
from vllm.v1.sample.tpu.sampler import Sampler as TPUSampler
44
from vllm.v1.utils import bind_kv_cache
45
-from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
45
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
46
+from vllm.v1.worker.tpu_input_batch import CachedRequestState, InputBatch
47
48
from .utils import (initialize_kv_cache_for_kv_sharing,
49
sanity_check_mm_encoder_outputs)