[Feature]: Support for multiple embedding types in a single inference call

[staugust]

Purpose

As issue #35190 proposed, support both sparse&dense embedding in a single inference call.

Test Plan

./tests/plugins_tests/test_bge_m3_sparse_io_processor_plugins.py

Test Result

All test cases passed. Both sparse&dense embedding is ok.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

The pull request introduces support for multiple embedding types in a single inference call, which is a significant feature enhancement. The changes involve updating various components to handle a list of pooling tasks instead of a single task. This includes modifications to request mixins, pooling parameters, IO processor logic, and model runner output structures. The test cases have also been updated to cover the new multi-task scenarios.

However, there are critical issues identified in the handling of request queues within the BgeM3SparseEmbeddingsProcessor and a list initialization bug in DispatchPooler that could lead to incorrect behavior or crashes. Additionally, the parameter validation for multiple pooling tasks in PoolingParams needs to be more robust.

[mergify]

Hi @staugust, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[DarkLight1337]

Hmm, at the end of the day we are still having multiple tasks in PoolingParams. So IOProcessor cannot really avoid this complexity. If that's the case then I suggest splling up the PR into two parts: the first part supports multiple tasks in PoolingParams while the second part updates the plugin.

[staugust]

@DarkLight1337 ok, shall we update pooling api to support multi pooling tasks?

[DarkLight1337]

Sure

[mergify]

Hi @staugust, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[noooop]

I don't think it's worth making such a big change for this feature.

I think it would be better to add one or more combination of tasks, rather than supporting multiple tasks.
For example, adding a task : "embed+token_embed+token_classify"

For this specific requirement, we only need to add one task to output the results of embed + token_embed + token_classify. Do we really need to broadly support multiple tasks?

[mergify]

Hi @staugust, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy or markdownlint failing?

mypy and markdownlint are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10
# For markdownlint
pre-commit run --hook-stage manual markdownlint

[staugust]

@noooop I've modify PoolingTask and two related PoolingRequest data class. For sparse&dense embedding, we only need to add 'token_classify+embed'. Anyway, core.py, scheduler.py and gpu_model_runner have to be updated to process two pooling_output of one request properly, PTAL, thanks.

[staugust]

Hi @noooop, @DarkLight1337 — just wanted to kindly check if you have any updates or feedback on this PR when you have a moment. Thanks!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @staugust.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mergify]

Hi @staugust, the pre-commit checks have failed. Please run:

uv pip install pre-commit
pre-commit install
pre-commit run --all-files

Then, commit the changes and push to your branch.

For future commits, pre-commit will run automatically on changed files before each commit.

Tip

Is mypy failing?

mypy is run differently in CI. If the failure is related to this check, please use the following command to run it locally:

# For mypy (substitute "3.10" with the failing version if needed)
pre-commit run --hook-stage manual mypy-3.10

[noooop]

@staugust

We are very sorry, but we plan to deprecate support for multi-task, (for v2 runner), meaning users will need to select only one of them at the beginning.

Please modify this bge-m3 plugin to use only the dense & sparse tasks, meaning that for all requests, the dense & sparse tasks will be computed. When returning results to users, filter the content they need.

This may increase overhead, but we have no other choice.

[noooop]

Sorry, I'm not 100% sure how you implemented support for dense, sparse, and dense&sparse simultaneously, but you need to change it to only use "embed&token_classify": BgeM3Pooler.

[staugust]

@staugust

We are very sorry, but we plan to deprecate support for multi-task, (for v2 runner), meaning users will need to select only one of them at the beginning.

Please modify this bge-m3 plugin to use only the dense & sparse tasks, meaning that for all requests, the dense & sparse tasks will be computed. When returning results to users, filter the content they need.

This may increase overhead, but we have no other choice.

@noooop Can a pooling model still support multi pooling tasks by /pooling api with different task configured in pooling request body ? Or pooling task is specified by command line?

[noooop]

pooling task is specified by command line

[noooop]

We had some discussions, summaries:

Make all supported pooling tasks available simultaneously (multi-task support) vs selecting only one at the start.

1. Typically, model training focuses on a single task — 99.9% of usage scenarios only require picking one of them at the beginning.

2. Having all supported pooling tasks ready at once is certainly fancier and gives users everything they want in one go.

3. In terms of UX, if users have to pick just one at the start, they'd need to manually select token_embed for multivector retrieval. e.g. jina_embeddings_v4 https://github.com/vllm-project/vllm/blob/main/examples/pooling/token_embed/jina_embeddings_v4_offline.py

4. BGE-M3 supports three tasks: embed for dense retrieval; token_embed for multi-vector retrieval; and token_classify for sparse retrieval.
   This PR ([Feature]: Support for multiple embedding types in a single inference call #35829) updates the BGE-M3 plugin to support dense, sparse, and dense+sparse modes, making it more convenient for BGE-M3 users.

[noooop]

If you have any important information, feel free to add.

[staugust]

@noooop I’d like to understand the reasoning behind v2 runner only supporting a single specific task at model serving instance startup.
If this is mainly for user experience, my assumption is that users calling the pooling API already know which pooling task their model supports, so it may not be necessary to configure the task via the command line.

[noooop]

This will simplify the design of the v2 runner; otherwise, supporting all pooling tasks running simultaneously (multi-task support) would make the runner complex.

[staugust]

@noooop ok, I'll update this plugin when v2 runner is ready.

[noooop]

It’s because the block v2 feature is being rolled out, so we need to deprecate support for multi-task first. That means you can proceed with implementing this feature now.

[staugust]

ok, I'll issue a pr to update plugin.