[Feature][Hardware][TPU]:Reduce the compile time

[yaochengji]

🚀 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.

[yaochengji]

Fixed by #16275