PaperlayerView exact revision on arXiv

Source-linked AI summary

SparDA: Sparse Decoupled Attention for Efficient Long-Context LLM Inference

Yaosheng Fu, Guangxuan Xiao, Xin Dong, Song Han, Oreste Villa

arXiv:2606.04511v1cs.CLcs.LG

TL;DR

Long-context sparse attention still faces KV-cache and selection bottlenecks. SparDA decouples selection through a Forecast projection, matching or slightly improving accuracy while delivering up to 1.25× prefill, 1.7× decode, and 5.3× decode-throughput gains.

  • Problem

    Long-context sparse attention still faces growing KV-cache capacity pressure, CPU-offload PCIe bottlenecks, and selection costs that increase with sequence length.

    Source: §1 · p. 2, §5.3 · p. 7

  • Method

    SparDA adds a decoupled Forecast projection that predicts next-layer KV blocks, enabling lookahead prefetching and a compact GQA-level selection indexer.

    Source: §1 · p. 2, §4.1 · p. 4

  • Results

    Up to 1.25× prefill speedup, 1.7× decode speedup, and 5.3× higher decode throughput are reported while maintaining comparable or slightly better accuracy.

    Source: §1 · p. 2

  • Takeaways & Limitations

    SparDA treats sparse selection as a schedulable signal, making sparse attention more compatible with overlapped KV-cache offloading.

    Source: §7 · p. 9

    SparDA’s accuracy remains bounded by the quality of its underlying sparse-attention backbone.

    Source: §6 · p. 9

Abstract

from arXiv · show

Sparse attention reduces compute and memory bandwidth for long-context LLM inference. However, two key challenges remain: (1) KV cache capacity still grows with sequence length, and offloading to CPU memory introduces a PCIe transfer bottleneck; (2) the sparse selection step itself retains $O(T^2)$ complexity and can dominate attention cost at long contexts. We propose SparDA, a decoupled sparse attention architecture that introduces a fourth per-layer projection, the Forecast, alongside Query, Key, and Value. The Forecast predicts the KV blocks needed by the next layer, enabling lookahead selection that overlaps CPU-to-GPU prefetch with current-layer execution. Because Forecast is decoupled from the attention query, our GQA implementation uses one Forecast head per GQA group, reducing selection overhead versus the original multi-head selector. SparDA adds $<$0.5% parameters and trains only the Forecast projections by matching the original selector's attention distribution. On two sparse-pretrained 8B models, SparDA matches or slightly improves accuracy and delivers up to 1.25$\times$ prefill speedup and 1.7$\times$ decode speedup over the sparse-attention offload baseline. By enabling larger feasible batch sizes on a single GPU, SparDA further reaches up to 5.3$\times$ higher decode throughput than the non-offload sparse baseline. Our source code is available at https://github.com/NVlabs/SparDA.

1. Introduction

SparDA addresses long-context inference bottlenecks from KV-cache offloading and sparse selection by decoupling next-layer selection from current-layer attention. It uses a lightweight Forecast projection for lookahead prefetching and reports speedups over sparse-attention offloading while preserving comparable accuracy.

  • Long-context serving faces high prefill attention compute, decode memory-bandwidth pressure, and growing decode KV-cache capacity demands.Source: §1 · pp. 1–2
  • CPU KV-cache fetching over PCIe becomes a bottleneck because it is slower than GPU-memory access during long-context decoding.Source: §1 · p. 2
  • Sparse selection remains O(T^2) even after sparse attention reduces complexity from O(T^2) to O(T), potentially dominating attention as context grows.Source: §1 · p. 2
  • SparDA adds a Forecast projection that predicts the next layer’s sparse selection, allowing CPU-to-GPU KV-cache prefetching to overlap current-layer execution.Source: §1 · p. 2
  • Up to 1.25× prefill and 1.7× decode speedups over the sparse-attention offload baseline accompany comparable or slightly better accuracy on evaluated sparse-pretrained 8B models.Source: §1 · p. 2

2. Related Work

Prior work reduces long-context attention cost through KV cache compression, training-free or trainable sparse selection, and KV-cache offloading. These approaches improve efficiency but retain accuracy, cache-capacity, or selection-latency tradeoffs.

KV cache compression and training-free sparse attention

  • KV cache compression discards or merges tokens, whereas training-free sparse attention selects query-specific tokens or blocks without removing the full cache.Dynamic sparse patterns in MInference, FlexPrefill, and XAttention accelerate prefill.Source: §2 · p. 2

