xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-06 01:25:42 +08:00
Code Issues Packages Projects Releases Wiki Activity

Fix return type error

Browse Source
This commit is contained in:
Woosuk Kwon 2023-02-16 01:33:03 +00:00
parent 8edcabc737
commit 9e68a6827e
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
cacheflow/worker/cache_engine.py
View File

@ -42,7 +42,7 @@ class CacheEngine:
# Initialize the events for stream synchronization. # Initialize the events for stream synchronization.
self.events = [torch.cuda.Event() for _ in range(self.num_layers)] self.events = [torch.cuda.Event() for _ in range(self.num_layers)]
def get_key_block_shape(self) -> Tuple[int, int, int, int, int]: def get_key_block_shape(self) -> Tuple[int, int, int, int]:
element_size = torch.tensor([], dtype=self.dtype).element_size() element_size = torch.tensor([], dtype=self.dtype).element_size()
x = 16 // element_size x = 16 // element_size
return ( return (
@ -52,7 +52,7 @@ class CacheEngine:
x, x,
) )
def get_value_block_shape(self) -> Tuple[int, int, int, int]: def get_value_block_shape(self) -> Tuple[int, int, int]:
return ( return (
self.num_heads, self.num_heads,
self.block_size, self.block_size,