Repository navigation
how to return hidden_states? #4414
Description
use gather_context_logits and gather_generation_logits can help to return logits. but how to return hidden_states? plz help, thx
{{ message }}
use gather_context_logits and gather_generation_logits can help to return logits. but how to return hidden_states? plz help, thx
Activity
brb-nv commentedon May 29, 2025
I don't think there's a way to return intermediate
hidden_statesright now. @hchings, can you please confirm?hchings commentedon May 29, 2025
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 commentedon Jun 18, 2025
Issue has not received an update in over 14 days. Adding stale label.
xxchauncey commentedon Jun 26, 2025
same problem, hidden_states for additional lm_head needed.
FateForever0222 commentedon Dec 15, 2025
same problem, do we now support returning hidden states? 🙏 @hchings
hchings commentedon Jan 9, 2026
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 commentedon Feb 24, 2026
same problem, I need the last hidden states before logits since I don't need the generate stage
hous-lab commentedon Aug 7, 2026
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', ...)beforegather_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 fromself.model()beforeLogitsProcessorcompresses it — works withadditional_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!
[docs] Add VLA full hidden_states export example
[docs] Add VLA full hidden_states export example
[docs] Add VLA full hidden_states export example
[docs] Add VLA full hidden_states export example