Clear cache every now and then

[awni]

As we step the KV cache.. the buffer cache fills up which causes the machine to use more RAM than is really needed.

For example:

mlx_lm.generate --model mlx-community/Meta-Llama-3.1-8B-Instruct-4bit -m 2048 --prompt - < prompt.txt

Pre:

Prompt: 32188 tokens, 430.339 tokens-per-sec
Generation: 892 tokens, 32.480 tokens-per-sec
Peak memory: 11.496 GB
Cache memory: 22.795 GB

Post:

Prompt: 32188 tokens, 424.646 tokens-per-sec
Generation: 892 tokens, 32.211 tokens-per-sec
Peak memory: 11.496 GB
Cache memory: 4.034 GB

No difference on M2 Ultra with:

mlx_lm.generate --model mlx-community/Mistral-7B-Instruct-v0.3-4bit --prompt "Write a story about Einstein"  --temp 0.0 --max-tokens 512

Still hits about 120.5 toks/sec

[angeloskath]

Nice :-)