Feature/vllm/input embedding completion api

tests/entrypoints/openai/test_completion.py

@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: Apache-2.0
 
 # imports for guided decoding tests
+import base64
+import io
 import json
 import re
 import shutil
@@ -11,10 +13,11 @@
 import openai  # use the official client for correctness check
 import pytest
 import pytest_asyncio
+import torch
 # downloading lora to test lora requests
 from huggingface_hub import snapshot_download
 from openai import BadRequestError
-from transformers import AutoTokenizer
+from transformers import AutoConfig, AutoTokenizer
 
 from vllm.transformers_utils.tokenizer import get_tokenizer
 
@@ -31,6 +34,7 @@
 PA_NUM_VIRTUAL_TOKENS = 8
 
 GUIDED_DECODING_BACKENDS = ["outlines", "lm-format-enforcer", "xgrammar"]
+CONFIG = AutoConfig.from_pretrained(MODEL_NAME)
 
 
 @pytest.fixture(scope="module")
@@ -107,6 +111,14 @@ async def client(server):
         yield async_client
 
 
+def create_dummy_embeds(num_tokens: int = 5) -> str:
+    """Create dummy embeddings and return them as base64 encoded string."""
+    dummy_embeds = torch.randn(num_tokens, CONFIG.hidden_size)
+    buffer = io.BytesIO()
+    torch.save(dummy_embeds, buffer)
+    return base64.b64encode(buffer.getvalue()).decode('utf-8')
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
     # first test base model, then test loras, then test prompt adapters
@@ -143,6 +155,45 @@ async def test_single_completion(client: openai.AsyncOpenAI, model_name: str,
     assert len(completion.choices[0].text) >= 1
     assert completion.choices[0].prompt_logprobs is None
 
+    # test using prompt_embeds
+    encoded_embeds = create_dummy_embeds()
+    completion = await client.completions.create(
+        model=model_name,
+        prompt="",  # Add empty prompt as required parameter
+        max_tokens=5,
+        temperature=0.0,
+        extra_body={"prompt_embeds": encoded_embeds})
+    assert len(completion.choices[0].text) >= 1
+    assert completion.choices[0].prompt_logprobs is None
+
+    # test batch completion with prompt_embeds
+    encoded_embeds2 = create_dummy_embeds()
+    completion = await client.completions.create(
+        model=model_name,
+        prompt="",  # Add empty prompt as required parameter
+        max_tokens=5,
+        temperature=0.0,
+        extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]})
+    assert len(completion.choices) == 2
+    assert len(completion.choices[0].text) >= 1
+    assert len(completion.choices[1].text) >= 1
+
+    # test error case: neither prompt nor prompt_embeds provided
+    with pytest.raises(BadRequestError):
+        await client.completions.create(
+            model=model_name,
+            max_tokens=5,
+            temperature=0.0,
+        )
+
+    # test error case: invalid prompt_embeds
+    with pytest.raises(BadRequestError):
+        await client.completions.create(
+            model=model_name,
+            max_tokens=5,
+            temperature=0.0,
+            extra_body={"prompt_embeds": "invalid_base64"})
+
 
 @pytest.mark.asyncio
 async def test_added_lora_tokens(client: openai.AsyncOpenAI):
