[RFC]: Support Returning Prompt Hidden States

[charlotte12l]

Motivation.

Hidden states are a valuable feature for many downstream applications, and their support has been requested and discussed in the vLLM community for a while(#15434 #4435 #6165 )

Currently, users who need both the prompt hidden states and the generated text must make two separate calls: LLM.encode (to obtain hidden states) and LLM.generate (to obtain generated text). This workflow is inconvenient and inefficient, especially for applications that require both outputs. The inability to return hidden states alongside generated text has also been highlighted in #12249

Proposed Change.

This is a very initial draft PR: #24202. Will polish it after RFC.

Configuration

Add return_prompt_hidden_states in Sampling Config.

There is prior discussion #15434 (comment) mentioning using such request-level parameter can't be handled well. In local testing, this approach appears feasible for prompt hidden states, but feedback on potential disadvantages is welcome.

How to build up prompt_hidden_states_tensor

The process for building the prompt hidden states tensor is similar to the existing _get_prompt_logprobs_dict implementation

vllm/vllm/v1/worker/gpu_model_runner.py

Line 203 in d0944b2

The main difference is that we construct a CPU tensor of shape (num_prompt_tokens, hidden_size). During chunked prefill, we copy the relevant chunk slices into it.

Hidden States Aggregation

To support flexible downstream use, we propose adding a plugin interface that allows users to perform custom aggregations on the hidden states (e.g., pooling, etc.).

Output

When return_prompt_hidden_states is enabled, the output will include:

• prompt_hidden_states: a tensor of shape (num_prompt_tokens, hidden_size)

Perf Concerns

• Default (disabled): No performance regression is expected.
• Enabled: Storing and serializing a large (num_prompt_tokens, hidden_size) tensor may impact performance, especially with long prompts. We will benchmark performance with large input lengths to assess the impact and determine if the regression is acceptable.

Feedback Period.

1-3 weeks

CC List.

@luccafong @zhuohan123

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[WoosukKwon]

As we discussed briefly offline, I'm wondering if it makes more sense to have an API to apply desired operations to hidden states and return the outputs, instead of returning the raw hidden states. This may save the serialization cost a lot and actually better serve the needs.

[charlotte12l]

As we discussed briefly offline, I'm wondering if it makes more sense to have an API to apply desired operations to hidden states and return the outputs, instead of returning the raw hidden states. This may save the serialization cost a lot and actually better serve the needs.

@WoosukKwon Thanks for the advice! Regarding "apply desired operations":
Are you suggesting that users could implement their own plugins to perform these operations (as mentioned in the Hidden States Aggregation section), or that we should implement a set of supported operations within the framework itself?

• In the former case, there’s a chance that raw hidden states could still be returned if a user’s plugin does so.
• In the latter case, we could enforce that raw hidden states are never returned (e.g., by raising errors if attempted).

[luccafong]

we can follow #22820 to add entry point for plugin to extend output post processor logic in gpu model runner to compute the desired processed hidden states. the other parts can be similar as current IO plugin. @WoosukKwon @charlotte12l

[noooop]

Currently, users who need both the prompt hidden states and the generated text must make two separate calls: LLM.encode (to obtain hidden states) and LLM.generate (to obtain generated text).

Overall I think it would be more reasonable to make two separate calls. Letting the generate api return embedding will make generate api unnecessary complex

#25370 is expected to land after the release. It can be used to efficiently extract prompt hidden states.

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.

(Submitting LLM.generate first and then LLM.encode requires recalculating the prompt, and Automatic Prefix Caching will not help. This will lead to inefficiency.)

---

Another detail is that chunked prefill does not support all pooling (and mean pooling).

#20930 makes the LAST Pooling model use chunked prefill by default and disables the encode API.

This caused us a lot of trouble PTAL #24853 #25165

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

[charlotte12l]

@noooop Appreciate the inputs!

Trying to understand how to use pooling models after the recent PRs: to confirm, with the current pooling models API, if we want to get each prompt token's hidden states, we need to use LLM.reward(...)? I saw #25370 was primarily for token_embed and token_classify.

Besides, does all pooling and mean pooling supported for the reward API with/without chunked prefill?