[blindTissue]

PR description

This PR adds Afterglow (paper), a training-free plug-in for the scoring step of query-aware presses. It follows up on #292 (QuerySpanPress), which was closed.

Query-aware scorers such as SnapKV give the highest scores to the first tokens of the span which the question asks about (a value, an identifier, a title). Eviction thus then retains many of the starting span cache but evicts most of the remaining span, leading to truncated or corrupted answer. Afterglow lets each token inherit the decayed scores of the tokens before it, separately for each layer and KV head:

s'(t) = max( s(t), max_{n_sink <= j < t} gamma^(t - j) * s(j) )      gamma = 0.95, n_sink = 4

The wrapped press then selects from s' with the same budget.

New presses:

• AfterglowPress(press) (afterglow_press.py): wraps any ScorerPress and composes with AdaKVPress and CriticalAdaKVPress, which then allocate their budgets from the propagated scores. CriticalKVPress sets the scores it protects to the largest float, so Afterglow goes inside it: CriticalAdaKVPress(AfterglowPress(SnapKVPress())). Afterglow works in log space and raises an error for negative scores (e.g. KnormPress).
• QueryAwareKVzipPlusPress (query_aware_kvzip_plus_press.py): the best configuration of the paper. It inherits from KVzipPress and replaces the reconstruction pass with the attention of the question (the last window_size tokens). The KVzip+ scores then go through Afterglow, with the inherited part scaled by a predictability gate u(t)^alpha, u(t) = 1 - p(x_t | x_<t) from the prefill's own final hidden states, so that predictable continuations inherit less. The budget allocation across layers and heads (compress_post), the ||W_O v|| normalization and the causal mask are KVzip's.
• PRESS_REGISTRY entries: snapkv_afterglow, adakv_snapkv_afterglow, critical_adakv_snapkv_afterglow, kvzip_plus_query_aware. All of them are meant for query-aware compression (query_aware=True).

Results

Each cell is without -> with Afterglow, from paper's runs.

Model	Press	50%	75%	87.5%	93.75%
Qwen3-8B	SnapKV	87.6 -> 95.1	74.2 -> 93.2	55.3 -> 83.4	23.5 -> 63.5
Qwen3-8B	Ada-KV	89.8 -> 95.2	78.5 -> 94.4	57.3 -> 85.9	26.7 -> 70.0
Qwen3-8B	CriticalAdaKV	95.0 -> 95.1	86.4 -> 94.8	67.6 -> 90.9	35.0 -> 74.6
Qwen3-8B	KVzip+ (query-aware)	93.1 -> 95.2	81.0 -> 94.8	62.6 -> 92.4	31.0 -> 82.4
Llama-3.1-8B	SnapKV	88.9 -> 95.3	78.6 -> 93.6	68.0 -> 87.7	49.2 -> 75.7
Llama-3.1-8B	Ada-KV	91.2 -> 95.5	82.6 -> 94.7	71.3 -> 90.3	49.1 -> 79.9
Llama-3.1-8B	CriticalAdaKV	94.5 -> 95.5	83.7 -> 95.1	74.6 -> 91.7	56.1 -> 82.0
Llama-3.1-8B	KVzip+ (query-aware)	94.8 -> 95.6	87.8 -> 95.2	74.0 -> 92.8	56.5 -> 85.9

kvzip_plus_query_aware, (Afterglow + predictability gate), run with the PR's code through evaluation/evaluate.py. Within 0.15pp of the paper at every ratio.

Model	50%	75%	87.5%	93.75%
Qwen3-8B	95.20	94.95	93.57	88.13
Llama-3.1-8B	95.64	95.38	94.38	89.00

Tests

tests/presses/test_afterglow_press.py checks afterglow() against a naive double loop, AfterglowPress alone and inside AdaKVPress / CriticalAdaKVPress, and the error for negative scores. QueryAwareKVzipPlusPress is in default_presses.

Checklist

Before submitting a PR, please make sure:

• Tests are working (make test)

• Code is formatted correctly (make style, on errors try fix with make format)

• Copyright header is included

• All commits are signed-off using git commit -s

• (new press) mypress_press.py is in the presses directory

• (new press) MyPress is in __init__.py

• (new press) README.md is updated with a 1 liner about the new press in the Available presses section

• (new press) New press is in the default_presses list in tests/default_presses.py

• (new press) A docstring is provided that follows the same structure as the existing ones

[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.