[Model Runner v2] Enable BGE M3 pooling embed token_classify (#50661)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Aug 2, 2026 at 11:46 UTC
55c98e370aa058f567a9e682dc0652bdfba6b0bb
2 files changed
+29
-9
tests/models/language/pooling/test_splade_sparse_pooler.py
+23
-7
@@ -192,19 +192,15 @@ def test_encoder_cache_reset_clears_late_interaction_state() -> None:
192
193
def test_pooling_runner_rejects_unsupported_selected_task() -> None:
194
model = MagicMock()
195
- model.pooler.get_supported_tasks.return_value = {
196
- "embed",
197
- "embed&token_classify",
198
- "token_classify",
199
- }
195
+ model.pooler.get_supported_tasks.return_value = {"embed", "plugin"}
196
vllm_config = MagicMock()
197
vllm_config.scheduler_config.max_num_seqs = 2
198
vllm_config.model_config.attn_type = "encoder_only"
203
- vllm_config.model_config.get_pooling_task.return_value = "embed&token_classify"
199
+ vllm_config.model_config.get_pooling_task.return_value = "plugin"
200
201
with (
202
patch.object(PoolingRunner, "get_supported_tasks", return_value=["embed"]),
207
- pytest.raises(ValueError, match="selects 'embed&token_classify'"),
203
+ pytest.raises(ValueError, match="selects 'plugin'"),
204
):
205
PoolingRunner(model, vllm_config)
206
@@ -273,3 +269,23 @@ def test_pooling_runner_filters_decoder_token_embedding() -> None:
269
runner = PoolingRunner(model, vllm_config)
270
271
assert runner.supported_tasks == {"embed"}
272
+
273
+
274
+def test_pooling_runner_gates_embed_token_classification_to_encoder() -> None:
275
+ # BGE-M3's combined task emits per-token weights alongside the embedding,
276
+ # so it is enabled only where prefill is unchunked.
277
+ model = MagicMock()
278
+ model.pooler.get_supported_tasks.return_value = {"embed", "embed&token_classify"}
279
+ vllm_config = MagicMock()
280
+ vllm_config.scheduler_config.max_num_seqs = 2
281
+
282
+ vllm_config.model_config.attn_type = "encoder_only"
283
+ vllm_config.model_config.get_pooling_task.return_value = "embed&token_classify"
284
+ assert PoolingRunner(model, vllm_config).supported_tasks == {
285
+ "embed",
286
+ "embed&token_classify",
287
+ }
288
+
289
+ vllm_config.model_config.attn_type = "decoder"
290
+ vllm_config.model_config.get_pooling_task.return_value = "embed"
291
+ assert PoolingRunner(model, vllm_config).supported_tasks == {"embed"}
vllm/v1/worker/gpu/pool/pooling_runner.py
+6
-2
@@ -18,9 +18,13 @@ from vllm.v1.worker.gpu.input_batch import InputBatch
18
from vllm.v1.worker.gpu.states import RequestState
19
20
_SUPPORTED_TASKS: frozenset[PoolingTask] = frozenset(
21
- {"embed", "classify", "token_embed", "token_classify"}
21
+ {"embed", "classify", "token_embed", "token_classify", "embed&token_classify"}
22
+)
23
+# `embed&token_classify` concatenates a fixed-size embedding with per-token
24
+# weights, so its output length scales with the prompt like the others here.
25
+_TOKEN_TASKS: frozenset[PoolingTask] = frozenset(
26
+ {"token_embed", "token_classify", "embed&token_classify"}
27
)
23
-_TOKEN_TASKS: frozenset[PoolingTask] = frozenset({"token_embed", "token_classify"})
28
29
30
class PoolingRunner: