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, 12:15 UTC
Update tpu_worker.py 's typo (#17288)
idouba committed
Apr 28, 2025 at 19:01 UTC
72c5b97231f252152fe012661c1980e38ef2a098
1 file changed
+2
-2
vllm/worker/tpu_worker.py
+2
-2
@@ -163,8 +163,8 @@ class TPUWorker(LoRANotSupportedWorkerBase, LocalOrDistributedWorkerBase):
163
usable_memory_size = int(total_memory_size *
164
self.cache_config.gpu_memory_utilization)
165
tpu_kv_cache_bytes = max(usable_memory_size - profiled, 0)
166
- dtype_btyes = get_dtype_size(self.cache_dtype)
167
- block_size_bytes = (dtype_btyes * self.cache_config.block_size *
166
+ dtype_bytes = get_dtype_size(self.cache_dtype)
167
+ block_size_bytes = (dtype_bytes * self.cache_config.block_size *
168
num_layers * 2 * head_size * num_kv_heads)
169
num_tpu_blocks = tpu_kv_cache_bytes // block_size_bytes
170
num_tpu_blocks = (num_tpu_blocks // 8) * 8 # Round down to 8.