通过前瞻预测减少长文本推理中的计算开销,提升速度。
SparDA: Sparse Decoupled Attention for Efficient Long-Context LLM Inference
- 引入预测头提前选定所需键值块,实现计算与数据预加载并行。
- 在80亿参数模型上,预填充速度提升1.25倍,解码速度提升1.7倍。
- 仅增加不到0.5%参数,适合大上下文长序列推理场景使用。
稀疏注意力能降低长上下文大模型推理的计算与内存带宽开销。但仍有两大挑战:(1) 键值缓存容量随序列长度增长,而将其移至CPU内存会引入PCIe传输瓶颈;(2) 稀疏选择步骤仍具$O(T^2)$复杂度,在长序列下可能成为主要开销。我们提出SparDA,一种解耦稀疏注意力架构,每层新增一个第四投影——预测头(Forecast),与查询、键、值并列。该预测头可预判下一层所需的键值块,实现前瞻选择,使CPU到GPU的数据预取与当前层计算重叠。由于预测头与注意力查询解耦,本方案采用每组通用查询注意力(GQA)对应一个预测头,相比原始多头选择器显著降低选择开销。SparDA仅增加小于0.5%的参数,且仅训练预测投影,使其注意力分布匹配原选择器。在两个稀疏预训练的80亿参数模型上,SparDA保持或略微提升准确率,较稀疏注意力+卸载基线实现最高1.25倍预填充加速和1.7倍解码加速。通过支持单张显卡更大批量处理,其解码吞吐量最高可达非卸载稀疏基线的5.3倍。源代码已开源:https://github.com/NVlabs/SparDA。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。