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.

  • Method

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

  • 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.

  • Takeaways & Limitations

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

  • Takeaways & Limitations

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

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.
  • CPU KV-cache fetching over PCIe becomes a bottleneck because it is slower than GPU-memory access during long-context decoding.
  • 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.
  • 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.
  • 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.

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.
  • 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.
  • Trainable sparse attention: 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.
  • 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.

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.
  • Unified sparse attention: 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.
  • 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.
  • Three-stage block representation: 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.
  • Three-stage block representation: 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.

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.
  • 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.
  • 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.
  • 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.
  • 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.

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.
  • 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.
  • 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.
  • 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.
  • 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.

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.
  • 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.
  • 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.

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.
  • Conclusion: 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.
  • 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.

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.
  • Prefill: 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).
  • 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.
  • 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.

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.

  • Only the Forecast projection weights are trained, while all backbone parameters remain frozen.
  • 32 NVIDIA H100 GPUs are used for indexer training in both model experiments.
  • 48 hours is the completion time for MiniCPM4.1-8B at 64K, versus 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.
  • Model configuration: 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.
  • 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.
  • 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.
  • 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.

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.
  • 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.
  • 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.
  • A100 throughput results: 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.
  • 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.
  • Per-task accuracy breakdowns: Per-task breakdowns provide HELMET, LongBench, RULER, and reasoning scores that complement the aggregated benchmark results.RULER is evaluated at each model’s native maximum context, while reasoning accuracy uses a GPT-5.2 judge.
Loading 2606.04511v1…