Add vLLM v0.18.1 source tree with KV transfer abort fix

third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:

  vllm/v1/core/sched/scheduler.py:
    Replace fatal assert with graceful skip when KV transfer callback
    arrives for an already-aborted request during PD disaggregated serving.

Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-05-22 00:30:38 +08:00
parent b6591950bc
commit 445e491123
4285 changed files with 1111303 additions and 1 deletions

View File

View File

@@ -0,0 +1,92 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref
import pytest
from vllm import LLM, PoolingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.platforms import current_platform
MODEL_NAME = "intfloat/multilingual-e5-small"
PROMPTS = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
TOKEN_IDS = [
# Using ID={0, 1, 2, 3} results in NaN values,
# so we add this offset of 1000
[1000],
[1000, 1001],
[1000, 1002, 1001],
[1000, 1003, 1001, 1002],
]
@pytest.fixture(scope="module")
def llm():
# ROCm: Use FLEX_ATTENTION backend as it's the only attention backend
# that supports encoder-only models on ROCm.
attention_config = None
if current_platform.is_rocm():
attention_config = {"backend": "FLEX_ATTENTION"}
# pytest caches the fixture so we use weakref.proxy to
# enable garbage collection
llm = LLM(
model=MODEL_NAME,
max_num_batched_tokens=32768,
tensor_parallel_size=1,
gpu_memory_utilization=0.75,
enforce_eager=True,
seed=0,
attention_config=attention_config,
)
yield weakref.proxy(llm)
del llm
cleanup_dist_env_and_memory()
@pytest.mark.skip_global_cleanup
def test_multiple_pooling_params(llm: LLM):
pooling_params = [
PoolingParams(),
PoolingParams(),
PoolingParams(),
PoolingParams(),
]
# Multiple PoolingParams should be matched with each prompt
outputs = llm.encode(PROMPTS, pooling_params=pooling_params, pooling_task="embed")
assert len(PROMPTS) == len(outputs)
# Exception raised, if the size of params does not match the size of prompts
with pytest.raises(ValueError):
outputs = llm.encode(
PROMPTS, pooling_params=pooling_params[:3], pooling_task="embed"
)
# Single PoolingParams should be applied to every prompt
single_pooling_params = PoolingParams()
outputs = llm.encode(
PROMPTS, pooling_params=single_pooling_params, pooling_task="embed"
)
assert len(PROMPTS) == len(outputs)
# pooling_params is None, default params should be applied
outputs = llm.encode(PROMPTS, pooling_params=None, pooling_task="embed")
assert len(PROMPTS) == len(outputs)
def test_right_side_truncation(llm: LLM):
# Embeddings models should truncate the end of the prompt
tokenizer = llm.get_tokenizer()
assert tokenizer.truncation_side == "right"

View File

@@ -0,0 +1,104 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Any
import openai
import pytest
import pytest_asyncio
from tests.utils import RemoteOpenAIServer
from vllm.platforms import current_platform
MODEL_NAME = "sentence-transformers/all-MiniLM-L12-v2"
max_model_len = 128
input = """Immerse yourself in the enchanting chronicle of calculus, a
mathematical domain that has radically transformed our comprehension of
change and motion. Despite its roots in ancient civilizations, the
formal birth of calculus predominantly occurred in the 17th century,
primarily under the influential guidance of Sir Isaac Newton and Gottfried
Wilhelm Leibniz. The earliest traces of calculus concepts are found in
ancient Greek mathematics,most notably in the works of Eudoxus and
Archimedes, around 300 BCE. They utilized the 'method of exhaustion'—a
technique for computing areas and volumes through the use of finite sums.
This methodology laid crucial foundational work for integral calculus.
In the 17th century, both Newton and Leibniz independently pioneered
calculus, each contributing unique perspectives that would shape this new
field."""
@pytest.fixture(scope="module")
def server():
args = [
"--runner",
"pooling",
"--dtype",
"bfloat16",
"--enforce-eager",
"--max-model-len",
str(max_model_len),
]
# ROCm: Use Flex Attention to support encoder-only self-attention.
if current_platform.is_rocm():
args.extend(["--attention-backend", "FLEX_ATTENTION"])
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def client(server):
async with server.get_async_client() as async_client:
yield async_client
@pytest.mark.asyncio
async def test_smaller_truncation_size(client: openai.AsyncOpenAI):
truncation_size = 10
kwargs: dict[str, Any] = {
"model": MODEL_NAME,
"input": input,
"truncate_prompt_tokens": truncation_size,
}
response = await client.post(path="embeddings", cast_to=object, body={**kwargs})
assert response["usage"]["prompt_tokens"] == truncation_size
@pytest.mark.asyncio
async def test_bigger_truncation_size(client: openai.AsyncOpenAI):
truncation_size = max_model_len + 1
kwargs: dict[str, Any] = {
"model": MODEL_NAME,
"input": input,
"truncate_prompt_tokens": truncation_size,
}
with pytest.raises(openai.BadRequestError) as err:
await client.post(path="embeddings", cast_to=object, body={**kwargs})
assert err.value.status_code == 400
error_details = err.value.response.json()["error"]
assert error_details["type"] == "BadRequestError"
expected_message = (
"truncate_prompt_tokens value is "
"greater than max_model_len."
" Please, select a smaller truncation size."
)
assert error_details["message"] == expected_message
@pytest.mark.asyncio
async def test_max_truncation_size(client: openai.AsyncOpenAI):
truncation_size = -1
kwargs: dict[str, Any] = {
"model": MODEL_NAME,
"input": input,
"truncate_prompt_tokens": truncation_size,
}
response = await client.post(path="embeddings", cast_to=object, body={**kwargs})
assert response["usage"]["prompt_tokens"] == max_model_len

View File

@@ -0,0 +1,114 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref
import pytest
import torch
from tests.models.utils import softmax
from vllm import LLM, ClassificationRequestOutput, PoolingParams, PoolingRequestOutput
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.tasks import PoolingTask
MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
prompt = "The chef prepared a delicious meal."
prompt_token_ids = [785, 29706, 10030, 264, 17923, 15145, 13]
num_labels = 2
@pytest.fixture(scope="module")
def llm():
# pytest caches the fixture so we use weakref.proxy to
# enable garbage collection
llm = LLM(
model=MODEL_NAME,
max_num_batched_tokens=32768,
tensor_parallel_size=1,
gpu_memory_utilization=0.75,
enforce_eager=True,
seed=0,
)
yield weakref.proxy(llm)
del llm
cleanup_dist_env_and_memory()
@pytest.mark.skip_global_cleanup
def test_str_prompts(llm: LLM):
outputs = llm.classify(prompt, use_tqdm=False)
assert len(outputs) == 1
assert isinstance(outputs[0], ClassificationRequestOutput)
assert outputs[0].prompt_token_ids == prompt_token_ids
assert len(outputs[0].outputs.probs) == num_labels
@pytest.mark.skip_global_cleanup
def test_token_ids_prompts(llm: LLM):
outputs = llm.classify([prompt_token_ids], use_tqdm=False)
assert len(outputs) == 1
assert isinstance(outputs[0], ClassificationRequestOutput)
assert outputs[0].prompt_token_ids == prompt_token_ids
assert len(outputs[0].outputs.probs) == num_labels
@pytest.mark.skip_global_cleanup
def test_list_prompts(llm: LLM):
outputs = llm.classify([prompt, prompt_token_ids], use_tqdm=False)
assert len(outputs) == 2
for i in range(len(outputs)):
assert isinstance(outputs[i], ClassificationRequestOutput)
assert outputs[i].prompt_token_ids == prompt_token_ids
assert len(outputs[i].outputs.probs) == num_labels
@pytest.mark.skip_global_cleanup
def test_token_classify(llm: LLM):
outputs = llm.encode(prompt, pooling_task="token_classify", use_tqdm=False)
assert len(outputs) == 1
assert isinstance(outputs[0], PoolingRequestOutput)
assert outputs[0].prompt_token_ids == prompt_token_ids
assert outputs[0].outputs.data.shape == (len(prompt_token_ids), num_labels)
@pytest.mark.skip_global_cleanup
def test_pooling_params(llm: LLM):
def get_outputs(use_activation):
outputs = llm.classify(
prompt,
pooling_params=PoolingParams(use_activation=use_activation),
use_tqdm=False,
)
return torch.tensor([x.outputs.probs for x in outputs])
default = get_outputs(use_activation=None)
w_activation = get_outputs(use_activation=True)
wo_activation = get_outputs(use_activation=False)
assert torch.allclose(default, w_activation, atol=1e-2), (
"Default should use activation."
)
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
"wo_activation should not use activation."
)
assert torch.allclose(softmax(wo_activation), w_activation, atol=1e-2), (
"w_activation should be close to activation(wo_activation)."
)
@pytest.mark.skip_global_cleanup
def test_score_api(llm: LLM):
err_msg = "Score API is only enabled for num_labels == 1."
with pytest.raises(ValueError, match=err_msg):
llm.score("ping", "pong", use_tqdm=False)
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
def test_unsupported_tasks(llm: LLM, task: PoolingTask):
err_msg = f"Unsupported task: '{task}' Supported tasks.+"
with pytest.raises(ValueError, match=err_msg):
llm.encode(prompt, pooling_task=task, use_tqdm=False)

View File

@@ -0,0 +1,472 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import requests
import torch
import torch.nn.functional as F
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.pooling.classify.protocol import ClassificationResponse
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
MODEL_NAME = "jason9693/Qwen2.5-1.5B-apeach"
DTYPE = "float32" # Use float32 to avoid NaN issue
input_text = "This product was excellent and exceeded my expectations"
input_tokens = [1986, 1985, 572, 9073, 323, 33808, 847, 16665]
@pytest.fixture(scope="module")
def server():
args = [
"--enforce-eager",
"--max-model-len",
"512",
"--dtype",
DTYPE,
]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_basic(server: RemoteOpenAIServer, model_name: str):
# test /v1/models
response = requests.get(server.url_for("/v1/models"))
served_model = response.json()["data"][0]["id"]
assert served_model == MODEL_NAME
# test /tokenize
response = requests.post(
server.url_for("/tokenize"),
json={"model": model_name, "prompt": input_text},
)
assert response.json()["tokens"] == input_tokens
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_completion_request(server: RemoteOpenAIServer, model_name: str):
# test input: str
classification_response = requests.post(
server.url_for("classify"),
json={"model": model_name, "input": input_text},
)
classification_response.raise_for_status()
output = ClassificationResponse.model_validate(classification_response.json())
assert output.object == "list"
assert output.model == MODEL_NAME
assert len(output.data) == 1
assert hasattr(output.data[0], "label")
assert hasattr(output.data[0], "probs")
# test input: list[int]
classification_response = requests.post(
server.url_for("classify"),
json={"model": model_name, "input": input_tokens},
)
classification_response.raise_for_status()
output = ClassificationResponse.model_validate(classification_response.json())
assert output.object == "list"
assert output.model == MODEL_NAME
assert len(output.data) == 1
assert hasattr(output.data[0], "label")
assert hasattr(output.data[0], "probs")
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_completion_request_batched(server: RemoteOpenAIServer, model_name: str):
N = 10
# test input: list[str]
classification_response = requests.post(
server.url_for("classify"),
json={"model": model_name, "input": [input_text] * N},
)
output = ClassificationResponse.model_validate(classification_response.json())
assert len(output.data) == N
for i, item in enumerate(output.data):
assert item.index == i
assert hasattr(item, "label")
assert hasattr(item, "probs")
assert len(item.probs) == item.num_classes
assert item.label in ["Default", "Spoiled"]
# test input: list[list[int]]
classification_response = requests.post(
server.url_for("classify"),
json={"model": model_name, "input": [input_tokens] * N},
)
output = ClassificationResponse.model_validate(classification_response.json())
assert len(output.data) == N
for i, item in enumerate(output.data):
assert item.index == i
assert hasattr(item, "label")
assert hasattr(item, "probs")
assert len(item.probs) == item.num_classes
assert item.label in ["Default", "Spoiled"]
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_empty_input_error(server: RemoteOpenAIServer, model_name: str):
classification_response = requests.post(
server.url_for("classify"),
json={"model": model_name, "input": ""},
)
error = classification_response.json()
assert classification_response.status_code == 400
assert "error" in error
classification_response = requests.post(
server.url_for("classify"),
json={"model": model_name, "input": []},
)
error = classification_response.json()
assert classification_response.status_code == 400
assert "error" in error
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_truncate_prompt_tokens(server: RemoteOpenAIServer, model_name: str):
long_text = "hello " * 600
classification_response = requests.post(
server.url_for("classify"),
json={"model": model_name, "input": long_text, "truncate_prompt_tokens": 5},
)
classification_response.raise_for_status()
output = ClassificationResponse.model_validate(classification_response.json())
assert len(output.data) == 1
assert output.data[0].index == 0
assert hasattr(output.data[0], "probs")
assert output.usage.prompt_tokens == 5
assert output.usage.total_tokens == 5
# invalid_truncate_prompt_tokens
classification_response = requests.post(
server.url_for("classify"),
json={"model": model_name, "input": "test", "truncate_prompt_tokens": 513},
)
error = classification_response.json()
assert classification_response.status_code == 400
assert "truncate_prompt_tokens" in error["error"]["message"]
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_add_special_tokens(server: RemoteOpenAIServer, model_name: str):
# The add_special_tokens parameter doesn't seem to be working with this model.
# working with papluca/xlm-roberta-base-language-detection
response = requests.post(
server.url_for("classify"),
json={"model": model_name, "input": input_text, "add_special_tokens": False},
)
response.raise_for_status()
ClassificationResponse.model_validate(response.json())
response = requests.post(
server.url_for("classify"),
json={"model": model_name, "input": input_text, "add_special_tokens": True},
)
response.raise_for_status()
ClassificationResponse.model_validate(response.json())
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_chat_request(server: RemoteOpenAIServer, model_name: str):
messages = [
{
"role": "user",
"content": "The cat sat on the mat.",
},
{
"role": "assistant",
"content": "A feline was resting on a rug.",
},
{
"role": "user",
"content": "Stars twinkle brightly in the night sky.",
},
]
# test chat request basic usage
response = requests.post(
server.url_for("classify"),
json={"model": model_name, "messages": messages},
)
response.raise_for_status()
output = ClassificationResponse.model_validate(response.json())
assert output.object == "list"
assert output.model == MODEL_NAME
assert len(output.data) == 1
assert hasattr(output.data[0], "label")
assert hasattr(output.data[0], "probs")
assert output.usage.prompt_tokens == 51
# test add_generation_prompt
response = requests.post(
server.url_for("classify"),
json={"model": model_name, "messages": messages, "add_generation_prompt": True},
)
response.raise_for_status()
output = ClassificationResponse.model_validate(response.json())
assert output.object == "list"
assert output.model == MODEL_NAME
assert len(output.data) == 1
assert hasattr(output.data[0], "label")
assert hasattr(output.data[0], "probs")
assert output.usage.prompt_tokens == 54
# test continue_final_message
response = requests.post(
server.url_for("classify"),
json={
"model": model_name,
"messages": messages,
"continue_final_message": True,
},
)
response.raise_for_status()
output = ClassificationResponse.model_validate(response.json())
assert output.object == "list"
assert output.model == MODEL_NAME
assert len(output.data) == 1
assert hasattr(output.data[0], "label")
assert hasattr(output.data[0], "probs")
assert output.usage.prompt_tokens == 49
# test add_special_tokens
# The add_special_tokens parameter doesn't seem to be working with this model.
response = requests.post(
server.url_for("classify"),
json={"model": model_name, "messages": messages, "add_special_tokens": True},
)
response.raise_for_status()
output = ClassificationResponse.model_validate(response.json())
assert output.object == "list"
assert output.model == MODEL_NAME
assert len(output.data) == 1
assert hasattr(output.data[0], "label")
assert hasattr(output.data[0], "probs")
assert output.usage.prompt_tokens == 51
# test continue_final_message with add_generation_prompt
response = requests.post(
server.url_for("classify"),
json={
"model": model_name,
"messages": messages,
"continue_final_message": True,
"add_generation_prompt": True,
},
)
assert (
"Cannot set both `continue_final_message` and `add_generation_prompt` to True."
in response.json()["error"]["message"]
)
@pytest.mark.asyncio
async def test_invocations_completion_request(server: RemoteOpenAIServer):
request_args = {
"model": MODEL_NAME,
"input": input_text,
}
classification_response = requests.post(
server.url_for("classify"), json=request_args
)
classification_response.raise_for_status()
invocation_response = requests.post(
server.url_for("invocations"), json=request_args
)
invocation_response.raise_for_status()
classification_output = classification_response.json()
invocation_output = invocation_response.json()
assert classification_output.keys() == invocation_output.keys()
for classification_data, invocation_data in zip(
classification_output["data"], invocation_output["data"]
):
assert classification_data.keys() == invocation_data.keys()
assert classification_data["probs"] == pytest.approx(
invocation_data["probs"], rel=0.01
)
@pytest.mark.asyncio
async def test_invocations_chat_request(server: RemoteOpenAIServer):
messages = [
{
"role": "user",
"content": "The cat sat on the mat.",
},
{
"role": "assistant",
"content": "A feline was resting on a rug.",
},
{
"role": "user",
"content": "Stars twinkle brightly in the night sky.",
},
]
request_args = {"model": MODEL_NAME, "messages": messages}
classification_response = requests.post(
server.url_for("classify"), json=request_args
)
classification_response.raise_for_status()
invocation_response = requests.post(
server.url_for("invocations"), json=request_args
)
invocation_response.raise_for_status()
classification_output = classification_response.json()
invocation_output = invocation_response.json()
assert classification_output.keys() == invocation_output.keys()
for classification_data, invocation_data in zip(
classification_output["data"], invocation_output["data"]
):
assert classification_data.keys() == invocation_data.keys()
assert classification_data["probs"] == pytest.approx(
invocation_data["probs"], rel=0.01
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
async def get_outputs(use_activation):
response = requests.post(
server.url_for("classify"),
json={
"model": model_name,
"input": input_text,
"use_activation": use_activation,
},
)
outputs = response.json()
return torch.tensor([x["probs"] for x in outputs["data"]])
default = await get_outputs(use_activation=None)
w_activation = await get_outputs(use_activation=True)
wo_activation = await get_outputs(use_activation=False)
assert torch.allclose(default, w_activation, atol=1e-2), (
"Default should use activation."
)
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
"wo_activation should not use activation."
)
assert torch.allclose(F.softmax(wo_activation, dim=-1), w_activation, atol=1e-2), (
"w_activation should be close to activation(wo_activation)."
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_score(server: RemoteOpenAIServer, model_name: str):
# score api is only enabled for num_labels == 1.
response = requests.post(
server.url_for("score"),
json={
"model": model_name,
"queries": "ping",
"documents": "pong",
},
)
assert response.json()["detail"] == "Not Found"
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_rerank(server: RemoteOpenAIServer, model_name: str):
# rerank api is only enabled for num_labels == 1.
response = requests.post(
server.url_for("rerank"),
json={
"model": model_name,
"query": "ping",
"documents": ["pong"],
},
)
assert response.json()["detail"] == "Not Found"
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": input_text,
"encoding_format": "float",
"task": "classify",
},
)
poolings = PoolingResponse.model_validate(response.json())
assert len(poolings.data) == 1
assert len(poolings.data[0].data) == 2
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
task = "token_classify"
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": input_text,
"encoding_format": "float",
"task": task,
},
)
poolings = PoolingResponse.model_validate(response.json())
assert len(poolings.data) == 1
assert len(poolings.data[0].data) == 8
assert len(poolings.data[0].data[0]) == 2
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
async def test_pooling_not_supported(
server: RemoteOpenAIServer, model_name: str, task: str
):
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": input_text,
"encoding_format": "float",
"task": task,
},
)
assert response.json()["error"]["type"] == "BadRequestError"
assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")

