lmstudio-ai / lmstudio-ai/mlx-engine
MLX `RotatingKVCache` trim behavior causes context overflow policies to always erase the whole cache
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.2k
- Forks
- 133
- Avg merge
- 21h 6m
- Merged PRs (30d)
- 1
Description
We defer KV cache creation to `mlx_lm.models.cache.make_prompt_cache`, which itself defers to the model for implementation. Most models (with some exceptions, e.g. the DeepSeeks) don't implement their own KV caches, so MLX defaults to a `mlx_lm.models.cache.RotatingKVCache` that uses a circular buffer to manage an arbitrary length generation. This avoids complexities introduced by manually shifting around a fixed-size linear cache (cf. llama.cpp), but introduces some of its own problems: namely, if we have generated `n` tokens where `n > max_kv_size`, the cache will no longer let us `trim` from the end of it.
This is a problem because we want to trim from the cache **exactly when** we have generated more than `max_kv_size` tokens, since that's the whole point of a **context overflow** policy! What occurs in practice is that `mlx_engine.cache_wrapper._get_unprocessed_tokens` attempts to trim the cache in accordance with the context overflow policy, but then fails because the cache rejects the trim request on account of being over capacity. The trim then falls back to erasing the entire cache, which requires usually thousands to tens of thousands of tokens to be unnecessarily recomputed at great expense. This is not ideal.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reading mlx_lm.models.cache.make_prompt_cache and RotatingKVCache, then follow mlx_engine.cache_wrapper._get_unprocessed_tokens where context-overflow trimming is attempted. Reproduce generation beyond max_kv_size and verify that trimming preserves the reusable portion of the cache instead of falling back to erasing the entire cache.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 38/100