use KVCache instead of RotatingKVCache

661035f0e0192a01fbf323ca35b7d8ae390e607a · 2024-11-03 16:56:35 -0500 · Alex Cheema

Files touched

M exo/inference/mlx/sharded_model.py

Diff

commit 661035f0e0192a01fbf323ca35b7d8ae390e607a
Author: Alex Cheema <alexcheema123@gmail.com>
Date:   Sun Nov 3 16:56:35 2024 -0500

    use KVCache instead of RotatingKVCache
---
 exo/inference/mlx/sharded_model.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/exo/inference/mlx/sharded_model.py b/exo/inference/mlx/sharded_model.py
index 37dfbaf2..9c9c8ffa 100644
--- a/exo/inference/mlx/sharded_model.py
+++ b/exo/inference/mlx/sharded_model.py
@@ -77,7 +77,8 @@ 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)
     if self.max_kv_size is not None:
-      cache = [RotatingKVCache(self.model.head_dim, n, max_size=self.max_kv_size, keep=4) for n in kv_heads]
+      # cache = [RotatingKVCache(self.model.head_dim, n, max_size=self.max_kv_size, keep=4) for n in kv_heads]
+      cache = [KVCache(self.model.head_dim, n) for n in kv_heads]
     else:
       cache = [KVCache(self.model.head_dim, n) for n in kv_heads]

← d0b7154f set fallback llama model to use smallest one also · back to Exo · update mlx to 0.19.3, mlx-lm to 0.19.2 bc7acfd3 →