[Feature][Hardware][TPU]:Reduce the compile time #14582
Description
🚀 The feature, motivation and pitch
After the fix of #14310,
We have num_token_bucket compilations for the main model and num_token_bucket x num_reqs_bucket for the logits processor.
We can make some improvement on this, as the num_token_bucket x num_reqs_bucket only happens on hidden_states[logits_indices], where we select part of the hidden states. Therefore, we can partition the graph to 3 parts:
main model: num_token_bucket
hidden_states[logits_indices]: num_token_bucket x num_reqs_bucket
logits: num_reqs_bucket
Alternatives
No response
Additional context
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.