Skip to content

GitLab

Unverified Commit 72c5b972 authored by idouba's avatar idouba Committed by GitHub
Browse files

Update tpu_worker.py 's typo (#17288)

parent fa93cd9f
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
vllm/worker/tpu_worker.py
View file @ 72c5b972
...@@ -163,8 +163,8 @@ class TPUWorker(LoRANotSupportedWorkerBase, LocalOrDistributedWorkerBase): ...@@ -163,8 +163,8 @@ class TPUWorker(LoRANotSupportedWorkerBase, LocalOrDistributedWorkerBase):
usable_memory_size = int(total_memory_size * usable_memory_size = int(total_memory_size *
self.cache_config.gpu_memory_utilization) self.cache_config.gpu_memory_utilization)
tpu_kv_cache_bytes = max(usable_memory_size - profiled, 0) tpu_kv_cache_bytes = max(usable_memory_size - profiled, 0)
dtype_btyes = get_dtype_size(self.cache_dtype) dtype_bytes = get_dtype_size(self.cache_dtype)
block_size_bytes = (dtype_btyes * self.cache_config.block_size * block_size_bytes = (dtype_bytes * self.cache_config.block_size *
num_layers * 2 * head_size * num_kv_heads) num_layers * 2 * head_size * num_kv_heads)
num_tpu_blocks = tpu_kv_cache_bytes // block_size_bytes num_tpu_blocks = tpu_kv_cache_bytes // block_size_bytes
num_tpu_blocks = (num_tpu_blocks // 8) * 8 # Round down to 8. num_tpu_blocks = (num_tpu_blocks // 8) * 8 # Round down to 8.
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment