SparDA
稀疏解耦注意力——长上下文LLM推理的效率革命

🏔️问题背景：长上下文推理的两座大山

现代LLM应用——尤其是agentic和长推理场景——越来越需要长上下文窗口。这给底层推理基础设施带来三大效率挑战：(1) prefill阶段注意力的高计算需求；(2) decode阶段注意力的高显存带宽压力；(3) decode阶段KV Cache的高容量压力。稀疏注意力已被前沿模型采用来改善前两项，但论文揭示了一个被忽视的事实：稀疏注意力主要减少了注意力计算和带宽，但KV Cache容量压力随上下文长度增长持续存在。

论文精确刻画了两座"大山"：

1.1 第一座山：KV Cache容量与PCIe传输瓶颈

直接的解决方案是将KV Cache卸载到CPU内存，decode时仅将选中的KV block动态拉取到GPU。但CPU通过PCIe接口的拉取远慢于GPU显存读取，在长上下文解码中成为瓶颈。已有工作InfiniGen提出了前瞻预取来隐藏部分延迟，但它依赖原始隐状态作为未来注意力的代理，当相邻层相似性失效时不准确——这是一个未被训练的启发式假设。

1.2 第二座山：稀疏选择本身的O(T²)复杂度

虽然稀疏注意力的复杂度已从O(T²)降至O(T)，但稀疏选择步骤本身仍保持O(T²)，在长上下文下可轻易主导注意力模块的开销。IndexCache和HISA等近期工作减少了选择开销，但通常引入精度-效率权衡，不能完全保持原始精度。

SparDA的核心洞察是：这两座山可以通过单一架构变更同时解决——将稀疏选择从注意力Query中解耦，使其可以提前一层计算，从而既减少选择开销（紧凑索引器），又实现KV预取与当前层执行的overlap（前瞻调度）。

🏗️SparDA核心架构：Forecast投影与跨层前瞻

SparDA在标准Transformer线性投影中增加第四个输出——Forecast F^l，与Q^l、K^l、V^l并列：

(Q^l, K^l, V^l, F^l) = φ_l(X^l)

2.1 解耦稀疏选择

在原始InfLLM-V2稀疏注意力中，Query Q^l同时驱动top-k稀疏选择和稀疏注意力计算，且两者在同一层内发生。SparDA将这两个角色分离：F^l（来自层l）驱动层l+1的top-k选择，Q^l+1仅执行稀疏注意力。具体地：

B^l+1 = B_init ∪ B_local ∪ f_top(F^l · K̃^l+1⊤, k)
O^l+1 = Attn(Q^l+1, K^l+1[B^l+1], V^l+1[B^l+1])

对于第一层（l=0），无前一层Forecast可用，SparDA通过单独投影产生当前层Forecast F₀^cur，仍使用每GQA组一个Forecast头，选择成本低于基线。对于最后一层，F^L-1无用（无后续层可预测）。

2.2 紧凑Forecast索引器

关键观察：一旦Forecast从注意力Query中解耦，稀疏选择不再需要使用与注意力相同的头结构。标准block-sparse选择器（如InfLLM-V2）用GQA组的所有G个Query头评分，因为注意力Query本身驱动选择。SparDA的F^l使用每GQA组一个Forecast头（即每个KV头一个），消除了逐Query头评分循环，显著减少top-k选择开销。它还自然跳过了标准block-sparse选择器在top-k排序前应用的softmax操作——因为无需在GQA组内多个Query头间进行分数求和。

这一设计灵感来自DeepSeek DSA的lightning indexer，它在token级别已利用解耦使用比注意力Query更少的头。SparDA将同样的设计带到block-sparse注意力。

2.3 Prefill与Decode的不同收益

解耦选择在两个推理阶段收益不同：

• Prefill：所有Key已在GPU上，唯一节省是降低稀疏选择成本（紧凑索引器替代多头选择器）。
• Decode：KV Cache卸载到CPU内存，一层前瞻预测变得关键——F^l在层l仍在执行时产生B^l+1，运行时可从CPU预取选中的KV条目，将PCIe传输与层执行overlap。压缩Key K̃缓存保留在GPU上并增量更新；由于稀疏选择排除初始和局部block，B^l+1不依赖下一层新追加的Key。

🎓训练策略：仅训练Forecast，冻结骨干网络

SparDA可添加到现有稀疏预训练模型上，仅训练Forecast投影，无需重训骨干网络。在MiniCPM4.1-8B和NOSA-8B上，这些投影仅增加33.5M参数（8B总量的0.41%），是一个可忽略模型大小开销的轻量附加模块。

3.1 KL散度训练目标

受DeepSeek DSA启发，Forecast索引器通过KL散度训练以匹配目标block-注意力分数。由于所用模型已用稀疏注意力训练，SparDA仅训练Forecast投影，跳过DSA两阶段流程中的全模型稀疏训练阶段。

训练损失在top-k分区分布上计算KL散度：保留k个目标选中block的分数 individually，将其余质量聚合到单个rest bucket，重归一化为1。这使索引器聚焦于S_l内的相对排序，同时约束非选中block上的总质量——集外logit通过rest bucket获得非平凡梯度。

