Extend renderer with embedding support and integrate completion endpoint (vllm-project#24405)

sfeng33 · FeiDaLI · commit 512aca502ae7 · 2025-09-25T18:54:14.000+08:00
Signed-off-by: sfeng33 &lt;4florafeng@gmail.com&gt;
diff --git a/tests/entrypoints/openai/test_prompt_validation.py b/tests/entrypoints/openai/test_prompt_validation.py
@@ -10,7 +10,7 @@
 import regex as re
 import torch
 
-from vllm.entrypoints.openai.serving_engine import OpenAIServing
+from vllm.entrypoints.renderer import BaseRenderer
 
 from ...utils import RemoteOpenAIServer
 
@@ -27,12 +27,16 @@ async def test_empty_prompt():
     with RemoteOpenAIServer(model_name, server_args) as remote_server:
         client = remote_server.get_async_client()
 
-        with pytest.raises(openai.BadRequestError,
-                           match="decoder prompt cannot be empty"):
+        with pytest.raises(
+                openai.BadRequestError,
+                match=
+                "Either prompt or prompt_embeds must be provided and non-empty."
+        ):
             await client.completions.create(model=model_name,
                                             prompt="",
                                             max_tokens=5,
-                                            temperature=0.0)
+                                            temperature=0.0,
+                                            extra_body={"prompt_embeds": []})
 
 
 @pytest.mark.asyncio
