[Feature]: AutoDeploy: push hidden states (logits) gather to right after last token mixer layer

🚀 The feature, motivation and pitch

Right now, we gather hidden_states before the final lm_head to save on compute. However, in principle we can push back the gathering to just after the final token mixer layer.

This would for example enable us to save on massive prefill requests when they are expensive layers like MoE after the final token mixer layer.

Moreover, for chunked prefill we can also just skip subsequent layers if we don't need context logits

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.