Skip to content

Apply max_kv_size to models that build their own cache - #1995

Open
Lazarus-931 wants to merge 3 commits into
Blaizzy:mainfrom
Lazarus-931:honor-max-kv-size
Open

Apply max_kv_size to models that build their own cache#1995
Lazarus-931 wants to merge 3 commits into
Blaizzy:mainfrom
Lazarus-931:honor-max-kv-size

Conversation

@Lazarus-931

@Lazarus-931 Lazarus-931 commented Aug 21, 2026

Copy link
Copy Markdown
Collaborator

make_prompt_cache deferred entirely to a model's make_cache(), so max_kv_size was silently dropped on 97 of 173 models.

Verification

e2e on cached checkpoints, a 32-token bound over a ~200-token prompt

model before after
Qwen3.5-0.8B identical output, still answers "Tuesday" answers "Monday" — context genuinely lost
LFM2.5-1.2B identical output, still answers "Tuesday" output changes
Qwen2.5-0.5B (no make_cache) already worked unchanged

Across 200 model directories, of the 72 instantiable from a generic config: 44 now apply the bound, 3 were already bounded tighter and are left alone, 1 correctly raises UnboundedKVCache, and none silently ignore it.

`make_prompt_cache` deferred entirely to a model's `make_cache()` when it had
one, so `max_kv_size` was dropped without a word. 97 of 173 model directories
define `make_cache`, and every one of them takes no arguments, so
`--max-kv-size` was a silent no-op on most models. Two models in the same
family disagreed: `qwen3_5` ignored the bound, `qwen3_5_moe` honoured it only
because it happens not to define `make_cache`.

The bound is now applied either way. A `make_cache` that accepts a maximum size
is given it directly -- `diffusion_gemma` already had that signature -- and
otherwise the caches it returns are bounded here: anything that grows with the
sequence is rewritten as a `RotatingKVCache`, and an existing rotating cache is
tightened. Fixed-size state is left exactly as the model built it, so on a
hybrid only the attention layers are bounded and recurrent state is untouched.
`CacheList` and tuple entries are recursed into.

Two cases are deliberately not bounds:
- a cache the model already bounded at or below the request is left alone
- if nothing in the cache grows with the sequence, `UnboundedKVCache` is raised
  rather than the request being ignored

Verified across 200 model directories: the bound now applies to 44 of the 72
that could be instantiated from a generic config, 3 were already bounded
tighter, 1 correctly refuses, and none silently ignore it. End to end on cached
checkpoints, a 32-token bound over a 200-token prompt now changes the answer on
Qwen3.5-0.8B, LFM2.5-1.2B and Qwen2.5-0.5B; before this change the first two
answered as though unbounded.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@Lazarus-931
Lazarus-931 marked this pull request as ready for review August 21, 2026 14:17
Signed-off-by: Alazer Manakelew <alazermanakelewb@gmail.com>
Comment thread mlx_vlm/models/cache.py
Comment thread mlx_vlm/models/cache.py

@Blaizzy Blaizzy Aug 21, 2026

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks Alazar, this indeed addresses a real issue!

However, these changes look like a hack, basically too much policy in one helper: signature introspection, recursive cache rewriting, status bookkeeping, and a new exception.

We could solve this in a much simpler way.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

would a

-    def make_cache(self):
-        return self.language_model.make_cache()
+    def make_cache(self, max_kv_size=None):
+        return self.language_model.make_cache(max_kv_size)

fix across all model definitions make sense?

@Blaizzy Blaizzy Aug 22, 2026

Copy link
Copy Markdown
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why? When make_cache makes no use of max_kv_size?

qwen3/language.py

def make_cache(self):
   from ..cache import KVCache

   return [KVCache() for _ in self.layers]

qwen3_5/language.py

def make_cache(self):
    return [ArraysCache(size=2) if l.is_linear else KVCache() for l in self.layers]

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants