[Frontend][4/N] Improve all pooling task | Add plugin pooling task

[noooop]

Improve all pooling task

• [Model][0/N] Improve all pooling task | clean up #25817
• [Frontend][1/N] Improve all pooling task | Support FP16 Embedding Base64 (Still uses fp32 by default). #26414
• [Model][2/N] Improve all pooling task | Support multi-vector retrieval #25370
• [Frontend][3/N] Improve all pooling task | Support binary embedding response #27066
• [Frontend][4/N] Improve all pooling task | Add plugin pooling task #26973
• [Frontend][Doc][5/N] Improve all pooling task | Polish encode (pooling) api & Document. #25524

These PRs are mostly conflicting with each other, so combining them into a series would better inform reviewers about what happened. And what else needs to be done after that?

Purpose

plugin task uses io_processor.parse_request to verify inputs, skipping PoolingParams verify

Test Plan

Test Result

---

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.

[mergify]

Documentation preview: https://vllm--26973.org.readthedocs.build/en/26973/

[gemini-code-assist]

Code Review

This pull request correctly introduces a new 'plugin' pooling task, which is a valuable addition for models that utilize custom IO processors for pooling, such as Terratorch. The changes are well-implemented across the codebase, including updates to examples and tests. However, I have identified one critical bug in the new DummyPooler implementation that must be fixed.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR.

• @codex fix this CI failure
• @codex address that feedback

[noooop]

When it comes to the IO processor plugins, the type of pooling activation function depends on the combination of model and plugin. As an example, for PrithviMAE we instantiate a All pooler and disable softmax because the plugin is expecting to receive the raw output from the pooler. Other models in the future might have a different behavior and require a different pooling strategy and activation function, which is fine and should still work.

However, I agree with @maxdebayser regarding keeping the encode task to support more generic cases like this one, without having to resort to using the remaining tasks that can be confusing for people.

We need to change at least the name to avoid confusion with the previous encode task.

@christian-pinto @maxdebayser @DarkLight1337

After careful consideration, we still need a plugin pooling task, which uses io_processor.parse_request to verify inputs, skipping the verification of PoolingParams.

[DarkLight1337]

parse_request is for OpenAI serving only. I think we might need to introduce a new method in the plugin explicitly to verify sampling/pooling params, or pass the params to pre_process method

[christian-pinto]

parse_request is for OpenAI serving only. I think we might need to introduce a new method in the plugin explicitly to verify sampling/pooling params, or pass the params to pre_process method

The request parsing is done in both offline and online mode, it's meant to verify that the input is appropriate for the plugin. I am fine with either methods to be honest. At the moment, only for the online case, the Pooling parameters are created in the IOProcessorRequest, while in the offline case the parameters are just coming from the request or defaults are given.

I can take care of this change and have the plugin generating/validating pooling parameters.

[noooop]

parse_request is for OpenAI serving only. I think we might need to introduce a new method in the plugin explicitly to verify sampling/pooling params, or pass the params to pre_process method

The request parsing is done in both offline and online mode, it's meant to verify that the input is appropriate for the plugin. I am fine with either methods to be honest. At the moment, only for the online case, the Pooling parameters are created in the IOProcessorRequest, while in the offline case the parameters are just coming from the request or defaults are given.

I can take care of this change and have the plugin generating/validating pooling parameters.

sorry, I can't run prithvi_geospatial_mae locally, so please help fix all possible errors in examples and/or all other places .

Will it be more convenient if I invite you to collaborate on this PR?

The devil is in the details; perhaps when we fix these details, we may discover bigger problems.

[christian-pinto]

parse_request is for OpenAI serving only. I think we might need to introduce a new method in the plugin explicitly to verify sampling/pooling params, or pass the params to pre_process method

The request parsing is done in both offline and online mode, it's meant to verify that the input is appropriate for the plugin. I am fine with either methods to be honest. At the moment, only for the online case, the Pooling parameters are created in the IOProcessorRequest, while in the offline case the parameters are just coming from the request or defaults are given.
I can take care of this change and have the plugin generating/validating pooling parameters.

sorry, I can't run prithvi_geospatial_mae locally, so please help fix all possible errors in examples and/or all other places .

Will it be more convenient if I invite you to collaborate on this PR?

The devil is in the details; perhaps when we fix these details, we will discover bigger problems.

Sure, invite me to this PR. I will be able to work on it starting on Monday though. Right now I have another task that I need to finish

[DarkLight1337]

The request parsing is done in both offline and online mode, it's meant to verify that the input is appropriate for the plugin.

You're right, sorry I misremembered this (was afk when I sent the previous message so I couldn't do code search to check 😓 )

[maxdebayser]

I think this extra task makes sense. Should we also add extra_kwargs to the request classes in protocol.py and in PoolingParams? In this way we would be able to pass custom args to the plugin.

[noooop]

I think this extra task makes sense. Should we also add extra_kwargs to the request classes in protocol.py and in PoolingParams? In this way we would be able to pass custom args to the plugin.

IOProcessor uses a data-in, data-out, where the data can be any thing

Do we still need a separate extra_kwargs, or can extra_kwargs be part of the data? The user can actually use any data format.

[noooop]

@christian-pinto PTAL #27063 #27066

I hope io_processor_plugins can support binary response, it will definitely be more efficient compared to base64.

[maxdebayser]

IOProcessor uses a data-in, data-out, where the data can be any thing

Yes, you're right.

[noooop]

On the compression bit, I assume this is something that happens after the plugin was applied? If we compress the output of the pooler right away, then the plugin would have to de-compress in order to be applied and this is not ideal as we would just be wasting time. If the compression happens after the post_process is applied in the /pooling endpoint then the plugins are not affected. Only the client will have to know the payload is compressed. But this I assume would be either something requested by the client itself, or a feature of the deployment where users know that data is being compressed.

Compression occurs after post-processing is applied in the /pooling endpoint,

I just want to confirm if there are any compatibility issues or what improvements are needed to make the plugin task + binary response more efficient.

[noooop]

@christian-pinto #27066 (binary response) has been merged.

Please help check if it can be used together with plugin task.

It would be best to have an example to tell users how to use binary response with prithvi_geospatial_mae .

[christian-pinto]

I have added a new function to the plugins interface that can be used for validating or generating params. There were a few other changes needed to make the tests pass - i.e., the new DummyPooler return one less dimension in the output and the test plugin needed to be changed to account for that.

[christian-pinto]

mypy should be happy now 🤞

[noooop]

@christian-pinto Thanks for your help

Please confirm if it can still run after merging #27204

PTAL:

• [Frontend] Require flag for loading text and image embeds #27204 (comment)
• https://buildkite.com/vllm/ci/builds/35962/steps/canvas?sid=019a0df9-dd02-4486-bf12-b506da36bc73

---

Please update the documentation for IO Processor Plugins

https://docs.vllm.ai/en/latest/design/io_processor_plugins.html

[christian-pinto]

After #27393 gets merged this one will pass fine too. Let me fix the documentation.

[noooop]

The new plugin pooling task looks great. @christian-pinto Thanks for your help!

cc @DarkLight1337

Are there any more modifications needed for this PR?