@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 15cb047e25

Fork 0

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 14, 05:30 UTC

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

Signed-off-by: sfeng33 <4florafeng@gmail.com>

Flora Feng committed Sep 9, 2025 at 10:46 UTC 15cb047e25e58e79273567d3448563cf43bb39e7

9 files changed +411 -310

tests/entrypoints/openai/test_prompt_validation.py

+9 -5

@@ -10,7 +10,7 @@ import pytest

10 import regex as re

11 import torch

13 -from vllm.entrypoints.openai.serving_engine import OpenAIServing

13 +from vllm.entrypoints.renderer import BaseRenderer

15 from ...utils import RemoteOpenAIServer

@@ -27,12 +27,16 @@ async def test_empty_prompt():

27 with RemoteOpenAIServer(model_name, server_args) as remote_server:

28 client = remote_server.get_async_client()

30 - with pytest.raises(openai.BadRequestError,

31 - match="decoder prompt cannot be empty"):

30 + with pytest.raises(

31 + openai.BadRequestError,

32 + match=

33 + "Either prompt or prompt_embeds must be provided and non-empty."

34 + ):

35 await client.completions.create(model=model_name,

36 prompt="",

37 max_tokens=5,

35 - temperature=0.0)

38 + temperature=0.0,

39 + extra_body={"prompt_embeds": []})

42 @pytest.mark.asyncio