View File

@@ -0,0 +1,146 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import pytest
import requests
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.pooling.classify.protocol import ClassificationResponse
from vllm.multimodal.utils import encode_image_url, fetch_image
MODEL_NAME = "muziyongshixin/Qwen2.5-VL-7B-for-VideoCls"
MAXIMUM_VIDEOS = 1
HF_OVERRIDES = {"architectures": ["Qwen2_5_VLForSequenceClassification"]}
input_text = "This product was excellent and exceeded my expectations"
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
image_base64 = {"url": encode_image_url(fetch_image(image_url))}
video_url = "https://www.bogotobogo.com/python/OpenCV_Python/images/mean_shift_tracking/slow_traffic_small.mp4"
@pytest.fixture(scope="module")
def server():
args = [
"--runner",
"pooling",
"--max-model-len",
"5000",
"--enforce-eager",
"--limit-mm-per-prompt",
json.dumps({"video": MAXIMUM_VIDEOS}),
]
with RemoteOpenAIServer(
MODEL_NAME, args, override_hf_configs=HF_OVERRIDES
) as remote_server:
yield remote_server
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_chat_text_request(server: RemoteOpenAIServer, model_name: str):
messages = [
{
"role": "assistant",
"content": "Please classify this text request.",
},
{
"role": "user",
"content": input_text,
},
]
response = requests.post(
server.url_for("classify"),
json={"model": model_name, "messages": messages},
)
response.raise_for_status()
output = ClassificationResponse.model_validate(response.json())
assert output.object == "list"
assert output.model == model_name
assert len(output.data) == 1
assert len(output.data[0].probs) == 2
assert output.usage.prompt_tokens == 35
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_chat_image_url_request(server: RemoteOpenAIServer, model_name: str):
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Please classify this image."},
{"type": "image_url", "image_url": {"url": image_url}},
],
}
]
response = requests.post(
server.url_for("classify"),
json={"model": model_name, "messages": messages},
)
response.raise_for_status()
output = ClassificationResponse.model_validate(response.json())
assert output.object == "list"
assert output.model == model_name
assert len(output.data) == 1
assert len(output.data[0].probs) == 2
assert output.usage.prompt_tokens == 47
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_chat_image_base64_request(server: RemoteOpenAIServer, model_name: str):
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Please classify this image."},
{"type": "image_url", "image_url": image_base64},
],
}
]
response = requests.post(
server.url_for("classify"),
json={"model": model_name, "messages": messages},
)
response.raise_for_status()
output = ClassificationResponse.model_validate(response.json())
assert output.object == "list"
assert output.model == model_name
assert len(output.data) == 1
assert len(output.data[0].probs) == 2
assert output.usage.prompt_tokens == 47
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_chat_video_url_request(server: RemoteOpenAIServer, model_name: str):
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Please classify this video."},
{"type": "video_url", "video_url": {"url": video_url}},
],
}
]
response = requests.post(
server.url_for("classify"),
json={"model": model_name, "messages": messages},
)
response.raise_for_status()
output = ClassificationResponse.model_validate(response.json())
assert output.object == "list"
assert output.model == model_name
assert len(output.data) == 1
assert len(output.data[0].probs) == 2
assert output.usage.prompt_tokens == 4807

View File

