add io_process_plugin for sparse embedding

[staugust]

Purpose

Update sparse embedding output for bge-m3, which contains both token_id and corresponding weights, a draft pr for #33882

Test Plan

run vllm in offline and online serving mode, and call llm.encode, /pooling api to fetch sparse embedding result.

Test Result

online mode

pooling result for ["I love China", "hi, bge-m3"].

+ curl -X POST http://127.0.0.1:41001/pooling -H 'Content-Type: application/json' -d '{
  "model": "auto",
  "task": "token_classify",
  "data" : {"input":["I love China", "hi, bge-m3"], "return_tokens": true }
}'
# response
{
  "request_id": "pool-9b3c403bbdad8437",
  "created_at": 1771054438,
  "data": {
    "data": [
      {
        "index": 0,
        "object": "sparse-embedding",
        "sparse_embedding": [
          {
            "token_id": 87,
            "weight": 0.1947021484375,
            "token": "I"
          },
          {
            "token_id": 5161,
            "weight": 0.286865234375,
            "token": "love"
          },
          {
            "token_id": 9098,
            "weight": 0.3173828125,
            "token": "China"
          }
        ]
      },
      {
        "index": 1,
        "object": "sparse-embedding",
        "sparse_embedding": [
          {
            "token_id": 1274,
            "weight": 0.1895751953125,
            "token": "hi"
          },
          {
            "token_id": 4,
            "weight": 0.1082763671875,
            "token": ","
          },
          {
            "token_id": 876,
            "weight": 0.11322021484375,
            "token": "b"
          },
          {
            "token_id": 429,
            "weight": 0.22265625,
            "token": "ge"
          },
          {
            "token_id": 9,
            "weight": 0.0894775390625,
            "token": "-"
          },
          {
            "token_id": 39,
            "weight": 0.1475830078125,
            "token": "m"
          },
          {
            "token_id": 363,
            "weight": 0.258544921875,
            "token": "3"
          }
        ]
      }
    ],
    "usage": {
      "prompt_tokens": 14,
      "total_tokens": 14,
      "completion_tokens": 0,
      "prompt_tokens_details": null
    }
  }
}

vllm serve command line:

vllm serve \
  --model "BAAI/bge-m3" \
  --served-model-name auto \
  --max-num-batched-tokens 524288 \
  --max-num-seqs 1024 \
  --runner pooling \
  --enable-prompt-tokens-details \
  --no-enable-prefix-caching \
  --dtype float16 \
  --hf_overrides '{"architectures": ["BgeM3EmbeddingModel"]}' \
  --optimization-level 2 \
  --io-processor-plugin bge_m3_sparse_plugin \
  --trust-remote-code  

offline mode

import vllm
import json
from vllm.pooling_params import PoolingParams

llm = vllm.LLM(model='/home/local/workspace/vllm-embeddings/models/BAAI/bge-m3', hf_overrides={"architectures": ["BgeM3EmbeddingModel"]}, io_processor_plugin="bge_m3_sparse_plugin")

pooling_params = PoolingParams(task='token_classify')

req = {"data": {"input":"I love China"}}

print('r1 is: ')
r1 = llm.encode(req, pooling_params=pooling_params, pooling_task='token_classify')
for out in r1:
    print(json.dumps(out.outputs.model_dump(), indent=2))

req["data"]["input"] = [ "I love China", "hi, bge-m3"]
req["data"]["return_tokens"] = True

r2 = llm.encode(req, pooling_params=pooling_params, pooling_task='token_classify')
print('r2 is: ')
for out in r2:
    print(json.dumps(out.outputs.model_dump(), indent=2))

The output is exactly the same with FlagEmbedding.

r1 is:
{
  "data": [
    {
      "index": 0,
      "object": "sparse-embedding",
      "sparse_embedding": [
        {
          "token_id": 87,
          "weight": 0.19468513131141663,
          "token": null
        },
        {
          "token_id": 5161,
          "weight": 0.28681907057762146,
          "token": null
        },
        {
          "token_id": 9098,
          "weight": 0.3174937665462494,
          "token": null
        }
      ]
    }
  ],
  "usage": {
    "prompt_tokens": 5,
    "total_tokens": 5,
    "completion_tokens": 0,
    "prompt_tokens_details": null
  }
}
r2 is:
{
  "data": [
    {
      "index": 0,
      "object": "sparse-embedding",
      "sparse_embedding": [
        {
          "token_id": 87,
          "weight": 0.19468513131141663,
          "token": "I"
        },
        {
          "token_id": 5161,
          "weight": 0.28681907057762146,
          "token": "love"
        },
        {
          "token_id": 9098,
          "weight": 0.3174937665462494,
          "token": "China"
        }
      ]
    },
    {
      "index": 1,
      "object": "sparse-embedding",
      "sparse_embedding": [
        {
          "token_id": 1274,
          "weight": 0.18952831625938416,
          "token": "hi"
        },
        {
          "token_id": 4,
          "weight": 0.10824987292289734,
          "token": ","
        },
        {
          "token_id": 876,
          "weight": 0.11320684850215912,
          "token": "b"
        },
        {
          "token_id": 429,
          "weight": 0.22263310849666595,
          "token": "ge"
        },
        {
          "token_id": 9,
          "weight": 0.08946191519498825,
          "token": "-"
        },
        {
          "token_id": 39,
          "weight": 0.1476268619298935,
          "token": "m"
        },
        {
          "token_id": 363,
          "weight": 0.25863635540008545,
          "token": "3"
        }
      ]
    }
  ],
  "usage": {
    "prompt_tokens": 14,
    "total_tokens": 14,
    "completion_tokens": 0,
    "prompt_tokens_details": null
  }
}