@@ -83,7 +87,7 @@ def test_load_prompt_embeds(dtype: torch.dtype, layout: torch.layout,

87 buffer.seek(0)

88 encoded_tensor = pybase64.b64encode(buffer.getvalue())

86 - loaded_prompt_embeds = OpenAIServing._load_prompt_embeds(encoded_tensor)

90 + loaded_prompt_embeds = BaseRenderer.load_prompt_embeds(encoded_tensor)

91 assert len(loaded_prompt_embeds) == 1

92 loaded_tensor = loaded_prompt_embeds[0]["prompt_embeds"]

93 assert loaded_tensor.device.type == "cpu"

tests/entrypoints/test_renderer.py

+133

@@ -1,13 +1,17 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 +import io

5 from dataclasses import dataclass

6 from typing import Optional

7 from unittest.mock import AsyncMock, MagicMock

9 +import pybase64

10 import pytest

11 +import torch

13 from vllm.entrypoints.renderer import CompletionRenderer

14 +from vllm.inputs.data import is_embeds_prompt

17 @dataclass

@@ -178,3 +182,132 @@ class TestRenderPrompt:

182 with pytest.raises(ValueError, match="No tokenizer available"):

183 await renderer_no_tokenizer.render_prompt(

184 prompt_or_prompts="Hello world", max_length=100)

185 +

186 + @pytest.mark.asyncio

187 + async def test_token_input_with_needs_detokenization(

188 + self, renderer, mock_async_tokenizer):

189 + # When needs_detokenization=True for token inputs, renderer should

190 + # use the async tokenizer to decode and include the original text

191 + # in the returned prompt object.

192 + mock_async_tokenizer.decode = AsyncMock(return_value="decoded text")

193 + renderer.async_tokenizer_pool[

194 + renderer.tokenizer] = mock_async_tokenizer

195 +

196 + tokens = [1, 2, 3, 4]

197 + results = await renderer.render_prompt(

198 + prompt_or_prompts=tokens,

199 + needs_detokenization=True,

200 + )

201 +

202 + assert len(results) == 1

203 + assert results[0]["prompt_token_ids"] == tokens

204 + assert results[0]["prompt"] == "decoded text"

205 + mock_async_tokenizer.decode.assert_awaited_once()

206 +

207 +

208 +class TestRenderEmbedPrompt:

209 +

210 + def _create_test_embed_bytes(self, tensor: torch.Tensor) -> bytes:

211 + """Helper to create base64-encoded tensor bytes"""

212 + buffer = io.BytesIO()

213 + torch.save(tensor, buffer)

214 + buffer.seek(0)

215 + return pybase64.b64encode(buffer.read())

216 +

217 + @pytest.mark.asyncio

218 + async def test_single_prompt_embed(self, renderer):

219 + # Create a test tensor

220 + test_tensor = torch.randn(10, 768, dtype=torch.float32)

221 + embed_bytes = self._create_test_embed_bytes(test_tensor)

222 +

223 + results = await renderer.render_prompt_and_embeds(

224 + prompt_embeds=embed_bytes, cache_salt="test_salt")

225 +

226 + assert len(results) == 1

227 + assert is_embeds_prompt(results[0])

228 + assert torch.allclose(results[0]["prompt_embeds"], test_tensor)

229 + assert results[0]["cache_salt"] == "test_salt"

230 +

231 + @pytest.mark.asyncio

232 + async def test_multiple_prompt_embeds(self, renderer):

233 + # Create multiple test tensors

234 + test_tensors = [

235 + torch.randn(8, 512, dtype=torch.float32),

236 + torch.randn(12, 512, dtype=torch.float32),

237 + ]

238 + embed_bytes_list = [

239 + self._create_test_embed_bytes(t) for t in test_tensors

240 + ]

241 +

242 + results = await renderer.render_prompt_and_embeds(

243 + prompt_embeds=embed_bytes_list)

244 +

245 + assert len(results) == 2

246 + for i, result in enumerate(results):

247 + assert is_embeds_prompt(result)

248 + assert torch.allclose(result["prompt_embeds"], test_tensors[i])

249 +

250 + @pytest.mark.asyncio

251 + async def test_prompt_embed_truncation(self, renderer):

252 + # Create tensor with more tokens than truncation limit

253 + test_tensor = torch.randn(20, 768, dtype=torch.float32)

254 + embed_bytes = self._create_test_embed_bytes(test_tensor)

255 +

256 + results = await renderer.render_prompt_and_embeds(

257 + prompt_embeds=embed_bytes, truncate_prompt_tokens=10)

258 +

259 + assert len(results) == 1

260 + # Should keep last 10 tokens

261 + expected = test_tensor[-10:]

262 + assert torch.allclose(results[0]["prompt_embeds"], expected)

263 +

264 + @pytest.mark.asyncio

265 + async def test_prompt_embed_different_dtypes(self, renderer):

266 + # Test different supported dtypes

267 + dtypes = [torch.float32, torch.float16, torch.bfloat16]

268 +

269 + for dtype in dtypes:

270 + test_tensor = torch.randn(5, 256, dtype=dtype)

271 + embed_bytes = self._create_test_embed_bytes(test_tensor)

272 +

273 + results = await renderer.render_prompt_and_embeds(

274 + prompt_embeds=embed_bytes)

275 +

276 + assert len(results) == 1

277 + assert results[0]["prompt_embeds"].dtype == dtype

278 +

279 + @pytest.mark.asyncio

280 + async def test_prompt_embed_squeeze_batch_dim(self, renderer):

281 + # Test tensor with batch dimension gets squeezed

282 + test_tensor = torch.randn(1, 10, 768, dtype=torch.float32)

283 + embed_bytes = self._create_test_embed_bytes(test_tensor)

284 +

285 + results = await renderer.render_prompt_and_embeds(

286 + prompt_embeds=embed_bytes)

287 +

288 + assert len(results) == 1

289 + # Should be squeezed to 2D

290 + assert results[0]["prompt_embeds"].shape == (10, 768)

291 +

292 + @pytest.mark.asyncio

293 + async def test_both_prompts_and_embeds(self, renderer,

294 + mock_async_tokenizer):

295 + # Set up text tokenization

296 + mock_async_tokenizer.return_value = MockTokenizerResult(

297 + [101, 102, 103])

298 + renderer.async_tokenizer_pool[

299 + renderer.tokenizer] = mock_async_tokenizer

300 +

301 + # Create embed

302 + test_tensor = torch.randn(5, 256, dtype=torch.float32)

303 + embed_bytes = self._create_test_embed_bytes(test_tensor)

304 +

305 + results = await renderer.render_prompt_and_embeds(

306 + prompt_or_prompts="Hello world", prompt_embeds=embed_bytes)

307 +

308 + assert len(results) == 2

309 + # First should be embed prompt

310 + assert is_embeds_prompt(results[0])

311 + # Second should be tokens prompt

312 + assert "prompt_token_ids" in results[1]

313 + assert results[1]["prompt_token_ids"] == [101, 102, 103]

tests/v1/entrypoints/openai/test_completion.py

+1 -1

@@ -686,7 +686,7 @@ async def test_invalid_grammar(client: openai.AsyncOpenAI, model_name: str):

686 async def test_completion_with_empty_prompt_embeds(

687 client: openai.AsyncOpenAI) -> None:

688 """Test completion with empty prompt embeds."""

689 - payload: dict[str, list] = {"prompt_embeds": []}

689 + payload: dict[str, object] = {"prompt": "Hello", "prompt_embeds": []}

690 headers: dict[str, str] = {"Content-Type": "application/json"}

691 # base_url = http://localhost:8000/v1/completions

692 response = requests.post(f"{client.base_url}completions",

vllm/entrypoints/openai/protocol.py

+13 -2

@@ -1270,9 +1270,20 @@ class CompletionRequest(OpenAIBaseModel):

1270 @model_validator(mode="before")

1271 @classmethod

1272 def validate_prompt_and_prompt_embeds(cls, data):

1273 - if data.get("prompt") is None and data.get("prompt_embeds") is None:

1273 + prompt = data.get("prompt")

1274 + prompt_embeds = data.get("prompt_embeds")

1275 +

1276 + prompt_is_empty = (prompt is None

1277 + or (isinstance(prompt, str) and prompt == ""))

1278 + embeds_is_empty = (prompt_embeds is None

1279 + or (isinstance(prompt_embeds, list)

1280 + and len(prompt_embeds) == 0))

1281 +

1282 + if prompt_is_empty and embeds_is_empty:

1283 raise ValueError(

1275 - "At least one of `prompt` or `prompt_embeds` must be set.")

1284 + "Either prompt or prompt_embeds must be provided and non-empty."

1285 + )

1286 +

1287 return data

1288

1289 @model_validator(mode="before")

vllm/entrypoints/openai/serving_completion.py

+28 -27

@@ -26,12 +26,8 @@ from vllm.entrypoints.openai.protocol import (CompletionLogProbs,

26 PromptTokenUsageInfo,

27 RequestResponseMetadata,

28 UsageInfo)

29 -from vllm.entrypoints.openai.serving_engine import (

30 - EmbedsPrompt as ServingEngineEmbedsPrompt)

29 from vllm.entrypoints.openai.serving_engine import (OpenAIServing,

32 - TextTokensPrompt,

33 - clamp_prompt_logprobs,

34 - is_text_tokens_prompt)

30 + clamp_prompt_logprobs)

31 # yapf: enable

32 from vllm.entrypoints.openai.serving_models import OpenAIServingModels

33 from vllm.entrypoints.utils import get_max_tokens

