fix: complete the n_kept floor sweep deferred in #264

[SuperMarioYL]

PR description

#264 introduced compute_n_kept (kvpress/utils.py) — 0 by design at compression_ratio >= 1.0, else max(1, int(...)) — and applied it to the k_len-form budget sites, but deferred the remaining budget derivations. Three unguarded sites survive on main, each silently misbehaving on short contexts with no error raised:

file:line	budget form	degenerate failure (probed on MaxJeblick/llama2-0b-unit-test, ctx=8)
kvpress/presses/kvcompose_press.py:277,284	int(numel * (1 - compression_ratio)) (numel form)	floors to 0 → topk(0) silently empties the whole cache (structured form, ratio 0.95: cache.get_seq_length()=0, keys=(1,2,0,6)); even keep_token_lower_bound-boosted tokens (+1e9) are dropped by topk(0)
kvpress/presses/criticalkv_press.py:83	int((1 - ratio) * k_len * first_stage_ratio)	floors to 0 → topk(scores, 0) is legal → stage-1 protection silently skipped (0 positions boosted to finfo.max)
kvpress/presses/fastkvzip_press.py:256	window_size = int(ctx_len * window_ratio)	floors to 0 → scores[:, :, -0:] = 1.0 sets all scores to 1.0 (the -0: slice is the whole tensor) → full-protection inversion on any context shorter than 1/window_ratio (50 tokens at the default 0.02)

Additionally, on the kvcompose structured path the per-layer distribution of the global budget can still leave a layer at 0 even after the global floor (probed: per-layer [1, 0] at ratio 0.99 — one layer's cache empties while compression_ratio < 1.0).

This PR completes the sweep:

• kvcompose: both numel forms route through compute_n_kept; the structured per-layer counts are clamped to max(1, keep_token_lower_bound) so no layer empties while compression_ratio < 1.0 (per-head zeros are left alone — head pruning is a legitimate outcome there).
• criticalkv: selection_budget gets a floor only when first_stage_ratio > 0 and compression_ratio < 1.0 (mirroring the n_safe guard shape from the fix: clamp n_kept to >=1 to prevent silent cache emptying on short contexts 🤖🤖🤖 #264 review; first_stage_ratio = 0 means stage 1 is disabled by design and must stay at 0).
• fastkvzip: window_size = max(1, int(ctx_len * self.window_ratio)).

The change is purely additive: every floor only activates when the derived budget is 0. Verified byte-identical behaviour at normal ratios (kvcompose/criticalkv/fastkvzip at 0.5 and 0.8, ctx=128, fixed seed) — the clamp is inactive whenever the budget covers the layers. compression_ratio >= 1.0 still means "evict everything" everywhere (compute_n_kept contract). The guarded sites from #264 and the hand-rolled chunk/pyramidkv floors are untouched.

Tests

tests/presses/test_budget_floor_sweep.py (10 cases, all on the real presses via the unit-test model, no mocks of the fixed code):

• kvcompose structured form never empties a layer at ratio 0.95 / 0.99 (every layer keeps >= 1 token), including bsz=2
• keep_token_lower_bound tokens survive at ratio 0.95 / 0.99 (per-layer >= lower bound)
• unstructured form keeps >= 1 token globally at ratio 0.98
• criticalkv boosts >= 1 position when first_stage_ratio > 0 (ratio 0.95 / 0.99), and stays at 0 when first_stage_ratio = 0 (design behavior, green on main too)
• fastkvzip short-context scores are not all 1.0

The 9 degenerate cases are red on main (emptying / skipping / flattening) and green with this change. Full tests/presses/ passes (529 passed; the one test_ea_with_stats.py::test_load_stats failure is a pre-existing gated-repo download error, identical on main).

Checklist

• Tests are working (make test) — tests/presses/ 529 passed / 4 skipped; the single test_ea_with_stats.py::test_load_stats failure is a pre-existing gated-repo download error that fails identically on main (verified in a clean checkout)
• Code is formatted correctly (make style) — flake8 clean, mypy clean (83 files); note black/isort report pre-existing drift on kvcompose_press.py that is identical on main (verified) and untouched by this PR
• Copyright header is included (SPDX on the new test file)
• All commits are signed-off using git commit -s

[copy-pr-bot]

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.

[alessiodevoto]

LGTM, thanks !