Reducing the KV cache during long reasoning
An LLM stores previously computed keys and values in its KV cache and reuses them when generating the next token. Longer reasoning means a larger cache. Removing earlier entries saves memory, but it can also remove information needed to finish the solution.
CASK keeps states needed to form the answer and consolidates repeated intermediate reasoning. We measured how closely the compressed model follows the original model.
The limits of token ranking
We examined methods such as TriAttention that score tokens and retain a subset. In some settings, changing the score barely changed which tokens survived. A more elaborate score alone did not reliably preserve the information needed for reasoning.
CASK separates generated states by their role. States supporting answer formation or subsequent reasoning remain in the protected core. Similar states from repeated exploration and self-checking become mergeable scratch.
Only scratch is merged. Each group contributes to a representative key and value according to its weights. The merge geometry accounts for relevance to future tokens, and diagnostics track information lost during consolidation. Core states are excluded from merging.
Input and decoding
A long prompt can consume nearly the entire cache budget before generation starts. That leaves no room for a policy designed to compress generated reasoning. CASK handles the two intervals separately.
- Evict selected prefix entries to make room for generation.
- Preserve the core and selectively consolidate scratch during decoding.
Experiment setup and results
The reasoning experiments used Qwen3-8B with bfloat16 and SDPA on an H100 PCIe, with six examples each from AIME24 and AIME25.
We first generated a continuation without cache compression, then fed those same tokens to each compressed model one at a time. We recorded next-token agreement, the probability assigned to the reference token, and the first mismatch.
| Dataset | Method | Next-token agreement | Mean NLL |
|---|---|---|---|
| AIME24 | TriAttention | 88.2% | 0.383 |
| AIME24 | CASK | 90.7% | 0.268 |
| AIME25 | TriAttention | 89.1% | 0.357 |
| AIME25 | CASK | 90.3% | 0.313 |
At a budget of 384, CASK had higher agreement and lower mean NLL on both datasets. Its agreement also exceeded TriAttention at a budget of 512: 89.4% on AIME24 and 89.9% on AIME25.
Limits with long prompts
In LongBench, multi_news and vcsum exercised the decode-stage merge and improved reference fidelity. The gains on qasper were mainly attributable to prefix processing.
qmsum and gov_report still exhausted the budget at the prefix. CASK also had lower agreement on 2wikimqa at budget 384.
Paper and implementation
CASK: Core-Aware Selective KV Compression for Reasoning Traces
Buseong Kim · Heejun Gwon
Published on arXiv on . arXiv:2604.10900.
The public repository includes CASK and TriAttention implementations, experiment runners, recorded results and paper figures.