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 5, 21:00 UTC
[BugFix][Frontend] apply task instruction as system prompt in cohere v2/embed (#38362)
Signed-off-by: walterbm <walter.beller.morales@gmail.com>
Walter Beller-Morales committed
Mar 28, 2026 at 14:30 UTC
fafca38adc1ce65d2c9e2857138c3c0d65b0905e
3 files changed
+242
-39
tests/entrypoints/pooling/embed/test_cohere_openai_parity.py
+23
-4
@@ -57,16 +57,25 @@ def _openai_embed(
57
return [item["embedding"] for item in resp.json()["data"]]
58
59
60
+def _cosine_sim(a: list[float], b: list[float]) -> float:
61
+ va, vb = np.array(a), np.array(b)
62
+ return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb)))
63
+
64
+
65
def test_single_text_parity(server: RemoteOpenAIServer):
61
- """A single text should produce identical embeddings via both APIs."""
66
+ """A single text should produce equivalent embeddings via both APIs."""
67
texts = ["the quick brown fox jumps over the lazy dog"]
68
v2 = _cohere_embed(server, texts)
69
v1 = _openai_embed(server, texts)
65
- np.testing.assert_allclose(v2[0], v1[0], rtol=1e-5)
70
+ # Full-suite BF16 runs can introduce tiny numerical drift even when both
71
+ # endpoints are functionally equivalent, so compare semantic equivalence
72
+ # instead of exact elementwise equality.
73
+ cos = _cosine_sim(v2[0], v1[0])
74
+ assert cos > 0.9999, f"single-text parity failed, cosine={cos}"
75
76
77
def test_batch_parity(server: RemoteOpenAIServer):
69
- """A batch of texts should produce identical embeddings via both APIs,
78
+ """A batch of texts should produce equivalent embeddings via both APIs,
79
in the same order."""
80
texts = [
81
"machine learning",
@@ -76,8 +85,18 @@ def test_batch_parity(server: RemoteOpenAIServer):
85
v2 = _cohere_embed(server, texts)
86
v1 = _openai_embed(server, texts)
87
assert len(v2) == len(v1) == 3
88
+
89
+ similarities = np.array(
90
+ [[_cosine_sim(v2_emb, v1_emb) for v1_emb in v1] for v2_emb in v2]
91
+ )
92
for i in range(3):
80
- np.testing.assert_allclose(v2[i], v1[i], rtol=1e-5, err_msg=f"index {i}")
93
+ assert int(np.argmax(similarities[i])) == i, (
94
+ f"batch parity order mismatch at index {i}: "
95
+ f"similarities={similarities[i].tolist()}"
96
+ )
97
+ assert similarities[i, i] > 0.9999, (
98
+ f"batch parity failed at index {i}, cosine={similarities[i, i]}"
99
+ )
100
101
102
def test_token_count_parity(server: RemoteOpenAIServer):
tests/entrypoints/pooling/embed/test_io_processor.py
+116
@@ -6,8 +6,11 @@ import pytest
6
7
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
8
from vllm.entrypoints.pooling.embed.protocol import (
9
+ CohereEmbedContent,
10
+ CohereEmbedInput,
11
CohereEmbedRequest,
12
)
13
+from vllm.entrypoints.pooling.typing import PoolingServeContext
14
15
16
class TestResolveTruncation:
@@ -206,3 +209,116 @@ class TestValidateInputType:
209
handler = self._make_handler({"a": "", "b": ""})
210
with pytest.raises(ValueError, match="Supported values: a, b"):
211
handler._validate_input_type("z")
212
+
213
+
214
+class TestPreProcessCohereOnline:
215
+ """Unit tests for EmbedIOProcessor._pre_process_cohere_online."""
216
+
217
+ @staticmethod
218
+ def _make_context(**request_kwargs) -> PoolingServeContext[CohereEmbedRequest]:
219
+ return PoolingServeContext(
220
+ request=CohereEmbedRequest(model="test", **request_kwargs),
221
+ model_name="test",
222
+ request_id="embd-test",
223
+ )
224
+
225
+ @staticmethod
226
+ def _make_handler():
227
+ handler = object.__new__(EmbedIOProcessor)
228
+ handler._validate_input_type = lambda _input_type: None
229
+ return handler
230
+
231
+ def test_text_only_without_task_prefix_uses_completion_path(self):
232
+ handler = self._make_handler()
233
+ ctx = self._make_context(texts=["hello"])
234
+ calls: list[tuple[str, object]] = []
235
+
236
+ def preprocess_completion(request, prompt_input, prompt_embeds):
237
+ calls.append(("completion", prompt_input))
238
+ return ["completion"]
239
+
240
+ handler._get_task_instruction_prefix = lambda _input_type: None
241
+ handler._has_chat_template = lambda: False
242
+ handler._preprocess_completion_online = preprocess_completion
243
+ handler._batch_render_chat = lambda *_args, **_kwargs: (
244
+ pytest.fail("text-only request should not require chat rendering")
245
+ )
246
+
247
+ handler._pre_process_cohere_online(ctx)
248
+
249
+ assert ctx.engine_inputs == ["completion"]
250
+ assert calls == [("completion", ["hello"])]
251
+
252
+ def test_text_only_falls_back_to_prefixed_completion_without_template(self):
253
+ handler = self._make_handler()
254
+ ctx = self._make_context(texts=["hello"], input_type="query")
255
+ calls: list[tuple[str, object]] = []
256
+
257
+ def preprocess_completion(request, prompt_input, prompt_embeds):
258
+ calls.append(("completion", prompt_input))
259
+ return ["fallback"]
260
+
261
+ handler._get_task_instruction_prefix = lambda _input_type: "query: "
262
+ handler._has_chat_template = lambda: False
263
+ handler._batch_render_chat = lambda *_args, **_kwargs: (
264
+ pytest.fail("chat rendering should be skipped without a template")
265
+ )
266
+ handler._preprocess_completion_online = preprocess_completion
267
+
268
+ handler._pre_process_cohere_online(ctx)
269
+
270
+ assert ctx.engine_inputs == ["fallback"]
271
+ assert calls == [("completion", ["query: hello"])]
272
+
273
+ def test_text_only_with_template_uses_chat_path(self):
274
+ handler = self._make_handler()
275
+ ctx = self._make_context(texts=["hello"], input_type="query")
276
+ calls: list[tuple[str, object]] = []
277
+
278
+ def batch_render_chat(
279
+ request,
280
+ all_messages,
281
+ truncate_prompt_tokens,
282
+ truncation_side,
283
+ ):
284
+ calls.append(
285
+ (
286
+ "chat",
287
+ {
288
+ "request": request,
289
+ "all_messages": all_messages,
290
+ "truncate_prompt_tokens": truncate_prompt_tokens,
291
+ "truncation_side": truncation_side,
292
+ },
293
+ )
294
+ )
295
+ return ["chat"]
296
+
297
+ handler._get_task_instruction_prefix = lambda _input_type: "query: "
298
+ handler._has_chat_template = lambda: True
299
+ handler._batch_render_chat = batch_render_chat
300
+ handler._preprocess_completion_online = lambda *_args, **_kwargs: (
301
+ pytest.fail("completion path should be skipped when a template exists")
302
+ )
303
+
304
+ handler._pre_process_cohere_online(ctx)
305
+
306
+ assert ctx.engine_inputs == ["chat"]
307
+ assert calls == [
308
+ (
309
+ "chat",
310
+ {
311
+ "request": ctx.request,
312
+ "all_messages": [
313
+ handler._mixed_input_to_messages(
314
+ CohereEmbedInput(
315
+ content=[CohereEmbedContent(type="text", text="hello")]
316
+ ),
317
+ task_prefix="query: ",
318
+ )
319
+ ],
320
+ "truncate_prompt_tokens": -1,
321
+ "truncation_side": None,
322
+ },
323
+ )
324
+ ]
vllm/entrypoints/pooling/embed/io_processor.py
+103
-35
@@ -18,6 +18,7 @@ from vllm.entrypoints.chat_utils import (
18
)
19
from vllm.entrypoints.pooling.base.io_processor import PoolingIOProcessor
20
from vllm.entrypoints.pooling.embed.protocol import (
21
+ CohereEmbedContent,
22
CohereEmbedInput,
23
CohereEmbedRequest,
24
EmbeddingChatRequest,
@@ -28,6 +29,7 @@ from vllm.inputs import EngineInput, tokens_input
29
from vllm.logger import init_logger
30
from vllm.outputs import PoolingOutput, PoolingRequestOutput
31
from vllm.renderers import merge_kwargs
32
+from vllm.renderers.hf import resolve_chat_template
33
from vllm.utils.collection_utils import chunk_list
34
from vllm.utils.mistral import is_mistral_tokenizer
35
@@ -284,13 +286,27 @@ class EmbedIOProcessor(PoolingIOProcessor):
286
) -> list[ChatCompletionMessageParam]:
287
"""Build chat messages from a mixed text+image input.
288
287
- When *task_prefix* is given, it is prepended to each text part.
289
+ When *task_prefix* is given, it is used as the system prompt.
290
"""
291
+ messages: list[ChatCompletionMessageParam] = []
292
+ if task_prefix is not None:
293
+ messages.append(
294
+ CustomChatCompletionMessageParam(
295
+ role="system",
296
+ content=[
297
+ ChatCompletionContentPartTextParam(
298
+ type="text", text=task_prefix
299
+ )
300
+ ],
301
+ )
302
+ )
303
+
304
parts: list[ChatCompletionContentPartParam] = []
305
for item in inp.content:
306
if item.type == "text" and item.text is not None:
292
- text = task_prefix + item.text if task_prefix else item.text
293
- parts.append(ChatCompletionContentPartTextParam(type="text", text=text))
307
+ parts.append(
308
+ ChatCompletionContentPartTextParam(type="text", text=item.text)
309
+ )
310
elif item.type == "image_url" and item.image_url is not None:
311
parts.append(
312
ChatCompletionContentPartImageParam(
@@ -298,7 +314,8 @@ class EmbedIOProcessor(PoolingIOProcessor):
314
image_url=ImageURL(url=item.image_url["url"]),
315
)
316
)
301
- return [CustomChatCompletionMessageParam(role="user", content=parts)]
317
+ messages.append(CustomChatCompletionMessageParam(role="user", content=parts))
318
+ return messages
319
320
@staticmethod
321
def _check_cohere_max_tokens(
@@ -346,9 +363,11 @@ class EmbedIOProcessor(PoolingIOProcessor):
363
def _pre_process_cohere_online(self, ctx: PoolingServeContext) -> None:
364
"""Convert a ``CohereEmbedRequest`` into engine prompts.
365
349
- For texts, a single batched completion request path is used.
350
- For images and mixed inputs, conversations are batch-rendered
351
- through the chat template in one ``render_chat`` call.
366
+ If a model has a chat template the task instruction are rendered
367
+ as a system prompt. Otherwise they are just prepended to the input text.
368
+
369
+ Images and mixed inputs are always batch-rendered through the chat
370
+ template in one ``render_chat`` call.
371
"""
372
request = ctx.request
373
assert isinstance(request, CohereEmbedRequest)
@@ -363,42 +382,91 @@ class EmbedIOProcessor(PoolingIOProcessor):
382
self._validate_input_type(input_type)
383
384
if request.images is not None:
366
- all_messages: list[list[ChatCompletionMessageParam]] = [
367
- [
368
- CustomChatCompletionMessageParam(
369
- role="user",
370
- content=[{"type": "image_url", "image_url": {"url": uri}}],
371
- )
372
- ]
385
+ input: list[CohereEmbedInput] = [
386
+ CohereEmbedInput(
387
+ content=[
388
+ CohereEmbedContent(type="image_url", image_url={"url": uri})
389
+ ]
390
+ )
391
for uri in request.images
392
]
375
- ctx.engine_inputs = self._batch_render_chat(
376
- request, all_messages, truncate_prompt_tokens, truncation_side
377
- )
378
-
393
elif request.inputs is not None:
394
+ input = request.inputs
395
+ else:
396
+ texts = request.texts or []
397
task_prefix = self._get_task_instruction_prefix(input_type)
398
+
399
+ if task_prefix is None:
400
+ ctx.engine_inputs = self._preprocess_cohere_text_completion(
401
+ request,
402
+ texts,
403
+ truncate_prompt_tokens,
404
+ truncation_side,
405
+ )
406
+ return
407
+
408
all_messages = [
382
- self._mixed_input_to_messages(inp, task_prefix=task_prefix)
383
- for inp in request.inputs
409
+ self._mixed_input_to_messages(
410
+ CohereEmbedInput(
411
+ content=[CohereEmbedContent(type="text", text=text)]
412
+ ),
413
+ task_prefix=task_prefix,
414
+ )
415
+ for text in texts
416
]
385
- ctx.engine_inputs = self._batch_render_chat(
386
- request, all_messages, truncate_prompt_tokens, truncation_side
387
- )
417
+ if self._has_chat_template():
418
+ ctx.engine_inputs = self._batch_render_chat(
419
+ request,
420
+ all_messages,
421
+ truncate_prompt_tokens,
422
+ truncation_side,
423
+ )
424
+ else:
425
+ ctx.engine_inputs = self._preprocess_cohere_text_completion(
426
+ request,
427
+ self._apply_task_instruction(texts, input_type),
428
+ truncate_prompt_tokens,
429
+ truncation_side,
430
+ )
431
+ return
432
389
- else:
390
- prefixed = self._apply_task_instruction(request.texts or [], input_type)
391
- proxy = EmbeddingCompletionRequest(
392
- model=request.model,
393
- input=prefixed,
394
- dimensions=request.output_dimension,
395
- encoding_format="float",
396
- truncate_prompt_tokens=truncate_prompt_tokens,
397
- truncation_side=truncation_side,
398
- )
399
- ctx.engine_inputs = self._preprocess_completion_online(
400
- proxy, prompt_input=proxy.input, prompt_embeds=None
433
+ task_prefix = self._get_task_instruction_prefix(input_type)
434
+ all_messages = [
435
+ self._mixed_input_to_messages(inp, task_prefix=task_prefix) for inp in input
436
+ ]
437
+ ctx.engine_inputs = self._batch_render_chat(
438
+ request, all_messages, truncate_prompt_tokens, truncation_side
439
+ )
440
+
441
+ def _has_chat_template(self) -> bool:
442
+ return (
443
+ resolve_chat_template(
444
+ self.renderer.tokenizer,
445
+ chat_template=self.chat_template,
446
+ tools=None,
447
+ model_config=self.model_config,
448
)
449
+ is not None
450
+ )
451
+
452
+ def _preprocess_cohere_text_completion(
453
+ self,
454
+ request: CohereEmbedRequest,
455
+ texts: list[str],
456
+ truncate_prompt_tokens: int | None,
457
+ truncation_side: Literal["left", "right"] | None,
458
+ ) -> list[EngineInput]:
459
+ proxy = EmbeddingCompletionRequest(
460
+ model=request.model,
461
+ input=texts,
462
+ dimensions=request.output_dimension,
463
+ encoding_format="float",
464
+ truncate_prompt_tokens=truncate_prompt_tokens,
465
+ truncation_side=truncation_side,
466
+ )
467
+ return self._preprocess_completion_online(
468
+ proxy, prompt_input=proxy.input, prompt_embeds=None
469
+ )
470
471
def _batch_render_chat(
472
self,