zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-11-02 20:48:11 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add the ability to load the KV cache from a file (#956)

Browse Source
This commit is contained in:
Angelos Katharopoulos
2024-08-28 22:11:45 -07:00
committed by GitHub
parent 7f8c961287
commit 1003a8b2dd
5 changed files with 250 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
llms/mlx_lm/models/base.py
View File

@@ -46,6 +46,7 @@ class KVCache:
self.values[..., prev : self.offset, :] = values
return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
@property
def state(self):
return self.keys, self.values
@@ -137,6 +138,7 @@ class RotatingKVCache:
return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
return self.keys, self.values
@property
def state(self):
return self.keys, self.values