[Cleanup] Remove vestigial `tpu_input_batch.py` (#56898)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Sep 14, 2026 at 21:29 UTC
7dbc3d62c388df67d2764f28c90c51cc11ff0329
4 files changed
+9
-574
tests/v1/e2e/general/test_mamba_prefix_cache.py
+3
-4
@@ -31,9 +31,8 @@ from vllm.v1.request import Request
31
from vllm.v1.sample.metadata import SamplingMetadata
32
from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
33
from vllm.v1.worker import mamba_utils
34
-from vllm.v1.worker.gpu_input_batch import CachedRequestState
34
+from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
35
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
36
-from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
36
from vllm.v1.worker.mamba_utils import get_mamba_groups
37
38
@@ -315,7 +314,7 @@ def get_fake_process_mamba_fn(
314
action: tuple[int, int],
315
kv_cache_config: KVCacheConfig,
316
forward_context: dict[str, Any],
318
- input_batch: GPUInputBatch,
317
+ input_batch: InputBatch,
318
):
319
assert copy_info is not None
320
if action == (-1, -1):
@@ -371,7 +370,7 @@ def get_fake_process_mamba_fn(
370
kv_cache_config: KVCacheConfig,
371
cache_config: CacheConfig,
372
mamba_state_idx: dict[str, int],
374
- input_batch: GPUInputBatch,
373
+ input_batch: InputBatch,
374
requests: dict[str, CachedRequestState],
375
forward_context: dict[str, Any],
376
mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
vllm/v1/worker/lora_model_runner_mixin.py
+1
-5
@@ -6,7 +6,6 @@ Define LoRA functionality mixin for model runners.
6
7
from collections.abc import Callable
8
from contextlib import contextmanager
9
-from typing import TypeAlias
9
10
import numpy as np
11
import torch
@@ -19,10 +18,7 @@ from vllm.lora.layers import LoRAMapping, LoRAMappingType
18
from vllm.lora.request import LoRARequest
19
from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager
20
from vllm.model_executor.models import supports_lora
22
-from vllm.v1.worker.gpu_input_batch import InputBatch as GPUInputBatch
23
-from vllm.v1.worker.tpu_input_batch import InputBatch as TPUInputBatch
24
-
25
-InputBatch: TypeAlias = TPUInputBatch | GPUInputBatch
21
+from vllm.v1.worker.gpu_input_batch import InputBatch
22
23
logger = init_logger(__name__)
24
vllm/v1/worker/mamba_utils.py
+5
-6
@@ -26,8 +26,7 @@ from vllm.v1.kv_cache_interface import (
26
UniformTypeKVCacheSpecs,
27
)
28
from vllm.v1.utils import CpuGpuBuffer
29
-from vllm.v1.worker.gpu_input_batch import CachedRequestState
30
-from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
29
+from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
30
31
logger = init_logger(__name__)
32
@@ -1436,7 +1435,7 @@ def preprocess_mamba(
1435
kv_cache_config: KVCacheConfig,
1436
cache_config: CacheConfig,
1437
mamba_state_idx: dict[str, int],
1439
- input_batch: GPUInputBatch,
1438
+ input_batch: InputBatch,
1439
requests: dict[str, CachedRequestState],
1440
forward_context: dict[str, Any],
1441
mamba_state_copy_funcs: MambaStateCopyFuncsByType,
@@ -1541,7 +1540,7 @@ def preprocess_mamba(
1540
def postprocess_mamba_all(
1541
scheduler_output: SchedulerOutput,
1542
kv_cache_config: KVCacheConfig,
1544
- input_batch: GPUInputBatch,
1543
+ input_batch: InputBatch,
1544
requests: dict[str, CachedRequestState],
1545
mamba_state_idx: dict[str, int],
1546
num_spec_tokens: int,
@@ -1572,7 +1571,7 @@ def postprocess_mamba_all(
1571
1572
def preprocess_mamba_all_specdec(
1573
scheduler_output: SchedulerOutput,
1575
- input_batch: GPUInputBatch,
1574
+ input_batch: InputBatch,
1575
mamba_state_idx: dict[str, int],
1576
num_reqs: int,
1577
prev_last_scheduled_idx_buf: CpuGpuBuffer,
@@ -1591,7 +1590,7 @@ def postprocess_mamba_align_gpu(
1590
num_reqs: int,
1591
num_accepted_tokens_gpu: torch.Tensor,
1592
num_accepted_tokens_cpu_tensor: torch.Tensor,
1594
- input_batch: GPUInputBatch,
1593
+ input_batch: InputBatch,
1594
kv_cache_config: KVCacheConfig,
1595
forward_context: dict[str, Any],
1596
mamba_state_copy_funcs: MambaStateCopyFuncsByType,
vllm/v1/worker/tpu_input_batch.py
deleted
-559
@@ -1,559 +0,0 @@
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 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 length_from_prompt_token_ids_or_embeds
13
-from vllm.utils.collection_utils import swap_dict_values
14
-from vllm.v1.outputs import LogprobsTensors
15
-from vllm.v1.worker.block_table import MultiGroupBlockTable
16
-from vllm.v1.worker.gpu_input_batch import CachedRequestState
17
-
18
-_SAMPLING_EPS = 1e-5
19
-
20
-
21
-class InputBatch:
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
- kernel_block_sizes: list[int],
32
- max_num_blocks_per_req: list[int],
33
- ):
34
- self.max_num_reqs = max_num_reqs
35
- self.max_model_len = max_model_len
36
- self.max_num_batched_tokens = max_num_batched_tokens
37
- self.device = device
38
- self.pin_memory = pin_memory
39
- self.vocab_size = vocab_size
40
-
41
- self._req_ids: list[str | None] = []
42
- self.req_id_to_index: dict[str, int] = {}
43
-
44
- # TODO(woosuk): This buffer could be too large if max_model_len is big.
45
- # Find a way to reduce the CPU memory usage.
46
- # This buffer is not directly transferred to the GPU, so it does not
47
- # need to be pinned.
48
- self.token_ids_cpu_tensor = torch.zeros(
49
- (max_num_reqs, max_model_len),
50
- device="cpu",
51
- dtype=torch.int32,
52
- pin_memory=False,
53
- )
54
- self.token_ids_cpu = self.token_ids_cpu_tensor.numpy()
55
- self.num_tokens_no_spec = np.zeros(max_num_reqs, dtype=np.int32)
56
- self.num_prompt_tokens = np.zeros(max_num_reqs, dtype=np.int32)
57
- self.num_computed_tokens_cpu_tensor = torch.zeros(
58
- (max_num_reqs,),
59
- device="cpu",
60
- dtype=torch.int32,
61
- pin_memory=pin_memory,
62
- )
63
- self.num_computed_tokens_cpu = 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_num_batched_tokens=max_num_batched_tokens,
69
- pin_memory=pin_memory,
70
- device=device,
71
- block_sizes=block_sizes,
72
- kernel_block_sizes=kernel_block_sizes,
73
- max_num_blocks=max_num_blocks_per_req,
74
- )
75
-
76
- # Sampling-related.
77
- self.temperature = torch.empty(
78
- (max_num_reqs,), dtype=torch.float32, device=device
79
- )
80
- self.temperature_cpu_tensor = torch.empty(
81
- (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory
82
- )
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,), dtype=torch.float32, device=device)
88
- self.top_p_cpu_tensor = torch.empty(
89
- (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory
90
- )
91
- self.top_p_cpu = self.top_p_cpu_tensor.numpy()
92
- self.top_p_reqs: set[str] = set()
93
-
94
- self.top_k = torch.empty((max_num_reqs,), dtype=torch.int32, device=device)
95
- self.top_k_cpu_tensor = torch.empty(
96
- (max_num_reqs,), dtype=torch.int32, device="cpu", pin_memory=pin_memory
97
- )
98
- self.top_k_cpu = self.top_k_cpu_tensor.numpy()
99
- self.top_k_reqs: set[str] = set()
100
-
101
- self.min_p = torch.empty((max_num_reqs,), dtype=torch.float32, device=device)
102
- self.min_p_cpu_tensor = torch.empty(
103
- (max_num_reqs,), dtype=torch.float32, device="cpu", pin_memory=pin_memory
104
- )
105
- self.min_p_cpu = self.min_p_cpu_tensor.numpy()
106
- self.min_p_reqs: set[str] = set()
107
-
108
- # Frequency penalty related data structures
109
- self.frequency_penalties = torch.empty(
110
- (max_num_reqs,), dtype=torch.float, device=device
111
- )
112
- self.frequency_penalties_cpu_tensor = torch.empty(
113
- (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory
114
- )
115
- self.frequency_penalties_cpu = self.frequency_penalties_cpu_tensor.numpy()
116
- self.frequency_penalties_reqs: set[str] = set()
117
-
118
- # Presence penalty related data structures
119
- self.presence_penalties = torch.empty(
120
- (max_num_reqs,), dtype=torch.float, device=device
121
- )
122
- self.presence_penalties_cpu_tensor = torch.empty(
123
- (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory
124
- )
125
- self.presence_penalties_cpu = self.presence_penalties_cpu_tensor.numpy()
126
- self.presence_penalties_reqs: set[str] = set()
127
-
128
- # Repetition penalty related data structures
129
- self.repetition_penalties = torch.empty(
130
- (max_num_reqs,), dtype=torch.float, device=device
131
- )
132
- self.repetition_penalties_cpu_tensor = torch.empty(
133
- (max_num_reqs,), dtype=torch.float, device="cpu", pin_memory=pin_memory
134
- )
135
- self.repetition_penalties_cpu = self.repetition_penalties_cpu_tensor.numpy()
136
- self.repetition_penalties_reqs: set[str] = set()
137
-
138
- # req_index -> (min_tokens, stop_token_ids)
139
- self.min_tokens: dict[int, tuple[int, set[int]]] = {}
140
-
141
- # lora related
142
- self.request_lora_mapping = np.zeros((self.max_num_reqs,), dtype=np.int64)
143
- self.lora_id_to_request_ids: dict[int, set[str]] = {}
144
- self.lora_id_to_lora_request: dict[int, LoRARequest] = {}
145
-
146
- # req_index -> generator
147
- # NOTE(woosuk): The indices of the requests that do not have their own
148
- # generator should not be included in the dictionary.
149
- self.generators: dict[int, torch.Generator] = {}
150
-
151
- self.num_logprobs: dict[str, int] = {}
152
-
153
- # To accumulate prompt logprobs tensor chunks across prefill steps.
154
- self.in_progress_prompt_logprobs_cpu: dict[str, LogprobsTensors] = {}
155
-
156
- self.logit_bias: list[dict[int, float] | None] = [None] * max_num_reqs
157
- self.has_allowed_token_ids: set[str] = set()
158
- # NOTE(lufang): In the mask tensor, if the corresponding token allowed,
159
- # the value is False. Since we use masked_fill_ to set -inf.
160
- self.allowed_token_ids_mask: torch.Tensor | None = None
161
- self.allowed_token_ids_mask_cpu_tensor: torch.Tensor | None = None
162
-
163
- # req_index -> bad_words_token_ids
164
- self.bad_words_token_ids: dict[int, list[list[int]]] = {}
165
-
166
- self.req_output_token_ids: list[list[int] | None] = []
167
-
168
- @property
169
- def req_ids(self) -> list[str]:
170
- # None elements should only be present transiently
171
- # while performing state updates to the batch.
172
- return cast(list[str], self._req_ids)
173
-
174
- def add_request(
175
- self,
176
- request: "CachedRequestState",
177
- req_index: int | None = None,
178
- ) -> None:
179
- if req_index is None:
180
- req_index = self.num_reqs
181
- assert req_index < self.max_num_reqs
182
-
183
- req_id = request.req_id
184
- if req_index == len(self._req_ids):
185
- self._req_ids.append(req_id)
186
- self.req_output_token_ids.append(request.output_token_ids)
187
- else:
188
- self._req_ids[req_index] = req_id
189
- self.req_output_token_ids[req_index] = request.output_token_ids
190
-
191
- self.req_id_to_index[req_id] = req_index
192
-
193
- # Copy the prompt token ids and output token ids.
194
- num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
195
- request.prompt_token_ids, request.prompt_embeds
196
- )
197
- # TODO: copy prompt_embeds
198
- self.num_prompt_tokens[req_index] = num_prompt_tokens
199
- self.token_ids_cpu[req_index, :num_prompt_tokens] = request.prompt_token_ids
200
- start_idx = num_prompt_tokens
201
- end_idx = start_idx + len(request.output_token_ids)
202
- self.token_ids_cpu[req_index, start_idx:end_idx] = request.output_token_ids
203
- # Number of tokens without spec decode tokens.
204
- self.num_tokens_no_spec[req_index] = request.num_tokens
205
-
206
- self.num_computed_tokens_cpu[req_index] = request.num_computed_tokens
207
- self.block_table.add_row(request.block_ids, req_index)
208
-
209
- sampling_params = request.sampling_params
210
- assert sampling_params is not None, "pooling requests not supported yet"
211
- if sampling_params.sampling_type == SamplingType.GREEDY:
212
- # Should avoid division by zero later when apply_temperature.
213
- self.temperature_cpu[req_index] = 0.0
214
- self.greedy_reqs.add(req_id)
215
- else:
216
- self.temperature_cpu[req_index] = sampling_params.temperature
217
- self.random_reqs.add(req_id)
218
-
219
- self.top_p_cpu[req_index] = sampling_params.top_p
220
- if sampling_params.top_p < 1:
221
- self.top_p_reqs.add(req_id)
222
- top_k = sampling_params.top_k
223
- if 0 < top_k < self.vocab_size:
224
- self.top_k_reqs.add(req_id)
225
- else:
226
- top_k = self.vocab_size
227
- self.top_k_cpu[req_index] = top_k
228
- self.min_p_cpu[req_index] = sampling_params.min_p
229
- self.frequency_penalties_cpu[req_index] = sampling_params.frequency_penalty
230
- if sampling_params.min_p > _SAMPLING_EPS:
231
- self.min_p_reqs.add(req_id)
232
- if sampling_params.frequency_penalty != 0.0:
233
- self.frequency_penalties_reqs.add(req_id)
234
- self.presence_penalties_cpu[req_index] = sampling_params.presence_penalty
235
- if sampling_params.presence_penalty != 0.0:
236
- self.presence_penalties_reqs.add(req_id)
237
- self.repetition_penalties_cpu[req_index] = sampling_params.repetition_penalty
238
- if sampling_params.repetition_penalty != 1.0:
239
- self.repetition_penalties_reqs.add(req_id)
240
- if sampling_params.min_tokens:
241
- self.min_tokens[req_index] = (
242
- sampling_params.min_tokens,
243
- sampling_params.all_stop_token_ids,
244
- )
245
-
246
- # NOTE(woosuk): self.generators should not include the requests that
247
- # do not have their own generator.
248
- if request.generator is not None:
249
- self.generators[req_index] = request.generator
250
-
251
- if sampling_params.logprobs is not None:
252
- self.num_logprobs[req_id] = sampling_params.logprobs
253
- if sampling_params.logit_bias is not None:
254
- self.logit_bias[req_index] = sampling_params.logit_bias
255
-
256
- if sampling_params.allowed_token_ids:
257
- self.has_allowed_token_ids.add(req_id)
258
- if self.allowed_token_ids_mask_cpu_tensor is None:
259
- # Lazy allocation for this tensor, which can be large.
260
- # False means we don't fill with -inf.
261
- self.allowed_token_ids_mask = torch.zeros(
262
- self.max_num_reqs,
263
- self.vocab_size,
264
- dtype=torch.bool,
265
- device=self.device,
266
- )
267
- self.allowed_token_ids_mask_cpu_tensor = torch.zeros(
268
- self.max_num_reqs, self.vocab_size, dtype=torch.bool, device="cpu"
269
- )
270
- self.allowed_token_ids_mask_cpu_tensor[req_index] = True
271
- # False means we don't fill with -inf.
272
- self.allowed_token_ids_mask_cpu_tensor[req_index][
273
- sampling_params.allowed_token_ids
274
- ] = False
275
-
276
- if sampling_params.bad_words_token_ids:
277
- self.bad_words_token_ids[req_index] = sampling_params.bad_words_token_ids
278
-
279
- # Add request lora ID
280
- if request.lora_request:
281
- lora_id = request.lora_request.lora_int_id
282
- if lora_id not in self.lora_id_to_request_ids:
283
- self.lora_id_to_request_ids[lora_id] = set()
284
-
285
- self.request_lora_mapping[req_index] = lora_id
286
- self.lora_id_to_request_ids[lora_id].add(request.req_id)
287
- self.lora_id_to_lora_request[lora_id] = request.lora_request
288
- else:
289
- # No LoRA
290
- self.request_lora_mapping[req_index] = 0
291
-
292
- def remove_request(self, req_id: str) -> int | None:
293
- """This method must always be followed by a call to condense()."""
294
-
295
- req_index = self.req_id_to_index.pop(req_id, None)
296
- if req_index is None:
297
- return None
298
- self._req_ids[req_index] = None
299
- self.req_output_token_ids[req_index] = None
300
-
301
- self.greedy_reqs.discard(req_id)
302
- self.random_reqs.discard(req_id)
303
- self.top_p_reqs.discard(req_id)
304
- self.top_k_reqs.discard(req_id)
305
- self.min_p_reqs.discard(req_id)
306
- self.min_tokens.pop(req_index, None)
307
- self.frequency_penalties_reqs.discard(req_id)
308
- self.presence_penalties_reqs.discard(req_id)
309
- self.repetition_penalties_reqs.discard(req_id)
310
- self.generators.pop(req_index, None)
311
- self.num_logprobs.pop(req_id, None)
312
- self.in_progress_prompt_logprobs_cpu.pop(req_id, None)
313
-
314
- # LoRA
315
- lora_id = self.request_lora_mapping[req_index]
316
- if lora_id != 0:
317
- self.lora_id_to_request_ids[lora_id].discard(req_id)
318
- if len(self.lora_id_to_request_ids[lora_id]) == 0:
319
- self.lora_id_to_request_ids.pop(lora_id)
320
- self.lora_id_to_lora_request.pop(lora_id)
321
- self.request_lora_mapping[req_index] = 0
322
-
323
- self.logit_bias[req_index] = None
324
- self.has_allowed_token_ids.discard(req_id)
325
- if self.allowed_token_ids_mask_cpu_tensor is not None:
326
- # False means we don't fill with -inf.
327
- self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)
328
- self.bad_words_token_ids.pop(req_index, None)
329
- return req_index
330
-
331
- def swap_states(self, i1: int, i2: int) -> None:
332
- old_id_i1 = self._req_ids[i1]
333
- old_id_i2 = self._req_ids[i2]
334
- self._req_ids[i1], self._req_ids[i2] = self._req_ids[i2], self._req_ids[i1] # noqa
335
- self.req_output_token_ids[i1], self.req_output_token_ids[i2] = (
336
- self.req_output_token_ids[i2],
337
- self.req_output_token_ids[i1],
338
- )
339
- assert old_id_i1 is not None and old_id_i2 is not None
340
- self.req_id_to_index[old_id_i1], self.req_id_to_index[old_id_i2] = (
341
- self.req_id_to_index[old_id_i2],
342
- self.req_id_to_index[old_id_i1],
343
- )
344
- self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2] = (
345
- self.num_tokens_no_spec[i2],
346
- self.num_tokens_no_spec[i1],
347
- )
348
- self.num_prompt_tokens[i1], self.num_prompt_tokens[i2] = (
349
- self.num_prompt_tokens[i2],
350
- self.num_prompt_tokens[i1],
351
- )
352
- self.num_computed_tokens_cpu[i1], self.num_computed_tokens_cpu[i2] = (
353
- self.num_computed_tokens_cpu[i2],
354
- self.num_computed_tokens_cpu[i1],
355
- )
356
- self.temperature_cpu[i1], self.temperature_cpu[i2] = (
357
- self.temperature_cpu[i2],
358
- self.temperature_cpu[i1],
359
- )
360
- self.top_p_cpu[i1], self.top_p_cpu[i2] = self.top_p_cpu[i2], self.top_p_cpu[i1]
361
- self.top_k_cpu[i1], self.top_k_cpu[i2] = self.top_k_cpu[i2], self.top_k_cpu[i1]
362
- self.frequency_penalties_cpu[i1], self.frequency_penalties_cpu[i2] = (
363
- self.frequency_penalties_cpu[i2],
364
- self.frequency_penalties_cpu[i1],
365
- )
366
- self.presence_penalties_cpu[i1], self.presence_penalties_cpu[i2] = (
367
- self.presence_penalties_cpu[i2],
368
- self.presence_penalties_cpu[i1],
369
- )
370
- self.repetition_penalties_cpu[i1], self.repetition_penalties_cpu[i2] = (
371
- self.repetition_penalties_cpu[i2],
372
- self.repetition_penalties_cpu[i1],
373
- )
374
- self.min_p_cpu[i1], self.min_p_cpu[i2] = self.min_p_cpu[i2], self.min_p_cpu[i1]
375
-
376
- # NOTE: the following is unsafe
377
- # self.token_ids_cpu[i1, ...], self.token_ids_cpu[i2, ...], =\
378
- # self.token_ids_cpu[i2, ...], self.token_ids_cpu[i1, ...]
379
- # instead, we need to temporarily copy the data for one of the indices
380
- n = max(self.num_tokens_no_spec[i1], self.num_tokens_no_spec[i2])
381
- tmp = self.token_ids_cpu[i1, :n].copy()
382
- self.token_ids_cpu[i1, :n] = self.token_ids_cpu[i2, :n]
383
- self.token_ids_cpu[i2, :n] = tmp
384
-
385
- swap_dict_values(self.generators, i1, i2)
386
- swap_dict_values(self.min_tokens, i1, i2)
387
- swap_dict_values(self.bad_words_token_ids, i1, i2)
388
-
389
- self.request_lora_mapping[i1], self.request_lora_mapping[i2] = (
390
- self.request_lora_mapping[i2],
391
- self.request_lora_mapping[i1],
392
- )
393
- self.logit_bias[i1], self.logit_bias[i2] = (
394
- self.logit_bias[i2],
395
- self.logit_bias[i1],
396
- )
397
-
398
- if self.allowed_token_ids_mask_cpu_tensor is not None:
399
- (
400
- self.allowed_token_ids_mask_cpu_tensor[i1],
401
- self.allowed_token_ids_mask_cpu_tensor[i2],
402
- ) = (
403
- self.allowed_token_ids_mask_cpu_tensor[i2],
404
- self.allowed_token_ids_mask_cpu_tensor[i1],
405
- )
406
- self.block_table.swap_row(i1, i2)
407
-
408
- def condense(self, empty_req_indices: list[int]) -> None:
409
- """Move non-empty requests down into lower, empty indices.
410
-
411
- Args:
412
- empty_req_indices: empty batch indices, sorted descending.
413
- """
414
- num_reqs = self.num_reqs
415
- if num_reqs == 0:
416
- # The batched states are empty.
417
- self._req_ids.clear()
418
- self.req_output_token_ids.clear()
419
- return
420
-
421
- # NOTE(woosuk): This function assumes that the empty_req_indices
422
- # is sorted in descending order.
423
- last_req_index = num_reqs + len(empty_req_indices) - 1
424
- while empty_req_indices:
425
- # Find the largest non-empty index.
426
- while last_req_index in empty_req_indices:
427
- last_req_index -= 1
428
-
429
- # Find the smallest empty index.
430
- empty_index = empty_req_indices.pop()
431
- if empty_index >= last_req_index:
432
- break
433
-
434
- # Swap the states.
435
- req_id = self._req_ids[last_req_index]
436
- output_token_ids = self.req_output_token_ids[last_req_index]
437
- assert req_id is not None
438
- self._req_ids[empty_index] = req_id
439
- self._req_ids[last_req_index] = None
440
- self.req_output_token_ids[empty_index] = output_token_ids
441
- self.req_output_token_ids[last_req_index] = None
442
- self.req_id_to_index[req_id] = empty_index
443
-
444
- num_tokens = self.num_tokens_no_spec[last_req_index]
445
- self.token_ids_cpu[empty_index, :num_tokens] = self.token_ids_cpu[
446
- last_req_index, :num_tokens
447
- ]
448
- self.num_tokens_no_spec[empty_index] = self.num_tokens_no_spec[
449
- last_req_index
450
- ]
451
- self.num_prompt_tokens[empty_index] = self.num_prompt_tokens[last_req_index]
452
- self.num_computed_tokens_cpu[empty_index] = self.num_computed_tokens_cpu[
453
- last_req_index
454
- ]
455
- self.block_table.move_row(last_req_index, empty_index)
456
- self.temperature_cpu[empty_index] = self.temperature_cpu[last_req_index]
457
- self.top_p_cpu[empty_index] = self.top_p_cpu[last_req_index]
458
- self.top_k_cpu[empty_index] = self.top_k_cpu[last_req_index]
459
- self.frequency_penalties_cpu[empty_index] = self.frequency_penalties_cpu[
460
- last_req_index
461
- ]
462
- self.presence_penalties_cpu[empty_index] = self.presence_penalties_cpu[
463
- last_req_index
464
- ]
465
- self.repetition_penalties_cpu[empty_index] = self.repetition_penalties_cpu[
466
- last_req_index
467
- ]
468
- self.min_p_cpu[empty_index] = self.min_p_cpu[last_req_index]
469
- generator = self.generators.pop(last_req_index, None)
470
- if generator is not None:
471
- self.generators[empty_index] = generator
472
-
473
- min_token = self.min_tokens.pop(last_req_index, None)
474
- if min_token is not None:
475
- self.min_tokens[empty_index] = min_token
476
-
477
- self.request_lora_mapping[empty_index] = self.request_lora_mapping[
478
- last_req_index
479
- ]
480
-
481
- self.logit_bias[empty_index] = self.logit_bias[last_req_index]
482
-
483
- if self.allowed_token_ids_mask_cpu_tensor is not None:
484
- self.allowed_token_ids_mask_cpu_tensor[empty_index] = (
485
- self.allowed_token_ids_mask_cpu_tensor[last_req_index]
486
- )
487
-
488
- bad_words_token_ids = self.bad_words_token_ids.pop(last_req_index, None)
489
- if bad_words_token_ids is not None:
490
- self.bad_words_token_ids[empty_index] = bad_words_token_ids
491
- # Decrement last_req_index since it is now empty.
492
- last_req_index -= 1
493
-
494
- # Trim lists to the batch size.
495
- del self._req_ids[self.num_reqs :]
496
- del self.req_output_token_ids[self.num_reqs :]
497
-
498
- def make_lora_inputs(
499
- self, num_scheduled_tokens: np.ndarray, num_sampled_tokens: np.ndarray
500
- ) -> tuple[tuple[int, ...], tuple[int, ...], set[LoRARequest]]:
501
- """
502
- Given the num_scheduled_tokens for each request in the batch, return
503
- datastructures used to activate the current LoRAs.
504
- Returns:
505
- 1. prompt_lora_mapping: A tuple of size self.num_reqs where,
506
- prompt_lora_mapping[i] is the LoRA id to use for the ith prompt.
507
- 2. token_lora_mapping: A tuple of size np.sum(num_scheduled_tokens)
508
- where, token_lora_mapping[i] is the LoRA id to use for ith token.
509
- 3. lora_requests: Set of relevant LoRA requests.
510
- """
511
-
512
- req_lora_mapping = self.request_lora_mapping[: self.num_reqs]
513
- prompt_lora_mapping = tuple(req_lora_mapping)
514
- token_lora_mapping = tuple(req_lora_mapping.repeat(num_scheduled_tokens))
515
- active_lora_requests: set[LoRARequest] = set(
516
- self.lora_id_to_lora_request.values()
517
- )
518
-
519
- return prompt_lora_mapping, token_lora_mapping, active_lora_requests
520
-
521
- @property
522
- def num_reqs(self) -> int:
523
- return len(self.req_id_to_index)
524
-
525
- @property
526
- def all_greedy(self) -> bool:
527
- return len(self.random_reqs) == 0
528
-
529
- @property
530
- def all_random(self) -> bool:
531
- return len(self.greedy_reqs) == 0
532
-
533
- @property
534
- def no_top_p(self) -> bool:
535
- return len(self.top_p_reqs) == 0
536
-
537
- @property
538
- def no_top_k(self) -> bool:
539
- return len(self.top_k_reqs) == 0
540
-
541
- @property
542
- def no_min_p(self) -> bool:
543
- return len(self.min_p_reqs) == 0
544
-
545
- @property
546
- def no_penalties(self) -> bool:
547
- return (
548
- len(self.presence_penalties_reqs) == 0
549
- and len(self.frequency_penalties_reqs) == 0
550
- and len(self.repetition_penalties_reqs) == 0
551
- )
552
-
553
- @property
554
- def max_num_logprobs(self) -> int | None:
555
- return max(self.num_logprobs.values()) if self.num_logprobs else None
556
-
557
- @property
558
- def no_allowed_token_ids(self) -> bool:
559
- return len(self.has_allowed_token_ids) == 0