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

# Learning how to Forget: Fine-tuning for Long-Context Sparse Attention
Source: [https://arxiv.org/html/2608.19920](https://arxiv.org/html/2608.19920)
Matthias SeegerNote:Correspondence tomseeger@gmail\.comAffiliation:Amazon Web ServicesEmail:[mseeger@gmail\.com](mailto:)Vihang PatilAffiliation:AmazonEmail:[pvihang@amazon\.de](mailto:)Konstantinos BenidisAffiliation:Amazon Web ServicesEmail:[kbenidis@amazon\.de](mailto:)Sebastian SchelterAffiliation:Technical University BerlinEmail:[schelter@tu\-berlin\.de](mailto:)

###### Abstract

A lot of prior work addressed key\-value \(KV\) cache selection and compression by sparse attention to enable long\-context inference for transformer language models 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\.[𝙺𝚎𝚢𝚜𝙰𝚗𝚍𝚅𝚊𝚕𝚞𝚎𝚜\\mathtt\{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\.

## 1Introduction

Modern large language models need to process very long contexts \(i\.e\., number of tokens\) for calling many tools with sizable outputs\([51](https://arxiv.org/html/2608.19920#bib.bib43);[17](https://arxiv.org/html/2608.19920#bib.bib15)\), running chain of thought reasoning\([66](https://arxiv.org/html/2608.19920#bib.bib57)\), or sustaining multi\-turn conversations\. While naive transformer implementations scale quadratically in compute and linearly in memory with context width, a lot of progress has been made on approximations with essentially linear time and constant memory scaling\. A particularly fruitful direction issparse attention, where key\-value \(KV\) information is stored in a fixed\-size KV cache, slots of which are evicted once it is full, and many different eviction policies have been proposed\.

In this paper, we address the problem of how topost\-traina transformer language model with sparse attention on a moderate hardware budget \(our experiments are run computing gradients for a 4B weights model on a single111We use 4 GPUs with distributed data parallel training to attain a larger batch size\.Nvidia A100 GPU with 40 GB RAM\)\. Our novel method works for any KV cache policy and requires no further approximations beyond sparse attention\. As we demonstrate in experiments on a range of long\-context benchmarks, our training algorithm allows the model to co\-adapt with the KV cache policy, often outperforming models trained with exact attention \(sequence parallelism\)\. Moreover, our fine\-tuning method runs on resources comparable to sparse attention inference\. It can be combined with orthogonal KV cache compression strategies such as grouped query attention\([1](https://arxiv.org/html/2608.19920#bib.bib1)\)or quantization\([28](https://arxiv.org/html/2608.19920#bib.bib25);[44](https://arxiv.org/html/2608.19920#bib.bib37)\)\.

The heavy\-hitter oracle \(H2O\)\([76](https://arxiv.org/html/2608.19920#bib.bib69)\)is one of the most prominent sparse attention policies\. We demonstrate several improvements to H2O, leading to a much more efficient implementation with dedicated scaled dot product attention \(SDPA\) kernel support\. Variants of H2O outperform other KV cache policies in our experiments, and our fast implementation takes a big step towards latencies competitive with SotA inference libraries such as vLLM\([33](https://arxiv.org/html/2608.19920#bib.bib28)\), which use context or sequence parallelism almost exclusively\. In summary, our contributions are:

- •New method for fine\-tuning transformer language models with sparse attention and arbitrary KV cache policy in place\. This method runs on resources comparable to sparse attention inference\. It combines nested activation checkpointing and CPU offloading with exploiting a linear KV cache buffer recurrence by way of autograd saved tensors packing\. Our method can process sequences of arbitrary length with constant resources\.
- •Methodological and implementation improvements of heavy\-hitter oracle \(H2O\) cache policy\([76](https://arxiv.org/html/2608.19920#bib.bib69)\)\. In particular, we provideTritoncode to return summed attention weights alongside a FlashInfer SDPA kernel\([72](https://arxiv.org/html/2608.19920#bib.bib63)\)\.
- •A comprehensive evaluation on a range of long\-context benchmarks\. Our training algorithm often outperforms models trained with sequence parallelism\([36](https://arxiv.org/html/2608.19920#bib.bib30)\)when sparse attention inference is used\.
- •𝙺𝚎𝚢𝚜𝙰𝚗𝚍𝚅𝚊𝚕𝚞𝚎𝚜\\mathtt\{KeysAndValues\}, a novel open source library for long\-context inference and fine\-tuning \([https://github\.com/awslabs/keys\_values](https://github.com/awslabs/keys_values)\)\.

## 2Related Work

There is a large body of work on long\-context inference by way of KV cache compression\. A simple idea is to group heads, so that less key and value vectors need to be stored\([54](https://arxiv.org/html/2608.19920#bib.bib46);[1](https://arxiv.org/html/2608.19920#bib.bib1);[7](https://arxiv.org/html/2608.19920#bib.bib5)\), or to impose a low\-rank structure in the query\-by\-key matrix\([15](https://arxiv.org/html/2608.19920#bib.bib10)\)\. These are modifications to be used during pre\-training already\. Cache buffers can be quantized to 8 or 4 bits, or even below\([28](https://arxiv.org/html/2608.19920#bib.bib25);[44](https://arxiv.org/html/2608.19920#bib.bib37);[79](https://arxiv.org/html/2608.19920#bib.bib71);[39](https://arxiv.org/html/2608.19920#bib.bib32);[56](https://arxiv.org/html/2608.19920#bib.bib48);[34](https://arxiv.org/html/2608.19920#bib.bib29);[58](https://arxiv.org/html/2608.19920#bib.bib49);[78](https://arxiv.org/html/2608.19920#bib.bib68)\)\. Sparse attention is a powerful general idea \(discussed in Section[3\.1](https://arxiv.org/html/2608.19920#S3.SS1)\) with many instantiations\. Big Bird\([77](https://arxiv.org/html/2608.19920#bib.bib67)\)prescribes fixed attention sparsity patterns\. The heavy\-hitter oracle \(H2O\)\([76](https://arxiv.org/html/2608.19920#bib.bib69)\)is discussed in Section[3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1)\. Q\-Hitter\([81](https://arxiv.org/html/2608.19920#bib.bib70)\)combines H2O with quantization, steering decisions by quantizability as well\. SnapKV\([40](https://arxiv.org/html/2608.19920#bib.bib31)\)uses summed attention weights like H2O, but makes decisions only at one point during generation\. Expected Attention\([12](https://arxiv.org/html/2608.19920#bib.bib12)\)tries to estimate future relevance of KV cache information \(under some strong assumptions\)\. FlexGen\([55](https://arxiv.org/html/2608.19920#bib.bib47)\)shows how to maximize throughput by using a cache hierarchy\. FastGen\([21](https://arxiv.org/html/2608.19920#bib.bib18)\)provides a meta\-strategy voting between a number of different cache policies\. CAKE\([49](https://arxiv.org/html/2608.19920#bib.bib42)\)runs sparse inference with a H2O\-related score, but also distributes an overall memory budget between layers\. Other sparse attention techniques include\([24](https://arxiv.org/html/2608.19920#bib.bib21);[69](https://arxiv.org/html/2608.19920#bib.bib59);[63](https://arxiv.org/html/2608.19920#bib.bib53);[8](https://arxiv.org/html/2608.19920#bib.bib6);[68](https://arxiv.org/html/2608.19920#bib.bib60);[62](https://arxiv.org/html/2608.19920#bib.bib54);[18](https://arxiv.org/html/2608.19920#bib.bib16);[71](https://arxiv.org/html/2608.19920#bib.bib62);[65](https://arxiv.org/html/2608.19920#bib.bib56)\)\. KVPop\([25](https://arxiv.org/html/2608.19920#bib.bib22)\)learns a cache policy against a future\-attention target, computed efficiently using FlexAttention\([14](https://arxiv.org/html/2608.19920#bib.bib14)\)\. Policies are parameterized as small xLSTMs[4](https://arxiv.org/html/2608.19920#bib.bib4)\. qTTT\([3](https://arxiv.org/html/2608.19920#bib.bib3)\)uses a few gradient updates at test time in order to improve inference results\. MInference\([30](https://arxiv.org/html/2608.19920#bib.bib27)\)and KVPress \([https://github\.com/NVIDIA/kvpress](https://github.com/NVIDIA/kvpress)\) are open source libraries providing several sparse attention methods\. ShadowKV\([60](https://arxiv.org/html/2608.19920#bib.bib52)\)is a high\-throughput long\-context inference system including KV cache selection\. SCBench\([41](https://arxiv.org/html/2608.19920#bib.bib33)\)provides a comprehensive empirical analysis of long\-context inference methods\.

Optimized scaled dot product attention \(SDPA\) kernels are essential for fast inference and training, pioneered by FlashAttention\([10](https://arxiv.org/html/2608.19920#bib.bib8);[52](https://arxiv.org/html/2608.19920#bib.bib44)\)\. FlashInfer\([72](https://arxiv.org/html/2608.19920#bib.bib63)\)is optimized for the inference case1≪Nq≪Nk1\\ll N\_\{q\}\\ll N\_\{k\}\(notation from Section[3\.1](https://arxiv.org/html/2608.19920#S3.SS1)\)\. FlexAttention\([14](https://arxiv.org/html/2608.19920#bib.bib14)\)allows to specify mask and score modification code, using𝚝𝚘𝚛𝚌𝚑\.𝚌𝚘𝚖𝚙𝚒𝚕𝚎\\mathtt\{torch\.compile\}under the hood\.

Our main contribution is on long\-contextfine\-tuning\. Prior work can be ordered into two groups\. Proposals in the first group modify multi\-head attention in ways which remedy the difficulties detailed in Section[3\.2](https://arxiv.org/html/2608.19920#S3.SS2)\. LongLoRA\([9](https://arxiv.org/html/2608.19920#bib.bib7)\)uses permutation and reshaping of keys and values, essentially trading\(B,N\)\(B,N\)for\(B⁡\(N/NC\),NC\)\(B\(N/N\_\{C\}\),N\_\{C\}\), whereBBis batch size,NNis sequence length,NCN\_\{C\}is cache length\. This speeds up MHA, but does not reduce KV memory, and works only ifN/NCN/N\_\{C\}is small\. Native sparse attention \(NSA\)\([74](https://arxiv.org/html/2608.19920#bib.bib65)\)bakes a mixture of some sparse attention kernels with different fixed policies into the model architecture\. DeepSeek sparse attention \(DSA\)\([16](https://arxiv.org/html/2608.19920#bib.bib11)\)is a refined variant of this idea\. Different toselectivesparse attention, this needs to be used during pre\-training already, whose cost is significantly increased\. IndexCache\([2](https://arxiv.org/html/2608.19920#bib.bib2)\)speeds up DSA somewhat by sharing the indexer \(i\.e\., the cache logic\) between several layers\. Also, by fixing the selection policy as part of the model choice, DSA offers considerable less flexibility during post\-training\. Finally, LSA and DSA still require storing the complete KV cache \(even though each attention call uses only a part of it\)\. LongGen\([20](https://arxiv.org/html/2608.19920#bib.bib19)\)uses static sparsity patterns for the top and bottom third of layers and full attention for the middle third, which speeds up training a bit, but does not reduce its memory requirements\. DMC\([46](https://arxiv.org/html/2608.19920#bib.bib40)\)uses a form of KV cache compression where new information is either appended or accumulated with the last recent slot\. They propose heuristics to train a model with this mechanism in place\. This approach seems restricted to KV cache updates during generation, it is not clear what is done when a large prompt needs to be processed\. Starting from Mamba\([23](https://arxiv.org/html/2608.19920#bib.bib20)\), there have been several attempts to resurrect LSTMs\([27](https://arxiv.org/html/2608.19920#bib.bib24)\), e\.g\.\([11](https://arxiv.org/html/2608.19920#bib.bib9);[4](https://arxiv.org/html/2608.19920#bib.bib4)\)\. When compared against the long\-context transformer SotA, none of them have been competitive enough in order to warrant costly pre\-training efforts\. YOCO\([61](https://arxiv.org/html/2608.19920#bib.bib51)\)proposes an architecture different to the transformer, where a single KV cache block serves all layers\. Note that KV buffer scaling with layers is not a major problem in practice, since at any time, all but one can be offloaded to CPU \(see Section[A\.4\.1](https://arxiv.org/html/2608.19920#A1.SS4.SSS1)\)\.

In the second group, KV cache buffers are not compressed, but both storage and computation are distributed across several devices\. This can be done with RingAttention\([42](https://arxiv.org/html/2608.19920#bib.bib36)\)or more recent variants\([43](https://arxiv.org/html/2608.19920#bib.bib38)\), in what is called context parallelism \(CP\), or with sequence parallelism \(SP\)\([36](https://arxiv.org/html/2608.19920#bib.bib30)\)\. While[38](https://arxiv.org/html/2608.19920#bib.bib34)observe that long sequences can be split into chunks, and that the computation graph factorizes along the chunk axis, the factor for the final chunk depends onallKV cache buffers of all layers, so cannot be represented on a single device222In fact, RingAttention’s way to compute gradients makes use of this graph structure implicitly\.\(attempts to sparsify this computation are heuristics, and experiments are done only on rather short sequences of up to 16k tokens\)\. OOMB\([37](https://arxiv.org/html/2608.19920#bib.bib35)\)shares properties with our work, such as chunk\-level processing, activation checkpointing, and efforts to compress KV cache buffers forautograd\. While for exact KV caches, they run into the same issues as\([38](https://arxiv.org/html/2608.19920#bib.bib34)\), their implementation supports NSA and DSA as well\. Their way of hiding KV cache buffers fromautogradrequires a number of complex dedicated CUDA kernels, which need to be hardcoded for every sparse attention variant\. We manage KV buffer size inautogradvia delta encoding, which renders our implementation agnostic to the KV cache policy\. Details are given in Section[A\.8](https://arxiv.org/html/2608.19920#A1.SS8)\. LongRoPE\([13](https://arxiv.org/html/2608.19920#bib.bib13);[53](https://arxiv.org/html/2608.19920#bib.bib45)\)combines SP with a search for non\-uniform RoPE and several lifting stages\. Other work tuning position encoding, data mix and fine\-tuning recipes, but not compressing KV cache, includes\([67](https://arxiv.org/html/2608.19920#bib.bib58);[80](https://arxiv.org/html/2608.19920#bib.bib72);[19](https://arxiv.org/html/2608.19920#bib.bib17)\)\. LongStraw\([84](https://arxiv.org/html/2608.19920#bib.bib75)\)is a system designed for long context reinforcement learning, with a specific emphasis on sharing prompt graphs and KV caches between different roll\-outs\. It uses OOMB for gradient computations, but could likely be configured with our method as well\. Highly optimized implementations of CP/SP are the state of the art for long\-context inference, e\.g\. vLLM\([33](https://arxiv.org/html/2608.19920#bib.bib28)\), SGLang\([83](https://arxiv.org/html/2608.19920#bib.bib74)\), and fine\-tuning, e\.g\. Nvidia NeMo RL \([https://github\.com/NVIDIA\-NeMo/RL](https://github.com/NVIDIA-NeMo/RL)\), MS\-SWIFT\([82](https://arxiv.org/html/2608.19920#bib.bib73)\)\. In Section[3\.3](https://arxiv.org/html/2608.19920#S3.SS3), we comment on reasons why sparse attention methods are less frequently used in practice, and how this could be changed\.

## 3Long\-Context Fine\-Tuning

In this section, we first introduce sparse attention and key\-value caching, providing several improvements to the heavy\-hitter oracle \(H2O\) cache policy\([76](https://arxiv.org/html/2608.19920#bib.bib69)\), leading to a more efficient implementation with dedicated SDPA kernel support\. Then, we detail our main contribution: a novel method to fine\-tune models with sparse attention in place, using resources comparable to sparse attention inference\. While in SotA CP/SP techniques, GPUs need to be used for sharding along the context, they can be used to increase throughput \(i\.e\., larger batch sizes\) or to handle larger models in our method\.

### 3\.1Sparse Attention\. Key\-Value Caching

Multi\-head attention \(MHA\)is the most important mechanism in modern transformer architectures\([64](https://arxiv.org/html/2608.19920#bib.bib55)\)\. At its core liesscaled dot product attention \(SDPA\):

𝒀=𝚂𝙳𝙿𝙰⁡\(𝑸,𝑲,𝑽\),𝒀,𝑸∈ℝ\(B,Hq,Nq,dh\),𝑲,𝑽∈ℝ\(B,Hk,Nk,dh\)\.\\boldsymbol\{Y\}=\\mathtt\{SDPA\}\(\\boldsymbol\{Q\},\\boldsymbol\{K\},\\boldsymbol\{V\}\),\\quad\\boldsymbol\{Y\},\\boldsymbol\{Q\}\\in\\mathbb\{R\}^\{\(B,H\_\{q\},N\_\{q\},d\_\{h\}\)\},\\;\\boldsymbol\{K\},\\boldsymbol\{V\}\\in\\mathbb\{R\}^\{\(B,H\_\{k\},N\_\{k\},d\_\{h\}\)\}\.\(1\)Here,𝑸\\boldsymbol\{Q\}\(queries\),𝑲\\boldsymbol\{K\}\(keys\),𝑽\\boldsymbol\{V\}\(values\) are 4D arrays,BBis the batch size,dhd\_\{h\}theper\-head embedding dimension,Hq,HkH\_\{q\},H\_\{k\}are numbers of heads, andNq,NkN\_\{q\},N\_\{k\}are sequence lengths \(i\.e\., the third dimension is mapping to token positions in the model context\)\. Themodel embedding dimensionisd=Hq⋅dhd=H\_\{q\}\\cdot d\_\{h\}\. Let us first assume thatHq=HkH\_\{q\}=H\_\{k\}and drop the first two dimensions\. Then:

𝒀=𝑴𝑽,𝑴=𝚜𝚘𝚏𝚝𝚖𝚊𝚡\(𝚖𝚊𝚜𝚔\(dh−1/2𝑸𝑲T\),𝚍𝚒𝚖=𝟷\)\.\\boldsymbol\{Y\}=\\boldsymbol\{M\}\\boldsymbol\{V\},\\quad\\boldsymbol\{M\}=\\mathtt\{softmax\}\\left\(\\mathtt\{mask\}\\left\(d\_\{h\}^\{\-1/2\}\\boldsymbol\{Q\}\\boldsymbol\{K\}^\{T\}\\right\),\\mathtt\{dim=1\}\\right\)\.\(2\)𝒀\\boldsymbol\{Y\}are weighted combinations of values𝑽\\boldsymbol\{V\}withattention weights𝑴∈ℝ\(B,Hq,Nq,Nk\)\\boldsymbol\{M\}\\in\\mathbb\{R\}^\{\(B,H\_\{q\},N\_\{q\},N\_\{k\}\)\},𝚜𝚘𝚏𝚝𝚖𝚊𝚡\\mathtt\{softmax\}applies𝒙↦exp⁡\(𝒙\)/\(𝟏T​exp⁡\(𝒙\)\)\\boldsymbol\{x\}\\mapsto\\exp\(\\boldsymbol\{x\}\)/\(\\boldsymbol\{1\}^\{T\}\\exp\(\\boldsymbol\{x\}\)\)along rows\.𝚖𝚊𝚜𝚔\\mathtt\{mask\}implements causal masking:\(𝚖𝚊𝚜𝚔\(𝑿\)\)i,j=xi,j−∞I\{P\+i<t\(j\)\}\(\\mathtt\{mask\}\(\\boldsymbol\{X\}\)\)\_\{i,j\}=x\_\{i,j\}\-\\infty\\mathrm\{I\}\_\{\\\{P\+i<t\(j\)\\\}\}, wherePPandt⁡\(j\)t\(j\)are token positions \(see Section[3\.3\.1](https://arxiv.org/html/2608.19920#S3.SS3.SSS1)for details\)\. If arrays are indexed by\(b,h,j,k\)\(b,h,j,k\), SDPA operates on\(j,k\)\(j,k\)in the same way for all\(b,h\)\(b,h\), computations are parallelized over batch and head positions\.

Inference in transformers switches between prompt processing and token generation\. Generating a token after a prompt of sizeNNrequires SDPA withNk=N,Nq=1N\_\{k\}=N,N\_\{q\}=1, with keys and values of size\(B,Hk,N,dh\)\(B,H\_\{k\},N,d\_\{h\}\)in GPU memory, for each ofLLmodel layers\. Exact transformer inference therefore requires𝒪⁡\(L⋅N⋅B​Hk​dh\)\\mathcal\{O\}\(L\\cdot N\\cdot BH\_\{k\}d\_\{h\}\)GPU memory for the fullkey\-value \(KV\) cache\. Even for moderate context lengthsNNof several hundred thousands, the KV cache far surpasses the model weights in size and cannot be stored in GPU memory as is\.

A large amount of prior work confronts this problem \(see also Section[2](https://arxiv.org/html/2608.19920#S2)\)\. In grouped query attention \(GQA\)\([1](https://arxiv.org/html/2608.19920#bib.bib1)\), we setHq=Hk⋅qgH\_\{q\}=H\_\{k\}\\cdot q\_\{g\},qg\>1q\_\{g\}\>1, so thatqgq\_\{g\}heads map to the same query group, which reduces KV cache size by a factor ofqgq\_\{g\}\. Most relevant to our work issparse attention\(orselective KV caching; e\.g\.[76](https://arxiv.org/html/2608.19920#bib.bib69)\), where the KV cache is represented by fixed\-size buffersindependentof the context length of the model\. Once all slots are filled, new information overwrites \(orevicts\) existing ones\. Formally, the cache \(for one model layer\) is represented by arrays𝚔𝚎𝚢𝚜,𝚟𝚊𝚕𝚞𝚎𝚜:\(B,Hk,NC,dh\)\\mathtt\{keys\},\\mathtt\{values\}:\(B,H\_\{k\},N\_\{C\},d\_\{h\}\)and𝚝𝚘𝚔𝚎𝚗​\_​𝚙𝚘𝚜:\(B,Hk,NC\)\\mathtt\{token\\\_pos\}:\(B,H\_\{k\},N\_\{C\}\)\. Here,NCN\_\{C\}is thecache length, which is chosen as large as GPU memory permits\. For up toNCN\_\{C\}tokens, the cache is filled from left to right\. After that, theKV cache policyπl​\(b,h,t\)∈\{0,…,NC−1\}\\pi\_\{l\}\(b,h,t\)\\in\\\{\{0\},\\dots,\{N\_\{C\}\-1\}\\\}dictates where additional key\-value information is written for token positiont≥NCt\\geq N\_\{C\}, batch positionb∈\{0,…,B−1\}b\\in\\\{\{0\},\\dots,\{B\-1\}\\\}and head \(or query group\)h∈\{0,…,Hk−1\}h\\in\\\{\{0\},\\dots,\{H\_\{k\}\-1\}\\\}\. Importantly,πl\\pi\_\{l\}can depend onb,hb,h: a token may be in the cache for some batch positions and heads, and not for others\.t⁡\(b,h,j\)=𝚝𝚘𝚔𝚎𝚗​\_​𝚙𝚘𝚜​\[b,h,j\]t\(b,h,j\)=\\mathtt\{token\\\_pos\}\[b,h,j\]lists the token position of what is stored in\(b,h,j\)\(b,h,j\)\. We do not require complex memory layouts and dedicated kernels for this setup, as for example PagedAttention\([33](https://arxiv.org/html/2608.19920#bib.bib28)\)needs\. In fact, \([2](https://arxiv.org/html/2608.19920#S3.E2)\) depends on absolute token positions only via𝚖𝚊𝚜𝚔\\mathtt\{mask\}\. This is as sparse as the conventional triangular one, but depends on token positionst⁡\(⋅\)=𝚝𝚘𝚔𝚎𝚗​\_​𝚙𝚘𝚜t\(\\cdot\)=\\mathtt\{token\\\_pos\}, which through cache evictions becomes non\-monotonic in general\. Moreover,𝚝𝚘𝚛𝚌𝚑\.𝚐𝚊𝚝𝚑𝚎𝚛\\mathtt\{torch\.gather\}and𝚝𝚘𝚛𝚌𝚑\.𝚜𝚌𝚊𝚝𝚝𝚎𝚛\\mathtt\{torch\.scatter\}provide fast read and write access for these buffers \(see Section[A\.1\.1](https://arxiv.org/html/2608.19920#A1.SS1.SSS1)\)\.

With sparse attention, we can run inference for any context widthNN\. First, we process up to the firstNCN\_\{C\}tokens with a single SDPA call \(Nq=Nk=NCN\_\{q\}=N\_\{k\}=N\_\{C\}\), this is known asprefilling\. The remainingN−NCN\-N\_\{C\}tokens are processed inchunksof sizeS<NCS<N\_\{C\}, using SDPA calls withNq=S,Nk=NCN\_\{q\}=S,N\_\{k\}=N\_\{C\}\(see Section[A\.4\.1](https://arxiv.org/html/2608.19920#A1.SS4.SSS1)for details\)\. Token generation usesNq=1,Nk=NCN\_\{q\}=1,N\_\{k\}=N\_\{C\}\. Memory requirements are independent ofNN\. For each chunk, the policy\[πl\]\[\\pi\_\{l\}\]is used to determineB⋅Hk⋅SB\\cdot H\_\{k\}\\cdot Spositions\(b,h,j\)\(b,h,j\)which are overwritten by the new keys and values\. WhileNCN\_\{C\}is chosen as large as memory permits, the choice ofSSis more subtle\. The largerSS, the fewer chunks, and less sequential computation results in faster processing\. The smallerSS, the more fine\-grained the cache policy is used, which can lead to better decisions \(see also Section[3\.3](https://arxiv.org/html/2608.19920#S3.SS3)\)\.

#### 3\.1\.1Variants of Heavy\-Hitter Oracle

The key idea behind the heayy\-hitter oracle \(H2O\)\([76](https://arxiv.org/html/2608.19920#bib.bib69)\)is to make use of the attention weights𝑴=\[mi,j\]\\boldsymbol\{M\}=\[m\_\{i,j\}\], a by\-product of SDPA \([2](https://arxiv.org/html/2608.19920#S3.E2)\)\. Dropping\(b,h\)\(b,h\)for the moment, we have that𝒀i,:=∑jmi,j𝑽j,:\\boldsymbol\{Y\}\_\{i,:\}=\\sum\_\{j\}m\_\{i,j\}\\boldsymbol\{V\}\_\{j,:\}and∑jmi,j=1\\sum\_\{j\}m\_\{i,j\}=1\.mi,jm\_\{i,j\}quantifies how much values𝑽j,:\\boldsymbol\{V\}\_\{j,:\}are used to create𝒀i,:\\boldsymbol\{Y\}\_\{i,:\}\. The cumulative sum∑imi,j\\sum\_\{i\}m\_\{i,j\}can be used to score the usefulness of vectors\(𝑲j,:,𝑽j,:\)\(\\boldsymbol\{K\}\_\{j,:\},\\boldsymbol\{V\}\_\{j,:\}\)in the KV cache\. Bringing\(b,h\)\(b,h\)back, we define the H2O score after having processedtttokens as

ϕh2ot​\(b,h,j\)=∑t⁡\(b,h,j\)≤s<tmb,h,s,j,\\phi^\{t\}\_\{\\text\{h2o\}\}\(b,h,j\)=\\sum\_\{t\(b,h,j\)\\leq s<t\}m\_\{b,h,s,j\},\(3\)wheret⁡\(b,h,j\)=𝚝𝚘𝚔𝚎𝚗​\_​𝚙𝚘𝚜​\[b,h,j\]t\(b,h,j\)=\\mathtt\{token\\\_pos\}\[b,h,j\]is the position represented at\(b,h,j\)\(b,h,j\)right now\. We sum overs∈\{t⁡\(b,h,j\),…,t\}s\\in\\\{\{t\(b,h,j\)\},\\dots,\{t\}\\\}because the slot is occupied from KV information corresponding to token positiont⁡\(b,h,j\)t\(b,h,j\), which entered the cache only then\. The largerϕh2ot​\(b,h,j\)\\phi^\{t\}\_\{\\text\{h2o\}\}\(b,h,j\), the more valuable this information has been so far\. When asked to insert new content forSStokens, for each\(b,h\)\(b,h\), we overwrite theseSSslotsjjfor whichϕh2ot​\(b,h,j\)\\phi^\{t\}\_\{\\text\{h2o\}\}\(b,h,j\)is smallest\.

In this paper, we modify the original H2O policy\([76](https://arxiv.org/html/2608.19920#bib.bib69)\)\(as provided by their implementation\) in several ways\. First, their code selects the same cache slots for each batch positionbb, using the scoreϕh2o\-origt​\(h,j\)=∑bϕh2ot​\(b,h,j\)\\phi^\{t\}\_\{\\text\{h2o\-orig\}\}\(h,j\)=\\sum\_\{b\}\\phi^\{t\}\_\{\\text\{h2o\}\}\(b,h,j\)\. The rationale for this restriction is unclear, we implement H2O without it as well\. Second, the cumulative H2O score seems to favour entries\(b,h,j\)\(b,h,j\)which have been in the cache for longer, since more terms in\[0,1\]\[0,1\]are summed then\. We introduce thenormalized H2O score:ϕh2o\-normt​\(b,h,j\)=\(t−t⁡\(b,h,j\)\)−1​ϕh2ot​\(b,h,j\)\\phi^\{t\}\_\{\\text\{h2o\-norm\}\}\(b,h,j\)=\(t\-t\(b,h,j\)\)^\{\-1\}\\phi^\{t\}\_\{\\text\{h2o\}\}\(b,h,j\)\.

Despite convincing empirical results of H2O, both in\([76](https://arxiv.org/html/2608.19920#bib.bib69)\)and Section[4](https://arxiv.org/html/2608.19920#S4), it is not widely used\. This is mostly because current implementations of H2O are much slower than the state of the art\. We need summed attention weights∑imb,h,i,j\\sum\_\{i\}m\_\{b,h,i,j\}for each\(b,h,j\)\(b,h,j\), as by\-product of SDPA \([2](https://arxiv.org/html/2608.19920#S3.E2)\), but none of the fast SDPA kernels derived from FlashAttention\([10](https://arxiv.org/html/2608.19920#bib.bib8)\)provide them\. Current H2O implementations use naive SDPA implementations, which are much too slow in practice\. Our implementation containsTritoncode to return summed attention weights alongside a FlashInfer SDPA kernel\([72](https://arxiv.org/html/2608.19920#bib.bib63)\)\. In Section[A\.6](https://arxiv.org/html/2608.19920#A1.SS6), we show how FlexAttention\([14](https://arxiv.org/html/2608.19920#bib.bib14)\)can be used to this end as well\. We come back to efficiency of sparse attention in Section[3\.3](https://arxiv.org/html/2608.19920#S3.SS3)\.

### 3\.2Fine\-Tuning for Sparse Attention

How should we train a model which uses sparse attention with some KV cache policy such as H2O? As noted in Section[2](https://arxiv.org/html/2608.19920#S2), all prior long\-context fine\-tuning methods either restrict the MHA approximation to a particular form, or use exact MHA with KV cache buffers distributed across several devices \(i\.e\., sequence or context parallelism\)\. However, the choice of KV cache policy, which dictates how the model’s short term memory is organized, should influence how the model is best trained\. Our results in Section[4](https://arxiv.org/html/2608.19920#S4)validate this hypothesis\.Sparse attention inference for a model trained with exact MHA and sequence parallelism \(which is the SotA\) often performs significantly worse than for a model trained with sparse attention and the desired policy in place\.

Fine\-tuning for models with sparse attention is difficult, because an enormous amount of memory is required, while GPU memory is on short supply\. We need to compute gradients for training loss functions on sequences of lengthN≫NCN\\gg N\_\{C\}, which is done by \(reverse mode\) automatic differentiation \(or error backpropagation,autograd\)\. Autograd works by creating a computation graph during the forward pass, whose nodes store arrays needed during the backward pass\. If the model hasLLlayers and the chunk size isSS, the training sequence is split into1\+⌈\(N−NC\)/S⌉1\+\\lceil\(N\-N\_\{C\}\)/S\\rceilchunks, the first \(prefill\) chunk of lengthNCN\_\{C\}and subsequent chunks of lengthSS\. SDPA is called for each layer and chunk, creating at least one node of the size of the KV cache, so we need at least𝒪⁡\(L⋅S−1​\(N−NC\)⋅NC⋅𝒟\)\\mathcal\{O\}\(L\\cdot S^\{\-1\}\(N\-N\_\{C\}\)\\cdot N\_\{C\}\\cdot\\mathcal\{D\}\)of memory, where𝒟=B​Hk​dh\\mathcal\{D\}=BH\_\{k\}d\_\{h\}\. AssumingS=a​NCS=aN\_\{C\}for some constantaa, this is𝒪⁡\(L⋅N⋅𝒟\)\\mathcal\{O\}\(L\\cdot N\\cdot\\mathcal\{D\}\): more than thefullKV cache would need, and far beyond what is tractable\.

We need several ideas in order to bring GPU memory requirements down to levels comparable to what inference needs\. First, we avoid differentiation through the KV cache policy \(which is often not even possible, and in general not tractable\)\. Along the forward pass, we store all KV cache policy decisions in areplay log, containing for each chunk \(and each layerll\) the tokens processed, and the decisions\{πl​\(b,h,t\)\}\\\{\\pi\_\{l\}\(b,h,t\)\\\}\. Later on, we usereplay caches, which act like normal KV caches, except that eviction decisions are replayed from the log\. If a cache policy is complex and expensive to compute, it needs to be run during the forward pass only\.

Next, we use activation checkpointing\([26](https://arxiv.org/html/2608.19920#bib.bib23)\)\. Even with moderate context widths, this technique is routinely used for models with many layers\. Gradients are computed in two passes: forward and backward\. Different from autograd, there is no computation graph built during the forward pass, only the input tensors for each transformer layer are stored to CPU memory\. The backward pass is split intoLLautograd calls, starting from the top\. Head gradients are supplied from the previous layer, inputs are loaded from CPU\. Computations graphs on GPU areLLtimes smaller, while forward computations have to be run twice\.

While standard activation checkpointing tackles largeLL, in long\-context situations the context widthNNis the more serious problem\. We cannot even keep complete inputs or head gradients for a single layer in GPU memory \(see also Section[A\.4\.1](https://arxiv.org/html/2608.19920#A1.SS4.SSS1)\), let alone KV buffers attached to each chunk in the computation graph\. We therefore use activation checkpointing twice, in a nested fashion\. To this end, we partition chunks intocells:\{\(B,S,d\)\}→\(B,k​S,d\)\\\{\(B,S,d\)\\\}\\to\(B,kS,d\), wherek=⌊α​NC/S⌋k=\\lfloor\\alpha N\_\{C\}/S\\rfloor, andα\>0\\alpha\>0is a hyperparameter which defaults toα=1\\alpha=1\. The first \(prefill\) chunk becomes the first cell\. As a rule of thumb, we group chunks into cells which occupy about the size of KV cache buffers\. The complete computation graph can be seen aslattice of cells: rows are layers, columns are cells \(i\.e\., groups of chunks\) along the context\. Our backward pass runs an outer loop over rows \(layers\), then inner loops over cells in each layer\. Each inner loop starts with a \(non\-autograd\) forward pass along the context, storing KV cache buffers \(the inner loop ”activations” to checkpoint\) going into each cell to CPU\. Next, autograd is run separately on each cell, starting from the right\. Inputs to a cell \(layer inputs from the bottom and KV cache buffers from the left\) and head gradients from the top are read from CPU, while head gradients from the right stay in GPU memory\. Gradients are accumulated\. We reuse the same CPU and GPU buffers for all inner loops\.333It is tempting to store all KV cache checkpoints to CPU during the initial forward pass, along with layer inputs\. However, this requires a lot of CPU memory, and the per\-layer forward passes to obtain KV cache checkpoints are subdominant to all other computations\.A detailed summary of our method is given in Section[A\.4\.2](https://arxiv.org/html/2608.19920#A1.SS4.SSS2)\.

Even with nested activation checkpointing, the autograd calls still need too much GPU memory\. Recall that a cell consists ofk=⌊α​NC/S⌋k=\\lfloor\\alpha N\_\{C\}/S\\rfloorchunks\. Autograd stores KV cache buffers for each chunk, so needs at least𝒪⁡\(k⋅NC⋅𝒟\)\\mathcal\{O\}\(k\\cdot N\_\{C\}\\cdot\\mathcal\{D\}\)memory, where𝒟=B​Hk​dh\\mathcal\{D\}=BH\_\{k\}d\_\{h\}\. As detailed in Section[3\.1](https://arxiv.org/html/2608.19920#S3.SS1),kkmust be sizable to allow cache policies to make good decisions\. In this section, we detail the last \(and maybe most important\) idea, which cuts GPU memory by a factor ofkk, to𝒪⁡\(NC⋅𝒟\)\\mathcal\{O\}\(N\_\{C\}\\cdot\\mathcal\{D\}\)per autograd call\. This is comparable to what is needed during inference alone\.

Consider KV cache buffers\(𝚔𝚎𝚢𝚜,𝚟𝚊𝚕𝚞𝚎𝚜\)\(\\mathtt\{keys\},\\mathtt\{values\}\),\(𝚔𝚎𝚢𝚜′,𝚟𝚊𝚕𝚞𝚎𝚜′\)\(\\mathtt\{keys\}^\{\\prime\},\\mathtt\{values\}^\{\\prime\}\)for neighboring chunks\. Their size is\(B,Hk,NC,dh\)\(B,H\_\{k\},N\_\{C\},d\_\{h\}\), but they only differ inS⋅𝒟S\\cdot\\mathcal\{D\}values, because a chunk consists ofSStokens only\. The relationship is simple:

𝚔𝚎𝚢𝚜′=𝚜𝚌𝚊𝚝𝚝𝚎𝚛⁡\(𝚔𝚎𝚢𝚜,𝚒𝚗𝚍𝚎𝚡,𝚔𝚎𝚢​\_​𝚗𝚎𝚠\),𝚟𝚊𝚕𝚞𝚎𝚜′=𝚜𝚌𝚊𝚝𝚝𝚎𝚛⁡\(𝚟𝚊𝚕𝚞𝚎𝚜,𝚒𝚗𝚍𝚎𝚡,𝚟𝚊𝚕𝚞𝚎​\_​𝚗𝚎𝚠\)\.\\mathtt\{keys\}^\{\\prime\}=\\mathtt\{scatter\}\(\\mathtt\{keys\},\\mathtt\{index\},\\mathtt\{key\\\_new\}\),\\;\\mathtt\{values\}^\{\\prime\}=\\mathtt\{scatter\}\(\\mathtt\{values\},\\mathtt\{index\},\\mathtt\{value\\\_new\}\)\.Here,𝚔𝚎𝚢​\_​𝚗𝚎𝚠,𝚟𝚊𝚕𝚞𝚎​\_​𝚗𝚎𝚠\\mathtt\{key\\\_new\},\\mathtt\{value\\\_new\}are KV vectors for new tokens with sizes\(B,Hq,S,dh\)\(B,H\_\{q\},S,d\_\{h\}\),𝚒𝚗𝚍𝚎𝚡\\mathtt\{index\}is based on the cache policyπl\\pi\_\{l\}, determining which slots are overwritten, and𝚜𝚌𝚊𝚝𝚝𝚎𝚛,𝚐𝚊𝚝𝚑𝚎𝚛\\mathtt\{scatter\},\\mathtt\{gather\}are linear𝚝𝚘𝚛𝚌𝚑\\mathtt\{torch\}operators defined in Section[A\.1\.1](https://arxiv.org/html/2608.19920#A1.SS1.SSS1)\.

This is a linear recurrence, which is easily inverted:

𝚔𝚎𝚢𝚜=𝚜𝚌𝚊𝚝𝚝𝚎𝚛⁡\(𝚔𝚎𝚢𝚜′,𝚒𝚗𝚍𝚎𝚡,𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\),𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢=𝚐𝚊𝚝𝚑𝚎𝚛⁡\(𝚔𝚎𝚢𝚜,𝚒𝚗𝚍𝚎𝚡\),\\mathtt\{keys\}=\\mathtt\{scatter\}\(\\mathtt\{keys\}^\{\\prime\},\\mathtt\{index\},\\mathtt\{delta\\\_key\}\),\\quad\\mathtt\{delta\\\_key\}=\\mathtt\{gather\}\(\\mathtt\{keys\},\\mathtt\{index\}\),\(4\)and the same for𝚟𝚊𝚕𝚞𝚎𝚜\\mathtt\{values\}\. Instead of storing𝚔𝚎𝚢𝚜,𝚟𝚊𝚕𝚞𝚎𝚜\\mathtt\{keys\},\\mathtt\{values\}for each chunk in the compute graph, it suffices to store444This is a form ofdelta encoding\([https://en\.wikipedia\.org/wiki/Delta\_encoding](https://en.wikipedia.org/wiki/Delta_encoding)\)\.𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢,𝚍𝚎𝚕𝚝𝚊​\_​𝚟𝚊𝚕𝚞𝚎\\mathtt\{delta\\\_key\},\\mathtt\{delta\\\_value\}\. Since memory requirements of autograd calls are dominated by the KV cache buffers, they are reduced by a factor ofkk\.

While the linear recurrence relation between neighboring cache buffers is simple, implementing it in the context ofPyTorch autogradis not\. We use a mechanism calledautograd saved tensors hooks555[https://docs\.pytorch\.org/tutorials/intermediate/autograd\_saved\_tensors\_hooks\_tutorial\.html](https://docs.pytorch.org/tutorials/intermediate/autograd_saved_tensors_hooks_tutorial.html), originally intended to implement activation checkpointing by CPU offloading, which can be shaped to ours needs\. In a nutshell, we use thePyTorchmechanism to store\(𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢,𝚍𝚎𝚕𝚝𝚊​\_​𝚟𝚊𝚕𝚞𝚎\)\(\\mathtt\{delta\\\_key\},\\mathtt\{delta\\\_value\}\)in the autograd graph in place of\(𝚔𝚎𝚢𝚜,𝚟𝚊𝚕𝚞𝚎𝚜\)\(\\mathtt\{keys\},\\mathtt\{values\}\)\(called ”packing”\), reconstructing the latter from the former and subsequent\(𝚔𝚎𝚢𝚜′,𝚟𝚊𝚕𝚞𝚎𝚜′\)\(\\mathtt\{keys\}^\{\\prime\},\\mathtt\{values\}^\{\\prime\}\)during the backward pass over chunks \(called ”unpacking”\)\. The key difficulty is the non\-selectiveness of the mechanism: it provides a𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔\\mathtt\{pack\\\_hook\}function called for all arraysPyTorch autograddecides to place into its graph\. There is no way to tag tensors in the forward code, so they can be recognized as𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔\\mathtt\{pack\\\_hook\}arguments\. Our solution is to createannotationsalongside the forward pass code, storing\(𝚒𝚗𝚍𝚎𝚡,𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\)\(\\mathtt\{index\},\\mathtt\{delta\\\_key\}\)for𝚔𝚎𝚢𝚜\\mathtt\{keys\},\(𝚒𝚗𝚍𝚎𝚡,𝚍𝚎𝚕𝚝𝚊​\_​𝚟𝚊𝚕𝚞𝚎\)\(\\mathtt\{index\},\\mathtt\{delta\\\_value\}\)for𝚟𝚊𝚕𝚞𝚎𝚜\\mathtt\{values\}\. In a𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔​\(𝒙\)\\mathtt\{pack\\\_hook\}\(\\boldsymbol\{x\}\)call, we relate𝒙\\boldsymbol\{x\}to current annotations:𝒙\\boldsymbol\{x\}matches\(𝚒𝚗𝚍𝚎𝚡,𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\)\(\\mathtt\{index\},\\mathtt\{delta\\\_key\}\)\(say\) if𝚐𝚊𝚝𝚑𝚎𝚛⁡\(𝒙,𝚒𝚗𝚍𝚎𝚡\)=𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\\mathtt\{gather\}\(\\boldsymbol\{x\},\\mathtt\{index\}\)=\\mathtt\{delta\\\_key\}\. A match leads to𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔​\(𝒙\)\\mathtt\{pack\\\_hook\}\(\\boldsymbol\{x\}\)returning a reference to\(𝚒𝚗𝚍𝚎𝚡,𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\)\(\\mathtt\{index\},\\mathtt\{delta\\\_key\}\), which is removed from the annotation list\. Note that failing to match an annotation does not lead to errors, but at most to a bit more memory being used\. More details are given in Section[A\.4\.3](https://arxiv.org/html/2608.19920#A1.SS4.SSS3)\.

### 3\.3Sparse Attention and Sequence Parallelism

While with context or sequence parallelism, the context width is strictly limited by the number and memory size of GPUs available, sparse attention inference can be run for any context width on moderate GPU resources\. Moreover, redundancies which exist in multi\-head attention, can be exploited by way of cache compression, and experimental results with H2O are in general not worse than with exact attention even if the KV cache is compressed to 20% or less\([76](https://arxiv.org/html/2608.19920#bib.bib69);[81](https://arxiv.org/html/2608.19920#bib.bib70)\)\. Even if many GPUs are available, using sparse attention allows us to increase batch size by way of distributed data parallel, or to keep more layers in GPU memory\. A large number of sparse attention variants have been proposed \(see Section[2](https://arxiv.org/html/2608.19920#S2)\)\.Why is it then that sparse attention methods are hardly used in SotA inference libraries such as vLLM\([33](https://arxiv.org/html/2608.19920#bib.bib28)\)? The short answer is that latency is significantly higher with existing sparse attention implementations\. Further comments are in Section[A\.7](https://arxiv.org/html/2608.19920#A1.SS7)\.

Some of the gap is due to less low level implementation support for sparse attention\. Highly optimized SDPA kernels are vital for fast inference\([10](https://arxiv.org/html/2608.19920#bib.bib8);[14](https://arxiv.org/html/2608.19920#bib.bib14);[72](https://arxiv.org/html/2608.19920#bib.bib63)\)\. However, as noted in Section[3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1), existing kernels do not cater for sparse attention inference \(see details in Section[3\.3\.1](https://arxiv.org/html/2608.19920#S3.SS3.SSS1)\)\. Other reasons for the gap are more difficult to address\. A long sequence is split into chunks, the first \(prefill\) chunk of lengthNCN\_\{C\}\(cache length\), subsequent chunks of lengthSS\. LargerSSmeans fewer chunks and lower latency\. But KV cache policies can make useful eviction decisions only ifSSis much less thanNCN\_\{C\}\(in our experiments in Section[4](https://arxiv.org/html/2608.19920#S4), we useNC=32768N\_\{C\}=32768andS∈\{1024,2048\}S\\in\\\{1024,2048\\\}\)\. For the extreme choiceS=NCS=N\_\{C\}, the whole cache is overwritten by new content for every chunk, and the KV cache policy plays no role at all\! Information cannot be kept in the cache beyond a chunk if we do not allow for substantial overlap\.

For example, suppose we use 8 devices, each supporting a cache lengthNCN\_\{C\}, and the sequence length isN=8​NCN=8N\_\{C\}\. With RingAttention, each device holdsNCN\_\{C\}slots, and the sequence is processed in 8 sequential chunks\. But for sparse attention, we need1\+7​NC/S1\+7N\_\{C\}/Schunks, which can be substantially larger\. Despite sparse attention supporting a 8 times larger batch size via distributed data parallel \(DDP\), inference tends to still be slower than with RingAttention\. In future work, we plan to improve sparse attention latency by appropriate kernel fusion\. However, the sequential nature of decision making in sparse attention may be an inherent disadvantage over sequence or context parallelism, which may remain the best choice if a large hardware budget can be afforded for inference\.

Apart from large hardware requirements to even work on long sequences, RingAttention requires𝒪⁡\(L⋅D\)\\mathcal\{O\}\(L\\cdot D\)synchronizations between all devices per gradient update, while sparse attention DDP only needs a single gradient averaging reduction\. While memory transfer between devices can be run in parallel with computations, this needs double buffering666On each device, one set of𝚔𝚎𝚢𝚜,𝚟𝚊𝚕𝚞𝚎𝚜\\mathtt\{keys\},\\mathtt\{values\}are read for computation and transfer, another set is written to\([42](https://arxiv.org/html/2608.19920#bib.bib36)\)\.in RingAttention, doubling the GPU memory needed\. The peer\-to\-peer memory transfer is more brittle than DDP used with sparse attention, and robust training code is more difficult to implement\. Finally, the ”waste by pre\-allocation” issues which motivate the fairly complex PagedAttention[33](https://arxiv.org/html/2608.19920#bib.bib28), have a simpler solution with sparse inference: KV cache buffers are of a fixed length, there is no need to split777If we frequently encounter batches shorter than the cache length, we can build up buffers in chunks\. But most modern agentic AI applications come with long prompts anyway\.them into pages along the sequence axis, and no custom SDPA code is needed\.

#### 3\.3\.1Discussion: SDPA Kernels for Sparse Attention

Here, we list some ideas for SDPA kernel developers to better support sparse attention\. First, we consider causal masking for sparse attention\. In standard MHA \(the ”training case”\),\(𝚖𝚊𝚜𝚔\(𝑿\)\)i,j=xi,j−∞I\{i<j\}\(\\mathtt\{mask\}\(\\boldsymbol\{X\}\)\)\_\{i,j\}=x\_\{i,j\}\-\\infty\\mathrm\{I\}\_\{\\\{i<j\\\}\}\. For sparse attention, KV information is stored in cache buffers in an ordering given by token positionst⁡\(b,h,j\)t\(b,h,j\), and the causal mask is given by\(b,h,i,j\)↦\(−∞\)I\{P\+i<t\(b,h,j\)\}\(b,h,i,j\)\\mapsto\(\-\\infty\)\\mathrm\{I\}\_\{\\\{P\+i<t\(b,h,j\)\\\}\}, wherePPis the number of tokens processed before the current MHA call \(so the new information is for tokens\{P,…,P\+Nq−1\}\\\{\{P\},\\dots,\{P\+N\_\{q\}\-1\}\\\}\)\. The new key\-value information has been written into the cache already, so that\{P,…,P\+Nq−1\}\\\{\{P\},\\dots,\{P\+N\_\{q\}\-1\}\\\}is part of\{t⁡\(b,h,j\)\}\\\{t\(b,h,j\)\\\}for each\(b,h\)\(b,h\)\.

Unfortunately, none of the fast SDPA kernel codes we know of support such variants of causal masking in an implicitly888We cannot pass an explicit mask matrix, which would be huge and defy the purpose of fast SDPA\.defined way\. In our implementation, we sort the token positions and reorder keys and values according to this index, separate for each\(b,h\)\(b,h\), after which we can use standard causal masking, where queries are right\-aligned with keys and values\. This needs extra computation and memory which could be saved with better SDPA kernel support\. Based on our experience, the following simple extensions of fast SDPA kernel libraries could make a major difference for sparse attention:

- •Return summed attention weights∑imb,h,i,j\\sum\_\{i\}m\_\{b,h,i,j\}\(see Section[3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1)\), an array of size\(B,Hq,Nk\)\(B,H\_\{q\},N\_\{k\}\), based on attention weights which are computed anyway\. This allows for H2O and related scores to be computed, driving advanced KV cache policies\.
- •Allow for implicitly defined causal masks of the form\(b,h,i,j\)↦\(−∞\)I\{P\+i<t\(b,h,j\)\}\(b,h,i,j\)\\mapsto\(\-\\infty\)\\mathrm\{I\}\_\{\\\{P\+i<t\(b,h,j\)\\\}\}, wheret⁡\(⋅\)t\(\\cdot\)is an\(B,Hk,Nk\)\(B,H\_\{k\},N\_\{k\}\)integer array\. While FlexAttention\([14](https://arxiv.org/html/2608.19920#bib.bib14)\)supports custom mask patterns, they need to be written in terms of scalar index variables, so cannot have a 3D array as input to the compute graph\.

## 4Experiments

The key question addressed in our experiments is:if long context inference uses sparse attention with a particular KV cache policy, how much is gained by fine\-tuning the model with the same policy in place \(using our novel method\) over training it with state of the art libraries using sequence or context parallelism? We run comparisons on the Helmet\([73](https://arxiv.org/html/2608.19920#bib.bib64)\)benchmarks, with context widths of 64k and 128k, covering a range of cache policies:

- •𝚕𝚊𝚜𝚝𝚛𝚎𝚌\\mathtt\{lastrec\}\(lr\):π\(b,h,t\)=tI\{t<NC\}\+\(mod\(t−NC,NC−β\)\+β\)I\{t≥NC\}\\pi\(b,h,t\)=t\\mathrm\{I\}\_\{\\\{t<N\_\{C\}\\\}\}\+\(\\mathrm\{mod\}\(t\-N\_\{C\},N\_\{C\}\-\\beta\)\+\\beta\)\\mathrm\{I\}\_\{\\\{t\\geq N\_\{C\}\\\}\}, whereβ∈\[0,NC\)\\beta\\in\[0,N\_\{C\}\)\. Keeps the last recentNC−βN\_\{C\}\-\\betaand firstβ\\betatokens in the cache\. It is important to chooseβ\>0\\beta\>0as default ”attention sink”\([69](https://arxiv.org/html/2608.19920#bib.bib59)\)\. Our default isβ=min⁡\(16,⌈NC/8⌉\)\\beta=\\min\(16,\\lceil N\_\{C\}/8\\rceil\)\.
- •𝚜𝚖𝚊𝚛𝚝​\_​𝚕𝚊𝚜𝚝𝚛𝚎𝚌\\mathtt\{smart\\\_lastrec\}\(slr\): Variant of𝚕𝚊𝚜𝚝𝚛𝚎𝚌\\mathtt\{lastrec\}, where the numberβ\\betaof initial tokens is chosen dependent on content \(see Section[A\.2\.1](https://arxiv.org/html/2608.19920#A1.SS2.SSS1)for details\)\. A simple version of this heuristic appeared in\([24](https://arxiv.org/html/2608.19920#bib.bib21)\)\.
- •𝚑𝟸𝚘\\mathtt\{h2o\}\(h2o\\mathrm\{h2o\}\),𝚑𝟸𝚘​\_​𝚗𝚘𝚛𝚖\\mathtt\{h2o\\\_norm\}\(h2ono\\mathrm\{h2o\}^\{\\text\{no\}\}\),𝚑𝟸𝚘​\_​𝚘𝚛𝚒𝚐\\mathtt\{h2o\\\_orig\}\(h2oor\\mathrm\{h2o\}^\{\\text\{or\}\}\): Variants of H2O\([76](https://arxiv.org/html/2608.19920#bib.bib69)\)\(see Section[3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1)for details\)\.𝚑𝟸𝚘​\_​𝚘𝚛𝚒𝚐\\mathtt\{h2o\\\_orig\}is equivalent to their code released, whereπ⁡\(b,h,t\)\\pi\(b,h,t\)does not depend on batch positionbb\.

We train aQwen3\-4B\-Instruct\-2507999This checkpoint has been post\-trained to process long contexts\. We still need to adjust the RoPE position encoding, which we do with YaRN\.model\([70](https://arxiv.org/html/2608.19920#bib.bib61)\), usingAdamW\([45](https://arxiv.org/html/2608.19920#bib.bib39)\)with base learning rate0\.00050\.0005for up to 5 epochs\. We train LoRA weights only\([29](https://arxiv.org/html/2608.19920#bib.bib26)\)\(rankr=16r=16,α=16\\alpha=16, on all linear blocks\)\. We run on four Nvidia A100 40 GB devices with a per\-device batch size of 2, using RoPE\([59](https://arxiv.org/html/2608.19920#bib.bib50)\)and YaRN\([48](https://arxiv.org/html/2608.19920#bib.bib41)\)for position encoding\. Fine\-tuning is done in different ways:

- •Sequence parallelism \(sp\): We use MS\-SWIFT\([82](https://arxiv.org/html/2608.19920#bib.bib73)\)for fine\-tuning with DeepSpeed ZeRO\-3 offload, FlashAttention, and Liger kernels enabled, running on four Nvidia A100 40GB GPUs\. We use a per\-device batch size of 2 and sequential gradient accumulation \(effective batch size 8\)101010For 64k \(128k\), 2 \(4\) devices cover the context, so we need 2 \(4\) sequential gradient accumulation steps\.\.
- •Our method \(us\): We use a cache lengthNc=32768N\_\{c\}=32768, batch sizeS∈\{1024,2048\}S\\in\\\{1024,2048\\\}, and chunks per cell multiplierα=1\\alpha=1forS=2048S=2048,α=0\.75\\alpha=0\.75forS=1024S=1024\(see also Section[A\.4\.1](https://arxiv.org/html/2608.19920#A1.SS4.SSS1)\)\. KV cache buffers are quantized to 8 bits using𝚝𝚘𝚛𝚌𝚑𝚊𝚘\\mathtt\{torchao\}\. We run distributed data parallel optimization on 4 devices to obtain an effective batch size of 8\. We evaluate the model on a heldout validation set every 10 gradient steps \(5 forpop\_qa\), and choose the checkpoint with the lowest validation loss for testing\.

64k datasets128k datasetsnqtri\_qahot\_qapop\_qanqtri\_qahot\_qapop\_qausspusspusspusspusspusspusspusspexact\\mathrm\{exact\}\-​50\.7\-​79\.8\-​60\.0\-​62\.7\-​50\.7\-​68\.7\-​46\.3\-​57\.0lr2​k\\mathrm\{lr\}\_\{2k\}​33\.5​57\.2​75\.3​57\.5​53\.3​62\.7​43\.7​60\.5​26\.0​33\.0​50\.8​52\.3​31\.0​46\.3​34\.0​25\.0slr2​k\\mathrm\{slr\}\_\{2k\}​47\.3​56\.5​74\.5​60\.8​50\.0​67\.3​44\.0​56\.7​26\.0​33\.7​61\.2​51\.8​34\.0​42\.0​37\.7​22\.2h2o2​k\\mathrm\{h2o\}\_\{2k\}​47\.2​70\.8​78\.0​72\.0​53\.0​68\.7​57\.5​44\.7​24\.2​40\.7​47\.8​63\.7​19\.3​26\.0​53\.3​49\.8h2o2​kno\\mathrm\{h2o\}\_\{2k\}^\{\\text\{no\}\}​47\.8​68\.2​63\.2​54\.5​58\.3​70\.0​53\.0​39\.8​43\.5​51\.3​66\.7​55\.3​37\.3​51\.0​50\.2​25\.2h2o2​kor\\mathrm\{h2o\}\_\{2k\}^\{\\text\{or\}\}​49\.5​73\.3​66\.3​65\.7​57\.3​68\.7​62\.2​45\.5​45\.3​58\.8​71\.2​71\.0​36\.7​44\.3​50\.2​33\.3lr1​k\\mathrm\{lr\}\_\{1k\}​59\.7​57\.0​73\.0​60\.7​47\.7​65\.3​41\.7​59\.3​23\.5​32\.7​59\.5​50\.0​28\.7​45\.0​34\.5​25\.2slr1​k\\mathrm\{slr\}\_\{1k\}​37\.8​57\.0​59\.8​59\.7​52\.7​65\.0​46\.8​57\.0​29\.3​36\.2​58\.2​49\.7​33\.7​44\.7​34\.3​21\.2h2o1​k\\mathrm\{h2o\}\_\{1k\}​62\.0​72\.8​79\.7​72\.0​55\.0​68\.0​59\.3​45\.8​23\.2​41\.3​51\.7​58\.7​25\.3​24\.3​51\.0​50\.5h2o1​kno\\mathrm\{h2o\}\_\{1k\}^\{\\text\{no\}\}​47\.5​71\.3​62\.3​59\.3​61\.7​73\.0​51\.0​43\.5​42\.7​53\.3​70\.3​58\.2​42\.3​54\.7​44\.1​26\.8h2o1​kor\\mathrm\{h2o\}\_\{1k\}^\{\\text\{or\}\}​49\.0​72\.2​75\.7​68\.7​57\.0​66\.7​55\.0​47\.3​44\.3​60\.5​72\.0​74\.8​31\.0​41\.3​51\.0​30\.8Table 1:Results for long\-context inference with 5 KV cache policies and chunk sizes2048=2​k,1024=1​k2048=2k,1024=1k\(rows\)\. The first rowexactis for exact inference \(sequence parallelism\)\. We showSubEMvalues on test splits for different Helmet datasetsnq, trivia\_qa, hotpot\_qa, pop\_qa, limiting sequence lengths to 64k or 128k tokens\. Columnsusare for models trained using our novel method with the same cache policy in place, columnsspare for models trained with sequence parallelism\.datasettrnexactslr1​k\\mathrm\{slr\}\_\{1k\}h2o1​kno\\mathrm\{h2o\}\_\{1k\}^\{\\text\{no\}\}h2o1​kor\\mathrm\{h2o\}\_\{1k\}^\{\\text\{or\}\}trec\_coarseus\-​96\.0​96\.4​96\.2sp​97\.8​30\.0​23\.2​77\.6no\-​28\.2​19\.8​36\.0nluus\-​90\.0​87\.4​79\.8sp​90\.2​28\.6​32\.8​74\.0no\-​24\.8​30\.0​21\.2clinc150us\-​97\.4​96\.8​94\.0sp​97\.6​64\.2​61\.6​68\.0no\-​62\.6​54\.0​34\.8inf\_qaus\-​26\.6​32\.2​36\.8sp​40\.8​2\.2​2\.2​3\.3no\-​2\.5​2\.9​3\.4inf\_mcus\-​40\.0​42\.0​54\.0sp​66\.0​25\.0​29\.0​39\.0no\-​36\.0​41\.0​40\.0json\_kvus\-​49\.0​50\.0​3\.0sp​100\.0​0\.0​0\.0​1\.0no\-​0\.0​0\.0​0\.0Table 2:Results for 6 additional Helmet datasets not featured in Table[1](https://arxiv.org/html/2608.19920#S4.T1)\(context width 128k\)\. Inference under 3 KV cache policies \(chunk size1024=1​k1024=1k\),exactuses sequence parallelism \(column\)\.trndenotes model checkpoint being used:ususes our novel method with the same cache policy in place,spis using sequence parallelism,nois the base checkpointQwen3\-4B\-Instruct\-2507\(no fine\-tuning\)\. Note that metrics are different, depending on the dataset \(see Table[3](https://arxiv.org/html/2608.19920#A1.T3)\)\.Results on 4 Helmet datasets \(nq, trivia\_qa, hotpot\_qa, pop\_qa\)\([73](https://arxiv.org/html/2608.19920#bib.bib64)\)and 10 different cache setups \(5 policies, 2 chunk sizes\) are provided in Table[1](https://arxiv.org/html/2608.19920#S4.T1)\. Respective results for the base checkpoint \(no fine\-tuning\) are given in Table[4](https://arxiv.org/html/2608.19920#A1.T4)in the Appendix \(referred to asnoelsewhere\)\. For further 6 Helmet datasets \(trec\_coarse, nlu, clinc150, inf\_qa, inf\_mc, json\_kv\), Table[2](https://arxiv.org/html/2608.19920#S4.T2)provides results for 3 cache policies and chunk sizeS=1024S=1024\. Details about Helmet datasets and metrics are given in Section[A\.3\.1](https://arxiv.org/html/2608.19920#A1.SS3.SSS1)\.

For the datasets in Table[1](https://arxiv.org/html/2608.19920#S4.T1), results are mixed and inconclusive:spis best fornqandhotpot\_qa,nofortrivia\_qa\(fine\-tuning does not help\), andusforpop\_qa\. However, for the datasets in Table[2](https://arxiv.org/html/2608.19920#S4.T2),usstrongly outperformsspandno\. A closer look at generated samples \(which can be up to 128 tokens\) reveals a major failure mode ofsp\(see Section[A\.5\.3](https://arxiv.org/html/2608.19920#A1.SS5.SSS3)\):its outputs are far too long and contain mostly random nonsense\. For most datasets in Table[2](https://arxiv.org/html/2608.19920#S4.T2), targets are single numerical values, and theAccuracymetric \(see Section[A\.3\.1](https://arxiv.org/html/2608.19920#A1.SS3.SSS1)\) compares this to the most frequently occuring number in the output\. Poor results are due to the output forspoften containing many numbers\. In contrast,uslearns how to stop properly and usually outputs a single number\. In fact,the same failure mode dominates outcomes in Table[1](https://arxiv.org/html/2608.19920#S4.T1)just as well, but the metricSubEMused for the 4 datasets ignores content or length of output, as long as the target string is contained in it\. Finally,spandnofail forjson\_kvas well, despite this using theSubEMmetric\. Targets are UUIDs of length 32 tokens\. While variants ofusidentify them about half the time, they are hardly ever contained in the outputs ofspandno\. In Section[A\.5\.3](https://arxiv.org/html/2608.19920#A1.SS5.SSS3), we quantify the failure mode in Table[7](https://arxiv.org/html/2608.19920#A1.T7)\. While outputs forusare close in length to true targets, they are longer by large factors forsp,no, which often \(but not always\) span the full 128 tokens \(despite true targets being much shorter\)\. We also provide randomly chosen examples for outputs there, showcasing their nonsense content forsp\.

The shortness of desired targets is a clear signal in the training data, expressed not only byus, but also by thespcheckpoints if exact inference is used with them \(see Table[8](https://arxiv.org/html/2608.19920#A1.T8)\)\. We should not be surprised byspexhibiting such failure modes\. When training with sequence parallelism \(SP\), each token can attend to any earlier one\. This property is cut during inference, when most KV information is evicted at some point according to a logic which SP was never aware of\. Clearly,models should be trained under the same conditions and restrictions which govern inference later on\. Our new method allows practitioners to do that even on a low budget, no matter what KV cache policy they like to use during inference\.

While not our main focus here, the different KV cache policies exhibit variable performance across the different datasets\. Ideally, the best policy is chosen for each task\. Our results are inconclusive when it comes to ranking the different H2O variants \(see Section[3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1)\)\. However, one concerning datapoint is the poor performance ofh2o1​kor\\text\{h2o\}\_\{1k\}^\{\\text\{or\}\}onjson\_kvin Table[2](https://arxiv.org/html/2608.19920#S4.T2)\. In Table[7](https://arxiv.org/html/2608.19920#A1.T7), we see thatR=3\.7R=3\.7andp128=85%p\_\{128\}=85\\%for this logic, hinting to a similar failure mode thanspandno\. At least in this case, the decision in\([76](https://arxiv.org/html/2608.19920#bib.bib69)\)to score and evict batch dimensions together works much less well than the alternatives\.

## 5Conclusions

We showed how transformer language models with sparse attention can be fine\-tuned on a moderate hardware budget \(e\.g\., a single Nvidia A100 GPU with 40 GB RAM\)\. Our method works for any KV cache selection or compression policy and allows the model to co\-adapt with the policy, often outperforming models trained with exact attention \(sequence parallelism\)\. We also provide a much more efficient implementation of H2O sparse attention \(the leading policy in our experiments\) with dedicated scaled dot product attention \(SDPA\) kernel support\. By simplifying KV cache structure and clarifying the requirements on SDPA, we hope to direct more attention of the fast inference community on sparse attention \(see Section[3\.3\.1](https://arxiv.org/html/2608.19920#S3.SS3.SSS1)\), which despite its major potential for post\-training specialization via cache selection or compression policy design does not currently play a significant role in long\-context inference or fine\-tuning practice\.

In future work, we will combine context parallelism with sparse attention\. We are also considering kernel fusion ideas in order to narrow the latency gap further\. An important direction will be multi\-stream asynchronous implementations which allow for on\-the\-fly CPU offloading\([75](https://arxiv.org/html/2608.19920#bib.bib66)\)\. We believe that once the host memory of a system can be used without much synchronization overhead and less double buffering, many current difficulties with KV caching will be much diminished\. Finally, we hope that𝙺𝚎𝚢𝚜𝙰𝚗𝚍𝚅𝚊𝚕𝚞𝚎𝚜\\mathtt\{KeysAndValues\}\([https://github\.com/awslabs/keys\_values](https://github.com/awslabs/keys_values)\), the open source library with which most experiments were run here \(see Section[A\.9](https://arxiv.org/html/2608.19920#A1.SS9)\), will make it easier to use, compare and extend sparse attention policies, work on which so far is somewhat cluttered when it comes to implementations\.

## References

- Ainslieet al\.\(2023\)J\. Ainslie, J\. Lee\-Thorp, M\. de Jong, Y\. Zemlyanskiy, Lebrón, and F\. SanghaiGQA: training generalized multi\-query transformer models from multi\-head checkpoints\.InProceedings of the 2023 Conference on Empirical Methods in Natural Language Processing \(EMNLP\),pp\. 4895–4901\.Cited by:[§1](https://arxiv.org/html/2608.19920#S1.p2.1),[§2](https://arxiv.org/html/2608.19920#S2.p1.1),[§3\.1](https://arxiv.org/html/2608.19920#S3.SS1.p3.1)\.
- Baiet al\.\(2026\)Y\. Bai, Q\. Dong, T\. Jiang, X\. Lv, Z\. Du, A\. Zeng, J\. Tang, and J\. LiIndexCache: accelerating sparse attention via cross\-layer index reuse\.Technical reportTechnical ReportarXiv:2603\.12201 \[cs\.CL\]\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p3.1)\.
- Bansalet al\.\(2026\)R\. Bansal, A\. Zhang, R\. Tiwari, L\. Madaan, S\. Duvvuri, F\. Devvrit, D\. Brandfonbrener, D\. Alvarez\-Melis, P\. Bhargava, M\. Kale, and S\. JelassiLet’s \(not\) just put things in context: test\-time training for long\-context LLMs\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- \[4\]M\. Beck, K\. Pöppel, M\. Spanring, A\. Auer, O\. Prudnikova, M\. Kopp, G\. Klambauer, J\. Brandstetter, and S\. HochreiterxLSTM: extended long short\-term memory\.See[22](https://arxiv.org/html/2608.19920#bib.bib83),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1),[§2](https://arxiv.org/html/2608.19920#S2.p3.1)\.
- D\. Belgrave, C\. Zhang, H\. Lin, R\. Pascanu, P\. Koniusz, M\. Ghassemi, and N\. Chen \(Eds\.\) \(2025\)D\. Belgrave, C\. Zhang, H\. Lin, R\. Pascanu, P\. Koniusz, M\. Ghassemi, and N\. Chen \(Eds\.\)Advances in neural information processing systems38\.Curran Associates\.Cited by:[18](https://arxiv.org/html/2608.19920#bib.bib16),[34](https://arxiv.org/html/2608.19920#bib.bib29),[40](https://arxiv.org/html/2608.19920#bib.bib31),[43](https://arxiv.org/html/2608.19920#bib.bib38)\.
- S\. Bengio, H\. Wallach, H\. Larochelle, K\. Grauman, N\. Cesa\-Bianchi, and R\. Garnett \(Eds\.\) \(2018\)S\. Bengio, H\. Wallach, H\. Larochelle, K\. Grauman, N\. Cesa\-Bianchi, and R\. Garnett \(Eds\.\)Advances in neural information processing systems31\.Curran Associates\.Cited by:[64](https://arxiv.org/html/2608.19920#bib.bib55)\.
- \[7\]W\. Brandon, M\. Mishra, A\. Nrusimha, R\. Panda, and J\. Ragan\-KelleyReducing transformer key\-value cache size with cross\-layer attention\.See[22](https://arxiv.org/html/2608.19920#bib.bib83),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Caiet al\.\(2025\)Z\. Cai, Y\. Zhang, B\. Gao, Y\. Liu, Y\. Li, T\. Liu, K\. Lu, W\. Xiong, Y\. Dong, J\. Hu, and W\. XiaoPyramidKV: dynamic KV cache compression based on pyramidal information funneling\.InConference on Language Modeling,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Chenet al\.\(2024\)Y\. Chen, S\. Qian, H\. Tang, X\. Lai, Z\. Liu, S\. Han, and J\. JiaLongLoRA: efficient fine\-tuning of long\-context large language models\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p3.1)\.
- \[10\]T\. Dao, D\. Fu, S\. Ermon, A\. Rudra, and C\. RéFlashAttention: fast and memory\-efficient exact attention with IO\-awareness\.See[47](https://arxiv.org/html/2608.19920#bib.bib82),Cited by:[4th item](https://arxiv.org/html/2608.19920#A1.I8.i4.p1.1),[§A\.6](https://arxiv.org/html/2608.19920#A1.SS6.p1.1),[§2](https://arxiv.org/html/2608.19920#S2.p2.1),[§3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1.p3.1),[§3\.3](https://arxiv.org/html/2608.19920#S3.SS3.p2.1)\.
- \[11\]T\. Dao and A\. GuTransformers are SSMs: generalized models and efficient algorithms through structured state space duality\.See[50](https://arxiv.org/html/2608.19920#bib.bib77),pp\. 10041–10071\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p3.1)\.
- Devotoet al\.\(2025\)A\. Devoto, M\. Jeblick, and S\. JégouExpected attention: KV cache compression by estimating attention from future queries distribution\.Technical reportTechnical ReportarXiv:2510\.00636 \[cs\.AI\]\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- \[13\]Y\. Ding, L\. Zhang, C\. Zhang, Y\. Xu, N\. Shang, J\. Xu, F\. Yang, and M\. YangLongRoPE: extending LLM context window beyond 2 million tokens\.See[50](https://arxiv.org/html/2608.19920#bib.bib77),pp\. 11091–11104\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p4.1)\.
- Donget al\.\(2025\)J\. Dong, B\. Feng, D\. Guessous, Y\. Liang, and H\. HeFlexAttention: a programming model for generating optimized attention kernels\.InProceedings of the 8th MLSys Conference,pp\. 381–394\.Cited by:[4th item](https://arxiv.org/html/2608.19920#A1.I8.i4.p1.1),[§A\.6\.2](https://arxiv.org/html/2608.19920#A1.SS6.SSS2.p1.1),[§A\.6](https://arxiv.org/html/2608.19920#A1.SS6.p1.1),[§A\.6](https://arxiv.org/html/2608.19920#A1.SS6.p2.1),[§2](https://arxiv.org/html/2608.19920#S2.p1.1),[§2](https://arxiv.org/html/2608.19920#S2.p2.1),[2nd item](https://arxiv.org/html/2608.19920#S3.I1.i2.p1.1),[§3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1.p3.1),[§3\.3](https://arxiv.org/html/2608.19920#S3.SS3.p2.1)\.
- etal\. \(2024\)D\. etal\.DeepSeek\-V2: a strong, economical, and efficient mixture\-of\-experts language model\.Technical reportTechnical ReportarXiv:2405\.04434 \[cs\.CL\]\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- etal\. \(2025\)D\. etal\.DeepSeek\-V3\.2: pushing the frontier of open large language models\.Technical reportTechnical ReportarXiv:2512\.02556 \[cs\.CL\]\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p3.1)\.
- Fenget al\.\(2025\)J\. Feng, S\. Huang, X\. Qu, G\. Zhang, Y\. Qin, B\. Zhong, C\. Jiang, J\. Chi, and W\. ZhongReTool: reinforcement learning for strategic tool use in LLMs\.Technical reportTechnical ReportarXiv:2504\.11536 \[cs\.CL\]\.Cited by:[§1](https://arxiv.org/html/2608.19920#S1.p1.1)\.
- \[18\]Y\. Feng, J\. Lv, Y\. Cao, X\. Xie, and S\. ZhouAda\-KV: optimizing KV cache eviction by adaptive budget allocation for efficient LLM inference\.See[5](https://arxiv.org/html/2608.19920#bib.bib84),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Gaoet al\.\(2023\)T\. Gao, A\. Wettig, H\. Yen, and D\. ChenHow to train long\-context language models \(effectively\)\.InProceedings of the 63rd Annual Meeting of the Association for Computational Linguistics \(ACL\),pp\. 2391–2404\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p4.1)\.
- Geet al\.\(2025\)S\. Ge, X\. Lin, Y\. Zhang, J\. Han, and H\. PengA little goes a long way: efficient long context training and inference with partial contexts\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p3.1)\.
- Geet al\.\(2024\)S\. Ge, Y\. Zhang, L\. Liu, M\. Zhang, J\. Han, and J\. GaoModel tells you what to discard: adaptive KV cache compression for LLMs\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- A\. Globerson, L\. Mackey, D\. Belgrave, A\. Fan, U\. Paquet, J\. Tomczak, and C\. Zhang \(Eds\.\) \(2024\)A\. Globerson, L\. Mackey, D\. Belgrave, A\. Fan, U\. Paquet, J\. Tomczak, and C\. Zhang \(Eds\.\)Advances in neural information processing systems37\.Curran Associates\.Cited by:[4](https://arxiv.org/html/2608.19920#bib.bib4),[7](https://arxiv.org/html/2608.19920#bib.bib5),[28](https://arxiv.org/html/2608.19920#bib.bib25),[30](https://arxiv.org/html/2608.19920#bib.bib27),[52](https://arxiv.org/html/2608.19920#bib.bib44),[61](https://arxiv.org/html/2608.19920#bib.bib51),[67](https://arxiv.org/html/2608.19920#bib.bib58),[79](https://arxiv.org/html/2608.19920#bib.bib71),[80](https://arxiv.org/html/2608.19920#bib.bib72),[83](https://arxiv.org/html/2608.19920#bib.bib74)\.
- Gu and Dao \(2023\)A\. Gu and T\. DaoMamba: linear\-time sequence modeling with selective state spaces\.Technical reportTechnical ReportarXiv:2312\.00752 \[cs\.LG\]\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p3.1)\.
- Hanet al\.\(2024\)C\. Han, Q\. Wang, H\. Peng, W\. Xiong, Y\. Chen, H\. Ji1, and S\. WangLM\-Infinite: zero\-shot extreme length generalization for large language models\.InProceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies \(NAACL\-HLT\),pp\. 3991–4008\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1),[2nd item](https://arxiv.org/html/2608.19920#S4.I1.i2.p1.1)\.
- Hauzenbergeret al\.\(2026\)L\. Hauzenberger, N\. Schmidinger, A\. Hartl, D\. Stap, T\. Schmied, Böck, S\. Klambauer, and S\. HochreiterKVpop – key\-value cache compression with predictive online pruning\.Technical reportTechnical ReportarXiv:2607\.05061 \[cs\.LG\]\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Herrmannet al\.\(2019\)J\. Herrmann, O\. Beaumont, L\. Eyraud\-Dubois, J\. Hermann, A\. Joly, and A\. ShilovaOptimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory\.Technical reportTechnical ReportarXiv:1911\.13214 \[cs\.LG\]\.Cited by:[§A\.4\.2](https://arxiv.org/html/2608.19920#A1.SS4.SSS2.p2.1),[§3\.2](https://arxiv.org/html/2608.19920#S3.SS2.p4.1)\.
- Hochreiter and Schmidhuber \(1997\)S\. Hochreiter and J\. SchmidhuberLong short\-term memory\.Neural Computation9\(8\),pp\. 1735–1780\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p3.1)\.
- \[28\]C\. Hooper, S\. Kim, H\. Mohammadzadeh, M\. Mahoney, S\. Shao, K\. Keutzer, and A\. GholamiKVQuant: towards 10 million context length LLM inference with KV cache quantization\.See[22](https://arxiv.org/html/2608.19920#bib.bib83),Cited by:[§1](https://arxiv.org/html/2608.19920#S1.p2.1),[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Huet al\.\(2022\)E\. Hu, Y\. Shen, P\. Wallis, Z\. Allen\-Zhu, Y\. Li, S\. Wang, L\. Wang, and W\. ChenLoRA: low\-rank adaptation of large language models\.InInt\. Conf\. Learning Representations,Cited by:[§4](https://arxiv.org/html/2608.19920#S4.p1.2)\.
- \[30\]H\. Jiang, Y\. Li, C\. Zhang, Q\. Wu, X\. Luo, S\. Ahn, Z\. Han, A\. Abdi, D\. Li, C\. Lin, Y\. Yang, and L\. QiuMInference 1\.0: accelerating pre\-filling for long\-context LLMs via dynamic sparse attention\.See[22](https://arxiv.org/html/2608.19920#bib.bib83),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- S\. Koyejo, S\. Mohamed, A\. Agarwal, D\. Belgrave, K\. Cho, and A\. Oh \(Eds\.\) \(2022\)S\. Koyejo, S\. Mohamed, A\. Agarwal, D\. Belgrave, K\. Cho, and A\. Oh \(Eds\.\)Advances in neural information processing systems35\.Curran Associates\.Cited by:[66](https://arxiv.org/html/2608.19920#bib.bib57)\.
- A\. Krause, E\. Brunskill, K\. Cho, B\. Engelhardt, S\. Sabato, and J\. Scarlett \(Eds\.\) \(2023\)A\. Krause, E\. Brunskill, K\. Cho, B\. Engelhardt, S\. Sabato, and J\. Scarlett \(Eds\.\)International conference on machine learning40\.Vol\.202,Proceedings of Machine Learning Research\.Cited by:[55](https://arxiv.org/html/2608.19920#bib.bib47)\.
- Kwonet al\.\(2023\)W\. Kwon, Z\. Li, S\. Zhuang, Y\. Sheng, L\. Zheng, C\. Yu, J\. Gonzalez, H\. Zhang, and I\. StoicaEfficient memory management for large language model serving with PagedAttention\.InProceedings of the 29th Symposium on Operating Systems Principles,pp\. 611–626\.Cited by:[1st item](https://arxiv.org/html/2608.19920#A1.I7.i1.p1.1),[§A\.7](https://arxiv.org/html/2608.19920#A1.SS7.p2.1),[§A\.7](https://arxiv.org/html/2608.19920#A1.SS7.p3.1),[§A\.9](https://arxiv.org/html/2608.19920#A1.SS9.p3.1),[§1](https://arxiv.org/html/2608.19920#S1.p3.1),[§2](https://arxiv.org/html/2608.19920#S2.p4.1),[§3\.1](https://arxiv.org/html/2608.19920#S3.SS1.p3.1),[§3\.3](https://arxiv.org/html/2608.19920#S3.SS3.p1.1),[§3\.3](https://arxiv.org/html/2608.19920#S3.SS3.p4.1)\.
- \[34\]A\. Lańcucki, K\. Staniszewski, P\. Nawrot, and E\. PontiInference\-time hyper\-scaling with kv cache compression\.See[5](https://arxiv.org/html/2608.19920#bib.bib84),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- H\. Larochelle, M\. Ranzato, R\. Hadsell, M\.F\. Balcan, and H\. Lin \(Eds\.\) \(2020\)H\. Larochelle, M\. Ranzato, R\. Hadsell, M\.F\. Balcan, and H\. Lin \(Eds\.\)Advances in neural information processing systems33\.Curran Associates\.Cited by:[77](https://arxiv.org/html/2608.19920#bib.bib67)\.
- Liet al\.\(2023\)S\. Li, F\. Xue, C\. Baranwal, Y\. Li, and Y\. YouSequence parallelism: long sequence training from system perspective\.InProceedings of the 61st Annual Meeting of the Association for Computational Linguistics \(ACL\),pp\. 7376–7399\.Cited by:[3rd item](https://arxiv.org/html/2608.19920#S1.I1.i3.p1.1),[§2](https://arxiv.org/html/2608.19920#S2.p4.1)\.
- Liet al\.\(2026\)W\. Li, D\. Yu, G\. Luo, Y\. Zhang, Y\. Wu, J\. Liu, Z\. Gong, Z\. Liao, F\. Chao, and R\. JiOut of the memory barrier: a highly memory\-efficient training system for LLMs with million\-token contexts\.InInt\. Conf\. Learning Representations,Cited by:[§A\.8](https://arxiv.org/html/2608.19920#A1.SS8.p1.1),[§2](https://arxiv.org/html/2608.19920#S2.p4.1)\.
- Liet al\.\(2025a\)W\. Li, Y\. Zhang, G\. Luo, D\. Yu, and R\. JiTraining long\-context LLMs efficiently via chunk\-wise optimization\.InFindings of the Association for Computational Linguistics \(ACL\),pp\. 2691–2700\.Cited by:[1st item](https://arxiv.org/html/2608.19920#A1.I7.i1.p1.1),[§2](https://arxiv.org/html/2608.19920#S2.p4.1)\.
- \[39\]X\. Li, Z\. Xing, M\. Li, L\. Qu, H\. Zhen, Y\. Yao, W\. Liu, S\. Pan, and M\. YuanKVTuner: sensitivity\-aware layer\-wise mixed\-precision KV cache quantization for efficient and nearly lossless LLM inference\.See[57](https://arxiv.org/html/2608.19920#bib.bib78),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- \[40\]Y\. Li, Y\. Huang, B\. Yang, B\. Venkitesh, A\. Locatelli, H\. Ye, T\. Cai, P\. Lewis, and D\. ChenSnapKV: LLM knows what you are looking for before generation\.See[5](https://arxiv.org/html/2608.19920#bib.bib84),pp\. 529–536\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Liet al\.\(2025b\)Y\. Li, H\. Jiang, Q\. Wu, X\. Luo, S\. Ahn, C\. Zhang, A\. Abdi, D\. Li, J\. Gao, Y\. Yang, and L\. QiuSCBench: a KV cache\-centric analysis of long\-context methods\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Liuet al\.\(2024\)H\. Liu, M\. Zaharia, and P\. AbbeelRingAttention with blockwise transformers for near\-infinite context\.InInt\. Conf\. Learning Representations,Cited by:[1st item](https://arxiv.org/html/2608.19920#A1.I7.i1.p1.1),[§2](https://arxiv.org/html/2608.19920#S2.p4.1),[footnote 6](https://arxiv.org/html/2608.19920#footnote6)\.
- \[43\]Z\. Liu, S\. Wang, S\. Cheng, Z\. Zhao, K\. Wang, X\. Zhao, J\. Demmel, and Y\. YouStarTrail: concentric ring sequence parallelism for efficient near\-infinite\-context transformer model training\.See[5](https://arxiv.org/html/2608.19920#bib.bib84),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p4.1)\.
- \[44\]Z\. Liu, J\. Yuan, H\. Jin, S\. Zhong, Z\. Xu, V\. Braverman, D\. Chen, and X\. HuKIVI: a tuning\-free asymmetric 2bit quantization for KV cache\.See[50](https://arxiv.org/html/2608.19920#bib.bib77),pp\. 32332–32344\.Cited by:[§1](https://arxiv.org/html/2608.19920#S1.p2.1),[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Loshchilov and Hutter \(2017\)I\. Loshchilov and F\. HutterDecoupled weight decay regularization\.Technical reportTechnical ReportarXiv:1711\.05101 \[cs\.LG\]\.External Links:[Link](https://arxiv.org/abs/1711.05101)Cited by:[§4](https://arxiv.org/html/2608.19920#S4.p1.2)\.
- \[46\]P\. Nawrot, A\. Lancucki, M\. Chochowski, D\. Tarjan, and E\. PontiDynamic memory compression: retrofitting LLMs for accelerated inference\.See[50](https://arxiv.org/html/2608.19920#bib.bib77),pp\. 37396–37412\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p3.1)\.
- A\. Oh, T\. Naumann, A\. Globerson, K\. Saenko, M\. Hardt, and S\. Levine \(Eds\.\) \(2023\)A\. Oh, T\. Naumann, A\. Globerson, K\. Saenko, M\. Hardt, and S\. Levine \(Eds\.\)Advances in neural information processing systems36\.Curran Associates\.Cited by:[10](https://arxiv.org/html/2608.19920#bib.bib8),[51](https://arxiv.org/html/2608.19920#bib.bib43),[76](https://arxiv.org/html/2608.19920#bib.bib69)\.
- Penget al\.\(2024\)B\. Peng, J\. Quesnelle, H\. Fan, and E\. ShippoleYaRN: efficient context window extension of large language models\.InInt\. Conf\. Learning Representations,Cited by:[§4](https://arxiv.org/html/2608.19920#S4.p1.2)\.
- Qinet al\.\(2025\)Z\. Qin, Y\. Cao, M\. Lin, W\. Hu, S\. Fan, K\. Cheng, W\. Lin, and J\. LiCAKE: cascading and adaptive KV cache eviction with layer preferences\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- R\. Salakhutdinov, Z\. Kolter, K\. Heller, A\. Weller, N\. Oliver, J\. Scarlett, and F\. Berkenkamp \(Eds\.\) \(2024\)R\. Salakhutdinov, Z\. Kolter, K\. Heller, A\. Weller, N\. Oliver, J\. Scarlett, and F\. Berkenkamp \(Eds\.\)International conference on machine learning41\.Proceedings of Machine Learning Research\.Cited by:[11](https://arxiv.org/html/2608.19920#bib.bib9),[13](https://arxiv.org/html/2608.19920#bib.bib13),[44](https://arxiv.org/html/2608.19920#bib.bib37),[46](https://arxiv.org/html/2608.19920#bib.bib40),[63](https://arxiv.org/html/2608.19920#bib.bib53)\.
- \[51\]T\. Schick, J\. Dwivedi\-Yu, R\. Dessi, R\. Raileanu, M\. Lomeli, E\. Hambro, L\. Zettlemoyer, N\. Cancedda, and T\. ScialomToolformer: language models can teach themselves to use tools\.See[47](https://arxiv.org/html/2608.19920#bib.bib82),Cited by:[§1](https://arxiv.org/html/2608.19920#S1.p1.1)\.
- \[52\]J\. Shah, G\. Bikshandi, Y\. Zhang, V\. Thakkar, P\. Ramani, and T\. DaoFlashAttention\-3: fast and accurate attention with asynchrony and low\-precision\.See[22](https://arxiv.org/html/2608.19920#bib.bib83),pp\. 68658–68685\.Cited by:[4th item](https://arxiv.org/html/2608.19920#A1.I8.i4.p1.1),[§2](https://arxiv.org/html/2608.19920#S2.p2.1)\.
- \[53\]N\. Shang, L\. Zhang, S\. Wang, G\. Zhang, G\. Lopez, F\. Yang, W\. Chen, and M\. YangLongRoPE2: near\-lossless LLM context window scaling\.See[57](https://arxiv.org/html/2608.19920#bib.bib78),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p4.1)\.
- Shazeer \(2019\)N\. ShazeerFast transformer decoding: one write\-head is all you need\.Technical reportTechnical ReportarXiv:1911\.02150 \[cs\.NE\]\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- \[55\]Y\. Sheng, L\. Zheng, B\. Yuan, Z\. Li, M\. Ryabinin, D\. Fu, Z\. Xie, B\. Chen, C\. Barrett, J\. Gonzalez, P\. Liang, C\. Re, I\. Stoica, and C\. ZhangFlexGen: high\-throughput generative inference of large language models with a single GPU\.See[32](https://arxiv.org/html/2608.19920#bib.bib76),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- \[56\]A\. Shutova, V\. Malinovskii, V\. Egiazarian, D\. Kuznedelev, D\. Mazur, S\. Nikita, I\. Ermakov, and D\. AlistarhCache me if you must: adaptive key\-value quantization for large language models\.See[57](https://arxiv.org/html/2608.19920#bib.bib78),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- A\. Singh, M\. Fazel, D\. Hsu, S\. Lacoste\-Julien, F\. Berkenkamp, T\. Maharaj, K\. Wagstaff, and J\. Zhu \(Eds\.\) \(2025\)A\. Singh, M\. Fazel, D\. Hsu, S\. Lacoste\-Julien, F\. Berkenkamp, T\. Maharaj, K\. Wagstaff, and J\. Zhu \(Eds\.\)International conference on machine learning42\.Proceedings of Machine Learning Research\.Cited by:[39](https://arxiv.org/html/2608.19920#bib.bib32),[53](https://arxiv.org/html/2608.19920#bib.bib45),[56](https://arxiv.org/html/2608.19920#bib.bib48),[60](https://arxiv.org/html/2608.19920#bib.bib52)\.
- Staniszewski and Lańcucki \(2026\)K\. Staniszewski and A\. LańcuckiKV cache transform coding for compact storage in LLM inference\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Suet al\.\(2024\)J\. Su, M\. Ahmed, Y\. Lu, S\. Pan, W\. Bo, and Y\. LiuRoFormer: enhanced transformer with rotary position embedding\.Neurocomputing568\(C\)\.Cited by:[§4](https://arxiv.org/html/2608.19920#S4.p1.2)\.
- \[60\]H\. Sun, L\. Chang, W\. Bao, S\. Zheng, N\. Zheng, X\. Liu, H\. Dong, Y\. Chi, and B\. ChenShadowKV: KV cache in shadows for high\-throughput long\-context LLM inference\.See[57](https://arxiv.org/html/2608.19920#bib.bib78),pp\. 57355–57373\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- \[61\]Y\. Sun, L\. Dong, Y\. Zhu, S\. Huang, W\. Wang, S\. Ma, Q\. Zhang, Z\. Wang, and F\. WeiYou only cache once: decoder\-decoder architectures for language models\.See[22](https://arxiv.org/html/2608.19920#bib.bib83),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p3.1)\.
- Tanget al\.\(2025\)H\. Tang, Y\. Lin, J\. Lin, Q\. Han, D\. Ke, S\. Hong, Y\. Yao, and G\. WangRazorAttention: efficient KV cache compression through retrieval heads\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- \[63\]J\. Tang, Y\. Zhao, K\. Zhu, G\. Xiao, B\. Kasikci, and S\. HanQuest: query\-aware sparsity for efficient long\-context llm inference\.See[50](https://arxiv.org/html/2608.19920#bib.bib77),pp\. 47901–47911\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- \[64\]A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. Gomez, L\. Kaiser, and I\. PolosukhinAttention is all you need\.See[6](https://arxiv.org/html/2608.19920#bib.bib79),pp\. 6000–6010\.Cited by:[§3\.1](https://arxiv.org/html/2608.19920#S3.SS1.p1.1)\.
- Wanget al\.\(2025\)G\. Wang, S\. Upasani, C\. Wu, D\. Gandhi, J\. Li, C\. Hu, B\. Li, and U\. ThakkerLLMs know what to drop: self\-attention guided KV cache eviction for efficient long\-context inference\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- \[66\]J\. Wei, X\. Wang, D\. Schuurmans, M\. Bosma, B\. Ichter, F\. Xia, W\. Chi, Q\. Le, and D\. ZhouChain\-of\-thought prompting elicits reasoning in large language models\.See[31](https://arxiv.org/html/2608.19920#bib.bib81),Cited by:[§1](https://arxiv.org/html/2608.19920#S1.p1.1)\.
- \[67\]T\. Wu, Y\. Zhao, and Z\. ZhengAn efficient recipe for long context extension via middle\-focused positional encoding\.See[22](https://arxiv.org/html/2608.19920#bib.bib83),Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p4.1)\.
- Xiaoet al\.\(2025\)G\. Xiao, J\. Tang, J\. Zuo, J\. Guo, S\. Yang, H\. Tang, Y\. Fu, and S\. HanDuoAttention: efficient long\-context LLM inference with retrieval and streaming heads\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Xiaoet al\.\(2024\)G\. Xiao, Y\. Tian, B\. Chen, S\. Han, and M\. LewisEfficient streaming language models with attention sinks\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1),[1st item](https://arxiv.org/html/2608.19920#S4.I1.i1.p1.1),[footnote 12](https://arxiv.org/html/2608.19920#footnote12)\.
- Yanget al\.\(2025a\)A\. Yang, A\. Li, B\. Yang, B\. Zhang, B\. Hui, B\. Zheng, B\. Yu, C\. Gao, C\. Huang, C\. Lv, C\. Zheng, D\. Liu, F\. Zhou, F\. Huang, F\. Hu, H\. Ge, H\. Wei, H\. Lin, J\. Tang, J\. Yang, J\. Tu, J\. Zhang, J\. Yang, J\. Yang, J\. Zhou, J\. Zhou, J\. Lin, K\. Dang, K\. Bao, K\. Yang, L\. Yu, L\. Deng, M\. Li, M\. Xue, M\. Li, P\. Zhang, P\. Wang, Q\. Zhu, R\. Men, R\. Gao, S\. Liu, S\. Luo, T\. Li, T\. Tang, W\. Yin, X\. Ren, X\. Wang, X\. Zhang, X\. Ren, Y\. Fan, Y\. Su, Y\. Zhang, Y\. Zhang, Y\. Wan, Y\. Liu, Z\. Wang, Z\. Cui, Z\. Zhang, Z\. Zhou, and Z\. QiuQwen3 technical report\.Technical reportTechnical ReportarXiv:2505\.09388 \[cs\.CL\]\.Cited by:[§4](https://arxiv.org/html/2608.19920#S4.p1.2)\.
- Yanget al\.\(2025b\)S\. Yang, J\. Guo, H\. Tang, Q\. Hu, G\. Xiao, J\. Tang, Y\. Lin, Z\. Liu, Y\. Lu, and S\. HanLServe: efficient long\-sequence LLM serving with unified sparse attention\.InProceedings of the 8th MLSys Conference,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Yeet al\.\(2025\)Z\. Ye, L\. Chen, R\. Lai, W\. Lin, Y\. Zhang, S\. Wang, T\. Chen, B\. Kasikci, V\. Grover, A\. Krishnamurthy, and L\. CezeFlashInfer: efficient and customizable attention engine for LLM inference serving\.Technical reportTechnical ReportarXiv:2501\.01005 \[cs\.DC\]\.Cited by:[4th item](https://arxiv.org/html/2608.19920#A1.I8.i4.p1.1),[§A\.6](https://arxiv.org/html/2608.19920#A1.SS6.p2.1),[2nd item](https://arxiv.org/html/2608.19920#S1.I1.i2.p1.1),[§2](https://arxiv.org/html/2608.19920#S2.p2.1),[§3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1.p3.1),[§3\.3](https://arxiv.org/html/2608.19920#S3.SS3.p2.1)\.
- Yenet al\.\(2025\)H\. Yen, T\. Gao, M\. Hou, K\. Ding, D\. Fleischer, P\. Izsak, M\. Wasserblat, and D\. ChenHelmet: how to evaluate long\-context models effectively and thoroughly\.InInt\. Conf\. Learning Representations,Cited by:[§A\.3\.1](https://arxiv.org/html/2608.19920#A1.SS3.SSS1.p1.1),[§4](https://arxiv.org/html/2608.19920#S4.p1.1),[§4](https://arxiv.org/html/2608.19920#S4.p2.1)\.
- Yuanet al\.\(2025\)J\. Yuan, H\. Gao, D\. Dai, J\. Luo, L\. Zhao, Z\. Zhang, Z\. Xie1, Y\. Wei, L\. Wang, Z\. Xiao, Y\. Wang, C\. Ruan, M\. Zhang, W\. Liang, and W\. ZengNative sparse attention: hardware\-aligned and natively trainable sparse attention\.Technical reportTechnical ReportarXiv:2502\.11089 \[cs\.CL\]\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p3.1)\.
- Yuanet al\.\(2026\)Z\. Yuan, H\. Sun, L\. Sun, and Y\. YeMegaTrain: full precision training of 100b\+ parameter large language models on a single GPU\.Technical reportTechnical ReportarXiv:2604\.05091 \[cs\.CL\]\.Cited by:[4th item](https://arxiv.org/html/2608.19920#A1.I7.i4.p1.1),[§5](https://arxiv.org/html/2608.19920#S5.p2.1)\.
- \[76\]Z\. Z\., Y\. Sheng, T\. Zhou\., T\. Chen, L\. Zheng, R\. Cai, Z\. Song, Y\. Tian, C\. Re, C\. Barrett, Z\. Wang, and B\. ChenH2O: heavy\-hitter oracle for efficient generative inference of large language models\.See[47](https://arxiv.org/html/2608.19920#bib.bib82),pp\. 34661–34710\.Cited by:[2nd item](https://arxiv.org/html/2608.19920#A1.I8.i2.p1.1),[2nd item](https://arxiv.org/html/2608.19920#S1.I1.i2.p1.1),[§1](https://arxiv.org/html/2608.19920#S1.p3.1),[§2](https://arxiv.org/html/2608.19920#S2.p1.1),[§3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1.p1.1),[§3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1.p2.1),[§3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1.p3.1),[§3\.1](https://arxiv.org/html/2608.19920#S3.SS1.p3.1),[§3\.3](https://arxiv.org/html/2608.19920#S3.SS3.p1.1),[§3](https://arxiv.org/html/2608.19920#S3.p1.1),[3rd item](https://arxiv.org/html/2608.19920#S4.I1.i3.p1.1),[§4](https://arxiv.org/html/2608.19920#S4.p5.1)\.
- \[77\]M\. Zaheer, G\. Guruganesh, A\. Dubey, J\. Ainslie, C\. Alberti, S\. Ontanon, P\. Pham, A\. Ravula, Q\. Wang, L\. Yang, and A\. AhmedBig Bird: transformers for longer sequences\.See[35](https://arxiv.org/html/2608.19920#bib.bib80),pp\. 17283–17297\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- Zandiehet al\.\(2026\)A\. Zandieh, M\. Daliri, M\. Hadian, and V\. MirrokniTurboQuant: online vector quantization with near\-optimal distortion rate\.InInt\. Conf\. Learning Representations,Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- \[79\]T\. Zhang, J\. Yi, Z\. Xu, and A\. ShrivastavaKV cache is 1 bit per channel: efficient large language model inference with coupled quantization\.See[22](https://arxiv.org/html/2608.19920#bib.bib83),pp\. 3304–3331\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p1.1)\.
- \[80\]Z\. Zhang, R\. Chen, S\. Liu, Z\. Yao, O\. Ruwase, B\. Chen, X\. Wu, and Z\. WangFound in the middle: how language models use long contexts better via plug\-and\-play positional encoding\.See[22](https://arxiv.org/html/2608.19920#bib.bib83),pp\. 60755–60775\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p4.1)\.
- Zhanget al\.\(2024\)Z\. Zhang, S\. Liu, R\. Chen, B\. Kailkhura, B\. Chen, and A\. WangQ\-Hitter: a better token oracle for efficient LLM inference via sparse\-quantized KV cache\.InProceedings of Machine Learning and Systems \(MLSys\),P\. Gibbons, G\. Pekhimenko, and C\. D\. Sa \(Eds\.\),Vol\.6,pp\. 381–394\.Cited by:[§A\.5\.1](https://arxiv.org/html/2608.19920#A1.SS5.SSS1.p2.1),[Table 5](https://arxiv.org/html/2608.19920#A1.T5),[§2](https://arxiv.org/html/2608.19920#S2.p1.1),[§3\.3](https://arxiv.org/html/2608.19920#S3.SS3.p1.1)\.
- Zhaoet al\.\(2025\)Y\. Zhao, J\. Huang, J\. Hu, X\. Wang, Y\. Mao, D\. Zhang, Z\. Jiang, Z\. Wu, B\. Ai, A\. Wang, W\. Zhou, and Y\. ChenSWIFT: a scalable lightweight infrastructure for fine\-tuning\.InProceedings of the 39th Conference on Artificial Intelligence \(AAAI\),pp\. 29733–29735\.Cited by:[1st item](https://arxiv.org/html/2608.19920#A1.I7.i1.p1.1),[§2](https://arxiv.org/html/2608.19920#S2.p4.1),[1st item](https://arxiv.org/html/2608.19920#S4.I2.i1.p1.1)\.
- \[83\]L\. Zheng, L\. Yin, Z\. Xie, C\. Sun, J\. Huang, C\. Yu, S\. Cao, C\. Kozyrakis, I\. Stoica, J\. Gonzalez, C\. Barrett, and Y\. ShengSGLang: efficient execution of structured language model programs\.See[22](https://arxiv.org/html/2608.19920#bib.bib83),pp\. 62557–62583\.Cited by:[§A\.9](https://arxiv.org/html/2608.19920#A1.SS9.p3.1),[§2](https://arxiv.org/html/2608.19920#S2.p4.1)\.
- Zhouet al\.\(2026\)C\. Zhou, K\. Liu, Y\. Zhou, Q\. Qiao, J\. Gao, H\. Zhang, I\. Lu, N\. Ho, L\. Li, A\. Lei, C\. Cheng, S\. Chiang, Y\. Zeng, D\. Zhang, R\. Yang, K\. Chen, A\. Chen, P\. Ma, W\. Zhang, and C\. JinLongStraw: long\-context RL beyond 2M tokens under a fixed GPU budget\.Technical reportTechnical ReportarXiv:2607\.14952 \[cs\.LG\]\.Cited by:[§2](https://arxiv.org/html/2608.19920#S2.p4.1)\.

## Appendix AAppendix

### A\.1Notation\. Definitions

Here, we add some details missing in the main text\.

#### A\.1\.1Buffer Read/Write Access by𝚝𝚘𝚛𝚌𝚑\.𝚜𝚌𝚊𝚝𝚝𝚎𝚛,𝚝𝚘𝚛𝚌𝚑\.𝚐𝚊𝚝𝚑𝚎𝚛\\mathtt\{torch\.scatter\},\\mathtt\{torch\.gather\}

These linear operations are defined in[https://docs\.pytorch\.org/docs/2\.12/generated/torch\.Tensor\.scatter\_\.html](https://docs.pytorch.org/docs/2.12/generated/torch.Tensor.scatter_.html)\.𝚜𝚌𝚊𝚝𝚝𝚎𝚛​\_\\mathtt\{scatter\\\_\}assigns values to entries in certain positions,𝚐𝚊𝚝𝚑𝚎𝚛\\mathtt\{gather\}extracts values of entries at certain positions\. For reference,𝚊𝚛𝚛\.𝚜𝚌𝚊𝚝𝚝𝚎𝚛​\_​\(𝚍𝚒𝚖,𝚒𝚗𝚍𝚎𝚡,𝚜𝚛𝚌\)\\mathtt\{arr\.scatter\\\_\(dim,index,src\)\}requires𝚊𝚛𝚛\.𝚗𝚍𝚒𝚖==𝚒𝚗𝚍𝚎𝚡\.𝚗𝚍𝚒𝚖==𝚜𝚛𝚌\.𝚗𝚍𝚒𝚖\\mathtt\{arr\.ndim==index\.ndim==src\.ndim\}and𝚒𝚗𝚍𝚎𝚡\.𝚜𝚑𝚊𝚙𝚎==𝚜𝚛𝚌\.𝚜𝚑𝚊𝚙𝚎\\mathtt\{index\.shape==src\.shape\}, whereas𝚊𝚛𝚛\.𝚜𝚑𝚊𝚙𝚎\\mathtt\{arr\.shape\}can differ on position𝚍𝚒𝚖\\mathtt\{dim\}\. Say that𝚊𝚛𝚛\.𝚗𝚍𝚒𝚖==𝟹\\mathtt\{arr\.ndim==3\}\. Then:

𝚊𝚛𝚛⁡\[𝚒𝚗𝚍𝚎𝚡⁡\[𝚒,𝚓,𝚔\],𝚓,𝚔\]=𝚜𝚛𝚌⁡\[𝚒,𝚓,𝚔\]\|dim==0𝚊𝚛𝚛⁡\[𝚒,𝚒𝚗𝚍𝚎𝚡⁡\[𝚒,𝚓,𝚔\],𝚔\]=𝚜𝚛𝚌⁡\[𝚒,𝚓,𝚔\]\|dim==1𝚊𝚛𝚛⁡\[𝚒,𝚓,𝚒𝚗𝚍𝚎𝚡⁡\[𝚒,𝚓,𝚔\]\]=𝚜𝚛𝚌⁡\[𝚒,𝚓,𝚔\]\|dim==2\\begin\{split\}\\mathtt\{arr\[index\[i,j,k\],j,k\]=src\[i,j,k\]\}\\quad&\|\\;dim==0\\\\ \\mathtt\{arr\[i,index\[i,j,k\],k\]=src\[i,j,k\]\}\\quad&\|\\;dim==1\\\\ \\mathtt\{arr\[i,j,index\[i,j,k\]\]=src\[i,j,k\]\}\\quad&\|\\;dim==2\\\\ \\end\{split\}Also, if𝚛𝚎𝚜=𝚊𝚛𝚛\.𝚐𝚊𝚝𝚑𝚎𝚛⁡\(𝚍𝚒𝚖,𝚒𝚗𝚍𝚎𝚡\)\\mathtt\{res=arr\.gather\(dim,index\)\}, then:

𝚛𝚎𝚜⁡\[𝚒,𝚓,𝚔\]=𝚊𝚛𝚛⁡\[𝚒𝚗𝚍𝚎𝚡⁡\[𝚒,𝚓,𝚔\],𝚓,𝚔\]\|dim==0𝚛𝚎𝚜⁡\[𝚒,𝚓,𝚔\]=𝚊𝚛𝚛⁡\[𝚒,𝚒𝚗𝚍𝚎𝚡⁡\[𝚒,𝚓,𝚔\],𝚔\]\|dim==1𝚛𝚎𝚜⁡\[𝚒,𝚓,𝚔\]=𝚊𝚛𝚛⁡\[𝚒,𝚓,𝚒𝚗𝚍𝚎𝚡⁡\[𝚒,𝚓,𝚔\]\]\|dim==2\\begin\{split\}\\mathtt\{res\[i,j,k\]=arr\[index\[i,j,k\],j,k\]\}\\quad&\|\\;dim==0\\\\ \\mathtt\{res\[i,j,k\]=arr\[i,index\[i,j,k\],k\]\}\\quad&\|\\;dim==1\\\\ \\mathtt\{res\[i,j,k\]=arr\[i,j,index\[i,j,k\]\]\}\\quad&\|\\;dim==2\\\\ \\end\{split\}In our use case, we apply these operations to 4D arrays with𝚍𝚒𝚖=𝟸\\mathtt\{dim=2\}, so that𝚊𝚛𝚛,𝚜𝚛𝚌\\mathtt\{arr,src\}are 4D, but𝚒𝚗𝚍𝚎𝚡\\mathtt\{index\}is 3D\. This is done by broadcasting𝚒𝚗𝚍𝚎𝚡\\mathtt\{index\}along the final axis:𝚒𝚗𝚍𝚎𝚡\.𝚞𝚗𝚜𝚚𝚞𝚎𝚎𝚣𝚎⁡\(−𝟷\)\.𝚎𝚡𝚝𝚎𝚗𝚍⁡\(−𝟷,−𝟷,−𝟷,𝚍𝚑\)\\mathtt\{index\.unsqueeze\(\-1\)\.extend\(\-1,\-1,\-1,d\_\{h\}\)\}\.

### A\.2Key\-Value Cache Policies

In this section, we present additional details about KV cache policies used in our experiments\.

#### A\.2\.1Policy𝚜𝚖𝚊𝚛𝚝​\_​𝚕𝚊𝚜𝚝𝚛𝚎𝚌\\mathtt\{smart\\\_lastrec\}\(slr\)

Recall the simple𝚕𝚊𝚜𝚝𝚛𝚎𝚌\\mathtt\{lastrec\}\(lr\) policy from Section[4](https://arxiv.org/html/2608.19920#S4), which keeps the last recentNC−βN\_\{C\}\-\\betaand firstβ\\betatokens in the cache\. One drawback of this policy is thatβ\\betais fixed, while prompts often start with initial important control information ofvariablelength\. Our𝚜𝚖𝚊𝚛𝚝​\_​𝚕𝚊𝚜𝚝𝚛𝚎𝚌\\mathtt\{smart\\\_lastrec\}policy is defined in terms of a regular expression for the end of this task control prefix, as well as some maximum prefix lengthMprefix<NCM\_\{\\text\{prefix\}\}<N\_\{C\}\. When processing the first \(prefill\) chunk, we search for the first match \(separately for each batch positionbb\)\. If this results in a prefix lengthM⁡\(b\)≤MprefixM\(b\)\\leq M\_\{\\text\{prefix\}\}, this is used, otherwiseM⁡\(b\)=MprefixM\(b\)=M\_\{\\text\{prefix\}\}\. For subsequent chunks and token positionst≥NCt\\geq N\_\{C\}, the policy isπ⁡\(b,h,t\)=M⁡\(b\)\+mod⁡\(t−NC,NC−M⁡\(b\)\)\\pi\(b,h,t\)=M\(b\)\+\\mathrm\{mod\}\(t\-N\_\{C\},N\_\{C\}\-M\(b\)\)\.

We also implemented a generalization where arange\[M0​\(b\),M1​\(b\)\)\[M\_\{0\}\(b\),M\_\{1\}\(b\)\)is protected from eviction\.111111We allow forM1​\(b\)<M0​\(b\)M\_\{1\}\(b\)<M\_\{0\}\(b\), in which case the protected area is\[0,M1​\(b\)\)∪\[M0​\(b\),NC\)\[0,M\_\{1\}\(b\)\)\\cup\[M\_\{0\}\(b\),N\_\{C\}\)\.Here,M0​\(b\)M\_\{0\}\(b\)is chosen as the position of the first non\-padding token, andM1​\(b\)M\_\{1\}\(b\)is chosen as above\. The idea is that initial padding tokens do not carry information and should not be attended to, so they can be evicted as soon as the cache is full\. In this variant, if all tokens in the prefill chunk are padding for somebb, the search for\[M0​\(b\),M1​\(b\)\)\[M\_\{0\}\(b\),M\_\{1\}\(b\)\)is shifted to subsequent chunks\. The prefix case above is obtained withM0​\(b\)=0M\_\{0\}\(b\)=0\. Surprisingly, in our experiments, the general variant did not improve over the prefix variant,121212This needs further analysis\. Maybe some initial padding tokens are used as so\-called ”attention sinks”\[[69](https://arxiv.org/html/2608.19920#bib.bib59)\]\.so that𝚜𝚖𝚊𝚛𝚝​\_​𝚕𝚊𝚜𝚝𝚛𝚎𝚌\\mathtt\{smart\\\_lastrec\}in Section[4](https://arxiv.org/html/2608.19920#S4)is the prefix variant throughout\.

### A\.3Long\-Context Benchmarks

#### A\.3\.1Helmet

Our training and evaluation suite is derived from Helmet\[[73](https://arxiv.org/html/2608.19920#bib.bib64)\], a benchmark designed for inference\-time evaluation of long\-context language models\. Helmet covers five capability categories across five context\-length scales \(8k to 128k tokens\)\. Helmet provides only a small number of instances per task\. For supervised fine\-tuning, we adapted it as follows\.

##### Instance Construction and Task Scope\.

We reconstruct each task from its upstream source data, following the original Helmet logic for forming contexts and controlling sequence length\. We focus on the 64k and 128k context\-length settings and include 10 tasks spanning five capability categories\. Table[3](https://arxiv.org/html/2608.19920#A1.T3)provides a summary of each task\.

##### Split Separation\.

For each task we produce two non\-overlapping partitions\. The instances used in the original Helmet evaluation are reserved as a held\-out*evaluation \(test\) set*\. All remaining instances are collected into a*development set*used for training\. For RAG tasks, we sample a single depth \(difficulty\) variant per query in the development set to prevent the model from memorising the same question paired with multiple distractor configurations\. For InfiniteBench QA/MC, we remove in\-context demonstrations from development instances because the demonstrations are drawn from the same small pool and would otherwise create data leakage during training\.

##### Evaluation Metrics\.

Our evaluation metrics are taken from the code coming with Helmet \([https://github\.com/princeton\-nlp/HELMET/blob/main/utils\.py](https://github.com/princeton-nlp/HELMET/blob/main/utils.py)\)\. Here, the output is a string generated by the model, the target is a string or a list of strings:

- •SubEM: Depending on the dataset, the target can be a list of strings\. We normalize the output \(strip whitespace, quotes, and common phrases such as “Answer:”\), map output and target\(s\) to lower\-case\. The value is 1 if at least one of the targets is a substring in the output, 0 otherwise\.
- •Accuracy: The target is a numerical value\. We extract all numerical values from the output and find the value which occurs most often \(with ties, the value is chosen which appears first\)\. We then use exact match between this value and the target\.
- •ROUGE\-F1: The target is a string\. We compute ROUGE\-N precision/recall/F1 between output \(after normalization as in SubEM\) and target\.

CategoryIDSourceMetricDevEvalRAGnqNatural QuestionsSubEM893600trivia\_qaTriviaQASubEM876600pop\_qaPopQASubEM192600hotpot\_qaHotpotQASubEM787300Many\-shot ICLtrec\_coarseTRECAccuracy1000500nluSNIPS NLUAccuracy2094500clinc150CLINC150Accuracy2600500Long\-doc QAinf\_qaInfiniteBench QAROUGE\-F1251100inf\_mcInfiniteBench MCAccuracy129100Synthetic Recalljson\_kvJSON\-KVSubEM500100Table 3:Overview of the 10 Helmet tasks\.*Dev*and*Eval*denote the number of instances in the training and evaluation partitions, respectively, at a single context\-length setting\.In what follows, we provide details for the different tasks\.

##### Retrieval\-augmented Generation \(RAG\)\.

Each instance consists of a natural\-language question, one or more gold passages, and a pool of hard\-negative distractors\. The context is formed by inserting gold passages at a random depth among the distractors, and the whole context is truncated to the target length\.Natural Questions\(NQ\) uses real Google search queries paired with Wikipedia passages\.TriviaQAuses trivia questions authored with independently collected evidence\.PopQAfocuses on long\-tail, entity\-centric questions generated from Wikidata triples; we filter the evaluation set to queries whose subject entities fall below a popularity threshold of 3\.HotpotQArequires multi\-hop reasoning across two gold passages\. Each of the six depth variants of a query places the gold passage\(s\) at a different relative position in the distractor pool\. All four tasks are evaluated with substring exact match \(SubEM\)\.

##### Many\-shot In\-context Learning\.

Three intent/question\-type classification datasets test the model’s ability to exploit many labelled demonstrations placed entirely within the context window\. Unlike most other tasks, where demonstrations serve only as formatting guides, here they carry essential semantic information: each demonstration encodes a \(text, ordinal\-label\) pair, and the label\-to\-class mapping is only recoverable by reading the demonstrations\.TREC Coarsehas 6 question\-type classes;SNIPS NLUhas 68 intent classes;CLINC150has 151 intent classes\. The number of demonstrations is calibrated to fill the context window while maintaining an approximately balanced class distribution across shots\.

##### Long\-document QA\.

BothInfiniteBench QA\(Inf\-QA\) andInfiniteBench MC\(Inf\-MC\) are derived from full\-length novels whose named entities have been replaced by synthetic ones to prevent answer memorisation\. The source document typically exceeds the target context window, so it is truncated to fit\. Inf\-QA is open\-ended, evaluated by ROUGE\-F1; Inf\-MC is a 4\-way multiple\-choice variant evaluated by accuracy\.

##### Synthetic Recall\.

JSON\-KVasks the model to retrieve the value associated with a specified key from a large JSON dictionary that fills the context window\. This task is evaluated with SubEM and serves as controlled probes of the model’s ability to locate and copy specific information over very long spans\.

### A\.4Gradient Computation

In this section, we provide additional details about our long\-context gradient computation method from Section[3\.2](https://arxiv.org/html/2608.19920#S3.SS2)\.

#### A\.4\.1Chunks, Cells, CPU Offloading

Recall that for long\-context inference or fine\-tuning with cache lengthNCN\_\{C\}, we split a sequence of lengthN\>NCN\>N\_\{C\}into1\+⌈\(N−NC\)/S⌉1\+\\lceil\(N\-N\_\{C\}\)/S\\rceilchunks, the first \(prefill\) chunk of lengthNCN\_\{C\}, subsequent chunks of lengthS<NCS<N\_\{C\}\. The chunk sizeSSis chosen according to a latency\-vs\-accuracy trade\-off, it is in general much shorter thanNCN\_\{C\}\(Section[3\.3](https://arxiv.org/html/2608.19920#S3.SS3)\)\. We also group chunks intocells\. The first cell consists of the prefill chunk alone, subsequent cells groupk=⌊α​NC/S⌋k=\\lfloor\\alpha N\_\{C\}/S\\rfloorofSS\-length chunks, whereα\>0\\alpha\>0is a hyperparameter which defaults toα=1\\alpha=1\.

In our gradient computation method,PyTorch autogradis run on cells\. WhileSSis chosen also with accuracy in mind, the choice ofkkandα\\alphais determined by efficiency \(both runtime and memory\) only\. The idea is that theautogradGPU memory requirements are on the order of one KV cache buffer, if we exploit the linear recurrence \(see Section[3\.2](https://arxiv.org/html/2608.19920#S3.SS2)and below\)\. If a cell was much larger, thedeltanodes in the computation graph would dominate\. Theα\\alphaparameter is adjusted so to not run out of memory\. Empirically, a smaller chunk sizeSSnecessitates a smallerα\\alpha, likely due to overhead inautograd\(for constantα\\alpha, smallerSSmeans largerkk, so a larger computation graph\)\. In our experiments, we choseα=1\\alpha=1forS=2048S=2048,α=0\.75\\alpha=0\.75forS=1024S=1024, andα=0\.1\\alpha=0\.1forS=128S=128, for a cache length ofNC=32768N\_\{C\}=32768and 40 GB of GPU memory\.

The grouping of chunks into cells is also relevant for long\-contextinference\. Recall thatLLdenotes the number of layers of our model, and layer inputs are of shape\(B,N,d\)\(B,N,d\), whered=Hq⋅dhd=H\_\{q\}\\cdot d\_\{h\}is the model embedding dimension\. For largeNN, not even the inputs to one layer can be kept in GPU memory\. This has implications for how inference is computed and what can be offloaded to CPU memory at which point during the process \(our library supports CPU offloading of KV cache buffers during the forward passes, as well as CPU offloading of weights during the backward pass, but this is not used in the experiments reported here\)\.

Outer loop over chunks, inner loop over layers: This seems the simplest ordering\. Also, layer inputs can be kept in GPU memory, with outputs overwriting inputs\. A major drawback is that either all KV cache buffers need to be kept in GPU memory, or they need to be read from and written back to CPU memory frequently\. When KV cache buffers are stored in quantized form, the quantization computation can be considerable\. For long\-context inference, this ordering is not suitable\.

Outer loop over layers, inner loop over chunks: With this ordering, KV cache buffers \(and even model weights\) can be offloaded to CPU for all but the currently active layer\. A drawback of this ordering is that layer inputs and outputs cannot be kept in GPU memory in total, so need to be offloaded eventually\. Still, this ordering is much better than the previous one\.

Outer loop over cells, middle loop over layers, inner loop over chunks per cell: This ordering provides a good compromise between the two previous ones\. Layer inputs and outputs can be kept in GPU memory, while KV cache buffers can be quantized and/or offloaded to CPU, which happens much less frequently\. We use this ordering in our implementation, also because CPU offloading of KV cache buffers is needed during gradient computation, so we can just reuse this code\.

#### A\.4\.2Summary of Method

Here, we present a detailed summary of our gradient computation technique\. Recall that we process a batch ofBBsequences of token lengthNN, and that caches in each layer have lengthNCN\_\{C\}\. ForN≤NCN\\leq N\_\{C\}, our method reduces to standard training, so assume thatN\>NCN\>N\_\{C\}\. For simplicity, we assume that caches in all layers have the same lengthNCN\_\{C\}\. This is easy to relax, and our implementation does so\.

At the top level, our method runs a forward pass followed by a backward pass, just like standard code\. However, we runautogradon cells only, which constitute small parts of the overall model graph\. As with activation checkpointing\[[26](https://arxiv.org/html/2608.19920#bib.bib23)\], this means we need to run forward passes over the model three times \(instead of just once\)\. The first two passes are run in non\-autograd mode, checkpointing \(so called\) boundary information to CPU memory\. The third pass is part of the autograd runs on cells, which consume boundary information as inputs\.

We need some notation\. Letℒ\\mathcal\{L\}denote the training loss for the current batch\. Layers are indexed byl=0,…,L−1l=\{0\},\\dots,\{L\-1\}, cells byc=0,…,Ncells−1c=\{0\},\\dots,\{N\_\{\\text\{cells\}\}\-1\}\.𝑿l\\boldsymbol\{X\}\_\{l\}are inputs to layerll, of shape\(B,N,d\)\(B,N,d\), and𝑿L\\boldsymbol\{X\}\_\{L\}is the top layer output\. Moreover,𝑿l,c\\boldsymbol\{X\}\_\{l,c\}denotes the slice of𝑿l\\boldsymbol\{X\}\_\{l\}along axis 1 corresponding to the cell\. Finally, let𝑲l,c\\boldsymbol\{K\}\_\{l,c\}denote the KV cache buffers131313These are really two tensors \(keys, values\), but we concatenate them into one for notational simplicity\.of shape\(2,B,Hk,NC,dh\)\(2,B,H\_\{k\},N\_\{C\},d\_\{h\}\)at the input of cellc\>0c\>0in layerll\.

Forward pass 1\. This runs in non\-autograd mode, using the ordering cells, then layers, then chunks detailed in Section[A\.4\.1](https://arxiv.org/html/2608.19920#A1.SS4.SSS1)\. Alongside:

- •Store KV cache replay log in each layer, containing the decisions\{π⁡\(b,h,t\)\}\\\{\\pi\(b,h,t\)\\\}\.
- •Checkpoint layer inputs𝑿l\\boldsymbol\{X\}\_\{l\}to CPU memory for each layerl=0,…,L−1l=\{0\},\\dots,\{L\-1\}\. Our implementation allows to quantize them in order to save CPU memory and CPU\-GPU transfer time, but this is not activated in our experiments, since the time is subdominant\. We also checkpoint top layer outputs𝑿L\\boldsymbol\{X\}\_\{L\}\.

Backward pass\. This runs backwards over layers\. In each step, we compute gradients for weights in layerllusing activation checkpointing\. We start with computing head gradients∂ℒ/∂𝑿L\\partial\\mathcal\{L\}/\\partial\\boldsymbol\{X\}\_\{L\}based on top layer outputs𝑿L\\boldsymbol\{X\}\_\{L\}, writing them to CPU \(in fact, head gradients∂ℒ/∂𝑿l\\partial\\mathcal\{L\}/\\partial\\boldsymbol\{X\}\_\{l\}overwrite𝑿l\\boldsymbol\{X\}\_\{l\}on CPU\)\. Next, we iterate over layersl=L−1,…,0l=\{L\-1\},\\dots,\{0\}:

- •Forward pass 2 for layerll\(non\-autograd mode\)\. Runs over cellsc=1,…,Ncells−1c=\{1\},\\dots,\{N\_\{\\text\{cells\}\}\-1\}, computing the cache buffers𝑲l,c\\boldsymbol\{K\}\_\{l,c\}and storing them to CPU\. These are quantized in order to save CPU memory and CPU\-GPU transfer time, and we overwrite the cache buffer checkpoints from previous layerl\+1l\+1\. Cache decisions are replayed from the log\. Note that we could checkpointall𝑲l,c\\boldsymbol\{K\}\_\{l,c\}during forward pass 1, but this would requireLLtimes more CPU memory, and the extra time for forward pass 2 is subdominant\. Also note that we load inputs𝑿l,c\\boldsymbol\{X\}\_\{l,c\}to forward pass 2 to GPU cell by cell: the whole𝑿l\\boldsymbol\{X\}\_\{l\}does not fit in GPU memory \(see Section[A\.4\.1](https://arxiv.org/html/2608.19920#A1.SS4.SSS1)\)\.
- •Runautogradon each cell, iterating from right to left,c=Ncells−1,…,0c=\{N\_\{\\text\{cells\}\}\-1\},\\dots,\{0\}\. For each cellcc, we load layer inputs𝑿l,c\\boldsymbol\{X\}\_\{l,c\}\(bottom\), layer head gradients∂ℒ/∂𝑿l\+1,c\\partial\\mathcal\{L\}/\\partial\\boldsymbol\{X\}\_\{l\+1,c\}\(top\) and incoming cache buffers𝑲l,c\\boldsymbol\{K\}\_\{l,c\}\(left; only forc\>0c\>0\) from CPU, while cache buffer head gradients∂ℒ/∂𝑲l,c\+1\\partial\\mathcal\{L\}/\\partial\\boldsymbol\{K\}\_\{l,c\+1\}\(right; only forc<Ncells−1c<N\_\{\\text\{cells\}\}\-1\) are kept in GPU memory\. Cache decisions are replayed from the log\.autogradworks as follows: - –Forward pass 3 for cell\(l,c\)\(l,c\), in autograd mode\. Whenever a KV cache buffer node \(𝚔𝚎𝚢𝚜\\mathtt\{keys\},𝚟𝚊𝚕𝚞𝚎𝚜\\mathtt\{values\}\) is created, we store anannotationin a list\. In the𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔\\mathtt\{pack\\\_hook\}function, we match arguments of the right shape against all annotation\. For a match, we replace the argument with itsdeltaencoding, which is stored in the computation graph instead\. See Section[A\.4\.3](https://arxiv.org/html/2608.19920#A1.SS4.SSS3)for details\. - –Backward: WhenPyTorchtraverses the computation graph in reverse order, it calls𝚞𝚗𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔\\mathtt\{unpack\\\_hook\}for each node\. For eachdeltaencoding, we play the recurrence \([4](https://arxiv.org/html/2608.19920#S3.E4)\) backwards, replacing𝚔𝚎𝚢𝚜′\\mathtt\{keys\}^\{\\prime\}by𝚔𝚎𝚢𝚜\\mathtt\{keys\}or𝚟𝚊𝚕𝚞𝚎𝚜′\\mathtt\{values\}^\{\\prime\}by𝚟𝚊𝚕𝚞𝚎𝚜\\mathtt\{values\}\. The outcome ofautogradon cell\(l,c\)\(l,c\)are gradients w\.r\.t\. layer weights \(which are accumulated\), a head gradient∂ℒ/∂𝑿l,c\\partial\\mathcal\{L\}/\\partial\\boldsymbol\{X\}\_\{l,c\}\(written to CPU, overwriting𝑿l,c\\boldsymbol\{X\}\_\{l,c\}\), and a head gradient∂ℒ/∂𝑲l,c\\partial\\mathcal\{L\}/\\partial\\boldsymbol\{K\}\_\{l,c\}which replaces the previous one \(forc\>0c\>0\)\. At the end of the loop over cells, gradients w\.r\.t\. layer weights are complete, and the head gradient∂ℒ/∂𝑲l\\partial\\mathcal\{L\}/\\partial\\boldsymbol\{K\}\_\{l\}is on CPU, so the layer below can be addressed \(or, forl=0l=0, gradients w\.r\.t\. input embeddings can be computed based on∂ℒ/∂𝑲0\\partial\\mathcal\{L\}/\\partial\\boldsymbol\{K\}\_\{0\}\)\.

#### A\.4\.3Exploiting Linear Recurrence of KV Cache Buffers

Recall the recurrence between KV cache buffers for subsequent chunks from Section[3\.2](https://arxiv.org/html/2608.19920#S3.SS2)\. We can exploit this recurrence by askingPyTorchto store𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\\mathtt\{delta\\\_key\}instead of𝚔𝚎𝚢𝚜\\mathtt\{keys\}and𝚍𝚎𝚕𝚝𝚊​\_​𝚟𝚊𝚕𝚞𝚎\\mathtt\{delta\\\_value\}instead of𝚟𝚊𝚕𝚞𝚎𝚜\\mathtt\{values\}in the computation graph, restoring the latter from the former during backward using \([4](https://arxiv.org/html/2608.19920#S3.E4)\)\.

While this is simple in principle, we need to do it insidePyTorch autograd\. To this end, we use a mechanism calledautograd saved tensors hooks\([https://docs\.pytorch\.org/tutorials/intermediate/autograd\_saved\_tensors\_hooks\_tutorial\.html](https://docs.pytorch.org/tutorials/intermediate/autograd_saved_tensors_hooks_tutorial.html)\)\. This was designed in order to implement activation checkpointing by CPU offloading, but can be used for our purposes as well\. It works by allowing the specification of two functions:

- •𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔​\(𝒙\)→𝒑⁡\(𝒙\)\\mathtt\{pack\\\_hook\}\(\\boldsymbol\{x\}\)\\to\\boldsymbol\{p\}\(\\boldsymbol\{x\}\): When building its computation graph during forward, this function is called for every array𝒙\\boldsymbol\{x\}PyTorch plans to store in the computation graph\. It then stores𝒑⁡\(𝒙\)\\boldsymbol\{p\}\(\\boldsymbol\{x\}\)in the graph instead of𝒙\\boldsymbol\{x\}\.
- •𝚞𝚗𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔​\(𝒑\)→𝒙⁡\(𝒑\)\\mathtt\{unpack\\\_hook\}\(\\boldsymbol\{p\}\)\\to\\boldsymbol\{x\}\(\\boldsymbol\{p\}\): When traversing the computation graph in reverse order during backward, PyTorch calls this function for every array𝒑\\boldsymbol\{p\}stored in the graph\. It then uses𝒙⁡\(𝒑\)\\boldsymbol\{x\}\(\\boldsymbol\{p\}\)instead of𝒑\\boldsymbol\{p\}\.

A major difficulty for us is the non\-selectiveness of this mechanism\. We do not want to pack all arrays stored in the graph, but only specific ones: the KV cache buffers \(for all other nodes, we just pass through𝒑⁡\(𝒙\)=𝒙\\boldsymbol\{p\}\(\\boldsymbol\{x\}\)=\\boldsymbol\{x\}and𝒙⁡\(𝒑\)=𝒑\\boldsymbol\{x\}\(\\boldsymbol\{p\}\)=\\boldsymbol\{p\}\)\. But𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔​\(𝒙\)\\mathtt\{pack\\\_hook\}\(\\boldsymbol\{x\}\)just takes a𝚝𝚘𝚛𝚌𝚑\.𝚃𝚎𝚗𝚜𝚘𝚛\\mathtt\{torch\.Tensor\}argument, there is no obvious way for tagging the nodes we want\. Also, there is some delay between a node being created in the forward pass and𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔​\(𝒙\)\\mathtt\{pack\\\_hook\}\(\\boldsymbol\{x\}\)being called for it, we even detected some differences in the relative ordering\. Finally, due to internal operator fusion, we cannot even be sure whether any node appearing in the forward code is indeed stored in the graph\.

Our implementation maintains anannotation list, which is appended to during the forward code, while entries are removed during𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔\\mathtt\{pack\\\_hook\}calls\. Whenever a KV cache buffer update in the form of a statement𝚔𝚎𝚢𝚜′=𝚜𝚌𝚊𝚝𝚝𝚎𝚛⁡\(𝚔𝚎𝚢𝚜,𝚒𝚗𝚍𝚎𝚡,𝚔𝚎𝚢​\_​𝚗𝚎𝚠\)\\mathtt\{keys\}^\{\\prime\}=\\mathtt\{scatter\}\(\\mathtt\{keys\},\\mathtt\{index\},\\mathtt\{key\\\_new\}\)is passed in the forward code, we compute𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢=𝚐𝚊𝚝𝚑𝚎𝚛⁡\(𝚔𝚎𝚢𝚜,𝚒𝚗𝚍𝚎𝚡\)\\mathtt\{delta\\\_key\}=\\mathtt\{gather\}\(\\mathtt\{keys\},\\mathtt\{index\}\), appending anannotationcontaining\(𝚒𝚗𝚍𝚎𝚡,𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\)\(\\mathtt\{index\},\\mathtt\{delta\\\_key\}\)and some meta\-data to the list\. Here,𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\\mathtt\{delta\\\_key\}serves a double role\. First, it is needed to reconstruct𝚔𝚎𝚢𝚜\\mathtt\{keys\}from𝚔𝚎𝚢𝚜′\\mathtt\{keys\}^\{\\prime\}in𝚞𝚗𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔\\mathtt\{unpack\\\_hook\}\. Second, it serves141414If𝚒𝚗𝚍𝚎𝚡\\mathtt\{index\}is small, we extend the fingerprint by additional random positions in order to avoid false matches\.as a ”fingerprint” of𝚔𝚎𝚢𝚜\\mathtt\{keys\}\. Namely, when𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔​\(𝒙\)\\mathtt\{pack\\\_hook\}\(\\boldsymbol\{x\}\)is called, we need to match the argument𝒙\\boldsymbol\{x\}against annotations\. This is first done by shape, filtering out most calls\. Next, for any annotation\(𝚒𝚗𝚍𝚎𝚡,𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\)\(\\mathtt\{index\},\\mathtt\{delta\\\_key\}\), we check whether𝚐𝚊𝚝𝚑𝚎𝚛⁡\(𝒙,𝚒𝚗𝚍𝚎𝚡\)=𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\\mathtt\{gather\}\(\\boldsymbol\{x\},\\mathtt\{index\}\)=\\mathtt\{delta\\\_key\}\. If so, we return𝒑⁡\(𝒙\)\\boldsymbol\{p\}\(\\boldsymbol\{x\}\)containing𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\\mathtt\{delta\\\_key\}and remove the annotation from the list\. If there is no match, we return𝒑⁡\(𝒙\)=𝒙\\boldsymbol\{p\}\(\\boldsymbol\{x\}\)=\\boldsymbol\{x\}\.

For𝚞𝚗𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔​\(𝒑\)\\mathtt\{unpack\\\_hook\}\(\\boldsymbol\{p\}\), we reconstruct the sequences𝚔𝚎𝚢𝚜\\mathtt\{keys\}and𝚟𝚊𝚕𝚞𝚎𝚜\\mathtt\{values\}in reverse order\. If𝒑\\boldsymbol\{p\}is not a packed object, we return𝒙⁡\(𝒑\)=𝒑\\boldsymbol\{x\}\(\\boldsymbol\{p\}\)=\\boldsymbol\{p\}\. Otherwise, we check the chunk number stored with𝒑\\boldsymbol\{p\}against the current state\(𝚔𝚎𝚢𝚜′,𝚟𝚊𝚕𝚞𝚎𝚜′\)\(\\mathtt\{keys\}^\{\\prime\},\\mathtt\{values\}^\{\\prime\}\)\. If this fits, we reconstruct𝚔𝚎𝚢𝚜\\mathtt\{keys\}from\(𝚔𝚎𝚢𝚜′,𝚍𝚎𝚕𝚝𝚊​\_​𝚔𝚎𝚢\)\(\\mathtt\{keys\}^\{\\prime\},\\mathtt\{delta\\\_key\}\)or𝚟𝚊𝚕𝚞𝚎𝚜\\mathtt\{values\}from\(𝚟𝚊𝚕𝚞𝚎𝚜′,𝚍𝚎𝚕𝚝𝚊​\_​𝚟𝚊𝚕𝚞𝚎\)\(\\mathtt\{values\}^\{\\prime\},\\mathtt\{delta\\\_value\}\), using \([4](https://arxiv.org/html/2608.19920#S3.E4)\)\. The new buffer overwrites the old one\.

Our implementation tracks which pack hook arguments of the right shape are not matched by annotations, and which annotations are not matched\. Both events do happen, but at a low rate\. It is important to note that we still obtain correct results even if some arguments are not matched\. This just means that a bit more GPU memory is being used\. What we need to avoid, however, are false matches, which we do by keeping fingerprints large enough\.

An important direction for future work is to simplify and robustify the mechanism for exploiting the linear recurrences\. We tried several simplifications\. One assumes that KV buffer nodes are created in exactly the same order as𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔​\(𝒙\)\\mathtt\{pack\\\_hook\}\(\\boldsymbol\{x\}\)calls\. If true, matching and managing the annotation list would be much simplified\. Unfortunately, this does not hold true, likely due to internals of PyTorch autograd we have no influence over\. In fact, the instruction𝚔𝚎𝚢𝚜′=𝚜𝚌𝚊𝚝𝚝𝚎𝚛⁡\(𝚔𝚎𝚢𝚜,𝚒𝚗𝚍𝚎𝚡,𝚔𝚎𝚢​\_​𝚗𝚎𝚠\)\\mathtt\{keys\}^\{\\prime\}=\\mathtt\{scatter\}\(\\mathtt\{keys\},\\mathtt\{index\},\\mathtt\{key\\\_new\}\)need not even trigger a call of𝚙𝚊𝚌𝚔​\_​𝚑𝚘𝚘𝚔\\mathtt\{pack\\\_hook\}with𝒙=𝚔𝚎𝚢𝚜′\\boldsymbol\{x\}=\\mathtt\{keys\}^\{\\prime\}\. The array𝚔𝚎𝚢𝚜′\\mathtt\{keys\}^\{\\prime\}is processed further inside the SDPA code, and PyTorch may use some form of operator fusion\. The simplest solution would be to tag each KV buffer node during creation in a way that allows us to recognize the tag in a pack hook argument𝒙\\boldsymbol\{x\}\. However, we did not find a way to do that yet\.

### A\.5Experimental Results

In this section, we provide additional experimental results beyond what is shown in the main text, as well as timing figures\. We also explain a failure mode we consistently observed when fine\-tuning a model with RingAttention to be used for inference with sparse attention\.

#### A\.5\.1Additional Results

64k datasets128k datasetsnqtri\_qahot\_qapop\_qanqtri\_qahot\_qapop\_qaexact\\mathrm\{exact\}\-\-\-\-\-\-\-\-lr2​k\\mathrm\{lr\}\_\{2k\}​47\.5​80\.2​57\.7​53\.8​35\.3​70\.5​37\.0​40\.2slr2​k\\mathrm\{slr\}\_\{2k\}​45\.3​80\.2​55\.3​52\.0​37\.5​69\.7​34\.7​34\.8h2o2​k\\mathrm\{h2o\}\_\{2k\}​42\.8​70\.8​44\.7​59\.0​20\.5​63\.2​12\.3​27\.3h2o2​kno\\mathrm\{h2o\}\_\{2k\}^\{\\text\{no\}\}​46\.0​82\.2​52\.7​59\.3​42\.2​78\.5​34\.3​39\.3h2o2​kor\\mathrm\{h2o\}\_\{2k\}^\{\\text\{or\}\}​44\.2​74\.7​47\.7​61\.0​43\.7​79\.8​25\.0​43\.3lr1​k\\mathrm\{lr\}\_\{1k\}​47\.5​80\.2​59\.7​54\.0​35\.8​70\.2​37\.7​37\.5slr1​k\\mathrm\{slr\}\_\{1k\}​49\.0​79\.0​55\.0​54\.3​37\.3​67\.3​32\.0​35\.7h2o1​k\\mathrm\{h2o\}\_\{1k\}​43\.2​71\.7​41\.7​59\.8​21\.0​62\.8​13\.0​27\.0h2o1​kno\\mathrm\{h2o\}\_\{1k\}^\{\\text\{no\}\}​47\.0​81\.2​50\.0​56\.7​40\.0​79\.8​32\.7​42\.3h2o1​kor\\mathrm\{h2o\}\_\{1k\}^\{\\text\{or\}\}​47\.5​76\.3​46\.7​58\.2​41\.8​79\.8​23\.7​39\.3Table 4:Results for long\-context inference with 5 KV cache policies and chunk sizes2048=2​k,1024=1​k2048=2k,1024=1k\(rows\)\. Here, the base checkpointQwen3\-4B\-Instruct\-2507is used without fine\-tuning\. The first rowexactis for exact inference \(sequence parallelism\)\. We show𝚜𝚞𝚋​\_​𝚎𝚡𝚊𝚌𝚝​\_​𝚖𝚊𝚝𝚌𝚑\\mathtt\{sub\\\_exact\\\_match\}values on test splits for different Helmet datasetsnq, trivia\_qa, hotpot\_qa, pop\_qa, limiting sequence lengths to 64k or 128k tokens\.In Table[4](https://arxiv.org/html/2608.19920#A1.T4), we provide results on Helmet 64k and 128k datasets for base checkpointsQwen3\-4B\-Instruct\-2507\(no fine\-tuning\)\. They should be related to results in Table[1](https://arxiv.org/html/2608.19920#S4.T1), where the base model was trained by our method \(columnsus\) or by sequence parallelism \(columnssp\)\.

64k datasetsnqtri\_qahot\_qapop\_qaslr128\\mathrm\{slr\}\_\{128\}​38\.8​66\.2​51\.7​41\.0h2o128\\mathrm\{h2o\}\_\{128\}​44\.5​65\.3​55\.7​54\.3h2o128no\\mathrm\{h2o\}\_\{128\}^\{\\text\{no\}\}​49\.0​74\.3​55\.0​56\.3h2o128or\\mathrm\{h2o\}\_\{128\}^\{\\text\{or\}\}​46\.8​67\.3​48\.3​52\.3qh2o2​k\\mathrm\{qh2o\}\_\{2k\}​37\.5​64\.0​40\.3​53\.7qh2o2​kno\\mathrm\{qh2o\}\_\{2k\}^\{\\text\{no\}\}​40\.2​65\.5​46\.3​53\.5Table 5:Results for long\-context inference with setups not covered in the main text\. We show𝚜𝚞𝚋​\_​𝚎𝚡𝚊𝚌𝚝​\_​𝚖𝚊𝚝𝚌𝚑\\mathtt\{sub\\\_exact\\\_match\}values on test splits for different Helmet datasetsnq, trivia\_qa, hotpot\_qa, pop\_qa, limiting sequence lengths to 64k or 128k tokens\.
slr128\\mathrm\{slr\}\_\{128\},h2o128\\mathrm\{h2o\}\_\{128\},h2o128no\\mathrm\{h2o\}\_\{128\}^\{\\text\{no\}\},h2o128or\\mathrm\{h2o\}\_\{128\}^\{\\text\{or\}\}use chunk sizeS=128S=128\.qh2o2​k\\mathrm\{qh2o\}\_\{2k\}andqh2o2​kno\\mathrm\{qh2o\}\_\{2k\}^\{\\text\{no\}\}are variants of Q\-Hitter\[[81](https://arxiv.org/html/2608.19920#bib.bib70)\]\.In Table[5](https://arxiv.org/html/2608.19920#A1.T5), we provide results on Helmet 64k and 128k datasets for setups not covered in the main text \(see Table[1](https://arxiv.org/html/2608.19920#S4.T1)\)\.slr128\\mathrm\{slr\}\_\{128\},h2o128\\mathrm\{h2o\}\_\{128\},h2o128no\\mathrm\{h2o\}\_\{128\}^\{\\text\{no\}\},h2o128or\\mathrm\{h2o\}\_\{128\}^\{\\text\{or\}\}use chunk sizeS=128S=128\. This runs significantly longer thanS∈\{1024,2048\}S\\in\\\{1024,2048\\\}used in the main text experiments, but allows the KV cache policy to make decisions 8 or 16 times more frequently\. However, at least in the experiments here, this does not lead to better results, justifying our choice of larger chunk sizes above\. We also ran experiments with Q\-Hitter\[[81](https://arxiv.org/html/2608.19920#bib.bib70)\], where KV cache buffers are quantized \(to 8 bits in our experiments\) and the decision score is a convex combination of \([3](https://arxiv.org/html/2608.19920#S3.E3)\) and a term quantifying the quantization error\. The results are consistently worse than for the H2O variants\.

#### A\.5\.2Running Time Figures

nqtri\_qahot\_qapop\_qaexact\\mathrm\{exact\}​258\.38 \(15\.14\)​266\.05 \(11\.31\)​262\.53 \(8\.18\)​236\.74 \(21\.44\)lr2​k\\mathrm\{lr\}\_\{2k\}​326\.07 \(21\.86\)​333\.28 \(16\.56\)​330\.76 \(22\.49\)​312\.29 \(27\.71\)slr2​k\\mathrm\{slr\}\_\{2k\}​323\.84 \(21\.19\)​333\.03 \(16\.19\)​330\.29 \(22\.55\)​310\.73 \(24\.46\)h2o2​k\\mathrm\{h2o\}\_\{2k\}​330\.83 \(21\.98\)​344\.83 \(16\.70\)​336\.79 \(23\.13\)​316\.02 \(27\.66\)h2o2​kno\\mathrm\{h2o\}\_\{2k\}^\{\\text\{no\}\}​331\.48 \(22\.10\)​341\.71 \(16\.65\)​337\.77 \(23\.12\)​317\.60 \(25\.16\)h2o2​kor\\mathrm\{h2o\}\_\{2k\}^\{\\text\{or\}\}​331\.83 \(22\.25\)​344\.85 \(16\.89\)​338\.46 \(23\.25\)​316\.94 \(25\.34\)lr1​k\\mathrm\{lr\}\_\{1k\}​364\.10 \(23\.28\)​374\.80 \(18\.62\)​375\.38 \(26\.38\)​343\.96 \(27\.99\)slr1​k\\mathrm\{slr\}\_\{1k\}​362\.84 \(23\.20\)​370\.55 \(18\.66\)​371\.18 \(25\.77\)​344\.77 \(27\.65\)h2o1​k\\mathrm\{h2o\}\_\{1k\}​378\.85 \(24\.88\)​385\.19 \(20\.21\)​382\.83 \(26\.94\)​359\.51 \(32\.00\)h2o1​kno\\mathrm\{h2o\}\_\{1k\}^\{\\text\{no\}\}​378\.61 \(24\.99\)​385\.29 \(19\.32\)​388\.77 \(27\.69\)​357\.95 \(29\.02\)h2o1​kor\\mathrm\{h2o\}\_\{1k\}^\{\\text\{or\}\}​378\.65 \(24\.79\)​385\.64 \(19\.46\)​382\.52 \(26\.91\)​358\.99 \(31\.74\)Table 6:Running time figures for training update step, for Helmet 128k datasets \(columns\), 5 KV cache policies and chunk sizes2048=2​k,1024=1​k2048=2k,1024=1k\(rows\)\. Batch size 8, running on 4 devices\. The step from 2k to 1k is 11% to 13% more expensive for𝚕𝚛,𝚜𝚕𝚛\\mathtt\{lr\},\\mathtt\{slr\}, 12% to 14% more expensive for𝚑𝟸𝚘\\mathtt\{h2o\}variants\. The step from𝚕𝚛,𝚜𝚕𝚛\\mathtt\{lr\},\\mathtt\{slr\}to𝚑𝟸𝚘\\mathtt\{h2o\}variants is 2% to 3% more expensive for 2k, 3% to 4% more expensive for 1k\.In this section, we present running time figures\. First, we consider training updates \(batch size 8; four Nvidia A100s with 40 GB each\)\. For Helmet 128k datasets,exactuses sequence parallelism with batch size 2, processing 4 micro\-batches sequentially, whereas our method \(for different cache logics\) processes 4 micro\-batches in parallel\.

First, our method is about 30% more expensive thanexactfor chunk sizeS=2048S=2048\(2k\)\. Given that our method computes gradients on a single GPU independent of the sequence length, using advanced cache logics such as H2O, nested activation checkpointing and delta encoding of cache buffers, this overhead is surprisingly small\. Reasons for the overhead are explained in Section[3\.3](https://arxiv.org/html/2608.19920#S3.SS3)\. The gap can likely be narrowed further by operator fusion, increasing the chunk sizeautogradis operating with\. Next, we would expectS=1024S=1024\(1k\) to run slower thanS=2048S=2048\(2k\), because more chunks need to be processed sequentially; and H2O policies to run slower than𝚕𝚛,𝚜𝚕𝚛\\mathtt\{lr\},\\mathtt\{slr\}, because summed attention weights are required, and scores need to be computed and sorted\. We see that chunk size 1k variants run between 11% and 14% longer than 2k variants, which is substantial\. On the other hand, H2O variants are only between 2% to 4% slower than𝚕𝚛,𝚜𝚕𝚛\\mathtt\{lr\},\\mathtt\{slr\}\. At least with our fast SDPA implementation, there is no penalty for using more advanced policies over simple baselines\.

#### A\.5\.3Analysis of Errors

In this paper, we compare different ways of fine\-tuning a model to be used with sparse attention inference under different KV cache logics: training by sequence parallelism \(sp\) versus training with the new, resource\-efficient technique developed here \(us\)\. While for datasetsnq, trivia\_qa, hotpot\_qa, pop\_qacoming with theSubEMmetric, results are inconclusive \(see Table[1](https://arxiv.org/html/2608.19920#S4.T1)\),usstrongly outperformsspon datasetstrec\_coarse,nlu,clinc150,inf\_qa,inf\_mc,json\_kv, where the metric is mostlyAccuracy\(see Table[2](https://arxiv.org/html/2608.19920#S4.T2)\)\. In this section, we identify a consistent failure mode ofsp, both via randomly chosen examples and statistics on the generated samples\.

Recall details about Helmet datasets from Section[A\.3\.1](https://arxiv.org/html/2608.19920#A1.SS3.SSS1)\. Fortrec\_coarse,nlu,clinc150,inf\_mc, targets are single integers, and theAccuracymetric requires the correct number to appear most frequently in the output\. Fornq, trivia\_qa, hotpot\_qa, pop\_qa, json\_kv, targets are lists of noun phrases of one of few words, and theSubEMmetric requires one of the targets to be a substring of the output\. As with all sample\-based metrics in our experiments, we allow the model to generate up to 128 tokens, or otherwise terminate by emitting the<eos\>token\. A closer look at generated samples reveals a major failure mode ofsp:its outputs are far too long and contain mostly random nonsense\. Here are some randomly chosen examples of tasks and outputs underusandsprespectively\.

trec\_coarse / smart\_lastrec\_1024:

\-idx:10

raw\_target:”2”

exact\_match\_us:1

output\_us:”2”

exact\_match\_sp:0

output\_sp:”2

WhatistheaveragelifeexpectancyofafemaleinIrelandin1996?

1

WhatisthenameoftheAmericanwhowascapturedwhenhisplanewentdownover

Syrian\-heldLebanon?

1

0

WhatisthefullnameofthePLO?

1

3

Whatisthedifferencebetweenacollegeandauniversity?

1

4

Whatistheoriginoftheterm“TheBrooklynSide””inbowling?

1

4

Whatistheoriginoftheword“assassinate””?

1

4

Whatdoestheabbreviationcwt\.?mean?

1

3

Whatdoesthe”

\-idx:318

raw\_target:”5”

exact\_match\_us:1

output\_us:”5”

exact\_match\_sp:0

output\_sp:”1

WhatdoestheacronymCPRmean?

label:3

Whatdoestheabbreviationcwt\.?

3

Whatistheoriginofthename‘Mabel”?

3

Whatisthefullformof\.com?

3

Whatisthenameofthecompanythatmanufacturesthe“AmericanGirl””doll

collection?

4

Howmanypeoplewerekilledinthe1974fireatthePentagon?

4

Whatisthepurposeofacarbra?

4

Whatisthedifferencebetweenaclassicalandanacousticguitar?

1

Whatisthenameofthecompanythatmanufacturedthefirstphonograph?

4”

clinc150 / smart\_lastrec\_1024:

\-idx:25

raw\_target:”27”

exact\_match\_us:1

output\_us:”27”

exact\_match\_sp:1

output\_sp:”27

howlonguntilmynextdayoff

label\|15”

\-idx:78

raw\_target:”8”

exact\_match\_us:1

output\_us:”8”

exact\_match\_us:0

output\_sp:”8\\n\\ncanyoufindmearecipeforchocolatechipcookies\\nlabelle66\\

\\\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\

\\\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n\\

66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n66\\n”

nlu / smart\_lastrec\_1024:

\-idx:63

raw\_target:”61”

exact\_match\_us:1

output\_us:”61”

exact\_match\_sp:0

output\_sp:”61

whatistheweatherinthecity

graph

57

howfarisitfromheretonewyork

graph

25

howmanycaloriesinabanana

graph:6

whatisthedefinitionofthewordpda

graph:11

whatisthebestmovietowatchthisweekend

graph:2

whatisthetimedifferencebetweeneasternandwestern

graph:49

turnonthewemoplug

graph:29

what”stheweatherinboston

graph:57

what”stheweatherinnewyork

graph:”

\-idx:245

raw\_target:”27”

exact\_match\_us:1

output\_us:”27”

exact\_match\_sp:1

output\_sp:”27

whatistheexchangeratebetweenusandcanada

graph:5

what”sthenameofthepersonbehindthesongoftaylorswiftbacktodecember

graph:52

what”stheweatherintampa

graph:27

whatistheweatherforecastfortheweek

graph:27

what”stheweathergoingtobeliketodayintampa

graph:27

iwouldliketoknowtheweatherinnewyork

graph:27

whatisthetimeinthepacifictimezone

graph:63

whatisthe”

json\_kv / h2onorm\_1024:

\-idx:81

raw\_target:d823df41\-5af0\-4f11\-925d\-d0e8018cc5be

sub\_exact\_match\_us:1

output\_us:d823df41\-5af0\-4f11\-925d\-d0e8018cc5be

sub\_exact\_match\_sp:0

output\_sp:’923928d\-2e95\-426c\-8893\-b5e80880a88c

json

\{”c”:”11984809639058896”,”a”:”923928d\-2e95\-426c\-8893\-b5e80880a88c”,”b”:”0e2b980e\-da28\-4af1\-8896\-3f0c’

\-idx:51

raw\_target:d8cdc4a6\-e37e\-4cd3\-b603\-ad33233518e9

sub\_exact\_match\_us:0

output\_us:d8cdc4a6\-e37e\-4cd3\-b6d3\-ad33233b18e9

sub\_exact\_match\_sp:0

output\_sp:”033a8782\-23d2\-488c\-ae99\-488f80bbc7d6\\n\\nKey:9c71e7e3\-9f60\-47f4\-baaa\-42dbca3e2715:\\

\\\\”d477445a\-0f1b\-4729\-b586\-10cbca0a5ba3\\”,\\n\\”9f71e7e3\-9f60\-47f4\-baaa\-4””

nq:

\-idx:418

raw\_target:

\-six

\-eight

sub\_exact\_match\_us:1

output\_us:”six”

sub\_exact\_match\_sp:0

output\_sp:”4hoopsareusedinagameofcroquet\.\(2blue,1red”

\-idx:419

raw\_target:

\-six

\-eight

sub\_exact\_match\_us:0

output\_us:”four”

sub\_exact\_match\_sp:0

output\_sp:”20hoops\(10perside\)areusedinagameofcroquet\.”

pop\_qa:

\-idx:275

raw\_target:

\-Paraguay

\-RepublicofParaguay

\-py

\-”\\U0001F1F5\\U0001F1FE”

\-HeartofSouthAmerica

sub\_exact\_match\_us:0

output\_us:”Peru”

sub\_exact\_match\_sp:0

output\_sp:”Peru

Question:WhatisthecapitalofPeru?

Answer:Lima

Question:Inwhatregion”

\-idx:273

raw\_target:

\-Paraguay

\-RepublicofParaguay

\-py

\-”\\U0001F1F5\\U0001F1FE”

\-HeartofSouthAmerica

sub\_exact\_match\_us:0

output\_us:”Peru”

sub\_exact\_match\_sp:0

output\_sp:”Peru

Question:WhatisthenameofthePeruviancitywheretheNationalLibraryis

located?”

Most other examples we inspected reveal the same failure mode\. Whileuslearns to output exactly the numerical answer or noun phrase and nothing else, sparse attention inference forsptends to output a lot of random content\. Recall that duringsptraining, each token can attend to any earlier one in principle\. Plugging in a cache eviction logic afterwards seems todiminish the model’s ability to correctly stop generation\.151515For most cache logics, inference withspcheckpoints does stop with<eos\>for some samples, emitting sometimes as few as 5 or 7 tokens\. Moreover, for some logics anduscheckpoints, generation overshoots to 128 tokens as well \(albeit rarely\)\.Sometimes, the first number in the output is correct, but is followed by many others in the output\. Forclinc150, idx:25, we haveexact\_match\_sp = 1despite the output being partly random and containing another number\. Fornlu, idx:245, the correct answer 27 appears most frequently in nonsense output\. In thepop\_qaexample, whileusgets the country wrong \(”Peru” instead of ”Paraguay”\),spalso outputs nonsense extra content after the single word\. Forjson\_kv, idx:51, the output forusis only off by two letters, while that forspis nonsense, containing several UUIDs completely different from the target\.

trnslr1​k\\mathrm\{slr\}\_\{1k\}h2o1​kno\\mathrm\{h2o\}\_\{1k\}^\{\\text\{no\}\}h2o1​kor\\mathrm\{h2o\}\_\{1k\}^\{\\text\{or\}\}RRp128p\_\{128\}RRp128p\_\{128\}RRp128p\_\{128\}nqus​1\.1±\\pm​0\.8​0\.0±\\pm​0\.0​1\.1±\\pm​1\.0​0\.0±\\pm​0\.0​1\.1±\\pm​2\.7​0\.2±\\pm​4\.1sp​35\.5±\\pm​24\.1​99\.5±\\pm​7\.1​35\.5±\\pm​21\.6​100\.0±\\pm​0\.0​36\.2±\\pm​22\.1​99\.3±\\pm​8\.1no​35\.2±\\pm​22\.0​97\.7±\\pm​15\.1​35\.8±\\pm​22\.8​99\.2±\\pm​9\.1​34\.9±\\pm​22\.0​97\.8±\\pm​14\.6trivia\_qaus​1\.2±\\pm​0\.8​0\.0±\\pm​0\.0​1\.3±\\pm​1\.0​0\.0±\\pm​0\.0​1\.1±\\pm​0\.7​0\.0±\\pm​0\.0sp​35\.4±\\pm​24\.1​97\.3±\\pm​16\.1​38\.2±\\pm​25\.2​99\.2±\\pm​9\.1​42\.0±\\pm​24\.4​96\.7±\\pm​18\.0no​43\.4±\\pm​28\.5​87\.8±\\pm​32\.7​46\.5±\\pm​28\.9​96\.0±\\pm​19\.6​48\.2±\\pm​30\.4​95\.7±\\pm​20\.4hotpot\_qaus​1\.0±\\pm​0\.5​0\.0±\\pm​0\.0​1\.0±\\pm​0\.6​0\.0±\\pm​0\.0​1\.3±\\pm​1\.4​1\.0±\\pm​9\.9sp​39\.2±\\pm​31\.8​93\.0±\\pm​25\.5​40\.9±\\pm​31\.4​98\.3±\\pm​12\.8​41\.4±\\pm​31\.1​99\.0±\\pm​9\.9no​41\.0±\\pm​30\.9​91\.7±\\pm​27\.6​41\.3±\\pm​31\.1​98\.0±\\pm​14\.0​41\.4±\\pm​31\.1​98\.7±\\pm​11\.5pop\_qaus​1\.1±\\pm​0\.6​0\.0±\\pm​0\.0​1\.1±\\pm​0\.5​0\.0±\\pm​0\.0​1\.0±\\pm​0\.4​0\.0±\\pm​0\.0sp​53\.8±\\pm​24\.9​98\.8±\\pm​10\.7​54\.4±\\pm​24\.7​99\.5±\\pm​7\.1​55\.0±\\pm​24\.7​98\.7±\\pm​11\.5no​58\.9±\\pm​29\.3​93\.5±\\pm​24\.7​60\.0±\\pm​30\.9​98\.8±\\pm​10\.7​58\.2±\\pm​28\.9​96\.8±\\pm​17\.5trec\_coarseus​1\.0±\\pm​0\.0​0\.0±\\pm​0\.0​1\.0±\\pm​0\.0​0\.0±\\pm​0\.0​1\.0±\\pm​0\.0​0\.0±\\pm​0\.0sp​127\.6±\\pm​6\.1​99\.6±\\pm​6\.3​127\.3±\\pm​8\.4​99\.0±\\pm​9\.9​128\.0±\\pm​0\.1​99\.6±\\pm​6\.3no​128\.0±\\pm​0\.0​100\.0±\\pm​0\.0​128\.0±\\pm​0\.0​100\.0±\\pm​0\.0​128\.0±\\pm​0\.1​99\.4±\\pm​7\.7nluus​1\.0±\\pm​0\.1​0\.0±\\pm​0\.0​1\.0±\\pm​0\.1​0\.0±\\pm​0\.0​1\.0±\\pm​0\.2​0\.0±\\pm​0\.0sp​70\.4±\\pm​21\.1​70\.6±\\pm​45\.6​69\.0±\\pm​21\.7​45\.6±\\pm​49\.8​14\.6±\\pm​18\.8​8\.4±\\pm​27\.7no​71\.7±\\pm​20\.8​100\.0±\\pm​0\.0​71\.7±\\pm​20\.8​100\.0±\\pm​0\.0​71\.7±\\pm​20\.8​99\.8±\\pm​4\.5clinc150us​1\.0±\\pm​0\.1​0\.0±\\pm​0\.0​1\.0±\\pm​0\.1​0\.0±\\pm​0\.0​1\.0±\\pm​0\.1​0\.0±\\pm​0\.0sp​30\.2±\\pm​28\.6​43\.8±\\pm​49\.6​32\.6±\\pm​31\.0​49\.8±\\pm​50\.0​57\.6±\\pm​32\.2​87\.6±\\pm​33\.0no​59\.9±\\pm​19\.7​100\.0±\\pm​0\.0​59\.9±\\pm​19\.7​100\.0±\\pm​0\.0​60\.1±\\pm​20\.5​100\.0±\\pm​0\.0inf\_qaus​1\.1±\\pm​0\.8​0\.0±\\pm​0\.0​1\.2±\\pm​0\.8​0\.0±\\pm​0\.0​1\.2±\\pm​0\.8​0\.0±\\pm​0\.0sp​45\.7±\\pm​32\.4​100\.0±\\pm​0\.0​45\.7±\\pm​32\.4​99\.0±\\pm​9\.9​45\.8±\\pm​32\.3​97\.0±\\pm​17\.1no​45\.5±\\pm​32\.3​95\.0±\\pm​21\.8​45\.7±\\pm​32\.4​99\.0±\\pm​9\.9​45\.7±\\pm​32\.4​96\.0±\\pm​19\.6inf\_mcus​1\.0±\\pm​0\.0​0\.0±\\pm​0\.0​2\.3±\\pm​12\.6​1\.0±\\pm​9\.9​1\.0±\\pm​0\.0​0\.0±\\pm​0\.0sp​128\.0±\\pm​0\.0​100\.0±\\pm​0\.0​128\.0±\\pm​0\.0​100\.0±\\pm​0\.0​127\.9±\\pm​1\.2​98\.0±\\pm​14\.0no​128\.0±\\pm​0\.0​100\.0±\\pm​0\.0​128\.0±\\pm​0\.0​100\.0±\\pm​0\.0​128\.3±\\pm​3\.6​97\.0±\\pm​17\.1json\_kvus​1\.1±\\pm​0\.1​0\.0±\\pm​0\.0​1\.0±\\pm​0\.0​0\.0±\\pm​0\.0​3\.7±\\pm​1\.0​85\.0±\\pm​35\.7sp​4\.1±\\pm​0\.3​100\.0±\\pm​0\.0​4\.1±\\pm​0\.3​100\.0±\\pm​0\.0​4\.1±\\pm​0\.3​98\.0±\\pm​14\.0no​4\.1±\\pm​0\.3​100\.0±\\pm​0\.0​4\.1±\\pm​0\.3​100\.0±\\pm​0\.0​4\.1±\\pm​0\.3​100\.0±\\pm​0\.0Table 7:Token length statistics of generated samples for 10 Helmet datasets \(of context width 128k\) and 3 cache logics\.trndenotes model checkpoint being used:ususes our novel method with the same cache policy in place,spis using sequence parallelism,nois the base checkpointQwen3\-4B\-Instruct\-2507\(no fine\-tuning\)\.RRis based on the ratio of output length to target length \(in tokens\),p128p\_\{128\}\(in percent\) is the fraction of outputs of maximal size 128 \(means, and stddevs over all test set samples\)\.nqtri\_qahot\_qapop\_qatrec\_cRRp128p\_\{128\}RRp128p\_\{128\}RRp128p\_\{128\}RRp128p\_\{128\}RRp128p\_\{128\}​1\.1±\\pm​1\.1​0\.0±\\pm​0\.0​1\.0±\\pm​0\.6​0\.0±\\pm​0\.0​1\.1±\\pm​0\.6​0\.0±\\pm​0\.0​1\.0±\\pm​0\.4​0\.0±\\pm​0\.0​1\.0±\\pm​0\.0​0\.0±\\pm​0\.0nluclc150inf\_qainf\_mcjson\_kvRRp128p\_\{128\}RRp128p\_\{128\}RRp128p\_\{128\}RRp128p\_\{128\}RRp128p\_\{128\}​1\.0±\\pm​0\.1​0\.0±\\pm​0\.0​1\.0±\\pm​0\.0​0\.0±\\pm​0\.0​1\.2±\\pm​0\.9​0\.0±\\pm​0\.0​1\.0±\\pm​0\.0​0\.0±\\pm​0\.0​1\.0±\\pm​0\.0​0\.0±\\pm​0\.0Table 8:Token length statistics of generated samples for 10 Helmet datasets \(of context width 128k\) for training and inference with exact attention \(sequence parallelism\)\.In order to quantify the prevalence of this failure mode across all datasets and setups, we use two statistics, estimated over all samples generated for each dataset and cache logic:

R=len⁡\(𝚘𝚞𝚝𝚙𝚞𝚝\)len⁡\(𝚝𝚊𝚛𝚐𝚎𝚝\),p128=I\{len\(𝚘𝚞𝚝𝚙𝚞𝚝\)=128\}\.R=\\frac\{\\mathrm\{len\}\(\\mathtt\{output\}\)\}\{\\mathrm\{len\}\(\\mathtt\{target\}\)\},\\quad p\_\{128\}=\\mathrm\{I\}\_\{\\\{\\mathrm\{len\}\(\\mathtt\{output\}\)=128\\\}\}\.Here,len⁡\(⋅\)\\mathrm\{len\}\(\\cdot\)denotes length in tokens, and samples are capped at 128 tokens\. For some datasets, the targets are a list, in which case𝚝𝚊𝚛𝚐𝚎𝚝\\mathtt\{target\}is the longest entry appearing as substring in𝚘𝚞𝚝𝚙𝚞𝚝\\mathtt\{output\}, or the longest entry otherwise\. Statistics for 10 Helmet datasets and 3 setups are shown in Table[7](https://arxiv.org/html/2608.19920#A1.T7)\. Withus, we haveR≈1R\\approx 1across datasets and setups: outputs are close in length to targets, the model learned the desired output type and stops generation properly\. But withsp,RRtends to be large, andp128p\_\{128\}is often close to 100%\. This is not a property of thespcheckpoints\. As seen in Table[8](https://arxiv.org/html/2608.19920#A1.T8),R≈1R\\approx 1andp128≈0p\_\{128\}\\approx 0if exact inference is used\. Instead, failures come from the inconsistency between training and inference\.

If we relate numbers in Table[7](https://arxiv.org/html/2608.19920#A1.T7)with good results in Table[1](https://arxiv.org/html/2608.19920#S4.T1)forspand in Table[4](https://arxiv.org/html/2608.19920#A1.T4)forno, this points to a shortcoming of theSubEMmetric used for these datasets\. Insensitive to any type and amount of nonsense extra output,it only requires the target to be contained in the output, without requiring a definite way to extract the substring\. Once such a requirement is added, as inAccuracy, performance forspandnoplummets\. In any case, by not even providing succinct outputs \(a very clear signal in the data\),spclearly does not behave satisfactory\. While the tolerance ofSubEM\(and alsoAccuracy, to a lesser extent\) to any amount of extra output is intended to not disadvantage LLMs \(which ”sugar\-coat” answers in longer sentences\), it makes the metrics blind to extra nonsense content returned\. We should at least ask for a deterministic way to extract the response from the output\.

### A\.6Computing Summed Attention Weights in SDPA

As noted in Section[3\.1\.1](https://arxiv.org/html/2608.19920#S3.SS1.SSS1)and Section[3\.3\.1](https://arxiv.org/html/2608.19920#S3.SS3.SSS1), KV cache policies like H2O or related ones need summed attention weights∑imb,h,i,j\\sum\_\{i\}m\_\{b,h,i,j\}for each\(b,h,j\)\(b,h,j\)\. This array of shape\(B,Hq,Nk\)\(B,H\_\{q\},N\_\{k\}\)can be obtained as byproduct of SDPA, which returns𝒀\\boldsymbol\{Y\}of shape\(B,Hq,Nq,dh\)\(B,H\_\{q\},N\_\{q\},d\_\{h\}\)\. Note that summed attention weights are smaller than attention outputs, so there is a priori no reason for not returning them\. However, all fast SDPA codes we know of, do not return this information\. FlexAttention\[[14](https://arxiv.org/html/2608.19920#bib.bib14)\]can return log\-sum\-exp valueslog\(𝑨𝟏\)Nk\\log\(\\boldsymbol\{A\}\\boldsymbol\{1\}\{\}\_\{N\_\{k\}\}\), where𝑨=𝚖𝚊𝚜𝚔\(dh−1/2𝑸𝑲T\)\\boldsymbol\{A\}=\\mathtt\{mask\}\(d\_\{h\}^\{\-1/2\}\\boldsymbol\{Q\}\\boldsymbol\{K\}^\{T\}\)is the argument of𝚜𝚘𝚏𝚝𝚖𝚊𝚡\\mathtt\{softmax\}, likely because this is directly computed during FlashAttention\[[10](https://arxiv.org/html/2608.19920#bib.bib8)\]\.

Our implementation containsTritoncode for computing summed attention weights alongside a FlashInfer SDPA kernel\[[72](https://arxiv.org/html/2608.19920#bib.bib63)\]\. This is an add\-on, and it would be better if leading SDPA codes returned summed attention weights directly\. In this section, we detail how this can be done \(even though we have not implemented this\)\. We also show how to compute them with FlexAttention\[[14](https://arxiv.org/html/2608.19920#bib.bib14)\], using two calls instead of one\. This is contained in our implementation as baseline\.

#### A\.6\.1FlashAttention for Summed Attention Weights

FlashAttention works by essentially computing the attention weights tensor𝑴\\boldsymbol\{M\}in blocks, using a lattice tiling along the query and the key axes\. It can be understood asmap\-reduce, wheremapis independent per cell\. More precisely, the full attention weights have shape\(B,Hq,Nq,Nk\)\(B,H\_\{q\},N\_\{q\},N\_\{k\}\)\. In the following, we drop\(B,Hq\)\(B,H\_\{q\}\), treating them as ”batch” dimensions\. Use cell indices\(r,s\)\(r,s\)and index rangesI⁡\(r\)I\(r\),J⁡\(s\)J\(s\), so that the union of allI⁡\(r\)I\(r\)covers\{0,…,Nq−1\}\\\{\{0\},\\dots,\{N\_\{q\}\-1\}\\\}and the union of allJ⁡\(s\)J\(s\)covers\{0,…,Nk−1\}\\\{\{0\},\\dots,\{N\_\{k\}\-1\}\\\}\. When computing the attention outputs,reduceoperates along the key axis\. Define an additional auxiliary tensor of shape\(B,Hq,Nq\)\(B,H\_\{q\},N\_\{q\}\), with values

𝝀r,s:=log\(exp\(𝑨I⁡\(r\),J⁡\(s\)\)𝟏\|J⁡\(s\)\|\),𝑨I,J:=𝚖𝚊𝚜𝚔\(dh−1/2𝑸I,⋅𝑲J,⋅T\)\.\\boldsymbol\{\\lambda\}\_\{r,s\}:=\\log\\left\(\\exp\\left\(\\boldsymbol\{A\}\_\{I\(r\),J\(s\)\}\\right\)\\boldsymbol\{1\}\_\{\|J\(s\)\|\}\\right\),\\quad\\boldsymbol\{A\}\_\{I,J\}:=\\mathtt\{mask\}\\left\(d\_\{h\}^\{\-1/2\}\\boldsymbol\{Q\}\_\{I,\\cdot\}\\boldsymbol\{K\}\_\{J,\\cdot\}^\{T\}\\right\)\.Reduction works as:

𝝀r,s1⊕s2=max⁡\{𝝀r,s1,𝝀r,s2\}\+log1p⁡\(exp⁡\(−\|𝝀r,s1−𝝀r,s2\|\)\),𝒀r,s1⊗s2=\(diag⁡exp⁡\(𝝀r,s1−𝝀r,s1⊗s2\)\)​𝒀r,s1\+\(diag⁡exp⁡\(𝝀r,s2−𝝀r,s1⊗s2\)\)​𝒀r,s2\.\\begin\{split\}\\boldsymbol\{\\lambda\}\_\{r,s\_\{1\}\\oplus s\_\{2\}\}&=\\max\\left\\\{\\boldsymbol\{\\lambda\}\_\{r,s\_\{1\}\},\\boldsymbol\{\\lambda\}\_\{r,s\_\{2\}\}\\right\\\}\+\\mathrm\{log1p\}\\left\(\\exp\\left\(\-\\left\|\\boldsymbol\{\\lambda\}\_\{r,s\_\{1\}\}\-\\boldsymbol\{\\lambda\}\_\{r,s\_\{2\}\}\\right\|\\right\)\\right\),\\\\ \\boldsymbol\{Y\}\_\{r,s\_\{1\}\\otimes s\_\{2\}\}&=\\left\(\\diag\\exp\\left\(\\boldsymbol\{\\lambda\}\_\{r,s\_\{1\}\}\-\\boldsymbol\{\\lambda\}\_\{r,s\_\{1\}\\otimes s\_\{2\}\}\\right\)\\right\)\\boldsymbol\{Y\}\_\{r,s\_\{1\}\}\+\\left\(\\diag\\exp\\left\(\\boldsymbol\{\\lambda\}\_\{r,s\_\{2\}\}\-\\boldsymbol\{\\lambda\}\_\{r,s\_\{1\}\\otimes s\_\{2\}\}\\right\)\\right\)\\boldsymbol\{Y\}\_\{r,s\_\{2\}\}\.\\end\{split\}We can now runmapindependently for all\(r,s\)\(r,s\), thenreducealongssfor allrr\.

Summed attention weights are given by𝒘∈ℝNk\\boldsymbol\{w\}\\in\\mathbb\{R\}^\{N\_\{k\}\}:

𝒘T=𝟏NqT​exp⁡\(𝑨−𝝀​𝟏NkT\)\.\\boldsymbol\{w\}^\{T\}=\\boldsymbol\{1\}\_\{N\_\{q\}\}^\{T\}\\exp\\left\(\\boldsymbol\{A\}\-\\boldsymbol\{\\lambda\}\\boldsymbol\{1\}\_\{N\_\{k\}\}^\{T\}\\right\)\.Define

𝑭~r,s=exp⁡\(𝑨I⁡\(r\),J⁡\(s\)−𝝀r,s​𝟏\|J⁡\(s\)\|T\),𝑭r,s=\(diag⁡exp⁡\(𝝀r,s−𝝀r\)\)​𝑭~r,s=exp⁡\(𝑨I⁡\(r\),J⁡\(s\)−𝝀r​𝟏\|J⁡\(s\)\|T\)\.\\begin\{split\}\\tilde\{\\boldsymbol\{F\}\}\_\{r,s\}&=\\exp\\left\(\\boldsymbol\{A\}\_\{I\(r\),J\(s\)\}\-\\boldsymbol\{\\lambda\}\_\{r,s\}\\boldsymbol\{1\}\_\{\|J\(s\)\|\}^\{T\}\\right\),\\\\ \\boldsymbol\{F\}\_\{r,s\}&=\\left\(\\diag\\exp\\left\(\\boldsymbol\{\\lambda\}\_\{r,s\}\-\\boldsymbol\{\\lambda\}\_\{r\}\\right\)\\right\)\\tilde\{\\boldsymbol\{F\}\}\_\{r,s\}=\\exp\\left\(\\boldsymbol\{A\}\_\{I\(r\),J\(s\)\}\-\\boldsymbol\{\\lambda\}\_\{r\}\\boldsymbol\{1\}\_\{\|J\(s\)\|\}^\{T\}\\right\)\.\\end\{split\}If

𝒘r,sT=𝟏\|I⁡\(r\)\|T​exp⁡\(𝑨I⁡\(r\),J⁡\(s\)−𝝀r​𝟏\|J⁡\(s\)\|T\)=𝟏\|I⁡\(r\)\|T​𝑭r,s,\\boldsymbol\{w\}\_\{r,s\}^\{T\}=\\boldsymbol\{1\}\_\{\|I\(r\)\|\}^\{T\}\\exp\\left\(\\boldsymbol\{A\}\_\{I\(r\),J\(s\)\}\-\\boldsymbol\{\\lambda\}\_\{r\}\\boldsymbol\{1\}\_\{\|J\(s\)\|\}^\{T\}\\right\)=\\boldsymbol\{1\}\_\{\|I\(r\)\|\}^\{T\}\\boldsymbol\{F\}\_\{r,s\},then𝒘=\[∑r𝒘r,s\]\\boldsymbol\{w\}=\[\\sum\_\{r\}\\boldsymbol\{w\}\_\{r,s\}\]\. We can compute𝒘\\boldsymbol\{w\}and𝒀\\boldsymbol\{Y\}with an outer loop overrr, inner loop overss\. Initialize𝒘=𝟎Nk\\boldsymbol\{w\}=\\boldsymbol\{0\}\_\{N\_\{k\}\}\. The iterationrrworks as follows:

- •Map: Compute\[𝝀r,s\],\[𝑭~r,s\],\[𝒀r,s\]\[\\boldsymbol\{\\lambda\}\_\{r,s\}\],\[\\tilde\{\\boldsymbol\{F\}\}\_\{r,s\}\],\[\\boldsymbol\{Y\}\_\{r,s\}\]in parallel\.
- •Reduce:\(𝝀r,𝒀r\)=𝚛𝚎𝚍𝚞𝚌𝚎⁡\(\[𝝀r,s\],\[𝒀r,s\]\)\(\\boldsymbol\{\\lambda\}\_\{r\},\\boldsymbol\{Y\}\_\{r\}\)=\\mathtt\{reduce\}\\left\(\[\\boldsymbol\{\\lambda\}\_\{r,s\}\],\[\\boldsymbol\{Y\}\_\{r,s\}\]\\right\)\.
- •Compute𝒘r,sT=𝟏\|I⁡\(r\)\|T​𝑭r,s=exp⁡\(𝝀r,s−𝝀r\)T​𝑭~r,s\\boldsymbol\{w\}\_\{r,s\}^\{T\}=\\boldsymbol\{1\}\_\{\|I\(r\)\|\}^\{T\}\\boldsymbol\{F\}\_\{r,s\}=\\exp\(\\boldsymbol\{\\lambda\}\_\{r,s\}\-\\boldsymbol\{\\lambda\}\_\{r\}\)^\{T\}\\tilde\{\\boldsymbol\{F\}\}\_\{r,s\}\. Add𝒘r=\[𝒘r,s\]\\boldsymbol\{w\}\_\{r\}=\[\\boldsymbol\{w\}\_\{r,s\}\]to𝒘\\boldsymbol\{w\}\.

Compared to standard FlashAttention, we need to first reduce alongssin order to obtain𝝀r\\boldsymbol\{\\lambda\}\_\{r\}, keeping𝑭~r,s\\tilde\{\\boldsymbol\{F\}\}\_\{r,s\}and𝝀r,s\\boldsymbol\{\\lambda\}\_\{r,s\}around\.

#### A\.6\.2Summed Attention Weights with FlexAttention

FlexAttention\[[14](https://arxiv.org/html/2608.19920#bib.bib14)\]stands out among fast SDPA codes by allowing the user to configure the computation in several ways \(see also[https://pytorch\.org/blog/flexattention/](https://pytorch.org/blog/flexattention/)\)\. Here we describe how to compute summed attention weights𝒘\\boldsymbol\{w\}alongside the attention output𝒀\\boldsymbol\{Y\}, by calling FlexAttention twice\.

Recall that SDPA computes

𝒀=exp⁡\(𝑨−𝝀​𝟏T\)​𝑽\.\\boldsymbol\{Y\}=\\exp\\left\(\\boldsymbol\{A\}\-\\boldsymbol\{\\lambda\}\\boldsymbol\{1\}^\{T\}\\right\)\\boldsymbol\{V\}\.The summed attention weights are

𝒘=exp⁡\(𝑨−𝝀​𝟏T\)​𝟏T=exp⁡\(𝑨T−𝟏​𝝀T\)​𝟏=exp⁡\(𝑨T\)​𝒗~,𝒗~:=exp⁡\(−𝝀\)\.\\boldsymbol\{w\}=\\exp\\left\(\\boldsymbol\{A\}\-\\boldsymbol\{\\lambda\}\\boldsymbol\{1\}^\{T\}\\right\)\{\}^\{T\}\\boldsymbol\{1\}=\\exp\\left\(\\boldsymbol\{A\}^\{T\}\-\\boldsymbol\{1\}\\boldsymbol\{\\lambda\}^\{T\}\\right\)\\boldsymbol\{1\}=\\exp\\left\(\\boldsymbol\{A\}^\{T\}\\right\)\\tilde\{\\boldsymbol\{v\}\},\\quad\\tilde\{\\boldsymbol\{v\}\}:=\\exp\(\-\\boldsymbol\{\\lambda\}\)\.Up to softmax normalization, we can obtain this by calling a variant of SDPA again, flipping𝑸\\boldsymbol\{Q\}and𝑲\\boldsymbol\{K\}, reverting the attention masking, and passingexp⁡\(−𝝀\)\\exp\(\-\\boldsymbol\{\\lambda\}\)as values\. Importantly, FlexAttention returns𝝀\\boldsymbol\{\\lambda\}with the optionreturn\_aux = AuxRequest\(lse=True\)\. Now, if𝝀~\\tilde\{\\boldsymbol\{\\lambda\}\}denoteslsefor the second call \(with𝑸\\boldsymbol\{Q\}and𝑲\\boldsymbol\{K\}flipped\), then:

𝒘=exp⁡\(𝑨T−𝝀~​𝟏T\+𝝀~​𝟏T\)​𝒗~=\(diag⁡exp⁡\(𝝀~\)\)​exp⁡\(𝑨T−𝝀~​𝟏T\)​𝒗~=exp⁡\(𝝀~\)∘𝒚~\.\\boldsymbol\{w\}=\\exp\\left\(\\boldsymbol\{A\}^\{T\}\-\\tilde\{\\boldsymbol\{\\lambda\}\}\\boldsymbol\{1\}^\{T\}\+\\tilde\{\\boldsymbol\{\\lambda\}\}\\boldsymbol\{1\}^\{T\}\\right\)\\tilde\{\\boldsymbol\{v\}\}=\\left\(\\diag\\exp\(\\tilde\{\\boldsymbol\{\\lambda\}\}\)\\right\)\\exp\\left\(\\boldsymbol\{A\}^\{T\}\-\\tilde\{\\boldsymbol\{\\lambda\}\}\\boldsymbol\{1\}^\{T\}\\right\)\\tilde\{\\boldsymbol\{v\}\}=\\exp\(\\tilde\{\\boldsymbol\{\\lambda\}\}\)\\circ\\tilde\{\\boldsymbol\{y\}\}\.Finally, trying to minimize numerical errors \(we are using 16 bit data types\), we useexp⁡\(−\(𝝀−λ¯​𝟏\)\)\\exp\(\-\(\\boldsymbol\{\\lambda\}\-\\bar\{\\lambda\}\\boldsymbol\{1\}\)\)andexp⁡\(𝝀~−λ¯​𝟏\)\\exp\(\\tilde\{\\boldsymbol\{\\lambda\}\}\-\\bar\{\\lambda\}\\boldsymbol\{1\}\), whereλ¯=Nq−1​𝟏T​𝝀\\bar\{\\lambda\}=N\_\{q\}^\{\-1\}\\boldsymbol\{1\}^\{T\}\\boldsymbol\{\\lambda\}is the mean of𝝀\\boldsymbol\{\\lambda\}\. All in all:

- •\(𝒀,𝝀\)=𝚂𝙳𝙿𝙰⁡\(𝑸,𝑲,𝑽\)\(\\boldsymbol\{Y\},\\boldsymbol\{\\lambda\}\)=\\mathtt\{SDPA\}\(\\boldsymbol\{Q\},\\boldsymbol\{K\},\\boldsymbol\{V\}\),λ¯=Nq−1​𝟏T​𝝀\\bar\{\\lambda\}=N\_\{q\}^\{\-1\}\\boldsymbol\{1\}^\{T\}\\boldsymbol\{\\lambda\}\.
- •\(𝒚~,𝝀~\)=𝚂𝙳𝙿𝙰​\_​𝚛𝚎𝚟​\(𝑲,𝑸,exp⁡\(−\(𝝀−λ¯​𝟏\)\)\)\(\\tilde\{\\boldsymbol\{y\}\},\\tilde\{\\boldsymbol\{\\lambda\}\}\)=\\mathtt\{SDPA\\\_rev\}\(\\boldsymbol\{K\},\\boldsymbol\{Q\},\\exp\(\-\(\\boldsymbol\{\\lambda\}\-\\bar\{\\lambda\}\\boldsymbol\{1\}\)\)\), then𝒘=exp⁡\(𝝀~−λ¯​𝟏\)∘𝒚~\\boldsymbol\{w\}=\\exp\(\\tilde\{\\boldsymbol\{\\lambda\}\}\-\\bar\{\\lambda\}\\boldsymbol\{1\}\)\\circ\\tilde\{\\boldsymbol\{y\}\}\.

Here,𝚂𝙳𝙿𝙰​\_​𝚛𝚎𝚟\\mathtt\{SDPA\\\_rev\}differs from𝚂𝙳𝙿𝙰\\mathtt\{SDPA\}by the attention masking being reversed\. FlexAttention allows to specify the attention mask as𝚋𝚕𝚘𝚌𝚔​\_​𝚖𝚊𝚜𝚔​\(𝚋,𝚑,𝚚​\_​𝚒𝚍𝚡,𝚔𝚟​\_​𝚒𝚍𝚡\)\\mathtt\{block\\\_mask\(b,h,q\\\_idx,kv\\\_idx\)\}\. The mask for𝚂𝙳𝙿𝙰​\_​𝚛𝚎𝚟\\mathtt\{SDPA\\\_rev\}is given by flipping𝚚​\_​𝚒𝚍𝚡\\mathtt\{q\\\_idx\}and𝚔𝚟​\_​𝚒𝚍𝚡\\mathtt\{kv\\\_idx\}in the code for𝚂𝙳𝙿𝙰\\mathtt\{SDPA\}\. Note that FlexAttention supports𝑽\\boldsymbol\{V\}to have a different \(final\) embedding dimension that𝑸,𝑲\\boldsymbol\{Q\},\\boldsymbol\{K\}\. Maybe this even translates in the second call being faster than the first\. All in all, compared to a single FlexAttention call and no attention weights, this is at most twice as expensive\.

### A\.7Sparse Attention and SotA Inference Libraries

In Section[3\.3](https://arxiv.org/html/2608.19920#S3.SS3), we discuss the \(somewhat surprising\) fact that as of today, sparse attention is not much used in real\-world practice, because existing implementations are too slow to be competitive with the state of the art\. While a part of the latency gap between sparse attention and sequence or context parallelism is probably inherent, we argued in Section[3\.3\.1](https://arxiv.org/html/2608.19920#S3.SS3.SSS1)some shortcomings of current sparse attention implementations are easy to eliminate by minor extensions of fast SDPA kernel codes\.

Here, we comment on why sparse attention policies, such as H2O, are not supported in vLLM\[[33](https://arxiv.org/html/2608.19920#bib.bib28)\], the leading fast inference library\. Details are found in[https://github\.com/vllm\-project/vllm/issues/10646](https://github.com/vllm-project/vllm/issues/10646),[https://github\.com/vllm\-project/vllm/issues/12254](https://github.com/vllm-project/vllm/issues/12254),[https://github\.com/vllm\-project/vllm/issues/5751](https://github.com/vllm-project/vllm/issues/5751)\. In vLLM, KV caches are maintained as set of fixed\-sized pages \(or blocks\)\. The main issue is that they require a page to store KV contentacross all heads: KV information for a token is stored in the cache for all heads or for none\. The cited RFCs mention that it would require significant changes to the memory layout and block manager abstractions to change that\. However, for modern sparse attention \(such as H2O\), policiesπ⁡\(b,h,t\)\\pi\(b,h,t\)depend on\(h,t\)\(h,t\)in general: they select different tokens per head\.

As detailed in Section[3\.1](https://arxiv.org/html/2608.19920#S3.SS1), our implementation has no problems with this\. We simply maintain dense buffers of shape\(B,Hk,NC,dh\)\(B,H\_\{k\},N\_\{C\},d\_\{h\}\), where the cache lengthNCN\_\{C\}is fixed independent of context width, and then use𝚝𝚘𝚛𝚌𝚑\.𝚐𝚊𝚝𝚑𝚎𝚛\\mathtt\{torch\.gather\}and𝚝𝚘𝚛𝚌𝚑\.𝚜𝚌𝚊𝚝𝚝𝚎𝚛\\mathtt\{torch\.scatter\}for read and write access\. Whereas PagedAttention requires specific SDPA kernels, we can use existing dense SDPA codes, as long as we cater for causal masking \(see Section[3\.3\.1](https://arxiv.org/html/2608.19920#S3.SS3.SSS1)\)\. While not supported in our current implementation, we could build up KV cache buffers in chunks to cater for sequence lengths shorter thanNCN\_\{C\}, thereby solving the issue of unnecessary pre\-allocations\[[33](https://arxiv.org/html/2608.19920#bib.bib28)\]\. Finally, while𝚝𝚘𝚛𝚌𝚑\.𝚐𝚊𝚝𝚑𝚎𝚛\\mathtt\{torch\.gather\}and𝚝𝚘𝚛𝚌𝚑\.𝚜𝚌𝚊𝚝𝚝𝚎𝚛\\mathtt\{torch\.scatter\}access the buffer in a non\-contiguous way, this is very subdominant to SDPA computations in our experience\. In fact, when calling𝚜𝚌𝚊𝚝𝚝𝚎𝚛⁡\(𝚔𝚎𝚢𝚜,𝚒𝚗𝚍𝚎𝚡,𝚔𝚎𝚢​\_​𝚗𝚎𝚠\)\\mathtt\{scatter\}\(\\mathtt\{keys\},\\mathtt\{index\},\\mathtt\{key\\\_new\}\), the final axis of𝚒𝚗𝚍𝚎𝚡\\mathtt\{index\}is always constant \(so that the final buffer axis of sizedhd\_\{h\}is accessed contiguously\), and optimized𝚜𝚌𝚊𝚝𝚝𝚎𝚛\\mathtt\{scatter\}and𝚐𝚊𝚝𝚑𝚎𝚛\\mathtt\{gather\}kernels could easily be implemented for this case if the PyTorch implementations do not already cater for this special case\. One advantage of PagedAttention over our approach is that they can in principle represent different numbers of tokens per head or batch dimension, which can render sparse attention a bit more flexible\. However, since vLLM requires each page to extend over all heads, this extra flexibility is not supported there\.

As long as highly optimized and widely used inference libraries do not support sparse attention, it may remain underused\. We hope that our work sparks some renewed interest in this direction\.

### A\.8Details on Related Work

Here, we provide additional details about relations of our method with prior work\. OOMB\[[37](https://arxiv.org/html/2608.19920#bib.bib35)\]shares properties with our work, such as chunk\-level processing, activation checkpointing, and efforts to compress KV cache buffers forautograd\. Details on the relationship are as follows:

- •Their implementation is better suited for representing KV caches exactly \(no selection or compression\) that ours\. They implement a paged memory management like\[[33](https://arxiv.org/html/2608.19920#bib.bib28)\], which we do not \(but see Section[A\.7](https://arxiv.org/html/2608.19920#A1.SS7)and comments in Section[3\.1](https://arxiv.org/html/2608.19920#S3.SS1)\)\. However, despite all efforts in CPU offloading and activation checkpointing, they run into the same barrier as\[[38](https://arxiv.org/html/2608.19920#bib.bib34)\], in that the factor for the final chunk depends onallKV cache buffers of all layers, so cannot be represented byautogradon a single device\. At this point, RingAttention\[[42](https://arxiv.org/html/2608.19920#bib.bib36)\]is the method of choice, and it is not clear why their library would improve on implementations such as MS\-SWIFT\[[82](https://arxiv.org/html/2608.19920#bib.bib73)\]\.
- •They deal with activation memory for GPU by activation recomputation, while we use activation checkpointing\. In the former, activations are recomputed during the backward passfrom the start, whereas in the latter, recomputation starts from the most recent checkpoint\. The former is too slow to be useful, so our guess is their code actually uses activation checkpointing\.
- •The most important difference is how they deal with KV cache buffers as nodes in theautogradgraphs, and what this implies for generality\. This is also the biggest challenge we face, and we deal with it by a combination of nested checkpointing, delta encoding of KV cache buffers, and integration intoPyTorchby way of autograd saved tensor hooks \(Section[A\.4\.3](https://arxiv.org/html/2608.19920#A1.SS4.SSS3)\)\. Together with recording and replaying KV cache decisions, thisrenders our implementation fully agnostic to the KV cache policy: it works with any selection or compression policy \(see Section[2](https://arxiv.org/html/2608.19920#S2)for many references\)\. In contrast, they try to hide all nodes representing KV cache content fromautogradaltogether, so that none of this information can be placed in the computation graph\. This is possible only by implementing a number of complex CUDA kernels, in which all inner derivatives w\.r\.t\. these “KV cache nodes” are made explicit\. Apart from substantial derivation and implementation complexity, their approachmust be specialized to the KV cache policy being used\. In fact, their paper only provides results for two specific sparse attention policies \(LSA and DSA\)\. Moreover, their paper is sparse on details how the hiding of KV cache buffers fromautogradworks in practice, since KV cache updates are tightly coupled with SDPA calls\. In the end, their implementation may not be agnostic to SDPA kernels, which given the speed of development of SDPA would be a major drawback\.
- •Both their and our implementation make use of CPU offloading of activations, KV cache buffers, and head gradients\. They claim to have done this asynchronously, as in\[[75](https://arxiv.org/html/2608.19920#bib.bib66)\], which hides latency\. We have also experimented with this, but did no so far achieve significant speedups\. Moreover, asynchronous transfer requires double buffering, which drives up GPU memory requirements\. Still, more effort in this direction is warranted\.

### A\.9Open Source Library KeysAndValues\. Experiments

For all experiments above, fine\-tuning with our method \(columnusin Table[1](https://arxiv.org/html/2608.19920#S4.T1)\) and inference with sparse attention \(all policies\) were done with a new open source library for long context fine\-tuning and inference:𝙺𝚎𝚢𝚜𝙰𝚗𝚍𝚅𝚊𝚕𝚞𝚎𝚜\\mathtt\{KeysAndValues\}\([https://github\.com/awslabs/keys\_values](https://github.com/awslabs/keys_values)\)\.

Apart from efficient code for our fine\-tuning method, the library provides clean and simple abstractions for sparse attention and key\-value caches of limited size\. Among its features are:

- •Long context fine\-tuning on a single GPU \(this work\)\.
- •Several variants of the H2O KV cache policy\[[76](https://arxiv.org/html/2608.19920#bib.bib69)\]\. The library provides a generic implementation for any cache logic of the formπ⁡\(b,h,t\)\\pi\(b,h,t\)which makes use of summed attention weights\.
- •Quantization of KV cache buffers\.
- •Integration of FlexAttention\[[14](https://arxiv.org/html/2608.19920#bib.bib14)\], FlashInfer\[[72](https://arxiv.org/html/2608.19920#bib.bib63)\], FlashAttention\[[10](https://arxiv.org/html/2608.19920#bib.bib8),[52](https://arxiv.org/html/2608.19920#bib.bib44)\]and eager SDPA behind a common multi\-head self\-attention interface\. This includes summed attention weights \(for H2O\-like policies\), as well as a proper𝚋𝚊𝚌𝚔𝚠𝚊𝚛𝚍\\mathtt\{backward\}implementation\.
- •Model implementations and inference code is from LitGPT \([https://github\.com/lightning\-ai/litgpt](https://github.com/lightning-ai/litgpt)\), which allows for almost any Hugging Face checkpoint to be used\. However, while bringing modern KV caching to Hugging Face would require hacking several code filesseparately for every single model, you can apply your KV cache policy or attention approximation to almost all models with few changes of common code\.
- •Support of CPU offloading of KV cache buffers and model weights\.
- •Support of distributed training \(distributed data parallel, CPU offloading of model weights optional\)\. Support of distributed evaluation\.

With this library, we do not intend to compete with vLLM\[[33](https://arxiv.org/html/2608.19920#bib.bib28)\]or SGLang\[[83](https://arxiv.org/html/2608.19920#bib.bib74)\], which include more low level optimizations and support of latest GPU architectures\. Instead, we make it easy for researchers to explore new KV cache policies, post\-time training ideas, or unusual multi\-head self\-attention approximations, providing clean abstractions of these concepts which can be used and extended without having to deal with intricate implementation details of existing high\-performance libraries\.

#### A\.9\.1Running Our Experiments

Once𝙺𝚎𝚢𝚜𝙰𝚗𝚍𝚅𝚊𝚕𝚞𝚎𝚜\\mathtt\{KeysAndValues\}has been properly installed, the training runs for our method can be reproduced as follows\. You need to be on an instance with at least four Nvidia A100 GPUs with 40 GB of RAM\. We usedAWS EC2 p4d\.24xlargeinstances, which have 8 A100 GPUs, running two experiments in parallel on each instance\.

exportDATASET\_KEY=”nq”;\\

exportDATASET\_SIZE=”128k”;\\

exportPOLICY\_NAME=”h2o\-orig”;\\

exportCACHE\_LENGTH=”32768”;\\

exportCHUNK\_SIZE=”2048”;\\

exportEVAL\_STEPS=10;\\

CUDA\_VISIBLE\_DEVICES=”0,1,2,3”\\

PYTORCH\_ALLOC\_CONF=expandable\_segments:True\\

KEYSVALS\_LOG\_DIR=”\./finetune/helmet\_$\{DATASET\_KEY\}\_$\{DATASET\_SIZE\}/$\{POLICY\_NAME\}\_cs$\{CHUNK\_SIZE\}/logs”\\

python3keys\_values/\_\_main\_\_\.pyfinetune\_long\_lora\\

Qwen/Qwen3\-4B\-Instruct\-2507\\

–out\_dir\./finetune/helmet\_$\{DATASET\_KEY\}\_$\{DATASET\_SIZE\}/$\{POLICY\_NAME\}\_cs$\{CHUNK\_SIZE\}\\

–precisionbf16\-true\\

–verbosesome\\

–devices4\\

–dataHelmet\\

–data\.dataset\_key$\{DATASET\_KEY\}\\

–data\.max\_length$\{DATASET\_SIZE\}\\

–data\.metadata\_dir\./data\\

–data\.trainloader\_longest\_firstTrue\\

–train\.save\_interval$\{EVAL\_STEPS\}\\

–train\.micro\_batch\_size2\\

–train\.epochs5\\

–train\.average\_loss\_per\_batchTrue\\

–eval\.interval$\{EVAL\_STEPS\}\\

–eval\.initial\_validationTrue\\

–eval\.use\_sample\_metricFalse\\

–kv\_cache\.cache\_length$\{CACHE\_LENGTH\}\\

–kv\_cache\.chunk\_size$\{CHUNK\_SIZE\}\\

–kv\_cache\.name$\{POLICY\_NAME\}\-torch\-quantized8\\

–grad\.layers\_per\_cell1\\

–grad\.layercp\_qnamedefault\\

–grad\.cachecp\_qnametorch\-quantized8\\

–grad\.chunks\_per\_cell\_multiplier1\\

–optimizer\.nameAdamW\\

–optimizer\.learning\_rate0\.0005

Once all desired training runs have finished, evaluations \(on the test sets\) can be run as follows\.

CUDA\_VISIBLE\_DEVICES=”0,1,2,3”\\

PYTORCH\_ALLOC\_CONF=expandable\_segments:True\\

KEYSVALS\_LOG\_DIR=”\./finetune/evaluation/myruns/logs”\\

python3keys\_values/\_\_main\_\_\.pyeval\_long\_ext\\

\./myruns\.yaml\\

–verbosesome\\

–devices4\\

–batch\_size2\\

–use\_sample\_metricTrue\\

–sample\_metric\_max\_generated\_tokens20\\

–num\_store\_generated\_samples1000

Here,myruns\.yamlis a YAML file containing entries of this form:

\-out\_dir:\./finetune/helmet\_nq\_64k/h2o\_cs2048

model\_type:lora

eval\_tasks:

\-step\-000420

For each setup, evaluations can be run for different checkpointsstep\-000\*\*\*stored alongside training\. In our experiments, for each setup, we select the checkpoint which minimize validation loss\. We refer toREADME\.mdfor further details on how to aggregate evaluation results and create result tables\.