@@ -83,7 +87,7 @@ def test_load_prompt_embeds(dtype: torch.dtype, layout: torch.layout,
     buffer.seek(0)
     encoded_tensor = pybase64.b64encode(buffer.getvalue())
 
-    loaded_prompt_embeds = OpenAIServing._load_prompt_embeds(encoded_tensor)
+    loaded_prompt_embeds = BaseRenderer.load_prompt_embeds(encoded_tensor)
     assert len(loaded_prompt_embeds) == 1
     loaded_tensor = loaded_prompt_embeds[0]["prompt_embeds"]
     assert loaded_tensor.device.type == "cpu"
diff --git a/tests/entrypoints/test_renderer.py b/tests/entrypoints/test_renderer.py
@@ -1,13 +1,17 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import io
 from dataclasses import dataclass
 from typing import Optional
 from unittest.mock import AsyncMock, MagicMock
 
+import pybase64
 import pytest
+import torch
 
 from vllm.entrypoints.renderer import CompletionRenderer
+from vllm.inputs.data import is_embeds_prompt
 
 
 @dataclass
@@ -178,3 +182,132 @@ async def test_no_tokenizer_for_text(self, mock_model_config):
         with pytest.raises(ValueError, match="No tokenizer available"):
             await renderer_no_tokenizer.render_prompt(
                 prompt_or_prompts="Hello world", max_length=100)
+
+    @pytest.mark.asyncio
+    async def test_token_input_with_needs_detokenization(
+            self, renderer, mock_async_tokenizer):
+        # When needs_detokenization=True for token inputs, renderer should
+        # use the async tokenizer to decode and include the original text
+        # in the returned prompt object.
+        mock_async_tokenizer.decode = AsyncMock(return_value="decoded text")
+        renderer.async_tokenizer_pool[
+            renderer.tokenizer] = mock_async_tokenizer
+
+        tokens = [1, 2, 3, 4]
+        results = await renderer.render_prompt(
+            prompt_or_prompts=tokens,
+            needs_detokenization=True,
+        )
+
+        assert len(results) == 1
+        assert results[0]["prompt_token_ids"] == tokens
+        assert results[0]["prompt"] == "decoded text"
+        mock_async_tokenizer.decode.assert_awaited_once()
+
+
+class TestRenderEmbedPrompt:
+
+    def _create_test_embed_bytes(self, tensor: torch.Tensor) -> bytes:
+        """Helper to create base64-encoded tensor bytes"""
+        buffer = io.BytesIO()
+        torch.save(tensor, buffer)
+        buffer.seek(0)
+        return pybase64.b64encode(buffer.read())
+
+    @pytest.mark.asyncio
+    async def test_single_prompt_embed(self, renderer):
+        # Create a test tensor
+        test_tensor = torch.randn(10, 768, dtype=torch.float32)
+        embed_bytes = self._create_test_embed_bytes(test_tensor)
+
+        results = await renderer.render_prompt_and_embeds(
+            prompt_embeds=embed_bytes, cache_salt="test_salt")
+
+        assert len(results) == 1
+        assert is_embeds_prompt(results[0])
+        assert torch.allclose(results[0]["prompt_embeds"], test_tensor)
+        assert results[0]["cache_salt"] == "test_salt"
+
+    @pytest.mark.asyncio
+    async def test_multiple_prompt_embeds(self, renderer):
+        # Create multiple test tensors
+        test_tensors = [
+            torch.randn(8, 512, dtype=torch.float32),
+            torch.randn(12, 512, dtype=torch.float32),
+        ]
+        embed_bytes_list = [
+            self._create_test_embed_bytes(t) for t in test_tensors
+        ]
+
+        results = await renderer.render_prompt_and_embeds(
+            prompt_embeds=embed_bytes_list)
+
+        assert len(results) == 2
+        for i, result in enumerate(results):
+            assert is_embeds_prompt(result)
+            assert torch.allclose(result["prompt_embeds"], test_tensors[i])
+
+    @pytest.mark.asyncio
+    async def test_prompt_embed_truncation(self, renderer):
+        # Create tensor with more tokens than truncation limit
+        test_tensor = torch.randn(20, 768, dtype=torch.float32)
+        embed_bytes = self._create_test_embed_bytes(test_tensor)
+
+        results = await renderer.render_prompt_and_embeds(
+            prompt_embeds=embed_bytes, truncate_prompt_tokens=10)
+
+        assert len(results) == 1
+        # Should keep last 10 tokens
+        expected = test_tensor[-10:]
+        assert torch.allclose(results[0]["prompt_embeds"], expected)
+
+    @pytest.mark.asyncio
+    async def test_prompt_embed_different_dtypes(self, renderer):
+        # Test different supported dtypes
+        dtypes = [torch.float32, torch.float16, torch.bfloat16]
+
+        for dtype in dtypes:
+            test_tensor = torch.randn(5, 256, dtype=dtype)
+            embed_bytes = self._create_test_embed_bytes(test_tensor)
+
+            results = await renderer.render_prompt_and_embeds(
+                prompt_embeds=embed_bytes)
+
+            assert len(results) == 1
+            assert results[0]["prompt_embeds"].dtype == dtype
+
+    @pytest.mark.asyncio
+    async def test_prompt_embed_squeeze_batch_dim(self, renderer):
+        # Test tensor with batch dimension gets squeezed
+        test_tensor = torch.randn(1, 10, 768, dtype=torch.float32)
+        embed_bytes = self._create_test_embed_bytes(test_tensor)
+
+        results = await renderer.render_prompt_and_embeds(
+            prompt_embeds=embed_bytes)
+
+        assert len(results) == 1
+        # Should be squeezed to 2D
+        assert results[0]["prompt_embeds"].shape == (10, 768)
+
+    @pytest.mark.asyncio
+    async def test_both_prompts_and_embeds(self, renderer,
+                                           mock_async_tokenizer):
+        # Set up text tokenization
+        mock_async_tokenizer.return_value = MockTokenizerResult(
+            [101, 102, 103])
+        renderer.async_tokenizer_pool[
+            renderer.tokenizer] = mock_async_tokenizer
+
+        # Create embed
+        test_tensor = torch.randn(5, 256, dtype=torch.float32)
+        embed_bytes = self._create_test_embed_bytes(test_tensor)
+
+        results = await renderer.render_prompt_and_embeds(
+            prompt_or_prompts="Hello world", prompt_embeds=embed_bytes)
+
+        assert len(results) == 2
+        # First should be embed prompt
+        assert is_embeds_prompt(results[0])
+        # Second should be tokens prompt
+        assert "prompt_token_ids" in results[1]
+        assert results[1]["prompt_token_ids"] == [101, 102, 103]
diff --git a/tests/v1/entrypoints/openai/test_completion.py b/tests/v1/entrypoints/openai/test_completion.py
@@ -686,7 +686,7 @@ async def test_invalid_grammar(client: openai.AsyncOpenAI, model_name: str):
 async def test_completion_with_empty_prompt_embeds(
         client: openai.AsyncOpenAI) -> None:
     """Test completion with empty prompt embeds."""
-    payload: dict[str, list] = {"prompt_embeds": []}
+    payload: dict[str, object] = {"prompt": "Hello", "prompt_embeds": []}
     headers: dict[str, str] = {"Content-Type": "application/json"}
     # base_url = http://localhost:8000/v1/completions
     response = requests.post(f"{client.base_url}completions",
diff --git a/vllm/entrypoints/openai/protocol.py b/vllm/entrypoints/openai/protocol.py
@@ -1270,9 +1270,20 @@ def validate_stream_options(cls, 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:
+        prompt = data.get("prompt")
+        prompt_embeds = data.get("prompt_embeds")
+
+        prompt_is_empty = (prompt is None
+                           or (isinstance(prompt, str) and prompt == ""))
+        embeds_is_empty = (prompt_embeds is None
+                           or (isinstance(prompt_embeds, list)
+                               and len(prompt_embeds) == 0))
+
+        if prompt_is_empty and embeds_is_empty:
             raise ValueError(
-                "At least one of `prompt` or `prompt_embeds` must be set.")
+                "Either prompt or prompt_embeds must be provided and non-empty."
+            )
+
         return data
 
     @model_validator(mode="before")
diff --git a/vllm/entrypoints/openai/serving_completion.py b/vllm/entrypoints/openai/serving_completion.py
@@ -26,12 +26,8 @@
                                               PromptTokenUsageInfo,
                                               RequestResponseMetadata,
                                               UsageInfo)
-from vllm.entrypoints.openai.serving_engine import (
-    EmbedsPrompt as ServingEngineEmbedsPrompt)
 from vllm.entrypoints.openai.serving_engine import (OpenAIServing,
-                                                    TextTokensPrompt,
-                                                    clamp_prompt_logprobs,
-                                                    is_text_tokens_prompt)
+                                                    clamp_prompt_logprobs)
 # yapf: enable
 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
 from vllm.entrypoints.utils import get_max_tokens
