Skip to content

fix: clamp n_kept to >=1 to prevent silent cache emptying on short contexts 🤖🤖🤖 - #264

Open
SuperMarioYL wants to merge 1 commit into
NVIDIA:mainfrom
SuperMarioYL:fix/n-kept-floor-guard
Open

fix: clamp n_kept to >=1 to prevent silent cache emptying on short contexts 🤖🤖🤖#264
SuperMarioYL wants to merge 1 commit into
NVIDIA:mainfrom
SuperMarioYL:fix/n-kept-floor-guard

Conversation

@SuperMarioYL

Copy link
Copy Markdown

PR description

ScorerPress.compress (and the wrapper presses that re-derive their own n_kept) compute the
number of KV pairs to keep as int(k_len * (1 - compression_ratio)) with no floor guard. On a
short context (k_len of 1 or 2) with any non-zero compression_ratio, this floors to
n_kept = 0, so scores.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 — producing
NaN / divergent output with no error raised.

This is the same "silent eviction" family as #227 (which fixed the pad-value +1), but on the
n_kept floor axis: chunk_press.py:77 and finch_press.py:107 already guard with
max(1, int(...)), so the maintainers clearly intend the floor — it was simply never backported
to the base ScorerPress.compress nor to the wrappers that re-derive n_kept. Notably
finch_press.py guards the chunk path (:107) but not the non-chunk path (:100), even
though 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:

file:line scope
kvpress/presses/scorer_press.py:94 base ScorerPress.compress (propagates to all ScorerPress subclasses)
kvpress/presses/block_press.py:66 BlockPress re-derives n_kept
kvpress/presses/merging_press.py:86 MergingPress (docstring: "Identical to ScorerPress.compress except…")
kvpress/presses/key_rerotation_press.py:146 KeyRerotationPress (q_len form)
kvpress/presses/finch_press.py:100 FinchPress non-chunk path (the chunk path :107 is already guarded)
kvpress/presses/adakv_press.py:64 AdaKVPress (# ScorerPress definition copy)
kvpress/presses/criticalkv_press.py:149 CriticalAdaKVPress (# ScorerPress definition copy)

The change is purely additive: max(1, X) is identical to X whenever X >= 1, so for any
context long enough that int(k_len * (1 - ratio)) >= 1 the behaviour is unchanged. It only ever
lifts n_kept from 0 to 1 on 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 ScorerPress idiom:

  • decoding_press.py carries the same pattern but the decoding subsystem has active in-flight
    branches; it is left untouched here to avoid collisions and can be handled in a follow-up.
  • kvcompose_press.py uses a composite_scores.numel() form rather than k_len; noted for a
    follow-up rather than mixed into this idiom-consistency sweep.

Refs #227.

Checklist

Before submitting a PR, please make sure:

  • Tests are working — added tests/presses/test_scorer_press.py (parametrized, CPU-only,
    no model download); make test runs in CI.
  • Code is formatted correctly (make style: flake8 clean on all touched files).
  • Copyright header is included (SPDX header on the new test file).
  • All commits are signed-off using git commit -s (DCO).

(n/a — not a new press: no __init__/README/default_presses/docstring changes.)

🤖🤖🤖

…ntexts 🤖🤖🤖

Signed-off-by: supermario_leo <leo.stack@outlook.com>
@copy-pr-bot

copy-pr-bot Bot commented Aug 14, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

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.

1 participant