A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 14:30 UTC
[Feature]: Support for multiple embedding types in a single inference call (#35829)
Signed-off-by: augusto.yjh <augusto.yjh@antgroup.com>
Augusto Yao committed
Mar 17, 2026 at 17:05 UTC
9c7cab5ebb0f8a15e632e7ea2cfeebcca1d3628f
7 files changed
+226
-36
tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/sparse_embeddings_processor.py
+105
-19
@@ -3,10 +3,10 @@
3
4
from collections.abc import Sequence
5
6
-from vllm.config import VllmConfig
6
+from vllm.config import ModelConfig, PoolerConfig, VllmConfig
7
from vllm.entrypoints.openai.engine.protocol import UsageInfo
8
+from vllm.entrypoints.pooling.base.protocol import EmbedRequestMixin
9
from vllm.inputs.data import PromptType
9
-from vllm.logger import init_logger
10
from vllm.outputs import PoolingRequestOutput
11
from vllm.plugins.io_processors.interface import (
12
IOProcessor,
@@ -16,14 +16,13 @@ from vllm.renderers import BaseRenderer
16
from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens
17
18
from .types import (
19
+ EMBED_TASKS,
20
SparseEmbeddingCompletionRequestMixin,
21
SparseEmbeddingResponse,
22
SparseEmbeddingResponseData,
23
SparseEmbeddingTokenWeight,
24
)
25
25
-logger = init_logger(__name__)
26
-
26
27
class BgeM3SparseEmbeddingsProcessor(
28
IOProcessor[SparseEmbeddingCompletionRequestMixin, SparseEmbeddingResponse]
@@ -33,6 +32,22 @@ class BgeM3SparseEmbeddingsProcessor(
32
self.offline_requests: list[SparseEmbeddingCompletionRequestMixin] = []
33
self.online_requests: dict[str, SparseEmbeddingCompletionRequestMixin] = {}
34
self.renderer: BaseRenderer = renderer
35
+ self.default_pooling_params = {}
36
+ pooler_config: PoolerConfig = vllm_config.model_config.pooler_config
37
+ if pooler_config is not None:
38
+ for param in ["use_activation", "dimensions"]:
39
+ if getattr(pooler_config, param, None) is None:
40
+ continue
41
+ self.default_pooling_params[param] = getattr(pooler_config, param)
42
+ self.embed_dimensions = vllm_config.model_config.embedding_size
43
+ self.embed_request_queue: list[EmbedRequestMixin] = []
44
+
45
+ def __repr__(self) -> str:
46
+ return (
47
+ f"BgeM3SparseEmbeddingsProcessor("
48
+ f"embed_dimensions={self.embed_dimensions}, "
49
+ f"default_pooling_params={self.default_pooling_params})"
50
+ )
51
52
def merge_pooling_params(
53
self,
@@ -41,7 +56,57 @@ class BgeM3SparseEmbeddingsProcessor(
56
if params is None:
57
params = PoolingParams()
58
# refer to PoolingCompletionRequest.to_pooling_params
44
- params.task = "token_classify"
59
+ # set and verify pooling params
60
+ params.skip_reading_prefix_cache = True
61
+
62
+ raw_embed_request = self.embed_request_queue.pop(0)
63
+ if raw_embed_request.embed_task not in EMBED_TASKS:
64
+ raise ValueError(
65
+ f"Unsupported task {raw_embed_request}, "
66
+ f"Supported tasks are {EMBED_TASKS}"
67
+ )
68
+ has_dense_embed = True
69
+ if raw_embed_request.embed_task == "dense":
70
+ params.task = "embed"
71
+ params.skip_reading_prefix_cache = False
72
+ elif raw_embed_request.embed_task == "sparse":
73
+ params.task = "token_classify"
74
+ has_dense_embed = False
75
+ else:
76
+ params.task = "embed&token_classify"
77
+ params.use_activation = raw_embed_request.use_activation
78
+ if params.use_activation is None:
79
+ params.use_activation = True
80
+ if not has_dense_embed:
81
+ params.dimensions = None
82
+ return params
83
+
84
+ params.dimensions = raw_embed_request.dimensions
85
+
86
+ model_config: ModelConfig = self.vllm_config.model_config
87
+ for param in self.default_pooling_params:
88
+ if getattr(params, param, None) is None:
89
+ setattr(params, param, self.default_pooling_params[param])
90
+
91
+ if params.dimensions is not None:
92
+ if not model_config.is_matryoshka:
93
+ raise ValueError(
94
+ f'Model "{model_config.served_model_name}" does not '
95
+ f"support matryoshka representation, "
96
+ f"changing output dimensions will lead to poor results."
97
+ )
98
+
99
+ mds = model_config.matryoshka_dimensions
100
+ if mds is not None:
101
+ if params.dimensions not in mds:
102
+ raise ValueError(
103
+ f"Model {model_config.served_model_name!r} "
104
+ f"only supports {str(mds)} matryoshka dimensions, "
105
+ f"use other output dimensions will "
106
+ f"lead to poor results."
107
+ )
108
+ elif params.dimensions < 1:
109
+ raise ValueError("Dimensions must be greater than 0")
110
return params
111
112
def parse_request(
@@ -61,14 +126,16 @@ class BgeM3SparseEmbeddingsProcessor(
126
if request_id is not None:
127
assert request_id not in self.online_requests, "request_id duplicated"
128
self.online_requests[request_id] = prompt
129
+ self.embed_request_queue.extend(prompt.to_embed_requests_online())
130
else:
131
self.offline_requests.append(prompt)
132
+ self.embed_request_queue.extend(prompt.to_embed_requests_offline())
133
return prompt.input
134
135
def _get_sparse_embedding_request(self, request_id: str | None = None):
136
if request_id:
137
return self.online_requests.pop(request_id, None)
71
- return self.offline_requests.pop()
138
+ return self.offline_requests.pop(0)
139
140
def _build_sparse_embedding_token_weights(
141
self,
@@ -100,26 +167,45 @@ class BgeM3SparseEmbeddingsProcessor(
167
) -> SparseEmbeddingResponse:
168
num_prompt_tokens = 0
169
response_data = []
103
- return_tokens = self._get_sparse_embedding_request(request_id).return_tokens
170
+ raw_request = self._get_sparse_embedding_request(request_id)
171
+ has_dense_embed = raw_request.embed_task in ["dense", "dense&sparse"]
172
+ has_sparse_embed = raw_request.embed_task in ["sparse", "dense&sparse"]
173
+ embed_dimensions = 0
174
+ if has_dense_embed:
175
+ embed_dimensions = (
176
+ self.embed_dimensions
177
+ if raw_request.dimensions is None
178
+ else raw_request.dimensions
179
+ )
180
for idx in range(len(model_output)):
181
mo = model_output[idx]
106
- sparse_embedding: dict[int, float] = {}
182
+ sparse_embedding_dict: dict[int, float] = {}
183
num_prompt_tokens += len(mo.prompt_token_ids)
108
- if len(mo.prompt_token_ids) != len(mo.outputs.data):
109
- # this is the case that add_special_tokens is True,
110
- # which means first token and last token are special tokens
111
- mo.prompt_token_ids = mo.prompt_token_ids[1:]
112
- for token_id, weight in zip(mo.prompt_token_ids, mo.outputs.data.tolist()):
113
- sparse_embedding[token_id] = max(
114
- weight, sparse_embedding.get(token_id, 0.0)
184
+ dense_embedding: list[float] | None = None
185
+ sparse_embedding: list[SparseEmbeddingTokenWeight] | None = None
186
+ if has_dense_embed:
187
+ dense_embedding = mo.outputs.data[:embed_dimensions].tolist()
188
+ if has_sparse_embed:
189
+ sparse_weights = mo.outputs.data[embed_dimensions:].tolist()
190
+ if len(mo.prompt_token_ids) != len(sparse_weights):
191
+ # this is the case that add_special_tokens is True,
192
+ # which means first token and last token are special tokens
193
+ mo.prompt_token_ids = mo.prompt_token_ids[1:]
194
+ for token_id, weight in zip(mo.prompt_token_ids, sparse_weights):
195
+ sparse_embedding_dict[token_id] = max(
196
+ weight, sparse_embedding_dict.get(token_id, 0.0)
197
+ )
198
+ sparse_embedding = self._build_sparse_embedding_token_weights(
199
+ sparse_embedding_dict,
200
+ raw_request.return_tokens,
201
)
202
+
203
response_data.append(
204
SparseEmbeddingResponseData(
205
index=idx,
119
- sparse_embedding=self._build_sparse_embedding_token_weights(
120
- sparse_embedding,
121
- return_tokens,
122
- ),
206
+ object=raw_request.embed_task,
207
+ sparse_embedding=sparse_embedding,
208
+ dense_embedding=dense_embedding,
209
)
210
)
211
tests/plugins/bge_m3_sparse_plugin/bge_m3_sparse_processor/types.py
+31
-4
@@ -1,18 +1,44 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from typing import Literal, get_args
5
+
6
from pydantic import BaseModel, Field
7
8
from vllm.entrypoints.openai.engine.protocol import UsageInfo
7
-from vllm.entrypoints.pooling.base.protocol import CompletionRequestMixin
9
+from vllm.entrypoints.pooling.base.protocol import (
10
+ CompletionRequestMixin,
11
+ EmbedRequestMixin,
12
+)
13
+
14
+EmbedTask = Literal[
15
+ "sparse",
16
+ "dense",
17
+ "dense&sparse",
18
+]
19
+
20
+EMBED_TASKS: tuple[EmbedTask, ...] = get_args(EmbedTask)
21
22
10
-class SparseEmbeddingCompletionRequestMixin(CompletionRequestMixin):
23
+class SparseEmbeddingCompletionRequestMixin(CompletionRequestMixin, EmbedRequestMixin):
24
return_tokens: bool | None = Field(
25
default=None,
26
description="Whether to return dict shows the mapping of token_id to text."
27
"`None` or False means not return.",
28
)
29
+ embed_task: EmbedTask = Field(
30
+ default="dense&sparse",
31
+ description="embed task, can be one of 'sparse', 'dense' , 'dense&sparse', "
32
+ "default to 'dense&sparse'",
33
+ )
34
+
35
+ def to_embed_requests_offline(self) -> list[EmbedRequestMixin]:
36
+ if isinstance(self.input, list):
37
+ return [self] * len(self.input)
38
+ return [self]
39
+
40
+ def to_embed_requests_online(self) -> list[EmbedRequestMixin]:
41
+ return [self]
42
43
44
class SparseEmbeddingTokenWeight(BaseModel):
@@ -23,8 +49,9 @@ class SparseEmbeddingTokenWeight(BaseModel):
49
50
class SparseEmbeddingResponseData(BaseModel):
51
index: int
26
- object: str = "sparse-embedding"
27
- sparse_embedding: list[SparseEmbeddingTokenWeight]
52
+ object: str = "dense&sparse"
53
+ sparse_embedding: list[SparseEmbeddingTokenWeight] | None
54
+ dense_embedding: list[float] | None
55
56
57
class SparseEmbeddingResponse(BaseModel):
tests/plugins_tests/test_bge_m3_sparse_io_processor_plugins.py
+24
-1
@@ -19,6 +19,12 @@ model_config = {
19
),
20
}
21
22
+dense_embedding_sum = [
23
+ -0.7214539647102356, # "What is the capital of France?"
24
+ -0.6926871538162231, # "What is the capital of Germany?"
25
+ -0.7129564881324768, # "What is the capital of Spain?"
26
+]
27
+
28
29
def _float_close(expected: object, result: object):
30
assert isinstance(expected, float) and isinstance(result, float), (
@@ -33,6 +39,12 @@ def _get_attr_or_val(obj: object | dict, key: str):
39
return getattr(obj, key, None)
40
41
42
+def _check_dense_embedding(data, index=0):
43
+ assert _float_close(sum(data), dense_embedding_sum[index]), (
44
+ "dense-embedding result not match"
45
+ )
46
+
47
+
48
def _check_sparse_embedding(data, check_tokens=False):
49
expected_weights = [
50
{"token_id": 32, "weight": 0.0552978515625, "token": "?"},
@@ -109,7 +121,7 @@ async def test_bge_m3_sparse_plugin_online(
121
assert len(_get_attr_or_val(parsed_response, "data")) > 0
122
123
data_entry = _get_attr_or_val(parsed_response, "data")[0]
112
- assert _get_attr_or_val(data_entry, "object") == "sparse-embedding"
124
+ assert _get_attr_or_val(data_entry, "object") == "dense&sparse"
125
assert _get_attr_or_val(data_entry, "sparse_embedding")
126
127
# Verify sparse embedding format
@@ -117,6 +129,11 @@ async def test_bge_m3_sparse_plugin_online(
129
assert isinstance(sparse_embedding, list)
130
_check_sparse_embedding(sparse_embedding, return_tokens)
131
132
+ # Verify dense embedding format
133
+ dense_embedding = _get_attr_or_val(data_entry, "dense_embedding")
134
+ assert isinstance(dense_embedding, list)
135
+ _check_dense_embedding(dense_embedding)
136
+
137
# Verify usage information
138
usage = _get_attr_or_val(parsed_response, "usage")
139
assert usage, f"usage not found for {parsed_response}"
@@ -164,6 +181,9 @@ def test_bge_m3_sparse_plugin_offline(vllm_runner, return_tokens: bool):
181
sparse_embedding = output.sparse_embedding
182
assert isinstance(sparse_embedding, list)
183
_check_sparse_embedding(sparse_embedding, return_tokens)
184
+ dense_embedding = output.dense_embedding
185
+ assert isinstance(dense_embedding, list)
186
+ _check_dense_embedding(dense_embedding)
187
188
# Verify usage
189
assert response.usage.prompt_tokens > 0
@@ -206,6 +226,9 @@ def test_bge_m3_sparse_plugin_offline_multiple_inputs(vllm_runner):
226
# Each output should have sparse embeddings
227
sparse_embedding = output.sparse_embedding
228
assert isinstance(sparse_embedding, list)
229
+ dense_embedding = output.dense_embedding
230
+ assert isinstance(dense_embedding, list)
231
+ _check_dense_embedding(dense_embedding, i)
232
233
# Verify usage
234
assert response.usage.prompt_tokens > 0
vllm/model_executor/layers/pooler/special.py
+39
-1
@@ -170,4 +170,42 @@ class BOSEOSFilter(Pooler):
170
return pooled_outputs
171
172
173
-__all__ = ["BOSEOSFilter", "DispatchPooler", "IdentityPooler"]
173
+class BgeM3Pooler(Pooler):
174
+ def __init__(self, token_classify_pooler: Pooler, embed_pooler: Pooler) -> None:
175
+ super().__init__()
176
+ self.token_classify_pooler = token_classify_pooler
177
+ self.embed_pooler = embed_pooler
178
+
179
+ def forward(
180
+ self, hidden_states: torch.Tensor, pooling_metadata: PoolingMetadata
181
+ ) -> PoolerOutput:
182
+ embed_outputs = self.embed_pooler(hidden_states, pooling_metadata)
183
+ token_classify_outputs = self.token_classify_pooler(
184
+ hidden_states, pooling_metadata
185
+ )
186
+ pooler_outputs: list[torch.Tensor] = []
187
+ for embed_output, token_classify_output in zip(
188
+ embed_outputs, token_classify_outputs
189
+ ):
190
+ pooler_outputs.append(
191
+ torch.cat(
192
+ [embed_output.view(-1), token_classify_output.view(-1)], dim=-1
193
+ )
194
+ )
195
+
196
+ return pooler_outputs
197
+
198
+ def get_supported_tasks(self) -> Set[PoolingTask]:
199
+ return {"embed&token_classify"}
200
+
201
+ def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
202
+ return self.embed_pooler.get_pooling_updates(
203
+ "embed"
204
+ ) | self.token_classify_pooler.get_pooling_updates("token_classify")
205
+
206
+ def extra_repr(self) -> str:
207
+ s = f"supported_task={self.get_supported_tasks()}"
208
+ return s
209
+
210
+
211
+__all__ = ["BOSEOSFilter", "DispatchPooler", "IdentityPooler", "BgeM3Pooler"]
vllm/model_executor/models/roberta.py
+16
-10
@@ -10,6 +10,7 @@ from transformers import RobertaConfig
10
11
from vllm.config import ModelConfig, PoolerConfig, VllmConfig
12
from vllm.model_executor.layers.pooler import (
13
+ BgeM3Pooler,
14
BOSEOSFilter,
15
DispatchPooler,
16
Pooler,
@@ -216,24 +217,29 @@ class BgeM3EmbeddingModel(RobertaEmbeddingModel):
217
self.colbert_linear = nn.Linear(
218
self.hidden_size, self.hidden_size, dtype=self.head_dtype
219
)
220
+ embed_pooler = pooler_for_embed(pooler_config)
221
+ token_classify_pooler = BOSEOSFilter(
222
+ pooler_for_token_classify(
223
+ pooler_config,
224
+ pooling=AllPool(),
225
+ classifier=self.sparse_linear,
226
+ act_fn=torch.relu,
227
+ ),
228
+ self.bos_token_id,
229
+ self.eos_token_id,
230
+ )
231
232
return DispatchPooler(
233
{
222
- "embed": pooler_for_embed(pooler_config),
234
+ "embed": embed_pooler,
235
"token_embed": BOSEOSFilter(
236
pooler_for_token_embed(pooler_config, self.colbert_linear),
237
self.bos_token_id,
238
# for some reason m3 only filters the bos for colbert vectors
239
),
228
- "token_classify": BOSEOSFilter(
229
- pooler_for_token_classify(
230
- pooler_config,
231
- pooling=AllPool(),
232
- classifier=self.sparse_linear,
233
- act_fn=torch.relu,
234
- ),
235
- self.bos_token_id,
236
- self.eos_token_id,
240
+ "token_classify": token_classify_pooler,
241
+ "embed&token_classify": BgeM3Pooler(
242
+ token_classify_pooler, embed_pooler
243
),
244
}
245
)
vllm/pooling_params.py
+4
@@ -96,6 +96,10 @@ class PoolingParams(
96
self.skip_reading_prefix_cache = True
97
return
98
99
+ # skipping verify, let plugins configure and validate pooling params
100
+ if self.task not in self.valid_parameters:
101
+ return
102
+
103
# NOTE: Task validation needs to done against the model instance,
104
# which is not available in model config. So, it's not included
105
# in this method
vllm/tasks.py
+7
-1
@@ -6,7 +6,13 @@ GenerationTask = Literal["generate", "transcription", "realtime"]
6
GENERATION_TASKS: tuple[GenerationTask, ...] = get_args(GenerationTask)
7
8
PoolingTask = Literal[
9
- "embed", "classify", "score", "token_embed", "token_classify", "plugin"
9
+ "embed",
10
+ "classify",
11
+ "score",
12
+ "token_embed",
13
+ "token_classify",
14
+ "plugin",
15
+ "embed&token_classify",
16
]
17
POOLING_TASKS: tuple[PoolingTask, ...] = get_args(PoolingTask)
18