fix(pooling): validate BGE-M3 combined task ownership (#51823)
Signed-off-by: Zhe Li <2843409461@qq.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Ola committed
Aug 17, 2026 at 18:30 UTC
95901ce70a625cd152e465d1b15441fd199bde24
4 files changed
+193
-18
docs/models/pooling_models/specific_models.md
+79
-16
@@ -365,36 +365,99 @@ curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
365
366
## BAAI/bge-m3
367
368
-The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json`
369
-the architecture is declared as `XLMRobertaModel`, which makes `vLLM` load it as a vanilla ROBERTA model without the
370
-extra weights. To load the full model weights, override its architecture like this:
368
+`BAAI/bge-m3` supports dense retrieval, lexical matching, and ColBERT-style
369
+multi-vector retrieval. Its `config.json` declares `XLMRobertaModel`, so vLLM
370
+otherwise loads it as a vanilla RoBERTa model without the extra sparse and
371
+ColBERT weights. The examples below therefore override the architecture with
372
+`BgeM3EmbeddingModel`.
373
+
374
+The three retrieval modes map to concrete pooling tasks as follows:
375
+
376
+| Retrieval mode | Pooling task | Output |
377
+| -------------- | ------------ | ------ |
378
+| Dense | `embed` | One embedding vector per input |
379
+| Lexical/sparse | `token_classify` | One scalar weight per non-special token |
380
+| ColBERT multi-vector | `token_embed` | One embedding vector per non-special token |
381
+
382
+Serve one concrete mode by selecting its task at load time:
383
+
384
+```shell
385
+vllm serve BAAI/bge-m3 \
386
+ --runner pooling \
387
+ --hf-overrides '{"architectures": ["BgeM3EmbeddingModel"]}' \
388
+ --pooler-config.task <task>
389
+```
390
+
391
+For dense embeddings, replace `<task>` with `embed` and use the Embeddings API:
392
393
```shell
373
-vllm serve BAAI/bge-m3 --hf-overrides '{"architectures": ["BgeM3EmbeddingModel"]}'
394
+curl -s http://localhost:8000/v1/embeddings \
395
+ -H "Content-Type: application/json" -d '{
396
+ "model": "BAAI/bge-m3",
397
+ "input": ["What is BGE M3?", "Definition of BM25"]
398
+ }'
399
```
400
376
-Then you obtain the sparse embeddings like this:
401
+For lexical weights, replace `<task>` with `token_classify` and use the
402
+Pooling API:
403
404
```shell
405
curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
380
- "model": "BAAI/bge-m3",
381
- "task": "token_classify",
382
- "input": ["What is BGE M3?", "Definition of BM25"]
406
+ "model": "BAAI/bge-m3",
407
+ "task": "token_classify",
408
+ "input": ["What is BGE M3?", "Definition of BM25"]
409
}'
410
```
411
412
Due to limitations in the output schema, the output consists of a list of
387
-token scores for each token for each input. This means that you'll have to call
388
-`/tokenize` as well to be able to pair tokens with scores.
389
-Refer to the tests in `tests/models/language/pooling/test_bge_m3.py` to see how
390
-to do that.
413
+token scores for each input. Call `/tokenize` as well to pair token IDs with
414
+their scores. See
415
+[`test_bge_m3.py`](../../../tests/models/language/pooling/test_bge_m3.py) for a
416
+complete example that also combines repeated token IDs.
417
+
418
+For ColBERT vectors, replace `<task>` with `token_embed` and use the Pooling API:
419
+
420
+```shell
421
+curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
422
+ "model": "BAAI/bge-m3",
423
+ "task": "token_embed",
424
+ "input": ["What is BGE M3?", "Definition of BM25"]
425
+}'
426
+```
427
+
428
+### Dense and sparse output through an IO processor plugin
429
+
430
+The source tree includes a reference
431
+[BGE-M3 IO processor plugin](../../../tests/plugins/bge_m3_sparse_plugin) that
432
+formats dense embeddings, sparse token weights, or both in one response. From
433
+a source checkout, install it in the vLLM environment and load it as follows:
434
392
-You can obtain the colbert embeddings like this:
435
+```shell
436
+uv pip install ./tests/plugins/bge_m3_sparse_plugin
437
+
438
+vllm serve BAAI/bge-m3 \
439
+ --runner pooling \
440
+ --hf-overrides '{"architectures": ["BgeM3EmbeddingModel"]}' \
441
+ --io-processor-plugin bge_m3_sparse_plugin
442
+```
443
+
444
+The plugin selects the internal `embed&token_classify` task so the model
445
+computes dense and lexical outputs together. Public requests must use task
446
+`plugin` and put the plugin-specific fields under `data`:
447
448
```shell
449
curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
396
- "model": "BAAI/bge-m3",
397
- "task": "token_embed",
398
- "input": ["What is BGE M3?", "Definition of BM25"]
450
+ "model": "BAAI/bge-m3",
451
+ "task": "plugin",
452
+ "data": {
453
+ "input": ["What is BGE M3?", "Definition of BM25"],
454
+ "embed_task": "dense&sparse",
455
+ "return_tokens": true
456
+ }
457
}'
458
```
459
+
460
+`embed_task` accepts `dense`, `sparse`, or `dense&sparse`. The combined
461
+`embed&token_classify` task is an internal execution contract for this plugin,
462
+not a generic Pooling API response format. Without the plugin, select one of
463
+the three concrete tasks above.
tests/entrypoints/pooling/test_factories.py
new
+94
@@ -0,0 +1,94 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from unittest.mock import MagicMock
5
+
6
+import pytest
7
+
8
+from vllm.entrypoints.pooling import factories
9
+from vllm.entrypoints.pooling.factories import init_pooling_io_processors
10
+from vllm.entrypoints.pooling.pooling.io_processor import (
11
+ PluginWithIOProcessorPlugins,
12
+ UnsupportedCombinedTaskIOProcessor,
13
+)
14
+
15
+
16
+def _bge_m3_config(io_processor_plugin=None):
17
+ model_config = MagicMock()
18
+ model_config.get_pooling_task.return_value = "embed&token_classify"
19
+ model_config.io_processor_plugin = io_processor_plugin
20
+ model_config.hf_config.to_dict.return_value = {}
21
+ model_config.architecture = "BgeM3EmbeddingModel"
22
+
23
+ vllm_config = MagicMock(model_config=model_config)
24
+ renderer = MagicMock()
25
+ renderer._executor = MagicMock()
26
+ chat_template_config = MagicMock(
27
+ chat_template=None,
28
+ chat_template_content_format="auto",
29
+ trust_request_chat_template=False,
30
+ )
31
+ return vllm_config, renderer, chat_template_config
32
+
33
+
34
+def test_combined_task_without_plugin_uses_rejection_processor():
35
+ vllm_config, renderer, chat_template_config = _bge_m3_config()
36
+
37
+ processors = init_pooling_io_processors(
38
+ supported_tasks=("embed", "embed&token_classify"),
39
+ vllm_config=vllm_config,
40
+ renderer=renderer,
41
+ chat_template_config=chat_template_config,
42
+ )
43
+
44
+ assert processors.keys() == {"embed&token_classify"}
45
+ assert isinstance(
46
+ processors["embed&token_classify"], UnsupportedCombinedTaskIOProcessor
47
+ )
48
+
49
+
50
+def test_combined_task_with_plugin_uses_plugin_processor(monkeypatch):
51
+ vllm_config, renderer, chat_template_config = _bge_m3_config("bge_m3_sparse_plugin")
52
+ monkeypatch.setattr(factories, "has_io_processor", lambda *_: True)
53
+ monkeypatch.setattr(
54
+ "vllm.entrypoints.pooling.pooling.io_processor.get_io_processor",
55
+ lambda *_: MagicMock(),
56
+ )
57
+
58
+ processors = init_pooling_io_processors(
59
+ supported_tasks=("embed", "embed&token_classify"),
60
+ vllm_config=vllm_config,
61
+ renderer=renderer,
62
+ chat_template_config=chat_template_config,
63
+ )
64
+
65
+ assert processors.keys() == {"embed&token_classify", "plugin"}
66
+ assert isinstance(processors["plugin"], PluginWithIOProcessorPlugins)
67
+
68
+
69
+def test_combined_task_plain_pooling_request_has_actionable_error(monkeypatch):
70
+ from vllm.entrypoints.pooling.pooling.protocol import PoolingCompletionRequest
71
+ from vllm.entrypoints.pooling.pooling.serving import ServingPooling
72
+
73
+ vllm_config, renderer, chat_template_config = _bge_m3_config("bge_m3_sparse_plugin")
74
+ monkeypatch.setattr(factories, "has_io_processor", lambda *_: True)
75
+ monkeypatch.setattr(
76
+ "vllm.entrypoints.pooling.pooling.io_processor.get_io_processor",
77
+ lambda *_: MagicMock(),
78
+ )
79
+
80
+ engine_client = MagicMock(renderer=renderer, vllm_config=vllm_config)
81
+ models = MagicMock(model_config=vllm_config.model_config)
82
+ serving = ServingPooling(
83
+ engine_client,
84
+ models,
85
+ supported_tasks=("embed", "embed&token_classify"),
86
+ request_logger=None,
87
+ chat_template_config=chat_template_config,
88
+ )
89
+ request = PoolingCompletionRequest(model="BAAI/bge-m3", input=["hola"])
90
+
91
+ assert serving.io_processors.keys() == {"embed&token_classify", "plugin"}
92
+ io_processor = serving.get_io_processor(request)
93
+ with pytest.raises(ValueError, match="plugin request with a 'data' field"):
94
+ io_processor.create_pooling_params(request)
vllm/entrypoints/pooling/factories.py
+8
-2
@@ -65,10 +65,16 @@ def init_pooling_io_processors(
65
66
processors["token_embed"] = TokenEmbedIOProcessor
67
68
- if has_io_processor(
68
+ if pooling_task == "embed&token_classify":
69
+ from .pooling.io_processor import UnsupportedCombinedTaskIOProcessor
70
+
71
+ processors[pooling_task] = UnsupportedCombinedTaskIOProcessor
72
+
73
+ has_plugin = has_io_processor(
74
vllm_config,
75
model_config.io_processor_plugin,
71
- ):
76
+ )
77
+ if has_plugin:
78
from .pooling.io_processor import PluginWithIOProcessorPlugins
79
80
processors["plugin"] = PluginWithIOProcessorPlugins
vllm/entrypoints/pooling/pooling/io_processor.py
+12
@@ -24,6 +24,18 @@ from .protocol import IOProcessorRequest, IOProcessorResponse
24
logger = init_logger(__name__)
25
26
27
+class UnsupportedCombinedTaskIOProcessor(PoolingIOProcessor):
28
+ name = "embed&token_classify"
29
+
30
+ def create_pooling_params(self, request):
31
+ raise ValueError(
32
+ "The 'embed&token_classify' pooling task is only available "
33
+ "through an IO processor plugin. Send a plugin request with "
34
+ "a 'data' field, "
35
+ "or select a concrete task with --pooler-config.task."
36
+ )
37
+
38
+
39
class PluginWithoutIOProcessorPlugins(PoolingIOProcessor):
40
# Some models, such as Terratorch (tests/models/test_terratorch.py),
41
# use plugin tasks in the pooler but do not use IO Processor plugins.