how to return hidden_states?

[lss15151161]

use gather_context_logits and gather_generation_logits can help to return logits. but how to return hidden_states? plz help, thx

[brb-nv]

I don't think there's a way to return intermediate hidden_states right now. @hchings, can you please confirm?

[hchings]

I don't think we have flags for that as @brb-nv mentioned in either TRT or PyTorch backend.

@lss15151161 you could share more context on use cases for our team's awareness. Thanks. cc @laikhtewari

[github-actions]

Issue has not received an update in over 14 days. Adding stale label.

[xxchauncey]

same problem, hidden_states for additional lm_head needed.

[FateForever0222]

same problem, do we now support returning hidden states? 🙏 @hchings

[hchings]

Currently, we don't. @FateForever0222 @xxchauncey - is the ask to be able to get the hidden state of any modules/layer, or just the llm_haed?

[CASHADEr]

same problem, I need the last hidden states before logits since I don't need the generate stage

[hous-lab]

Hi all — I've been working on deploying a VLA model (Orion, ICCV'25 autonomous driving) with TensorRT-LLM and encountered the same need for full hidden_states export.

I've opened PR #17420 with a complete example showing two solutions:

Solution A (TRT backend, v0.7–v0.21): Insert hidden_states.mark_output('full_hidden_states', ...) before gather_last_token_logits. Verified in production: hidden_states CosSim=0.9994 (INT8), end-to-end plan_L2=0.686.

Solution B (PyTorch backend, v1.x): In DecoderModelForCausalLM.forward, attach the raw 3D tensor from self.model() before LogitsProcessor compresses it — works with additional_model_outputs=["hidden_states"].

The key insight is that VLA models need the complete 3D [batch, seq_len, hidden_dim] tensor (at a specific token position like a "waypoint" token), not just the last token. This is different from standard LLM interpretability use cases.

Would love feedback from @hchings and the team on the preferred integration approach!