[TPU] fix kv cache update kernel block size

Signed-off-by: Chengji Yao <chengjiyao@google.com>

Author: yaochengji

vllm/v1/attention/backends/pallas.py

@@ -329,4 +329,6 @@ def kv_cache_update_op_non_xla(kv: torch.Tensor, slot_mapping: torch.Tensor,
 def get_page_size_bytes(block_size: int, num_kv_heads: int, head_size: int,
                         kv_cache_dtype: torch.dtype) -> int:
     """Returns the size in bytes of one page of the KV cache."""
-    return block_size * num_kv_heads * head_size * kv_cache_dtype.itemsize
+    padded_head_size = cdiv(head_size,
+                            TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
+    return block_size * num_kv_heads * padded_head_size * kv_cache_dtype.itemsize

vllm/v1/worker/tpu_model_runner.py

@@ -1863,8 +1863,9 @@ def _get_num_slices_per_kv_cache_update_block(page_size_bytes: int) -> int:
     out of scalar registers. Thus this function will limit the number of
     slices to 64.
     """
-    # Conservative VMEM usage limit: 32 MiB
-    vmem_limit = 32 * 1024 * 1024
+    # The default vmem_limit_bytes of a pallas kernel is 32MB. Here we
+    # calculate num_slices_per_block based on 16MB in case any register spills.
+    vmem_limit = 16 * 1024 * 1024
     num_slices_per_block = vmem_limit // page_size_bytes
     assert num_slices_per_block > 0, "Number of slices should be positive"
     num_slices_per_block = prev_power_of_2(num_slices_per_block)