zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-29 16:58:45 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix cache.py to support non-top level layers

Browse Source
This commit is contained in:
Shunta Saito 2025-02-13 13:44:31 +09:00
parent 58686bbcac
commit 9a6e6541de
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
llms/mlx_lm/models/cache.py
View File

@ -26,7 +26,10 @@ def make_prompt_cache(
if hasattr(model, "make_cache"): if hasattr(model, "make_cache"):
return model.make_cache() return model.make_cache()
num_layers = len(model.layers) if hasattr(model, "layers"):
num_layers = len(model.layers)
else:
num_layers = len(model.model.layers)
if max_kv_size is not None: if max_kv_size is not None:
return [ return [
RotatingKVCache(max_size=max_kv_size, keep=4) for _ in range(num_layers) RotatingKVCache(max_size=max_kv_size, keep=4) for _ in range(num_layers)