lmstudio-ai / lmstudio-ai/mlx-engine

MLX `RotatingKVCache` trim behavior causes context overflow policies to always erase the whole cache

Open
#177 4 comments 2 reactions 0 assignees View on GitHub

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.