@@ -132,12 +128,19 @@ async def create_completion(
             else:
                 tokenizer = await self.engine_client.get_tokenizer(lora_request
                                                                    )
-
-            request_prompts, engine_prompts = await self._preprocess_completion(
-                request,
-                tokenizer,
-                request.prompt,
+            renderer = self._get_renderer(tokenizer)
+            max_input_tokens_len = self.max_model_len - (request.max_tokens
+                                                         or 0)
+
+            engine_prompts = await renderer.render_prompt_and_embeds(
+                prompt_or_prompts=request.prompt,
+                prompt_embeds=request.prompt_embeds,
+                max_length=max_input_tokens_len,
+                truncate_prompt_tokens=request.truncate_prompt_tokens,
                 add_special_tokens=request.add_special_tokens,
+                cache_salt=request.cache_salt,
+                needs_detokenization=bool(request.echo
+                                          and not request.return_token_ids),
             )
         except ValueError as e:
             logger.exception("Error in preprocessing prompt inputs")
@@ -198,7 +201,7 @@ async def create_completion(
 
                 self._log_inputs(
                     request_id_item,
-                    request_prompts[i],
+                    engine_prompt,
                     params=sampling_params,
                     lora_request=lora_request,
                 )
@@ -249,7 +252,7 @@ async def create_completion(
         if stream:
             return self.completion_stream_generator(
                 request,
-                request_prompts,
+                engine_prompts,
                 result_generator,
                 request_id,
                 created_time,
@@ -273,11 +276,9 @@ async def create_completion(
                 # We did not pass it into vLLM engine to avoid being redundant
                 # with the inputs token IDs
                 if final_res.prompt is None:
-                    request_prompt = request_prompts[i]
-                    if is_text_tokens_prompt(request_prompt):
-                        final_res.prompt = request_prompt["prompt"]
-                    else:
-                        final_res.prompt = None
+                    engine_prompt = engine_prompts[i]
+                    final_res.prompt = None if is_embeds_prompt(
+                        engine_prompt) else engine_prompt.get("prompt")
 
             final_res_batch_checked = cast(list[RequestOutput],
                                            final_res_batch)
@@ -313,8 +314,7 @@ async def fake_stream_generator() -> AsyncGenerator[str, None]:
     async def completion_stream_generator(
         self,
         request: CompletionRequest,
-        request_prompts: list[Union[TextTokensPrompt,
-                                    ServingEngineEmbedsPrompt]],
+        engine_prompts: list[Union[TokensPrompt, EmbedsPrompt]],
         result_generator: AsyncIterator[tuple[int, RequestOutput]],
         request_id: str,
         created_time: int,
@@ -350,14 +350,11 @@ async def completion_stream_generator(
                     num_cached_tokens = res.num_cached_tokens
                     first_iteration = False
 
-                if res.prompt is not None:
-                    prompt_text = res.prompt
-                else:
-                    request_prompt = request_prompts[prompt_idx]
-                    if is_text_tokens_prompt(request_prompt):
-                        prompt_text = request_prompt["prompt"]
-                    else:
-                        prompt_text = None
+                prompt_text = res.prompt
+                if prompt_text is None:
+                    engine_prompt = engine_prompts[prompt_idx]
+                    prompt_text = None if is_embeds_prompt(
+                        engine_prompt) else engine_prompt.get("prompt")
 
                 # Prompt details are excluded from later streamed outputs
                 if prompt_token_ids is not None:
@@ -378,6 +375,8 @@ async def completion_stream_generator(
                     assert request.max_tokens is not None
                     if request.echo and not has_echoed[i]:
                         assert prompt_token_ids is not None
+                        if request.return_token_ids:
+                            prompt_text = ""
                         assert prompt_text is not None
                         if request.max_tokens == 0:
                             # only return the prompt
@@ -525,6 +524,8 @@ def request_output_to_completion_response(
             for output in final_res.outputs:
                 assert request.max_tokens is not None
                 if request.echo:
+                    if request.return_token_ids:
+                        prompt_text = ""
                     assert prompt_text is not None
                     if request.max_tokens == 0:
                         token_ids = prompt_token_ids
diff --git a/vllm/entrypoints/openai/serving_embedding.py b/vllm/entrypoints/openai/serving_embedding.py
@@ -28,7 +28,6 @@
                                                     TextTokensPrompt)
 # yapf: enable
 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
-from vllm.inputs.data import EmbedsPrompt as EngineEmbedsPrompt
 from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
 from vllm.logger import init_logger
 from vllm.outputs import (EmbeddingOutput, EmbeddingRequestOutput,
@@ -290,7 +289,7 @@ def _is_text_tokens_prompt(self, prompt) -> bool:
     async def _create_single_prompt_generator(
         self,
         ctx: EmbeddingServeContext,
-        engine_prompt: Union[EngineTokensPrompt, EngineEmbedsPrompt],
+        engine_prompt: EngineTokensPrompt,
         pooling_params: PoolingParams,
         trace_headers: Optional[Mapping[str, str]],
         prompt_index: int,
@@ -303,12 +302,6 @@ async def _create_single_prompt_generator(
                          params=pooling_params,
                          lora_request=ctx.lora_request)
 
-        # Mypy has an existing bug related to inferring the variance
-        # of TypedDicts with `builtins.enumerate`:
-        # https://github.com/python/mypy/issues/8586#issuecomment-2867698435
-        engine_prompt = cast(Union[EngineTokensPrompt, EngineEmbedsPrompt],
-                             engine_prompt)
-
         # Return the original generator without wrapping
         return self.engine_client.encode(
             engine_prompt,
@@ -375,12 +368,8 @@ async def _prepare_generators(
                         continue
 
                 # Normal processing for short prompts or non-token prompts
-                # Cast engine_prompt to the expected type for mypy
-                engine_prompt_typed = cast(
-                    Union[EngineTokensPrompt, EngineEmbedsPrompt],
-                    engine_prompt)
                 generator = await self._create_single_prompt_generator(
-                    ctx, engine_prompt_typed, pooling_params, trace_headers, i)
+                    ctx, engine_prompt, pooling_params, trace_headers, i)
                 generators.append(generator)
 
             from vllm.utils import merge_async_iterators
diff --git a/vllm/entrypoints/openai/serving_engine.py b/vllm/entrypoints/openai/serving_engine.py
diff --git a/vllm/entrypoints/renderer.py b/vllm/entrypoints/renderer.py
diff --git a/vllm/inputs/data.py b/vllm/inputs/data.py
