Learning how to Forget: Fine-tuning for Long-Context Sparse Attention

Abstract: A lot of prior work addressed key-value (KV) cache selection and compression by sparse attention to enable long-context inference for transformer LLMs without excessive hardware budgets. We provide a new method for fine-tuning models with sparse attention. It works for any KV cache policy, runs on a moderate hardware budget (e.g., a single Nvidia A100 GPU with 40 GB RAM), and allows the model to co-adapt with the policy, often outperforming models trained with exact attention (sequence parallelism). We also provide an efficient implementation of H2O sparse attention (the leading policy in our experiments) with dedicated scaled dot product attention kernel support. KeysAndValues (https://github.com/awslabs/keys_values), a new open source library for long-context inference and fine-tuning, provides easy-to-use and performant code for all methods discussed here.

• The paper introduces a policy-aware fine-tuning method using replay caches, nested checkpointing, and delta-encoded buffers to train 4B-parameter models with sparse attention on a single A100 GPU.
• Experiments show that models trained with their deployment-time cache policy preserve output termination and substantially outperform sequence-parallel baselines on strict tasks, including 49% versus 0% on json_kv.
• The work improves H2O cache eviction with per-position and age-normalized scoring plus faster SDPA support, while identifying inference latency, kernel limitations, and limited single-model evaluation as key barriers.

Overview

This paper addresses a gap between two strands of long-context transformer research: sparse attention methods that compress the key-value (KV) cache to a fixed size, and long-context fine-tuning methods that rely on exact attention distributed across devices via sequence or context parallelism (SP/CP). The authors observe that models fine-tuned with exact attention and then deployed with a KV cache eviction policy suffer from a train–inference mismatch, and they propose a fine-tuning method that trains the model with an arbitrary KV cache policy in place, on hardware comparable to sparse-attention inference itself — gradients for a 4B-parameter model on a single Nvidia A100 40 GB GPU (2608.19920). The work also contributes improved variants of the heavy-hitter oracle (H2O) policy with dedicated scaled dot product attention (SDPA) kernel support, and an open-source library, KeysAndValues.

The train–inference mismatch

The central hypothesis is that the choice of KV cache policy — which determines how the model's short-term memory is organized — should shape how the model is trained. When a model is trained under sequence parallelism, every token can attend to every earlier token; at inference time, most KV information is evicted according to logic the model never saw. The paper's experiments confirm this: sparse-attention inference for SP-trained checkpoints often performs significantly worse than for checkpoints trained with the same policy in place.

A concrete failure mode dominates the results: SP-trained (and base) checkpoints, when run with sparse attention, produce outputs that are far too long and largely nonsensical — they frequently fail to emit <eos> and generate up to the 128-token cap. Measured by the ratio $R$ of output length to target length and the fraction $p_{128}$ of maximal-length outputs, SP-trained models show $R$ values of 30–128 and $p_{128}$ near 100% across many Helmet tasks, whereas policy-co-trained ("us") checkpoints yield $R \approx 1$ and $p_{128} \approx 0$. Crucially, this is not a property of the checkpoints themselves: under exact inference, the same SP checkpoints produce well-calibrated output lengths. The failure stems purely from the inconsistency between training conditions and inference-time cache eviction.

Fine-tuning method

Fine-tuning with sparse attention naively requires storing KV cache buffers for every chunk in the autograd graph, costing $\mathcal{O}(L \cdot N \cdot \mathcal{D})$ memory — more than the full uncompressed cache would need. The method combines four ideas:

• Replay caches: cache policy decisions are recorded during the forward pass into a replay log; backward passes replay evictions rather than differentiating through the (often non-differentiable or expensive) policy.
• Nested activation checkpointing: chunks are grouped into cells of roughly one KV-cache-buffer size; backward runs an outer loop over layers and inner loops over cells, checkpointing layer inputs and incoming cache buffers to CPU. This handles both large depth and long context.
• Delta encoding of cache buffers: neighboring chunk buffers differ in only $S \cdot \mathcal{D}$ entries, related by a linear scatter recurrence that is trivially inverted via gather. Instead of storing full buffers per chunk in the graph, only deltas are stored, cutting autograd memory by a factor of $k$ (chunks per cell), down to $\mathcal{O}(N_C \cdot \mathcal{D})$ — comparable to inference alone.
• Autograd saved tensors hooks: delta encoding is implemented inside PyTorch autograd via pack/unpack hooks, using an annotation list and fingerprint matching (checking whether gather(x, index) equals the stored delta). The authors candidly note implementation fragility: node creation order does not reliably match hook call order due to operator fusion, unmatched arguments occur at low rates (harmless, just extra memory), and false matches must be avoided by keeping fingerprints large enough.

The result is policy-agnostic: unlike OOMB, which hides KV nodes from autograd via hand-written CUDA kernels specialized per sparse-attention variant, this approach works with any selection or compression policy and composes with GQA and quantization (KV buffers are quantized to 8 bits in the experiments).

H2O improvements

The paper revisits H2O, which scores cache slots by cumulative summed attention weights $p_{128}$0. Three changes are introduced: per-batch-position eviction decisions (the original code aggregates over batch positions, restricting $p_{128}$1 to depend only on $p_{128}$2); a normalized score dividing by slot age $p_{128}$3, correcting a bias toward longer-resident entries; and a fast implementation returning summed attention weights alongside a FlashInfer SDPA kernel via Triton code — since no FlashAttention-derived kernel provides these weights, prior H2O implementations relied on slow naive SDPA. The paper also details how FlexAttention can compute summed weights via a second, transposed SDPA call using returned log-sum-exp values, at most doubling cost. Notably, the ranking among H2O variants remains inconclusive in the experiments, though the original batch-aggregated variant shows a concerning failure on json_kv ($p_{128}$4, $p_{128}$5).

Experimental findings

Experiments use Qwen3-4B-Instruct-2507 with LoRA (rank 16), trained on adapted Helmet tasks at 64k and 128k context widths, comparing co-training ("us", cache length $p_{128}$6, chunk sizes 1024/2048) against MS-SWIFT sequence parallelism ("sp") and the base checkpoint ("no"). Results split sharply:

Task group	Metric	Outcome
nq, trivia_qa, hotpot_qa, pop_qa	SubEM	Mixed/inconclusive
trec_coarse, nlu, clinc150, inf_qa, inf_mc	Accuracy	"us" strongly outperforms
json_kv	SubEM	"us" ~50%, sp/no ~0%

On the first group, SubEM tolerates arbitrary trailing content as long as the target string appears anywhere, masking the verbosity failure mode; on stricter metrics, sp collapses (e.g., clinc150: 97.4 vs 64.2; inf_mc: 54.0 vs 25.0; json_kv: 49.0 vs 0.0). The paper argues this exposes a shortcoming of SubEM-style evaluation, which should require deterministic extraction of the answer. Ablations show that smaller chunk sizes ($p_{128}$7), giving policies finer-grained decisions, do not improve accuracy, and Q-Hitter variants perform consistently worse than plain H2O here. Timing figures show the method costs about 30% more than exact SP training per step at $p_{128}$8 — modest given constant-memory single-GPU gradient computation — with H2O policies adding only 2–4% over simple recency baselines thanks to the fast SDPA path.

Limitations and open questions

Several limitations are stated plainly. Latency remains the main obstacle to adoption: sparse attention processes $p_{128}$9 sequential chunks where RingAttention uses 8, so even with an 8× larger feasible batch via DDP, inference tends to be slower than CP/SP; the authors acknowledge the sequential nature of eviction decisions may be an inherent disadvantage when large hardware budgets are available. Current SDPA kernels lack support for implicitly defined causal masks over non-monotonic token positions and for returning summed attention weights, forcing costly sorting/reordering workarounds. vLLM's page structure, which spans all heads per page, cannot represent head-dependent eviction policies without significant refactoring. On the training side, the autograd saved-tensors-hook mechanism relies on fingerprint matching that the authors could not make fully robust, and asynchronous CPU offloading did not yet yield speedups in their implementation. Evaluation is limited to LoRA fine-tuning of a single 4B model on Helmet-derived tasks; generalization to full fine-tuning, other architectures, and other benchmarks is not established.

Conclusion

The paper demonstrates that fine-tuning a LLM with its deployment-time KV cache policy in place is feasible on modest hardware, via replay caches, nested activation checkpointing, and delta-encoded cache buffers integrated into PyTorch autograd. The empirical evidence supports the claim that train–inference consistency matters: co-trained models learn to terminate generation properly under eviction, while SP-trained models degenerate into verbose nonsense under sparse inference despite performing well with exact attention. Whether sparse attention can close its latency gap with sequence parallelism — through kernel fusion, better SDPA kernel support, and asynchronous host-memory use — remains the decisive open question for practical adoption.