[Kernel] optimize kv cache update kernel block size

[yaochengji]

Description

• Automatically adjust kv cache update kernel's num_slices_per_block, porting from [TPU] Optimize kv cache update kernel vllm#20415 and [TPU] fix kv_cache_update kernel block size choosing logic vllm#21007
• change RPA kernel's vmem_limit_bytes to 100MB to avoid vmem OOM issue

Tests

1. pytest -s -v tests/kernels/ragged_kv_cache_update_test.py
2. server: TPU_BACKEND_TYPE=torchax VLLM_TORCHAX_ENABLED=1 vllm serve meta-llama/Llama-3.1-8B-Instruct --disable-log-requests --gpu-memory-utilization 0.98 --max-num-batched-tokens 2048 --max-num-seqs 128 --max-model-len 2048 --no-enable-prefix-caching --tensor_parallel_size=1
   client: python3 ./benchmarks/benchmark_serving.py --model meta-llama/Llama-3.1-8B-Instruct --dataset-name sonnet --dataset-path benchmarks/sonnet_4x.txt --sonnet-input-len 1800 --sonnet-output-len 128 --ignore_eos

We can observe the throughput increases from 8.11 req/s to 8.14 req/s.

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.

[yaochengji]

cc @lsy323 @xiangxu-google

[xiangxu-google]

Thanks!

[bythew3i]

Nit: better to test the throughput before and after this change. Kernel won't be affected but next op's prefetch will be affected.

[yaochengji]

It would crash due to vmem OOM before the change.

[bythew3i]

How is this calculated?

[yaochengji]

Basically I'd like to have 32MB vmem to be used for scratch buffer, and it's round up to 40MB.

[bythew3i]

I would suggest - have an autotune and benchmarking in google workspace internally first. Like we do for RPA and quantized_matmul

[yaochengji]

Thanks for the suggestion! Since I'd like to make the kv cache update kernel on par with vLLM torch/xla path, we can add the auto-tuning and benchmarking later in google internally later?

[bythew3i]

LGTM - if the current change blocks you