Fix rotating kv cache size

[angeloskath]

Currently the rotating KV cache can be both max_size and max_size - 1 depending on whether we filled it by generating or by prompt token processing. This PR makes sure the max is always max_size and not max_size - 1.

The following two simple repros exemplify the problem a bit:

Generate past max size

from mlx_lm.models.cache import make_prompt_cache
from mlx_lm.utils import load
from mlx_lm.utils import stream_generate

PROMPT = "Tell me a story"

model, tokenizer = load("mlx-community/Llama-3.2-3B-Instruct-4bit")

cache = make_prompt_cache(model, max_kv_size=100)
for text in stream_generate(model, tokenizer, PROMPT, prompt_cache=cache, max_tokens=100):
    pass
for text in stream_generate(model, tokenizer, PROMPT, prompt_cache=cache, max_tokens=100):
    pass

Prompt process past max size

from mlx_lm.models.cache import make_prompt_cache
from mlx_lm.utils import load
from mlx_lm.utils import stream_generate

PROMPT = "Tell me a story "

model, tokenizer = load("mlx-community/Llama-3.2-3B-Instruct-4bit")

cache = make_prompt_cache(model, max_kv_size=512)
for text in stream_generate(model, tokenizer, PROMPT * 384, prompt_cache=cache, max_tokens=50):
    pass

The first breaks on main. Then changing line 45 on base.py breaks the 2nd test and changing the trim size makes sure everything works. Btw these also occur in mlx_lm.chat and mlx_lm.cache_prompt it is just more explicit with the tests above.

[awni]

Looks good thanks for digging into that and fixing it!