Prefix Caching with HBM and latency test

[yuyanpeng-google]

Description

Implements Prefix Caching in HBM and latency test in inference_microbenchmark.

Stores prefix tokens as a trie for fast lookup index of PrefixCache store in cache.

Insert longer Key replace shorter key to be the longest common prefix key.
The shorter key will never be returned even if longer key is erased, and should got evicted in the future.

Assume Key is equal length to tokens, which can be used to slice prompt and cache Value.
Should check the return key common prefix length by the caller.

If erase the Key not the leaf, nothing will happen.
If erased key match at a leaf, delete the node and ancestors would be the leaf after deleted.

Value will be moved to the cache, which means cannot used the same value reference after add_to_cache.

The jax may modified the value even stored in another python reference.
If the value need to be used after add_to_cache, make sure copy them before add_to_cache.

Return value copied from cache to avoid modified value in the cache, always copied the value before return.

Add PrefixCaching benchmark test in inference_microbenchmark.

Using half of the prefill_length as the common prefix and save 100 prefix in the cache.

Loading the cache (including jax.array.copy) appears to be independent of the prefill_length (tested with 128 and 1024), even though the saved cache sizes are different.

Using jax.profiler shows that the copy operation consumes a similar amount of time on TPU. This might be because the sizes aren't large or different enough to see a significant impact.

Part of results below

Prefix Cache benchmark results for prefill length 128:

PrefixCaching results:
	Per prefix size bytes: 0.124 GB
	Average save cache time: 12.142 ms
	Average fetch longest prefix time: 0.029 ms
	Average load cache time: 5.589 ms


Prefix Cache benchmark results for prefill length 1024:

PrefixCaching results:
	Per prefix size bytes: 0.220 GB
	Average save cache time: 12.987 ms
	Average fetch longest prefix time: 0.218 ms
	Average load cache time: 5.143 ms

FIXES: b/389788256
TESTED: unittest

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[vipannalla]

Looks good, I left few comments.

[yuyanpeng-google]

Modify the recursive erase to prevent stack recursive limit in long context.
Fix the comments.
Rebase to main.

[mailvijayasingh]

LGTM!

[liurupeng]

Synced with Yuyan offline, I think the current approach looks good

[yuyanpeng-google]

Rebase and squash

[vipannalla]

looks good