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

Summary (Overview)

• Core contribution: SparDA introduces a fourth per-layer projection called the Forecast (alongside Query, Key, Value) that predicts the KV blocks needed by the next layer, decoupling sparse selection from attention computation.
• Key innovation: The one-layer-ahead prediction enables lookahead prefetching of KV cache blocks from CPU memory, overlapping PCIe transfers with current-layer execution to hide latency.
• Efficiency gains: Achieves up to 1.25× prefill speedup and 1.7× decode speedup over the sparse-attention offload baseline, and up to 5.3× higher decode throughput than the non-offload sparse baseline by enabling larger feasible batch sizes.
• Lightweight integration: Adds only <0.5% parameters (33.5M for 8B models) and trains only the Forecast projections via KL divergence, without retraining the base model.
• Accuracy preservation: Matches or slightly improves accuracy over the sparse baseline on long-context benchmarks (HELMET, LongBench, RULER, and long-reasoning suites) across two sparse-pretrained 8B models.

Introduction and Theoretical Foundation

Background and Motivation

Long-context LLM inference faces three major efficiency challenges:

1. High compute demand for attention in the prefill phase
2. High memory bandwidth pressure for attention in the decode phase
3. High KV cache capacity pressure in the decode phase

Sparse attention reduces compute and memory bandwidth, but two key challenges persist:

• KV cache capacity: Still grows with sequence length; offloading to CPU memory introduces a PCIe transfer bottleneck
• Sparse selection overhead: The selection step itself retains $O(T^2)$ complexity and can dominate attention cost at long contexts

Theoretical Foundation

The paper builds on InfLLM-V2, a block-sparse attention method with three components:

$$\mathcal{B}_l(i) = \mathcal{B}_{\text{init}} \cup \mathcal{B}_{\text{local}}(i) \cup \mathcal{B}_{\text{topk}}(i) \tag{1}$$

where $\mathcal{B}_{\text{init}}$ covers initial blocks, $\mathcal{B}_{\text{local}}(i)$ covers local sliding window blocks, and $\mathcal{B}_{\text{topk}}(i)$ contains top-𝑘 blocks selected by compression scores.

InfLLM-V2 uses a three-stage block representation with mean-pooled compressed keys:

$$\widetilde{\mathbf{K}}_{l,j} = \text{Mean}(\mathbf{K}_{l, j \cdot s_{C_1}: j \cdot s_{C_1} + l_{C_1}}) \tag{2}$$

with kernel size $l_{C_1} = 32$ and stride $s_{C_1} = 16$.

Key Insight

Prior work (InfiniGen) used raw hidden states as a proxy for future attention, which can be inaccurate when adjacent-layer similarity breaks down. SparDA instead learns a cross-layer Forecast projection that directly predicts the sparse selector's next-layer block choices.

Methodology

SparDA Architecture

SparDA adds a fourth per-layer projection, the Forecast $\mathbf{F}_l$, alongside the standard $\mathbf{Q}_l, \mathbf{K}_l, \mathbf{V}_l$:

$$(\mathbf{Q}_l, \mathbf{K}_l, \mathbf{V}_l, \mathbf{F}_l) = \phi_l(\mathbf{X}_l) \tag{3}$$

The Forecast drives top-𝑘 selection for layer $l+1$:

$$\mathcal{B}_{l+1} = \mathcal{B}_{\text{init}} \cup \mathcal{B}_{\text{local}} \cup f_{\text{top}}(\mathbf{F}_l \widetilde{\mathbf{K}}_{l+1}^\top, k) \tag{4}$$

while the attention query still performs sparse attention:

$$\mathbf{O}_{l+1} = \text{Attn}(\mathbf{Q}_{l+1}, \mathbf{K}_{l+1}[\mathcal{B}_{l+1}], \mathbf{V}_{l+1}[\mathcal{B}_{l+1}]) \tag{5}$$

Compact Forecast Indexer

Key observation: once Forecast is decoupled from the attention query, sparse selection no longer needs the full query-head structure. SparDA uses one Forecast head per GQA group (one per KV head) rather than one per query head, eliminating the per-query-head scoring loop and skipping softmax normalization entirely.

Indexer Training

The Forecast projections are trained via KL divergence against the original selector's block attention distribution:

$$\mathbf{S}_{l,m}^{\text{tgt}} = \sum_{h=1}^{G} \text{softmax}\left(\mathbf{Q}_{l,m,h} \widetilde{\mathbf{K}}_{l,m}^{\text{tgt}\top} / \tau\right), \quad \mathbf{S}_{l,m}^{\text{pred}} = \text{softmax}\left(\mathbf{F}_{l-1,m} \widetilde{\mathbf{K}}_{l,m}^{\text{pred}\top} / \tau\right) \tag{6}$$

The training loss is computed over a top-𝑘 partitioned distribution:

