[Model][2/N] Improve all pooling task | Support multi-vector retrieval

[noooop]

Break change

• Split the encode task into two tasks: token_embed and token_classify
  • token_embed is the same as embed, using normalize as activation
  • token_classify is the same as classify, default using softmax as activation (we actually allow classify and token_classify to use any activation function by setting act_fn. )
• Complete remove softmax from PoolingParams and perfer using activation, since we actually allow classify and token_classify to use any activation function.
• pooling_task required for llm.encode and /pooling <- This will take effect entirely in the next PR ([Frontend][Doc][Last/N] Improve all pooling task | Polish encode (pooling) api & Document. #25524)

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
• [Model][5/N] Improve all pooling task | Support chunked prefill with ALL pooling #27145
• [Frontend][3/N] Improve all pooling task | Support binary embedding response #27066
• [Frontend][4/N] Improve all pooling task | Add plugin pooling task #26973
• Support bge-m3 sparse embeddings (lexical weights) #14526 <- It is best to be after [Frontend][4/N] Improve all pooling task | Add plugin pooling task #26973, before [Frontend][Doc][Last/N] Improve all pooling task | Polish encode (pooling) api & Document. #25524
• [Model][6/N] Improve all pooling task | Improve enable chunked_prefill & prefix_caching logic. #26623
• [Model][7/N] Improve all pooling task | Deprecation as_reward_model. Extract hidden states prefer using new multi-vector retrieval API #26686
• [Frontend][Doc][Last/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

After #21227 landed, we hope that pooling models can always use all pooling, and users don’t need to manually set using all pooling.

The current encode api (/pooling api ) mainly targets the classify for each token scenario (e.g. TokenClassification #24872 & reward models), overlooked the embed for each token scenario.

Let's support embed for each token scenario (multi-vector retrieval)

Partial_Fix #25165

We are stepping closer to support ColBERT & ColPali

cc @DarkLight1337 @maxdebayser

Test Plan

tests/models/language/pooling/test_multi_vector_retrieval.py
tests/test_pooling_params.py

Test Result

pass

Known Issues

• Maybe we should find a way to support chunked prefill + all pooling (and mean pooling))
• support ColBERT & ColPali

---

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.

[noooop]

@jupyterjazz

try:

from vllm import LLM

llm = LLM(
    model="jinaai/jina-embeddings-v4-vllm-text-matching",
    enforce_eager=True,
    max_model_len=1024,
    enable_chunked_prefill=False,  # <- In order to use the encode api
    runner="pooling")

prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "The future of AI is",
]

outputs = llm.embed(prompts)

for prompt, output in zip(prompts, outputs):
    embeds = output.outputs.embedding
    print(len(embeds))

outputs = llm.encode(prompts, pooling_task="encode")

for prompt, output in zip(prompts, outputs):
    multi_vector = output.outputs.data
    print(multi_vector.shape)

Do you ok with this api and outputs?

---

There are still some broken features that need to be fixed. But the multi_vector feature is now testable.

[noooop]

@DarkLight1337

Ready for review

(Slight) Break change

• Split the encode task into two tasks: token_embed and token_classify
  • token_embed is the same as embed, using normalize as activation
  • token_classify is the same as classify, default using softmax as activation (we actually allow classify and token_classify to use any activation function by setting act_fn. )
  • Use the following code for compatibility:

def encode2pooling_task(supported_tasks):
    # Currently no model supports both token_embed and token_classify.
    if "token_embed" in supported_tasks:
        return "token_embed"
    elif "token_classify" in supported_tasks:
        return "token_classify"
    else:
        raise ValueError(f"pooling_task must be one of {supported_tasks}.")

• Complete remove softmax from PoolingParams and perfer using activation, since we actually allow classify and token_classify to use any activation function.

[jupyterjazz]

Hi @noooop,

I just tested and it works fine. The only thing it was missing is override_pooler_config=PoolerConfig(normalize=False). I don't have a strong opinion on this but setting normalization to False by default during encode could make sense because oftentimes you would want the actual last hidden layer at this step. Other than this, the API looks good to me. Thank you for such a quick fix!

[DarkLight1337]

Please update this page https://docs.vllm.ai/en/latest/models/pooling_models.html#model-conversion to not use encode task anymore

[DarkLight1337]

I'll delay the merge of this PR until after the release so we don't have to worry about back-compatibility issues which further complicate future PRs

[maxdebayser]

Overall the changes in this PR look good to me but I would prefer to still keep the generic "encode" task around for uses cases that don't cleanly fit with the old and new tasks introduced in this PR. As an example, the use case introduced by this PR: #22820 can't really be described as "token_embed", "token_classify", "embed", "classify" or "score".

[mergify]

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

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

[noooop]

Overall the changes in this PR look good to me but I would prefer to still keep the generic "encode" task around for uses cases that don't cleanly fit with the old and new tasks introduced in this PR. As an example, the use case introduced by this PR: #22820 can't really be described as "token_embed", "token_classify", "embed", "classify" or "score".

Now only mypy will check whether pooling_task belongs to PoolingTask, so pooling_task in encode api can accept any str.

In fact, we basically support custom pooling task and pooler. As long as user implement a OOT model with pooler, and use the corresponding pooling_task in encode. LOL

---

Do we need to move towards allowing users to use pooling task plugins?

[DarkLight1337]

cc @christian-pinto

[noooop]

Can I temporarily skip this test?
( I think there is a conflict in dependencies between terratorch and vllm

No, we cannot skip it

It seems this pr can pass the test, but I don't know how it happened.

[DarkLight1337]

Maybe you can try to pip install the specific version of terratorch

[noooop]

Maybe you can try to pip install the specific version of terratorch

vllm/requirements/test.in

Line 57 in b8a4572

terratorch @ git+https://github.com/IBM/terratorch.git@1.1.rc3 # required for PrithviMAE test

pip install git+https://github.com/IBM/terratorch.git@1.1.rc3

It seems I can only rely on ci.

[christian-pinto]

probably this is happening because on CI we use the compiled list of dependencies, rather than manually pip installing the various packages?

Have you tried installing all dependencies from here: https://github.com/vllm-project/vllm/blob/main/requirements/test.in

We are in the process of releasing a new version of TerraTorch and then we will be able to just pip install from pypy rather than a specific tag from git. I will post a PR to fix that in vLLM soon.

[noooop]

By the way, this PR splits the encode task into two tasks: token_embed and token_classify.

• token_embed is the same as embed, using normalize as activation.
• token_classify is the same as classify, default using softmax as activation (we actually allow classify and token_classify to use any activation function by setting act_fn.)

prithvi_mae and test_io_processor_plugins are not one of "token_embed", "token_classify", "embed", "classify" or "score".

We no longer have the encode task that puts all the others into it.

@christian-pinto How do you think about the plugin pooling task?

[noooop]

@DarkLight1337

Are there any more modifications needed for this PR?

[DarkLight1337]

As long as the test passes then I'm fine with it

[christian-pinto]

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.

[noooop]

As long as the test passes then I'm fine with it

(You forgot to approve it