output of FlagEmbeddings

{
  "dense_vecs": null,
  "lexical_weights": [
    {
      "87": 0.19470215,
      "5161": 0.2866211,
      "9098": 0.31762695
    },
    {
      "1274": 0.18969727,
      "4": 0.10821533,
      "876": 0.11315918,
      "429": 0.22290039,
      "9": 0.089660645,
      "39": 0.14746094,
      "363": 0.25854492
    }
  ],
  "colbert_vecs": null
}

---

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

This pull request introduces a new IO processor plugin for sparse embeddings with bge-m3. The changes look good overall, but I have a few points of feedback. There's a critical bug in post_process where special tokens are not handled correctly, which will lead to incorrect embeddings. I've also identified a potential TypeError in validate_or_generate_params due to a possible None return, and suggested using isinstance instead of type in parse_request for more robust type checking.

[staugust]

@maxdebayser @noooop Would you like to take a look? I've implemented an io_processor_plugin to generate sparse output, and it seems ok for both offline and online mode.

[maxdebayser]

This is looking good!

Here is an idea that I was exploring a couple of weeks ago: to move the plugin code under model_executor/models so that it would be loaded automatically for BgeM3EmbeddingModel.

The idea is as follows:

1. Modify vllm/plugins/io_processors/__init__.py::get_io_processor to first try to load the plugin as a classname so that we don't need to actually install a module.
2. Implement verify_and_update_model_config for BgeM3EmbeddingModel as in vllm/model_executor/models/config.py so that it modifies vllm_config to always load the plugin for BgeM3EmbeddingModel (if no other plugin is set)

[maxdebayser]

@staugust , if you like the idea, feel free to cherry-pick this commit to load classes as plugins: maxdebayser@52eb0bd

[staugust]

This is looking good!

Here is an idea that I was exploring a couple of weeks ago: to move the plugin code under model_executor/models so that it would be loaded automatically for BgeM3EmbeddingModel.

The idea is as follows:

1. Modify vllm/plugins/io_processors/__init__.py::get_io_processor to first try to load the plugin as a classname so that we don't need to actually install a module.
2. Implement verify_and_update_model_config for BgeM3EmbeddingModel as in vllm/model_executor/models/config.py so that it modifies vllm_config to always load the plugin for BgeM3EmbeddingModel (if no other plugin is set)

@maxdebayser I really like this idea. Not requiring users to install an extra plugin package would make it much easier for them to use the vLLM io_processor_plugin.
We might also need to think about how to manage custom io_processor_plugins inside vLLM. For example, we could:
Create a new directory under vllm/plugins/io_processors to store all custom io_processor_plugins, or
Introduce a git submodule in the vLLM repo to host plugins in a separate project, and then include these custom io_processor_plugins when building the vLLM distribution.
What do you think about this approach? @noooop @DarkLight1337

[DarkLight1337]

Let's create a vllm/io_processors directory and put the plugins there. I'm planning to eventually treat it similar to vllm/renderers.

[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

[staugust]

@DarkLight1337 Please let me know once the vllm/io_processors directory has been created. Thanks.

[DarkLight1337]

I need to move Renderer outside of InputPreprocessor before I can do this, so it's going to take a while.

[staugust]

@DarkLight1337 Got it, I'd hold this pr until the prerequisite work is completed.

[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 @DarkLight1337 @christian-pinto As discussed in #34419 , I've updated the implementation to follow io_process_plugin design, PTAL, 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 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

[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

[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

[christian-pinto]

LGTM

I will open a PR to separate the terratorch io processor plugin tests from the generic io processor plugin tests.

[noooop]

There are few H100s available currently. Please wait patiently.

[staugust]

@noooop ok, there was a check failed with info

Mergify / Rule: auto-rebase if approved, ready, and 40 commits behind main (rebase)
 Base branch update has failed
staugust does not have write access to the forked repository.

I guessed the code might be behind the main branch, so I tried merging the main branch.