笔记日期: 2026-06-24 笔记作者: Zhongzhu Zhou 论文标题: SparDA: Sparse Decoupled Attention for Efficient Long-Context LLM Inference 作者: Yaosheng Fu, Guangxuan Xiao, Xin Dong, Song Han, Oreste Villa arXiv: 2606.04511 状态 / Venue: arXiv 预印本,2026 年 6 月(NVIDIA, Thinking Machines Lab, ByteDance Seed, MIT)
一句话总结
SparDA 是一个轻量级的注意力机制增强方案——它在每层的线性投影中新增了一个”Forecast”输出,用来提前一层预测下一层需要哪些 KV 块,从而在当前层执行时异步将这些块从 CPU 内存预取到 GPU,将 PCIe 传输延迟隐藏在 GPU 计算后面。同时,Forecast 只需每个 GQA 组一个头(而不是每个查询头一个),大幅降低了块选择开销。整套方案只需训练 0.41% 的新参数,在 128K 上下文、H100 上实现了最高 1.25× 预填充加速、1.7× 解码加速,以及相比非卸载稀疏注意力 5.3× 的解码吞吐量提升。
前置知识
在进入论文正文之前,我想先把几个关键概念讲清楚——包括 KV 缓存、稀疏注意力、GQA、KV 卸载和 PCIe 瓶颈。如果你对这些已经很熟了,可以直接跳到下一节。
1. 自回归解码与 KV 缓存
大语言模型的推理是逐 token 进行的:每次生成一个新 token,模型都要对所有历史 token 做一次注意力。如果每步都从头计算 Key 和 Value,代价是 ,显然不现实。KV 缓存的做法是:每个 token 的 K 和 V 投影结果在计算后就保存起来,下次直接复用,这样每步注意力代价降为 。
KV 缓存的内存占用是:
其中 是层数, 是 KV 头数, 是头维度, 是序列长度, 是批大小。以一个 70B 的模型为例:8 个 KV 头,,序列长 128K,批大小 8,BF16 精度:
这已经占掉了 H100 80GB HBM 的一半以上。随着上下文窗口越来越长,KV 缓存往往成为内存的主要瓶颈。
2. 分组查询注意力(GQA)
GQA 的思路是让多个查询头共享一对 K/V 头。如果每组有 个查询头共享一个 KV 头,那么 KV 头数从 降到 ,KV 缓存也跟着减少 倍。LLaMA-3、DeepSeek-V3 等主流大模型都在用 GQA, 通常取 4–8。
3. 块稀疏注意力:原理与开销
即便有了 GQA,每步解码仍然要读遍整个 KV 缓存——带宽压力巨大。块稀疏注意力的洞察是:注意力权重高度集中,大多数查询只和一小部分历史 token 有实质性的交互。
做法:把历史 token 划分成若干 block(比如每块 64 个 token),每个查询只选 top- 个最相关的 block 来做注意力,其余的直接跳过:
- :固定的初始 block(前几百个 token,几乎所有查询都会关注)
- :位置 附近的滑动窗口
- :按相关性评分选出的 top- 个 block
注意力计算的代价从 降到 (固定常数,与 无关)。
但选择的代价呢? 要找到 top-,需要对所有 block 打分。InfLLM-V2 先把 Key 做均值池化压缩:
然后用查询头向量和这些压缩 Key 做内积打分,再做 top- 选取。这个打分步骤是 ( 是每 GQA 组的查询头数, 是压缩 block 数),在 128K 上下文下非常耗时,甚至超过注意力本身。
4. KV 卸载与 PCIe 瓶颈
当 KV 缓存放不进 GPU 显存时,可以把它卸载到 CPU 内存,每步只把需要的 block 拉回 GPU。问题是,PCIe 5.0 x16 的带宽约 64 GB/s,而 HBM3 是 3.4 TB/s,差了 50 倍。拉一次数据的延迟可能比注意力本身还长。
怎么绕过这个瓶颈? 唯一的办法是提前预测好下一步需要哪些 block,在 GPU 还在算当前层的时候就提前把它们传过来。这就是”预取”(prefetch)的思路,也是 SparDA 的核心。
5. 现有预取方案的问题
InfiniGen(OSDI 2024)尝试用当前层的隐状态 作为下一层查询的代理,提前预测需要的 block。问题在于:相邻层之间的隐状态经过 FFN 和 Layer Norm 的变换后,结构差异很大,代理信号不可靠——InfiniGen 在好几个 benchmark 上精度大幅下降就是因为这个。
SparDA 的思路是:不用”代理”,直接训练一个专门用于跨层预测的向量,让它学会”预测下一层需要哪些 block”这件事。
论文概述
SparDA(Sparse Decoupled Attention,稀疏解耦注意力)的核心问题意识是:块选择和块注意力是两件不同的事,没必要在同一层同时做。把它们解耦,就可以让选择提前发生,从而为异步预取打开时间窗口。
论文贡献可以分三层:
- Forecast 投影:一个轻量的第四投影 ,专门用于预测下一层的 block 选择
- 紧凑 Forecast 索引器:每个 GQA 组只用一个 Forecast 头(而非 个查询头),选择开销降低约 倍
- 持久化 UVA 预取内核:一个基于 Triton + UVA 的异步 DMA 内核,把 CPU→GPU 的数据传输完全覆盖在 GPU 计算窗口里
三者协同工作:Forecast 提前算好下一层的 block 索引(贡献 1),紧凑索引器保证这个计算足够快(贡献 2),UVA 内核保证传输与计算真正并行(贡献 3)。
方法详解
4.1 Forecast 投影:解耦选择与注意力
改动只有一行:把每层的线性投影 从输出三个向量改为输出四个:
其中 ,维度与单个 KV 头相同。这相当于在原有投影矩阵旁边多加了一个 。
关键: 用于预测下一层( 层)需要的 block,而不是当前层。它打的分是 Forecast 向量与下一层压缩 Key 的内积:
下一层的注意力照样用 来算,不受影响:
这就是”解耦”:选择(用 )发生在第 层,注意力计算发生在第 层。两者之间的时间差——整整一个 层的注意力和 FFN 计算时间——就是异步预取可以利用的窗口。
graph TD
subgraph "传统稀疏注意力(InfLLM-V2)"
Xl_a[X_l] --> ProjA[线性投影]
ProjA --> Ql_a[Q_l]
ProjA --> Kl_a[K_l]
ProjA --> Vl_a[V_l]
Ql_a -->|"驱动 top-k 选择"| SelA["块选择(当前层,在关键路径上)"]
SelA --> Bla["块集合 B_l(i)"]
Bla --> SAA["稀疏注意力"]
Ql_a --> SAA
Kl_a --> SAA
Vl_a --> SAA
SAA --> Xl1_a["X_{l+1}"]
end
图 1(论文图 1a):标准块稀疏注意力。 既要驱动块选择,又要做注意力,选择在关键路径上——所有 KV 数据必须在注意力开始前就到位,没有预取的时间窗口。
graph LR
subgraph "SparDA 第 l 层"
Xl_s[X_l] --> ProjS["线性 phi_l"]
ProjS --> Ql_s[Q_l]
ProjS --> Kl_s[K_l]
ProjS --> Vl_s[V_l]
ProjS --> Fl_s["F_l Forecast 头"]
Fl_s -->|"打分 K_tilde_{l+1}"| SelS["top-k 选择 -> B_{l+1}"]
SelS -->|"异步预取"| UVA["UVA 持久化内核\nCPU 内存 -> GPU HBM"]
end
subgraph "SparDA 第 l+1 层"
UVA -->|"数据已就位"| SAl1["稀疏注意力\n用 KV_{l+1}[B_{l+1}]"]
Xl1s["X_{l+1}"] --> ProjS1["线性 phi_{l+1}"]
ProjS1 --> Ql1["Q_{l+1}"]
ProjS1 --> Fl1["F_{l+1}(预测第 l+2 层)"]
Ql1 --> SAl1
SAl1 --> Xl2s["X_{l+2}"]
end
图 2(论文图 1b-c、图 2):SparDA 数据流。 在第 层执行期间就算出了第 层需要的 block,并触发异步预取。等第 层的稀疏注意力开始时,数据早已到 GPU 了。
4.2 伪代码:SparDA 解码单步详解
为了让流程更直观,我把一步解码(新生成一个 token,历史 个 token 的 KV 都在 CPU 内存里)的完整步骤写成伪代码:
算法 1:SparDA 解码单步(第 l 层,生成第 T+1 个 token)
输入:
X_l 新 token 在第 l 层的隐状态 [d_model]
K_l_cpu, V_l_cpu 第 l 层完整 KV 缓存(CPU 内存,T 个 block)
K_tilde_l_gpu 第 l 层压缩 Key 缓存(GPU,N_b 个 block)
F_prev 上一层的 Forecast 向量(GPU)
k top-k block 数
步骤 1 — 线性投影(GPU):
Q_l, K_new, V_new, F_l = phi_l(X_l)
将 K_new, V_new 追加到 K_l_cpu, V_l_cpu(CPU 固定内存)
增量更新 K_tilde_l_gpu(均值池化,O(d_kv) 代价)
步骤 2 — Forecast 索引器(GPU,用上一层的 F_prev):
scores = F_prev @ K_tilde_{l+1}_gpu.T # [N_b] 内积
B_init_local = 固定初始 block + 滑动窗口 block
B_topk = top_k(scores, k, 排除 B_init_local)
B_{l+1} = B_init_local UNION B_topk
步骤 3 — 下发异步预取(UVA 持久化内核):
uva_queue.enqueue(layer=l+1, block_indices=B_{l+1})
# 控制权立即返回;持久化 CTA 池后台执行 PCIe 传输
步骤 4 — 稀疏注意力(用上上层预测的 B_l,数据已在 GPU):
O_l = sparse_attention(Q_l, GPU_buffer[l][B_l])
步骤 5 — FFN:
X_{l+1} = FFN(O_l) + O_l
步骤 6 — 同步:等待步骤 3 中的预取完成
# 实际上步骤 4-5 期间传输早已结束,等待时间接近零
输出:
X_{l+1},F_l(传给第 l+1 层用于预测第 l+2 层)
核心不变量:每层稀疏注意力开始时(步骤 4),所需的 KV block 已由上一步的预取放在 GPU 上。“先预测 → 边算边传 → 用时已到”三步构成一个流水线,只有第一层和最后一层例外(见 4.6 节)。
4.3 紧凑 Forecast 索引器
在 InfLLM-V2 中,块选择器用每个 GQA 组的所有 个查询头打分:
然后把 个头的得分加总,再做 max-pooling,得到 block 级别的重要性分数。整个过程代价正比于 ( 是压缩 block 数), 越大越贵, 越长越贵。
SparDA 的解法:因为 Forecast 解耦了选择和注意力,它不再需要保持注意力的多头结构——每个 GQA 组只用一个 Forecast 头:
其中 只有一个头(维度 ,而非 )。这带来两个收益:
- 倍降低内积计算量:一个头 vs. 个头,选择循环直接减少 倍
- 省掉 softmax 归一化和跨头求和:只有一个头,不需要对多个头的分数做归一化再加总
论文在 128K 上下文下实测,块选择代价比 InfLLM-V2 基线降低了 2.50×(论文图 3a)。在解码阶段,标准 InfLLM-V2 的块选择时间随序列长度线性增长,成为最大瓶颈;而 SparDA 的 Forecast 索引器的代价只取决于 ,与 无关,解码时选择时间几乎保持恒定。
graph LR
subgraph "InfLLM-V2 选择器(每组 G 个查询头)"
Qheads["Q_{l,m,1}...Q_{l,m,G}\n(G 个头)"] --> Scores["G 次内积计算"]
Scores --> Softmax["每头 softmax"]
Softmax --> Sum["跨 G 头求和"]
Sum --> MaxPool["max-pool -> block 分数"]
end
subgraph "SparDA Forecast 索引器(每组 1 个头)"
Fhead["F_{l-1,m}\n(1 个 Forecast 头)"] --> Score1["1 次内积计算"]
Score1 --> BlockScores["block 分数\n(无 softmax,无求和)"]
end
图 3:选择器代价对比。 InfLLM-V2 的多头选择器需要 次内积加 softmax 加跨头求和。SparDA 的 Forecast 索引器只需 1 次内积——代价大约 倍便宜。
4.4 KL 散度训练目标
Forecast 投影是在已经做完稀疏预训练的基模型上附加的,只训练 Forecast 权重,基模型冻结。训练目标是让 Forecast 索引器的 block 打分分布,尽量逼近原始查询向量 的 block 注意力分布。
设 为 top- 目标 block 集合(经过因果掩码和初始/局部 block 排除后)。目标分布(来自原始查询,所有 头的贡献加总):
预测分布(来自上一层的 Forecast):
这里有一个刻意的不对称:目标用的是细粒度压缩(核大小 ,步长 ),预测用的是推理时标准压缩(,)。细粒度 Key 代表更少的 token,分数分辨率更高,监督信号更”锋利”。由于两者 block 数不同,目标分数先做 max-pooling 对齐到粗粒度网格。
损失函数只作用在被选中的 block 上,加一个”其余所有 block 聚合”的 rest bucket:
其中 是一个 维分布: 个被选 block 的分数各自保留,其余所有 block 的分数汇总到一个 rest bucket,归一化后求 KL。
为什么设计 rest bucket? 如果只对被选 block 算 KL,模型对未被选 block 的打分就没有约束——它们可能学成任意值。下次训练换一批序列,这些 block 可能就变成了被选 block,打分混乱会导致训练不稳定。rest bucket 聚合了未选 block 的总概率质量,保证它们也获得有意义的梯度信号。
为什么用细粒度监督? 每个压缩 Key 代表的 token 越少,该 block 的重要性分数就越”有区分度”——两个相邻 block 之间的分数差异更大,Forecast 因此必须学会做细粒度的排序判断,而不是满足于”大概对”。消融实验(附录 D 表 6)证实了这个训练时的不对称设计确实有效。
graph TD
subgraph "Forecast 训练信号流"
A["Q_{l,m,h} 与 K_tgt 内积(细粒度,核大小 2)\n所有 G 头求和 + softmax"] --> Stgt["S_tgt(细粒度监督)"]
Stgt -->|"max-pool 对齐到粗粒度"| Stgt_c["S_tgt_coarse(与推理网格一致)"]
B["F_{l-1,m} 与 K_tilde_pred 内积(粗粒度,核大小 32)\nsoftmax"] --> Spred["S_pred(推理分辨率)"]
Stgt_c -->|"top-k 限制 + rest bucket"| S_r1["S_tgt_restricted\n(k+1) 维"]
Spred -->|"top-k 限制 + rest bucket"| S_r2["S_pred_restricted\n(k+1) 维"]
S_r1 --> KL["KL 散度损失 L_KL"]
S_r2 --> KL
end
图 4:Forecast 索引器训练目标的信号流。 细粒度目标提供更锐利的监督,故意不与推理时的粗粒度压缩对齐。两者在映射到同一 block 网格后,通过 KL 散度(含 rest bucket)约束 Forecast 的分布。
4.5 高效实现:持久化 UVA 预取内核
预测做得再准,如果 CPU→GPU 传输本身是同步的,还是会卡在关键路径上。SparDA 用一个持久化 Triton 内核(基于 CUDA Unified Virtual Addressing)来实现真正的异步传输:
工作流程(逐步):
第 1 步——推理开始时启动内核:一小批 CUDA 线程块(CTA)在推理开始时被启动,整个推理期间一直保持存活,循环等待传输请求。
第 2 步——Forecast 算完 block 索引,立刻入队:在第 层执行期间,一旦算出 ,主 GPU 流立刻向 UVA 队列发送一条传输请求(layer = ,block 索引 = )。控制权立即返回到主流,不等传输完成。
第 3 步——持久化 CTA 池处理请求:空闲的 CTA 取出请求,通过 UVA 指针直接读取 CPU 固定内存中的 KV block,发起 DMA 传输到 GPU HBM 中的暂存区。固定内存(pinned memory)保证了高效的 PCIe 传输。
第 4 步——主流继续执行第 层的注意力和 FFN:这段计算时间正好覆盖 PCIe 传输时间。
第 5 步——同步:第 层稀疏注意力开始前,主流等待 UVA 流的事件。实际上传输在步骤 4 期间就已经完成,等待时间几乎为零。
sequenceDiagram
participant Main as 主 GPU 流
participant UVA as UVA 持久化 CTA 池
participant CPU as CPU 固定内存
Main->>Main: 第 l 层:Q/K/V/F_l 线性投影
Main->>Main: Forecast 索引器:F_l 打分 K_tilde_{l+1} -> B_{l+1}
Main->>UVA: 入队预取请求(layer=l+1, B_{l+1})
Note over Main: 第 l 层:稀疏注意力开始执行
UVA->>CPU: DMA 读取 KV_{l+1}[B_{l+1}](固定内存)
CPU-->>UVA: PCIe 传输中(约 64 GB/s)
Note over Main: 第 l 层:FFN 执行
UVA-->>Main: 传输完毕,发送信号
Main->>Main: 同步等待(通常为零,传输已完成)
Main->>Main: 第 l+1 层:稀疏注意力(KV_{l+1}[B_{l+1}] 已在 GPU HBM)
图 5:异步预取时序图。 Forecast 在第 层执行早期触发预取,UVA 内核在后台持续传输;等第 层的稀疏注意力要用 KV 数据时,数据早已到位,主流零等待。
批大小自适应的 CTA 数目:CTA 越多,PCIe 带宽越高,但占用的 SM 越多,会与注意力/FFN 内核争资源。批大小小时 GPU 利用率低,少量 CTA 即可;批大小大时 GPU 满载,多分配 CTA 用于传输才划算。论文在附录 D 表 7 给出了 H100 上不同批大小对应的 CTA 配置。
4.6 首尾层的特殊处理
第 0 层:没有前一层的 Forecast。SparDA 用当前层隐状态的一个独立投影 做当前层的块选择——同样是每 GQA 组一个头,代价低于 InfLLM-V2 的多头选择器,只是没有 lookahead 的好处。
最后一层(): 预测了一个不存在的第 层的 block,直接丢弃即可。
压缩 Key 缓存的内存: 需要在第 层就可以访问到。这些压缩 Key 常驻 GPU(尺寸远小于全量 KV 缓存,每头 )并在每次生成新 token 时增量更新——内存开销可忽略。
实验
实验设置
评测模型:
- MiniCPM4.1-8B:稀疏骨干为 InfLLM-V2,评测最大上下文 64K
- NOSA-8B:在 InfLLM-V2 上叠加了一个 query-agnostic eviction head,评测最大上下文 32K
两个模型都已完成 InfLLM-V2 的两阶段稀疏预训练,然后在冻结基模型的基础上训练 Forecast 投影(33.5M 参数,占 8B 总量的 0.41%)。
基线对比:
- Dense†:全密集注意力,不卸载(长上下文 OOM)
- Sparse†:稀疏注意力,不卸载(大批 / 长上下文 OOM)
- Sparse:稀疏注意力 + CPU 卸载(同步取块)
- InfiniGen:基于原始隐状态做 lookahead 预取
- SparDA:本文方法
精度 benchmark: HELMET(长上下文理解、召回、生成)、LongBench(多任务长文档理解)、RULER(合成长上下文召回测试,32K–128K)、Reasoning(MATH-500, AIME 2024, AIME 2025)
效率硬件: NVIDIA H100(80GB HBM3);附录 D 补充了 A100 结果。
精度结果
表 1(来自论文):各 benchmark 平均精度汇总
MiniCPM4.1-8B NOSA-8B
方法 HELMET LongBench RULER 推理 Avg HELMET LongBench RULER 推理 Avg
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
图 6(论文表 1):各方法精度汇总。 SparDA 在两个模型上均与稀疏注意力基线持平或略好。InfiniGen 因隐状态代理不可靠,在 HELMET 和 RULER 上大幅下滑。
几点有趣的观察:
MiniCPM4.1-8B:SparDA 在 RULER(+0.5)和推理(+1.1)上优于 Sparse,在 LongBench 上持平(+0.1),HELMET 微降(-0.6)。推理任务的提升说明精细 KL 训练让 Forecast 更能识别”长程推理需要哪些 block”,而不仅仅是局部相关性。
NOSA-8B:平均提升更大(+2.3),推理任务大涨 +6.5——这是论文中最令人印象深刻的一个数字。可能的解释:NOSA 的 query-agnostic eviction head 基于全局 token 流行度统计做 block 淘汰,不针对具体查询。对于多步推理问题(数学、AIME),真正重要的上下文往往是某一特定推理步骤所需的”定义”或”中间结果”,它们在全局意义上可能并不”流行”,但对当前查询向量非常关键。SparDA 的 Forecast 通过学习逼近实际注意力查询的 block 分布,能够识别出这些查询相关的 block 并把它们”找回来”,因此推理分数大涨。这是一个很好的例子,说明了学习式选择在结构化、查询相关的上下文依赖任务上胜过启发式选择的原因。
InfiniGen 的失败:HELMET 下降了 8.2 点(MiniCPM4.1-8B)到 11.2 点(NOSA-8B)。相邻层之间隐状态经过 FFN 和 Norm 变换后,结构差异大,直接用来预测下一层的 block 选择是不可靠的。
Dense vs. Sparse 的精度差距:两个模型在 HELMET 和 RULER 上都落后 Dense 约 2–7 点,这是稀疏预训练时上下文偏短(32K/16K vs. 评测时 64K/128K)的锅,和 SparDA 无关。
RULER 长度泛化(表 2):
SparDA 在全部序列长度下都优于 Sparse,NOSA-8B 上优势随上下文增长持续扩大(最大 +4.3@128K),说明 Forecast 学到的 block 相关性表示能泛化到训练长度之外。
效率结果
注意力时间拆解(论文图 3):
在 MiniCPM4.1-8B 上把每层注意力时间分解成块选择(绿色)和块稀疏注意力计算(蓝色):
预填充阶段:块稀疏注意力计算时间主导且基本不随上下文增长,而选择时间随序列线性增长,在 128K 时接近注意力本身。SparDA 将选择时间压缩最多 2.50×,注意力计算时间不变。
解码阶段:每步只有一个新 token,块稀疏注意力极快,选择反而成了瓶颈——InfLLM-V2 的多头选择器随序列增长而增加,SparDA 的 Forecast 索引器几乎不随序列长度变化,在 128K 时把选择时间削减超过 2×。
预填充吞吐(H100,批大小 4):
128K 上下文,MiniCPM4.1-8B:SparDA 达到 17,087 tok/s,Sparse 为 13,661 tok/s(1.25× 加速),Dense 只有 8,085 tok/s(2.11× 加速,Dense 二次方扩展拖累了预填充)。预填充阶段卸载本身几乎不影响性能(Sparse† 和 Sparse 的预填充吞吐接近),所以加速完全来自 Forecast 索引器的更低选择代价。
解码吞吐(H100,128K 上下文):
128K 上下文解码吞吐,MiniCPM4.1-8B(tok/s)
方法 B4 B8 B16 B32 B64 B128
Dense† 108.6 — — — — —
Sparse† 189.5 — — — — — (B8+ OOM)
Sparse 167.8 279.5 447.9 618.6 788.9 —
InfiniGen 51.8 66.5 85.6 117.5 — —
SparDA 240.2 471.2 705.3 899.2 1000.1 —
图 7(论文表 4 节选):128K 上下文解码吞吐对比。 SparDA 在全部批大小下都领先。InfiniGen 因 CPU 侧聚合瓶颈,在 B32 时仅 117.5 tok/s,远不及 SparDA 的 899.2 tok/s。
几个关键数字:
- B4(交互式推理场景):SparDA 240.2 vs. Sparse 167.8,1.43× 加速——小批大小时 PCIe 传输主导,lookahead 带来的比例收益最大
- B16:SparDA 705.3 vs. Sparse 447.9,1.57× 加速
- B64:SparDA 1000.1 vs. Sparse 788.9,1.27× 加速
- 相比无卸载的 Sparse†(B4 时 189.5 tok/s,B8+ OOM):SparDA B64 达到 1000.1 tok/s,5.28× 更高吞吐(卸载使批大小上限从 B4 扩展到 B64+)
为什么 InfiniGen 这么慢? InfiniGen 在 CPU 端先做 top- block 的聚合,再发起传输——CPU 成了调度瓶颈。128K 加大批时,CPU 侧的循环串行化了所有层和批次元素的处理,严重拖慢了吞吐。SparDA 的全部调度在 GPU 上(Forecast 索引器),CPU 只是被动的 DMA 数据源,完全不在关键路径上。
批判性分析:不足与可改进之处
不好的地方:方法的弱点
1. 评测规模太小——仅两个 8B 模型
所有实验都在 MiniCPM4.1-8B 和 NOSA-8B 上进行,两者都用 InfLLM-V2 作为稀疏骨干。论文声称 Forecast 原则可以泛化到 DSA/CSA(DeepSeek-V3.2、GLM-5、DeepSeek-V4 使用的架构),但没有任何实验支撑。671B 参数的 MoE 模型与 8B 的密集模型在层间隐状态结构上差异巨大,Forecast 的训练效果、精度和效率表现都可能非常不同。这个”泛化性”的宣称缺乏实证基础。
2. Dense vs. Sparse 精度差距完全未被解决
两个模型在 HELMET 和 RULER 上持续落后 Dense 约 2–9 点。论文归因于”稀疏预训练时上下文偏短”,但没有进一步分析这个差距如何随上下文增长变化,也没有讨论是否能通过更长的稀疏预训练来缩小差距。对于精度要求严格的实际部署场景,这个未被解决的差距是一个实质性障碍。
3. 与 InfiniGen 的对比不够公平
InfiniGen 的低吞吐主要来自 CPU 侧的聚合瓶颈,而不一定反映其预测机制本身的优劣。如果给 InfiniGen 也配上同样的 UVA 持久化内核(把 CPU 侧聚合换成 GPU 侧调度),两者的吞吐差距会缩小多少?论文没有做这个实验,把内核工程优化和预测机制优化的贡献混在一起了,读者无法分清哪部分是 Forecast 设计本身的价值。
4. 缺少 Forecast 头数量的消融实验
SparDA 固定用每 GQA 组 1 个 Forecast 头。论文没有给出 1 头 vs. 2 头 vs. 4 头的对比——不清楚 1 头是不是帕累托最优,还是说 2 头能在很小的开销下带来有意义的精度提升,尤其是对长程推理任务而言。
5. 没有延迟指标(TTFT、TPOT)
论文只报告了吞吐量(批次整体 tok/s)。对于交互式应用——长上下文模型最主要的部署场景之一——首 token 延迟(TTFT)和单 token 生成时间(TPOT)才是关键指标。批大小 1 的延迟数据完全缺失。
6. CTA 配置调优是硬件相关的,指导不足
论文只提供了 H100 的 CTA 配置(附录 D 表 7),A100 有结果但没有配置指导。对于部署在非 H100 硬件(H200、MI300X)的团队,需要重新调优 CTA 数目,而论文没有提供分析模型或调优策略。
作者淡化或回避的局限
论文明确说 SparDA 的精度”受限于基础稀疏骨干的质量”——这是对的,但说得太轻描淡写。基础稀疏骨干的质量取决于 top- 比率、block 大小、预训练长度分布。SparDA 训练时只是让 Forecast 去逼近原始选择器,不会比原始选择器更好。如果骨干用的 很小(激进稀疏),精度上限就被压得很低,Forecast 无论学得多好也没用。
另外,SparDA 要求 KV 缓存存在**固定内存(pinned memory)**中,才能高效走 UVA DMA。固定内存不可被操作系统换出,在多租户服务器上,数十 GB 的固定内存可能和其他进程产生严重竞争。这个实际部署约束论文一字未提。
具体的改进建议
1. 至少用一个较大的 DSA 骨干模型做实验:哪怕只跑 RULER + MATH-500,在一个 70B 级别的 DSA 模型上跑通 SparDA,泛化性的声称就能站稳脚跟。
2. 补充 B1 和 B4 的 TTFT / TPOT 延迟数据:加一张延迟表,和吞吐表并列,覆盖交互式推理场景。
3. 做 Forecast 头数量的消融:1 vs. 2 vs. 4 头,在 RULER 上报告精度-效率权衡,确认单头是否为最优配置。
4. 隔离内核工程贡献:用 UVA 持久化内核重新实现 InfiniGen,单独对比 Forecast 机制(训练向量)vs. 隐状态代理的预测质量,给出更干净的精度和吞吐对比。
5. 量化固定内存需求:给出不同模型大小 × 上下文长度 × 批大小下的固定内存占用表,并讨论多租户场景下的内存管理策略。
深入分析:解耦为什么有效
标准稀疏注意力的调度死锁
在标准块稀疏注意力中,一步解码在第 层的计算图是:
X_l → [Q_l, K_l, V_l] → 块选择 (Q_l vs K̃_l) → 从 CPU 取 KV_l → 稀疏注意力 → X_{l+1}
每一步都严格依赖上一步。“从 CPU 取 KV_l” 是同步操作:GPU 必须等 PCIe 传输完成才能继续。唯一的出路是在需要数据之前就开始传输——但在标准图中,你无法在完成块选择之前知道需要哪些 block,而块选择需要 , 需要 , 需要上一层完成。
这是一个经典的生产者-消费者依赖:选择(生产者)和注意力(消费者)在同一层,形成硬性顺序约束,根本没有 lookahead 调度的空间。
SparDA 的突破口:跨层流水线
SparDA 的架构改变了这张依赖图:
第 l 层: X_l → [Q_l, K_l, V_l, F_l]
↓
F_l 对 K_tilde_{l+1} 打分 → 入队 B_{l+1} 预取请求
↓(异步 UVA 内核)
第 l+1 层:KV_{l+1}[B_{l+1}] 已在 GPU
X_{l+1} → [Q_{l+1}, ...] → 稀疏注意力(KV 数据已就位)
核心不变量:第 层需要 KV 数据时,该请求在第 层就已发出。PCIe 传输提前了整整一层,第 层的全部注意力 + FFN 计算时间都可以用来完成传输。
这本质上是把流水线并行的思想应用到了内存传输上:就像流水线并行用计算遮盖通信,SparDA 用当前层的计算遮盖下一层需要的内存传输。
传输时间和计算时间能对得上吗?
我们来做一个粗略的估算,验证这个 overlap 是否在实际中成立。
以 MiniCPM4.1-8B,128K 上下文,批大小 16 为例:
- 每层需要取的 KV block:top-,每 block 64 个 token, 头,,BF16:
以 PCIe 5.0 固定内存有效吞吐约 50 GB/s 估算:
对于每层计算时间(注意力 + FFN,128K 上下文,批大小 16,H100):
- 稀疏注意力:约 1–2 ms/层
- FFN(,中间维 14336):约 2–4 ms/层
- 合计约 3–6 ms/层
传输时间(1.3 ms)完全在计算窗口(3–6 ms)内,overlap 完全可行。这解释了为什么解码加速在中等批大小(B8–B16)最显著——这时 GPU 利用率适中,计算窗口刚好够把传输完全藏起来。
小批大小(B4)时 GPU 利用率低,计算时间短,但传输量也小( 少),两者都缩小,overlap 效果保持。大批大小(B64+)时 GPU 满载,FFN 时间主导,传输只占计算时间的一小部分——大批下的吞吐提升更多来自卸载使得可行批大小大幅扩展(从 B4 上限扩展到 B64+)。
相关工作:SparDA 在技术谱系中的位置
理解 SparDA 需要把它放到三条相关工作路线中来看。
路线一:KV 缓存永久剪裁(Permanent Eviction)
H2O、SnapKV、ScissorHands 等方法根据历史注意力分数或统计显著性,永久丢弃”不重要”的 KV 条目。优点是内存占用低;缺点是损失精度,且决策不可逆——如果早期丢掉的 token 后来被查询需要,就再也找不回来了(例如文档开头的定义或指代关系)。SparDA 保留了 CPU 上完整的 KV 缓存,只做按需取块,不做永久丢弃,精度损失完全来自稀疏骨干本身,而非 SparDA 新引入的。
路线二:无训练稀疏注意力(Training-Free Sparse Attention)
StreamingLLM、QUEST、MInference 等方法用实时的注意力查询向量选取一部分 KV 对做注意力,不修改模型权重。这些方法精度损失很小,但 KV 数据必须全部留在 GPU,或者做同步取块——无法支持上下文远超 GPU 显存的场景。
路线三:KV 卸载稀疏注意力(Sparse + Offload)
SparseServe 和 HiSparse 把 KV 缓存卸载到 CPU,解码时同步取回需要的 block,传输在关键路径上。InfiniGen(OSDI 2024)尝试用原始隐状态 预测下一层的 block 选择,做 lookahead 预取——方向正确,但不可靠。NOSA 增加了一个 query-agnostic eviction head,减少总传输量,但仍然是响应式(reactive)而非预测式。
SparDA 的独特定位:在”可训练稀疏选择”(InfLLM-V2 路线)和”KV 卸载”之间架了一座桥——把两者组合时的主要障碍(PCIe 传输延迟)通过 lookahead 预取解决了,而 lookahead 的准确性由训练的 Forecast 保证,不依赖不可靠的隐状态代理。
复杂度分析:加速从何而来
做一个明确的复杂度分析,有助于理解 SparDA 的收益来自哪里。
符号定义: 层, 个 KV 头,每 KV 头对应 个查询头,头维度 ,序列长度 ,block 大小 ,top- 选取 个 block,压缩后 block 数 。
InfLLM-V2 每层代价:
预填充阶段块选择:
解码阶段块选择(每步一个新 token):
块稀疏注意力(预填充或解码):
SparDA 每层代价:
Forecast 选择(预填充):
Forecast 选择(解码):
新增 Forecast 投影代价(每层相对 InfLLM-V2 的新增量):
这个代价和 、、 的投影同阶,是整层计算代价中的小头。
理论加速 vs. 实测加速:
对于 MiniCPM4.1-8B(),选择代价的理论加速为:
实测在 128K 上下文下是 2.50×。差距来自:一方面 InfLLM-V2 的多头选择器有 softmax 和跨头求和的额外开销(理论分析里没算进去,但提高了分母),另一方面 Forecast 单头的内存访问更连续,cache 利用率更高,所以实测反映了包含内存访问效率在内的综合加速。
解码阶段的特殊重要性:
解码时每步只有一个新 token 参与注意力,块稀疏注意力本身极快(计算量只有预填充的 )。这时选择代价和注意力代价的比值变得很高:选择需要对所有 个压缩 block 打分( 随序列增长),而注意力只需对 个 block 做计算( 固定)。SparDA 的 Forecast 索引器的代价只取决于 (与 无关),这在解码阶段意义最大,直接决定了”decoding bottleneck 是选择还是传输”这个核心问题。
局限性(论文自述)
论文明确承认了三点局限:
- SparDA 只是附加组件,不修改稀疏注意力计算本身,精度上限由基础稀疏骨干决定
- Forecast 原则预计可泛化到 token 级稀疏注意力(DSA/CSA),但尚未实验验证
- 未在 DeepSeek-V3.2、GLM-5、DeepSeek-V4 等真实部署规模最大的模型上做评测
可复现性
代码已开源:https://github.com/NVlabs/SparDA。训练只涉及 Forecast 投影(33.5M 参数),基模型冻结,训练门槛较低。压缩 Key 缓存 的增量维护是 InfLLM-V2 的标准组件。持久化 UVA Triton 内核需要支持 UVA 的 NVIDIA GPU(Kepler 以后全部支持)和 CPU 固定内存。论文提供了压缩窗口配置、CTA 数目表和训练超参数,信息基本充分。
与 CPU 硬件预取的类比
SparDA 的机制在计算机体系结构中有一个直接的类比:CPU 的硬件预取器(Hardware Prefetcher)。现代 CPU 会根据内存访问的步长模式(stride pattern)或链表跟随模式,提前把即将需要的 cache line 从 DRAM 预取到 L1/L2 缓存,让”实际加载”指令执行时数据已经在缓存里。
SparDA 做的事本质相同,只是在 GPU↔CPU 内存层级上:
| 类比维度 | CPU 硬件预取 | SparDA Forecast |
|---|---|---|
| 预测机制 | 步长/链表规律(规则式) | 训练的线性投影(学习式) |
| 预测粒度 | Cache line(64B) | KV Block(数 KB) |
| 带宽通道 | DRAM(~200GB/s) | PCIe(~50GB/s) |
| 目标 | 隐藏 DRAM 延迟 | 隐藏 PCIe 延迟 |
为什么注意力模式需要”学习式”预取而不是”规则式”:注意力的 block 选择是数据相关的(每个查询的”步长”都不同),没有固定的地址模式,规则式预取器根本不够用。只有通过训练,让 Forecast 学会”这类查询向量一般需要那类历史 block”,预测才能可靠。
这个类比也解释了为什么精度很重要:CPU 预取器取错 cache line 会浪费带宽、污染缓存。Forecast 选错 block 同样浪费 PCIe 带宽,甚至可能把正确的 block 从 GPU staging buffer 里挤出去。KL 训练的目的正是保证 Forecast 的准确性足够高,使得”带宽浪费”可以忽略不计。
工程落地:实际部署的考量
固定内存的配置要求
UVA 持久化内核要求 KV 缓存存在**固定内存(pinned memory)**中,原因是 GPU 的 DMA 引擎需要物理地址稳定,普通的 malloc 分配的内存可能被操作系统换页,导致 DMA 失败或需要额外的 IOMMU 映射开销。
在 PyTorch 中的实现:
kv_cache = torch.empty(L, H_kv, T, d_kv, dtype=torch.bfloat16).pin_memory()
代价:固定内存不可被操作系统换出,减少了物理内存调度的灵活性。对于一个 70B 模型,128K 上下文,批大小 16,固定内存 KV 缓存约 175 GB——需要专用高内存服务器(1TB+ DRAM),不适合消费级 GPU 或多租户云实例。
接入 vLLM / TGI 的关键改动点
把 SparDA 接入 vLLM 或 TGI 需要以下改动:
- KV 缓存分配器:改用
cudaMallocHost分配固定内存,替换原有的 GPU HBM 分配 - 模型权重:把 Forecast 投影矩阵加入模型 state dict,checkpoint 中保存和加载
- 调度器钩子:在 layer 执行循环中嵌入 UVA 预取的分发逻辑
- CTA 调优:按照论文附录 D 表 7 为目标 GPU 配置合适的 CTA 数量
论文代码仓库(https://github.com/NVlabs/SparDA)提供了独立的 InfLLM-V2 + SparDA 实现,移植到 vLLM/TGI 等生产框架需要额外的工程工作。
工程接入步骤(简明版):
- 取一个已完成 InfLLM-V2 两阶段稀疏预训练的模型
- 为每层每个 GQA 组添加 Forecast 投影矩阵
- 冻结所有基模型权重,只训练 矩阵
- 训练时:前向传播收集 、、;计算细粒度目标分数;计算推理时粗粒度预测分数;用 top- 限制的 KL 损失反传,梯度只流向
- 部署时:KV 缓存需分配在固定内存(pinned memory)中,可通过
torch.Tensor.pin_memory()或cudaMallocHost()实现
更广泛的意义:SparDA 告诉了我们什么
SparDA 不只是一篇解决了某个特定问题的论文,它的设计思路有几个更广泛的启示。
1. 调度信号与执行信号应该解耦
在几乎所有的 transformer 操作中,驱动计算的向量(查询 )同时也驱动数据选择(block 打分)。SparDA 的核心洞察是:这两个角色可以分开——一个轻量的”调度查询”(Forecast)提前一步告诉内存系统需要什么数据,而全精度的注意力查询在正常时间执行。这个解耦原则不只适用于稀疏注意力,还可以用到 MoE 的专家选择路由、KV 压缩路由、以及其他数据相关的内存访问模式。
2. 调度信号应该训练,而不是借用代理
InfiniGen 的失败是一个警示:用未经训练的信号做代理,结果往往不稳定。原始隐状态 包含了预测下一层注意力模式所需的信息,但不在一个对压缩 Key 打分好用的形式里。Forecast 投影学的正是这个线性变换:从”通用隐状态信息”到”对 block 级别 Key 打分准确的向量”。这个变换简单(一层线性)、训练成本低(0.41% 参数)、跨上下文长度泛化良好。
3. GPU↔CPU 内存层级是一个被低估的优化维度
LLM 推理优化工作大多集中在计算(FlashAttention、算子融合)或模型压缩(量化、剪枝)。GPU↔CPU 的内存层级管理相对被忽视。随着 serving 实例的 KV 缓存需求达到 10–100GB 量级,高效管理这个内存层级(预取、分层存储、压缩)将越来越重要。SparDA 是早期将模型架构设计与内存层级管理协同优化的清晰案例。
4. lookahead 窗口值得为它训练
SparDA 的一层 lookahead 已经带来了显著加速。更长的 lookahead 窗口(比如两层)能给传输更多时间,但对应的预测任务更难——需要预测两层变换后的注意力模式。SparDA 的成功鼓励了对更长 lookahead 的探索。
总结
SparDA 是一个设计简洁、动机清晰的工作。它解决的问题——如何让 KV 卸载在大批 + 长上下文下真正高效——是 LLM serving 系统的实际痛点,解法也足够优雅:把块选择从注意力中”解耦”出来,提前一层预测,异步传输,三件事互相配合。
论文最强的地方在效率侧:解码吞吐的实验设计合理,对比方法选择合适,分析层次清晰(从注意力时间拆解到吞吐数字,一步步说明加速来自哪里)。5.3× 的解码吞吐提升不是挑了一个好看的数字——它来自两个独立因素叠加:1.7× 的单层解码加速(来自 lookahead overlap 和紧凑索引器),加上 ~3× 的可行批大小扩大(来自卸载使 GPU 显存压力降低),两者相乘才得到 5.3×。这两个因素都有独立的实验支撑,不是凑数。精度实验的结论可信,但受制于两个 8B 模型的规模,泛化性说服力有限。NOSA-8B 推理任务上 +6.5 的结果暗示了一件有意思的事:训练式 block 选择在结构化任务上优于无训练基线,而不仅仅是与之持平——这值得进一步探索。
对于在 InfLLM-V2 或 NOSA 骨干上搭建长上下文 serving 系统的工程师来说,SparDA 的接入成本很低(训练 0.41% 的参数,加一个 UVA 内核),收益明确。对于研究者来说,这篇论文最有价值的洞察是:块选择是一个”可调度的、可训练的信号”,而不是必须绑定在当前注意力计算上的在线计算——未来稀疏注意力系统的设计都应该认真考虑这个原则。
常见问题解答
Q:为什么 SparDA 选择每 GQA 组一个 Forecast 头,而不是每 KV 头一个或每查询头一个?
每 KV 头一个(等价于整层只有一个 Forecast 头)最便宜,但无法区分不同 GQA 组的注意力偏好,精度太差。每查询头一个则等于复制了原始的多头选择器,代价优势全没了。每 GQA 组一个是自然的折中:不同组可以专门化自己的 block 预测,同时代价比每查询头低 倍。InfLLM-V2 的 DSA 在 token 级别也用了”每 KV 头一个头”的设计——SparDA 把同样的原则用在了 block 级别。
Q:Forecast 预测错了 block 会怎样?
如果 Forecast 漏选了 实际需要的一个 block,该 block 就不会参与注意力——等效于对那一步做了 block eviction。论文的实验结果表明,SparDA 的精度与 Sparse 基线持平或略好,说明 Forecast 的漏选率在可接受范围内。预测错误也会浪费 PCIe 带宽(传了用不上的 block),但论文显示这在实测中对吞吐影响可以忽略。
Q:SparDA 能用于标准密集注意力模型(比如 LLaMA-3、Qwen-3)吗?
不能直接用。SparDA 依赖模型已经做过稀疏预训练(InfLLM-V2 的两阶段训练),使得推理时注意力模式本身就是稀疏的——否则 Forecast 预测了 block 索引,但注意力还是会读全量 KV 缓存,预取没有意义。要把 SparDA 用于密集模型,必须先做稀疏预训练,这是一个不小的门槛。
Q:训练 Forecast 需要多少算力?
训练只涉及 33.5M 参数(0.41% 的 8B 模型),基模型冻结,可以用相对短的序列(不需要 128K 训练长度)。论文没有精确报告训练时间,但从参数规模和任务难度估算,应该在数百 GPU 小时量级(远低于稀疏预训练的代价),对于想在已有稀疏模型上叠加 SparDA 的团队来说,这是可接受的成本。
未来方向
基于论文的发现和局限性,我认为以下几个方向最值得探索:
1. 将 Forecast 原则扩展到 token 级别稀疏注意力(DSA/CSA)
DeepSeek-V3.2 使用的 DSA 是 token 级别的稀疏选择,而不是 block 级别。将 Forecast 原则用于 token 级别,意味着预测每个 token 的重要性分数而不是每个 block 的。架构变化小,但训练挑战更大——token 级别的监督需要更细粒度的数据集,且预测难度更高。这是论文明确指出的未来方向。
2. 多步 lookahead
当前设计提前一层预测。提前两层或三层可以给更长上下文下更长的 PCIe 传输窗口。挑战是多层变换后隐状态结构差异更大,Forecast 的预测准确性会下降。从系统角度看,更长的 lookahead 窗口可以完全隐藏长上下文(128K+)的 PCIe 传输延迟,这对于配置了 NVLink 的超大集群外的单机推理场景特别有价值。
3. 自适应 top-k 与 Forecast 置信度
当前设计对所有层和位置使用固定的 。一个置信度感知的 Forecast 可以在预测不确定时用更大的 (多取几个 block 作为保险),在高置信时用更小的 (节省带宽)。这与 Speculative Decoding 的可变草稿长度思路相似,但应用于 KV block 的选择而非 token 的生成。
4. 与前缀缓存(Prefix Caching)的结合
在多轮对话或共享前缀的场景中,同样的 KV blocks 会被多个请求需要。Forecast 感知的前缀缓存系统可以根据 Forecast 的预测,把高概率需要的 block 提前保留在 GPU 显存里,进一步减少每次请求的 PCIe 传输量。这是 SparDA 和 KV 共享缓存(如 Mooncake、SGLang 的 RadixAttention)的天然结合点。
5. Forecast 作为压缩引导信号
Forecast 的 block 重要性分数直接反映了每个 KV block 对下一层注意力的贡献程度。这些分数可以用来引导 KV 量化:重要性分数持续低的 block 可以用更低精度量化或合并,重要性高的保持全精度。这样可以把 SparDA(稀疏选择)和 KVQuant 等量化方法结合起来,形成一个由 Forecast 驱动的混合精度 KV 缓存管理方案。
这几个未来方向的共同主线是:Forecast 不只是一个调度提示,它是一个学习到的重要性信号,编码了”哪些历史 token 对未来计算有价值”的模型级知识。充分利用这个信号,可以统一驱动预取、压缩、淘汰和缓存——构建一个更完整的长上下文 LLM 推理内存管理系统。
论文总结一览表
| 维度 | SparDA |
|---|---|
| 核心创新 | Forecast 投影解耦块选择与注意力计算,实现提前一层的 KV 预取 |
| 参数开销 | +33.5M(占 8B 模型的 0.41%),仅 Forecast 权重矩阵 |
| 训练成本 | 低:基模型冻结,短序列训练,标准 KL 目标 |
| 预填充加速 | 最高 1.25× 于稀疏卸载基线(来自紧凑 Forecast 索引器) |
| 解码加速 | 最高 1.7× 于稀疏卸载基线(索引器 + 预取 overlap 共同贡献) |
| 解码吞吐 | 最高 5.3× 于无卸载稀疏注意力(通过更大可行批大小实现) |
| 精度 vs Sparse | 持平或略好(MiniCPM4.1-8B: +0.3 avg;NOSA-8B: +2.3 avg) |
| vs InfiniGen 优势 | 更准确(训练式 vs. 无训练代理);更快(GPU-native UVA vs. CPU 聚合) |
| 主要局限 | 精度受限于稀疏骨干;仅在 8B InfLLM-V2/NOSA 模型上验证 |
| 最大开放问题 | 能否泛化到 DSA/CSA 骨干(DeepSeek-V3.2、GLM-5、DeepSeek-V4) |
| 代码 | https://github.com/NVlabs/SparDA |