推論系統
KeysAndValues Fine-Tunes 4B Models for 128K Contexts with a Fixed 32K KV Cache
A new method lets models adapt directly to KV eviction strategies such as H2O during fine-tuning, raising accuracy on some 128K tasks from roughly 20%–64% to over 90%. However, it does not consistently lead on natural question-answering benchmarks, and the open-source implementation remains positioned as a research tool.
Researchers from AWS, the University of Amsterdam, and TU Berlin have introduced a sparse-attention fine-tuning method for long contexts and released [KeysAndValues](https://github.com/awslabs/keys_values) 0.2.0 under the Apache 2.0 license. The core problem is that when a model is trained with full attention but only given a fixed-capacity KV cache at deployment time, it has never learned which information will be evicted. Its output quality may degrade, and it may even fail to stop generating correctly.
The new method records overwrite decisions made by the cache policy during the forward pass and replays them during backpropagation. It controls memory usage through two-level activation checkpointing, CPU offload, and a differential representation of the KV buffer. Because it does not need to differentiate through the discrete eviction policy, it can theoretically work with any fixed-capacity policy. The team also rewrote the H2O path, combining FlashInfer SDPA with Triton kernels to accumulate attention weights for each KV slot directly during attention computation.
The [paper](https://arxiv.org/abs/2608.19920) evaluates Qwen3-4B-Instruct-2507 with LoRA rank 16 and a 32K cache on HELMET at 64K and 128K context lengths. On the 128K TREC coarse task, models jointly fine-tuned with the cache policy achieved 96.0%–96.4%, compared with only 23.2%–77.6% when models fine-tuned with full attention were switched to sparse inference. On CLINC150, accuracy likewise rose from 61.6%–68.0% to 94.0%–97.4%. The key engineering implication is that long-context training no longer necessarily requires splitting each sequence across multiple GPUs: the example can run on a single 40GB A100, while additional GPUs can instead be used to increase batch throughput.
The limitations are equally clear. Results across four natural question-answering datasets were mixed, and sequence parallelism remained superior in most NQ and HotpotQA configurations. The paper also primarily evaluates only one 4B model. The FlashInfer extension requires NVIDIA compute capability 8.0 or later, and the project explicitly notes that production deployment still requires substantial engineering. The next step is to see whether the method can integrate with paged batching in vLLM and SGLang while preserving its quality and latency advantages with larger models, quantized caches, and multi-tenant workloads.