3.2 细粒度训练监督（关键发现）

论文一个重要发现：目标压缩Key K̃^tgt和预测压缩Key K̃^pred可使用不同核大小和步长。K̃^pred始终使用标准InfLLM-V2压缩窗口（l_C1=32, s_C1=16）以匹配推理配置。但对目标使用更小的核和步长(2, 1)产生更好的索引器。

直觉：更细粒度的压缩提供更高分辨率的监督信号——每个压缩Key代表更少的token组，目标分数更具区分性，索引器学到更锐利的选择决策。由于更细的K̃^tgt产生比K̃^pred更多的压缩位置，训练时将目标分数张量max-pool到标准(32,16)网格后再计算KL损失。消融实验确认：细粒度窗口在MiniCPM4.1-8B上提升RULER +3.0、reasoning +2.2；在NOSA-8B上提升全部四个benchmark族，包括RULER +2.5、reasoning +1.8。

3.3 训练配置

• 32块NVIDIA H100 GPU
• 2000 optimizer steps，constant学习率5×10⁻⁴
• AdamW (β₁=0.9, β₂=0.95, ε=10⁻⁶)，weight decay 0.01
• BF16混合精度，gradient clipping max norm 0.5
• 训练数据：ProLong-64K，序列长度65536 tokens
• MiniCPM4.1-8B @ 64K：48小时完成；NOSA-8B @ 32K：24小时完成

⚡系统实现：异步预取与Triton内核

SparDA的效率不仅来自架构设计，还来自精心实现的系统管线。

4.1 持久UVA内核

一旦下一层的稀疏模式被预测，运行时从pinned CPU内存中获取选中的KV block，在专用CUDA stream上传输，使其与层执行overlap。SparDA不依赖多次小型不规则内存拷贝，而是使用基于Unified Virtual Addressing (UVA)的持久Triton内核：保持一组固定的GPU线程块（Cooperative Thread Arrays, CTAs）活跃，让它们在单次launch中持续处理block传输任务。这减少了launch开销，避免了频繁同步，并限制了对主计算stream的干扰。

4.2 批次自适应CTA分配

持久内核中CTA数量控制预取吞吐与层执行速度之间的权衡。更多CTA加速传输趋向PCIe带宽上限，但消耗本可运行注意力和FFN内核的Streaming Multiprocessor (SM)。

• 小batch：GPU未充分利用，层执行主导关键路径；适度CTA预算足够（总传输量小）。
• 大batch：层执行时间亚线性增长（更好的SM占用），预取量线性增长；超过交叉点后预取成为新瓶颈，分配更多CTA值得。

SparDA采用简单启发式：H100上batch size < 32时用16 CTAs，否则用32 CTAs。该自适应方案在每个batch size上匹配或保持在最佳固定配置的4%以内。A100上类似：batch < 64时16 CTAs，否则32 CTAs。

📊实验结果：精度与效率全方位验证

实验在两个8B稀疏预训练模型上评估：MiniCPM4.1-8B（InfLLM-V2骨干，原生64K）和NOSA-8B（InfLLM-V2 + eviction head，原生32K）。精度在HELMET、LongBench、RULER和长推理套件（MATH-500、AIME 2024、AIME 2025）上测量。效率在NVIDIA H100和A100 GPU上测量。比较四种配置：Dense、Sparse、InfiniGen、SparDA。

5.1 精度结果（Table 1）

MiniCPM4.1-8B					NOSA-8B
方法	HELMET	LongBench	RULER	Reasoning	Avg	HELMET	LongBench	RULER	Reasoning
Dense	41.7	44.8	85.3	82.3	63.5	39.3	42.5	86.2	41.6	52.4
Sparse	38.9	45.0	78.2	83.6	61.4	32.2	42.4	72.2	50.7	49.4
InfiniGen	33.5	45.1	68.4	83.7	57.7	28.1	41.6	65.2	47.6	45.6
SparDA	38.3	45.1	78.7	84.7	61.7	33.4	42.3	73.9	57.2	51.7

关键发现：

• MiniCPM4.1-8B：SparDA平均 +0.3 over Sparse，RULER +0.5，reasoning +1.1，LongBench持平，HELMET略低 -0.6
• NOSA-8B：SparDA平均 +2.3 over Sparse，reasoning +6.5（大幅提升），RULER +1.7，HELMET +1.2
• InfiniGen在两个模型上均显著退化——作为training-free方法，它依赖隐状态跨层相似性，这一假设不总成立

5.2 长度泛化（Table 2）

SparDA在RULER上每个序列长度上都优于Sparse：

模型	方法	32K	64K	96K	128K
MiniCPM4.1-8B	Sparse	86.1	78.2	68.7	67.7
	SparDA	87.6	78.7	70.8	68.8
	Δ	+1.5	+0.5	+2.1	+1.1
NOSA-8B	Sparse	72.2	56.6	48.8	40.7
	SparDA	73.9	60.5	52.9	45.0
	Δ	+1.7	+3.9	+4.1	+4.3

NOSA-8B上差距随序列长度稳步扩大（+1.7 → +4.3），表明学习到的Forecast泛化能力至少与training-free基线选择器相当，在某些场景下更好。