@@ -0,0 +1,28 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Pytest configuration for vLLM pooling embed tests."""
import warnings
import torch
from vllm.platforms import current_platform
def pytest_collection_modifyitems(config, items):
"""Configure ROCm-specific settings based on collected tests."""
if not current_platform.is_rocm():
return
# Disable Flash/MemEfficient SDP on ROCm to avoid HF Transformers
# accuracy issues: https://github.com/vllm-project/vllm/issues/30167
# TODO: Remove once ROCm SDP accuracy issues are resolved on HuggingFace
torch.backends.cuda.enable_flash_sdp(False)
torch.backends.cuda.enable_mem_efficient_sdp(False)
torch.backends.cuda.enable_math_sdp(True)
warnings.warn(
"ROCm: Disabled flash_sdp and mem_efficient_sdp, enabled math_sdp "
"to avoid HuggingFace Transformers accuracy issues",
UserWarning,
stacklevel=1,
)

View File

@@ -0,0 +1,310 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the Cohere /v2/embed API with generic (non-Cohere) models.
Validates that the Cohere v2 embed endpoint works correctly with standard
embedding models, covering text embedding, embedding type conversions,
response structure, batching, normalisation, and semantic similarity.
"""
import base64
import struct
import numpy as np
import pytest
import requests
from tests.utils import RemoteOpenAIServer
DTYPE = "bfloat16"
MODELS: list[tuple[str, list[str]]] = [
("intfloat/multilingual-e5-small", []),
(
"Snowflake/snowflake-arctic-embed-m-v1.5",
[
"--trust_remote_code",
"--hf_overrides",
'{"matryoshka_dimensions":[256]}',
],
),
]
@pytest.fixture(scope="module", params=MODELS, ids=lambda m: m[0])
def model_config(request):
return request.param
@pytest.fixture(scope="module")
def model_name(model_config):
return model_config[0]
@pytest.fixture(scope="module")
def server(model_config):
name, extra_args = model_config
args = [
"--runner",
"pooling",
"--dtype",
DTYPE,
"--enforce-eager",
"--max-model-len",
"512",
"--gpu-memory-utilization",
"0.02",
] + extra_args
with RemoteOpenAIServer(name, args) as remote_server:
yield remote_server
def _cohere_embed(
server: RemoteOpenAIServer,
model_name: str,
texts: list[str] | None = None,
images: list[str] | None = None,
input_type: str | None = None,
embedding_types: list[str] | None = None,
) -> dict:
body: dict = {"model": model_name}
if input_type is not None:
body["input_type"] = input_type
if texts is not None:
body["texts"] = texts
if images is not None:
body["images"] = images
if embedding_types is not None:
body["embedding_types"] = embedding_types
resp = requests.post(server.url_for("/v2/embed"), json=body)
resp.raise_for_status()
return resp.json()
def _openai_embed(
server: RemoteOpenAIServer, model_name: str, texts: list[str]
) -> dict:
body = {"model": model_name, "input": texts, "encoding_format": "float"}
resp = requests.post(server.url_for("/v1/embeddings"), json=body)
resp.raise_for_status()
return resp.json()
def _cosine_sim(a: list[float], b: list[float]) -> float:
va, vb = np.array(a), np.array(b)
return float(np.dot(va, vb) / (np.linalg.norm(va) * np.linalg.norm(vb)))
# -----------------------------------------------------------
# Text embedding tests
# -----------------------------------------------------------
def test_basic_embed(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(
server, model_name, texts=["hello world"], embedding_types=["float"]
)
assert "embeddings" in r
assert len(r["embeddings"]["float"]) == 1
assert len(r["embeddings"]["float"][0]) > 0
def test_unsupported_input_type_rejected(server: RemoteOpenAIServer, model_name: str):
"""An input_type not defined in the model's prompt config should be
rejected with a 400 error."""
body = {
"model": model_name,
"input_type": "nonexistent_type",
"texts": ["hello world"],
"embedding_types": ["float"],
}
resp = requests.post(server.url_for("/v2/embed"), json=body)
assert resp.status_code == 400
assert "Unsupported input_type" in resp.json()["error"]["message"]
def test_omitted_input_type_accepted(server: RemoteOpenAIServer, model_name: str):
"""Omitting input_type should always work (no prompt prefix applied)."""
body = {
"model": model_name,
"texts": ["hello world"],
"embedding_types": ["float"],
}
resp = requests.post(server.url_for("/v2/embed"), json=body)
assert resp.status_code == 200
data = resp.json()
assert len(data["embeddings"]["float"]) == 1
def test_v1_v2_parity(server: RemoteOpenAIServer, model_name: str):
"""v1 (OpenAI) and v2 (Cohere) endpoints should produce the same
float embeddings for a generic model."""
texts = ["hello world"]
v2 = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
v1 = _openai_embed(server, model_name, texts)
cos = _cosine_sim(v2["embeddings"]["float"][0], v1["data"][0]["embedding"])
assert cos > 0.9999, f"v1/v2 parity failed, cosine={cos}"
def test_embedding_types(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(
server,
model_name,
texts=["test"],
embedding_types=["float", "binary", "ubinary"],
)
dim = len(r["embeddings"]["float"][0])
assert len(r["embeddings"]["binary"][0]) == dim // 8
assert len(r["embeddings"]["ubinary"][0]) == dim // 8
def test_response_structure(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(server, model_name, texts=["test"], embedding_types=["float"])
assert "id" in r
assert "embeddings" in r
assert "texts" in r
assert r["texts"] == ["test"]
assert "meta" in r
assert r["meta"]["api_version"]["version"] == "2"
assert "billed_units" in r["meta"]
assert r["meta"]["billed_units"]["input_tokens"] > 0
assert r["meta"]["billed_units"]["image_tokens"] == 0
def test_batch(server: RemoteOpenAIServer, model_name: str):
texts = ["apple", "banana", "cherry"]
r = _cohere_embed(server, model_name, texts=texts, embedding_types=["float"])
assert len(r["embeddings"]["float"]) == 3
dim = len(r["embeddings"]["float"][0])
for emb in r["embeddings"]["float"]:
assert len(emb) == dim
def test_l2_normalized(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(
server, model_name, texts=["hello world"], embedding_types=["float"]
)
emb = np.array(r["embeddings"]["float"][0])
assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
def test_semantic_similarity(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(
server,
model_name,
texts=["machine learning", "deep learning", "chocolate cake recipe"],
embedding_types=["float"],
)
embs = r["embeddings"]["float"]
cos_related = _cosine_sim(embs[0], embs[1])
cos_unrelated = _cosine_sim(embs[0], embs[2])
assert cos_related > cos_unrelated
def test_missing_input_returns_error(server: RemoteOpenAIServer, model_name: str):
body = {"model": model_name}
resp = requests.post(server.url_for("/v2/embed"), json=body)
assert resp.status_code == 400
def test_base64_embedding_type(server: RemoteOpenAIServer, model_name: str):
r = _cohere_embed(
server,
model_name,
texts=["test encoding"],
embedding_types=["float", "base64"],
)
float_emb = r["embeddings"]["float"][0]
b64_str = r["embeddings"]["base64"][0]
decoded = struct.unpack(f"<{len(float_emb)}f", base64.b64decode(b64_str))
np.testing.assert_allclose(float_emb, decoded, rtol=1e-5)
# -----------------------------------------------------------
# Truncation tests
# -----------------------------------------------------------
def _cohere_embed_raw(
server: RemoteOpenAIServer,
body: dict,
) -> requests.Response:
return requests.post(server.url_for("/v2/embed"), json=body)
def test_truncate_end_succeeds(server: RemoteOpenAIServer, model_name: str):
"""truncate=END should silently truncate long input."""
long_text = " ".join(["word"] * 2000)
body = {
"model": model_name,
"texts": [long_text],
"embedding_types": ["float"],
"truncate": "END",
}
resp = _cohere_embed_raw(server, body)
assert resp.status_code == 200
data = resp.json()
assert len(data["embeddings"]["float"]) == 1
def test_truncate_start_succeeds(server: RemoteOpenAIServer, model_name: str):
"""truncate=START should silently truncate long input from the start."""
long_text = " ".join(["word"] * 2000)
body = {
"model": model_name,
"texts": [long_text],
"embedding_types": ["float"],
"truncate": "START",
}
resp = _cohere_embed_raw(server, body)
assert resp.status_code == 200
data = resp.json()
assert len(data["embeddings"]["float"]) == 1
def test_truncate_none_rejects_long_input(server: RemoteOpenAIServer, model_name: str):
"""truncate=NONE should error when input exceeds model context."""
long_text = " ".join(["word"] * 2000)
body = {
"model": model_name,
"texts": [long_text],
"embedding_types": ["float"],
"truncate": "NONE",
}
resp = _cohere_embed_raw(server, body)
assert resp.status_code == 400
def test_truncate_start_vs_end_differ(server: RemoteOpenAIServer, model_name: str):
"""START and END truncation should produce different embeddings
when the input is long enough to actually be truncated.
We construct input with distinct tokens at the start vs end
so that keeping different halves produces different embeddings.
"""
start_words = " ".join([f"alpha{i}" for i in range(300)])
end_words = " ".join([f"omega{i}" for i in range(300)])
long_text = start_words + " " + end_words
body_end = {
"model": model_name,
"texts": [long_text],
"embedding_types": ["float"],
"truncate": "END",
}
body_start = {
"model": model_name,
"texts": [long_text],
"embedding_types": ["float"],
"truncate": "START",
}
r_end = _cohere_embed_raw(server, body_end).json()
r_start = _cohere_embed_raw(server, body_start).json()
emb_end = r_end["embeddings"]["float"][0]
emb_start = r_start["embeddings"]["float"][0]
cos = _cosine_sim(emb_end, emb_start)
assert cos < 0.99, (
f"START and END truncation should produce different embeddings "
f"for long input, but cosine similarity was {cos}"
)

View File

@@ -0,0 +1,135 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for the Cohere /v2/embed API with a multimodal model (SigLIP).
Validates image embedding, batching, normalisation, and embedding type
conversions through the /v2/embed endpoint.
"""
import base64
import struct
import zlib
import numpy as np
import pytest
import requests
from tests.utils import RemoteOpenAIServer
MODEL_NAME = "google/siglip-so400m-patch14-384"
DTYPE = "bfloat16"
@pytest.fixture(scope="module")
def server():
args = [
"--runner",
"pooling",
"--dtype",
DTYPE,
"--enforce-eager",
"--max-model-len",
"64",
"--gpu-memory-utilization",
"0.3",
]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
def _make_tiny_png(r: int, g: int, b: int, w: int = 2, h: int = 2) -> str:
raw = b""
for _ in range(h):
raw += b"\x00" + bytes([r, g, b]) * w
compressed = zlib.compress(raw)
def chunk(ctype: bytes, cdata: bytes) -> bytes:
c = ctype + cdata
return (
struct.pack(">I", len(cdata))
+ c
+ struct.pack(">I", zlib.crc32(c) & 0xFFFFFFFF)
)
ihdr = struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)
png = (
b"\x89PNG\r\n\x1a\n"
+ chunk(b"IHDR", ihdr)
+ chunk(b"IDAT", compressed)
+ chunk(b"IEND", b"")
)
return "data:image/png;base64," + base64.b64encode(png).decode()
def _cohere_embed(
server: RemoteOpenAIServer,
texts: list[str] | None = None,
images: list[str] | None = None,
embedding_types: list[str] | None = None,
) -> dict:
body: dict = {"model": MODEL_NAME}
if texts is not None:
body["texts"] = texts
if images is not None:
body["images"] = images
if embedding_types is not None:
body["embedding_types"] = embedding_types
resp = requests.post(server.url_for("/v2/embed"), json=body)
resp.raise_for_status()
return resp.json()
def test_image_embed(server: RemoteOpenAIServer):
img_uri = _make_tiny_png(255, 0, 0)
r = _cohere_embed(
server,
images=[img_uri],
embedding_types=["float"],
)
assert "embeddings" in r
assert len(r["embeddings"]["float"]) == 1
assert len(r["embeddings"]["float"][0]) > 0
assert r["meta"]["billed_units"]["image_tokens"] > 0
assert r["meta"]["billed_units"]["input_tokens"] == 0
def test_image_batch(server: RemoteOpenAIServer):
red = _make_tiny_png(255, 0, 0)
blue = _make_tiny_png(0, 0, 255)
r = _cohere_embed(
server,
images=[red, blue],
embedding_types=["float"],
)
assert len(r["embeddings"]["float"]) == 2
def test_image_l2_normalized(server: RemoteOpenAIServer):
img_uri = _make_tiny_png(0, 255, 0)
r = _cohere_embed(
server,
images=[img_uri],
embedding_types=["float"],
)
emb = np.array(r["embeddings"]["float"][0])
assert abs(float(np.linalg.norm(emb)) - 1.0) < 0.01
def test_image_embedding_types(server: RemoteOpenAIServer):
img_uri = _make_tiny_png(128, 128, 128)
r = _cohere_embed(
server,
images=[img_uri],
embedding_types=["float", "binary", "ubinary"],
)
dim = len(r["embeddings"]["float"][0])
assert len(r["embeddings"]["binary"][0]) == dim // 8
assert len(r["embeddings"]["ubinary"][0]) == dim // 8
def test_text_embed_on_multimodal(server: RemoteOpenAIServer):
"""SigLIP also supports text-only embedding via /v2/embed."""
r = _cohere_embed(server, texts=["hello world"], embedding_types=["float"])
assert "embeddings" in r
assert len(r["embeddings"]["float"]) == 1
assert len(r["embeddings"]["float"][0]) > 0

View File

@@ -0,0 +1,102 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Parity test between Cohere /v2/embed and OpenAI /v1/embeddings.
Verifies that both endpoints produce identical float embeddings when
no prompt prefix is applied (input_type omitted for Cohere /v2/embed).
"""
import numpy as np
import pytest
import requests
from tests.utils import RemoteOpenAIServer
MODEL_NAME = "BAAI/bge-base-en-v1.5"
DTYPE = "bfloat16"
@pytest.fixture(scope="module")
def server():
args = [
"--runner",
"pooling",
"--dtype",
DTYPE,
"--enforce-eager",
"--max-model-len",
"512",
"--gpu-memory-utilization",
"0.02",
]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
def _cohere_embed(
server: RemoteOpenAIServer,
texts: list[str],
) -> list[list[float]]:
body = {
"model": MODEL_NAME,
"texts": texts,
"embedding_types": ["float"],
}
resp = requests.post(server.url_for("/v2/embed"), json=body)
resp.raise_for_status()
return resp.json()["embeddings"]["float"]
def _openai_embed(
server: RemoteOpenAIServer,
texts: list[str],
) -> list[list[float]]:
body = {"model": MODEL_NAME, "input": texts, "encoding_format": "float"}
resp = requests.post(server.url_for("/v1/embeddings"), json=body)
resp.raise_for_status()
return [item["embedding"] for item in resp.json()["data"]]
def test_single_text_parity(server: RemoteOpenAIServer):
"""A single text should produce identical embeddings via both APIs."""
texts = ["the quick brown fox jumps over the lazy dog"]
v2 = _cohere_embed(server, texts)
v1 = _openai_embed(server, texts)
np.testing.assert_allclose(v2[0], v1[0], rtol=1e-5)
def test_batch_parity(server: RemoteOpenAIServer):
"""A batch of texts should produce identical embeddings via both APIs,
in the same order."""
texts = [
"machine learning",
"deep learning",
"natural language processing",
]
v2 = _cohere_embed(server, texts)
v1 = _openai_embed(server, texts)
assert len(v2) == len(v1) == 3
for i in range(3):
np.testing.assert_allclose(v2[i], v1[i], rtol=1e-5, err_msg=f"index {i}")
def test_token_count_parity(server: RemoteOpenAIServer):
"""Both APIs should report the same prompt token count."""
texts = ["hello world"]
v2_resp = requests.post(
server.url_for("/v2/embed"),
json={
"model": MODEL_NAME,
"texts": texts,
"embedding_types": ["float"],
},
)
v1_resp = requests.post(
server.url_for("/v1/embeddings"),
json={"model": MODEL_NAME, "input": texts, "encoding_format": "float"},
)
v2_resp.raise_for_status()
v1_resp.raise_for_status()
v2_tokens = v2_resp.json()["meta"]["billed_units"]["input_tokens"]
v1_tokens = v1_resp.json()["usage"]["prompt_tokens"]
assert v2_tokens == v1_tokens

View File

@@ -0,0 +1,46 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
import pytest
from tests.models.language.pooling_mteb_test.mteb_embed_utils import (
MTEB_EMBED_TASKS,
MTEB_EMBED_TOL,
OpenAIClientMtebEncoder,
run_mteb_embed_task,
)
from tests.utils import RemoteOpenAIServer
from vllm.platforms import current_platform
os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
MODEL_NAME = "intfloat/e5-small"
MAIN_SCORE = 0.7422994752439667
@pytest.fixture(scope="module")
def server():
args = ["--runner", "pooling", "--enforce-eager", "--disable-uvicorn-access-log"]
# ROCm: Use Flex Attention to support encoder-only self-attention.
if current_platform.is_rocm():
args.extend(["--attention-backend", "FLEX_ATTENTION"])
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
def test_mteb_embed(server):
client = server.get_client()
encoder = OpenAIClientMtebEncoder(MODEL_NAME, client)
vllm_main_score = run_mteb_embed_task(encoder, MTEB_EMBED_TASKS)
st_main_score = MAIN_SCORE
print("VLLM main score: ", vllm_main_score)
print("SentenceTransformer main score: ", st_main_score)
print("Difference: ", st_main_score - vllm_main_score)
# We are not concerned that the vllm mteb results are better
# than SentenceTransformers, so we only perform one-sided testing.
assert st_main_score - vllm_main_score < MTEB_EMBED_TOL

View File

@@ -0,0 +1,208 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for EmbedIOProcessor."""
import pytest
from vllm.entrypoints.pooling.embed.io_processor import EmbedIOProcessor
from vllm.entrypoints.pooling.embed.protocol import (
CohereEmbedRequest,
)
class TestResolveTruncation:
"""Unit tests for EmbedIOProcessor._resolve_cohere_truncation."""
@staticmethod
def _make_request(**kwargs) -> CohereEmbedRequest:
defaults = {
"model": "test",
"input_type": "search_document",
"texts": ["hello"],
}
return CohereEmbedRequest(**(defaults | kwargs))
def test_truncate_end_default(self):
req = self._make_request()
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens == -1
assert side is None
def test_truncate_end_explicit(self):
req = self._make_request(truncate="END")
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens == -1
assert side is None
def test_truncate_end_with_max_tokens(self):
req = self._make_request(truncate="END", max_tokens=128)
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens == 128
assert side is None
def test_truncate_none(self):
req = self._make_request(truncate="NONE")
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens is None
assert side is None
def test_truncate_none_with_max_tokens(self):
"""truncate=NONE should NOT set truncate_prompt_tokens; the
max_tokens limit is enforced separately via _check_max_tokens."""
req = self._make_request(truncate="NONE", max_tokens=10)
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens is None
assert side is None
def test_truncate_start(self):
req = self._make_request(truncate="START")
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens == -1
assert side == "left"
def test_truncate_start_with_max_tokens(self):
req = self._make_request(truncate="START", max_tokens=64)
tokens, side = EmbedIOProcessor._resolve_cohere_truncation(req)
assert tokens == 64
assert side == "left"
class TestApplyStPrompt:
"""Unit tests for EmbedIOProcessor._apply_task_instruction."""
@staticmethod
def _make_handler(task_instructions: dict[str, str] | None):
handler = object.__new__(EmbedIOProcessor)
handler.task_instructions = task_instructions
return handler
def test_no_prompts_configured(self):
handler = self._make_handler(None)
texts = ["hello", "world"]
assert handler._apply_task_instruction(texts, "query") is texts
def test_matching_input_type(self):
handler = self._make_handler({"query": "search_query: "})
result = handler._apply_task_instruction(["hello"], "query")
assert result == ["search_query: hello"]
def test_non_matching_input_type(self):
handler = self._make_handler({"query": "search_query: "})
texts = ["hello"]
assert handler._apply_task_instruction(texts, "document") is texts
def test_multiple_texts(self):
handler = self._make_handler(
{"query": "Represent this sentence for searching: "}
)
result = handler._apply_task_instruction(["a", "b", "c"], "query")
assert result == [
"Represent this sentence for searching: a",
"Represent this sentence for searching: b",
"Represent this sentence for searching: c",
]
def test_empty_prefix_returns_unchanged(self):
handler = self._make_handler({"passage": ""})
texts = ["hello"]
assert handler._apply_task_instruction(texts, "passage") is texts
class TestLoadTaskInstructions:
"""Unit tests for EmbedIOProcessor._load_task_instructions."""
def test_no_attribute(self):
class FakeConfig:
pass
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
def test_with_task_instructions(self):
class FakeConfig:
task_instructions = {
"retrieval.query": "Represent the query: ",
"retrieval.passage": "",
}
result = EmbedIOProcessor._load_task_instructions(FakeConfig())
assert result == {
"retrieval.query": "Represent the query: ",
"retrieval.passage": "",
}
def test_empty_dict(self):
class FakeConfig:
task_instructions = {}
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
def test_non_dict(self):
class FakeConfig:
task_instructions = "not a dict"
assert EmbedIOProcessor._load_task_instructions(FakeConfig()) is None
class TestCheckMaxTokens:
"""Unit tests for EmbedIOProcessor._check_cohere_max_tokens."""
@staticmethod
def _fake_output(n_tokens: int):
class _Out:
def __init__(self, n: int):
self.prompt_token_ids = list(range(n))
return _Out(n_tokens)
def test_none_check_is_noop(self):
outs = [self._fake_output(100)]
EmbedIOProcessor._check_cohere_max_tokens(outs, None)
def test_within_limit(self):
outs = [self._fake_output(5), self._fake_output(3)]
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
def test_exceeds_limit(self):
outs = [self._fake_output(3), self._fake_output(10)]
with pytest.raises(ValueError, match="exceeds max_tokens=5"):
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
def test_exact_limit(self):
outs = [self._fake_output(5)]
EmbedIOProcessor._check_cohere_max_tokens(outs, 5)
class TestValidateInputType:
"""Unit tests for EmbedIOProcessor._validate_input_type."""
@staticmethod
def _make_handler(task_instructions: dict[str, str] | None):
handler = object.__new__(EmbedIOProcessor)
handler.task_instructions = task_instructions
return handler
def test_none_input_type_always_accepted(self):
handler = self._make_handler(None)
handler._validate_input_type(None)
handler_with = self._make_handler({"query": "q: "})
handler_with._validate_input_type(None)
def test_no_prompts_rejects(self):
handler = self._make_handler(None)
with pytest.raises(ValueError, match="does not define any input_type"):
handler._validate_input_type("anything")
def test_known_type_accepted(self):
handler = self._make_handler({"query": "q: ", "document": "d: "})
handler._validate_input_type("query")
handler._validate_input_type("document")
def test_unknown_type_rejected(self):
handler = self._make_handler({"query": "q: ", "document": "d: "})
with pytest.raises(ValueError, match="Unsupported input_type 'other'"):
handler._validate_input_type("other")
def test_error_lists_supported(self):
handler = self._make_handler({"a": "", "b": ""})
with pytest.raises(ValueError, match="Supported values: a, b"):
handler._validate_input_type("z")

View File

@@ -0,0 +1,72 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref
import pytest
import torch
import torch.nn.functional as F
from vllm import LLM, PoolingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.platforms import current_platform
MODEL_NAME = "intfloat/multilingual-e5-small"
prompts = ["The chef prepared a delicious meal."]
@pytest.fixture(scope="module")
def llm():
# ROCm: Use FLEX_ATTENTION backend as it's the only attention backend
# that supports encoder-only models on ROCm.
attention_config = None
if current_platform.is_rocm():
attention_config = {"backend": "FLEX_ATTENTION"}
# pytest caches the fixture so we use weakref.proxy to
# enable garbage collection
llm = LLM(
model=MODEL_NAME,
max_num_batched_tokens=32768,
tensor_parallel_size=1,
gpu_memory_utilization=0.75,
enforce_eager=True,
seed=0,
attention_config=attention_config,
)
yield weakref.proxy(llm)
del llm
cleanup_dist_env_and_memory()
@pytest.mark.skip_global_cleanup
def test_token_embed(llm: LLM):
outputs = llm.encode(prompts, pooling_task="token_embed", use_tqdm=False)
multi_vector = outputs[0].outputs.data
assert multi_vector.shape == (11, 384)
def test_pooling_params(llm: LLM):
def get_outputs(normalize):
outputs = llm.embed(
prompts,
pooling_params=PoolingParams(use_activation=normalize),
use_tqdm=False,
)
return torch.tensor([x.outputs.embedding for x in outputs])
default = get_outputs(normalize=None)
w_normal = get_outputs(normalize=True)
wo_normal = get_outputs(normalize=False)
assert torch.allclose(default, w_normal, atol=1e-2), "Default should use normal."
assert not torch.allclose(w_normal, wo_normal, atol=1e-2), (
"wo_normal should not use normal."
)
assert torch.allclose(w_normal, F.normalize(wo_normal, p=2, dim=-1), atol=1e-2), (
"w_normal should be close to normal(wo_normal)."
)

View File

@@ -0,0 +1,770 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import base64
import json
import numpy as np
import openai
import pytest
import pytest_asyncio
import requests
import torch
import torch.nn.functional as F
from tests.models.language.pooling.embed_utils import run_embedding_correctness_test
from tests.models.utils import check_embeddings_close
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
from vllm.entrypoints.pooling.utils import (
MetadataItem,
build_metadata_items,
decode_pooling_output,
)
from vllm.platforms import current_platform
from vllm.tokenizers import get_tokenizer
from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor
MODEL_NAME = "intfloat/multilingual-e5-small"
DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
DTYPE = "bfloat16"
input_text = "The best thing about vLLM is that it supports many different models"
input_tokens = [
0,
581,
2965,
13580,
1672,
81,
23708,
594,
83,
450,
442,
8060,
7,
5941,
12921,
115774,
2,
]
if current_platform.is_rocm():
# Disable Flash/MemEfficient SDP on ROCm to avoid HF Transformers
# accuracy issues: https://github.com/vllm-project/vllm/issues/30167
# TODO: Remove once ROCm SDP accuracy issues are resolved on HuggingFace
torch.backends.cuda.enable_flash_sdp(False)
torch.backends.cuda.enable_mem_efficient_sdp(False)
torch.backends.cuda.enable_math_sdp(True)
# On ROCm, floating-point reductions in attention and GEMM kernels are
# non-associative and sensitive to batch geometry. Force LLM instances
# into an identical, deterministic execution mode:
ROCM_DETERMINISM_ARGS: list[str] = (
["--max-num-seqs", "1"] if current_platform.is_rocm() else []
)
@pytest.fixture(scope="module")
def server():
args = [
"--runner",
"pooling",
"--dtype",
DTYPE,
"--enforce-eager",
"--max-model-len",
"512",
"--chat-template",
DUMMY_CHAT_TEMPLATE,
*ROCM_DETERMINISM_ARGS,
]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def client(server):
async with server.get_async_client() as async_client:
yield async_client
@pytest.fixture(scope="module")
def hf_model(hf_runner):
with hf_runner(MODEL_NAME, dtype=DTYPE, is_sentence_transformer=True) as hf_model:
yield hf_model
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_basic(
server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
):
# test /v1/models
response = requests.get(server.url_for("/v1/models"))
model = response.json()["data"][0]["id"]
assert model == MODEL_NAME
models = await client.models.list()
models = models.data
served_model = models[0]
assert served_model.id == MODEL_NAME
# test /tokenize
response = requests.post(
server.url_for("/tokenize"),
json={"model": model_name, "prompt": input_text},
)
assert response.json()["tokens"] == input_tokens
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_completion_request(
client: openai.AsyncOpenAI, model_name: str, hf_model
):
# test input: str
embedding_response = await client.embeddings.create(
model=model_name,
input=input_text,
encoding_format="float",
)
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == 1
assert len(embeddings.data[0].embedding) == 384
assert embeddings.usage.completion_tokens == 0
assert embeddings.usage.prompt_tokens == len(input_tokens)
assert embeddings.usage.total_tokens == len(input_tokens)
vllm_outputs = [d.embedding for d in embeddings.data]
run_embedding_correctness_test(hf_model, [input_text], vllm_outputs)
# test input: list[int]
embedding_response = await client.embeddings.create(
model=model_name,
input=input_tokens,
encoding_format="float",
)
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == 1
assert len(embeddings.data[0].embedding) == 384
assert embeddings.usage.completion_tokens == 0
assert embeddings.usage.prompt_tokens == len(input_tokens)
assert embeddings.usage.total_tokens == len(input_tokens)
vllm_outputs = [d.embedding for d in embeddings.data]
run_embedding_correctness_test(hf_model, [input_text], vllm_outputs)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_completion_request_batched(
client: openai.AsyncOpenAI, model_name: str, hf_model
):
N = 10
input_texts = [input_text] * N
# test input: list[str]
embedding_response = await client.embeddings.create(
model=model_name,
input=input_texts,
encoding_format="float",
)
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == N
assert len(embeddings.data[0].embedding) == 384
assert embeddings.usage.completion_tokens == 0
assert embeddings.usage.prompt_tokens == len(input_tokens) * N
assert embeddings.usage.total_tokens == len(input_tokens) * N
vllm_outputs = [d.embedding for d in embeddings.data]
run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
# test list[list[int]]
embedding_response = await client.embeddings.create(
model=model_name,
input=[input_tokens] * N,
encoding_format="float",
)
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == N
assert len(embeddings.data[0].embedding) == 384
assert embeddings.usage.completion_tokens == 0
assert embeddings.usage.prompt_tokens == len(input_tokens) * N
assert embeddings.usage.total_tokens == len(input_tokens) * N
vllm_outputs = [d.embedding for d in embeddings.data]
run_embedding_correctness_test(hf_model, input_texts, vllm_outputs)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_truncate_prompt_tokens(client: openai.AsyncOpenAI, model_name: str):
input_texts = [
"Como o Brasil pode fomentar o desenvolvimento de modelos de IA?",
]
# test single embedding
embedding_response = await client.embeddings.create(
model=model_name, input=input_texts, extra_body={"truncate_prompt_tokens": 10}
)
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == 1
assert len(embeddings.data[0].embedding) == 384
assert embeddings.usage.completion_tokens == 0
assert embeddings.usage.prompt_tokens == 10
assert embeddings.usage.total_tokens == 10
input_tokens = [
1,
24428,
289,
18341,
26165,
285,
19323,
283,
289,
26789,
3871,
28728,
9901,
340,
2229,
385,
340,
315,
28741,
28804,
2,
]
embedding_response = await client.embeddings.create(
model=model_name, input=input_tokens, extra_body={"truncate_prompt_tokens": 10}
)
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == 1
assert len(embeddings.data[0].embedding) == 384
assert embeddings.usage.completion_tokens == 0
assert embeddings.usage.prompt_tokens == 10
assert embeddings.usage.total_tokens == 10
# invalid_truncate_prompt_tokens
input_texts = [
"Como o Brasil pode fomentar o desenvolvimento de modelos de IA?",
]
with pytest.raises(openai.BadRequestError):
response = await client.embeddings.create(
model=model_name,
input=input_texts,
extra_body={"truncate_prompt_tokens": 8193},
)
assert "error" in response.object
assert (
"truncate_prompt_tokens value is greater than max_model_len. "
"Please, select a smaller truncation size." in response.message
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_chat_request(
server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
):
messages = [
{
"role": "user",
"content": "The cat sat on the mat.",
},
{
"role": "assistant",
"content": "A feline was resting on a rug.",
},
{
"role": "user",
"content": "Stars twinkle brightly in the night sky.",
},
]
# test chat request basic usage
chat_response = requests.post(
server.url_for("v1/embeddings"),
json={
"model": model_name,
"messages": messages,
"encoding_format": "float",
},
)
chat_response.raise_for_status()
chat_embeddings = EmbeddingResponse.model_validate(chat_response.json())
tokenizer = get_tokenizer(tokenizer_name=model_name)
prompt = tokenizer.apply_chat_template(
messages,
chat_template=DUMMY_CHAT_TEMPLATE,
add_generation_prompt=True,
continue_final_message=False,
tokenize=False,
)
completion_response = await client.embeddings.create(
model=model_name,
input=prompt,
encoding_format="float",
# To be consistent with chat
extra_body={"add_special_tokens": False},
)
completion_embeddings = EmbeddingResponse.model_validate(
completion_response.model_dump(mode="json")
)
assert chat_embeddings.id is not None
assert completion_embeddings.id is not None
assert chat_embeddings.created <= completion_embeddings.created
# Use tolerance-based comparison for embeddings
check_embeddings_close(
embeddings_0_lst=[d.embedding for d in chat_embeddings.data],
embeddings_1_lst=[d.embedding for d in completion_embeddings.data],
name_0="chat",
name_1="completion",
)
assert chat_embeddings.model_dump(exclude={"id", "created", "data"}) == (
completion_embeddings.model_dump(exclude={"id", "created", "data"})
)
# test add_generation_prompt
response = requests.post(
server.url_for("v1/embeddings"),
json={"model": model_name, "messages": messages, "add_generation_prompt": True},
)
response.raise_for_status()
output = EmbeddingResponse.model_validate(response.json())
assert output.object == "list"
assert len(output.data) == 1
assert output.model == MODEL_NAME
assert output.usage.prompt_tokens == 34
# test continue_final_message
response = requests.post(
server.url_for("v1/embeddings"),
json={
"model": model_name,
"messages": messages,
"continue_final_message": True,
},
)
response.raise_for_status()
output = EmbeddingResponse.model_validate(response.json())
assert output.object == "list"
assert len(output.data) == 1
assert output.model == MODEL_NAME
assert output.usage.prompt_tokens == 33
# test add_special_tokens
response = requests.post(
server.url_for("v1/embeddings"),
json={"model": model_name, "messages": messages, "add_special_tokens": True},
)
response.raise_for_status()
output = EmbeddingResponse.model_validate(response.json())
assert output.object == "list"
assert len(output.data) == 1
assert output.model == MODEL_NAME
assert output.usage.prompt_tokens == 36
# test continue_final_message with add_generation_prompt
response = requests.post(
server.url_for("v1/embeddings"),
json={
"model": model_name,
"messages": messages,
"continue_final_message": True,
"add_generation_prompt": True,
},
)
assert (
"Cannot set both `continue_final_message` and `add_generation_prompt` to True."
in response.json()["error"]["message"]
)
@pytest.mark.asyncio
async def test_invocations_completion_request(
server: RemoteOpenAIServer, client: openai.AsyncOpenAI
):
request_args = {
"model": MODEL_NAME,
"input": input_text,
"encoding_format": "float",
}
completion_response = await client.embeddings.create(**request_args)
invocation_response = requests.post(
server.url_for("invocations"), json=request_args
)
invocation_response.raise_for_status()
completion_output = completion_response.model_dump()
invocation_output = invocation_response.json()
assert completion_output.keys() == invocation_output.keys()
for completion_data, invocation_data in zip(
completion_output["data"], invocation_output["data"]
):
assert completion_data.keys() == invocation_data.keys()
check_embeddings_close(
embeddings_0_lst=[completion_data["embedding"]],
embeddings_1_lst=[invocation_data["embedding"]],
name_0="completion",
name_1="invocation",
)
@pytest.mark.asyncio
async def test_invocations_chat_request(server: RemoteOpenAIServer):
messages = [
{
"role": "user",
"content": "The cat sat on the mat.",
},
{
"role": "assistant",
"content": "A feline was resting on a rug.",
},
{
"role": "user",
"content": "Stars twinkle brightly in the night sky.",
},
]
request_args = {
"model": MODEL_NAME,
"messages": messages,
"encoding_format": "float",
}
chat_response = requests.post(server.url_for("v1/embeddings"), json=request_args)
chat_response.raise_for_status()
invocation_response = requests.post(
server.url_for("invocations"), json=request_args
)
invocation_response.raise_for_status()
chat_output = chat_response.json()
invocation_output = invocation_response.json()
assert chat_output.keys() == invocation_output.keys()
for chat_data, invocation_data in zip(
chat_output["data"], invocation_output["data"]
):
assert chat_data.keys() == invocation_data.keys()
check_embeddings_close(
embeddings_0_lst=[chat_data["embedding"]],
embeddings_1_lst=[invocation_data["embedding"]],
name_0="chat",
name_1="invocation",
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_base64_embedding(hf_model, client: openai.AsyncOpenAI, model_name: str):
input_texts = [
"Hello my name is",
"The best thing about vLLM is that it supports many different models",
]
responses_float = await client.embeddings.create(
input=input_texts, model=model_name, encoding_format="float"
)
float_data = [d.embedding for d in responses_float.data]
run_embedding_correctness_test(hf_model, input_texts, float_data)
responses_base64 = await client.embeddings.create(
input=input_texts, model=model_name, encoding_format="base64"
)
base64_data = []
for data in responses_base64.data:
base64_data.append(
np.frombuffer(base64.b64decode(data.embedding), dtype="float32").tolist()
)
run_embedding_correctness_test(hf_model, input_texts, base64_data)
# Default response is float32 decoded from base64 by OpenAI Client
responses_default = await client.embeddings.create(
input=input_texts, model=model_name
)
default_data = [d.embedding for d in responses_default.data]
run_embedding_correctness_test(hf_model, input_texts, default_data)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_base64_embed_dtype_and_endianness(
server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
):
input_texts = [input_text] * 3
responses_float = await client.embeddings.create(
input=input_texts, model=model_name, encoding_format="float"
)
float_data = [d.embedding for d in responses_float.data]
for embed_dtype in EMBED_DTYPES:
for endianness in ENDIANNESS:
responses_base64 = requests.post(
server.url_for("/v1/embeddings"),
json={
"model": model_name,
"input": input_texts,
"encoding_format": "base64",
"embed_dtype": embed_dtype,
"endianness": endianness,
},
)
base64_data = []
for data in responses_base64.json()["data"]:
binary = base64.b64decode(data["embedding"])
tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
base64_data.append(tensor.to(torch.float32).tolist())
check_embeddings_close(
embeddings_0_lst=float_data,
embeddings_1_lst=base64_data,
name_0="float_data",
name_1="base64_data",
tol=1e-2,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_bytes_embed_dtype_and_endianness(
server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
):
input_texts = [input_text] * 3
responses_float = await client.embeddings.create(
input=input_texts, model=model_name, encoding_format="float"
)
float_data = [d.embedding for d in responses_float.data]
for embed_dtype in EMBED_DTYPES:
for endianness in ENDIANNESS:
responses_bytes = requests.post(
server.url_for("/v1/embeddings"),
json={
"model": model_name,
"input": input_texts,
"encoding_format": "bytes",
"embed_dtype": embed_dtype,
"endianness": endianness,
},
)
metadata = json.loads(responses_bytes.headers["metadata"])
body = responses_bytes.content
items = [MetadataItem(**x) for x in metadata["data"]]
bytes_data = decode_pooling_output(items=items, body=body)
bytes_data = [x.to(torch.float32).tolist() for x in bytes_data]
check_embeddings_close(
embeddings_0_lst=float_data,
embeddings_1_lst=bytes_data,
name_0="float_data",
name_1="bytes_data",
tol=1e-2,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_bytes_only_embed_dtype_and_endianness(
server: RemoteOpenAIServer, client: openai.AsyncOpenAI, model_name: str
):
input_texts = [
"The best thing about vLLM is that it supports many different models",
] * 2
responses_float = await client.embeddings.create(
input=input_texts, model=model_name, encoding_format="float"
)
float_data = [d.embedding for d in responses_float.data]
embedding_size = len(float_data[0])
for embed_dtype in EMBED_DTYPES:
for endianness in ENDIANNESS:
responses_bytes = requests.post(
server.url_for("/v1/embeddings"),
json={
"model": model_name,
"input": input_texts,
"encoding_format": "bytes_only",
"embed_dtype": embed_dtype,
"endianness": endianness,
},
)
assert "metadata" not in responses_bytes.headers
body = responses_bytes.content
items = build_metadata_items(
embed_dtype=embed_dtype,
endianness=endianness,
shape=(embedding_size,),
n_request=len(input_texts),
)
bytes_data = decode_pooling_output(items=items, body=body)
bytes_data = [x.to(torch.float32).tolist() for x in bytes_data]
check_embeddings_close(
embeddings_0_lst=float_data,
embeddings_1_lst=bytes_data,
name_0="float_data",
name_1="bytes_data",
tol=1e-2,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("param_name", ["encoding_format", "embed_dtype", "endianness"])
async def test_params_not_supported(
server: RemoteOpenAIServer, model_name: str, param_name: str
):
responses_base64 = requests.post(
server.url_for("/v1/embeddings"),
json={
"model": model_name,
"input": input_text,
"encoding_format": "base64",
param_name: f"bad_{param_name}",
},
)
assert responses_base64.status_code == 400
assert "literal_error" in responses_base64.json()["error"]["message"]
assert f"bad_{param_name}" in responses_base64.json()["error"]["message"]
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
async def get_outputs(use_activation):
request_args = {
"model": MODEL_NAME,
"input": input_text,
"encoding_format": "float",
"use_activation": use_activation,
}
response = requests.post(server.url_for("v1/embeddings"), json=request_args)
outputs = response.json()
return torch.tensor([x["embedding"] for x in outputs["data"]])
default = await get_outputs(use_activation=None)
w_normal = await get_outputs(use_activation=True)
wo_normal = await get_outputs(use_activation=False)
assert torch.allclose(default, w_normal, atol=1e-2), "Default should use normal."
assert not torch.allclose(w_normal, wo_normal, atol=1e-2), (
"wo_normal should not use normal."
)
assert torch.allclose(w_normal, F.normalize(wo_normal, p=2, dim=-1), atol=1e-2), (
"w_normal should be close to normal(wo_normal)."
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_pooling_embed(server: RemoteOpenAIServer, model_name: str):
task = "embed"
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": input_text,
"encoding_format": "float",
"task": task,
},
)
poolings = PoolingResponse.model_validate(response.json())
assert len(poolings.data) == 1
assert len(poolings.data[0].data) == 384
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_pooling_token_embed(server: RemoteOpenAIServer, model_name: str):
task = "token_embed"
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": input_text,
"encoding_format": "float",
"task": task,
},
)
poolings = PoolingResponse.model_validate(response.json())
assert len(poolings.data) == 1
assert len(poolings.data[0].data) == len(input_tokens)
assert len(poolings.data[0].data[0]) == 384
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("task", ["classify", "token_classify", "plugin"])
async def test_pooling_not_supported(
server: RemoteOpenAIServer, model_name: str, task: str
):
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": "test",
"encoding_format": "float",
"task": task,
},
)
assert response.json()["error"]["type"] == "BadRequestError"
assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")

View File

@@ -0,0 +1,131 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Run `pytest tests/entrypoints/openai/test_embedding_dimensions.py`.
"""
import openai
import pytest
from tests.conftest import HfRunner
from tests.models.language.pooling.embed_utils import run_embedding_correctness_test
from tests.models.utils import EmbedModelInfo
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
from vllm.platforms import current_platform
MODELS = [
EmbedModelInfo("intfloat/multilingual-e5-small", is_matryoshka=False),
EmbedModelInfo(
"Snowflake/snowflake-arctic-embed-m-v1.5",
is_matryoshka=True,
matryoshka_dimensions=[256],
),
]
input_texts = [
"The chef prepared a delicious meal.",
]
@pytest.fixture(scope="module", params=MODELS)
def model_info(request):
return request.param
@pytest.fixture(scope="module", params=["bfloat16"])
def dtype(request):
return request.param
@pytest.fixture(scope="module")
def server(model_info, dtype: str):
args = [
"--runner",
"pooling",
# use half precision for speed and memory savings in CI environment
"--dtype",
dtype,
"--enforce-eager",
"--max-model-len",
"512",
]
if model_info.name == "Snowflake/snowflake-arctic-embed-m-v1.5":
# Manually enable Matryoshka Embeddings
args.extend(
["--trust_remote_code", "--hf_overrides", '{"matryoshka_dimensions":[256]}']
)
# ROCm: Use Flex Attention to support encoder-only self-attention.
if current_platform.is_rocm():
args.extend(["--attention-backend", "FLEX_ATTENTION"])
with RemoteOpenAIServer(model_info.name, args) as remote_server:
yield remote_server
@pytest.fixture(scope="module")
def hf_model(hf_runner, model_info, dtype: str):
with hf_runner(
model_info.name, dtype=dtype, is_sentence_transformer=True
) as hf_model:
yield hf_model
@pytest.mark.asyncio
async def test_matryoshka(
model_info: EmbedModelInfo, server: RemoteOpenAIServer, hf_model: HfRunner
):
client = server.get_async_client()
async def make_request_and_correctness_test(dimensions):
prompts = input_texts * 3
embedding_response = await client.embeddings.create(
model=model_info.name,
input=prompts,
dimensions=dimensions,
encoding_format="float",
)
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == 3
assert len(embeddings.data[0].embedding) > 0
assert embeddings.usage.completion_tokens == 0
assert embeddings.usage.prompt_tokens > 0
assert embeddings.usage.total_tokens > 0
if dimensions is not None:
assert len(embeddings.data[0].embedding) == dimensions
vllm_outputs = [d.embedding for d in embeddings.data]
run_embedding_correctness_test(hf_model, prompts, vllm_outputs, dimensions)
if model_info.is_matryoshka:
valid_dimensions: list[int | None] = [None]
if model_info.matryoshka_dimensions is not None:
valid_dimensions += model_info.matryoshka_dimensions[:2]
for dimensions in valid_dimensions:
await make_request_and_correctness_test(dimensions)
invalid_dimensions: list[int | None] = [-1]
if model_info.matryoshka_dimensions is not None:
assert 5 not in model_info.matryoshka_dimensions
invalid_dimensions.append(5)
for dimensions in invalid_dimensions:
with pytest.raises(openai.BadRequestError):
await make_request_and_correctness_test(dimensions)
else:
for dimensions in [None]:
await make_request_and_correctness_test(dimensions)
for dimensions in [-1, 16]:
with pytest.raises(openai.BadRequestError):
await make_request_and_correctness_test(dimensions)

View File

@@ -0,0 +1,457 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Test cases for long text embedding with automatic chunking mechanism.
This test suite validates vLLM's automatic chunking functionality for handling
text inputs that exceed the model's maximum token length, specifically targeting
the intfloat/multilingual-e5-small model (max token length: 512).
"""
import random
import openai
import pytest
import pytest_asyncio
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
from vllm.platforms import current_platform
def _generate_random_text(word_count: int) -> str:
"""Generate random text with approximately the specified word count."""
# Common English words with focus on verbs and nouns for realistic text
common_words = [
# Essential articles and pronouns (minimal)
"the",
"and",
"you",
"they",
"this",
"that",
"these",
"those",
# Action verbs
"create",
"build",
"develop",
"design",
"implement",
"execute",
"analyze",
"process",
"generate",
"calculate",
"evaluate",
"optimize",
"transform",
"integrate",
"configure",
"deploy",
"monitor",
"manage",
"discover",
"explore",
"investigate",
"research",
"study",
"examine",
"improve",
"enhance",
"upgrade",
"modify",
"update",
"maintain",
"solve",
"resolve",
"handle",
"address",
"tackle",
"overcome",
"communicate",
"collaborate",
"coordinate",
"organize",
"plan",
"achieve",
"accomplish",
"complete",
"finish",
"deliver",
"provide",
# Technology and science nouns
"system",
"application",
"software",
"hardware",
"network",
"database",
"algorithm",
"model",
"framework",
"platform",
"interface",
"protocol",
"architecture",
"infrastructure",
"component",
"module",
"service",
"technology",
"innovation",
"solution",
"methodology",
"approach",
"artificial",
"intelligence",
"machine",
"learning",
"neural",
"network",
"computer",
"processor",
"memory",
"storage",
"computation",
"data",
"information",
"knowledge",
"insight",
"pattern",
"trend",
"analysis",
"research",
"development",
"engineering",
"science",
"mathematics",
"statistics",
"probability",
"optimization",
"performance",
"efficiency",
# General nouns
"project",
"team",
"organization",
"company",
"business",
"industry",
"market",
"customer",
"user",
"client",
"product",
"feature",
"function",
"requirement",
"specification",
"documentation",
"report",
"result",
"outcome",
"impact",
"benefit",
"advantage",
"challenge",
"problem",
"opportunity",
"strategy",
"goal",
"objective",
"target",
"milestone",
"process",
"procedure",
"workflow",
"pipeline",
"operation",
"task",
"activity",
"event",
"session",
"meeting",
"discussion",
"decision",
]
words = []
for _ in range(word_count):
words.append(random.choice(common_words))
# Add some punctuation for more realistic text
text = " ".join(words)
# Add periods every 10-20 words
words_list = text.split()
result = []
for i, word in enumerate(words_list):
result.append(word)
if (i + 1) % random.randint(10, 20) == 0 and i < len(words_list) - 1:
result[-1] += "."
return " ".join(result)
MODEL_NAME = "intfloat/multilingual-e5-small"
DTYPE = "bfloat16"
# Test text: Generate text with approximately 1500 words to exceed 1024 tokens
LONG_TEXT_1500_WORDS = _generate_random_text(1500)
# Test text: Generate text with approximately 2500 words to exceed 2048 tokens
LONG_TEXT_2500_WORDS = _generate_random_text(2500)
@pytest.fixture(scope="module")
def server_with_chunked_processing():
"""Start server with automatic chunking processing enabled."""
args = [
"--runner",
"pooling",
"--dtype",
DTYPE,
"--enforce-eager",
"--max-model-len",
"512", # Set smaller max_model_len to trigger chunking mechanism
"--pooler-config",
(
'{"pooling_type": "MEAN", "use_activation": true, '
'"enable_chunked_processing": true, "max_embed_len": 10000}'
),
"--gpu-memory-utilization",
"0.8",
]
# ROCm: Use Flex Attention to support encoder-only self-attention.
if current_platform.is_rocm():
args.extend(["--attention-backend", "FLEX_ATTENTION"])
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
@pytest_asyncio.fixture
async def client_with_chunked_processing(server_with_chunked_processing):
"""Create async client with chunking processing support."""
async with server_with_chunked_processing.get_async_client() as async_client:
yield async_client
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_long_text_embedding_1500_chars(
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
):
"""Test embedding processing for ~1500 character long text
(~1028 tokens, exceeding 512 token limit)."""
# Verify text length
# Verify text has sufficient word count (approximately 1500 words)
word_count = len(LONG_TEXT_1500_WORDS.split())
assert word_count >= 1400, f"Test text word count insufficient: {word_count} words"
# Send embedding request
embedding_response = await client_with_chunked_processing.embeddings.create(
model=model_name,
input=[LONG_TEXT_1500_WORDS],
encoding_format="float",
)
# Verify response structure
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == 1
assert (
len(embeddings.data[0].embedding) == 384
) # multilingual-e5-small embedding dimension
assert embeddings.usage.completion_tokens == 0
# Due to chunked processing, token count should
# reflect actual processed tokens
# With ~1500 words, we expect roughly
# 1024+ tokens (exceeding 512 token limit)
# Should exceed single chunk limit of 512
assert embeddings.usage.prompt_tokens > 800
assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
# Verify embedding vector validity
embedding_vector = embeddings.data[0].embedding
assert all(isinstance(x, float) for x in embedding_vector), (
"Embedding vector should contain floats"
)
assert not all(x == 0 for x in embedding_vector), (
"Embedding vector should not be all zeros"
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_long_text_embedding_2500_chars(
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
):
"""Test embedding processing for ~2500 character long text
(~2048 tokens, requiring multiple chunks)."""
# Verify text length
# Verify text has sufficient word count (approximately 2500 words)
word_count = len(LONG_TEXT_2500_WORDS.split())
assert word_count >= 2300, f"Test text word count insufficient: {word_count} words"
# Send embedding request
embedding_response = await client_with_chunked_processing.embeddings.create(
model=model_name,
input=[LONG_TEXT_2500_WORDS],
encoding_format="float",
)
# Verify response structure
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == 1
assert (
len(embeddings.data[0].embedding) == 384
) # multilingual-e5-small embedding dimension
assert embeddings.usage.completion_tokens == 0
# Due to chunked processing, token count should
# reflect actual processed tokens
# With ~2500 words, we expect
# roughly 2048+ tokens (requiring multiple chunks)
# Should require multiple chunks for processing
assert embeddings.usage.prompt_tokens > 1500
assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
# Verify embedding vector validity
embedding_vector = embeddings.data[0].embedding
assert all(isinstance(x, float) for x in embedding_vector), (
"Embedding vector should contain floats"
)
assert not all(x == 0 for x in embedding_vector), (
"Embedding vector should not be all zeros"
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_batch_long_text_embedding(
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
):
"""Test batch long text embedding processing."""
input_texts = [
LONG_TEXT_1500_WORDS,
LONG_TEXT_2500_WORDS,
"This is a short text test.", # Short text for comparison
]
# Send batch embedding request
embedding_response = await client_with_chunked_processing.embeddings.create(
model=model_name,
input=input_texts,
encoding_format="float",
)
# Verify response structure
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == 3 # Three input texts
# Verify each embedding dimension
for i, embedding_data in enumerate(embeddings.data):
assert len(embedding_data.embedding) == 384
assert embedding_data.index == i
# Verify embedding vector validity
embedding_vector = embedding_data.embedding
assert all(isinstance(x, float) for x in embedding_vector)
assert not all(x == 0 for x in embedding_vector)
# Verify token usage
assert embeddings.usage.completion_tokens == 0
# Total token count should be very substantial
assert embeddings.usage.prompt_tokens > 1000
assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_chunked_vs_normal_consistency(
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
):
"""Test consistency between chunked and
normal processing (using short text)."""
# Use a short text within the 512 token limit
short_text = (
"Artificial intelligence technology is changing our world, "
"bringing unprecedented opportunities and challenges."
)
# Send embedding request
embedding_response = await client_with_chunked_processing.embeddings.create(
model=model_name,
input=[short_text],
encoding_format="float",
)
# Verify response structure
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == 1
assert len(embeddings.data[0].embedding) == 384
assert embeddings.usage.completion_tokens == 0
# Short text should not require chunked processing
assert embeddings.usage.prompt_tokens < 512
assert embeddings.usage.total_tokens == embeddings.usage.prompt_tokens
# 验证embedding向量的有效性
embedding_vector = embeddings.data[0].embedding
assert all(isinstance(x, float) for x in embedding_vector)
assert not all(x == 0 for x in embedding_vector)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_chunked_processing_response_format(
client_with_chunked_processing: openai.AsyncOpenAI, model_name: str
):
"""Test response format and structure during chunked processing."""
# Test with long text to trigger chunking
embedding_response = await client_with_chunked_processing.embeddings.create(
model=model_name,
input=[LONG_TEXT_1500_WORDS],
encoding_format="float",
)
# Verify response structure
embeddings = EmbeddingResponse.model_validate(
embedding_response.model_dump(mode="json")
)
assert embeddings.id is not None
assert len(embeddings.data) == 1
assert embeddings.data[0].object == "embedding"
assert embeddings.data[0].index == 0
# Verify embedding vector properties
embedding_vector = embeddings.data[0].embedding
import math
vector_norm = math.sqrt(sum(x * x for x in embedding_vector))
# Check that the vector is normalized
# (default behavior for most embedding models)
assert 0.8 < vector_norm < 1.2, (
f"Vector norm should be reasonable, actual: {vector_norm}"
)

View File

@@ -0,0 +1,210 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import pytest
import requests
from transformers import AutoProcessor
from tests.utils import VLLM_PATH, RemoteOpenAIServer
from vllm.entrypoints.pooling.embed.protocol import EmbeddingResponse
from vllm.multimodal.media import MediaWithBytes
from vllm.multimodal.utils import encode_image_url, fetch_image
MODEL_NAME = "TIGER-Lab/VLM2Vec-Full"
MAXIMUM_IMAGES = 2
vlm2vec_jinja_path = VLLM_PATH / "examples/pooling/embed/template/vlm2vec_phi3v.jinja"
assert vlm2vec_jinja_path.exists()
# Test different image extensions (JPG/PNG) and formats (gray/RGB/RGBA)
TEST_IMAGE_ASSETS = [
"2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg"
"Grayscale_8bits_palette_sample_image.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/Grayscale_8bits_palette_sample_image.png",
"1280px-Venn_diagram_rgb.svg.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/1280px-Venn_diagram_rgb.svg.png",
"RGBA_comp.png", # "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/RGBA_comp.png",
]
input_text = "The best thing about vLLM is that it supports many different models"
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
image_base64 = {"url": encode_image_url(fetch_image(image_url))}
@pytest.fixture(scope="module")
def server():
args = [
"--runner",
"pooling",
"--max-model-len",
"2048",
"--max-num-seqs",
"5",
"--enforce-eager",
"--trust-remote-code",
"--limit-mm-per-prompt",
json.dumps({"image": MAXIMUM_IMAGES}),
"--chat-template",
str(vlm2vec_jinja_path),
]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_chat_text_request(server: RemoteOpenAIServer, model_name: str):
messages = [
{
"role": "user",
"content": input_text,
},
]
# note: vlm2vec_phi3v.jinja
# Embedding models should only embed one message at a time.
response = requests.post(
server.url_for("v1/embeddings"),
json={"model": model_name, "messages": messages},
)
response.raise_for_status()
output = EmbeddingResponse.model_validate(response.json())
assert len(output.data) == 1
assert output.model == MODEL_NAME
assert len(output.data[0].embedding) == 3072
assert output.usage.prompt_tokens == 14
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_chat_image_url_request(server: RemoteOpenAIServer, model_name: str):
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Represent the user's input."},
{"type": "image_url", "image_url": {"url": image_url}},
],
}
]
response = requests.post(
server.url_for("v1/embeddings"),
json={"model": model_name, "messages": messages},
)
response.raise_for_status()
output = EmbeddingResponse.model_validate(response.json())
assert len(output.data) == 1
assert output.model == MODEL_NAME
assert len(output.data[0].embedding) == 3072
assert output.usage.prompt_tokens == 767
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_chat_image_base64_request(server: RemoteOpenAIServer, model_name: str):
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Represent the user's input."},
{"type": "image_url", "image_url": image_base64},
],
}
]
response = requests.post(
server.url_for("v1/embeddings"),
json={"model": model_name, "messages": messages},
)
response.raise_for_status()
output = EmbeddingResponse.model_validate(response.json())
assert len(output.data) == 1
assert output.model == MODEL_NAME
assert len(output.data[0].embedding) == 3072
assert output.usage.prompt_tokens == 767
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_chat_image_with_media_io_kwargs(server: RemoteOpenAIServer, model_name: str):
rgba_image_url = (
"https://vllm-public-assets.s3.us-west-2.amazonaws.com"
"/vision_model_images/RGBA_comp.png"
)
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "Represent the user's input."},
{"type": "image_url", "image_url": {"url": rgba_image_url}},
],
}
]
response = requests.post(
server.url_for("v1/embeddings"),
json={
"model": model_name,
"messages": messages,
"media_io_kwargs": {
"image": {"rgba_background_color": [0, 0, 0]},
},
},
)
response.raise_for_status()
output = EmbeddingResponse.model_validate(response.json())
assert len(output.data) == 1
assert len(output.data[0].embedding) == 3072
def get_hf_prompt_tokens(model_name, content, image_url):
processor = AutoProcessor.from_pretrained(
model_name, trust_remote_code=True, num_crops=4
)
placeholder = "<|image_1|> "
prompt = f"{placeholder}{content}"
image = fetch_image(image_url)
# Unwrap MediaWithBytes if present
if isinstance(image, MediaWithBytes):
image = image.media
images = [image]
inputs = processor(prompt, images, return_tensors="pt")
return inputs.input_ids.shape[1]
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("image_url", TEST_IMAGE_ASSETS, indirect=True)
async def test_image_embedding(
server: RemoteOpenAIServer, model_name: str, image_url: str
):
content_text = "Represent the given image."
messages = [
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": image_url}},
{"type": "text", "text": content_text},
],
}
]
response = requests.post(
server.url_for("v1/embeddings"),
json={"model": model_name, "messages": messages, "encoding_format": "float"},
)
response.raise_for_status()
embeddings = EmbeddingResponse.model_validate(response.json())
hf_prompt_tokens = get_hf_prompt_tokens(model_name, content_text, image_url)
assert embeddings.id is not None
assert len(embeddings.data) == 1
assert len(embeddings.data[0].embedding) == 3072
assert embeddings.usage.completion_tokens == 0
assert embeddings.usage.prompt_tokens == hf_prompt_tokens
assert embeddings.usage.total_tokens == hf_prompt_tokens

View File

@@ -0,0 +1,129 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for Cohere embed protocol: build_typed_embeddings and its
underlying packing helpers, plus Cohere-specific serving helpers."""
import base64
import struct
import numpy as np
import pytest
from vllm.entrypoints.pooling.embed.protocol import (
build_typed_embeddings,
)
@pytest.fixture
def sample_embeddings() -> list[list[float]]:
return [
[0.1, -0.2, 0.3, -0.4, 0.5, -0.6, 0.7, -0.8],
[-0.05, 0.15, -0.25, 0.35, -0.45, 0.55, -0.65, 0.75],
]
class TestBuildTypedEmbeddingsFloat:
def test_float_passthrough(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(sample_embeddings, ["float"])
assert result.float == sample_embeddings
assert result.binary is None
def test_empty_input(self):
result = build_typed_embeddings([], ["float"])
assert result.float == []
class TestBuildTypedEmbeddingsBinary:
def test_binary_packing(self):
# 8 values: positive->1, negative->0 => bits: 10101010 = 0xAA = 170
# signed: 170 - 128 = 42
embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
result = build_typed_embeddings(embs, ["binary"])
assert result.binary is not None
assert result.binary[0] == [42]
def test_ubinary_packing(self):
embs = [[1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0]]
result = build_typed_embeddings(embs, ["ubinary"])
assert result.ubinary is not None
assert result.ubinary[0] == [170] # 0b10101010
def test_binary_all_positive(self):
embs = [[0.1] * 8]
result = build_typed_embeddings(embs, ["binary"])
assert result.binary is not None
# all bits = 1 => 0xFF = 255, signed: 255 - 128 = 127
assert result.binary[0] == [127]
def test_binary_all_negative(self):
embs = [[-0.1] * 8]
result = build_typed_embeddings(embs, ["binary"])
assert result.binary is not None
# all bits = 0, signed: 0 - 128 = -128
assert result.binary[0] == [-128]
def test_binary_dimension_is_eighth(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(sample_embeddings, ["binary"])
assert result.binary is not None
for orig, packed in zip(sample_embeddings, result.binary):
assert len(packed) == len(orig) // 8
def test_zero_treated_as_positive(self):
embs = [[0.0] * 8]
result = build_typed_embeddings(embs, ["binary"])
assert result.binary is not None
# 0.0 >= 0 is True, so bit=1 for all => 127 (signed)
assert result.binary[0] == [127]
def test_non_multiple_of_8_raises(self):
embs = [[0.1] * 7]
with pytest.raises(ValueError, match="multiple of 8"):
build_typed_embeddings(embs, ["binary"])
def test_ubinary_non_multiple_of_8_raises(self):
embs = [[0.1] * 10]
with pytest.raises(ValueError, match="multiple of 8"):
build_typed_embeddings(embs, ["ubinary"])
class TestBuildTypedEmbeddingsBase64:
def test_base64_roundtrip(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(sample_embeddings, ["base64"])
assert result.base64 is not None
assert len(result.base64) == 2
for orig, b64_str in zip(sample_embeddings, result.base64):
decoded = base64.b64decode(b64_str)
n = len(orig)
values = struct.unpack(f"<{n}f", decoded)
np.testing.assert_allclose(orig, values, rtol=1e-5)
def test_base64_byte_length(self):
embs = [[0.1, 0.2, 0.3]]
result = build_typed_embeddings(embs, ["base64"])
assert result.base64 is not None
raw = base64.b64decode(result.base64[0])
assert len(raw) == 3 * 4 # 3 floats * 4 bytes each
class TestBuildTypedEmbeddingsMultiple:
def test_all_types_at_once(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(
sample_embeddings,
["float", "binary", "ubinary", "base64"],
)
assert result.float is not None
assert result.binary is not None
assert result.ubinary is not None
assert result.base64 is not None
def test_subset_types(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(sample_embeddings, ["float", "binary"])
assert result.float is not None
assert result.binary is not None
assert result.ubinary is None
assert result.base64 is None
def test_unknown_type_ignored(self, sample_embeddings: list[list[float]]):
result = build_typed_embeddings(sample_embeddings, ["float", "unknown_type"])
assert result.float is not None

View File

@@ -0,0 +1,569 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import base64
import json
import numpy as np
import pytest
import requests
import torch
from tests.models.utils import check_embeddings_close
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
from vllm.entrypoints.pooling.utils import (
MetadataItem,
build_metadata_items,
decode_pooling_output,
)
from vllm.tokenizers import get_tokenizer
from vllm.utils.serial_utils import EMBED_DTYPES, ENDIANNESS, binary2tensor
MODEL_NAME = "internlm/internlm2-1_8b-reward"
DUMMY_CHAT_TEMPLATE = """{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\\n'}}{% endfor %}""" # noqa: E501
input_text = "The chef prepared a delicious meal."
input_tokens = [1, 918, 29981, 10166, 395, 18067, 15265, 281]
@pytest.fixture(scope="module")
def server():
args = [
"--runner",
"pooling",
# use half precision for speed and memory savings in CI environment
"--dtype",
"bfloat16",
"--enforce-eager",
"--max-model-len",
"512",
"--chat-template",
DUMMY_CHAT_TEMPLATE,
"--trust-remote-code",
]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_basic(server: RemoteOpenAIServer, model_name: str):
# test /v1/models
response = requests.get(server.url_for("/v1/models"))
served_model = response.json()["data"][0]["id"]
assert served_model == MODEL_NAME
# test /tokenize
response = requests.post(
server.url_for("/tokenize"),
json={"model": model_name, "prompt": input_text},
)
assert response.json()["tokens"] == input_tokens
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_completion_request(server: RemoteOpenAIServer, model_name: str):
# test input: str
response = requests.post(
server.url_for("pooling"),
json={"model": model_name, "input": input_text, "encoding_format": "float"},
)
response.raise_for_status()
poolings = PoolingResponse.model_validate(response.json())
assert poolings.id is not None
assert len(poolings.data) == 1
assert len(poolings.data[0].data) == len(input_tokens)
assert poolings.usage.completion_tokens == 0
assert poolings.usage.prompt_tokens == len(input_tokens)
assert poolings.usage.total_tokens == len(input_tokens)
# test input: list[int]
response = requests.post(
server.url_for("pooling"),
json={"model": model_name, "input": input_tokens, "encoding_format": "float"},
)
response.raise_for_status()
poolings = PoolingResponse.model_validate(response.json())
assert poolings.id is not None
assert len(poolings.data) == 1
assert len(poolings.data[0].data) == len(input_tokens)
assert poolings.usage.completion_tokens == 0
assert poolings.usage.prompt_tokens == len(input_tokens)
assert poolings.usage.total_tokens == len(input_tokens)
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_completion_request_batched(server: RemoteOpenAIServer, model_name: str):
N = 10
input_texts = [input_text] * N
response = requests.post(
server.url_for("pooling"),
json={"model": model_name, "input": input_texts, "encoding_format": "float"},
)
response.raise_for_status()
poolings = PoolingResponse.model_validate(response.json())
assert poolings.id is not None
assert len(poolings.data) == N
assert len(poolings.data[0].data) == len(input_tokens)
assert poolings.usage.completion_tokens == 0
assert poolings.usage.prompt_tokens == len(input_tokens) * N
assert poolings.usage.total_tokens == len(input_tokens) * N
# test list[list[int]]
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": [input_tokens] * N,
"encoding_format": "float",
},
)
response.raise_for_status()
poolings = PoolingResponse.model_validate(response.json())
assert poolings.id is not None
assert len(poolings.data) == N
assert len(poolings.data[0].data) == len(input_tokens)
assert poolings.usage.completion_tokens == 0
assert poolings.usage.prompt_tokens == len(input_tokens) * N
assert poolings.usage.total_tokens == len(input_tokens) * N
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_chat_request(server: RemoteOpenAIServer, model_name: str):
messages = [
{
"role": "user",
"content": "The cat sat on the mat.",
},
{
"role": "assistant",
"content": "A feline was resting on a rug.",
},
{
"role": "user",
"content": "Stars twinkle brightly in the night sky.",
},
]
# test chat request basic usage
chat_response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"messages": messages,
"encoding_format": "float",
},
)
chat_response.raise_for_status()
chat_poolings = PoolingResponse.model_validate(chat_response.json())
tokenizer = get_tokenizer(tokenizer_name=model_name, trust_remote_code=True)
prompt = tokenizer.apply_chat_template(
messages,
chat_template=DUMMY_CHAT_TEMPLATE,
add_generation_prompt=True,
continue_final_message=False,
tokenize=False,
)
completions_response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": prompt,
"encoding_format": "float",
# To be consistent with chat
"add_special_tokens": False,
},
)
completions_response.raise_for_status()
completion_poolings = PoolingResponse.model_validate(completions_response.json())
assert chat_poolings.id is not None
assert completion_poolings.id is not None
assert chat_poolings.created <= completion_poolings.created
assert chat_poolings.model_dump(exclude={"id", "created"}) == (
completion_poolings.model_dump(exclude={"id", "created"})
)
# test add_generation_prompt
response = requests.post(
server.url_for("pooling"),
json={"model": model_name, "messages": messages, "add_generation_prompt": True},
)
response.raise_for_status()
output = PoolingResponse.model_validate(response.json())
assert output.object == "list"
assert len(output.data) == 1
assert output.model == MODEL_NAME
assert output.usage.prompt_tokens == 33
# test continue_final_message
# The continue_final_message parameter doesn't seem to be working with this model.
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"messages": messages,
"continue_final_message": True,
},
)
response.raise_for_status()
output = PoolingResponse.model_validate(response.json())
assert output.object == "list"
assert len(output.data) == 1
assert output.model == MODEL_NAME
assert output.usage.prompt_tokens == 33
# test add_special_tokens
response = requests.post(
server.url_for("pooling"),
json={"model": model_name, "messages": messages, "add_special_tokens": True},
)
response.raise_for_status()
output = PoolingResponse.model_validate(response.json())
assert output.object == "list"
assert len(output.data) == 1
assert output.model == MODEL_NAME
assert output.usage.prompt_tokens == 34
# test continue_final_message with add_generation_prompt
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"messages": messages,
"continue_final_message": True,
"add_generation_prompt": True,
},
)
assert (
"Cannot set both `continue_final_message` and `add_generation_prompt` to True."
in response.json()["error"]["message"]
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_batch_base64_pooling(server: RemoteOpenAIServer, model_name: str):
input_texts = [
"Hello my name is",
"The best thing about vLLM is that it supports many different models",
]
float_response = requests.post(
server.url_for("pooling"),
json={
"input": input_texts,
"model": model_name,
"encoding_format": "float",
},
)
float_response.raise_for_status()
responses_float = PoolingResponse.model_validate(float_response.json())
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
base64_response = requests.post(
server.url_for("pooling"),
json={
"input": input_texts,
"model": model_name,
"encoding_format": "base64",
},
)
base64_response.raise_for_status()
responses_base64 = PoolingResponse.model_validate(base64_response.json())
decoded_responses_base64_data = []
for data in responses_base64.data:
decoded_responses_base64_data.append(
np.frombuffer(base64.b64decode(data.data), dtype="float32").tolist()
)
check_embeddings_close(
embeddings_0_lst=float_data,
embeddings_1_lst=decoded_responses_base64_data,
name_0="float32",
name_1="base64",
)
# Default response is float32 decoded from base64 by OpenAI Client
default_response = requests.post(
server.url_for("pooling"),
json={
"input": input_texts,
"model": model_name,
},
)
default_response.raise_for_status()
responses_default = PoolingResponse.model_validate(default_response.json())
default_data = [
np.array(d.data).squeeze(-1).tolist() for d in responses_default.data
]
check_embeddings_close(
embeddings_0_lst=float_data,
embeddings_1_lst=default_data,
name_0="float32",
name_1="default",
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_base64_embed_dtype_and_endianness(
server: RemoteOpenAIServer, model_name: str
):
input_texts = [input_text] * 3
url = server.url_for("pooling")
float_response = requests.post(
url,
json={
"model": model_name,
"input": input_texts,
"encoding_format": "float",
},
)
responses_float = PoolingResponse.model_validate(float_response.json())
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
for embed_dtype in EMBED_DTYPES:
for endianness in ENDIANNESS:
responses_base64 = requests.post(
url,
json={
"model": model_name,
"input": input_texts,
"encoding_format": "base64",
"embed_dtype": embed_dtype,
"endianness": endianness,
},
)
base64_data = []
for data in responses_base64.json()["data"]:
binary = base64.b64decode(data["data"])
tensor = binary2tensor(binary, (-1,), embed_dtype, endianness)
base64_data.append(tensor.to(torch.float32).tolist())
check_embeddings_close(
embeddings_0_lst=float_data,
embeddings_1_lst=base64_data,
name_0="float_data",
name_1="base64_data",
tol=1e-2,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_bytes_embed_dtype_and_endianness(
server: RemoteOpenAIServer, model_name: str
):
input_texts = [input_text] * 3
url = server.url_for("pooling")
float_response = requests.post(
url,
json={
"model": model_name,
"input": input_texts,
"encoding_format": "float",
},
)
responses_float = PoolingResponse.model_validate(float_response.json())
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
for embed_dtype in EMBED_DTYPES:
for endianness in ENDIANNESS:
responses_bytes = requests.post(
url,
json={
"model": model_name,
"input": input_texts,
"encoding_format": "bytes",
"embed_dtype": embed_dtype,
"endianness": endianness,
},
)
metadata = json.loads(responses_bytes.headers["metadata"])
body = responses_bytes.content
items = [MetadataItem(**x) for x in metadata["data"]]
bytes_data = decode_pooling_output(items=items, body=body)
bytes_data = [x.to(torch.float32).view(-1).tolist() for x in bytes_data]
check_embeddings_close(
embeddings_0_lst=float_data,
embeddings_1_lst=bytes_data,
name_0="float_data",
name_1="bytes_data",
tol=1e-2,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_bytes_only_embed_dtype_and_endianness(
server: RemoteOpenAIServer, model_name: str
):
input_texts = [input_text] * 3
url = server.url_for("pooling")
float_response = requests.post(
url,
json={
"model": model_name,
"input": input_texts,
"encoding_format": "float",
},
)
responses_float = PoolingResponse.model_validate(float_response.json())
float_data = [np.array(d.data).squeeze(-1).tolist() for d in responses_float.data]
n_tokens = responses_float.usage.prompt_tokens // len(input_texts)
for embed_dtype in EMBED_DTYPES:
for endianness in ENDIANNESS:
responses_bytes = requests.post(
url,
json={
"model": model_name,
"input": input_texts,
"encoding_format": "bytes_only",
"embed_dtype": embed_dtype,
"endianness": endianness,
},
)
assert "metadata" not in responses_bytes.headers
body = responses_bytes.content
items = build_metadata_items(
embed_dtype=embed_dtype,
endianness=endianness,
shape=(n_tokens, 1),
n_request=len(input_texts),
)
bytes_data = decode_pooling_output(items=items, body=body)
bytes_data = [x.to(torch.float32).view(-1).tolist() for x in bytes_data]
check_embeddings_close(
embeddings_0_lst=float_data,
embeddings_1_lst=bytes_data,
name_0="float_data",
name_1="bytes_data",
tol=1e-2,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("param_name", ["encoding_format", "embed_dtype", "endianness"])
async def test_params_not_supported(
server: RemoteOpenAIServer, model_name: str, param_name: str
):
responses_base64 = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": input_text,
"encoding_format": "base64",
param_name: f"bad_{param_name}",
},
)
assert responses_base64.status_code == 400
assert "literal_error" in responses_base64.json()["error"]["message"]
assert f"bad_{param_name}" in responses_base64.json()["error"]["message"]
@pytest.mark.asyncio
async def test_invocations_chat_request(server: RemoteOpenAIServer):
request_args = {
"model": MODEL_NAME,
"input": input_text,
"encoding_format": "float",
}
completion_response = requests.post(server.url_for("pooling"), json=request_args)
completion_response.raise_for_status()
invocation_response = requests.post(
server.url_for("invocations"), json=request_args
)
invocation_response.raise_for_status()
completion_output = completion_response.json()
invocation_output = invocation_response.json()
assert completion_output.keys() == invocation_output.keys()
for completion_data, invocation_data in zip(
completion_output["data"], invocation_output["data"]
):
assert completion_data.keys() == invocation_data.keys()
check_embeddings_close(
embeddings_0_lst=completion_data["data"],
embeddings_1_lst=invocation_data["data"],
name_0="completion",
name_1="invocation",
)
@pytest.mark.asyncio
async def test_invocations_conversation_chat_request(server: RemoteOpenAIServer):
messages = [
{
"role": "user",
"content": "The cat sat on the mat.",
},
{
"role": "assistant",
"content": "A feline was resting on a rug.",
},
{
"role": "user",
"content": "Stars twinkle brightly in the night sky.",
},
]
request_args = {
"model": MODEL_NAME,
"messages": messages,
"encoding_format": "float",
}
chat_response = requests.post(server.url_for("pooling"), json=request_args)
chat_response.raise_for_status()
invocation_response = requests.post(
server.url_for("invocations"), json=request_args
)
invocation_response.raise_for_status()
chat_output = chat_response.json()
invocation_output = invocation_response.json()
assert chat_output.keys() == invocation_output.keys()
for chat_data, invocation_data in zip(
chat_output["data"], invocation_output["data"]
):
assert chat_data.keys() == invocation_data.keys()
check_embeddings_close(
embeddings_0_lst=chat_data["data"],
embeddings_1_lst=invocation_data["data"],
name_0="chat",
name_1="invocation",
)

View File

@@ -0,0 +1,67 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref
import pytest
import torch
from tests.models.utils import softmax
from vllm import LLM, PoolingParams
from vllm.distributed import cleanup_dist_env_and_memory
MODEL_NAME = "internlm/internlm2-1_8b-reward"
prompts = ["The chef prepared a delicious meal."]
@pytest.fixture(scope="module")
def llm():
# pytest caches the fixture so we use weakref.proxy to
# enable garbage collection
llm = LLM(
model=MODEL_NAME,
max_num_batched_tokens=32768,
tensor_parallel_size=1,
gpu_memory_utilization=0.75,
enforce_eager=True,
trust_remote_code=True,
seed=0,
)
yield weakref.proxy(llm)
del llm
cleanup_dist_env_and_memory()
@pytest.mark.skip_global_cleanup
def test_config(llm: LLM):
vllm_config = llm.llm_engine.vllm_config
assert vllm_config.cache_config.enable_prefix_caching
assert vllm_config.scheduler_config.enable_chunked_prefill
def test_pooling_params(llm: LLM):
def get_outputs(use_activation):
outputs = llm.reward(
prompts,
pooling_params=PoolingParams(use_activation=use_activation),
use_tqdm=False,
)
return torch.cat([x.outputs.data for x in outputs])
default = get_outputs(use_activation=None)
w_activation = get_outputs(use_activation=True)
wo_activation = get_outputs(use_activation=False)
assert torch.allclose(default, w_activation, atol=1e-2), (
"Default should use activation."
)
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
"wo_activation should not use activation."
)
assert torch.allclose(softmax(wo_activation), w_activation, atol=1e-2), (
"w_activation should be close to activation(wo_activation)."
)

View File

@@ -0,0 +1,61 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
import pytest
from tests.models.language.pooling_mteb_test.mteb_score_utils import (
MTEB_RERANK_LANGS,
MTEB_RERANK_TASKS,
MTEB_RERANK_TOL,
RerankClientMtebEncoder,
ScoreClientMtebEncoder,
run_mteb_rerank,
)
from tests.utils import RemoteOpenAIServer
from vllm.platforms import current_platform
os.environ["VLLM_LOGGING_LEVEL"] = "WARNING"
MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2"
st_main_score = 0.33457
@pytest.fixture(scope="module")
def server():
args = ["--runner", "pooling", "--enforce-eager", "--disable-uvicorn-access-log"]
# ROCm: Use Flex Attention to support encoder-only self-attention.
if current_platform.is_rocm():
args.extend(["--attention-backend", "FLEX_ATTENTION"])
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
def test_mteb_score(server):
url = server.url_for("score")
encoder = ScoreClientMtebEncoder(MODEL_NAME, url)
vllm_main_score = run_mteb_rerank(encoder, MTEB_RERANK_TASKS, MTEB_RERANK_LANGS)
print("VLLM main score: ", vllm_main_score)
print("SentenceTransformer main score: ", st_main_score)
print("Difference: ", st_main_score - vllm_main_score)
# We are not concerned that the vllm mteb results are better
# than SentenceTransformers, so we only perform one-sided testing.
assert st_main_score - vllm_main_score < MTEB_RERANK_TOL
def test_mteb_rerank(server):
url = server.url_for("rerank")
encoder = RerankClientMtebEncoder(MODEL_NAME, url)
vllm_main_score = run_mteb_rerank(encoder, MTEB_RERANK_TASKS, MTEB_RERANK_LANGS)
print("VLLM main score: ", vllm_main_score)
print("SentenceTransformer main score: ", st_main_score)
print("Difference: ", st_main_score - vllm_main_score)
# We are not concerned that the vllm mteb results are better
# than SentenceTransformers, so we only perform one-sided testing.
assert st_main_score - vllm_main_score < MTEB_RERANK_TOL

View File

@@ -0,0 +1,69 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import weakref
import pytest
import torch
from tests.models.utils import softmax
from vllm import LLM, PoolingParams
from vllm.distributed import cleanup_dist_env_and_memory
from vllm.platforms import current_platform
MODEL_NAME = "tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
@pytest.fixture(scope="module")
def llm():
# ROCm: Use FLEX_ATTENTION backend as it's the only attention backend
# that supports encoder-only models on ROCm.
attention_config = None
if current_platform.is_rocm():
attention_config = {"backend": "FLEX_ATTENTION"}
# pytest caches the fixture so we use weakref.proxy to
# enable garbage collection
llm = LLM(
model=MODEL_NAME,
max_num_batched_tokens=32768,
tensor_parallel_size=1,
gpu_memory_utilization=0.75,
enforce_eager=True,
seed=0,
attention_config=attention_config,
)
yield weakref.proxy(llm)
del llm
cleanup_dist_env_and_memory()
def test_pooling_params(llm: LLM):
def get_outputs(use_activation):
queries = "What is the capital of France?"
documents = "The capital of France is Paris."
outputs = llm.score(
queries,
documents,
pooling_params=PoolingParams(use_activation=use_activation),
use_tqdm=False,
)
return torch.tensor([x.outputs.score for x in outputs])
default = get_outputs(use_activation=None)
w_activation = get_outputs(use_activation=True)
wo_activation = get_outputs(use_activation=False)
assert torch.allclose(default, w_activation, atol=1e-2), (
"Default should use activation."
)
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
"wo_activation should not use activation."
)
assert torch.allclose(softmax(wo_activation), w_activation, atol=1e-2), (
"w_activation should be close to activation(wo_activation)."
)

View File

@@ -0,0 +1,142 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Online API tests for ColBERT late interaction scoring."""
import pytest
import requests
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
MODEL_NAME = "answerdotai/answerai-colbert-small-v1"
COLBERT_DIM = 96
MAX_MODEL_LEN = 512
@pytest.fixture(scope="module")
def server():
args = [
"--max-model-len",
str(MAX_MODEL_LEN),
]
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
class TestColBERTOnline:
def test_rerank(self, server: RemoteOpenAIServer):
"""Test ColBERT rerank endpoint."""
query = "What is the capital of France?"
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
]
rerank_response = requests.post(
server.url_for("rerank"),
json={
"model": MODEL_NAME,
"query": query,
"documents": documents,
},
)
rerank_response.raise_for_status()
rerank = RerankResponse.model_validate(rerank_response.json())
assert rerank.id is not None
assert rerank.results is not None
assert len(rerank.results) == 2
paris_result = next(r for r in rerank.results if r.index == 1)
brazil_result = next(r for r in rerank.results if r.index == 0)
assert paris_result.relevance_score > brazil_result.relevance_score
def test_rerank_top_n(self, server: RemoteOpenAIServer):
"""Test ColBERT rerank with top_n parameter."""
query = "What is the capital of France?"
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
"Machine learning is a field of AI.",
]
rerank_response = requests.post(
server.url_for("rerank"),
json={
"model": MODEL_NAME,
"query": query,
"documents": documents,
"top_n": 2,
},
)
rerank_response.raise_for_status()
rerank = RerankResponse.model_validate(rerank_response.json())
assert len(rerank.results) == 2
assert rerank.results[0].index == 1
def test_score(self, server: RemoteOpenAIServer):
"""Test ColBERT score endpoint."""
text_1 = "What is the capital of France?"
text_2 = ["The capital of France is Paris.", "Python is a language."]
score_response = requests.post(
server.url_for("score"),
json={
"model": MODEL_NAME,
"text_1": text_1,
"text_2": text_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
assert score.data[0].score > score.data[1].score
def test_token_embed(self, server: RemoteOpenAIServer):
"""Test ColBERT token_embed task via pooling endpoint."""
text = "What is the capital of France?"
pooling_response = requests.post(
server.url_for("pooling"),
json={
"model": MODEL_NAME,
"input": text,
"task": "token_embed",
},
)
pooling_response.raise_for_status()
pooling = pooling_response.json()
assert "data" in pooling
assert len(pooling["data"]) == 1
embeddings = pooling["data"][0]["data"]
assert isinstance(embeddings, list)
assert len(embeddings) > 0
assert len(embeddings[0]) == COLBERT_DIM
def test_embed_not_supported(self, server: RemoteOpenAIServer):
"""Test that ColBERT model does not support 'embed' task."""
task = "embed"
text = "What is the capital of France?"
response = requests.post(
server.url_for("pooling"),
json={
"model": MODEL_NAME,
"input": text,
"task": task,
},
)
assert response.json()["error"]["type"] == "BadRequestError"
assert response.json()["error"]["message"].startswith(
f"Unsupported task: {task!r}"
)

View File

@@ -0,0 +1,235 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import requests
import torch
import torch.nn.functional as F
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.pooling.pooling.protocol import PoolingResponse
from vllm.entrypoints.pooling.score.protocol import RerankResponse
from vllm.platforms import current_platform
MODEL_NAME = "BAAI/bge-reranker-base"
DTYPE = "bfloat16"
input_text = "This product was excellent and exceeded my expectations"
input_tokens = [0, 3293, 12996, 509, 40881, 136, 204839, 297, 759, 202702, 2]
@pytest.fixture(scope="module")
def server():
args = ["--enforce-eager", "--max-model-len", "100", "--dtype", DTYPE]
# ROCm: Use Flex Attention to support encoder-only self-attention.
if current_platform.is_rocm():
args.extend(["--attention-backend", "FLEX_ATTENTION"])
with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
yield remote_server
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_basic(server: RemoteOpenAIServer, model_name: str):
# test /v1/models
response = requests.get(server.url_for("/v1/models"))
served_model = response.json()["data"][0]["id"]
assert served_model == MODEL_NAME
# test /tokenize
response = requests.post(
server.url_for("/tokenize"),
json={"model": model_name, "prompt": input_text},
)
assert response.json()["tokens"] == input_tokens
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_rerank_texts(server: RemoteOpenAIServer, model_name: str):
query = "What is the capital of France?"
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
]
rerank_response = requests.post(
server.url_for("rerank"),
json={
"model": model_name,
"query": query,
"documents": documents,
},
)
rerank_response.raise_for_status()
rerank = RerankResponse.model_validate(rerank_response.json())
assert rerank.id is not None
assert rerank.results is not None
assert len(rerank.results) == 2
assert rerank.results[0].relevance_score >= 0.9
assert rerank.results[1].relevance_score <= 0.01
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_top_n(server: RemoteOpenAIServer, model_name: str):
query = "What is the capital of France?"
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
"Cross-encoder models are neat",
]
rerank_response = requests.post(
server.url_for("rerank"),
json={"model": model_name, "query": query, "documents": documents, "top_n": 2},
)
rerank_response.raise_for_status()
rerank = RerankResponse.model_validate(rerank_response.json())
assert rerank.id is not None
assert rerank.results is not None
assert len(rerank.results) == 2
assert rerank.results[0].relevance_score >= 0.9
assert rerank.results[1].relevance_score <= 0.01
@pytest.mark.parametrize("model_name", [MODEL_NAME])
def test_rerank_max_model_len(server: RemoteOpenAIServer, model_name: str):
query = "What is the capital of France?" * 100
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
]
rerank_response = requests.post(
server.url_for("rerank"),
json={"model": model_name, "query": query, "documents": documents},
)
assert rerank_response.status_code == 400
# Assert just a small fragments of the response
assert "Please reduce the length of the input." in rerank_response.text
def test_invocations(server: RemoteOpenAIServer):
query = "What is the capital of France?"
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
]
request_args = {
"model": MODEL_NAME,
"query": query,
"documents": documents,
}
rerank_response = requests.post(server.url_for("rerank"), json=request_args)
rerank_response.raise_for_status()
invocation_response = requests.post(
server.url_for("invocations"), json=request_args
)
invocation_response.raise_for_status()
rerank_output = rerank_response.json()
invocation_output = invocation_response.json()
assert rerank_output.keys() == invocation_output.keys()
for rerank_result, invocations_result in zip(
rerank_output["results"], invocation_output["results"]
):
assert rerank_result.keys() == invocations_result.keys()
assert rerank_result["relevance_score"] == pytest.approx(
invocations_result["relevance_score"], rel=0.05
)
# TODO: reset this tolerance to 0.01 once we find
# an alternative to flash_attn with bfloat16
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_use_activation(server: RemoteOpenAIServer, model_name: str):
async def get_outputs(use_activation):
query = "What is the capital of France?"
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
]
response = requests.post(
server.url_for("rerank"),
json={
"model": model_name,
"query": query,
"documents": documents,
"use_activation": use_activation,
},
)
outputs = response.json()
return torch.tensor([x["relevance_score"] for x in outputs["results"]])
default = await get_outputs(use_activation=None)
w_activation = await get_outputs(use_activation=True)
wo_activation = await get_outputs(use_activation=False)
assert torch.allclose(default, w_activation, atol=1e-2), (
"Default should use activation."
)
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
"wo_activation should not use activation."
)
assert torch.allclose(F.sigmoid(wo_activation), w_activation, atol=1e-2), (
"w_activation should be close to activation(wo_activation)."
)
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_pooling_classify(server: RemoteOpenAIServer, model_name: str):
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": input_text,
"encoding_format": "float",
"task": "classify",
},
)
poolings = PoolingResponse.model_validate(response.json())
assert len(poolings.data) == 1
assert len(poolings.data[0].data) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
async def test_pooling_token_classify(server: RemoteOpenAIServer, model_name: str):
response = requests.post(
server.url_for("pooling"),
json={"model": model_name, "input": input_text, "encoding_format": "float"},
)
poolings = PoolingResponse.model_validate(response.json())
assert len(poolings.data) == 1
assert len(poolings.data[0].data) == len(input_tokens)
assert len(poolings.data[0].data[0]) == 1
@pytest.mark.asyncio
@pytest.mark.parametrize("model_name", [MODEL_NAME])
@pytest.mark.parametrize("task", ["embed", "token_embed", "plugin"])
async def test_pooling_not_supported(
server: RemoteOpenAIServer, model_name: str, task: str
):
response = requests.post(
server.url_for("pooling"),
json={
"model": model_name,
"input": input_text,
"encoding_format": "float",
"task": task,
},
)
assert response.json()["error"]["type"] == "BadRequestError"
assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")

View File

@@ -0,0 +1,342 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Any
import pytest
import requests
import torch
import torch.nn.functional as F
from torch import tensor
from tests.utils import RemoteOpenAIServer
from vllm.entrypoints.pooling.score.protocol import ScoreResponse
from vllm.platforms import current_platform
MODELS = [
{"name": "BAAI/bge-reranker-v2-m3", "is_cross_encoder": True},
{"name": "BAAI/bge-base-en-v1.5", "is_cross_encoder": False},
]
DTYPE = "half"
def run_transformers(hf_model, model, text_pairs):
if model["is_cross_encoder"]:
return hf_model.predict(text_pairs).tolist()
else:
hf_embeddings = [hf_model.encode(text_pair) for text_pair in text_pairs]
return [
F.cosine_similarity(tensor(pair[0]), tensor(pair[1]), dim=0)
for pair in hf_embeddings
]
@pytest.fixture(scope="class", params=MODELS)
def model(request):
yield request.param
@pytest.fixture(scope="class")
def server(model: dict[str, Any]):
args = ["--enforce-eager", "--max-model-len", "100", "--dtype", DTYPE]
# ROCm: Use Flex Attention to support encoder-only self-attention.
if current_platform.is_rocm():
args.extend(["--attention-backend", "FLEX_ATTENTION"])
with RemoteOpenAIServer(model["name"], args) as remote_server:
yield remote_server
@pytest.fixture(scope="class")
def runner(model: dict[str, Any], hf_runner):
kwargs = {
"dtype": DTYPE,
"is_cross_encoder"
if model["is_cross_encoder"]
else "is_sentence_transformer": True,
}
with hf_runner(model["name"], **kwargs) as hf_model:
yield hf_model
class TestModel:
def test_queries_str_documents_str(
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
):
queries = "What is the capital of France?"
documents = "The capital of France is Paris."
score_response = requests.post(
server.url_for("score"),
json={
"model": model["name"],
"queries": queries,
"documents": documents,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 1
vllm_outputs = [d.score for d in score.data]
text_pairs = [[queries, documents]]
hf_outputs = run_transformers(runner, model, text_pairs)
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
def test_queries_str_items_str(
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
):
queries = "What is the capital of France?"
items = "The capital of France is Paris."
score_response = requests.post(
server.url_for("score"),
json={
"model": model["name"],
"queries": queries,
"items": items,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 1
vllm_outputs = [d.score for d in score.data]
text_pairs = [[queries, items]]
hf_outputs = run_transformers(runner, model, text_pairs)
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
def test_text_1_str_text_2_str(
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
):
text_1 = "What is the capital of France?"
text_2 = "The capital of France is Paris."
score_response = requests.post(
server.url_for("score"),
json={
"model": model["name"],
"text_1": text_1,
"text_2": text_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 1
vllm_outputs = [d.score for d in score.data]
text_pairs = [[text_1, text_2]]
hf_outputs = run_transformers(runner, model, text_pairs)
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
def test_data_1_str_data_2_str(
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
):
data_1 = "What is the capital of France?"
data_2 = "The capital of France is Paris."
score_response = requests.post(
server.url_for("score"),
json={
"model": model["name"],
"data_1": data_1,
"data_2": data_2,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 1
vllm_outputs = [d.score for d in score.data]
text_pairs = [[data_1, data_2]]
hf_outputs = run_transformers(runner, model, text_pairs)
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
def test_queries_str_documents_list(
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
):
queries = "What is the capital of France?"
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
]
score_response = requests.post(
server.url_for("score"),
json={
"model": model["name"],
"queries": queries,
"documents": documents,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
text_pairs = [[queries, documents[0]], [queries, documents[1]]]
hf_outputs = run_transformers(runner, model, text_pairs)
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
def test_queries_list_documents_list(
self, server: RemoteOpenAIServer, model: dict[str, Any], runner
):
queries = [
"What is the capital of the United States?",
"What is the capital of France?",
]
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
]
score_response = requests.post(
server.url_for("score"),
json={
"model": model["name"],
"queries": queries,
"documents": documents,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 2
vllm_outputs = [d.score for d in score.data]
text_pairs = [[queries[0], documents[0]], [queries[1], documents[1]]]
hf_outputs = run_transformers(runner, model, text_pairs)
for i in range(len(vllm_outputs)):
assert hf_outputs[i] == pytest.approx(vllm_outputs[i], rel=0.01)
def test_score_max_model_len(
self, server: RemoteOpenAIServer, model: dict[str, Any]
):
queries = "What is the capital of France?" * 20
documents = [
"The capital of Brazil is Brasilia.",
"The capital of France is Paris.",
]
score_response = requests.post(
server.url_for("score"),
json={
"model": model["name"],
"queries": queries,
"documents": documents,
},
)
assert score_response.status_code == 400
# Assert just a small fragments of the response
assert "Please reduce the length of the input." in score_response.text
# Test truncation
score_response = requests.post(
server.url_for("score"),
json={
"model": model["name"],
"queries": queries,
"documents": documents,
"truncate_prompt_tokens": 101,
},
)
assert score_response.status_code == 400
assert "Please request a smaller truncation size." in score_response.text
def test_invocations(self, server: RemoteOpenAIServer, model: dict[str, Any]):
queries = "What is the capital of France?"
documents = "The capital of France is Paris."
request_args = {
"model": model["name"],
"queries": queries,
"documents": documents,
}
score_response = requests.post(server.url_for("score"), json=request_args)
score_response.raise_for_status()
invocation_response = requests.post(
server.url_for("invocations"), json=request_args
)
invocation_response.raise_for_status()
score_output = score_response.json()
invocation_output = invocation_response.json()
assert score_output.keys() == invocation_output.keys()
for score_data, invocation_data in zip(
score_output["data"], invocation_output["data"]
):
assert score_data.keys() == invocation_data.keys()
assert score_data["score"] == pytest.approx(
invocation_data["score"], rel=0.05
)
# TODO: reset this tolerance to 0.01 once we find
# an alternative to flash_attn with bfloat16
def test_use_activation(self, server: RemoteOpenAIServer, model: dict[str, Any]):
def get_outputs(use_activation):
queries = "What is the capital of France?"
documents = "The capital of France is Paris."
response = requests.post(
server.url_for("score"),
json={
"model": model["name"],
"queries": queries,
"documents": documents,
"use_activation": use_activation,
},
)
outputs = response.json()
return torch.tensor([x["score"] for x in outputs["data"]])
default = get_outputs(use_activation=None)
w_activation = get_outputs(use_activation=True)
wo_activation = get_outputs(use_activation=False)
if model["is_cross_encoder"]:
assert torch.allclose(default, w_activation, atol=1e-2), (
"Default should use activation."
)
assert not torch.allclose(w_activation, wo_activation, atol=1e-2), (
"wo_activation should not use activation."
)
assert torch.allclose(F.sigmoid(wo_activation), w_activation, atol=1e-2), (
"w_activation should be close to activation(wo_activation)."
)

View File

@@ -0,0 +1,355 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import pytest
import requests
from tests.utils import VLLM_PATH, RemoteOpenAIServer
from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
from vllm.multimodal.utils import encode_image_url, fetch_image
from vllm.platforms import current_platform
MODEL_NAME = "Qwen/Qwen3-VL-Reranker-2B"
HF_OVERRIDES = {
"architectures": ["Qwen3VLForSequenceClassification"],
"classifier_from_token": ["no", "yes"],
"is_original_qwen3_reranker": True,
}
ROCM_ATTN_BACKENDS = [
"ROCM_ATTN",
"ROCM_AITER_FA",
"TRITON_ATTN",
"FLEX_ATTENTION",
]
ATTN_BACKENDS = ROCM_ATTN_BACKENDS if current_platform.is_rocm() else ["auto"]
# Per-backend tolerance with explicit entries; "default" is the fallback
BACKEND_TOL: dict[str, float] = {
"default": 0.05, # 5% tolerance for other backends (e.g. FLASH_ATTN)
# Relaxed tolerances for ROCm attn
# See: https://github.com/vllm-project/vllm/issues/35569
"ROCM_ATTN": 0.09, # gfx950:~8.45%, gfx942:~3.70%
"ROCM_AITER_FA": 0.045, # gfx950:~2.00%, gfx942:~0.80%
"TRITON_ATTN": 0.045, # gfx950:~3.00%, gfx942:~2.20%
"FLEX_ATTENTION": 0.045, # gfx950:~3.25%, gfx942:~1.10%
}
# ROCm: disable skinny GEMM to avoid non-deterministic results from
# atomic reductions in wvSplitKrc kernel.
# See: https://github.com/vllm-project/vllm/pull/33493#issuecomment-3906083975
ROCM_ENV_OVERRIDES = (
{"VLLM_ROCM_USE_SKINNY_GEMM": "0"} if current_platform.is_rocm() else {}
)
# ROCm: disable prefix caching and eliminate batch variance to reduce
# test flakiness.
ROCM_EXTRA_ARGS = (
["--no-enable-prefix-caching", "--max-num-seqs", "1"]
if current_platform.is_rocm()
else []
)
def get_tol(backend: str) -> float:
return BACKEND_TOL.get(backend, BACKEND_TOL["default"])
def assert_score(actual: float, expected: float, backend: str, label: str):
tol = get_tol(backend)
diff = abs(actual - expected)
rel_diff = diff / abs(expected) if expected != 0 else diff
print(
f"[{backend}] {label}: actual={actual:.6f} expected={expected:.6f} "
f"diff={diff:.6f} rel_diff={rel_diff:.4f} tol={tol}"
)
assert actual == pytest.approx(expected, rel=tol), (
f"[{backend}] {label}: score mismatch — "
f"actual={actual:.6f}, expected={expected:.6f}, "
f"rel_diff={rel_diff:.4f}, tol={tol}"
)
query = "A cat standing in the snow."
document = "This product was excellent and exceeded my expectations."
image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/cat_snow.jpg"
documents = [
{
"type": "text",
"text": document,
},
{
"type": "image_url",
"image_url": {"url": image_url},
},
{
"type": "image_url",
"image_url": {"url": encode_image_url(fetch_image(image_url))},
},
]
TEXT_VS_TEXT = 0.10040374100208282
TEXT_VS_IMAGE = 0.7423753142356873
TEXT_VS_TEXT_PLUS_IMAGE = 0.5298863053321838
@pytest.fixture(scope="module", params=ATTN_BACKENDS)
def server(request):
backend = request.param
print(f"\n=== Starting server with attention backend: {backend} ===")
args = [
"--enforce-eager",
"--max-model-len",
"8192",
"--chat-template",
str(VLLM_PATH / "examples/pooling/score/template/qwen3_vl_reranker.jinja"),
]
env = dict()
if backend != "auto":
args += ["--attention-config", json.dumps({"backend": backend})]
args += ROCM_EXTRA_ARGS
env = dict(ROCM_ENV_OVERRIDES)
if backend != "ROCM_AITER_FA":
env["VLLM_ROCM_USE_AITER"] = "0"
with RemoteOpenAIServer(
MODEL_NAME, args, override_hf_configs=HF_OVERRIDES, env_dict=env
) as remote_server:
print(f"=== Server ready with backend: {backend} ===")
yield remote_server, backend
def test_score_api_queries_str_documents_str(server: tuple[RemoteOpenAIServer, str]):
remote_server, backend = server
score_response = requests.post(
remote_server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": query,
"documents": document,
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 1
assert score.usage.prompt_tokens == 81
assert_score(score.data[0].score, TEXT_VS_TEXT, backend, "text_vs_text")
def test_score_api_queries_str_documents_text_content(
server: tuple[RemoteOpenAIServer, str],
):
remote_server, backend = server
score_response = requests.post(
remote_server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": query,
"documents": {"content": [documents[0]]},
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 1
assert score.usage.prompt_tokens == 81
assert_score(score.data[0].score, TEXT_VS_TEXT, backend, "text_vs_text")
def test_score_api_queries_str_documents_image_url_content(
server: tuple[RemoteOpenAIServer, str],
):
remote_server, backend = server
score_response = requests.post(
remote_server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": query,
"documents": {"content": [documents[1]]},
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 1
assert score.usage.prompt_tokens == 98
assert_score(score.data[0].score, TEXT_VS_IMAGE, backend, "text_vs_image")
def test_score_api_queries_str_documents_image_base64_content(
server: tuple[RemoteOpenAIServer, str],
):
remote_server, backend = server
score_response = requests.post(
remote_server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": query,
"documents": {"content": [documents[2]]},
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 1
assert score.usage.prompt_tokens == 98
assert_score(score.data[0].score, TEXT_VS_IMAGE, backend, "text_vs_image_base64")
def test_score_api_queries_str_documents_image_url_plus_text_content(
server: tuple[RemoteOpenAIServer, str],
):
remote_server, backend = server
score_response = requests.post(
remote_server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": query,
"documents": {"content": [documents[0], documents[1]]},
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 1
assert score.usage.prompt_tokens == 108
assert_score(
score.data[0].score, TEXT_VS_TEXT_PLUS_IMAGE, backend, "text_vs_text_plus_image"
)
def test_score_api_queries_str_documents_list(
server: tuple[RemoteOpenAIServer, str],
):
remote_server, backend = server
score_response = requests.post(
remote_server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": query,
"documents": [
document,
{"content": [documents[0]]},
{"content": [documents[1]]},
{"content": [documents[0], documents[1]]},
],
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 4
assert score.usage.prompt_tokens == 368
assert_score(score.data[0].score, TEXT_VS_TEXT, backend, "list[0]_text_vs_text")
assert_score(score.data[1].score, TEXT_VS_TEXT, backend, "list[1]_text_vs_text")
assert_score(score.data[2].score, TEXT_VS_IMAGE, backend, "list[2]_text_vs_image")
assert_score(
score.data[3].score,
TEXT_VS_TEXT_PLUS_IMAGE,
backend,
"list[3]_text_vs_text_plus_image",
)
def test_rerank_api_queries_str_documents_list(
server: tuple[RemoteOpenAIServer, str],
):
remote_server, backend = server
rerank_response = requests.post(
remote_server.url_for("rerank"),
json={
"model": MODEL_NAME,
"query": query,
"documents": [
document,
{"content": [documents[0]]},
{"content": [documents[1]]},
{"content": [documents[0], documents[1]]},
],
},
)
rerank_response.raise_for_status()
rerank = RerankResponse.model_validate(rerank_response.json())
assert rerank.id is not None
assert rerank.model is not None
assert rerank.usage is not None
assert len(rerank.results) == 4
rerank.results.sort(key=lambda x: x.index)
assert_score(
rerank.results[0].relevance_score,
TEXT_VS_TEXT,
backend,
"rerank[0]_text_vs_text",
)
assert_score(
rerank.results[1].relevance_score,
TEXT_VS_TEXT,
backend,
"rerank[1]_text_vs_text",
)
assert_score(
rerank.results[2].relevance_score,
TEXT_VS_IMAGE,
backend,
"rerank[2]_text_vs_image",
)
assert_score(
rerank.results[3].relevance_score,
TEXT_VS_TEXT_PLUS_IMAGE,
backend,
"rerank[3]_text_vs_text_plus_image",
)
def test_score_api_queries_list_documents_list(
server: tuple[RemoteOpenAIServer, str],
):
remote_server, backend = server
score_response = requests.post(
remote_server.url_for("score"),
json={
"model": MODEL_NAME,
"queries": [query] * 4,
"documents": [
document,
{"content": [documents[0]]},
{"content": [documents[1]]},
{"content": [documents[0], documents[1]]},
],
},
)
score_response.raise_for_status()
score = ScoreResponse.model_validate(score_response.json())
assert score.id is not None
assert score.data is not None
assert len(score.data) == 4
assert score.usage.prompt_tokens == 368
assert_score(score.data[0].score, TEXT_VS_TEXT, backend, "paired[0]_text_vs_text")
assert_score(score.data[1].score, TEXT_VS_TEXT, backend, "paired[1]_text_vs_text")
assert_score(score.data[2].score, TEXT_VS_IMAGE, backend, "paired[2]_text_vs_image")
assert_score(
score.data[3].score,
TEXT_VS_TEXT_PLUS_IMAGE,
backend,
"paired[3]_text_vs_text_plus_image",
)

View File

@@ -0,0 +1,353 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import patch
import pytest
from vllm.config import ModelConfig
from vllm.entrypoints.chat_utils import ChatTemplateResolutionError
from vllm.entrypoints.pooling.score.utils import (
get_score_prompt,
)
from vllm.inputs import TokensPrompt
from vllm.tokenizers import get_tokenizer
# A cross-encoder model for testing
CROSS_ENCODER_MODEL_ID = "cross-encoder/ms-marco-MiniLM-L-6-v2"
def assert_prompt_tokenization_consistent(
tokenizer, full_prompt, engine_prompt, add_special_tokens=True
):
"""Verify that engine_prompt token_ids match tokenizing full_prompt."""
expected_ids = tokenizer(full_prompt, add_special_tokens=add_special_tokens)[
"input_ids"
]
actual_ids = engine_prompt["prompt_token_ids"]
assert actual_ids == expected_ids, (
f"Token IDs don't match.\nExpected: {expected_ids}\nActual: {actual_ids}"
)
@pytest.fixture(scope="module")
def cross_encoder_model_config():
return ModelConfig(
CROSS_ENCODER_MODEL_ID,
runner="pooling",
)
@pytest.fixture(scope="module")
def cross_encoder_tokenizer(cross_encoder_model_config):
return get_tokenizer(
CROSS_ENCODER_MODEL_ID,
trust_remote_code=cross_encoder_model_config.trust_remote_code,
)
@pytest.fixture(scope="module")
def llm_reranker_model_config():
"""Model config for LLM-as-reranker style (no pad token)."""
config = ModelConfig(
CROSS_ENCODER_MODEL_ID,
runner="pooling",
)
# use_sep_token is a property that reads from hf_config,
# so we set it there to override the default (True)
config.hf_config.use_sep_token = False
return config
@pytest.fixture
def tokenization_kwargs():
"""Common tokenization kwargs used across tests."""
return {"add_special_tokens": True, "return_tensors": None}
@pytest.fixture
def mock_model_with_score_template():
"""Mock model class that supports score template and tracks post_process calls."""
class MockModelWithScoreTemplate:
supports_score_template = True
post_process_called: list[TokensPrompt] = []
@staticmethod
def get_score_template(p1: str, p2: str) -> str:
return f"[QUERY]{p1}[SEP][DOC]{p2}"
@staticmethod
def post_process_tokens(prompt: TokensPrompt) -> None:
MockModelWithScoreTemplate.post_process_called.append(prompt)
return MockModelWithScoreTemplate
@pytest.fixture
def mock_model_no_score_template():
"""Mock model class that does not support score template."""
class MockModelNoScoreTemplate:
supports_score_template = False
return MockModelNoScoreTemplate
class TestGetScorePrompt:
"""Tests for the get_score_prompt function."""
def test_tokenization_kwargs_passed_through(
self,
llm_reranker_model_config,
cross_encoder_tokenizer,
):
"""Test that tokenization kwargs are properly passed through."""
data_1 = "Query text"
data_2 = "Document text"
# Test with truncation - custom kwargs for this test
custom_tokenization_kwargs = {
"add_special_tokens": True,
"return_tensors": None,
"truncation": True,
"max_length": 20,
}
full_prompt, engine_prompt = get_score_prompt(
llm_reranker_model_config,
cross_encoder_tokenizer,
custom_tokenization_kwargs,
data_1,
data_2,
)
assert isinstance(full_prompt, str)
assert "prompt_token_ids" in engine_prompt
# With max_length=20 and truncation, should not exceed this
assert len(engine_prompt["prompt_token_ids"]) <= 20
# Since truncation was applied, token_ids should be a prefix of full encoding
full_ids = cross_encoder_tokenizer(full_prompt, add_special_tokens=True)[
"input_ids"
]
actual_ids = engine_prompt["prompt_token_ids"]
assert full_ids[: len(actual_ids)] == actual_ids, (
f"Token IDs are not a prefix of full encoding.\n"
f"Full IDs: {full_ids}\n"
f"Actual IDs: {actual_ids}"
)
def test_model_supports_score_template(
self,
cross_encoder_model_config,
cross_encoder_tokenizer,
tokenization_kwargs,
mock_model_with_score_template,
):
"""Test when model supports score template (no score_template arg)."""
with patch(
"vllm.model_executor.model_loader.get_model_cls",
return_value=mock_model_with_score_template,
):
full_prompt, engine_prompt = get_score_prompt(
cross_encoder_model_config,
cross_encoder_tokenizer,
tokenization_kwargs,
"query text",
"document text",
)
assert full_prompt == "[QUERY]query text[SEP][DOC]document text"
assert "prompt_token_ids" in engine_prompt
assert len(engine_prompt["prompt_token_ids"]) > 0
assert_prompt_tokenization_consistent(
cross_encoder_tokenizer, full_prompt, engine_prompt
)
def test_model_supports_score_template_but_custom_template_provided(
self,
cross_encoder_model_config,
cross_encoder_tokenizer,
tokenization_kwargs,
mock_model_with_score_template,
):
"""Test when model supports score template but custom template is provided."""
template = (
'TEMPLATE_USED {{ messages[0]["content"] }} {{ messages[1]["content"] }}'
)
with (
patch(
"vllm.model_executor.model_loader.get_model_cls",
return_value=mock_model_with_score_template,
),
):
full_prompt, engine_prompt = get_score_prompt(
cross_encoder_model_config,
cross_encoder_tokenizer,
tokenization_kwargs,
"query",
"doc",
score_template=template, # Providing a template
)
assert "prompt_token_ids" in engine_prompt
assert full_prompt == "TEMPLATE_USED query doc"
assert_prompt_tokenization_consistent(
cross_encoder_tokenizer, full_prompt, engine_prompt
)
def test_not_using_default_template(
self,
llm_reranker_model_config,
cross_encoder_tokenizer,
tokenization_kwargs,
mock_model_no_score_template,
):
# FIXME: For now, we only apply a template when one is explicitly provided.
# We cannot rely on the tokenizer's chat template because many models
# inherit junk templates from their base LLM, which breaks both the models
# and the tests that use them.
with (
patch(
"vllm.model_executor.model_loader.get_model_cls",
return_value=mock_model_no_score_template,
),
patch(
"vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
return_value="test querytest doc",
),
):
full_prompt, engine_prompt = get_score_prompt(
llm_reranker_model_config,
cross_encoder_tokenizer,
tokenization_kwargs,
"test query",
"test doc",
)
assert full_prompt == "test querytest doc"
assert "prompt_token_ids" in engine_prompt
assert_prompt_tokenization_consistent(
cross_encoder_tokenizer, full_prompt, engine_prompt
)
def test_fallback_with_sep_token(
self,
cross_encoder_model_config,
cross_encoder_tokenizer,
tokenization_kwargs,
mock_model_no_score_template,
):
"""Test fallback path when ChatTemplateResolutionError
and use_sep_token=True."""
with (
patch(
"vllm.model_executor.model_loader.get_model_cls",
return_value=mock_model_no_score_template,
),
patch(
"vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
side_effect=ChatTemplateResolutionError("No template"),
),
):
full_prompt, engine_prompt = get_score_prompt(
cross_encoder_model_config, # use_sep_token=True
cross_encoder_tokenizer,
tokenization_kwargs,
"query",
"document",
)
assert "prompt_token_ids" in engine_prompt
# Should have token_type_ids from text_pair encoding
assert "token_type_ids" in engine_prompt
assert "query" in full_prompt
assert "document" in full_prompt
assert full_prompt != "querydocument"
assert (
engine_prompt["prompt_token_ids"]
== cross_encoder_tokenizer(
"query", text_pair="document", add_special_tokens=True
)["input_ids"]
)
# FIXME(?): add_special_tokens=False is needed because in this case
# full_prompt is obtained by decoding the tokenized prompt, which includes
# special tokens and we would get duplicated special tokens otherwise.
# This is inconsistent with other cases.
assert_prompt_tokenization_consistent(
cross_encoder_tokenizer,
full_prompt,
engine_prompt,
add_special_tokens=False,
)
def test_fallback_without_sep_token(
self,
llm_reranker_model_config,
cross_encoder_tokenizer,
tokenization_kwargs,
mock_model_no_score_template,
):
"""Test fallback path when ChatTemplateResolutionError
and use_sep_token=False."""
with (
patch(
"vllm.model_executor.model_loader.get_model_cls",
return_value=mock_model_no_score_template,
),
patch(
"vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
side_effect=ChatTemplateResolutionError("No template"),
),
):
full_prompt, engine_prompt = get_score_prompt(
llm_reranker_model_config, # use_sep_token=False
cross_encoder_tokenizer,
tokenization_kwargs,
"query",
"document",
)
assert full_prompt == "querydocument"
assert "prompt_token_ids" in engine_prompt
assert_prompt_tokenization_consistent(
cross_encoder_tokenizer, full_prompt, engine_prompt
)
def test_post_process_tokens_called(
self,
cross_encoder_model_config,
cross_encoder_tokenizer,
tokenization_kwargs,
mock_model_with_score_template,
):
"""Test that post_process_tokens is called on the engine prompt."""
# Reset the call tracker
mock_model_with_score_template.post_process_called.clear()
with (
patch(
"vllm.model_executor.model_loader.get_model_cls",
return_value=mock_model_with_score_template,
),
patch(
"vllm.entrypoints.pooling.score.utils.safe_apply_chat_template",
side_effect=ChatTemplateResolutionError("No template"),
),
):
full_prompt, engine_prompt = get_score_prompt(
cross_encoder_model_config,
cross_encoder_tokenizer,
tokenization_kwargs,
"query",
"doc",
)
# post_process_tokens should have been called once
assert len(mock_model_with_score_template.post_process_called) == 1
assert mock_model_with_score_template.post_process_called[0] is engine_prompt
assert_prompt_tokenization_consistent(
cross_encoder_tokenizer, full_prompt, engine_prompt
)