| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import inspect |
| 5 | from collections.abc import Sequence |
| 6 | |
| 7 | import numpy as np |
| 8 | import pytest |
| 9 | import torch |
| 10 | |
| 11 | from vllm.platforms import current_platform |
| 12 | from vllm.sampling_params import SamplingParams |
| 13 | from vllm.utils.torch_utils import make_tensor_with_pad |
| 14 | from vllm.v1.pool.metadata import PoolingMetadata |
| 15 | from vllm.v1.sample.logits_processor import LogitsProcessors |
| 16 | from vllm.v1.sample.metadata import SamplingMetadata |
| 17 | from vllm.v1.utils import CpuGpuBuffer |
| 18 | from vllm.v1.worker.block_table import BlockTable, MultiGroupBlockTable |
| 19 | from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch |
| 20 | |
| 21 | VOCAB_SIZE = 1024 |
| 22 | NUM_OUTPUT_TOKENS = 20 |
| 23 | MAX_PROMPT_SIZE = 100 |
| 24 | DEVICE_TYPE = current_platform.device_type |
| 25 | DEVICES = [f"{DEVICE_TYPE}:{i}" for i in range(min(current_platform.device_count(), 2))] |
| 26 | MAX_NUM_PROMPT_TOKENS = 64 |
| 27 | |
| 28 | |
| 29 | def _compare_objs(obj1, obj2, skip: Sequence = ("logitsprocs", "batch_update_builder")): |
| 30 | attrs = inspect.getmembers(obj1, lambda a: not (inspect.isroutine(a))) |
| 31 | attr_names = set( |
| 32 | [a[0] for a in attrs if not (a[0].startswith("__") and a[0].endswith("__"))] |
| 33 | ) |
| 34 | for attr_name in attr_names: |
| 35 | if attr_name in skip: |
| 36 | continue |
| 37 | |
| 38 | a = getattr(obj1, attr_name) |
| 39 | b = getattr(obj2, attr_name) |
| 40 | |
| 41 | is_same = False |
| 42 | if isinstance(a, torch.Tensor): |
| 43 | if a.numel() == 0 or b.numel() == 0: |
| 44 | is_same = a.numel() == 0 and b.numel() == 0 |
| 45 | elif torch.allclose(a, b): |
| 46 | is_same = True |
| 47 | elif isinstance(a, np.ndarray): |
| 48 | if np.allclose(a, b): |
| 49 | is_same = True |
| 50 | elif isinstance(a, MultiGroupBlockTable): |
| 51 | for a_i, b_i in zip(a.block_tables, b.block_tables): |
| 52 | _compare_objs(a_i, b_i) |
| 53 | is_same = True |
| 54 | elif isinstance(a, (BlockTable, SamplingMetadata, PoolingMetadata)): |
| 55 | _compare_objs(a, b) |
| 56 | is_same = True # if we make it here must be same |
| 57 | elif a == b: |
| 58 | is_same = True |
| 59 | elif isinstance(a, CpuGpuBuffer): |
| 60 | is_same = np.allclose(a.np, b.np) and torch.allclose(a.gpu, b.gpu) |
| 61 | assert is_same, ( |
| 62 | f"Attribute {attr_name} is different in {obj1} and {obj2}: {a} != {b}" |
| 63 | ) |
| 64 | |
| 65 | |
| 66 | def _remove_requests( |
| 67 | input_batch: InputBatch, batch_size: int, reqs: list[CachedRequestState] |
| 68 | ) -> set[str]: |
| 69 | """Remove some requests randomly from the batch and returns |
| 70 | set of request removed |
| 71 | """ |
| 72 | num_reqs_to_remove = np.random.randint(0, batch_size) |
| 73 | req_indices_to_remove: set[int] = set() |
| 74 | for _ in range(num_reqs_to_remove): |
| 75 | req_index_to_remove = np.random.randint(0, batch_size) |
| 76 | req_indices_to_remove.add(req_index_to_remove) |
| 77 | |
| 78 | req_ids_to_remove: set[str] = set() |
| 79 | for index in req_indices_to_remove: |
| 80 | input_batch.remove_request(reqs[index].req_id) |
| 81 | req_ids_to_remove.add(reqs[index].req_id) |
| 82 | return req_ids_to_remove |
| 83 | |
| 84 | |
| 85 | def _construct_expected_sampling_metadata( |
| 86 | reqs: list[CachedRequestState], |
| 87 | req_ids_retained: set[int], |
| 88 | req_id_index_in_input_batch: dict[str, int], |
| 89 | device: torch.device, |
| 90 | ) -> SamplingMetadata: |
| 91 | """Constructs and returns the expected SamplingMetadata for this |
| 92 | batch. |
| 93 | """ |
| 94 | num_reqs = len(req_ids_retained) |
| 95 | output_token_ids: list[list[int]] = [list() for _ in range(num_reqs)] |
| 96 | prompt_token_ids: list[list[int]] = [list() for _ in range(num_reqs)] |
| 97 | presence_penalties = [0.0 for _ in range(num_reqs)] |
| 98 | frequency_penalties = [0.0 for _ in range(num_reqs)] |
| 99 | repetition_penalties = [1.0 for _ in range(num_reqs)] |
| 100 | top_k = [0 for _ in range(num_reqs)] |
| 101 | top_p = [0.0 for _ in range(num_reqs)] |
| 102 | temperature = [0.0 for _ in range(num_reqs)] |
| 103 | min_tokens = {} |
| 104 | logit_bias = [None] * num_reqs |
| 105 | allowed_token_ids_mask = torch.zeros( |
| 106 | num_reqs, VOCAB_SIZE, dtype=torch.bool, device=device |
| 107 | ) |
| 108 | bad_words_token_ids = {} |
| 109 | for req in reqs: |
| 110 | if req.req_id not in req_ids_retained: |
| 111 | continue |
| 112 | index_in_input_batch = req_id_index_in_input_batch[req.req_id] |
| 113 | output_token_ids[index_in_input_batch] = req.output_token_ids |
| 114 | prompt_token_ids[index_in_input_batch] = req.prompt_token_ids |
| 115 | presence_penalties[index_in_input_batch] = req.sampling_params.presence_penalty |
| 116 | frequency_penalties[index_in_input_batch] = ( |
| 117 | req.sampling_params.frequency_penalty |
| 118 | ) |
| 119 | repetition_penalties[index_in_input_batch] = ( |
| 120 | req.sampling_params.repetition_penalty |
| 121 | ) |
| 122 | top_k[index_in_input_batch] = req.sampling_params.top_k |
| 123 | top_p[index_in_input_batch] = req.sampling_params.top_p |
| 124 | temperature[index_in_input_batch] = req.sampling_params.temperature |
| 125 | min_tokens[index_in_input_batch] = ( |
| 126 | req.sampling_params.min_tokens, |
| 127 | req.sampling_params.all_stop_token_ids, |
| 128 | ) |
| 129 | logit_bias[index_in_input_batch] = req.sampling_params.logit_bias |
| 130 | if req.sampling_params.allowed_token_ids: |
| 131 | allowed_token_ids_mask[index_in_input_batch][ |
| 132 | req.sampling_params.allowed_token_ids |
| 133 | ] = True |
| 134 | if req.sampling_params.bad_words_token_ids: |
| 135 | bad_words_token_ids[index_in_input_batch] = ( |
| 136 | req.sampling_params.bad_words_token_ids |
| 137 | ) |
| 138 | |
| 139 | return SamplingMetadata( |
| 140 | temperature=torch.tensor(temperature, dtype=torch.float, device=device), |
| 141 | all_greedy=False, |
| 142 | all_random=True, |
| 143 | top_p=None |
| 144 | if all(x == 1.0 for x in top_p) |
| 145 | else torch.tensor(top_p, dtype=torch.float, device=device), |
| 146 | top_k=None |
| 147 | if all(x == 0 for x in top_k) |
| 148 | else torch.tensor(top_k, dtype=torch.int, device=device), |
| 149 | generators={}, |
| 150 | max_num_logprobs=0, |
| 151 | prompt_token_ids=make_tensor_with_pad( |
| 152 | prompt_token_ids, |
| 153 | pad=VOCAB_SIZE, |
| 154 | device=torch.device(device), |
| 155 | dtype=torch.int64, |
| 156 | ), |
| 157 | frequency_penalties=torch.tensor( |
| 158 | frequency_penalties, dtype=torch.float, device=device |
| 159 | ), |
| 160 | presence_penalties=torch.tensor( |
| 161 | presence_penalties, dtype=torch.float, device=device |
| 162 | ), |
| 163 | repetition_penalties=torch.tensor( |
| 164 | repetition_penalties, dtype=torch.float, device=device |
| 165 | ), |
| 166 | output_token_ids=output_token_ids, |
| 167 | spec_token_ids=[[] for _ in range(len(output_token_ids))], |
| 168 | no_penalties=( |
| 169 | all(x == 0 for x in presence_penalties) |
| 170 | and all(x == 0 for x in frequency_penalties) |
| 171 | and all(x == 1 for x in repetition_penalties) |
| 172 | ), |
| 173 | allowed_token_ids_mask=allowed_token_ids_mask, |
| 174 | bad_words_token_ids=bad_words_token_ids, |
| 175 | logitsprocs=LogitsProcessors(), |
| 176 | ) |
| 177 | |
| 178 | |
| 179 | def _create_sampling_params(): |
| 180 | return SamplingParams( |
| 181 | top_k=np.random.randint(1, 10), |
| 182 | top_p=np.random.uniform(0.0, 1.0), |
| 183 | presence_penalty=np.random.uniform(-2.0, 2.0), |
| 184 | repetition_penalty=np.random.uniform(0.0, 2.0), |
| 185 | frequency_penalty=np.random.uniform(-2.0, 2.0), |
| 186 | min_tokens=np.random.randint(1, 10), |
| 187 | stop_token_ids=[ |
| 188 | np.random.randint(0, VOCAB_SIZE) for _ in range(np.random.randint(10)) |
| 189 | ], |
| 190 | logit_bias={0: np.random.uniform(-3.0, 3.0)}, |
| 191 | ) |
| 192 | |
| 193 | |
| 194 | def _construct_cached_request_state(req_id_suffix: int): |
| 195 | prompt_token_ids = [ |
| 196 | np.random.randint(0, VOCAB_SIZE) |
| 197 | for _ in range(np.random.randint(0, MAX_PROMPT_SIZE)) |
| 198 | ] |
| 199 | output_token_ids = [ |
| 200 | np.random.randint(0, VOCAB_SIZE) |
| 201 | for _ in range(np.random.randint(0, NUM_OUTPUT_TOKENS)) |
| 202 | ] |
| 203 | return CachedRequestState( |
| 204 | req_id=f"req_id_{req_id_suffix}", |
| 205 | prompt_token_ids=prompt_token_ids, |
| 206 | sampling_params=_create_sampling_params(), |
| 207 | pooling_params=None, |
| 208 | mm_features=[], |
| 209 | block_ids=([],), |
| 210 | generator=None, |
| 211 | num_computed_tokens=len(output_token_ids), |
| 212 | output_token_ids=output_token_ids, |
| 213 | ) |
| 214 | |
| 215 | |
| 216 | @pytest.mark.parametrize("device", DEVICES) |
| 217 | @pytest.mark.parametrize("batch_size", [1, 2, 32, 64]) |
| 218 | def test_sampling_metadata_in_input_batch(device: str, batch_size: int): |
| 219 | """Tests the logic for managing sampling metadata in the InputBatch. |
| 220 | |
| 221 | This test involves adding a set of requests to the InputBatch, |
| 222 | followed by removing a subset of them. Afterward, the batch is compacted, |
| 223 | and the `make_sampling_metadata` method is invoked on the batch. The |
| 224 | output of `make_sampling_metadata` is then compared against the expected |
| 225 | results to ensure correctness. |
| 226 | |
| 227 | Note: Ignore logits processor logic, which is tested separately |
| 228 | """ |
| 229 | input_batch: InputBatch = InputBatch( |
| 230 | max_num_reqs=batch_size, |
| 231 | max_model_len=1024, |
| 232 | max_num_batched_tokens=1024, |
| 233 | device=torch.device(device), |
| 234 | vocab_size=1024, |
| 235 | block_sizes=[1], |
| 236 | kernel_block_sizes=[1], |
| 237 | max_num_blocks_per_req=[1024], |
| 238 | ) |
| 239 | reqs: list[CachedRequestState] = [] |
| 240 | req_id_reqs = {} |
| 241 | req_id_output_token_ids = {} |
| 242 | |
| 243 | # Add requests |
| 244 | for req_index in range(batch_size): |
| 245 | req: CachedRequestState = _construct_cached_request_state(req_index) |
| 246 | assigned_req_index = input_batch.add_request(req) |
| 247 | assert req_index == assigned_req_index |
| 248 | reqs.append(req) |
| 249 | req_id_reqs[req.req_id] = req |
| 250 | req_id_output_token_ids[req.req_id] = req.output_token_ids |
| 251 | |
| 252 | # Remove some requests |
| 253 | req_ids_to_remove = _remove_requests(input_batch, batch_size, reqs) |
| 254 | req_ids_retained = set(req_id_reqs.keys()) - req_ids_to_remove |
| 255 | |
| 256 | # Compact the input batch |
| 257 | input_batch.condense() |
| 258 | |
| 259 | # Generate the sampling metadata |
| 260 | sampling_metadata = input_batch._make_sampling_metadata() |
| 261 | |
| 262 | # Create expected output. |
| 263 | expected_sampling_metadata = _construct_expected_sampling_metadata( |
| 264 | reqs, req_ids_retained, input_batch.req_id_to_index, device=torch.device(device) |
| 265 | ) |
| 266 | |
| 267 | def same(t1: torch.Tensor | None, t2: torch.Tensor | None) -> bool: |
| 268 | return (t1 is None and t2 is None) or ( |
| 269 | t1 is not None and t2 is not None and torch.allclose(t1, t2) |
| 270 | ) |
| 271 | |
| 272 | # Assert the actual and expected output. |
| 273 | assert torch.allclose( |
| 274 | expected_sampling_metadata.temperature, sampling_metadata.temperature |
| 275 | ) |
| 276 | assert same(expected_sampling_metadata.top_p, sampling_metadata.top_p) |
| 277 | assert same(expected_sampling_metadata.top_k, sampling_metadata.top_k) |
| 278 | assert torch.allclose( |
| 279 | expected_sampling_metadata.frequency_penalties, |
| 280 | sampling_metadata.frequency_penalties, |
| 281 | ) |
| 282 | assert torch.allclose( |
| 283 | expected_sampling_metadata.presence_penalties, |
| 284 | sampling_metadata.presence_penalties, |
| 285 | ) |
| 286 | assert torch.allclose( |
| 287 | expected_sampling_metadata.repetition_penalties, |
| 288 | sampling_metadata.repetition_penalties, |
| 289 | ) |
| 290 | assert torch.allclose( |
| 291 | expected_sampling_metadata.prompt_token_ids, sampling_metadata.prompt_token_ids |
| 292 | ) |
| 293 | assert ( |
| 294 | expected_sampling_metadata.output_token_ids |
| 295 | == sampling_metadata.output_token_ids |
| 296 | ) |
| 297 | assert expected_sampling_metadata.no_penalties == sampling_metadata.no_penalties |
| 298 | if sampling_metadata.allowed_token_ids_mask: |
| 299 | assert torch.allclose( |
| 300 | expected_sampling_metadata.allowed_token_ids_mask, |
| 301 | sampling_metadata.allowed_token_ids_mask, |
| 302 | ) |
| 303 | assert ( |
| 304 | expected_sampling_metadata.bad_words_token_ids |
| 305 | == sampling_metadata.bad_words_token_ids |
| 306 | ) |
| 307 | |
| 308 | |
| 309 | @pytest.mark.parametrize("device", DEVICES) |
| 310 | @pytest.mark.parametrize("batch_size", [32]) |
| 311 | @pytest.mark.parametrize("swap_list", [((0, 1),)]) |
| 312 | def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: list): |
| 313 | """Tests the logic for managing sampling metadata in the InputBatch. |
| 314 | |
| 315 | This test involves adding a set of requests to the InputBatch, |
| 316 | followed by removing a subset of them. Afterward, the batch is compacted, |
| 317 | and the `make_sampling_metadata` method is invoked on the batch. The |
| 318 | output of `make_sampling_metadata` is then compared against the expected |
| 319 | results to ensure correctness. |
| 320 | |
| 321 | Note: Ignore logits processor logic, which is tested separately |
| 322 | """ |
| 323 | input_batch: InputBatch = InputBatch( |
| 324 | max_num_reqs=batch_size, |
| 325 | max_model_len=1024, |
| 326 | max_num_batched_tokens=1024, |
| 327 | device=torch.device(device), |
| 328 | vocab_size=1024, |
| 329 | block_sizes=[1], |
| 330 | kernel_block_sizes=[1], |
| 331 | max_num_blocks_per_req=[1024], |
| 332 | ) |
| 333 | ref_input_batch: InputBatch = InputBatch( |
| 334 | max_num_reqs=batch_size, |
| 335 | max_model_len=1024, |
| 336 | max_num_batched_tokens=1024, |
| 337 | device=torch.device(device), |
| 338 | vocab_size=1024, |
| 339 | block_sizes=[1], |
| 340 | kernel_block_sizes=[1], |
| 341 | max_num_blocks_per_req=[1024], |
| 342 | ) |
| 343 | |
| 344 | reqs: list[CachedRequestState] = [] |
| 345 | req_id_reqs = {} |
| 346 | req_id_output_token_ids = {} |
| 347 | # Add requests |
| 348 | for req_index in range(batch_size): |
| 349 | req: CachedRequestState = _construct_cached_request_state(req_index) |
| 350 | assigned_req_index = input_batch.add_request(req) |
| 351 | assert assigned_req_index == req_index |
| 352 | reqs.append(req) |
| 353 | req_id_reqs[req.req_id] = req |
| 354 | req_id_output_token_ids[req.req_id] = req.output_token_ids |
| 355 | |
| 356 | reordered_reqs = reqs.copy() |
| 357 | for swap_pair in swap_list: |
| 358 | reordered_reqs[swap_pair[0]], reordered_reqs[swap_pair[1]] = ( |
| 359 | reordered_reqs[swap_pair[1]], |
| 360 | reordered_reqs[swap_pair[0]], |
| 361 | ) |
| 362 | input_batch.swap_states(swap_pair[0], swap_pair[1]) |
| 363 | |
| 364 | for req_index in range(batch_size): |
| 365 | req = reordered_reqs[req_index] |
| 366 | assigned_req_index = ref_input_batch.add_request(req) |
| 367 | assert assigned_req_index == req_index |
| 368 | |
| 369 | input_batch.refresh_metadata() |
| 370 | ref_input_batch.refresh_metadata() |
| 371 | |
| 372 | _compare_objs(input_batch, ref_input_batch) |
| 373 | |
| 374 | |
| 375 | @pytest.mark.parametrize("device", DEVICES) |
| 376 | def test_condense_clears_stale_allowed_token_ids_mask(device: str): |
| 377 | """condense() must clear the mask row a constrained request is moved out |
| 378 | of. Otherwise a later request that reuses that row without |
| 379 | allowed_token_ids inherits the stale whitelist. |
| 380 | See https://github.com/vllm-project/vllm/issues/43894. |
| 381 | """ |
| 382 | batch_size = 4 |
| 383 | allowed_token_id = 13 |
| 384 | input_batch = InputBatch( |
| 385 | max_num_reqs=batch_size, |
| 386 | max_model_len=1024, |
| 387 | max_num_batched_tokens=1024, |
| 388 | device=torch.device(device), |
| 389 | vocab_size=VOCAB_SIZE, |
| 390 | block_sizes=[1], |
| 391 | kernel_block_sizes=[1], |
| 392 | max_num_blocks_per_req=[1024], |
| 393 | ) |
| 394 | |
| 395 | def _make_req(suffix: int, allowed_token_ids=None) -> CachedRequestState: |
| 396 | return CachedRequestState( |
| 397 | req_id=f"req_id_{suffix}", |
| 398 | prompt_token_ids=[1, 2, 3], |
| 399 | sampling_params=SamplingParams(allowed_token_ids=allowed_token_ids), |
| 400 | pooling_params=None, |
| 401 | mm_features=[], |
| 402 | block_ids=([],), |
| 403 | generator=None, |
| 404 | num_computed_tokens=0, |
| 405 | output_token_ids=[], |
| 406 | ) |
| 407 | |
| 408 | # Only req_id_2 is constrained, and it lands on the highest row. |
| 409 | assert input_batch.add_request(_make_req(0)) == 0 |
| 410 | assert input_batch.add_request(_make_req(1)) == 1 |
| 411 | assert input_batch.add_request(_make_req(2, [allowed_token_id])) == 2 |
| 412 | |
| 413 | mask = input_batch.allowed_token_ids_mask_cpu_tensor |
| 414 | assert mask is not None |
| 415 | # The constrained row masks every token except the single allowed id. |
| 416 | assert not mask[2][allowed_token_id].item() |
| 417 | assert int(mask[2].sum().item()) == VOCAB_SIZE - 1 |
| 418 | |
| 419 | # Free row 0, then condense: req_id_2 slides from row 2 down into row 0. |
| 420 | input_batch.remove_request("req_id_0") |
| 421 | input_batch.condense() |
| 422 | assert input_batch.req_id_to_index["req_id_2"] == 0 |
| 423 | assert int(mask[2].sum().item()) == 0 |
| 424 | |
| 425 | # A new unrestricted request reuses row 2 and must stay unconstrained. |
| 426 | assert input_batch.add_request(_make_req(3)) == 2 |
| 427 | assert "req_id_3" not in input_batch.has_allowed_token_ids |
| 428 | assert int(mask[2].sum().item()) == 0 |
| 429 | |
| 430 | |
| 431 | @pytest.mark.parametrize("device", DEVICES) |
| 432 | def test_swap_states_preserves_allowed_token_ids_mask(device: str): |
| 433 | """swap_states() must exchange the two allowed_token_ids_mask rows. |
| 434 | |
| 435 | Swapping row views via a tuple assignment |
| 436 | (t[i1], t[i2] = t[i2], t[i1]) aliases them, so both rows collapse to the |
| 437 | second row's contents and one request silently loses its token |
| 438 | restriction. This is the sibling of the condense() leak fixed for |
| 439 | https://github.com/vllm-project/vllm/issues/43894. |
| 440 | """ |
| 441 | allowed_token_id = 13 |
| 442 | input_batch = InputBatch( |
| 443 | max_num_reqs=2, |
| 444 | max_model_len=1024, |
| 445 | max_num_batched_tokens=1024, |
| 446 | device=torch.device(device), |
| 447 | vocab_size=VOCAB_SIZE, |
| 448 | block_sizes=[1], |
| 449 | kernel_block_sizes=[1], |
| 450 | max_num_blocks_per_req=[1024], |
| 451 | ) |
| 452 | |
| 453 | def _make_req(suffix: int, allowed_token_ids=None) -> CachedRequestState: |
| 454 | return CachedRequestState( |
| 455 | req_id=f"req_id_{suffix}", |
| 456 | prompt_token_ids=[1, 2, 3], |
| 457 | sampling_params=SamplingParams(allowed_token_ids=allowed_token_ids), |
| 458 | pooling_params=None, |
| 459 | mm_features=[], |
| 460 | block_ids=([],), |
| 461 | generator=None, |
| 462 | num_computed_tokens=0, |
| 463 | output_token_ids=[], |
| 464 | ) |
| 465 | |
| 466 | # Row 0 is constrained to a single token; row 1 is unconstrained. |
| 467 | assert input_batch.add_request(_make_req(0, [allowed_token_id])) == 0 |
| 468 | assert input_batch.add_request(_make_req(1)) == 1 |
| 469 | |
| 470 | mask = input_batch.allowed_token_ids_mask_cpu_tensor |
| 471 | assert mask is not None |
| 472 | assert int(mask[0].sum().item()) == VOCAB_SIZE - 1 |
| 473 | assert not mask[0][allowed_token_id].item() |
| 474 | assert int(mask[1].sum().item()) == 0 |
| 475 | |
| 476 | input_batch.swap_states(0, 1) |
| 477 | |
| 478 | # The constrained request now sits at row 1 and must keep its mask; the |
| 479 | # unconstrained request now sits at row 0 and must stay open. |
| 480 | assert int(mask[1].sum().item()) == VOCAB_SIZE - 1 |
| 481 | assert not mask[1][allowed_token_id].item() |
| 482 | assert int(mask[0].sum().item()) == 0 |
| 483 | |
| 484 | |
| 485 | def _construct_pooling_request(req_id_suffix: int, pooling_params=None): |
| 486 | from vllm.pooling_params import PoolingParams |
| 487 | |
| 488 | prompt_token_ids = [ |
| 489 | np.random.randint(0, VOCAB_SIZE) |
| 490 | for _ in range(np.random.randint(10, MAX_PROMPT_SIZE)) |
| 491 | ] |
| 492 | return CachedRequestState( |
| 493 | req_id=f"pool_req_{req_id_suffix}", |
| 494 | prompt_token_ids=prompt_token_ids, |
| 495 | sampling_params=None, |
| 496 | pooling_params=pooling_params or PoolingParams(task="classify"), |
| 497 | mm_features=[], |
| 498 | block_ids=([],), |
| 499 | generator=None, |
| 500 | num_computed_tokens=0, |
| 501 | output_token_ids=[], |
| 502 | ) |
| 503 | |
| 504 | |
| 505 | @pytest.mark.parametrize("device", DEVICES) |
| 506 | def test_pooling_prompt_lens_not_aliased(device: str): |
| 507 | """Verify that prompt_lens in PoolingMetadata does not share memory |
| 508 | with the internal num_prompt_tokens pinned buffer. Guards against possible |
| 509 | non-determinism in pooling metadata due to mutations to the internal buffer. |
| 510 | """ |
| 511 | batch_size = 4 |
| 512 | input_batch = InputBatch( |
| 513 | max_num_reqs=batch_size * 2, |
| 514 | max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS, |
| 515 | max_num_batched_tokens=batch_size * (MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS), |
| 516 | device=torch.device(device), |
| 517 | vocab_size=VOCAB_SIZE, |
| 518 | block_sizes=[16], |
| 519 | kernel_block_sizes=[16], |
| 520 | max_num_blocks_per_req=[64], |
| 521 | is_pooling_model=True, |
| 522 | ) |
| 523 | |
| 524 | reqs = [] |
| 525 | # Add requests |
| 526 | for i in range(batch_size): |
| 527 | req = _construct_pooling_request(i) |
| 528 | input_batch.add_request(req) |
| 529 | reqs.append(req) |
| 530 | input_batch.refresh_metadata() |
| 531 | |
| 532 | # prompt_lens must be a snapshot |
| 533 | metadata = input_batch.get_pooling_metadata() |
| 534 | prompt_lens_snapshot = metadata.prompt_lens.clone() |
| 535 | |
| 536 | # Mutate the internal buffer (simulates next batch adding new requests) |
| 537 | input_batch.num_prompt_tokens_cpu_tensor.fill_(999) |
| 538 | |
| 539 | # prompt_lens must be unaffected by the mutation |
| 540 | assert torch.equal(metadata.prompt_lens, prompt_lens_snapshot), ( |
| 541 | "prompt_lens shares memory with internal pinned buffer; " |
| 542 | "mutations to num_prompt_tokens_cpu_tensor corrupted prompt_lens. " |
| 543 | f"Expected {prompt_lens_snapshot}, got {metadata.prompt_lens}" |
| 544 | ) |
| 545 | |
| 546 | |
| 547 | def test_placeholder_spec_token_ids_written_verbatim(): |
| 548 | input_batch = InputBatch( |
| 549 | max_num_reqs=1, |
| 550 | max_model_len=8, |
| 551 | max_num_batched_tokens=8, |
| 552 | device=torch.device("cpu"), |
| 553 | vocab_size=VOCAB_SIZE, |
| 554 | block_sizes=[16], |
| 555 | kernel_block_sizes=[16], |
| 556 | max_num_blocks_per_req=[1], |
| 557 | ) |
| 558 | req = CachedRequestState( |
| 559 | req_id="req", |
| 560 | prompt_token_ids=[10, 11], |
| 561 | mm_features=[], |
| 562 | sampling_params=SamplingParams(), |
| 563 | block_ids=([],), |
| 564 | generator=None, |
| 565 | num_computed_tokens=3, |
| 566 | output_token_ids=[12], |
| 567 | ) |
| 568 | input_batch.add_request(req) |
| 569 | |
| 570 | input_batch.update_req_spec_token_ids( |
| 571 | req, |
| 572 | {"req": [13, -1, -1]}, |
| 573 | ) |
| 574 | |
| 575 | # Placeholders (-1) are kept verbatim in both the spec_token_ids list and |
| 576 | # the token buffer; they are clamped to 0 only at the embedding boundary |
| 577 | # (GPUModelRunner._preprocess). |
| 578 | assert input_batch.spec_token_ids[0] == [13, -1, -1] |
| 579 | assert input_batch.token_ids_cpu[0, 3:6].tolist() == [13, -1, -1] |
| 580 | |
| 581 | |
| 582 | @pytest.mark.parametrize( |
| 583 | ("pooling_params", "expect_cpu_prompt_token_ids"), |
| 584 | [ |
| 585 | ({"task": "classify"}, False), |
| 586 | ({"task": "classify", "requires_token_ids": True}, True), |
| 587 | ], |
| 588 | ) |
| 589 | def test_pooling_metadata_token_id_buffers( |
| 590 | pooling_params: dict[str, object], |
| 591 | expect_cpu_prompt_token_ids: bool, |
| 592 | ): |
| 593 | from vllm.pooling_params import PoolingParams |
| 594 | |
| 595 | input_batch = InputBatch( |
| 596 | max_num_reqs=1, |
| 597 | max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS, |
| 598 | max_num_batched_tokens=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS, |
| 599 | device=torch.device("cpu"), |
| 600 | vocab_size=VOCAB_SIZE, |
| 601 | block_sizes=[16], |
| 602 | kernel_block_sizes=[16], |
| 603 | max_num_blocks_per_req=[64], |
| 604 | is_pooling_model=True, |
| 605 | ) |
| 606 | req = _construct_pooling_request(0, PoolingParams(**pooling_params)) |
| 607 | input_batch.add_request(req) |
| 608 | input_batch.refresh_metadata() |
| 609 | |
| 610 | metadata = input_batch.get_pooling_metadata() |
| 611 | assert input_batch.sampling_metadata.prompt_token_ids is None |
| 612 | assert metadata.prompt_token_ids is None |
| 613 | |
| 614 | if expect_cpu_prompt_token_ids: |
| 615 | assert metadata.prompt_token_ids_cpu is not None |
| 616 | assert metadata.get_prompt_token_ids_cpu()[0].tolist() == req.prompt_token_ids |
| 617 | else: |
| 618 | assert metadata.prompt_token_ids_cpu is None |
| 619 | |
| 620 | |
| 621 | PROMPT_EMBEDS_HIDDEN = 8 |
| 622 | PROMPT_EMBEDS_LEN = 4 |
| 623 | |
| 624 | |
| 625 | def _make_input_batch(is_pooling_model: bool = False) -> InputBatch: |
| 626 | return InputBatch( |
| 627 | max_num_reqs=2, |
| 628 | max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS, |
| 629 | max_num_batched_tokens=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS, |
| 630 | device=torch.device("cpu"), |
| 631 | vocab_size=VOCAB_SIZE, |
| 632 | block_sizes=[16], |
| 633 | kernel_block_sizes=[16], |
| 634 | max_num_blocks_per_req=[64], |
| 635 | is_pooling_model=is_pooling_model, |
| 636 | ) |
| 637 | |
| 638 | |
| 639 | def _make_embeds_request(req_id: str, pooling: bool = False) -> CachedRequestState: |
| 640 | from vllm.pooling_params import PoolingParams |
| 641 | |
| 642 | return CachedRequestState( |
| 643 | req_id=req_id, |
| 644 | prompt_token_ids=None, |
| 645 | prompt_embeds=torch.randn(PROMPT_EMBEDS_LEN, PROMPT_EMBEDS_HIDDEN), |
| 646 | mm_features=[], |
| 647 | sampling_params=None if pooling else SamplingParams(), |
| 648 | pooling_params=PoolingParams(task="classify") if pooling else None, |
| 649 | block_ids=([],), |
| 650 | generator=None, |
| 651 | num_computed_tokens=0, |
| 652 | output_token_ids=[], |
| 653 | ) |
| 654 | |
| 655 | |
| 656 | def _make_token_request(req_id: str) -> CachedRequestState: |
| 657 | return CachedRequestState( |
| 658 | req_id=req_id, |
| 659 | prompt_token_ids=[10, 11, 12, 13], |
| 660 | mm_features=[], |
| 661 | sampling_params=SamplingParams(), |
| 662 | block_ids=([],), |
| 663 | generator=None, |
| 664 | num_computed_tokens=0, |
| 665 | output_token_ids=[], |
| 666 | ) |
| 667 | |
| 668 | |
| 669 | def test_remove_request_releases_prompt_embeds(): |
| 670 | """A finished prompt-embeds request must not keep its tensor referenced by |
| 671 | the persistent batch; an empty batch owns no prompt embeds.""" |
| 672 | input_batch = _make_input_batch() |
| 673 | req = _make_embeds_request("embeds-req") |
| 674 | input_batch.add_request(req) |
| 675 | req_index = input_batch.req_id_to_index[req.req_id] |
| 676 | assert req_index in input_batch.req_prompt_embeds |
| 677 | |
| 678 | input_batch.remove_request(req.req_id) |
| 679 | input_batch.condense() |
| 680 | |
| 681 | assert req_index not in input_batch.req_prompt_embeds |
| 682 | assert not input_batch.req_prompt_embeds |
| 683 | |
| 684 | |
| 685 | def test_reused_slot_does_not_retain_prompt_embeds(): |
| 686 | """Reusing a vacated slot with a token-id request must not keep the previous |
| 687 | occupant's prompt-embeds tensor alive.""" |
| 688 | input_batch = _make_input_batch() |
| 689 | embeds_req = _make_embeds_request("embeds-req") |
| 690 | input_batch.add_request(embeds_req) |
| 691 | slot = input_batch.req_id_to_index[embeds_req.req_id] |
| 692 | |
| 693 | input_batch.remove_request(embeds_req.req_id) |
| 694 | input_batch.condense() |
| 695 | |
| 696 | token_req = _make_token_request("token-req") |
| 697 | input_batch.add_request(token_req) |
| 698 | assert input_batch.req_id_to_index[token_req.req_id] == slot |
| 699 | assert slot not in input_batch.req_prompt_embeds |
| 700 | |
| 701 | |
| 702 | def test_pooling_model_releases_prompt_embeds(): |
| 703 | """The cleanup must run before remove_request's pooling-model early return.""" |
| 704 | input_batch = _make_input_batch(is_pooling_model=True) |
| 705 | req = _make_embeds_request("pool-embeds-req", pooling=True) |
| 706 | input_batch.add_request(req) |
| 707 | req_index = input_batch.req_id_to_index[req.req_id] |
| 708 | assert req_index in input_batch.req_prompt_embeds |
| 709 | |
| 710 | input_batch.remove_request(req.req_id) |
| 711 | input_batch.condense() |
| 712 | |
| 713 | assert req_index not in input_batch.req_prompt_embeds |