fix: clamp n_kept to >=1 to prevent silent cache emptying on short contexts 🤖🤖🤖 - #264
Open
SuperMarioYL wants to merge 1 commit into
Open
fix: clamp n_kept to >=1 to prevent silent cache emptying on short contexts 🤖🤖🤖#264SuperMarioYL wants to merge 1 commit into
SuperMarioYL wants to merge 1 commit into
Conversation
…ntexts 🤖🤖🤖 Signed-off-by: supermario_leo <leo.stack@outlook.com>
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.
PR description
ScorerPress.compress(and the wrapper presses that re-derive their ownn_kept) compute thenumber of KV pairs to keep as
int(k_len * (1 - compression_ratio))with no floor guard. On ashort context (
k_lenof 1 or 2) with any non-zerocompression_ratio, this floors ton_kept = 0, soscores.topk(0)returns empty indices,keys.gather(2, …)writes an empty(bsz, heads, 0, head_dim)cache, and the following decode step attends zero keys — producingNaN/ divergent output with no error raised.This is the same "silent eviction" family as #227 (which fixed the pad-value
+1), but on then_keptfloor axis:chunk_press.py:77andfinch_press.py:107already guard withmax(1, int(...)), so the maintainers clearly intend the floor — it was simply never backportedto the base
ScorerPress.compressnor to the wrappers that re-deriven_kept. Notablyfinch_press.pyguards the chunk path (:107) but not the non-chunk path (:100), eventhough both were added in the same PR #139 — an oversight, not a design choice.
This PR applies the existing
max(1, int(...))idiom in-place at every unguarded site:kvpress/presses/scorer_press.py:94ScorerPress.compress(propagates to allScorerPresssubclasses)kvpress/presses/block_press.py:66BlockPressre-derivesn_keptkvpress/presses/merging_press.py:86MergingPress(docstring: "Identical to ScorerPress.compress except…")kvpress/presses/key_rerotation_press.py:146KeyRerotationPress(q_lenform)kvpress/presses/finch_press.py:100FinchPressnon-chunk path (the chunk path:107is already guarded)kvpress/presses/adakv_press.py:64AdaKVPress(# ScorerPress definitioncopy)kvpress/presses/criticalkv_press.py:149CriticalAdaKVPress(# ScorerPress definitioncopy)The change is purely additive:
max(1, X)is identical toXwheneverX >= 1, so for anycontext long enough that
int(k_len * (1 - ratio)) >= 1the behaviour is unchanged. It only everlifts
n_keptfrom0to1on degenerate short inputs, preventing the silent empty cache.Two related sites are deliberately excluded from this PR to keep it a clean, reviewable
consistency fix of the
ScorerPressidiom:decoding_press.pycarries the same pattern but the decoding subsystem has active in-flightbranches; it is left untouched here to avoid collisions and can be handled in a follow-up.
kvcompose_press.pyuses acomposite_scores.numel()form rather thank_len; noted for afollow-up rather than mixed into this idiom-consistency sweep.
Refs #227.
Checklist
Before submitting a PR, please make sure:
tests/presses/test_scorer_press.py(parametrized, CPU-only,no model download);
make testruns in CI.make style: flake8 clean on all touched files).git commit -s(DCO).(n/a — not a new press: no
__init__/README/default_presses/docstring changes.)🤖🤖🤖