5.3 效率结果

Prefill吞吐量（Table 3，H100）

SparDA从64K起在MiniCPM4.1-8B上持续领先所有方法：

• 128K：17,087.6 tok/s，1.25× over Sparse（13,661.8），2.11× over Dense（8,085.8）
• NOSA-8B 128K：11,332.7 tok/s，1.16× over Sparse，1.40× over Dense
• 主要加速来源：Forecast索引器降低的block选择成本，随序列长度增长比基线多头选择器更优

Decode吞吐量（Table 4，H100）

Decode是前瞻设计收益最大的阶段——CPU→GPU KV预取与当前层计算overlap：

• 128K offload regime内：1.69× over Sparse（MiniCPM4.1-8B），1.40×（NOSA-8B）
• 最大加速出现在中等batch size（B8-B16），预取与层执行大致平衡
• 对比非offload基线（†，在长上下文B16后OOM）：5.28× over Sparse†，9.21× over Dense†（MiniCPM4.1-8B 128K）
• NOSA-8B加速较低：其query-agnostic eviction head已减少KV拉取流量，留给overlap的空间更少

Decode加速分解（Table 10，H100 128K）

论文隔离了SparDA decode加速的两个来源：

• 无预取变体（Forecast索引器但同步拉取）：B4时已优于Sparse（选择节省），但B64时低于Sparse（选择节省在大batch时递减，索引器投影成本仍在）
• 完整SparDA：B16起超越无预取变体，B64时约快40%——预取overlap对大batch优势至关重要

Block选择成本降低（Figure 3）

• Prefill 128K：block选择成本降低2.50×，block-sparse注意力时间与Sparse相当
• Decode 128K：选择开销降低2×以上，decode选择曲线近乎平坦——索引器计算足够小，GPU未充分利用，更长上下文几乎不影响decode延迟

A100结果

趋势与H100一致，加速略低：prefill 1.23× over Sparse（128K），decode 1.55× over Sparse（128K B8）。InfiniGen的CPU侧gather瓶颈在A100上更明显，大batch时吞吐量急剧退化。

⚖️技术对比与定位

6.1 与InfiniGen的对比

InfiniGen是与SparDA最接近的工作：它在每层前使用原始隐状态作为跨层代理预取CPU-resident KV子集到GPU。SparDA的关键差异：

• InfiniGen用未训练的隐状态相似性作为代理；SparDA训练专用Forecast投影直接预测选择器的下一层选择
• InfiniGen在CPU上gather top-k KV block后传输到GPU，CPU侧gather成为瓶颈；SparDA用持久UVA内核直接在GPU上管理传输
• 实验中InfiniGen在两个模型上均显著精度退化，SparDA持平或略优

6.2 与DeepSeek DSA的关系

DSA的lightning indexer已在token级别利用解耦使用比注意力Query更少的头。SparDA将同样的设计带到block-sparse注意力，并在其基础上增加了跨层前瞻预取能力。SparDA的KL训练目标直接受DSA启发，但跳过了DSA的全模型稀疏训练阶段（因为目标模型已稀疏预训练）。

6.3 与NOSA/SparseServe/HiSparse的对比

• NOSA：配对query-agnostic eviction head + query-aware selector，同步发出UVA传输，减少传输量而非隐藏延迟
• SparseServe/HiSparse：offload KV到host，按需reactive swap——均不针对层内传输-计算overlap
• SparDA：唯一同时实现选择开销降低 + 传输-计算overlap的工作

🌍对产业的影响与展望

7.1 核心范式转变

论文结论中最深刻的表述："sparse attention mechanisms should expose future memory-access patterns early enough for the serving system to act on them, turning sparsity from a compute-saving device into an offloading-friendly schedule." 这意味着稀疏注意力不仅应减少计算，还应提前暴露未来内存访问模式，使serving系统能据此调度——将稀疏性从计算节省手段转变为卸载友好的调度原语。

7.2 产业影响

• 长上下文推理成本大幅降低：5.3× decode吞吐量提升意味着同样硬件可服务更多并发长上下文请求，或同样负载可用更少GPU
• 无需重训骨干网络：<0.5%参数 + 32 H100 × 48h 的训练成本极低，对已有稀疏预训练模型是即插即用升级
• 更大batch size可行：offload使单GPU可运行远超显存限制的batch size，直接转化为吞吐量收益
• 开源代码可用：https://github.com/NVlabs/SparDA

7.3 局限与未来方向

论文诚实声明局限：

• SparDA本身不是稀疏注意力方法，是附加模块；精度受限于骨干稀疏注意力的质量
• 解耦索引器原则不限于block级选择——预期可扩展到token级（DSA，已被DeepSeek-V3.2和GLM-5采用）和CSA（DeepSeek-V4）
• 未来工作：将SparDA应用于DSA和CSA模型，以及远大于8B的稀疏预训练模型
• SparDA因保留layer-0 KV cache在GPU上，在长上下文下比Sparse略早OOM——但峰值吞吐量在每个上下文长度上都超过Sparse

📄关键数据速览与总结