[Model][-/N] Improve all pooling task | Support chunked prefill with ALL pooling

[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
• [Model][-/N] Improve all pooling task | Support chunked prefill with ALL pooling #27145
• [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][5/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][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

Address:

• [Bug]: "pooling_type='ALL' no longer supported for embeddings #25165 (comment)
• [RFC]: Support Returning Prompt Hidden States #24288 (comment)

Let's support chunked prefill with ALL pooling

We are moving further towards supporting both generation and Returning Prompt Hidden States simultaneously

Test Plan

tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py

Test Result

pass

---

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 enables chunked prefill for ALL pooling tasks, which is a great step towards improving performance for long prompts. The changes involve modifying the AllPool layer to cache hidden states during chunked prefill and updating related metadata and data structures. A new test is also added to validate this functionality.

I've found a couple of critical issues that need to be addressed. One is a mutable default argument in PoolingParams which can lead to shared state and incorrect behavior. The other is an incorrect return type in is_partial_prefill which will cause runtime errors. Please see my detailed comments for suggestions on how to fix these.

[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, like "@codex address that feedback".

[noooop]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for chunked prefill with ALL pooling, which is a valuable enhancement. The core logic in AllPool is updated to cache hidden states across chunks and process them upon completion of prefill. A corresponding test case has been added to validate this new functionality. While the overall approach is sound, I've identified a few critical issues. Notably, a change to is_partial_prefill appears to break other pooling methods. There's also a potential memory leak and correctness bug in AllPool due to the hidden states cache not being cleared. Lastly, a minor typing issue in a method could lead to unexpected behavior. Please see my detailed comments for suggestions on how to address these points.

[noooop]

/gemini review

[gemini-code-assist]

Code Review

This pull request enables chunked prefill for ALL pooling tasks, which was previously unsupported. The changes involve modifying the AllPool layer to cache hidden states across multiple prefill steps and return the complete hidden states only when a request is finished. This is a good step towards more flexible pooling operations. The related metadata and worker logic have been updated accordingly, and a new test case is added to verify the functionality. My review focuses on improving code clarity and maintainability.

[noooop]

/gemini review

[noooop]

cc @DarkLight1337 @maxdebayser

Ready for review

I came up with a simple method to support chunked prefill with ALL pooling, without modifying too much core code.

[gemini-code-assist]

Code Review

This pull request introduces support for chunked prefill with ALL pooling tasks, which was previously unsupported. The changes are well-structured and involve modifying the AllPool layer to cache hidden states across prefill chunks, updating PoolingParams and PoolingMetadata to support this, and adjusting the GPUModelRunner to handle the new logic. A new test case is also added to validate the implementation.

My main feedback is on a design choice in PoolingParams. I've left a comment with a suggestion for a more robust implementation that better separates data transfer objects from runtime state. Overall, this is a good step towards more flexible and efficient pooling operations.

[noooop]

Another strange problem was discovered at the same time

[noooop]

@heheda12345

I want all pooling requests to not use prefix_caching, even if enable_prefix_caching=True at the request level. Otherwise, the output hidden_states would be fewer than n_prompt_tokens.

However, I don't want to completely disable prefix_caching when using all pooling, because enabling prefix_caching can accelerate scenarios where generation and returning prompt hidden states are performed simultaneously.

PTAL #24288 (comment)

If you turn on Automatic Prefix Caching, submit LLM.encode(..., pooling_task="token_embed") first, then submit LLM.generate. I think it should be no cost in making two separate calls.

[noooop]

/gemini review

[gemini-code-assist]

Code Review

This pull request successfully enables chunked prefill for ALL pooling tasks, which was previously unsupported. This is a valuable enhancement for models utilizing token_embed or token_classify. The implementation, including the stateful handling of hidden states in AllPool and updates to related data structures, appears correct and well-integrated. The tests have been appropriately updated and new ones added to validate this new capability. I've identified one minor issue in a new test case where an assertion doesn't seem to match its intended purpose. Other than that, this is a solid contribution.

[DarkLight1337]

@maxdebayser can you help review this? Thanks

[maxdebayser]

Thanks for completing this. Back when added support in V1 for pooling and chunked prefill, I left the addition of support for the stateful poolers as a future exercise. My main concern in this PR is that the PoolingParams object is not a good place to store this state.

I have two different suggestions of places where the state could go:

1. PoolingMetadata

The pooling metadata object could have a state dict that could be used by the poolers. The GPU input batch would have to be modified to handle persistent pooling metadata objects as well as the pooling params object. The advantage here is that in the case of request cancellation we can easily remove the pooling metadata.

2. The pooler

If we add the request ID in the pooling metadata, then the pooler can have a dict of request ID to state. The only challenge would be to get notified of request cancellation to remove the state.

BTW, the it's easy to extend the chunked prefill support for mean pooling as well. It doesn't even need to store the full hidden states, only the current average and the count of tokens has to be stored.

[noooop]

My main concern in this PR is that the PoolingParams object is not a good place to store this state.

+1

I can’t seem to find a better place to store it either. I tried storing it in the request, but that would require too many code changes.

mainly because it needs to be notified to delete the stored data to prevent memory leaks.

The primary benefit of storing data in PoolingParams is that it can always be garbage-collected along with the request. and without requiring any (or almost any) modifications to the core code.

[maxdebayser]

I think the easiest of the suggestions would be to add the state to the PoolingMetadata. Currently the input batch has a dict of request ID to PoolingParams. And when get_pooling_metadata() is called, we build a PoolingMetadata object and initialize it with a reference to pooling params. But what if the input batch had a dict of request to PoolingMetadata instead, and each PoolingMetadata has a PoolingParams object?

[noooop]

I think the easiest of the suggestions would be to add the state to the PoolingMetadata.

vllm/vllm/v1/worker/gpu_input_batch.py

Lines 834 to 841 in bfe0b4b

	def get_pooling_metadata(self) -> PoolingMetadata:
	pooling_params = self.get_pooling_params()
	return PoolingMetadata(
	prompt_lens=torch.from_numpy(self.num_prompt_tokens[: self.num_reqs]),
	prompt_token_ids=self.sampling_metadata.prompt_token_ids,
	pooling_params=pooling_params,
	)

Each executor step creates a new PoolingMetadata. If hidden_states_cache is in PoolingMetadata, the position needs to be moved every step . More importantly, during two executions, some requests might not be scheduled, and these unscheduled requests' hidden_states_cache don't know where to place them.

---

Mean pooling + chunked prefill can be considered later; in fact, we don't have a model that uses mean pooling + chunked prefill yet.

[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