Accessing decoder only model's hidden states with Executor API

[tkhanipov]

Hello,

We have a decoder only Llama-like model and we need to access its hidden states. Currently for this purpose we patch both the Python LLaMAModel and the C++ GptSession classes. In the former we do hidden_states.mark_output(...) to make these visible and in the latter we access them with outputBuffer.at(...) here.

We would like to timely switch to the new Executor API (with the current API announced to be deprecated) but cannot find the proper way to achieve what we want there.

What would be the easiest way to access the hidden states in the new API? Is it possible to do within something like executorExampleBasic?

The returnEncoderOutput setting looked promising but it does not seem to work for a decoder only model (as expected from the comment): Result::encoderOutput is empty.

[Advait251206]

For decoder-only models, returnEncoderOutput is expected to be empty, since there is no encoder stage. That option is intended for encoder-decoder architectures and won't expose the decoder's hidden states.

From what you've described, your current approach of:

1. marking hidden_states as an engine output during model construction, and
2. retrieving it from the runtime output buffers,

is the standard way to expose intermediate tensors.

As far as I'm aware, the new Executor API does not currently provide a high-level mechanism to request arbitrary intermediate activations (such as decoder hidden states) through the public API. The Result object exposes generated tokens, log probabilities, encoder outputs (when applicable), and related metadata, but not internal decoder activations.

So unless I've overlooked a recently added feature, you would likely still need to:

• modify the model/network definition to mark the hidden states as outputs, and
• extend the Executor/runtime to propagate those additional outputs through to Result (or another user-visible interface).

I don't believe this can be achieved purely from an application like executorExampleBasic without corresponding engine/runtime changes, because the TensorRT engine itself only returns tensors that were explicitly marked as outputs during build time.

One question for the TensorRT-LLM team:

• Is there a supported mechanism (or roadmap) for exposing arbitrary intermediate tensors, such as decoder hidden states, through the Executor API without patching the runtime? For example, something analogous to Hugging Face's output_hidden_states=True that could be enabled at build or inference time.

That would be useful for applications such as feature extraction, embeddings, reranking, and custom downstream tasks.

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206