A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 01:00 UTC
[TPU] fix kv_cache_update kernel block size choosing logic (#21007)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Jul 15, 2025 at 21:39 UTC
85431bd9ad1646fd8741677da5b57b33f6ea6c64
2 files changed
+51
-3
vllm/v1/attention/backends/pallas.py
+48
-1
@@ -326,7 +326,54 @@ def kv_cache_update_op_non_xla(kv: torch.Tensor, slot_mapping: torch.Tensor,
326
return kv_cache
327
328
329
+# We can move this function to a common utils file if it's also useful for other
330
+# hardware.
331
+def dtype_bits(dtype: torch.dtype):
332
+ if dtype.is_floating_point:
333
+ try:
334
+ return torch.finfo(dtype).bits
335
+ except TypeError:
336
+ pass
337
+ elif dtype.is_complex:
338
+ if dtype is torch.complex32:
339
+ return 32
340
+ elif dtype is torch.complex64:
341
+ return 64
342
+ elif dtype is torch.complex128:
343
+ return 128
344
+ else:
345
+ try:
346
+ return torch.iinfo(dtype).bits
347
+ # torch.iinfo cannot support int4, int2, bits8...
348
+ except TypeError:
349
+ pass
350
+ str_dtype = str(dtype)
351
+ # support torch.int4, torch.int5, torch.uint5...
352
+ if str_dtype.startswith("torch.int") or str_dtype.startswith("torch.uint"):
353
+ return int(str_dtype[-1])
354
+ raise TypeError(f"Getting the bit width of {dtype} is not supported")
355
+
356
+
357
+def get_dtype_packing(dtype):
358
+ bits = dtype_bits(dtype)
359
+ if 32 % bits != 0:
360
+ raise ValueError(
361
+ f"The bit width must be divisible by 32, but got bits={bits}, "
362
+ "dtype={dtype}")
363
+ return 32 // bits
364
+
365
+
366
def get_page_size_bytes(block_size: int, num_kv_heads: int, head_size: int,
367
kv_cache_dtype: torch.dtype) -> int:
368
"""Returns the size in bytes of one page of the KV cache."""
332
- return block_size * num_kv_heads * head_size * kv_cache_dtype.itemsize
369
+ padded_head_size = cdiv(head_size,
370
+ TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
371
+ num_combined_kv_heads = num_kv_heads * 2
372
+
373
+ # NOTE: for the implicit padding in XLA
374
+ packing = get_dtype_packing(kv_cache_dtype)
375
+ num_combined_kv_heads = cdiv(num_combined_kv_heads, packing) * packing
376
+
377
+ kv_cache_dtype_bits = dtype_bits(kv_cache_dtype)
378
+ return (block_size * num_combined_kv_heads * padded_head_size *
379
+ kv_cache_dtype_bits // 8)
vllm/v1/worker/tpu_model_runner.py
+3
-2
@@ -1863,8 +1863,9 @@ def _get_num_slices_per_kv_cache_update_block(page_size_bytes: int) -> int:
1863
out of scalar registers. Thus this function will limit the number of
1864
slices to 64.
1865
"""
1866
- # Conservative VMEM usage limit: 32 MiB
1867
- vmem_limit = 32 * 1024 * 1024
1866
+ # The default vmem_limit_bytes of a pallas kernel is 32MB. Here we
1867
+ # calculate num_slices_per_block based on 16MB in case any register spills.
1868
+ vmem_limit = 16 * 1024 * 1024
1869
num_slices_per_block = vmem_limit // page_size_bytes
1870
assert num_slices_per_block > 0, "Number of slices should be positive"
1871
num_slices_per_block = prev_power_of_2(num_slices_per_block)