[TPU host offload] Setup precompile functions to TPU host offload

[saikat-royc]

Description

Start with a short description of what the PR does and how this is a change from
the past.

This PR introduces a bucketing strategy of the swap operations. This decomposes operations involving arbitrary numbers of KV cache blocks into smaller, predefined block size buckets. This approach leverages JAX's JIT compilation cache, avoiding costly recompilations for varying input shapes during the hot path of serving.

The rest of the description includes relevant details and context, examples:

• why is this change being made,

1. prepare precompile functions which will cycle through the load and save jitted functions
2. decompose the load and save util functions to be block buckets aligned. This would ensure that compilation is not triggered in the hot path of the serving

• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

• existing unit tests
• sglang tool benchmark. some preliminary test cases show x% improvement in the tokens/s and TTFT improvements, because we are avoiding any recompilation during the actual load and save operations

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.