Apply max_kv_size to models that build their own cache - #1995
Open
Lazarus-931 wants to merge 3 commits into
Open
Apply max_kv_size to models that build their own cache#1995Lazarus-931 wants to merge 3 commits into
Lazarus-931 wants to merge 3 commits into
Conversation
`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
marked this pull request as ready for review
August 21, 2026 14:17
lucasnewman
approved these changes
Aug 21, 2026
Signed-off-by: Alazer Manakelew <alazermanakelewb@gmail.com>
lucasnewman
reviewed
Aug 21, 2026
lucasnewman
approved these changes
Aug 21, 2026
Blaizzy
reviewed
Aug 21, 2026
Owner
There was a problem hiding this comment.
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.
Collaborator
Author
There was a problem hiding this comment.
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?
Owner
There was a problem hiding this comment.
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]
Blaizzy
requested changes
Aug 21, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
make_prompt_cachedeferred entirely to a model'smake_cache(), somax_kv_sizewas silently dropped on 97 of 173 models.Verification
e2e on cached checkpoints, a 32-token bound over a ~200-token prompt
make_cache)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.