🚀 The feature, motivation and pitch

copied from feature request in vllm upstream

Great first issue!

Should enable for both torchax and jax (if needed).

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.