Trainable sparse attention

  • Trainable sparse attention learns selection during pretraining using parameter-free, self-distilled, mean-pooled-key, or learnable-gating mechanisms.Representative methods include InfLLM-V2, MoBA, SeerAttention, and NSA.Source: §2 · pp. 2–3
  • DeepSeek Sparse Attention and DeepSeek-V4 advance token-level and compressed sparse attention, but longer contexts still leave absolute KV-cache capacity as a bottleneck.DeepSeek-V4 interleaves CSA with HCA and cuts per-token KV cache by approximately 10× over DeepSeek-V3.2.Source: §2 · p. 3

Sparse attention with KV cache offloading

  • KV-cache offloading methods manage GPU-resident entries through eviction, recall, or low-rank proxies, while NOSA performs synchronous per-layer CPU-to-GPU transfers.SparseServe and HiSparse reactively swap offloaded entries, whereas NOSA reduces transfer volume rather than hiding latency.Source: §2 · p. 3

3. Preliminary

InfLLM-V2 provides SparDA’s block-sparse foundation, combining initial tokens, a local sliding window, and top-k selected blocks. Its parameter-free coarse-to-fine selector scores compressed key representations and aggregates relevance within GQA groups before max-pooling to block scores.

Unified sparse attention

  • InfLLM-V2 combines initial blocks, local sliding-window blocks, and top-k selected blocks in a unified sparse-attention module.SparDA modifies only the top-k selection path, so the design is expected to extend to methods with the same initial/local/top-k structure.Source: §3.1 · p. 3
  • Full attention is computed only over the attended KV blocks, reducing per-query cost from O(T) to O(|B_l(i)| · B).T denotes sequence length, and B denotes block size.Source: §3.1 · p. 3

Three-stage block representation

  • The selector uses a parameter-free, three-stage coarse-to-fine pipeline that mean-pools keys into overlapping compressed representations.The first stage uses a kernel size and stride for compression before relevance scoring.Source: §3.1 · p. 3
  • Each GQA group shares one KV head, and per-head relevance scores are summed within the group to form a shared importance score.This group-level aggregation produces shared scores across the query heads in each GQA group.Source: §3.1 · p. 3
  • Max-pooling over each shared importance score produces the block-level score used for top-k selection.The process scores several overlapping sub-blocks rather than relying on a single coarse pooling operation, preserving finer-grained information.Source: §3.1 · p. 3

4. Method

SparDA decouples sparse selection from attention by adding a Forecast projection that predicts next-layer KV blocks, reducing selection overhead and enabling overlapped CPU-to-GPU prefetching. It trains only lightweight Forecast projections with a KL-based objective tailored to sparse block selection.

Decoupled sparse selection

  • SparDA adds a fourth per-layer Forecast projection alongside Q, K, and V to select next-layer KV blocks while Q performs attention.Forecast scores compressed keys from layer l+1, whose top-k blocks are merged with initial and local blocks.Source: §4.1 · p. 4

Forecast indexer

  • One Forecast head per GQA group reduces selector cost versus using all query heads, while preserving same-layer selection for the first layer.The final layer’s Forecast is unused because it predicts no subsequent layer.Source: §4.1 · p. 5

Decoupled sparse selection

  • During decode, Forecast predicts the next layer’s sparse pattern early enough to prefetch selected CPU-resident KV blocks while the current layer executes.During prefill, the same decoupling saves selection cost, but all keys already reside on the GPU.Source: §4.1 · p. 5, §4.3 · p. 6

Training

  • 33.5M parameters, or 0.41% of the 8B total, are added by training only Forecast projections on sparse-pretrained MiniCPM4.1-8B and NOSA-8B.This avoids retraining the base model and keeps SparDA lightweight.Source: §4.2 · p. 5

Training objective

  • KL divergence trains the Forecast indexer against top-k target block-attention distributions, retaining individual selected-block scores and aggregating remaining mass into a rest bucket.The target uses finer-grained (2, 1) compression than the inference-time (32, 16) window to provide sharper supervision.Source: §4.2 · pp. 5–6

5. Experiments

Across two sparse-pretrained 8B models, SparDA preserves or improves accuracy over Sparse, including stronger long-context RULER performance, while reducing selection overhead and improving H100 prefill and decode throughput through lookahead KV prefetch overlap.

Accuracy

  • SparDA matches or improves Sparse on aggregate accuracy across all four benchmark families, with average gains of +0.3 on MiniCPM4.1-8B and +2.3 on NOSA-8B.The gains include reasoning improvements of +1.1 and +6.5, respectively.Source: §5.2 · p. 7

Length generalization

  • At extended lengths, SparDA outperforms Sparse on RULER for both models, with NOSA-8B’s advantage widening from +1.7 at 32K to +4.3 at 128K.MiniCPM4.1-8B gains are +1.5 at 32K, +0.5 at 64K, +2.1 at 96K, and +1.1 at 128K.Source: §5.2 · p. 7

Attention breakdown

  • SparDA reduces block-selection cost during both prefill and decode, targeting the selection overhead that remains costly in sparse attention.The attention breakdown measures block selection separately from block-sparse attention on MiniCPM4.1-8B at batch size 4.Source: §5.3 · p. 7, Figure 3 · p. 8

Prefill throughput

  • From 64K onward, SparDA consistently outperforms all other methods on MiniCPM4.1-8B prefill, reaching up to 1.25× speedup.Prefill throughput is reported on an NVIDIA H100 GPU at batch size 4.Source: §5.3 · p. 8, Table 3 · p. 8

Decode throughput

  • At 128K, SparDA reaches up to 1.69× speedup over Sparse on MiniCPM4.1-8B and 1.40× on NOSA-8B within the offload regime.Lookahead overlaps CPU-to-GPU KV prefetch with current-layer compute; SparDA’s peak throughput exceeds Sparse’s at every context length, even when Sparse runs at a larger batch.Source: §5.3 · pp. 8–9

6. Limitations

SparDA is an add-on to an existing sparse attention backbone rather than a standalone sparse attention method. Its accuracy therefore depends on the quality of the underlying sparse attention method, although the decoupled indexer principle is expected to extend beyond block-level selection.

Method scope

  • SparDA improves inference efficiency by replacing the selection path without changing the sparse attention computation or underlying sparse pattern.The decoupled indexer is an add-on to an existing sparse attention backbone.Source: §6 · p. 9

Accuracy dependence

  • SparDA’s accuracy is bounded by the quality of the base sparse attention method.Because SparDA does not alter the underlying sparse attention method, its accuracy inherits that method’s limitations.Source: §6 · p. 9

Generalizability

  • The decoupled indexer principle is not specific to block-level selection and is expected to extend beyond it.The passage identifies broader applicability as an expectation rather than a demonstrated result.Source: §6 · p. 9

7. Conclusion

SparDA treats sparse selection as a trainable, schedulable signal decoupled from the current attention query. Its one-layer-ahead Forecast projection overlaps KV transfer with execution, reduces selector cost through GQA-level indexing, and motivates extensions to token-level sparse attention and larger sparse-pretrained models.

Conclusion

  • SparDA decouples sparse selection from the current attention query by using Forecast to predict block choices one layer ahead.This reframes sparse selection as a trainable, schedulable signal.Source: §7 · p. 9
  • Forecast enables CPU-to-GPU KV transfer to overlap with current-layer execution while a compact GQA-level indexer reduces selector cost.These mechanisms jointly address runtime scheduling and selection overhead.Source: §7 · p. 9

Future work

  • Extending decoupled lookahead selection to token-level sparse attention and larger sparse-pretrained models is identified as future work.The conclusion presents both extensions as natural directions.Source: §7 · p. 10

A. Algorithm Pseudocode

SparDA selects pooled KV blocks using Forecast representations and prepares the next layer’s data while executing the current layer. Its prefill and decode procedures handle initialization, one-layer-ahead selection, and incomplete prefetches explicitly.

Prefill

  • Prefill takes hidden states X_l and the prior layer’s Forecast F_{l-1}, then outputs X_{l+1} and Forecast F_l for the next layer.The algorithm input and output explicitly connect each layer’s Forecast to the following layer.Source: A. Algorithm Pseudocode · p. 13
  • Pooled prefill selection forms B_l from initial and local blocks plus the top-k blocks scored by F_{l-1}K_l^T.The selected set is B_init ∪ B_local ∪ f_top(F_{l−1}K_l^T, k).Source: A. Algorithm Pseudocode · p. 13

Initialization

  • For l=0, SparDA uses the current layer’s Forecast because no previous Forecast exists.In decode, the layer-0 KV cache stays on GPU, so prefetching starts from subsequent layers.Source: A. Algorithm Pseudocode · p. 13

