One-sentence Summary

NVIDIA et al. propose SparDA, a decoupled sparse attention architecture that introduces a Forecast projection for lookahead KV-block prefetch overlapping computation, where using one Forecast head per GQA group reduces selection overhead, adds $< 0 . 5 \%$<0.5% parameters, trains only the Forecast projections by matching the original selector's attention distribution, and on two sparse-pretrained 8B models matches or slightly improves accuracy while achieving up to $1.25\times$1.25× prefill speedup, $1.7\times$1.7× decode speedup over the sparse-attention offload baseline, and up to $5.3\times$5.3× higher decode throughput than the non-offload sparse baseline.

Key Contributions

• SparDA is a decoupled sparse attention architecture that adds a Forecast projection per layer to predict the KV blocks needed by the next layer, enabling lookahead prefetching that overlaps CPU-to-GPU data transfer with current-layer computation.
• The design uses one Forecast head per GQA group, reducing sparse selection overhead versus multi-head selectors while adding under 0.5% parameters and training only the Forecast projections by matching the original selector's attention distribution.
• On two sparse-pretrained 8B models, SparDA achieves up to 1.25× prefill and 1.7× decode speedup over a sparse-attention offload baseline, and by enabling larger batch sizes it reaches up to 5.3× higher decode throughput than a non-offload sparse baseline.

Introduction

Long-context LLM serving strains infrastructure with high compute for prefill attention, memory bandwidth pressure during decoding, and growing KV cache capacity demands. Sparse attention alleviates compute and bandwidth but leaves capacity pressure unresolved, and offloading the KV cache to CPU memory with selective fetching introduces PCIe latency bottlenecks. Prior lookahead prefetching methods like InfiniGen attempt to hide this latency by using the raw hidden state as a proxy for future attention, but they become inaccurate when layer-to-layer similarity breaks down. Meanwhile, sparse selection itself incurs O(T²) overhead that can dominate as context length grows, and existing optimizations often trade accuracy for speed. The authors propose SparDA, which decouples sparse selection from attention by training a lightweight Forecast projection to predict next-layer KV block selections directly, enabling accurate asynchronous prefetching and a compact indexer that reduces selection overhead without retraining the base model.

Method

The authors build SparDA on top of InfLLM-V2, which represents a block-sparse attention architecture. InfLLM-V2 decomposes attention into initial tokens, a local sliding window, and top-$k$k selected blocks. It merges selected attention and sliding-window attention into a single sparse attention module with shared KV projections. To compute block-level relevance scores, it employs a parameter-free, coarse-to-fine compression pipeline where keys are mean-pooled into overlapping compressed representations.

To enable KV cache prefetching that overlaps with layer execution for latency hiding and to reduce the overhead of sparse selection, the authors introduce a single architectural change: a fourth per-layer projection called the Forecast $\mathbf{F}_l$Fl​. This projection is produced alongside the standard $\mathbf{Q}_l$Ql​, $\mathbf{K}_l$Kl​, and $\mathbf{V}_l$Vl​ to drive a compact sparse-selection indexer. By decoupling sparse selection from the attention query, the Forecast can be computed one layer ahead.

As shown in the figure below:

In the baseline InfLLM-V2, the query $\mathbf{Q}_l$Ql​ drives both top-$k$k selection and sparse attention within the same layer, placing selection on the attention critical path. SparDA decouples these roles. The linear projection $\phi_l$ϕl​ in each layer produces the additional Forecast $\mathbf{F}_l$Fl​: $(\mathbf{Q}_l, \mathbf{K}_l, \mathbf{V}_l, \mathbf{F}_l) = \phi_l(\mathbf{X}_l).$(Ql​,Kl​,Vl​,Fl​)=ϕl​(Xl​). The Forecast $\mathbf{F}_l$Fl​ is used for top-$k$k sparse selection, while $\mathbf{Q}_l$Ql​ is reserved solely for the sparse attention computation. Specifically, $\mathbf{F}_l$Fl​ scores against the compressed keys $\widetilde{\mathbf{K}}_{l+1}$Kl+1​ of the next layer, and the resulting top-$k$k blocks are merged with the initial and local blocks to form the attended set for layer $l+1$l+1. During prefill, the Forecast replaces the expensive multi-head selector with an indexer that uses one Forecast head per GQA group. During decode, $\mathbf{F}_l$Fl​ predicts the blocks for the next layer early enough to asynchronously prefetch the selected KV blocks from CPU memory, effectively hiding the PCIe transfer behind the current layer compute.

For training, the authors add SparDA to existing sparse-pretrained models by training only the Forecast projections, which adds a negligible number of parameters. Inspired by DeepSeek DSA, the Forecast indexer is trained to match target block-attention scores via Kullback-Leibler divergence after top-$k$k restriction and renormalization. The target and predicted scores are defined as: $\mathbf{S}_{l,m}^{\mathrm{tgt}} = \sum_{h=1}^{G} \operatorname{softmax}\left(\mathbf{Q}_{l,m,h} \widetilde{\mathbf{K}}_{l,m}^{\mathrm{tgt}\top} / \tau\right), \quad \mathbf{S}_{l,m}^{\mathrm{pred}} = \operatorname{softmax}\left(\mathbf{F}_{l-1,m} \widetilde{\mathbf{K}}_{l,m}^{\mathrm{pred}\top} / \tau\right),$Sl,mtgt​=∑h=1G​softmax(Ql,m,h​Kl,mtgt⊤​/τ),Sl,mpred​=softmax(Fl−1,m​Kl,mpred⊤​/τ), where the predicted branch uses the Forecast from the previous layer indexer to produce a single score per KV head directly. The training loss computes the KL divergence over a top-$k$k partitioned distribution, focusing the indexer on the relative ranking within the selected set while constraining the total mass on non-selected blocks through a rest bucket.

To provide fine-grained training supervision, the authors use a smaller kernel size and stride for the target compressed keys compared to the inference-time configuration. This finer-grained compression provides a higher-resolution supervision signal, allowing the indexer to learn sharper selection. The target score tensor is then downsampled to the standard grid before computing the KL loss.

At inference, SparDA exploits the early availability of the Forecast through an asynchronous prefetch pipeline. The authors employ a persistent Triton kernel based on Unified Virtual Addressing to fetch selected KV blocks from pinned CPU memory on a dedicated CUDA stream. This kernel maintains a small fixed set of GPU thread blocks to continuously process block-transfer tasks within a single launch, reducing overhead and limiting interference with the main compute stream. To balance prefetch throughput and layer execution speed, SparDA adopts a batch-adaptive CTA allocation heuristic. This approach allocates more Cooperative Thread Arrays as the batch size grows to maximize overall throughput for the underlying hardware.

Experiment

The evaluation compares SparDA against dense, sparse, and InfiniGen baselines on two 8B models using long-context benchmarks, measuring accuracy and efficiency on NVIDIA H100 and A100 GPUs. SparDA’s learned Forecast indexer improves accuracy over the sparse baseline, especially on reasoning and length generalization, while also reducing block-selection overhead to accelerate prefill and decode. Its lookahead prefetch design overlaps KV cache transfers with computation, yielding substantial decode throughput gains at long contexts. In contrast, the training-free InfiniGen method suffers significant accuracy loss and is slower due to CPU-side gather operations.

SparDA improves average accuracy over the Sparse baseline on both models, with the largest gains on reasoning and RULER benchmarks. While dense attention still leads on HELMET and RULER, SparDA narrows the gap on RULER and surpasses dense on reasoning tasks. InfiniGen consistently underperforms Sparse across all metrics. SparDA achieves higher overall average accuracy than Sparse, with a modest gain on MiniCPM4.1-8B and a larger improvement on NOSA-8B. Reasoning accuracy benefits the most, with SparDA outperforming even dense attention on both models. RULER scores improve with SparDA, reducing the gap to dense attention, while LongBench remains essentially unchanged. HELMET accuracy slightly decreases on MiniCPM4.1-8B but increases on NOSA-8B. InfiniGen trails Sparse across all benchmarks, with the largest drops on HELMET and RULER.

SparDA consistently improves RULER accuracy over the Sparse baseline at every evaluated sequence length for both MiniCPM4.1-8B and NOSA-8B. The advantage is most pronounced on NOSA-8B, where the gap widens steadily from 32K to 128K, while MiniCPM4.1-8B shows a more variable but always positive margin. This indicates the learned Forecast indexer generalizes at least as well as the training-free selector, and in some regimes better. SparDA outperforms Sparse on RULER at all sequence lengths from 32K to 128K across both models. On NOSA-8B, the accuracy gap grows monotonically with context length, reaching +4.3 points at 128K.

Dense attention achieves the highest prefill throughput at short sequences (32K) but its quadratic scaling causes a steep decline as context length increases. SparDA surpasses all methods at longer contexts, leading from 64K on MiniCPM4.1-8B and from 96K on NOSA-8B, with speedups over both Sparse and Dense that grow with sequence length. The improvement is driven by a more efficient block selection step, while KV cache offloading itself introduces negligible prefill overhead. Dense attention leads at 32K but its throughput drops sharply with longer sequences, falling behind sparse methods. SparDA achieves the highest prefill throughput from 64K onward on MiniCPM4.1-8B and from 96K onward on NOSA-8B. On MiniCPM4.1-8B at 128K, SparDA delivers a 1.25× speedup over Sparse and a 2.11× speedup over Dense. The speedup comes from a block selection mechanism that scales more favorably with sequence length than the baseline multi-head selector. KV cache offloading has negligible impact on prefill throughput, as seen by the near-identical numbers for Sparse with and without offload. All sparse methods are slower on NOSA-8B because its query-agnostic eviction head adds extra work.

SparDA's offloading design allows decode throughput to scale to much larger batch sizes than non-offload baselines, which run out of memory beyond batch 16 at long contexts. At 32K context, SparDA overtakes Sparse at batch 128 on MiniCPM4.1-8B, while InfiniGen is consistently slower because its CPU-side KV gather becomes a bottleneck. The largest relative gains occur at middle batch sizes where prefetch overlap is most effective, and SparDA achieves up to 1.69× speedup over Sparse at 128K. At 32K, SparDA reaches 1358.2 tok/s at batch 128 on MiniCPM4.1-8B, exceeding Sparse's 1225.1 tok/s, while Dense without offload runs out of memory after batch 16. Offloading lets SparDA run batches up to 128, whereas non-offload baselines OOM past batch 16 at long contexts, leading to throughput gains of up to 9.21× over Dense without offload. InfiniGen's decode throughput is far lower (e.g., 222.4 tok/s at batch 128, 32K) because it gathers KV blocks on the CPU before transferring them to the GPU. The largest speedups appear at middle batch sizes (8–16), where the overlap of KV prefetch with layer execution is roughly balanced.

The evaluation compares SparDA against sparse and dense baselines, as well as InfiniGen, on two models (MiniCPM4.1-8B and NOSA-8B) using accuracy benchmarks, prefill throughput, and decode throughput. SparDA consistently improves accuracy over the sparse baseline, with the largest gains on reasoning and RULER tasks, and it even surpasses dense attention on reasoning while narrowing the gap on RULER. In throughput, SparDA achieves the highest prefill speeds at long contexts due to efficient block selection, and its KV cache offloading allows decoding at much larger batch sizes, yielding substantial speedups over non-offloading methods. InfiniGen trails in both accuracy and throughput across all settings.