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 6, 05:45 UTC
[Bugfix] Fix pooling non-determinism from pinned prompt_lens aliasing (#37775)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Mar 21, 2026 at 22:22 UTC
66f927f205fd86a933642ef46b84997dc69bb668
2 files changed
+63
-1
tests/v1/worker/test_gpu_input_batch.py
+62
@@ -378,3 +378,65 @@ def test_swap_states_in_input_batch(device: str, batch_size: int, swap_list: lis
378
ref_input_batch.refresh_metadata()
379
380
_compare_objs(input_batch, ref_input_batch)
381
+
382
+
383
+def _construct_pooling_request(req_id_suffix: int):
384
+ from vllm.pooling_params import PoolingParams
385
+
386
+ prompt_token_ids = [
387
+ np.random.randint(0, VOCAB_SIZE)
388
+ for _ in range(np.random.randint(10, MAX_PROMPT_SIZE))
389
+ ]
390
+ return CachedRequestState(
391
+ req_id=f"pool_req_{req_id_suffix}",
392
+ prompt_token_ids=prompt_token_ids,
393
+ sampling_params=None,
394
+ pooling_params=PoolingParams(task="classify"),
395
+ mm_features=[],
396
+ block_ids=([],),
397
+ generator=None,
398
+ num_computed_tokens=0,
399
+ output_token_ids=[],
400
+ )
401
+
402
+
403
+@pytest.mark.parametrize("device", CUDA_DEVICES)
404
+def test_pooling_prompt_lens_not_aliased(device: str):
405
+ """Verify that prompt_lens in PoolingMetadata does not share memory
406
+ with the internal num_prompt_tokens pinned buffer. Guards against possible
407
+ non-determinism in pooling metadata due to mutations to the internal buffer.
408
+ """
409
+ batch_size = 4
410
+ input_batch = InputBatch(
411
+ max_num_reqs=batch_size * 2,
412
+ max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
413
+ max_num_batched_tokens=batch_size * (MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS),
414
+ device=torch.device(device),
415
+ pin_memory=is_pin_memory_available(),
416
+ vocab_size=VOCAB_SIZE,
417
+ block_sizes=[16],
418
+ kernel_block_sizes=[16],
419
+ is_pooling_model=True,
420
+ )
421
+
422
+ reqs = []
423
+ # Add requests
424
+ for i in range(batch_size):
425
+ req = _construct_pooling_request(i)
426
+ input_batch.add_request(req)
427
+ reqs.append(req)
428
+ input_batch.refresh_metadata()
429
+
430
+ # prompt_lens must be a snapshot
431
+ metadata = input_batch.get_pooling_metadata()
432
+ prompt_lens_snapshot = metadata.prompt_lens.clone()
433
+
434
+ # Mutate the internal buffer (simulates next batch adding new requests)
435
+ input_batch.num_prompt_tokens_cpu_tensor.fill_(999)
436
+
437
+ # prompt_lens must be unaffected by the mutation
438
+ assert torch.equal(metadata.prompt_lens, prompt_lens_snapshot), (
439
+ "prompt_lens shares memory with internal pinned buffer; "
440
+ "mutations to num_prompt_tokens_cpu_tensor corrupted prompt_lens. "
441
+ f"Expected {prompt_lens_snapshot}, got {metadata.prompt_lens}"
442
+ )
vllm/v1/worker/gpu_input_batch.py
+1
-1
@@ -892,7 +892,7 @@ class InputBatch:
892
pooling_states = self.get_pooling_states()
893
894
return PoolingMetadata(
895
- prompt_lens=torch.from_numpy(self.num_prompt_tokens[: self.num_reqs]),
895
+ prompt_lens=self.num_prompt_tokens_cpu_tensor[: self.num_reqs].clone(),
896
prompt_token_ids=self.sampling_metadata.prompt_token_ids,
897
pooling_params=pooling_params,
898
pooling_states=pooling_states,