@@ -132,12 +128,19 @@ class OpenAIServingCompletion(OpenAIServing):

128 else:

129 tokenizer = await self.engine_client.get_tokenizer(lora_request

130 )

135 -

136 - request_prompts, engine_prompts = await self._preprocess_completion(

137 - request,

138 - tokenizer,

139 - request.prompt,

131 + renderer = self._get_renderer(tokenizer)

132 + max_input_tokens_len = self.max_model_len - (request.max_tokens

133 + or 0)

134 +

135 + engine_prompts = await renderer.render_prompt_and_embeds(

136 + prompt_or_prompts=request.prompt,

137 + prompt_embeds=request.prompt_embeds,

138 + max_length=max_input_tokens_len,

139 + truncate_prompt_tokens=request.truncate_prompt_tokens,

140 add_special_tokens=request.add_special_tokens,

141 + cache_salt=request.cache_salt,

142 + needs_detokenization=bool(request.echo

143 + and not request.return_token_ids),

144 )

145 except ValueError as e:

146 logger.exception("Error in preprocessing prompt inputs")

@@ -198,7 +201,7 @@ class OpenAIServingCompletion(OpenAIServing):

201

202 self._log_inputs(

203 request_id_item,

201 - request_prompts[i],

204 + engine_prompt,

205 params=sampling_params,

206 lora_request=lora_request,

207 )

@@ -249,7 +252,7 @@ class OpenAIServingCompletion(OpenAIServing):

252 if stream:

253 return self.completion_stream_generator(

254 request,

252 - request_prompts,

255 + engine_prompts,

256 result_generator,

257 request_id,

258 created_time,

@@ -273,11 +276,9 @@ class OpenAIServingCompletion(OpenAIServing):

276 # We did not pass it into vLLM engine to avoid being redundant

277 # with the inputs token IDs

278 if final_res.prompt is None:

276 - request_prompt = request_prompts[i]

277 - if is_text_tokens_prompt(request_prompt):

278 - final_res.prompt = request_prompt["prompt"]

279 - else:

280 - final_res.prompt = None

279 + engine_prompt = engine_prompts[i]

280 + final_res.prompt = None if is_embeds_prompt(

281 + engine_prompt) else engine_prompt.get("prompt")

282

283 final_res_batch_checked = cast(list[RequestOutput],

284 final_res_batch)

@@ -313,8 +314,7 @@ class OpenAIServingCompletion(OpenAIServing):

314 async def completion_stream_generator(

315 self,

316 request: CompletionRequest,

316 - request_prompts: list[Union[TextTokensPrompt,

317 - ServingEngineEmbedsPrompt]],

317 + engine_prompts: list[Union[TokensPrompt, EmbedsPrompt]],

318 result_generator: AsyncIterator[tuple[int, RequestOutput]],

319 request_id: str,

320 created_time: int,

@@ -350,14 +350,11 @@ class OpenAIServingCompletion(OpenAIServing):

350 num_cached_tokens = res.num_cached_tokens

351 first_iteration = False

352

353 - if res.prompt is not None:

354 - prompt_text = res.prompt

355 - else:

356 - request_prompt = request_prompts[prompt_idx]

357 - if is_text_tokens_prompt(request_prompt):

358 - prompt_text = request_prompt["prompt"]

359 - else:

360 - prompt_text = None

353 + prompt_text = res.prompt

354 + if prompt_text is None:

355 + engine_prompt = engine_prompts[prompt_idx]

356 + prompt_text = None if is_embeds_prompt(

357 + engine_prompt) else engine_prompt.get("prompt")

358

359 # Prompt details are excluded from later streamed outputs

360 if prompt_token_ids is not None:

@@ -378,6 +375,8 @@ class OpenAIServingCompletion(OpenAIServing):

375 assert request.max_tokens is not None

376 if request.echo and not has_echoed[i]:

377 assert prompt_token_ids is not None

378 + if request.return_token_ids:

379 + prompt_text = ""

380 assert prompt_text is not None

381 if request.max_tokens == 0:

382 # only return the prompt

@@ -525,6 +524,8 @@ class OpenAIServingCompletion(OpenAIServing):

524 for output in final_res.outputs:

525 assert request.max_tokens is not None

526 if request.echo:

527 + if request.return_token_ids:

528 + prompt_text = ""

529 assert prompt_text is not None

530 if request.max_tokens == 0:

531 token_ids = prompt_token_ids

vllm/entrypoints/openai/serving_embedding.py

+2 -13

@@ -28,7 +28,6 @@ from vllm.entrypoints.openai.serving_engine import (EmbeddingServeContext,

28 TextTokensPrompt)

29 # yapf: enable

30 from vllm.entrypoints.openai.serving_models import OpenAIServingModels

31 -from vllm.inputs.data import EmbedsPrompt as EngineEmbedsPrompt

31 from vllm.inputs.data import TokensPrompt as EngineTokensPrompt

32 from vllm.logger import init_logger

33 from vllm.outputs import (EmbeddingOutput, EmbeddingRequestOutput,

@@ -290,7 +289,7 @@ class EmbeddingMixin(OpenAIServing):

289 async def _create_single_prompt_generator(

290 self,

291 ctx: EmbeddingServeContext,

293 - engine_prompt: Union[EngineTokensPrompt, EngineEmbedsPrompt],

292 + engine_prompt: EngineTokensPrompt,

293 pooling_params: PoolingParams,

294 trace_headers: Optional[Mapping[str, str]],

295 prompt_index: int,

@@ -303,12 +302,6 @@ class EmbeddingMixin(OpenAIServing):

302 params=pooling_params,

303 lora_request=ctx.lora_request)

304

306 - # Mypy has an existing bug related to inferring the variance

307 - # of TypedDicts with `builtins.enumerate`:

308 - # https://github.com/python/mypy/issues/8586#issuecomment-2867698435

309 - engine_prompt = cast(Union[EngineTokensPrompt, EngineEmbedsPrompt],

310 - engine_prompt)

311 -

305 # Return the original generator without wrapping

306 return self.engine_client.encode(

307 engine_prompt,

@@ -375,12 +368,8 @@ class EmbeddingMixin(OpenAIServing):

368 continue

369

370 # Normal processing for short prompts or non-token prompts

378 - # Cast engine_prompt to the expected type for mypy

379 - engine_prompt_typed = cast(

380 - Union[EngineTokensPrompt, EngineEmbedsPrompt],

381 - engine_prompt)

371 generator = await self._create_single_prompt_generator(

383 - ctx, engine_prompt_typed, pooling_params, trace_headers, i)

372 + ctx, engine_prompt, pooling_params, trace_headers, i)

373 generators.append(generator)

374

375 from vllm.utils import merge_async_iterators

vllm/entrypoints/openai/serving_engine.py

+2 -215

@@ -1,7 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -import asyncio

4 -import io

3 import json

4 import sys

5 import time

@@ -9,10 +7,8 @@ import traceback

7 from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence

8 from concurrent.futures import ThreadPoolExecutor

9 from http import HTTPStatus

12 -from typing import (Annotated, Any, Callable, ClassVar, Generic, Optional,

13 - TypeVar, Union, cast, overload)

10 +from typing import Any, Callable, ClassVar, Generic, Optional, TypeVar, Union

15 -import pybase64

12 import torch

13 from fastapi import Request

14 from pydantic import BaseModel, ConfigDict, Field

@@ -64,10 +60,8 @@ from vllm.entrypoints.openai.serving_models import OpenAIServingModels

60 from vllm.entrypoints.openai.tool_parsers import ToolParser

61 from vllm.entrypoints.renderer import BaseRenderer, CompletionRenderer

62 # yapf: enable

67 -from vllm.inputs.data import EmbedsPrompt as EngineEmbedsPrompt

63 from vllm.inputs.data import PromptType

64 from vllm.inputs.data import TokensPrompt as EngineTokensPrompt

70 -from vllm.inputs.parse import parse_and_batch_prompt

65 from vllm.logger import init_logger

66 from vllm.logprobs import Logprob, PromptLogprobs

67 from vllm.lora.request import LoRARequest

@@ -149,8 +143,7 @@ class RequestProcessingMixin(BaseModel):

143 """

144

145 request_prompts: Optional[Sequence[RequestPrompt]] = []

152 - engine_prompts: Optional[Union[list[EngineTokensPrompt],

153 - list[EngineEmbedsPrompt]]] = []

146 + engine_prompts: Optional[list[EngineTokensPrompt]] = []

147

148 model_config = ConfigDict(arbitrary_types_allowed=True)

149

@@ -368,13 +361,6 @@ class OpenAIServing:

361 for i, engine_prompt in enumerate(ctx.engine_prompts):

362 request_id_item = f"{ctx.request_id}-{i}"

363

371 - # Mypy has an existing bug related to inferring the variance of

372 - # TypedDicts with `builtins.enumerate`:

373 - # https://github.com/python/mypy/issues/8586#issuecomment-2867698435

374 - engine_prompt = cast(

375 - Union[EngineTokensPrompt, EngineEmbedsPrompt],

376 - engine_prompt)

377 -

364 self._log_inputs(

365 request_id_item,

366 engine_prompt,

@@ -737,170 +723,6 @@ class OpenAIServing:

723 tokenizer=tokenizer,

724 )

725

740 - async def _tokenize_prompt_input_or_inputs_async(

741 - self,

742 - request: AnyRequest,

743 - tokenizer: Optional[AnyTokenizer],

744 - input_or_inputs: Optional[Union[str, list[str], list[int],

745 - list[list[int]]]],

746 - add_special_tokens: bool = True,

747 - ) -> tuple[list[TextTokensPrompt], list[EmbedsPrompt]]:

748 - """

749 - Tokenize/detokenize depending on the input format.

750 -

751 - According to `OpenAI API <https://platform.openai.com/docs/api-reference/embeddings/create>`_

752 - , each input can be a string or array of tokens. Note that each request

753 - can pass one or more inputs.

754 - """

755 - inputs_embeds = list[EmbedsPrompt]()

756 - inputs_text = list[TextTokensPrompt]()

757 -

758 - truncate_prompt_tokens = getattr(request, "truncate_prompt_tokens",

759 - None)

760 -

761 - if (truncate_prompt_tokens or 0) < 0:

762 - truncate_prompt_tokens = self.max_model_len

763 -

764 - if (isinstance(request, CompletionRequest)

765 - and request.prompt_embeds is not None):

766 - inputs_embeds.extend(

767 - self._load_prompt_embeds(request.prompt_embeds,

768 - truncate_prompt_tokens))

769 -

770 - # Empty prompts are okay as long as there are prompt embeddings

771 - if input_or_inputs is None or (inputs_embeds

772 - and input_or_inputs == ""):

773 - return [], inputs_embeds

774 -

775 - # Although our type checking is based on mypy,

776 - # VSCode Pyright extension should still work properly

777 - # "is False" is required for Pyright to perform type narrowing

778 - # See: https://github.com/microsoft/pyright/issues/7672

779 -

780 - # Parse and batch the input prompts

781 - batch_inputs = parse_and_batch_prompt(input_or_inputs)

782 -

783 - # Process each input in the batch concurrently

784 - tasks = []

785 - for prompt_input in batch_inputs:

786 - if prompt_input["is_tokens"] is False:

787 - assert tokenizer is not None, (

788 - "Tokenizer is required for text prompts")

789 - task = self._normalize_prompt_text_to_input(

790 - request,

791 - prompt_input["content"],

792 - tokenizer=tokenizer,

793 - add_special_tokens=add_special_tokens,

794 - )

795 - else:

796 - task = self._normalize_prompt_tokens_to_input(

797 - request, prompt_input["content"], tokenizer=tokenizer)

798 - tasks.append(task)

799 -

800 - # Wait for all tokenization tasks to complete

801 - results = await asyncio.gather(*tasks)

802 - inputs_text.extend(results)

803 -

804 - return inputs_text, inputs_embeds

805 -

806 - @overload

807 - async def _preprocess_completion(

808 - self,

809 - request: Union[

810 - DetokenizeRequest,

811 - EmbeddingCompletionRequest,

812 - RerankRequest,

813 - ClassificationRequest,

814 - ScoreRequest,

815 - TokenizeCompletionRequest,

816 - ],

817 - tokenizer: Optional[AnyTokenizer],

818 - input_or_inputs: Union[str, list[str], list[int], list[list[int]]],

819 - add_special_tokens: bool = ...,

820 - ) -> tuple[list[TextTokensPrompt], list[EngineTokensPrompt]]:

821 - ...

822 -

823 - @overload

824 - async def _preprocess_completion(

825 - self,

826 - request: CompletionRequest,

827 - tokenizer: Optional[AnyTokenizer],

828 - input_or_inputs: Optional[Union[str, list[str], list[int],

829 - list[list[int]]]],

830 - add_special_tokens: bool = ...,

831 - ) -> tuple[

832 - list[Union[TextTokensPrompt, EmbedsPrompt]],

833 - list[Union[EngineTokensPrompt, EngineEmbedsPrompt]],

834 - ]:

835 - ...

836 -

837 - async def _preprocess_completion(

838 - self,

839 - request: CompletionLikeRequest,

840 - tokenizer: Optional[AnyTokenizer],

841 - input_or_inputs: Optional[Union[str, list[str], list[int],

842 - list[list[int]]]],

843 - add_special_tokens: bool = True,

844 - ) -> tuple[

845 - Union[list[TextTokensPrompt], list[Union[TextTokensPrompt,

846 - EmbedsPrompt]]],

847 - Union[

848 - list[EngineTokensPrompt],

849 - list[Union[EngineTokensPrompt, EngineEmbedsPrompt]],

850 - ],

851 - ]:

852 - if (not isinstance(request, CompletionRequest)

853 - and input_or_inputs is None):

854 - raise ValueError(

855 - "Prompt embeds with non-completion requests is not"

856 - " currently supported.")

857 -

858 - (

859 - request_prompts_text,

860 - request_prompts_embeds,

861 - ) = await self._tokenize_prompt_input_or_inputs_async(

862 - request,

863 - tokenizer,

864 - input_or_inputs,

865 - add_special_tokens=add_special_tokens,

866 - )

867 -

868 - engine_prompts_text = [

869 - EngineTokensPrompt(

870 - prompt_token_ids=request_prompt_text["prompt_token_ids"])

871 - for request_prompt_text in request_prompts_text

872 - ]

873 - cache_salt = (request.cache_salt if

874 - (hasattr(request, "cache_salt")

875 - and request.cache_salt is not None) else None)

876 - if cache_salt:

877 - for prompt_text in engine_prompts_text:

878 - prompt_text["cache_salt"] = cache_salt

879 -

880 - # This check is equivalent to simply checking if

881 - # `request_prompts_embeds` is empty, but it's difficult to propagate

882 - # overloads to the private helper functions to enable this check.

883 - # This overload is needed because only TextPrompts are allowed for

884 - # non-completion requests and if we don't add the overload here,

885 - # everywhere this function is used outside of serving_completion will

886 - # need logic asserting that only text prompts are in the request.

887 - if (not isinstance(request, CompletionRequest)

888 - and input_or_inputs is not None):

889 - return request_prompts_text, engine_prompts_text

890 -

891 - engine_prompts_embeds = [

892 - EngineEmbedsPrompt(

893 - prompt_embeds=request_prompt_embeds["prompt_embeds"])

894 - for request_prompt_embeds in request_prompts_embeds

895 - ]

896 - if cache_salt:

897 - for prompt_embed in engine_prompts_embeds:

898 - prompt_embed["cache_salt"] = cache_salt

899 -

900 - request_prompts = request_prompts_embeds + request_prompts_text

901 - engine_prompts = engine_prompts_embeds + engine_prompts_text

902 - return request_prompts, engine_prompts

903 -

726 async def _preprocess_chat(

727 self,

728 request: Union[ChatLikeRequest, ResponsesRequest],

@@ -1073,41 +895,6 @@ class OpenAIServing:

895 # OPTIMIZATION

896 priority = orig_priority - 1

897

1076 - @staticmethod

1077 - def _load_prompt_embeds(

1078 - prompt_embeds: Optional[Union[bytes, list[bytes]]],

1079 - truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None,

1080 - ) -> list[EmbedsPrompt]:

1081 -

1082 - def _load_and_validate_embed(embed: bytes) -> EmbedsPrompt:

1083 - tensor = torch.load(

1084 - io.BytesIO(pybase64.b64decode(embed, validate=True)),

1085 - weights_only=True,

1086 - map_location=torch.device("cpu"),

1087 - )

1088 - assert isinstance(tensor, torch.Tensor) and tensor.dtype in (

1089 - torch.float32,

1090 - torch.bfloat16,

1091 - torch.float16,

1092 - )

1093 - tensor = tensor.to_dense()

1094 - if tensor.dim() > 2:

1095 - tensor = tensor.squeeze(0)

1096 - assert tensor.dim() == 2

1097 - if truncate_prompt_tokens is not None:

1098 - tensor = tensor[-truncate_prompt_tokens:]

1099 - return {"prompt_embeds": tensor}

1100 -

1101 - if prompt_embeds:

1102 - if isinstance(prompt_embeds, list):

1103 - return [

1104 - _load_and_validate_embed(embed) for embed in prompt_embeds

1105 - ]

1106 - else:

1107 - return [_load_and_validate_embed(prompt_embeds)]

1108 - else:

1109 - return []

1110 -

898 def _log_inputs(

899 self,

900 request_id: str,

vllm/entrypoints/renderer.py

+220 -47

@@ -2,12 +2,16 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import asyncio

5 +import io

6 from abc import ABC, abstractmethod

7 from typing import Annotated, Optional, Union

9 +import pybase64

10 +import torch

11 from pydantic import Field

13 from vllm.config import ModelConfig

14 +from vllm.inputs.data import EmbedsPrompt as EngineEmbedsPrompt

15 from vllm.inputs.data import TokensPrompt as EngineTokensPrompt

16 from vllm.inputs.parse import parse_and_batch_prompt

17 from vllm.transformers_utils.tokenizer import AnyTokenizer

@@ -49,37 +53,121 @@ class BaseRenderer(ABC):

53 truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,

54 add_special_tokens: Optional[bool] = True,

55 cache_salt: Optional[str] = None,

56 + needs_detokenization: Optional[bool] = False,

57 ) -> list[EngineTokensPrompt]:

58 """

54 - Convert input prompts into tokenized format for engine processing.

55 -

56 - This is the core method that transforms various input formats into

57 - standardized TokensPrompt objects. Implementations should handle

58 - tokenization, special token insertion, truncation, and validation

59 - according to model requirements.

60 -

59 + Convert text or token inputs into engine-ready TokensPrompt objects.

60 +

61 + This method accepts text or token inputs and produces a

62 + list of [`TokensPrompt`][vllm.inputs.data.TokensPrompt] objects

63 + for the engine.

64 +

65 + Args:

66 + prompt_or_prompts: One of:

67 + - ``str``: Single text prompt.

68 + - ``list[str]``: Batch of text prompts.

69 + - ``list[int]``: Single pre-tokenized sequence.

70 + - ``list[list[int]]``: Batch of pre-tokenized sequences.

71 + max_length: Maximum allowable total input token length. If provided,

72 + token inputs longer than this raise ``ValueError``.

73 + truncate_prompt_tokens: Number of tokens to keep. ``None`` means no

74 + truncation. ``0`` yields an empty list (and skips embeds).

75 + ``-1`` maps to ``model_config.max_model_len``.

76 + add_special_tokens: Whether to add model-specific special tokens

77 + during text tokenization.

78 + cache_salt: Optional string to disambiguate prefix cache entries.

79 + needs_detokenization: If True and ``prompt_or_prompts`` is token

80 + input, detokenize IDs back to text for inclusion in outputs.

81 +

82 + Returns:

83 + list[EngineTokensPrompt]: Engine-ready token prompts.

84 +

85 + Raises:

86 + ValueError: If input formats are invalid or length limits exceeded.

87 + """

88 + raise NotImplementedError

89 +

90 + @abstractmethod

91 + async def render_prompt_and_embeds(

92 + self,

93 + prompt_or_prompts: Optional[Union[str, list[str], list[int],

94 + list[list[int]]]] = None,

95 + prompt_embeds: Optional[Union[bytes, list[bytes]]] = None,

96 + max_length: Optional[int] = None,

97 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,

98 + add_special_tokens: Optional[bool] = True,

99 + cache_salt: Optional[str] = None,

100 + needs_detokenization: Optional[bool] = False,

101 + ) -> list[Union[EngineTokensPrompt, EngineEmbedsPrompt]]:

102 + """

103 + Convert text/token and/or base64-encoded embeddings inputs into

104 + engine-ready prompt objects.

105 +

106 + At least one of ``prompt_or_prompts`` or ``prompt_embeds`` must be

107 + provided and non-empty. If both are omitted or empty (e.g., empty

108 + string and empty list), a ``ValueError`` is raised.

109 +

110 Args:

62 - prompt_or_prompts: Input data in various formats:

63 - - str: Single text prompt

64 - - list[str]: Batch of text prompts

65 - - list[int]: Pre-tokenized sequence

66 - - list[list[int]]: Batch of pre-tokenized sequences

67 - max_length: Maximum sequence length (endpoint-specific behavior)

68 - truncate_prompt_tokens: Truncate to last N tokens

69 - (None=no truncation, 0=empty)

70 - add_special_tokens: Add model-specific tokens (e.g., [CLS], [SEP])

71 - to text inputs

72 - cache_salt: Optional string to disambiguate cached prompts

73 -

111 + prompt_or_prompts: Text or token inputs to include.

112 + prompt_embeds: Base64-encoded bytes (or list thereof) containing a

113 + torch-saved tensor to be used as prompt embeddings.

114 + max_length: Maximum allowable total input token length. If provided,

115 + inputs longer than this raise ``ValueError``.

116 + truncate_prompt_tokens: Number of tokens/rows to keep from the end

117 + of the sequence. ``-1`` maps to ``model_config.max_model_len``.

118 + add_special_tokens: Whether to add model-specific special tokens

119 + during text tokenization.

120 + cache_salt: Optional string to disambiguate prefix cache entries.

121 + needs_detokenization: If True and ``prompt_or_prompts`` is token

122 + input, detokenize IDs back to text for inclusion in outputs.

123 +

124 Returns:

75 - list[EngineTokensPrompt]: Tokenized prompts ready for engine

76 - consumption

77 -

125 + list[Union[EngineTokensPrompt, EngineEmbedsPrompt]]:

126 + Engine-ready prompt objects.

127 +

128 Raises:

79 - ValueError: If input format is invalid or length limits exceeded

129 + ValueError: If both ``prompt_or_prompts`` and ``prompt_embeds``

130 + are omitted or empty (decoder prompt cannot be empty), or if

131 + length limits are exceeded.

132 """

133 raise NotImplementedError

134

135 + @classmethod

136 + def load_prompt_embeds(

137 + cls,

138 + prompt_embeds: Union[bytes, list[bytes]],

139 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=0)]] = None,

140 + cache_salt: Optional[str] = None,

141 + ) -> list[EngineEmbedsPrompt]:

142 + """Load and validate base64-encoded embeddings into prompt objects."""

143 +

144 + def _load_and_validate_embed(embed: bytes) -> EngineEmbedsPrompt:

145 + tensor = torch.load(

146 + io.BytesIO(pybase64.b64decode(embed, validate=True)),

147 + weights_only=True,

148 + map_location=torch.device("cpu"),

149 + )

150 + assert isinstance(tensor, torch.Tensor) and tensor.dtype in (

151 + torch.float32,

152 + torch.bfloat16,

153 + torch.float16,

154 + )

155 + tensor = tensor.to_dense()

156 + if tensor.dim() > 2:

157 + tensor = tensor.squeeze(0)

158 + assert tensor.dim() == 2

159 + if truncate_prompt_tokens is not None:

160 + tensor = tensor[-truncate_prompt_tokens:]

161 + embeds_prompt = EngineEmbedsPrompt(prompt_embeds=tensor)

162 + if cache_salt is not None:

163 + embeds_prompt["cache_salt"] = cache_salt

164 + return embeds_prompt

165 +

166 + if isinstance(prompt_embeds, list):

167 + return [_load_and_validate_embed(embed) for embed in prompt_embeds]

168 + else:

169 + return [_load_and_validate_embed(prompt_embeds)]

170 +

171

172 class CompletionRenderer(BaseRenderer):

173

@@ -101,50 +189,110 @@ class CompletionRenderer(BaseRenderer):

189 truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,

190 add_special_tokens: Optional[bool] = True,

191 cache_salt: Optional[str] = None,

192 + needs_detokenization: Optional[bool] = False,

193 ) -> list[EngineTokensPrompt]:

194 """Implementation of prompt rendering for completion-style requests.

195

196 Uses async tokenizer pooling for improved performance. See base class

197 for detailed parameter documentation.

198 """

110 - if truncate_prompt_tokens is not None:

111 - if truncate_prompt_tokens == 0:

112 - return []

113 - if truncate_prompt_tokens < 0:

114 - truncate_prompt_tokens = self.model_config.max_model_len

115 - if max_length is not None and truncate_prompt_tokens > max_length:

116 - raise ValueError(

117 - f"truncate_prompt_tokens ({truncate_prompt_tokens}) "

118 - f"cannot be greater than max_length ({max_length}). "

119 - f"Please select a smaller truncation size.")

199 + truncate_prompt_tokens = self._validate_and_normalize_truncate_tokens(

200 + truncate_prompt_tokens, max_length)

201 + if truncate_prompt_tokens == 0:

202 + return []

203

204 # Parse and batch the input prompts

205 batch_inputs = parse_and_batch_prompt(prompt_or_prompts)

206

124 - rendered_prompts: list[EngineTokensPrompt] = []

125 - tokenize_tasks = []

207 + tasks = []

208 for prompt_input in batch_inputs:

209 if prompt_input["is_tokens"] is True:

210 # Token input

129 - token_ids = self._maybe_apply_truncation(

130 - prompt_input["content"], truncate_prompt_tokens)

131 - rendered_prompts.append(

132 - self._create_tokens_prompt(token_ids, max_length,

133 - cache_salt))

211 + detokenize_task = asyncio.create_task(

212 + # Note: detokenization is needed when echo is enabled,

213 + # where the input token IDs are decoded back to text.

214 + self._maybe_detokenize(prompt_input["content"], max_length,

215 + truncate_prompt_tokens, cache_salt,

216 + needs_detokenization))

217 + tasks.append(detokenize_task)

218 else:

219 # Text input

220 tokenize_task = asyncio.create_task(

221 self._tokenize(prompt_input["content"], max_length,

222 truncate_prompt_tokens, add_special_tokens,

223 cache_salt))

140 - tokenize_tasks.append(tokenize_task)

224 + tasks.append(tokenize_task)

225

226 # Wait for all text tokenization to finish

143 - if tokenize_tasks:

144 - tokenized_text_prompts = await asyncio.gather(*tokenize_tasks)

145 - rendered_prompts.extend(tokenized_text_prompts)

227 + if tasks:

228 + tokenized_text_prompts = await asyncio.gather(*tasks)

229 + return tokenized_text_prompts

230

147 - return rendered_prompts

231 + return []

232 +

233 + async def render_prompt_and_embeds(

234 + self,

235 + prompt_or_prompts: Optional[Union[str, list[str], list[int],

236 + list[list[int]]]] = None,

237 + prompt_embeds: Optional[Union[bytes, list[bytes]]] = None,

238 + max_length: Optional[int] = None,

239 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,

240 + add_special_tokens: Optional[bool] = True,

241 + cache_salt: Optional[str] = None,

242 + needs_detokenization: Optional[bool] = False,

243 + ) -> list[Union[EngineTokensPrompt, EngineEmbedsPrompt]]:

244 + """

245 + Render text/token prompts and/or precomputed embedding prompts. At

246 + least one of `prompt_or_prompts` or `prompt_embeds` must be provided.

247 + """

248 + truncate_prompt_tokens = self._validate_and_normalize_truncate_tokens(

249 + truncate_prompt_tokens, max_length)

250 + if truncate_prompt_tokens == 0:

251 + return []

252 +

253 + rendered: list[Union[EngineTokensPrompt, EngineEmbedsPrompt]] = []

254 +

255 + if prompt_embeds is not None:

256 + rendered.extend(

257 + self.load_prompt_embeds(prompt_embeds, truncate_prompt_tokens,

258 + cache_salt))

259 + if prompt_or_prompts is None or prompt_or_prompts == "":

260 + return rendered

261 +

262 + token_prompts = await self.render_prompt(

263 + prompt_or_prompts=prompt_or_prompts,

264 + max_length=max_length,

265 + truncate_prompt_tokens=truncate_prompt_tokens,

266 + add_special_tokens=add_special_tokens,

267 + cache_salt=cache_salt,

268 + needs_detokenization=needs_detokenization,

269 + )

270 + rendered.extend(token_prompts)

271 +

272 + return rendered

273 +

274 + def _validate_and_normalize_truncate_tokens(

275 + self,

276 + truncate_prompt_tokens: Optional[int],

277 + max_length: Optional[int],

278 + ) -> Optional[int]:

279 + """Validate and normalize truncate_prompt_tokens parameter."""

280 + if truncate_prompt_tokens is None:

281 + return None

282 +

283 + if truncate_prompt_tokens == 0:

284 + return 0

285 +

286 + if truncate_prompt_tokens < 0:

287 + truncate_prompt_tokens = self.model_config.max_model_len

288 +

289 + if max_length is not None and truncate_prompt_tokens > max_length:

290 + raise ValueError(

291 + f"truncate_prompt_tokens ({truncate_prompt_tokens}) "

292 + f"cannot be greater than max_length ({max_length}). "

293 + f"Please select a smaller truncation size.")

294 +

295 + return truncate_prompt_tokens

296

297 def _maybe_apply_truncation(

298 self, token_ids: list[int],

@@ -186,7 +334,29 @@ class CompletionRenderer(BaseRenderer):

334 max_length=truncate_prompt_tokens)

335

336 return self._create_tokens_prompt(encoded.input_ids, max_length,

189 - cache_salt)

337 + cache_salt, text)

338 +

339 + async def _maybe_detokenize(

340 + self,

341 + token_ids: list[int],

342 + max_length: Optional[int],

343 + truncate_prompt_tokens: Optional[int],

344 + cache_salt: Optional[str],

345 + needs_detokenization: Optional[bool] = False,

346 + ) -> EngineTokensPrompt:

347 + """Optionally detokenize token IDs and build a tokens prompt."""

348 + token_ids = self._maybe_apply_truncation(token_ids,

349 + truncate_prompt_tokens)

350 +

351 + prompt = None

352 + if needs_detokenization is True:

353 + async_tokenizer = self._get_async_tokenizer()

354 + prompt = await async_tokenizer.decode(token_ids)

355 +

356 + return self._create_tokens_prompt(token_ids=token_ids,

357 + max_length=max_length,

358 + cache_salt=cache_salt,

359 + prompt=prompt)

360

361 def _get_async_tokenizer(self) -> AsyncMicrobatchTokenizer:

362 """Get or create async tokenizer using shared pool."""

@@ -210,6 +380,7 @@ class CompletionRenderer(BaseRenderer):

380 token_ids: list[int],

381 max_length: Optional[int] = None,

382 cache_salt: Optional[str] = None,

383 + prompt: Optional[str] = None,

384 ) -> EngineTokensPrompt:

385 """Create validated EngineTokensPrompt."""

386 if max_length is not None and len(token_ids) > max_length:

@@ -221,4 +392,6 @@ class CompletionRenderer(BaseRenderer):

392 tokens_prompt = EngineTokensPrompt(prompt_token_ids=token_ids)

393 if cache_salt is not None:

394 tokens_prompt["cache_salt"] = cache_salt

224 - return tokens_prompt

395 + if prompt is not None:

396 + tokens_prompt["prompt"] = prompt

397 + return tokens_prompt

\ No newline at end of file

vllm/inputs/data.py

@@ -52,6 +52,9 @@ class TokensPrompt(TypedDict):

52 prompt_token_ids: list[int]

53 """A list of token IDs to pass to the model."""

55 + prompt: NotRequired[str]

56 + """The prompt text corresponding to the token IDs, if available."""

57 +

58 token_type_ids: NotRequired[list[int]]

59 """A list of token type IDs to pass to the cross encoder model."""