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, 04:45 UTC
[TPU][Bugfix] fix kv cache padding (#20048)
Signed-off-by: Chengji Yao <chengjiyao@google.com>
Chengji Yao committed
Jun 25, 2025 at 14:24 UTC
2cc206997012057152f194c0f25e19e3ab3297ea
2 files changed
+14
-9
vllm/v1/attention/backends/pallas.py
+1
-7
@@ -48,13 +48,7 @@ class PallasAttentionBackend(AttentionBackend):
48
) -> tuple[int, ...]:
49
padded_head_size = cdiv(
50
head_size, TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
51
- num_blocks = num_blocks * head_size // padded_head_size
52
- if padded_head_size != head_size:
53
- logger.warning_once(
54
- "head size is padded to %d, and num_blocks is adjusted to %d"
55
- " accordingly", padded_head_size, num_blocks)
56
- head_size = padded_head_size
57
- return (num_blocks, block_size, num_kv_heads * 2, head_size)
51
+ return (num_blocks, block_size, num_kv_heads * 2, padded_head_size)
52
53
@staticmethod
54
def swap_blocks(
vllm/v1/worker/tpu_worker.py
+13
-2
@@ -18,7 +18,8 @@ from vllm.distributed import (ensure_model_parallel_initialized,
18
from vllm.logger import init_logger
19
from vllm.lora.request import LoRARequest
20
from vllm.model_executor import set_random_seed
21
-from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE
21
+from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
22
+from vllm.v1.attention.backends.pallas import TPU_HEAD_SIZE_ALIGNMENT
23
from vllm.v1.core.sched.output import SchedulerOutput
24
from vllm.v1.kv_cache_interface import (AttentionSpec, KVCacheConfig,
25
KVCacheSpec)
@@ -221,7 +222,17 @@ class TPUWorker:
222
usable_memory_size = int(total_memory_size *
223
self.cache_config.gpu_memory_utilization)
224
tpu_kv_cache_bytes = max(usable_memory_size - profiled, 0)
224
-
225
+ head_size = self.model_config.get_head_size()
226
+ if head_size > 0:
227
+ padded_head_size = cdiv(
228
+ head_size, TPU_HEAD_SIZE_ALIGNMENT) * TPU_HEAD_SIZE_ALIGNMENT
229
+ if padded_head_size != head_size:
230
+ logger.warning_once("head size is padded to %d",
231
+ padded_head_size)
232
+ # We adjust the usable memory size for the KV cache to prevent OOM
233
+ # errors, even after padding the head_size.
234
+ tpu_kv_cache_bytes = (tpu_kv_cache_bytes * head_size //
235
+ padded_head_size)
236
return int(tpu_kv_cache_bytes)
237
238
def execute_model(