biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Hardware] Initial TPU integration (#5292)

Browse Source
This commit is contained in:
Woosuk Kwon
2024-06-12 11:53:03 -07:00
committed by GitHub
parent 847cdcca1c
commit 1a8bfd92d5
22 changed files with 1322 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
vllm/worker/cache_engine.py
View File

@@ -6,7 +6,8 @@ import torch
from vllm.attention import get_attn_backend
from vllm.config import CacheConfig, ModelConfig, ParallelConfig
from vllm.logger import init_logger
from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, is_pin_memory_available
from vllm.utils import (STR_DTYPE_TO_TORCH_DTYPE, get_dtype_size,
is_pin_memory_available)
logger = init_logger(__name__)
@@ -108,9 +109,5 @@ class CacheEngine:
dtype = model_config.dtype
else:
dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_config.cache_dtype]
dtype_size = _get_dtype_size(dtype)
dtype_size = get_dtype_size(dtype)
return dtype_size * total
def _get_dtype_size(dtype: torch.dtype) -> int:
return torch.tensor([], dtype=dtype).element_size()