Support running inference with KV cache on Kithara's MaxText models.

[wenxindongwork]

This PR integrates MaxEngine with Kithara's MaxTextModel to improve inference performance. Previously, model inference (i.e. MaxTextModel.generate()) is supported via a naive autoregressive for-loop without KV Cache, which results in slow tokens per second for longer sequences. Now, inference for MaxTextModel is backed by a KV Cache system.

Fast inference is required for supporting online and offline evaluation, as well as on-policy RLHF.

It should also be mentioned that Kithara's KerasHubModel.generate() is already backed by KV Cache.

Key Changes

• Added JetStream as a submodule
• Patched MaxText's MaxEngine to intake an parameterized model
• Redesigned the model generation API to support a wider range of input formats:
  • String inputs (single or batched)
  • Token inputs as integer lists or numpy arrays
• Added configurable max_prefill_predict_length parameter. This parameter determines the max prefill length.
• Updated documentation
• Support batched inference with progress tracking

Testing

• Added unit tests verifying compatibility with MaxText models

[lienchen0526]

I left some python readability stuff here. We can discuss which can be a good practice from both of our sides

[wenxindongwork]

thank you Jerry for the thorough code review! And thanks for catching the typo :)