@@ -343,6 +394,55 @@ async def test_completion_streaming(client: openai.AsyncOpenAI,
     assert chunk.choices[0].text
     assert "".join(chunks) == single_output
 
+    # test streaming with prompt_embeds
+    encoded_embeds = create_dummy_embeds()
+    single_completion = await client.completions.create(
+        model=model_name,
+        prompt="",  # Add empty prompt as required parameter
+        max_tokens=5,
+        temperature=0.0,
+        extra_body={"prompt_embeds": encoded_embeds})
+    single_output = single_completion.choices[0].text
+
+    stream = await client.completions.create(
+        model=model_name,
+        prompt="",  # Add empty prompt as required parameter
+        max_tokens=5,
+        temperature=0.0,
+        stream=True,
+        extra_body={"prompt_embeds": encoded_embeds})
+    chunks = []
+    finish_reason_count = 0
+    async for chunk in stream:
+        chunks.append(chunk.choices[0].text)
+        if chunk.choices[0].finish_reason is not None:
+            finish_reason_count += 1
+    assert finish_reason_count == 1
+    assert chunk.choices[0].finish_reason == "length"
+    assert chunk.choices[0].text
+    assert "".join(chunks) == single_output
+
+    # test batch streaming with prompt_embeds
+    encoded_embeds2 = create_dummy_embeds()
+    stream = await client.completions.create(
+        model=model_name,
+        prompt="",  # Add empty prompt as required parameter
+        max_tokens=5,
+        temperature=0.0,
+        stream=True,
+        extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]})
+    chunks = [[], []]
+    finish_reason_count = 0
+    async for chunk in stream:
+        chunks[chunk.choices[0].index].append(chunk.choices[0].text)
+        if chunk.choices[0].finish_reason is not None:
+            finish_reason_count += 1
+    assert finish_reason_count == 2
+    assert chunk.choices[0].finish_reason == "length"
+    assert chunk.choices[0].text
+    assert len(chunks[0]) > 0
+    assert len(chunks[1]) > 0
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize(
@@ -760,6 +860,53 @@ async def test_echo_logprob_completion(client: openai.AsyncOpenAI,
                        1) <= len(top_logprobs) <= logprobs_arg + 1
         assert len(logprobs.tokens) > 5
 
+    # test using prompt_embeds
+    encoded_embeds = create_dummy_embeds()
+    completion = await client.completions.create(
+        model=model_name,
+        prompt="",  # Add empty prompt as required parameter
+        max_tokens=5,
+        temperature=0.0,
+        echo=True,
+        logprobs=logprobs_arg,
+        extra_body={"prompt_embeds": encoded_embeds})
+
+    logprobs = completion.choices[0].logprobs
+    assert logprobs is not None
+    assert len(logprobs.text_offset) > 5
+    assert (len(logprobs.token_logprobs) > 5
+            and logprobs.token_logprobs[0] is None)
+    assert (len(logprobs.top_logprobs) > 5
+            and logprobs.top_logprobs[0] is None)
+    for top_logprobs in logprobs.top_logprobs[1:]:
+        assert max(logprobs_arg, 1) <= len(top_logprobs) <= logprobs_arg + 1
+    assert len(logprobs.tokens) > 5
+
+    # test batch completion with prompt_embeds
+    encoded_embeds2 = create_dummy_embeds()
+    completion = await client.completions.create(
+        model=model_name,
+        prompt="",  # Add empty prompt as required parameter
+        max_tokens=5,
+        temperature=0.0,
+        echo=True,
+        logprobs=logprobs_arg,
+        extra_body={"prompt_embeds": [encoded_embeds, encoded_embeds2]})
+
+    assert len(completion.choices) == 2
+    for choice in completion.choices:
+        logprobs = choice.logprobs
+        assert logprobs is not None
+        assert len(logprobs.text_offset) > 5
+        assert (len(logprobs.token_logprobs) > 5
+                and logprobs.token_logprobs[0] is None)
+        assert (len(logprobs.top_logprobs) > 5
+                and logprobs.top_logprobs[0] is None)
+        for top_logprobs in logprobs.top_logprobs[1:]:
+            assert max(logprobs_arg,
+                       1) <= len(top_logprobs) <= logprobs_arg + 1
+        assert len(logprobs.tokens) > 5
+
 
 @pytest.mark.asyncio
 @pytest.mark.parametrize("guided_decoding_backend", GUIDED_DECODING_BACKENDS)

vllm/entrypoints/logger.py

@@ -2,6 +2,8 @@
 
 from typing import Optional, Union
 
+import torch
+
 from vllm.logger import init_logger
 from vllm.lora.request import LoRARequest
 from vllm.pooling_params import PoolingParams
@@ -23,6 +25,7 @@ def log_inputs(
         request_id: str,
         prompt: Optional[str],
         prompt_token_ids: Optional[list[int]],
+        prompt_embeds: Optional[torch.Tensor],
         params: Optional[Union[SamplingParams, PoolingParams,
                                BeamSearchParams]],
         lora_request: Optional[LoRARequest],
@@ -39,6 +42,8 @@ def log_inputs(
         logger.info(
             "Received request %s: prompt: %r, "
             "params: %s, prompt_token_ids: %s, "
+            "prompt_embeds shape: %s, "
             "lora_request: %s, prompt_adapter_request: %s.", request_id,
-            prompt, params, prompt_token_ids, lora_request,
-            prompt_adapter_request)
+            prompt, params, prompt_token_ids,
+            prompt_embeds.shape if prompt_embeds is not None else None,
+            lora_request, prompt_adapter_request)

vllm/entrypoints/openai/protocol.py

@@ -754,8 +754,9 @@ def check_cache_salt_support(cls, data):
 class CompletionRequest(OpenAIBaseModel):
     # Ordered by official OpenAI API documentation
     # https://platform.openai.com/docs/api-reference/completions/create
-    model: Optional[str] = None
-    prompt: Union[list[int], list[list[int]], str, list[str]]
+    model: str
+    prompt: Optional[Union[list[int], list[list[int]], str, list[str]]] = None
+    prompt_embeds: Optional[Union[bytes, list[bytes]]] = None
     best_of: Optional[int] = None
     echo: Optional[bool] = False
     frequency_penalty: Optional[float] = 0.0
@@ -1029,6 +1030,14 @@ def validate_stream_options(cls, data):
 
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def validate_prompt_and_prompt_embeds(cls, data):
+        if data.get("prompt") is None and data.get("prompt_embeds") is None:
+            raise ValueError(
+                "At least one of `prompt` or `prompt_embeds` must be set.")
+        return data
+
 
 class EmbeddingCompletionRequest(OpenAIBaseModel):
     # Ordered by official OpenAI API documentation

vllm/entrypoints/openai/serving_completion.py

@@ -108,7 +108,7 @@
             request_prompts, engine_prompts = await self._preprocess_completion(
                 request,
                 tokenizer,
                 request.prompt,
                 truncate_prompt_tokens=request.truncate_prompt_tokens,
                 add_special_tokens=request.add_special_tokens,
             )
@@ -131,7 +131,9 @@
             for i, engine_prompt in enumerate(engine_prompts):
                 sampling_params: Union[SamplingParams, BeamSearchParams]
                 default_max_tokens = self.max_model_len - len(
-                    engine_prompt["prompt_token_ids"])
+                    engine_prompt.get("prompt_token_ids", [])
+                    or engine_prompt.get("prompt_embeds", []))
+
                 if request.use_beam_search:
                     sampling_params = request.to_beam_search_params(
                         default_max_tokens, self.default_sampling_params)
@@ -154,13 +156,13 @@
 
                 if isinstance(sampling_params, BeamSearchParams):
                     generator = self.engine_client.beam_search(
                         prompt=engine_prompt,
                         request_id=request_id,
                         params=sampling_params,
                     )
                 else:
                     generator = self.engine_client.generate(
                         engine_prompt,
                         sampling_params,
                         request_id_item,
                         lora_request=lora_request,
@@ -211,7 +213,7 @@
                 # We did not pass it into vLLM engine to avoid being redundant
                 # with the inputs token IDs
                 if final_res.prompt is None:
-                    final_res.prompt = request_prompts[i]["prompt"]
+                    final_res.prompt = request_prompts[i].get("prompt")
 
             final_res_batch_checked = cast(list[RequestOutput],
                                            final_res_batch)