Decode

  • Decode performs one-layer-ahead pooled selection for B_{l+1} using F_lK_{cache,l+1}^T, while KV prefetch begins during layer l−1.If B_l’s KV entries are incomplete, the algorithm waits for them before proceeding.Source: A. Algorithm Pseudocode · p. 13

B. Training Details

Training updates only the Forecast projection while freezing the backbone. Experiments use 32 NVIDIA H100 GPUs, completing within 48 hours for MiniCPM4.1-8B at 64K and 24 hours for NOSA-8B at 32K.

C. Evaluation Details

The evaluation specifies matched sparse-attention configurations for MiniCPM4.1-8B and NOSA-8B, benchmark-specific long-context extensions and decoding protocols, and throughput measurements across GPU–CPU offloading setups. SparDA changes only its Forecast projections while preserving the sparse baseline’s remaining parameters.

Model configuration

  • 64-token blocks, 1 initial block, compression kernel 32 with stride 16, and a 96-block top-k budget configure both models.MiniCPM4.1-8B uses a 2048-token local window, while NOSA-8B uses 1024 tokens and 24 query-aware blocks.Source: C. Evaluation Details · p. 13
  • The SparDA indexer loads trained Forecast projections, including the main and layer-0 same-layer projections; all other parameters match the sparse baseline.These Forecast projections provide the only model-parameter changes described for the indexer.Source: C. Evaluation Details · p. 14

Long-context extension

  • 64K and 32K are the native context limits for MiniCPM4.1-8B and NOSA-8B, respectively, with positional-encoding extensions applied beyond those limits.MiniCPM4.1-8B uses checkpoint-provided 128K-validated LongRoPE factors; NOSA-8B increases rope_theta from 10,000.Source: C. Evaluation Details · p. 14

Benchmark protocols

  • Reasoning evaluations comprise 500 MATH-500 problems and 30 problems each for AIME 2024 and AIME 2025, using sampling and long generation budgets.MiniCPM4.1-8B uses temperature 0.9 and top_p 0.95, while NOSA-8B uses temperature 1.0 and top_p 1.0.Source: C. Evaluation Details · p. 14

Efficiency

  • Throughput is measured on 80 GB H100 and A100 systems with PCIe Gen5×16 and Gen4×16 connections, respectively, using two measured runs after discarding one warmup.For offloading, the full KV cache resides in pinned CPU memory, while compressed keys and the layer-0 cache remain on GPU.Source: C. Evaluation Details · pp. 14–15

D. Additional Results

Additional ablations show that finer compression windows improve benchmark averages and adaptive prefetch CTA allocation remains near-optimal across batch sizes. On A100, SparDA preserves substantial prefill and decode speedups, while breakdowns attribute decode gains to lower selection cost and overlapped prefetch.

Compression window ablation

  • Finer compression windows improve average scores on both models, with the largest gains on RULER and Reasoning.On MiniCPM4.1-8B, RULER and reasoning improve by +3.0 and +2.2; on NOSA-8B, they improve by +2.5 and +1.8.Source: D. Additional Results · p. 15, Table 6 · p. 15

Prefetch CTA allocation ablation

  • Adaptive CTA allocation matches or stays within 4% of the best fixed configuration at every batch size.The heuristic uses 16 CTAs below batch size 32 and 32 CTAs otherwise; 16 is best at B8–B16, while 32 or 64 is preferred for larger batches.Source: D. Additional Results · p. 15, Table 7 · p. 15

A100 throughput results

  • 1.23× prefill speedup over Sparse and 1.81× over Dense is achieved on MiniCPM4.1-8B at 128K on A100.SparDA also reaches 1.16× over Sparse on NOSA-8B, and the gains are consistent with H100 results.Source: D. Additional Results · p. 15, Table 8 · p. 15
  • 1.55× decode speedup over Sparse is achieved on MiniCPM4.1-8B at 128K and B8, while NOSA-8B reaches 1.33× at 128K and B16.These gains are slightly lower than on H100, peak at middle batch sizes, and are lower for NOSA-8B because of reduced KV fetch traffic.Source: D. Additional Results · pp. 15–16, Table 9 · p. 16

Decode speedup breakdown

  • The Forecast indexer improves decode throughput even without prefetch by reducing block-selection latency, outperforming Sparse at B4 on both models.The breakdown isolates reduced block-selection cost and overlapped CPU-to-GPU prefetch as SparDA’s two decode-speedup sources.Source: D. Additional Results · p. 16, Table 10 · p. 16

Per-task accuracy breakdowns

Related papers

Loading 2606.04511v1…