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.
核心贡献 · Key contributions
引入 Forecast，即每层第四个投影，预测下一层所需的 KV 块，实现前瞻预取。 Introduces Forecast, a fourth per-layer projection that predicts KV blocks for the next layer, enabling lookahead prefetch.
将稀疏选择与注意力查询解耦，每个 GQA 组使用一个 Forecast 头，降低选择开销。 Decouples sparse selection from attention query, using one Forecast head per GQA group to reduce selection overhead.
通过持久化 UVA 内核实现异步 CPU 到 GPU 预取，使传输与层执行重叠。 Implements asynchronous CPU-to-GPU prefetch with a persistent UVA kernel, overlapping transfer with layer execution.
增加 <0.5% 参数，仅通过 KL 散度训练 Forecast 投影，保持基础模型精度。 Adds <0.5% parameters and trains only Forecast projections via KL divergence, preserving base model accuracy.
在 8B 模型上，相比稀疏注意力卸载基线，实现高达 1.25 倍预填充和 1.7 倍解码加速。 Achieves up to 1.25x prefill and 1.7x decode speedup over sparse-attention offload baseline on 8B models.
在单 GPU 上支持更大批大小，解码吞吐量比非卸载稀疏基线高 5.3 倍。 Enables larger batch sizes on a single GPU, reaching up to 5.3x higher decode throughput than non-offload sparse baseline.
局限 · Limitations
精度受限于基础稀疏注意力方法；SparDA 不改变稀疏模式。 Accuracy is bounded by the base sparse attention method; SparDA does not change the sparse pattern.
仅在 8B 模型和块稀疏注意力上评估；扩展到词元级稀疏性是未来工作。 Evaluated only on 8B models with block-sparse attention; extension to token-level sparsity is future work.
需要稀疏预训练模型；不适用于未重新训练的稠密模型。 Requires sparse-pretrained models; not applicable to dense models without retraining.
第 0 层 KV 缓存保留在 GPU 上，在极长上下文时略早出现内存不足。 Layer-0 KV cache remains on GPU, causing slightly earlier out-of-memory at very long contexts.
预取 CTA 数量启发式方法可能并非在所有硬件和批大小下最优。 Prefetch CTA count heuristic may not be optimal across all hardware and batch sizes.