Release prompt_embeds tensor when its InputBatch slot is freed (#57988)
Signed-off-by: khushali9 <khushali.desai9@gmail.com>
Khushali Desai committed
Sep 24, 2026 at 17:00 UTC
7833ff25dd6fadd5cdc4ee46ab67af64fbc30a94
2 files changed
+96
tests/v1/worker/test_gpu_input_batch.py
+95
@@ -506,3 +506,98 @@ def test_pooling_metadata_token_id_buffers(
506
assert metadata.get_prompt_token_ids_cpu()[0].tolist() == req.prompt_token_ids
507
else:
508
assert metadata.prompt_token_ids_cpu is None
509
+
510
+
511
+PROMPT_EMBEDS_HIDDEN = 8
512
+PROMPT_EMBEDS_LEN = 4
513
+
514
+
515
+def _make_input_batch(is_pooling_model: bool = False) -> InputBatch:
516
+ return InputBatch(
517
+ max_num_reqs=2,
518
+ max_model_len=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
519
+ max_num_batched_tokens=MAX_PROMPT_SIZE + NUM_OUTPUT_TOKENS,
520
+ device=torch.device("cpu"),
521
+ vocab_size=VOCAB_SIZE,
522
+ block_sizes=[16],
523
+ kernel_block_sizes=[16],
524
+ max_num_blocks_per_req=[64],
525
+ is_pooling_model=is_pooling_model,
526
+ )
527
+
528
+
529
+def _make_embeds_request(req_id: str, pooling: bool = False) -> CachedRequestState:
530
+ from vllm.pooling_params import PoolingParams
531
+
532
+ return CachedRequestState(
533
+ req_id=req_id,
534
+ prompt_token_ids=None,
535
+ prompt_embeds=torch.randn(PROMPT_EMBEDS_LEN, PROMPT_EMBEDS_HIDDEN),
536
+ mm_features=[],
537
+ sampling_params=None if pooling else SamplingParams(),
538
+ pooling_params=PoolingParams(task="classify") if pooling else None,
539
+ block_ids=([],),
540
+ generator=None,
541
+ num_computed_tokens=0,
542
+ output_token_ids=[],
543
+ )
544
+
545
+
546
+def _make_token_request(req_id: str) -> CachedRequestState:
547
+ return CachedRequestState(
548
+ req_id=req_id,
549
+ prompt_token_ids=[10, 11, 12, 13],
550
+ mm_features=[],
551
+ sampling_params=SamplingParams(),
552
+ block_ids=([],),
553
+ generator=None,
554
+ num_computed_tokens=0,
555
+ output_token_ids=[],
556
+ )
557
+
558
+
559
+def test_remove_request_releases_prompt_embeds():
560
+ """A finished prompt-embeds request must not keep its tensor referenced by
561
+ the persistent batch; an empty batch owns no prompt embeds."""
562
+ input_batch = _make_input_batch()
563
+ req = _make_embeds_request("embeds-req")
564
+ input_batch.add_request(req)
565
+ req_index = input_batch.req_id_to_index[req.req_id]
566
+ assert req_index in input_batch.req_prompt_embeds
567
+
568
+ input_batch.remove_request(req.req_id)
569
+ input_batch.condense()
570
+
571
+ assert req_index not in input_batch.req_prompt_embeds
572
+ assert not input_batch.req_prompt_embeds
573
+
574
+
575
+def test_reused_slot_does_not_retain_prompt_embeds():
576
+ """Reusing a vacated slot with a token-id request must not keep the previous
577
+ occupant's prompt-embeds tensor alive."""
578
+ input_batch = _make_input_batch()
579
+ embeds_req = _make_embeds_request("embeds-req")
580
+ input_batch.add_request(embeds_req)
581
+ slot = input_batch.req_id_to_index[embeds_req.req_id]
582
+
583
+ input_batch.remove_request(embeds_req.req_id)
584
+ input_batch.condense()
585
+
586
+ token_req = _make_token_request("token-req")
587
+ input_batch.add_request(token_req)
588
+ assert input_batch.req_id_to_index[token_req.req_id] == slot
589
+ assert slot not in input_batch.req_prompt_embeds
590
+
591
+
592
+def test_pooling_model_releases_prompt_embeds():
593
+ """The cleanup must run before remove_request's pooling-model early return."""
594
+ input_batch = _make_input_batch(is_pooling_model=True)
595
+ req = _make_embeds_request("pool-embeds-req", pooling=True)
596
+ input_batch.add_request(req)
597
+ req_index = input_batch.req_id_to_index[req.req_id]
598
+ assert req_index in input_batch.req_prompt_embeds
599
+
600
+ input_batch.remove_request(req.req_id)
601
+ input_batch.condense()
602
+
603
+ assert req_index not in input_batch.req_prompt_embeds
vllm/v1/worker/gpu_input_batch.py
+1
@@ -541,6 +541,7 @@ class InputBatch:
541
self.req_output_token_ids[req_index] = None
542
self.spec_token_ids[req_index].clear()
543
self.block_table.clear_row(req_index)
544
+ self.req_prompt_embeds.pop(req_index, None)
545
546
# LoRA
547
lora_id = self.request_lora_mapping[req_index]