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
Add renderer-based prompt processing for embedding and classification endpoints (#24356)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
Flora Feng committed
Sep 7, 2025 at 01:26 UTC
0661cb9df390a47dda0a5c201cf712dd2d943666
6 files changed
+59
-56
tests/entrypoints/openai/test_truncation.py
+4
-10
@@ -73,17 +73,11 @@ async def test_zero_truncation_size(client: openai.AsyncOpenAI):
73
"truncate_prompt_tokens": truncation_size
74
}
75
76
- with pytest.raises(openai.BadRequestError) as err:
77
- await client.post(path="embeddings", cast_to=object, body={**kwargs})
78
-
79
- assert err.value.status_code == 400
80
- error_details = err.value.response.json()["error"]
76
+ response = await client.post(path="embeddings",
77
+ cast_to=object,
78
+ body={**kwargs})
79
82
- assert error_details["type"] == "BadRequestError"
83
- assert "This model's maximum context length is" in error_details["message"]
84
- assert "tokens in the input for embedding generation" in error_details[
85
- "message"]
86
- assert "Please reduce the length of the input" in error_details["message"]
80
+ assert response["usage"]["prompt_tokens"] == truncation_size
81
82
83
@pytest.mark.asyncio
tests/entrypoints/test_renderer.py
+17
@@ -130,6 +130,23 @@ class TestRenderPrompt:
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_truncation_negative(self, renderer, mock_async_tokenizer):
135
+ # Test that negative truncation uses model's max_model_len
136
+ mock_async_tokenizer.return_value = MockTokenizerResult(
137
+ [101, 7592, 2088]) # Truncated to max_model_len
138
+ renderer.async_tokenizer_pool[
139
+ renderer.tokenizer] = mock_async_tokenizer
140
+
141
+ results = await renderer.render_prompt(prompt_or_prompts="Hello world",
142
+ max_length=200,
143
+ truncate_prompt_tokens=-1)
144
+
145
+ assert len(results) == 1
146
+ call_args = mock_async_tokenizer.call_args
147
+ assert call_args.kwargs["truncation"] is True
148
+ assert call_args.kwargs["max_length"] == 100 # model's max_model_len
149
+
150
@pytest.mark.asyncio
151
async def test_token_truncation_last_elements(self, renderer):
152
# Test that token truncation keeps the last N elements
vllm/entrypoints/openai/serving_classification.py
+5
-8
@@ -54,14 +54,11 @@ class ClassificationMixin(OpenAIServing):
54
ctx.tokenizer = await self.engine_client.get_tokenizer(
55
ctx.lora_request)
56
57
- (
58
- ctx.request_prompts,
59
- ctx.engine_prompts,
60
- ) = await self._preprocess_completion(
61
- ctx.request,
62
- ctx.tokenizer,
63
- ctx.request.input,
64
- )
57
+ renderer = self._get_renderer(ctx.tokenizer)
58
+ ctx.engine_prompts = await renderer.render_prompt(
59
+ prompt_or_prompts=ctx.request.input,
60
+ max_length=self.max_model_len,
61
+ truncate_prompt_tokens=ctx.request.truncate_prompt_tokens)
62
63
return None
64
vllm/entrypoints/openai/serving_embedding.py
+19
-26
@@ -24,7 +24,6 @@ from vllm.entrypoints.openai.protocol import (EmbeddingChatRequest,
24
ErrorResponse, UsageInfo)
25
from vllm.entrypoints.openai.serving_engine import (EmbeddingServeContext,
26
OpenAIServing,
27
- RequestPrompt,
27
ServeContext,
28
TextTokensPrompt)
29
# yapf: enable
@@ -79,11 +78,12 @@ class EmbeddingMixin(OpenAIServing):
78
79
tokenizer = await self.engine_client.get_tokenizer(ctx.lora_request
80
)
81
+ renderer = self._get_renderer(tokenizer)
82
83
if isinstance(ctx.request, EmbeddingChatRequest):
84
(
85
_,
86
- ctx.request_prompts,
86
+ _,
87
ctx.engine_prompts,
88
) = await self._preprocess_chat(
89
ctx.request,
@@ -98,13 +98,18 @@ class EmbeddingMixin(OpenAIServing):
98
add_special_tokens=ctx.request.add_special_tokens,
99
)
100
else:
101
- (ctx.request_prompts,
102
- ctx.engine_prompts) = await self._preprocess_completion(
103
- ctx.request,
104
- tokenizer,
105
- ctx.request.input,
106
- add_special_tokens=ctx.request.add_special_tokens,
107
- )
101
+ # Set max_length based on chunked processing capability
102
+ if self._should_use_chunked_processing(ctx.request):
103
+ max_length = None
104
+ else:
105
+ max_length = self.max_embed_len or self.max_model_len
106
+
107
+ ctx.engine_prompts = await renderer.render_prompt(
108
+ prompt_or_prompts=ctx.request.input,
109
+ max_length=max_length,
110
+ truncate_prompt_tokens=ctx.request.truncate_prompt_tokens,
111
+ add_special_tokens=ctx.request.add_special_tokens,
112
+ )
113
return None
114
except (ValueError, TypeError) as e:
115
logger.exception("Error in preprocessing prompt inputs")
@@ -286,7 +291,6 @@ class EmbeddingMixin(OpenAIServing):
291
self,
292
ctx: EmbeddingServeContext,
293
engine_prompt: Union[EngineTokensPrompt, EngineEmbedsPrompt],
289
- request_prompt: RequestPrompt,
294
pooling_params: PoolingParams,
295
trace_headers: Optional[Mapping[str, str]],
296
prompt_index: int,
@@ -295,7 +299,7 @@ class EmbeddingMixin(OpenAIServing):
299
request_id_item = f"{ctx.request_id}-{prompt_index}"
300
301
self._log_inputs(request_id_item,
298
- request_prompt,
302
+ engine_prompt,
303
params=pooling_params,
304
lora_request=ctx.lora_request)
305
@@ -353,20 +357,14 @@ class EmbeddingMixin(OpenAIServing):
357
return self.create_error_response(
358
"Engine prompts not available")
359
356
- if ctx.request_prompts is None:
357
- return self.create_error_response(
358
- "Request prompts not available")
359
-
360
max_pos_embeddings = self._get_max_position_embeddings()
361
362
for i, engine_prompt in enumerate(ctx.engine_prompts):
363
- request_prompt = ctx.request_prompts[i]
364
-
363
# Check if this specific prompt needs chunked processing
366
- if self._is_text_tokens_prompt(request_prompt):
364
+ if self._is_text_tokens_prompt(engine_prompt):
365
# Cast to TextTokensPrompt since we've verified
366
# prompt_token_ids
369
- text_tokens_prompt = cast(TextTokensPrompt, request_prompt)
367
+ text_tokens_prompt = cast(TextTokensPrompt, engine_prompt)
368
if (len(text_tokens_prompt["prompt_token_ids"])
369
> max_pos_embeddings):
370
# Use chunked processing for this prompt
@@ -382,8 +380,7 @@ class EmbeddingMixin(OpenAIServing):
380
Union[EngineTokensPrompt, EngineEmbedsPrompt],
381
engine_prompt)
382
generator = await self._create_single_prompt_generator(
385
- ctx, engine_prompt_typed, request_prompt, pooling_params,
386
- trace_headers, i)
383
+ ctx, engine_prompt_typed, pooling_params, trace_headers, i)
384
generators.append(generator)
385
386
from vllm.utils import merge_async_iterators
@@ -419,10 +416,6 @@ class EmbeddingMixin(OpenAIServing):
416
if not use_chunked:
417
return await super()._collect_batch(ctx=ctx)
418
422
- if ctx.request_prompts is None:
423
- return self.create_error_response(
424
- "Request prompts not available")
425
-
419
if ctx.result_generator is None:
420
return self.create_error_response(
421
"Result generator not available")
@@ -538,7 +531,7 @@ class EmbeddingMixin(OpenAIServing):
531
data=final_embedding)
532
533
# Get original prompt token IDs for this prompt
541
- original_prompt = ctx.request_prompts[prompt_idx]
534
+ original_prompt = ctx.engine_prompts[prompt_idx]
535
if not self._is_text_tokens_prompt(original_prompt):
536
return self.create_error_response(
537
f"Chunked prompt {prompt_idx} is not a "
vllm/entrypoints/openai/serving_engine.py
+7
-10
@@ -368,23 +368,20 @@ class OpenAIServing:
368
for i, engine_prompt in enumerate(ctx.engine_prompts):
369
request_id_item = f"{ctx.request_id}-{i}"
370
371
- if ctx.request_prompts is None:
372
- return self.create_error_response(
373
- "Request prompts not available")
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
378
self._log_inputs(
379
request_id_item,
377
- ctx.request_prompts[i],
380
+ engine_prompt,
381
params=pooling_params,
382
lora_request=ctx.lora_request,
383
)
384
382
- # Mypy has an existing bug related to inferring the variance of
383
- # TypedDicts with `builtins.enumerate`:
384
- # https://github.com/python/mypy/issues/8586#issuecomment-2867698435
385
- engine_prompt = cast(
386
- Union[EngineTokensPrompt, EngineEmbedsPrompt],
387
- engine_prompt)
385
generator = self.engine_client.encode(
386
engine_prompt,
387
pooling_params,
vllm/entrypoints/renderer.py
+7
-2
@@ -108,10 +108,15 @@ class CompletionRenderer(BaseRenderer):
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
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.")
120
121
# Parse and batch the input prompts
122
batch_inputs = parse_and_batch_prompt(prompt_or_prompts)