@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 712b273f65
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

[Refactor] Introduce basic Renderer for completion-style request (#24010)

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

Flora Feng committed Sep 3, 2025 at 22:21 UTC 712b273f655fe007e4efbaae00f1a343cfeb0742
5 files changed +416 -27
tests/entrypoints/test_renderer.py new
+163
@@ -0,0 +1,163 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from dataclasses import dataclass
5 +from typing import Optional
6 +from unittest.mock import AsyncMock, MagicMock
7 +
8 +import pytest
9 +
10 +from vllm.entrypoints.renderer import CompletionRenderer
11 +
12 +
13 +@dataclass
14 +class MockModelConfig:
15 + max_model_len: int = 100
16 + encoder_config: Optional[dict] = None
17 +
18 +
19 +class MockTokenizerResult:
20 +
21 + def __init__(self, input_ids):
22 + self.input_ids = input_ids
23 +
24 +
25 +@pytest.fixture
26 +def mock_model_config():
27 + return MockModelConfig()
28 +
29 +
30 +@pytest.fixture
31 +def mock_tokenizer():
32 + tokenizer = MagicMock()
33 + return tokenizer
34 +
35 +
36 +@pytest.fixture
37 +def mock_async_tokenizer():
38 + async_tokenizer = AsyncMock()
39 + return async_tokenizer
40 +
41 +
42 +@pytest.fixture
43 +def renderer(mock_model_config, mock_tokenizer):
44 + return CompletionRenderer(model_config=mock_model_config,
45 + tokenizer=mock_tokenizer,
46 + async_tokenizer_pool={})
47 +
48 +
49 +class TestRenderPrompt:
50 + """Test Category A: Basic Functionality Tests"""
51 +
52 + @pytest.mark.asyncio
53 + async def test_token_input(self, renderer):
54 + tokens = [101, 7592, 2088]
55 + results = await renderer.render_prompt(prompt_or_prompts=tokens,
56 + max_length=100)
57 +
58 + assert len(results) == 1
59 + assert results[0]["prompt_token_ids"] == tokens
60 +
61 + @pytest.mark.asyncio
62 + async def test_token_list_input(self, renderer):
63 + token_lists = [[101, 7592, 2088], [102, 1234, 5678, 9012], [103, 4567]]
64 + results = await renderer.render_prompt(prompt_or_prompts=token_lists,
65 + max_length=100)
66 +
67 + assert len(results) == 3
68 + assert results[0]["prompt_token_ids"] == [101, 7592, 2088]
69 + assert results[1]["prompt_token_ids"] == [102, 1234, 5678, 9012]
70 + assert results[2]["prompt_token_ids"] == [103, 4567]
71 +
72 + @pytest.mark.asyncio
73 + async def test_text_input(self, renderer, mock_async_tokenizer):
74 + mock_async_tokenizer.return_value = MockTokenizerResult(
75 + [101, 7592, 2088])
76 + renderer.async_tokenizer_pool[
77 + renderer.tokenizer] = mock_async_tokenizer
78 +
79 + results = await renderer.render_prompt(prompt_or_prompts="Hello world",
80 + max_length=100)
81 +
82 + assert len(results) == 1
83 + assert results[0]["prompt_token_ids"] == [101, 7592, 2088]
84 + mock_async_tokenizer.assert_called_once()
85 +
86 + @pytest.mark.asyncio
87 + async def test_text_list_input(self, renderer, mock_async_tokenizer):
88 + mock_async_tokenizer.return_value = MockTokenizerResult(
89 + [101, 7592, 2088])
90 + renderer.async_tokenizer_pool[
91 + renderer.tokenizer] = mock_async_tokenizer
92 +
93 + text_list_input = ["Hello world", "How are you?", "Good morning"]
94 + results = await renderer.render_prompt(
95 + prompt_or_prompts=text_list_input, max_length=100)
96 +
97 + assert len(results) == 3
98 + for result in results:
99 + assert result["prompt_token_ids"] == [101, 7592, 2088]
100 + assert mock_async_tokenizer.call_count == 3
101 +
102 + @pytest.mark.asyncio
103 + async def test_no_truncation(self, renderer, mock_async_tokenizer):
104 + mock_async_tokenizer.return_value = MockTokenizerResult(
105 + [101, 7592, 2088])
106 + renderer.async_tokenizer_pool[
107 + renderer.tokenizer] = mock_async_tokenizer
108 +
109 + results = await renderer.render_prompt(prompt_or_prompts="Hello world",
110 + max_length=100)
111 +
112 + assert len(results) == 1
113 + call_args = mock_async_tokenizer.call_args
114 + assert "truncation" not in call_args.kwargs or call_args.kwargs[
115 + "truncation"] is False
116 +
117 + @pytest.mark.asyncio
118 + async def test_truncation_positive(self, renderer, mock_async_tokenizer):
119 + mock_async_tokenizer.return_value = MockTokenizerResult(
120 + [101, 7592, 2088]) # Truncated
121 + renderer.async_tokenizer_pool[
122 + renderer.tokenizer] = mock_async_tokenizer
123 +
124 + results = await renderer.render_prompt(prompt_or_prompts="Hello world",
125 + max_length=100,
126 + truncate_prompt_tokens=50)
127 +
128 + assert len(results) == 1
129 + call_args = mock_async_tokenizer.call_args
130 + assert call_args.kwargs["truncation"] is True
131 + assert call_args.kwargs["max_length"] == 50
132 +
133 + @pytest.mark.asyncio
134 + async def test_token_truncation_last_elements(self, renderer):
135 + # Test that token truncation keeps the last N elements
136 + long_tokens = [100, 101, 102, 103, 104, 105, 106, 107, 108,
137 + 109] # 10 tokens
138 + results = await renderer.render_prompt(prompt_or_prompts=long_tokens,
139 + max_length=100,
140 + truncate_prompt_tokens=5)
141 +
142 + assert len(results) == 1
143 + # Should keep the last 5 tokens: [105, 106, 107, 108, 109]
144 + assert results[0]["prompt_token_ids"] == [105, 106, 107, 108, 109]
145 +
146 + @pytest.mark.asyncio
147 + async def test_max_length_exceeded(self, renderer):
148 + long_tokens = list(range(150)) # Exceeds max_model_len=100
149 +
150 + with pytest.raises(ValueError, match="maximum context length"):
151 + await renderer.render_prompt(prompt_or_prompts=long_tokens,
152 + max_length=100)
153 +
154 + @pytest.mark.asyncio
155 + async def test_no_tokenizer_for_text(self, mock_model_config):
156 + renderer_no_tokenizer = CompletionRenderer(
157 + model_config=mock_model_config,
158 + tokenizer=None,
159 + async_tokenizer_pool={})
160 +
161 + with pytest.raises(ValueError, match="No tokenizer available"):
162 + await renderer_no_tokenizer.render_prompt(
163 + prompt_or_prompts="Hello world", max_length=100)
vllm/entrypoints/openai/serving_engine.py
+15 -5
@@ -62,8 +62,10 @@ from vllm.entrypoints.openai.protocol import (ChatCompletionRequest,
62 TranslationRequest)
63 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
64 from vllm.entrypoints.openai.tool_parsers import ToolParser
65 +from vllm.entrypoints.renderer import BaseRenderer, CompletionRenderer
66 # yapf: enable
67 from vllm.inputs.data import EmbedsPrompt as EngineEmbedsPrompt
68 +from vllm.inputs.data import PromptType
69 from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
70 from vllm.inputs.parse import parse_and_batch_prompt
71 from vllm.logger import init_logger
@@ -243,6 +245,16 @@ class OpenAIServing:
245 AsyncMicrobatchTokenizer] = {}
246 self.log_error_stack = log_error_stack
247
248 + def _get_renderer(self, tokenizer: Optional[AnyTokenizer]) -> BaseRenderer:
249 + """
250 + Get a Renderer instance with the provided tokenizer.
251 + Uses shared async tokenizer pool for efficiency.
252 + """
253 + return CompletionRenderer(
254 + model_config=self.model_config,
255 + tokenizer=tokenizer,
256 + async_tokenizer_pool=self._async_tokenizer_pool)
257 +
258 def _get_async_tokenizer(self, tokenizer) -> AsyncMicrobatchTokenizer:
259 """
260 Return (and cache) an `AsyncMicrobatchTokenizer` bound to the
@@ -1098,7 +1110,7 @@ class OpenAIServing:
1110 def _log_inputs(
1111 self,
1112 request_id: str,
1101 - inputs: RequestPrompt,
1113 + inputs: Union[RequestPrompt, PromptType],
1114 params: Optional[Union[SamplingParams, PoolingParams,
1115 BeamSearchParams]],
1116 lora_request: Optional[LoRARequest],
@@ -1110,11 +1122,9 @@ class OpenAIServing:
1122 prompt = inputs
1123 elif isinstance(inputs, list):
1124 prompt_token_ids = inputs
1113 - elif "prompt_embeds" in inputs:
1114 - prompt_embeds = inputs.get("prompt_embeds")
1125 else:
1116 - prompt = inputs["prompt"]
1117 - prompt_token_ids = inputs["prompt_token_ids"]
1126 + prompt = getattr(inputs, 'prompt', None)
1127 + prompt_token_ids = getattr(inputs, 'prompt_token_ids', None)
1128
1129 self.request_logger.log_inputs(
1130 request_id,
vllm/entrypoints/openai/serving_pooling.py
+12 -14
@@ -4,7 +4,7 @@
4 import asyncio
5 import base64
6 import time
7 -from collections.abc import AsyncGenerator, Sequence
7 +from collections.abc import AsyncGenerator
8 from typing import Final, Literal, Optional, Union, cast
9
10 import jinja2
@@ -26,7 +26,7 @@ from vllm.entrypoints.openai.protocol import (ErrorResponse,
26 PoolingRequest, PoolingResponse,
27 PoolingResponseData, UsageInfo)
28 # yapf: enable
29 -from vllm.entrypoints.openai.serving_engine import OpenAIServing, RequestPrompt
29 +from vllm.entrypoints.openai.serving_engine import OpenAIServing
30 from vllm.entrypoints.openai.serving_models import OpenAIServingModels
31 from vllm.entrypoints.utils import _validate_truncation_size
32 from vllm.logger import init_logger
@@ -104,6 +104,7 @@ class OpenAIServingPooling(OpenAIServing):
104 else:
105 tokenizer = await self.engine_client.get_tokenizer(lora_request
106 )
107 + renderer = self._get_renderer(tokenizer)
108
109 if getattr(request, "dimensions", None) is not None:
110 return self.create_error_response(
@@ -126,14 +127,11 @@ class OpenAIServingPooling(OpenAIServing):
127
128 engine_prompts = await self.io_processor.pre_process_async(
129 prompt=validated_prompt, request_id=request_id)
129 - request_prompts: Sequence[RequestPrompt] = [
130 - ""
131 - ] * len(engine_prompts)
130
131 elif isinstance(request, PoolingChatRequest):
132 (
133 _,
136 - request_prompts,
134 + _,
135 engine_prompts,
136 ) = await self._preprocess_chat(
137 request,
@@ -149,13 +147,13 @@ class OpenAIServingPooling(OpenAIServing):
147 add_special_tokens=request.add_special_tokens,
148 )
149 elif isinstance(request, PoolingCompletionRequest):
152 - (request_prompts,
153 - engine_prompts) = await self._preprocess_completion(
154 - request,
155 - tokenizer,
156 - request.input,
157 - add_special_tokens=request.add_special_tokens,
158 - )
150 + engine_prompts = await renderer.render_prompt(
151 + prompt_or_prompts=request.input,
152 + max_length=self.max_model_len,
153 + truncate_prompt_tokens=truncate_prompt_tokens,
154 + add_special_tokens=request.add_special_tokens,
155 + cache_salt=getattr(request, 'cache_salt', None),
156 + )
157 else:
158 raise ValueError(
159 f"Unsupported request of type {type(request)}")
@@ -177,7 +175,7 @@ class OpenAIServingPooling(OpenAIServing):
175 request_id_item = f"{request_id}-{i}"
176
177 self._log_inputs(request_id_item,
180 - request_prompts[i],
178 + engine_prompt,
179 params=pooling_params,
180 lora_request=lora_request)
181
vllm/entrypoints/openai/serving_tokenization.py
+7 -8
@@ -65,6 +65,7 @@ class OpenAIServingTokenization(OpenAIServing):
65 lora_request = self._maybe_get_adapters(request)
66
67 tokenizer = await self.engine_client.get_tokenizer(lora_request)
68 + renderer = self._get_renderer(tokenizer)
69
70 if isinstance(request, TokenizeChatRequest):
71 tool_dicts = (None if request.tools is None else
@@ -87,13 +88,11 @@ class OpenAIServingTokenization(OpenAIServing):
88 add_special_tokens=request.add_special_tokens,
89 )
90 else:
90 - (request_prompts,
91 - engine_prompts) = await self._preprocess_completion(
92 - request,
93 - tokenizer,
94 - request.prompt,
95 - add_special_tokens=request.add_special_tokens,
96 - )
91 + engine_prompts = await renderer.render_prompt(
92 + prompt_or_prompts=request.prompt,
93 + add_special_tokens=request.add_special_tokens,
94 + cache_salt=getattr(request, 'cache_salt', None),
95 + )
96 except (ValueError, TypeError, jinja2.TemplateError) as e:
97 logger.exception("Error in preprocessing prompt inputs")
98 return self.create_error_response(f"{e} {e.__cause__}")
@@ -101,7 +100,7 @@ class OpenAIServingTokenization(OpenAIServing):
100 input_ids: list[int] = []
101 for i, engine_prompt in enumerate(engine_prompts):
102 self._log_inputs(request_id,
104 - request_prompts[i],
103 + engine_prompt,
104 params=None,
105 lora_request=lora_request)
106
vllm/entrypoints/renderer.py new
+219
@@ -0,0 +1,219 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import asyncio
5 +from abc import ABC, abstractmethod
6 +from typing import Annotated, Optional, Union
7 +
8 +from pydantic import Field
9 +
10 +from vllm.config import ModelConfig
11 +from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
12 +from vllm.inputs.parse import parse_and_batch_prompt
13 +from vllm.transformers_utils.tokenizer import AnyTokenizer
14 +from vllm.utils import AsyncMicrobatchTokenizer
15 +
16 +
17 +class BaseRenderer(ABC):
18 + """
19 + Base class for unified input processing and rendering.
20 +
21 + The Renderer serves as a unified input processor that consolidates
22 + tokenization, chat template formatting, and multimodal input handling
23 + into a single component.
24 + It converts high-level API requests (OpenAI-style JSON) into token IDs and
25 + multimodal features ready for engine consumption.
26 +
27 + Key responsibilities:
28 + - Convert text prompts to token sequences with proper special tokens
29 + - Apply chat templates and format conversations
30 + - Handle multimodal inputs (images, audio, etc.) when applicable
31 + - Manage prompt truncation and length validation
32 + - Provide clean separation between API layer and engine core
33 + """
34 +
35 + def __init__(
36 + self,
37 + model_config: ModelConfig,
38 + tokenizer: Optional[AnyTokenizer] = None,
39 + ):
40 + super().__init__()
41 + self.model_config = model_config
42 + self.tokenizer = tokenizer
43 +
44 + @abstractmethod
45 + async def render_prompt(
46 + self,
47 + prompt_or_prompts: Union[str, list[str], list[int], list[list[int]]],
48 + max_length: Optional[int] = None,
49 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,
50 + add_special_tokens: Optional[bool] = True,
51 + cache_salt: Optional[str] = None,
52 + ) -> list[EngineTokensPrompt]:
53 + """
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 +
61 + 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 +
74 + Returns:
75 + list[EngineTokensPrompt]: Tokenized prompts ready for engine
76 + consumption
77 +
78 + Raises:
79 + ValueError: If input format is invalid or length limits exceeded
80 + """
81 + raise NotImplementedError
82 +
83 +
84 +class CompletionRenderer(BaseRenderer):
85 +
86 + def __init__(
87 + self,
88 + model_config: ModelConfig,
89 + tokenizer: Optional[AnyTokenizer] = None,
90 + async_tokenizer_pool: Optional[dict[AnyTokenizer,
91 + AsyncMicrobatchTokenizer]] = None,
92 + ):
93 + super().__init__(model_config, tokenizer)
94 + self.async_tokenizer_pool = async_tokenizer_pool or {}
95 + self.async_tokenizer: Optional[AsyncMicrobatchTokenizer] = None
96 +
97 + async def render_prompt(
98 + self,
99 + prompt_or_prompts: Union[str, list[str], list[int], list[list[int]]],
100 + max_length: Optional[int] = None,
101 + truncate_prompt_tokens: Optional[Annotated[int, Field(ge=-1)]] = None,
102 + add_special_tokens: Optional[bool] = True,
103 + cache_salt: Optional[str] = None,
104 + ) -> list[EngineTokensPrompt]:
105 + """Implementation of prompt rendering for completion-style requests.
106 +
107 + Uses async tokenizer pooling for improved performance. See base class
108 + for detailed parameter documentation.
109 + """
110 + if truncate_prompt_tokens is not None:
111 + if max_length is not None:
112 + assert 0 <= truncate_prompt_tokens <= max_length
113 + if truncate_prompt_tokens == 0:
114 + return []
115 +
116 + # Parse and batch the input prompts
117 + batch_inputs = parse_and_batch_prompt(prompt_or_prompts)
118 +
119 + rendered_prompts: list[EngineTokensPrompt] = []
120 + tokenize_tasks = []
121 + for prompt_input in batch_inputs:
122 + if prompt_input["is_tokens"] is True:
123 + # Token input
124 + token_ids = self._maybe_apply_truncation(
125 + prompt_input["content"], truncate_prompt_tokens)
126 + rendered_prompts.append(
127 + self._create_tokens_prompt(token_ids, max_length,
128 + cache_salt))
129 + else:
130 + # Text input
131 + tokenize_task = asyncio.create_task(
132 + self._tokenize(prompt_input["content"], max_length,
133 + truncate_prompt_tokens, add_special_tokens,
134 + cache_salt))
135 + tokenize_tasks.append(tokenize_task)
136 +
137 + # Wait for all text tokenization to finish
138 + if tokenize_tasks:
139 + tokenized_text_prompts = await asyncio.gather(*tokenize_tasks)
140 + rendered_prompts.extend(tokenized_text_prompts)
141 +
142 + return rendered_prompts
143 +
144 + def _maybe_apply_truncation(
145 + self, token_ids: list[int],
146 + truncate_prompt_tokens: Optional[int]) -> list[int]:
147 + """Apply truncation to token sequence."""
148 + if truncate_prompt_tokens is None:
149 + return token_ids
150 + if truncate_prompt_tokens >= len(token_ids):
151 + return token_ids
152 +
153 + return token_ids[-truncate_prompt_tokens:]
154 +
155 + async def _tokenize(
156 + self,
157 + text: str,
158 + max_length: Optional[int],
159 + truncate_prompt_tokens: Optional[int],
160 + add_special_tokens: Optional[bool],
161 + cache_salt: Optional[str],
162 + ) -> EngineTokensPrompt:
163 + """Tokenize text input asynchronously."""
164 + async_tokenizer = self._get_async_tokenizer()
165 +
166 + # Handle encoder-specific preprocessing
167 + if (self.model_config.encoder_config is not None
168 + and self.model_config.encoder_config.get(
169 + "do_lower_case", False)):
170 + text = text.lower()
171 +
172 + # Tokenize texts
173 + if truncate_prompt_tokens is None:
174 + encoded = await async_tokenizer(
175 + text, add_special_tokens=add_special_tokens)
176 + else:
177 + encoded = await async_tokenizer(
178 + text,
179 + add_special_tokens=add_special_tokens,
180 + truncation=True,
181 + max_length=truncate_prompt_tokens)
182 +
183 + return self._create_tokens_prompt(encoded.input_ids, max_length,
184 + cache_salt)
185 +
186 + def _get_async_tokenizer(self) -> AsyncMicrobatchTokenizer:
187 + """Get or create async tokenizer using shared pool."""
188 + if self.async_tokenizer is not None:
189 + return self.async_tokenizer
190 + if self.tokenizer is None:
191 + raise ValueError(
192 + "No tokenizer available for text input processing")
193 +
194 + # Check shared pool first
195 + if self.tokenizer in self.async_tokenizer_pool:
196 + return self.async_tokenizer_pool[self.tokenizer]
197 +
198 + # Create new async tokenizer and add to pool
199 + self.async_tokenizer = AsyncMicrobatchTokenizer(self.tokenizer)
200 + self.async_tokenizer_pool[self.tokenizer] = self.async_tokenizer
201 + return self.async_tokenizer
202 +
203 + def _create_tokens_prompt(
204 + self,
205 + token_ids: list[int],
206 + max_length: Optional[int] = None,
207 + cache_salt: Optional[str] = None,
208 + ) -> EngineTokensPrompt:
209 + """Create validated EngineTokensPrompt."""
210 + if max_length is not None and len(token_ids) > max_length:
211 + raise ValueError(
212 + f"This maximum context length is {max_length} tokens. "
213 + f"However, your request has {len(token_ids)} input tokens. "
214 + "Please reduce the length of the input messages.")
215 +
216 + tokens_prompt = EngineTokensPrompt(prompt_token_ids=token_ids)
217 + if cache_salt is not None:
218 + tokens_prompt["cache_salt"] = cache_salt
219 + return tokens_prompt