[Pooling] Support prompt embeddings in MRV2 decoder pooling (#56639)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Sep 13, 2026 at 00:49 UTC
0cf266a469bd6f06e90fa267d466bbfe5b11dce5
2 files changed
+2
-2
vllm/v1/worker/gpu/model_runner.py
-1
@@ -1128,7 +1128,6 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1128
1129
if self.pooling_runner is not None:
1130
assert new_req_data.pooling_params is not None
1131
- assert new_req_data.prompt_token_ids is not None
1131
self.pooling_runner.add_request(
1132
req_id,
1133
req_index,
vllm/v1/worker/gpu/pool/pooling_runner.py
+2
-1
@@ -64,7 +64,7 @@ class PoolingRunner:
64
req_id: str,
65
req_index: int,
66
pooling_params: PoolingParams,
67
- prompt_token_ids: list[int],
67
+ prompt_token_ids: list[int] | None,
68
) -> None:
69
task = pooling_params.task
70
if task not in self.supported_tasks:
@@ -77,6 +77,7 @@ class PoolingRunner:
77
self.pooling_params[req_index] = pooling_params
78
self.pooling_states[req_index] = PoolingStates()
79
if pooling_params.requires_token_ids:
80
+ assert prompt_token_ids is not None
81
self.prompt_token_ids[req_index] = torch.tensor(
82
prompt_token_ids, dtype=torch.int64
83
)