keep 4 in RotatingKVCache

5101f03369b3abf49ba9ac76d4ac5e769b80a260 · 2024-08-22 13:14:28 +0100 · Alex Cheema

Files touched

M exo/inference/mlx/sharded_model.py

Diff

commit 5101f03369b3abf49ba9ac76d4ac5e769b80a260
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Thu Aug 22 13:14:28 2024 +0100

    keep 4 in RotatingKVCache
---
 exo/inference/mlx/sharded_model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/exo/inference/mlx/sharded_model.py b/exo/inference/mlx/sharded_model.py
index ff46b5c0..87b85a9d 100644
--- a/exo/inference/mlx/sharded_model.py
+++ b/exo/inference/mlx/sharded_model.py
@@ -75,7 +75,7 @@ class StatefulShardedModel:
 
   def init_cache(self, request_id: str):
     kv_heads = [self.model.n_kv_heads] * len(self.model.layers) if isinstance(self.model.n_kv_heads, int) else self.model.n_kv_heads
-    new_cache = [RotatingKVCache(self.model.head_dim, n, self.max_kv_size) for n in kv_heads]
+    new_cache = [RotatingKVCache(self.model.head_dim, n, self.max_kv_size, keep=4) for n in kv_heads]
 
     if len(self.caches) >= self.max_caches:
       self.caches.popitem(last=False)

← 6db73fab laptop gpu device capabilites · back to Exo · fix mistral nemo 417114fa →