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 8, 00:00 UTC
Support prompt_embeds for pooling requests in output processor (#34904)
Signed-off-by: Li Zhang <lzhanga@amazon.com> Co-authored-by: Li Zhang <lzhanga@amazon.com>
Li committed
Feb 20, 2026 at 22:57 UTC
59c62332978fcce318784df499713764f14c7bc1
1 file changed
+7
-8
vllm/v1/engine/output_processor.py
+7
-8
@@ -337,16 +337,20 @@ class RequestState:
337
finished: bool,
338
kv_transfer_params: dict[str, Any] | None = None,
339
) -> RequestOutput | PoolingRequestOutput:
340
+ # If prompt embeds were used, put placeholder prompt token ids
341
+ prompt_token_ids = self.prompt_token_ids
342
+ if prompt_token_ids is None and self.prompt_embeds is not None:
343
+ prompt_token_ids = [0] * len(self.prompt_embeds)
344
+ assert prompt_token_ids is not None
345
+
346
first_output = outputs[0]
347
if isinstance(first_output, PoolingOutput):
348
assert len(outputs) == 1
343
- # Prompt embeddings are currently not supported by pooling requests.
344
- assert self.prompt_token_ids is not None
349
return PoolingRequestOutput(
350
request_id=external_req_id,
351
outputs=first_output,
352
num_cached_tokens=self.num_cached_tokens,
349
- prompt_token_ids=self.prompt_token_ids,
353
+ prompt_token_ids=prompt_token_ids,
354
finished=finished,
355
)
356
assert self.logprobs_processor is not None
@@ -356,11 +360,6 @@ class RequestState:
360
else:
361
prompt_logprobs = self.logprobs_processor.prompt_logprobs
362
359
- # If prompt embeds were used, put placeholder prompt token ids
360
- prompt_token_ids = self.prompt_token_ids
361
- if prompt_token_ids is None and self.prompt_embeds is not None:
362
- prompt_token_ids = [0] * len(self.prompt_embeds)
363
-
363
return RequestOutput(
364
request_id=external_req_id, # request_id is what was provided externally
365
lora_request=self.lora_request,