[Model][Last/N] Improve all pooling task | Generate runner supports using embed and token_embed tasks.

[noooop]

Improve all pooling task

• [Frontend][Doc][5/N] Improve all pooling task | Polish encode (pooling) api & Document. #25524
• [Model] Add num_cached_tokens for PoolingRequestOutput #27378
• [Model] Allow users to control skip reading cache per request. #28194
• Improve enable chunked_prefill & prefix_caching logic. #26623
• [Model][6/N] Improve all pooling task | Support chunked prefill with ALL pooling #27145
• [Model][7/N] Improve all pooling task | Deprecation as_reward_model. Extract hidden states prefer using new multi-vector retrieval API #26686
• [Frontend] Binary embedding response does not return metadata by setting encoding_format to bytes_only. #30249
• [Deprecation] Deprecation --convert reward, use --convert embed instead. #30463
• [Model] Automatic conversion of TokenClassification model #30666
• [Model][Last/N] Improve all pooling task | Generate runner supports using embed and token_embed tasks. #30672

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

Generate runner supports using embed and token_embed tasks & End the Improve all pooling tasks series

FIX #11905
FIX #24288
FIX #6165
FIX #4435

Test Plan

tests/models/language/pooling/test_extract_hidden_states.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 updates a test for pooling tasks to align with recent improvements, including support for returning prompt hidden states during generation. The test now uses the 'generate' runner and includes a new case to verify text generation with prefix caching.

My main feedback is that the new test case, while intended to verify the return of prompt hidden states, lacks assertions to confirm their presence and correctness. I've suggested adding these assertions to ensure the feature is properly tested.

[noooop]

@DarkLight1337

Are we planning to implement generate runner support using embed and token_embed tasks?

With this very very dirty fix can make the tests/models/language/pooling/test_extract_hidden_states.py test pass, but we can't batch the generated request and pooling request together.

[DarkLight1337]

but we can't batch the generated request and pooling request together.

I think this limitation is ok if we can alternate between generative and pooling batches

[noooop]

but we can't batch the generated request and pooling request together.

I think this limitation is ok if we can alternate between generative and pooling batches

I will further refine this PR if it's decided to implement “Generate runner support using embed and token_embed tasks”

[noooop]

hello @breakices

PTAL: #24288 (comment)

token_embed can be used to extract Prompt Hidden States

• [Model][2/N] Improve all pooling task | Support multi-vector retrieval #25370
• [Model][7/N] Improve all pooling task | Deprecation as_reward_model. Extract hidden states prefer using new multi-vector retrieval API #26686

Could using "generate runner support with embed and token_embed tasks" as a form of "Returning Prompt Hidden States" help with RLVR?

[noooop]

hello @charlotte12l

A long time has passed since my last comment #24288 (comment), and I'm finally about to implement it.

Could using "generate runner support with embed and token_embed tasks" as a form of "Returning Prompt Hidden States" help with your use case.