$$\mathcal{L}_{\text{KL}} = \sum_{l} \text{KL}\left(\bar{\mathbf{S}}_{l,\mathcal{S}}^{\text{tgt}} \| \bar{\mathbf{S}}_{l,\mathcal{S}}^{\text{pred}}\right) \tag{7}$$

where $\bar{\mathbf{S}}_{\cdot,\mathcal{S}}$ keeps the scores of the 𝑘 target-selected blocks individually and aggregates the remaining mass into a single rest bucket.

Fine-grained supervision: The target uses a smaller compression window (kernel size 2, stride 1) than the prediction window (32, 16), providing higher-resolution supervision. The target scores are max-pooled down to the standard grid before computing the KL loss.

Efficient Implementation

Persistent UVA kernel: A persistent Triton kernel based on Unified Virtual Addressing (UVA) performs high-throughput host-to-device transfers in parallel with main GPU kernels, using a small fixed set of Cooperative Thread Arrays (CTAs) to reduce launch overhead.

Batch-adaptive CTA allocation: A heuristic chooses the CTA count that maximizes overall throughput, balancing prefetch throughput against layer execution speed.

Empirical Validation / Results

Experimental Setup

• Models: MiniCPM4.1-8B (InfLLM-V2 backbone) and NOSA-8B (adds query-agnostic eviction head)
• Hardware: NVIDIA H100 and A100 GPUs
• Benchmarks: HELMET, LongBench, RULER, and long-reasoning suite (MATH-500, AIME 2024, AIME 2025)
• Comparisons: Dense, Sparse, InfiniGen, and SparDA

Accuracy Results

Table 1 | Aggregated benchmark averages.

SparDA improves the average by +0.3 on MiniCPM4.1-8B and +2.3 on NOSA-8B over the Sparse baseline. InfiniGen suffers significant degradation due to its training-free hidden-state similarity assumption.

Length generalization (RULER): SparDA outperforms Sparse at every sequence length, with the gap widening on NOSA-8B from +1.7 at 32K to +4.3 at 128K.

Efficiency Results

Prefill throughput (Table 3, H100, batch size 4):

• SparDA achieves up to 1.25× speedup over Sparse and 2.11× over Dense at 128K on MiniCPM4.1-8B
• On NOSA-8B: up to 1.16× over Sparse and 1.40× over Dense at 128K

Decode throughput (Table 4, H100):

• At 128K context: SparDA achieves up to 1.69× speedup over Sparse on MiniCPM4.1-8B and 1.40× on NOSA-8B
• Against non-offload baselines: up to 5.28× over Sparse† and 9.21× over Dense† on MiniCPM4.1-8B via larger feasible batch sizes
• InfiniGen is consistently slower than Sparse due to CPU-side gather bottleneck

Attention breakdown (Figure 3): SparDA reduces block-selection cost by up to 2.50× at 128K during prefill and keeps decode-time selection nearly flat, cutting overhead by more than 2× at 128K.

Theoretical and Practical Implications

Theoretical Significance

1. Decoupling selection from attention: SparDA demonstrates that sparse selection can be treated as a trainable, schedulable signal rather than an operation tied to the current attention query, opening new architectural possibilities.

2. Sparsity as an offloading-friendly schedule: The results suggest sparse attention mechanisms should expose future memory-access patterns early enough for serving systems to act on them, turning sparsity from a compute-saving device into an offloading-friendly schedule.

3. Generalizability: The decoupled indexer principle extends beyond block-level selection to token-level sparse attention (e.g., DSA) and Compressed Sparse Attention (CSA) paths, though validation on larger models is left as future work.

Practical Implications

1. Immediate deployment: SparDA can be integrated into existing sparse-pretrained models with minimal overhead (<0.5% parameters) and no base model retraining.

2. Hardware efficiency: The persistent UVA kernel with batch-adaptive CTA allocation provides a practical recipe for hiding PCIe transfer latency in real serving systems.

3. Scalability: By enabling larger feasible batch sizes on a single GPU, SparDA improves throughput without requiring multi-GPU setups, which is particularly valuable for cost-sensitive serving scenarios.

Conclusion

SparDA introduces a decoupled sparse attention architecture where a trainable Forecast projection predicts KV block selections one layer ahead, enabling:

1. Lookahead prefetching: CPU-to-GPU KV transfers overlap with current-layer execution, hiding PCIe latency
2. Reduced selection overhead: A compact GQA-level indexer with one Forecast head per group eliminates per-query-head scoring and softmax

The method matches or slightly improves accuracy over sparse baselines while delivering up to 1.25× prefill and 1.7× decode speedups, with up to 5.3× higher decode throughput via larger feasible batches.

Limitations: SparDA's accuracy is bounded by the quality of the base sparse attention method; it does not change the underlying sparse pattern.

Future directions:

• Extending decoupled lookahead selection to token-level sparse attention (DSA, CSA)
• Validation on larger sparse-pretrained models (e.g., DeepSeek-V3.2, GLM-5, DeepSeek-V4)
• Applying the lookahead design to other sparse attention structures with initial/local/top-𝑘 composition