@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 328cbb2773
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 9, 19:30 UTC

[Frontend][2/n] Make pooling entrypoints request schema consensus | ChatRequest (#32574)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>

wang.yuqi committed Jan 22, 2026 at 18:32 UTC 328cbb2773d93d45d18dfd383d631de4fa276e69
24 files changed +456 -205
docs/design/io_processor_plugins.md
+1 -1
@@ -79,7 +79,7 @@ The `post_process*` methods take `PoolingRequestOutput` objects as input and gen
79 The `validate_or_generate_params` method is used for validating with the plugin any `SamplingParameters`/`PoolingParameters` received with the user request, or to generate new ones if none are specified. The function always returns the validated/generated parameters.
80 The `output_to_response` method is used only for online serving and converts the plugin output to the `IOProcessorResponse` type that is then returned by the API Server. The implementation of the `/pooling` serving endpoint is available here [vllm/entrypoints/openai/serving_pooling.py](../../vllm/entrypoints/pooling/pooling/serving.py).
81
82 -An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/IBM/terratorch/tree/main/terratorch/vllm/plugins/segmentation). Please, also refer to our online ([examples/pooling/plugin/prithvi_geospatial_mae_client.py](../../examples/pooling/plugin/prithvi_geospatial_mae_client.py)) and offline ([examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py](../../examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py)) inference examples.
82 +An example implementation of a plugin that enables generating geotiff images with the PrithviGeospatialMAE model is available [here](https://github.com/IBM/terratorch/tree/main/terratorch/vllm/plugins/segmentation). Please, also refer to our online ([examples/pooling/plugin/prithvi_geospatial_mae_online.py](../../examples/pooling/plugin/prithvi_geospatial_mae_online.py)) and offline ([examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py](../../examples/pooling/plugin/prithvi_geospatial_mae_io_processor.py)) inference examples.
83
84 ## Using an IO Processor plugin
85
docs/models/pooling_models.md
+2 -2
@@ -273,7 +273,7 @@ outputs = llm.embed(
273 print(outputs[0].outputs)
274 ```
275
276 -A code example can be found here: [examples/pooling/embed/embed_matryoshka_fy.py](../../examples/pooling/embed/embed_matryoshka_fy.py)
276 +A code example can be found here: [examples/pooling/embed/embed_matryoshka_fy_offline.py](../../examples/pooling/embed/embed_matryoshka_fy_offline.py)
277
278 ### Online Inference
279
@@ -303,7 +303,7 @@ Expected output:
303 {"id":"embd-5c21fc9a5c9d4384a1b021daccaf9f64","object":"list","created":1745476417,"model":"jinaai/jina-embeddings-v3","data":[{"index":0,"object":"embedding","embedding":[-0.3828125,-0.1357421875,0.03759765625,0.125,0.21875,0.09521484375,-0.003662109375,0.1591796875,-0.130859375,-0.0869140625,-0.1982421875,0.1689453125,-0.220703125,0.1728515625,-0.2275390625,-0.0712890625,-0.162109375,-0.283203125,-0.055419921875,-0.0693359375,0.031982421875,-0.04052734375,-0.2734375,0.1826171875,-0.091796875,0.220703125,0.37890625,-0.0888671875,-0.12890625,-0.021484375,-0.0091552734375,0.23046875]}],"usage":{"prompt_tokens":8,"total_tokens":8,"completion_tokens":0,"prompt_tokens_details":null}}
304 ```
305
306 -An OpenAI client example can be found here: [examples/pooling/embed/openai_embedding_matryoshka_fy.py](../../examples/pooling/embed/openai_embedding_matryoshka_fy.py)
306 +An OpenAI client example can be found here: [examples/pooling/embed/openai_embedding_matryoshka_fy_client.py](../../examples/pooling/embed/openai_embedding_matryoshka_fy_client.py)
307
308 ## Deprecated Features
309
docs/models/supported_models.md
+1 -1
@@ -619,7 +619,7 @@ These models primarily support the [`LLM.encode`](./pooling_models.md#llmencode)
619 | `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` | | |
620
621 !!! note
622 - Named Entity Recognition (NER) usage, please refer to [examples/pooling/token_classify/ner.py](../../examples/pooling/token_classify/ner.py), [examples/pooling/token_classify/ner_client.py](../../examples/pooling/token_classify/ner_client.py).
622 + Named Entity Recognition (NER) usage, please refer to [examples/pooling/token_classify/ner_offline.py](../../examples/pooling/token_classify/ner_offline.py), [examples/pooling/token_classify/ner_online.py](../../examples/pooling/token_classify/ner_online.py).
623
624 ## List of Multimodal Language Models
625
docs/serving/openai_compatible_server.md
+1 -1
@@ -551,7 +551,7 @@ Our Pooling API encodes input prompts using a [pooling model](../models/pooling_
551
552 The input format is the same as [Embeddings API](#embeddings-api), but the output data can contain an arbitrary nested list, not just a 1-D list of floats.
553
554 -Code example: [examples/pooling/pooling/openai_pooling_client.py](../../examples/pooling/pooling/openai_pooling_client.py)
554 +Code example: [examples/pooling/pooling/pooling_online.py](../../examples/pooling/pooling/pooling_online.py)
555
556 ### Classification API
557
examples/pooling/embed/embed_jina_embeddings_v3_offline.py renamed
examples/pooling/embed/embed_matryoshka_fy_offline.py renamed
examples/pooling/embed/embedding_requests_base64_online.py renamed
+18 -12
@@ -26,36 +26,42 @@ def post_http_request(prompt: dict, api_url: str) -> requests.Response:
26
27
28 def parse_args():
29 - parser = argparse.ArgumentParser()
30 - parser.add_argument("--host", type=str, default="localhost")
31 - parser.add_argument("--port", type=int, default=8000)
32 - parser.add_argument("--model", type=str, default="intfloat/e5-small")
33 -
34 - return parser.parse_args()
29 + parse = argparse.ArgumentParser()
30 + parse.add_argument("--host", type=str, default="localhost")
31 + parse.add_argument("--port", type=int, default=8000)
32 + return parse.parse_args()
33
34
35 def main(args):
38 - api_url = f"http://{args.host}:{args.port}/v1/embeddings"
39 - model_name = args.model
36 + base_url = f"http://{args.host}:{args.port}"
37 + models_url = base_url + "/v1/models"
38 + embeddings_url = base_url + "/v1/embeddings"
39 +
40 + response = requests.get(models_url)
41 + model = response.json()["data"][0]["id"]
42 +
43 + input_texts = [
44 + "The best thing about vLLM is that it supports many different models",
45 + ] * 2
46
47 # The OpenAI client does not support the embed_dtype and endianness parameters.
48 for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
49 for endianness in ENDIANNESS:
50 prompt = {
45 - "model": model_name,
46 - "input": "vLLM is great!",
51 + "model": model,
52 + "input": input_texts,
53 "encoding_format": "base64",
54 "embed_dtype": embed_dtype,
55 "endianness": endianness,
56 }
51 - response = post_http_request(prompt=prompt, api_url=api_url)
57 + response = post_http_request(prompt=prompt, api_url=embeddings_url)
58
59 embedding = []
60 for data in response.json()["data"]:
61 binary = base64.b64decode(data["embedding"])
62 tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
63 embedding.append(tensor.to(torch.float32))
58 - embedding = torch.cat(embedding)
64 + embedding = torch.stack(embedding)
65 print(embed_dtype, endianness, embedding.shape)
66
67
examples/pooling/embed/embedding_requests_bytes_online.py renamed
+11 -7
@@ -31,14 +31,18 @@ def parse_args():
31 parser = argparse.ArgumentParser()
32 parser.add_argument("--host", type=str, default="localhost")
33 parser.add_argument("--port", type=int, default=8000)
34 - parser.add_argument("--model", type=str, default="intfloat/e5-small")
34
35 return parser.parse_args()
36
37
38 def main(args):
40 - api_url = f"http://{args.host}:{args.port}/v1/embeddings"
41 - model_name = args.model
39 + base_url = f"http://{args.host}:{args.port}"
40 + models_url = base_url + "/v1/models"
41 + embeddings_url = base_url + "/v1/embeddings"
42 +
43 + response = requests.get(models_url)
44 + model = response.json()["data"][0]["id"]
45 +
46 embedding_size = 0
47
48 input_texts = [
@@ -50,13 +54,13 @@ def main(args):
54 for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
55 for endianness in ENDIANNESS:
56 prompt = {
53 - "model": model_name,
57 + "model": model,
58 "input": input_texts,
59 "encoding_format": "bytes",
60 "embed_dtype": embed_dtype,
61 "endianness": endianness,
62 }
59 - response = post_http_request(prompt=prompt, api_url=api_url)
63 + response = post_http_request(prompt=prompt, api_url=embeddings_url)
64 metadata = json.loads(response.headers["metadata"])
65 body = response.content
66 items = [MetadataItem(**x) for x in metadata["data"]]
@@ -73,13 +77,13 @@ def main(args):
77 for embed_dtype in EMBED_DTYPE_TO_TORCH_DTYPE:
78 for endianness in ENDIANNESS:
79 prompt = {
76 - "model": model_name,
80 + "model": model,
81 "input": input_texts,
82 "encoding_format": "bytes_only",
83 "embed_dtype": embed_dtype,
84 "endianness": endianness,
85 }
82 - response = post_http_request(prompt=prompt, api_url=api_url)
86 + response = post_http_request(prompt=prompt, api_url=embeddings_url)
87 body = response.content
88
89 items = build_metadata_items(
examples/pooling/embed/openai_embedding_matryoshka_fy_client.py renamed
examples/pooling/plugin/prithvi_geospatial_mae_online.py renamed
examples/pooling/pooling/pooling_online.py renamed
+10 -7
@@ -25,18 +25,21 @@ def parse_args():
25 parser = argparse.ArgumentParser()
26 parser.add_argument("--host", type=str, default="localhost")
27 parser.add_argument("--port", type=int, default=8000)
28 - parser.add_argument("--model", type=str, default="internlm/internlm2-1_8b-reward")
28
29 return parser.parse_args()
30
31
32 def main(args):
34 - api_url = f"http://{args.host}:{args.port}/pooling"
35 - model_name = args.model
33 + base_url = f"http://{args.host}:{args.port}"
34 + models_url = base_url + "/v1/models"
35 + pooing_url = base_url + "/pooling"
36 +
37 + response = requests.get(models_url)
38 + model = response.json()["data"][0]["id"]
39
40 # Input like Completions API
38 - prompt = {"model": model_name, "input": "vLLM is great!"}
39 - pooling_response = post_http_request(prompt=prompt, api_url=api_url)
41 + prompt = {"model": model, "input": "vLLM is great!"}
42 + pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
43 print("-" * 50)
44 print("Pooling Response:")
45 pprint.pprint(pooling_response.json())
@@ -44,7 +47,7 @@ def main(args):
47
48 # Input like Chat API
49 prompt = {
47 - "model": model_name,
50 + "model": model,
51 "messages": [
52 {
53 "role": "user",
@@ -52,7 +55,7 @@ def main(args):
55 }
56 ],
57 }
55 - pooling_response = post_http_request(prompt=prompt, api_url=api_url)
58 + pooling_response = post_http_request(prompt=prompt, api_url=pooing_url)
59 print("Pooling Response:")
60 pprint.pprint(pooling_response.json())
61 print("-" * 50)
examples/pooling/token_classify/ner_offline.py renamed
examples/pooling/token_classify/ner_online.py renamed
examples/pooling/token_embed/jina_embeddings_v4_offline.py renamed
examples/pooling/token_embed/multi_vector_retrieval_offline.py renamed
examples/pooling/token_embed/multi_vector_retrieval_online.py renamed
tests/entrypoints/pooling/classify/test_online.py
+148 -2
@@ -167,7 +167,8 @@ def test_truncate_prompt_tokens(server: RemoteOpenAIServer, model_name: str):
167
168 @pytest.mark.parametrize("model_name", [MODEL_NAME])
169 def test_add_special_tokens(server: RemoteOpenAIServer, model_name: str):
170 - # FIXME: The add_special_tokens parameter doesn't seem to be working.
170 + # The add_special_tokens parameter doesn't seem to be working with this model.
171 + # working with papluca/xlm-roberta-base-language-detection
172 response = requests.post(
173 server.url_for("classify"),
174 json={"model": model_name, "input": input_text, "add_special_tokens": False},
@@ -184,7 +185,110 @@ def test_add_special_tokens(server: RemoteOpenAIServer, model_name: str):
185
186
187 @pytest.mark.asyncio
187 -async def test_invocations(server: RemoteOpenAIServer):
188 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
189 +async def test_chat_request(server: RemoteOpenAIServer, model_name: str):
190 + messages = [
191 + {
192 + "role": "user",
193 + "content": "The cat sat on the mat.",
194 + },
195 + {
196 + "role": "assistant",
197 + "content": "A feline was resting on a rug.",
198 + },
199 + {
200 + "role": "user",
201 + "content": "Stars twinkle brightly in the night sky.",
202 + },
203 + ]
204 +
205 + # test chat request basic usage
206 + response = requests.post(
207 + server.url_for("classify"),
208 + json={"model": model_name, "messages": messages},
209 + )
210 +
211 + response.raise_for_status()
212 + output = ClassificationResponse.model_validate(response.json())
213 +
214 + assert output.object == "list"
215 + assert output.model == MODEL_NAME
216 + assert len(output.data) == 1
217 + assert hasattr(output.data[0], "label")
218 + assert hasattr(output.data[0], "probs")
219 + assert output.usage.prompt_tokens == 51
220 +
221 + # test add_generation_prompt
222 + response = requests.post(
223 + server.url_for("classify"),
224 + json={"model": model_name, "messages": messages, "add_generation_prompt": True},
225 + )
226 +
227 + response.raise_for_status()
228 + output = ClassificationResponse.model_validate(response.json())
229 +
230 + assert output.object == "list"
231 + assert output.model == MODEL_NAME
232 + assert len(output.data) == 1
233 + assert hasattr(output.data[0], "label")
234 + assert hasattr(output.data[0], "probs")
235 + assert output.usage.prompt_tokens == 54
236 +
237 + # test continue_final_message
238 + response = requests.post(
239 + server.url_for("classify"),
240 + json={
241 + "model": model_name,
242 + "messages": messages,
243 + "continue_final_message": True,
244 + },
245 + )
246 +
247 + response.raise_for_status()
248 + output = ClassificationResponse.model_validate(response.json())
249 +
250 + assert output.object == "list"
251 + assert output.model == MODEL_NAME
252 + assert len(output.data) == 1
253 + assert hasattr(output.data[0], "label")
254 + assert hasattr(output.data[0], "probs")
255 + assert output.usage.prompt_tokens == 49
256 +
257 + # test add_special_tokens
258 + # The add_special_tokens parameter doesn't seem to be working with this model.
259 + response = requests.post(
260 + server.url_for("classify"),
261 + json={"model": model_name, "messages": messages, "add_special_tokens": True},
262 + )
263 +
264 + response.raise_for_status()
265 + output = ClassificationResponse.model_validate(response.json())
266 +
267 + assert output.object == "list"
268 + assert output.model == MODEL_NAME
269 + assert len(output.data) == 1
270 + assert hasattr(output.data[0], "label")
271 + assert hasattr(output.data[0], "probs")
272 + assert output.usage.prompt_tokens == 51
273 +
274 + # test continue_final_message with add_generation_prompt
275 + response = requests.post(
276 + server.url_for("classify"),
277 + json={
278 + "model": model_name,
279 + "messages": messages,
280 + "continue_final_message": True,
281 + "add_generation_prompt": True,
282 + },
283 + )
284 + assert (
285 + "Cannot set both `continue_final_message` and `add_generation_prompt` to True."
286 + in response.json()["error"]["message"]
287 + )
288 +
289 +
290 +@pytest.mark.asyncio
291 +async def test_invocations_completion_request(server: RemoteOpenAIServer):
292 request_args = {
293 "model": MODEL_NAME,
294 "input": input_text,
@@ -213,6 +317,48 @@ async def test_invocations(server: RemoteOpenAIServer):
317 )
318
319
320 +@pytest.mark.asyncio
321 +async def test_invocations_chat_request(server: RemoteOpenAIServer):
322 + messages = [
323 + {
324 + "role": "user",
325 + "content": "The cat sat on the mat.",
326 + },
327 + {
328 + "role": "assistant",
329 + "content": "A feline was resting on a rug.",
330 + },
331 + {
332 + "role": "user",
333 + "content": "Stars twinkle brightly in the night sky.",
334 + },
335 + ]
336 +
337 + request_args = {"model": MODEL_NAME, "messages": messages}
338 +
339 + classification_response = requests.post(
340 + server.url_for("classify"), json=request_args
341 + )
342 + classification_response.raise_for_status()
343 +
344 + invocation_response = requests.post(
345 + server.url_for("invocations"), json=request_args
346 + )
347 + invocation_response.raise_for_status()
348 +
349 + classification_output = classification_response.json()
350 + invocation_output = invocation_response.json()
351 +
352 + assert classification_output.keys() == invocation_output.keys()
353 + for classification_data, invocation_data in zip(
354 + classification_output["data"], invocation_output["data"]
355 + ):
356 + assert classification_data.keys() == invocation_data.keys()
357 + assert classification_data["probs"] == pytest.approx(
358 + invocation_data["probs"], rel=0.01
359 + )
360 +
361 +
362 @pytest.mark.asyncio
363 @pytest.mark.parametrize("model_name", [MODEL_NAME])
364 async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
tests/entrypoints/pooling/embed/test_online.py
+124 -60
@@ -214,64 +214,6 @@ async def test_completion_request_batched(
214 run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
215
216
217 -@pytest.mark.asyncio
218 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
219 -async def test_conversation_embedding(
220 - server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
221 -):
222 - messages = [
223 - {
224 - "role": "user",
225 - "content": "The cat sat on the mat.",
226 - },
227 - {
228 - "role": "assistant",
229 - "content": "A feline was resting on a rug.",
230 - },
231 - {
232 - "role": "user",
233 - "content": "Stars twinkle brightly in the night sky.",
234 - },
235 - ]
236 -
237 - chat_response = requests.post(
238 - server.url_for("v1/embeddings"),
239 - json={
240 - "model": model_name,
241 - "messages": messages,
242 - "encoding_format": "float",
243 - },
244 - )
245 - chat_response.raise_for_status()
246 - chat_embeddings = EmbeddingResponse.model_validate(chat_response.json())
247 -
248 - tokenizer = get_tokenizer(tokenizer_name=model_name)
249 - prompt = tokenizer.apply_chat_template(
250 - messages,
251 - chat_template=DUMMY_CHAT_TEMPLATE,
252 - add_generation_prompt=True,
253 - continue_final_message=False,
254 - tokenize=False,
255 - )
256 - completion_response = await client.embeddings.create(
257 - model=model_name,
258 - input=prompt,
259 - encoding_format="float",
260 - # To be consistent with chat
261 - extra_body={"add_special_tokens": False},
262 - )
263 - completion_embeddings = EmbeddingResponse.model_validate(
264 - completion_response.model_dump(mode="json")
265 - )
266 -
267 - assert chat_embeddings.id is not None
268 - assert completion_embeddings.id is not None
269 - assert chat_embeddings.created <= completion_embeddings.created
270 - assert chat_embeddings.model_dump(exclude={"id", "created"}) == (
271 - completion_embeddings.model_dump(exclude={"id", "created"})
272 - )
273 -
274 -
217 @pytest.mark.asyncio
218 @pytest.mark.parametrize("model_name", [MODEL_NAME])
219 async def test_truncate_prompt_tokens(client: openai.AsyncOpenAI, model_name: str):
@@ -350,7 +292,129 @@ async def test_truncate_prompt_tokens(client: openai.AsyncOpenAI, model_name: st
292
293
294 @pytest.mark.asyncio
353 -async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenAI):
295 +@pytest.mark.parametrize("model_name", [MODEL_NAME])
296 +async def test_chat_request(
297 + server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
298 +):
299 + messages = [
300 + {
301 + "role": "user",
302 + "content": "The cat sat on the mat.",
303 + },
304 + {
305 + "role": "assistant",
306 + "content": "A feline was resting on a rug.",
307 + },
308 + {
309 + "role": "user",
310 + "content": "Stars twinkle brightly in the night sky.",
311 + },
312 + ]
313 +
314 + # test chat request basic usage
315 + chat_response = requests.post(
316 + server.url_for("v1/embeddings"),
317 + json={
318 + "model": model_name,
319 + "messages": messages,
320 + "encoding_format": "float",
321 + },
322 + )
323 + chat_response.raise_for_status()
324 + chat_embeddings = EmbeddingResponse.model_validate(chat_response.json())
325 +
326 + tokenizer = get_tokenizer(tokenizer_name=model_name)
327 + prompt = tokenizer.apply_chat_template(
328 + messages,
329 + chat_template=DUMMY_CHAT_TEMPLATE,
330 + add_generation_prompt=True,
331 + continue_final_message=False,
332 + tokenize=False,
333 + )
334 + completion_response = await client.embeddings.create(
335 + model=model_name,
336 + input=prompt,
337 + encoding_format="float",
338 + # To be consistent with chat
339 + extra_body={"add_special_tokens": False},
340 + )
341 + completion_embeddings = EmbeddingResponse.model_validate(
342 + completion_response.model_dump(mode="json")
343 + )
344 +
345 + assert chat_embeddings.id is not None
346 + assert completion_embeddings.id is not None
347 + assert chat_embeddings.created <= completion_embeddings.created
348 + assert chat_embeddings.model_dump(exclude={"id", "created"}) == (
349 + completion_embeddings.model_dump(exclude={"id", "created"})
350 + )
351 +
352 + # test add_generation_prompt
353 + response = requests.post(
354 + server.url_for("v1/embeddings"),
355 + json={"model": model_name, "messages": messages, "add_generation_prompt": True},
356 + )
357 +
358 + response.raise_for_status()
359 + output = EmbeddingResponse.model_validate(response.json())
360 +
361 + assert output.object == "list"
362 + assert len(output.data) == 1
363 + assert output.model == MODEL_NAME
364 + assert output.usage.prompt_tokens == 34
365 +
366 + # test continue_final_message
367 + response = requests.post(
368 + server.url_for("v1/embeddings"),
369 + json={
370 + "model": model_name,
371 + "messages": messages,
372 + "continue_final_message": True,
373 + },
374 + )
375 +
376 + response.raise_for_status()
377 + output = EmbeddingResponse.model_validate(response.json())
378 +
379 + assert output.object == "list"
380 + assert len(output.data) == 1
381 + assert output.model == MODEL_NAME
382 + assert output.usage.prompt_tokens == 33
383 +
384 + # test add_special_tokens
385 + response = requests.post(
386 + server.url_for("v1/embeddings"),
387 + json={"model": model_name, "messages": messages, "add_special_tokens": True},
388 + )
389 +
390 + response.raise_for_status()
391 + output = EmbeddingResponse.model_validate(response.json())
392 +
393 + assert output.object == "list"
394 + assert len(output.data) == 1
395 + assert output.model == MODEL_NAME
396 + assert output.usage.prompt_tokens == 36
397 +
398 + # test continue_final_message with add_generation_prompt
399 + response = requests.post(
400 + server.url_for("v1/embeddings"),
401 + json={
402 + "model": model_name,
403 + "messages": messages,
404 + "continue_final_message": True,
405 + "add_generation_prompt": True,
406 + },
407 + )
408 + assert (
409 + "Cannot set both `continue_final_message` and `add_generation_prompt` to True."
410 + in response.json()["error"]["message"]
411 + )
412 +
413 +
414 +@pytest.mark.asyncio
415 +async def test_invocations_completion_request(
416 + server: RemoteOpenAIServer, client: openai.AsyncOpenAI
417 +):
418 request_args = {
419 "model": MODEL_NAME,
420 "input": input_text,
@@ -381,7 +445,7 @@ async def test_invocations(server: RemoteOpenAIServer, client: openai.AsyncOpenA
445
446
447 @pytest.mark.asyncio
384 -async def test_invocations_conversation(server: RemoteOpenAIServer):
448 +async def test_invocations_chat_request(server: RemoteOpenAIServer):
449 messages = [
450 {
451 "role": "user",
tests/entrypoints/pooling/pooling/test_online.py
+66 -3
@@ -138,7 +138,7 @@ def test_completion_request_batched(server: RemoteOpenAIServer, model_name: str)
138
139 @pytest.mark.asyncio
140 @pytest.mark.parametrize("model_name", [MODEL_NAME])
141 -async def test_conversation_pooling(server: RemoteOpenAIServer, model_name: str):
141 +async def test_chat_request(server: RemoteOpenAIServer, model_name: str):
142 messages = [
143 {
144 "role": "user",
@@ -154,6 +154,7 @@ async def test_conversation_pooling(server: RemoteOpenAIServer, model_name: str)
154 },
155 ]
156
157 + # test chat request basic usage
158 chat_response = requests.post(
159 server.url_for("pooling"),
160 json={
@@ -193,6 +194,68 @@ async def test_conversation_pooling(server: RemoteOpenAIServer, model_name: str)
194 completion_poolings.model_dump(exclude={"id", "created"})
195 )
196
197 + # test add_generation_prompt
198 + response = requests.post(
199 + server.url_for("pooling"),
200 + json={"model": model_name, "messages": messages, "add_generation_prompt": True},
201 + )
202 +
203 + response.raise_for_status()
204 + output = PoolingResponse.model_validate(response.json())
205 +
206 + assert output.object == "list"
207 + assert len(output.data) == 1
208 + assert output.model == MODEL_NAME
209 + assert output.usage.prompt_tokens == 33
210 +
211 + # test continue_final_message
212 + # The continue_final_message parameter doesn't seem to be working with this model.
213 + response = requests.post(
214 + server.url_for("pooling"),
215 + json={
216 + "model": model_name,
217 + "messages": messages,
218 + "continue_final_message": True,
219 + },
220 + )
221 +
222 + response.raise_for_status()
223 + output = PoolingResponse.model_validate(response.json())
224 +
225 + assert output.object == "list"
226 + assert len(output.data) == 1
227 + assert output.model == MODEL_NAME
228 + assert output.usage.prompt_tokens == 33
229 +
230 + # test add_special_tokens
231 + response = requests.post(
232 + server.url_for("pooling"),
233 + json={"model": model_name, "messages": messages, "add_special_tokens": True},
234 + )
235 +
236 + response.raise_for_status()
237 + output = PoolingResponse.model_validate(response.json())
238 +
239 + assert output.object == "list"
240 + assert len(output.data) == 1
241 + assert output.model == MODEL_NAME
242 + assert output.usage.prompt_tokens == 34
243 +
244 + # test continue_final_message with add_generation_prompt
245 + response = requests.post(
246 + server.url_for("pooling"),
247 + json={
248 + "model": model_name,
249 + "messages": messages,
250 + "continue_final_message": True,
251 + "add_generation_prompt": True,
252 + },
253 + )
254 + assert (
255 + "Cannot set both `continue_final_message` and `add_generation_prompt` to True."
256 + in response.json()["error"]["message"]
257 + )
258 +
259
260 @pytest.mark.asyncio
261 @pytest.mark.parametrize("model_name", [MODEL_NAME])
@@ -430,7 +493,7 @@ async def test_params_not_supported(
493
494
495 @pytest.mark.asyncio
433 -async def test_invocations(server: RemoteOpenAIServer):
496 +async def test_invocations_chat_request(server: RemoteOpenAIServer):
497 request_args = {
498 "model": MODEL_NAME,
499 "input": input_text,
@@ -462,7 +525,7 @@ async def test_invocations(server: RemoteOpenAIServer):
525
526
527 @pytest.mark.asyncio
465 -async def test_invocations_conversation(server: RemoteOpenAIServer):
528 +async def test_invocations_conversation_chat_request(server: RemoteOpenAIServer):
529 messages = [
530 {
531 "role": "user",
vllm/entrypoints/pooling/base/protocol.py
+66 -2
@@ -2,10 +2,11 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5 -from typing import Annotated
5 +from typing import Annotated, Any
6
7 -from pydantic import Field
7 +from pydantic import Field, model_validator
8
9 +from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
10 from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel
11 from vllm.utils import random_uuid
12
@@ -44,3 +45,66 @@ class CompletionRequestMixin(OpenAIBaseModel):
45 "the prompt."
46 ),
47 )
48 +
49 +
50 +class ChatRequestMixin(OpenAIBaseModel):
51 + messages: list[ChatCompletionMessageParam]
52 +
53 + add_generation_prompt: bool = Field(
54 + default=False,
55 + description=(
56 + "If true, the generation prompt will be added to the chat template. "
57 + "This is a parameter used by chat template in tokenizer config of the "
58 + "model."
59 + ),
60 + )
61 +
62 + continue_final_message: bool = Field(
63 + default=False,
64 + description=(
65 + "If this is set, the chat will be formatted so that the final "
66 + "message in the chat is open-ended, without any EOS tokens. The "
67 + "model will continue this message rather than starting a new one. "
68 + 'This allows you to "prefill" part of the model\'s response for it. '
69 + "Cannot be used at the same time as `add_generation_prompt`."
70 + ),
71 + )
72 +
73 + add_special_tokens: bool = Field(
74 + default=False,
75 + description=(
76 + "If true, special tokens (e.g. BOS) will be added to the prompt "
77 + "on top of what is added by the chat template. "
78 + "For most models, the chat template takes care of adding the "
79 + "special tokens so this should be set to false (as is the "
80 + "default)."
81 + ),
82 + )
83 +
84 + chat_template: str | None = Field(
85 + default=None,
86 + description=(
87 + "A Jinja template to use for this conversion. "
88 + "As of transformers v4.44, default chat template is no longer "
89 + "allowed, so you must provide a chat template if the tokenizer "
90 + "does not define one."
91 + ),
92 + )
93 +
94 + chat_template_kwargs: dict[str, Any] | None = Field(
95 + default=None,
96 + description=(
97 + "Additional keyword args to pass to the template renderer. "
98 + "Will be accessible by the chat template."
99 + ),
100 + )
101 +
102 + @model_validator(mode="before")
103 + @classmethod
104 + def check_generation_prompt(cls, data):
105 + if data.get("continue_final_message") and data.get("add_generation_prompt"):
106 + raise ValueError(
107 + "Cannot set both `continue_final_message` and "
108 + "`add_generation_prompt` to True."
109 + )
110 + return data
vllm/entrypoints/pooling/classify/protocol.py
+2 -42
@@ -10,9 +10,9 @@ from pydantic import (
10
11 from vllm import PoolingParams
12 from vllm.config.pooler import get_use_activation
13 -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
13 from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
14 from vllm.entrypoints.pooling.base.protocol import (
15 + ChatRequestMixin,
16 CompletionRequestMixin,
17 PoolingBasicRequestMixin,
18 )
@@ -45,48 +45,8 @@ class ClassificationCompletionRequest(PoolingBasicRequestMixin, CompletionReques
45 )
46
47
48 -class ClassificationChatRequest(PoolingBasicRequestMixin):
49 - messages: list[ChatCompletionMessageParam]
50 -
48 +class ClassificationChatRequest(PoolingBasicRequestMixin, ChatRequestMixin):
49 # --8<-- [start:chat-classification-extra-params]
52 - add_generation_prompt: bool = Field(
53 - default=False,
54 - description=(
55 - "If true, the generation prompt will be added to the chat template. "
56 - "This is a parameter used by chat template in tokenizer config of the "
57 - "model."
58 - ),
59 - )
60 -
61 - add_special_tokens: bool = Field(
62 - default=False,
63 - description=(
64 - "If true, special tokens (e.g. BOS) will be added to the prompt "
65 - "on top of what is added by the chat template. "
66 - "For most models, the chat template takes care of adding the "
67 - "special tokens so this should be set to false (as is the "
68 - "default)."
69 - ),
70 - )
71 -
72 - chat_template: str | None = Field(
73 - default=None,
74 - description=(
75 - "A Jinja template to use for this conversion. "
76 - "As of transformers v4.44, default chat template is no longer "
77 - "allowed, so you must provide a chat template if the tokenizer "
78 - "does not define one."
79 - ),
80 - )
81 -
82 - chat_template_kwargs: dict[str, Any] | None = Field(
83 - default=None,
84 - description=(
85 - "Additional keyword args to pass to the template renderer. "
86 - "Will be accessible by the chat template."
87 - ),
88 - )
89 -
50 mm_processor_kwargs: dict[str, Any] | None = Field(
51 default=None,
52 description=("Additional kwargs to pass to the HF processor."),
vllm/entrypoints/pooling/classify/serving.py
+2 -2
@@ -86,8 +86,8 @@ class ClassificationMixin(OpenAIServing):
86 ChatTemplateContentFormatOption,
87 getattr(self, "chat_template_content_format", "auto"),
88 ),
89 - add_generation_prompt=False,
90 - continue_final_message=False,
89 + add_generation_prompt=chat_request.add_generation_prompt,
90 + continue_final_message=chat_request.continue_final_message,
91 add_special_tokens=chat_request.add_special_tokens,
92 )
93 ctx.engine_prompts = engine_prompts
vllm/entrypoints/pooling/embed/protocol.py
+2 -59
@@ -5,13 +5,12 @@ from typing import Any, TypeAlias
5
6 from pydantic import (
7 Field,
8 - model_validator,
8 )
9
10 from vllm import PoolingParams
12 -from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
11 from vllm.entrypoints.openai.engine.protocol import OpenAIBaseModel, UsageInfo
12 from vllm.entrypoints.pooling.base.protocol import (
13 + ChatRequestMixin,
14 CompletionRequestMixin,
15 PoolingBasicRequestMixin,
16 )
@@ -57,57 +56,11 @@ class EmbeddingCompletionRequest(PoolingBasicRequestMixin, CompletionRequestMixi
56 )
57
58
60 -class EmbeddingChatRequest(PoolingBasicRequestMixin):
61 - messages: list[ChatCompletionMessageParam]
62 -
59 +class EmbeddingChatRequest(PoolingBasicRequestMixin, ChatRequestMixin):
60 encoding_format: EncodingFormat = "float"
61 dimensions: int | None = None
62
63 # --8<-- [start:chat-embedding-extra-params]
67 - add_generation_prompt: bool = Field(
68 - default=False,
69 - description=(
70 - "If true, the generation prompt will be added to the chat template. "
71 - "This is a parameter used by chat template in tokenizer config of the "
72 - "model."
73 - ),
74 - )
75 - continue_final_message: bool = Field(
76 - default=False,
77 - description=(
78 - "If this is set, the chat will be formatted so that the final "
79 - "message in the chat is open-ended, without any EOS tokens. The "
80 - "model will continue this message rather than starting a new one. "
81 - 'This allows you to "prefill" part of the model\'s response for it. '
82 - "Cannot be used at the same time as `add_generation_prompt`."
83 - ),
84 - )
85 - add_special_tokens: bool = Field(
86 - default=False,
87 - description=(
88 - "If true, special tokens (e.g. BOS) will be added to the prompt "
89 - "on top of what is added by the chat template. "
90 - "For most models, the chat template takes care of adding the "
91 - "special tokens so this should be set to false (as is the "
92 - "default)."
93 - ),
94 - )
95 - chat_template: str | None = Field(
96 - default=None,
97 - description=(
98 - "A Jinja template to use for this conversion. "
99 - "As of transformers v4.44, default chat template is no longer "
100 - "allowed, so you must provide a chat template if the tokenizer "
101 - "does not define one."
102 - ),
103 - )
104 - chat_template_kwargs: dict[str, Any] | None = Field(
105 - default=None,
106 - description=(
107 - "Additional keyword args to pass to the template renderer. "
108 - "Will be accessible by the chat template."
109 - ),
110 - )
64 mm_processor_kwargs: dict[str, Any] | None = Field(
65 default=None,
66 description=("Additional kwargs to pass to the HF processor."),
@@ -134,16 +87,6 @@ class EmbeddingChatRequest(PoolingBasicRequestMixin):
87 )
88 # --8<-- [end:chat-embedding-extra-params]
89
137 - @model_validator(mode="before")
138 - @classmethod
139 - def check_generation_prompt(cls, data):
140 - if data.get("continue_final_message") and data.get("add_generation_prompt"):
141 - raise ValueError(
142 - "Cannot set both `continue_final_message` and "
143 - "`add_generation_prompt` to True."
144 - )
145 - return data
146 -
90 def to_pooling_params(self):
91 return PoolingParams(
92 truncate_prompt_tokens=self.truncate_prompt_tokens,
vllm/entrypoints/pooling/pooling/serving.py
+2 -4
@@ -144,10 +144,8 @@ class OpenAIServingPooling(OpenAIServing):
144 request.messages,
145 chat_template=request.chat_template or self.chat_template,
146 chat_template_content_format=self.chat_template_content_format,
147 - # In pooling requests, we are not generating tokens,
148 - # so there is no need to append extra tokens to the input
149 - add_generation_prompt=False,
150 - continue_final_message=False,
147 + add_generation_prompt=request.add_generation_prompt,
148 + continue_final_message=request.continue_final_message,
149 add_special_tokens=request.add_special_tokens,
150 )
151 elif isinstance(request, PoolingCompletionRequest):