Sparse-inference accuracy gap: position compaction (A1) + eviction sweep (A2)

[mmjerge]

Investigating the sparse-inference accuracy gap (A1 + A2)

Follow-up to the RL work in awslabs#142. Matthias' standing concern: compared to RL/inference with dense attention, sparse (evicting) attention "is not accurate enough ... that needs to be dealt with now." This PR is the investigation into why, and what actually helps.

The measured gap (base Qwen2.5-0.5B, HELMET @8k, h2o@4096, n=100): ~8pp EM on trivia_qa (0.62 dense -> 0.54), similar on nq. Two hypotheses tested.

A1 -- position compaction (negative)

After eviction, retained tokens keep the RoPE rotations of their original absolute positions, so attention sees a "holey" layout (0, 3, 47, 812, ...) never seen in pretraining. keys_values/kvcache/pos_compact.py re-expresses queries and cached keys at compacted (rank-among-survivors) positions via an exact delta-rotation (RoPE composes additively; relative order preserved so the causal mask is unchanged). Toggle with set_position_compaction(model, True) or grpo_helmet_crosseval.py --compact-positions.

Result: no recovery (1-2pp, within noise, both datasets/metrics). The holey layout is not the main cause -- the information in the evicted tokens is simply gone. Details: docs/POSITION_COMPACTION_A1.md.

A2 -- eviction-policy sweep (one actionable finding)

examples/eviction_sweep.py sweeps retention knobs at a fixed slot budget: quantized vs fp buffers, normalize_scores, keep_initial_fraction (sink), large grace, v-length scoring, lastrec control; plus a budget frontier and pure-quantization controls.

• No retention knob closes the gap at fixed slots (all within 1-2pp).
• 8-bit KV quantization is accuracy-free (h2o_q8 >= h2o_fp), so at equal memory it buys 2x the slots.
• Budget frontier: h2o_q8@8192 matches dense on both datasets (trivia 0.64/0.62, nq 0.26/0.25) at ~half the KV memory of bf16 dense.
• Controls (after internal skepticism): unquantized H2O@8192 reproduces dense bf16 exactly (no eviction when the cache covers the prompt); dense-q8 reproduces h2o-q8 exactly (the +2pp is 8-bit rounding perturbing greedy decode, i.e. noise, not a gain). Clean decomposition: retention determines accuracy; quantization is a +/-2pp numeric effect; H2O scoring only matters under active eviction.

Details + tables: docs/EVICTION_SWEEP_A2.md.

Takeaway

Size the cache to cover the prompt and quantize, rather than evicting at full precision -- dense-level accuracy at ~half the memory. When the cache is genuinely smaller than the prompt, the gap is real, degrades gracefully (retaining ~25% of the prompt costs ~9pp EM), and none of the inference-time knobs here fix it. That regime (context >> cache) is the open problem where positional/retention research -- including the spherical-KV direction -- would actually be stressed.

Contents

• keys_values/kvcache/pos_compact.py -- exact position-compaction transform + set_position_compaction
• keys_values/kvcache/base.py -- compact_positions flag, applied before the attention call (no-op for dense; off by default)
• examples/eviction_sweep.py, examples/grpo_helmet_crosseval.py --compact-positions -- the A2 sweep and A1 A/B arm
• test/kvcache/test_pos_compact.py -- delta-rotation exactness, dense-cache identity (no-op), h2o logit-change
• docs/POSITION_COMPACTION_A1.md, docs/EVICTION_SWEEP_A2.md -- findings (checked-in markdown)

Notes

• Base is grpo-upstream (Safe eviction defaults (#140), GRPO single-epoch cache reuse, HELMET GRPO driver awslabs/keys_values#142): the sweep reuses that PR's HELMET eval harness. Rebase onto main once Safe eviction defaults (#140), GRPO single-epoch cache reuse, HELMET GRPO driver awslabs/keys_values#142 lands.
• All results single-seed, n=100, 0.5B, 8k context -- flagged in the docs.