稀疏注意力最容易被一句 带偏。这个量级描述的是稠密注意力随序列长度增长的关系数,却没有告诉我们机器在训练、Prefill 和 Decode 阶段分别搬了多少数据,也没有说明所谓“稀疏”删掉的是注意力边、KV 状态、缓存读取,还是某一层的完整 softmax。两个方法都声称把复杂度降到 ,一个可能真的少读了 KV,另一个却先扫描完整历史再产生 top- 索引;前者能降低长序列 Decode 的带宽,后者可能只减少索引之后的点积。比较路线之前,先把成本对象拆开。
设 batch 为 ,序列长度为 ,query head 数为 ,KV head 数为 ,head dimension 为 ,每个元素占 字节。稠密训练和 Prefill 的注意力点积大致随 增长,KV cache 容量则约为 。Decode 每一步只产生一个 query,点积数对当前历史长度 呈线性增长;更麻烦的是历史 K/V 会被逐层重新读取。于是训练常由计算、激活和反向主导,长 Prompt 的 Prefill 同时受计算与 kernel 形状影响,长 Decode 更容易受 HBM 带宽约束。只报一个复杂度,无法判断优化发生在哪个阶段。
稀疏方法至少作用于四种不同对象。第一种是关系稀疏:每个 query 只允许访问窗口、块或候选 token,减少 QK 点积和 softmax 项。第二种是状态稀疏:只保留一部分历史 KV,直接缩小 cache,但被驱逐的信息通常无法恢复。第三种是读取稀疏:完整 KV 仍在显存或分层存储中,每步仅加载少量页或 token;它节省流量,不节省总容量。第四种是层稀疏:局部层、线性层与全局层交替,成本由各类层的比例共同决定。KDA 一类有限状态方法还会把历史压进固定大小矩阵,它省掉逐 token 地址,误差来源也从“候选漏召回”变成“有限状态碰撞”。
这四类对象不能用同一个稀疏率代替。关系稀疏率可以写成保留边数与因果稠密边数之比;状态压缩率比较保留 KV 字节与完整 KV 字节;读取比例统计一次请求真正从 HBM 取出的页或 token;物理工作量还要按 kernel 实际处理的 tile 向上取整。逻辑上选择 17 个 token,kernel 若必须装入两个 64-token tile,真实工作量仍接近 128 个 token。索引重复、候选排序、跨 head 对齐和变长 batch 填充也都要计入。
三个阶段要分别算账
训练阶段需要前向、反向和优化器更新。一个 mask 在数学上很稀疏,不等于训练实现真的跳过了被遮蔽位置;若底层仍计算完整分数矩阵再填入负无穷,FLOPs 和激活峰值并没有减少。训练原生稀疏还会引入选择器梯度、负载不均衡与跨设备通信,收益应与同规模稠密训练的 tokens/s、峰值显存和收敛曲线一起测。
Prefill 同时处理大量 query。规则窗口或块稀疏可以形成稳定 tile,较容易摊薄索引开销;细粒度 token top- 若产生离散地址,可能让计算量下降却使访存更加破碎。动态方法还要回答索引在哪一层生成、能否跨 head 或跨层复用、索引阶段是否读过完整 K。百万 token 场景中,离线模式搜索与在线候选生成的边界会直接决定首 token 延迟。
Decode 每步 query 数很少,矩阵乘容易退化成带宽受限的矩阵向量运算。此时减少 KV 读取通常比减少少量算术更重要。页级选择能让 DMA 和分页 cache 保持连续,却要接受页内无关 token 一起被加载;token 级选择更精确,却需要索引、聚集读取和变长归并。状态驱逐把容量直接封顶,但错删后没有回退路径。比较时至少同时记录 TPOT、每 token 读取字节、KV 总容量和候选召回。
从逻辑集合走到物理执行
用 表示位置 的逻辑候选集合,只能说明哪些边参与 softmax。真正执行还要把 映射为块表、页表或压缩索引。设物理 tile 大小为 ,某行访问的独立 tile 数为 ,则该行至少装入 个位置;当候选零散时, 可能远大于 。多 query 是否共享 tile、GQA 的 query heads 是否共用 KV 页、同一 KV 是否因地址乱序被重复读取,都会改变有效带宽。
因此,算法层的“选择正确”与系统层的“访问划算”是两个连续关口。候选集合先要覆盖高价值关系,再要能排列成硬件可以批量执行的工作表。块级方法牺牲选择精度换连续访存,细粒度索引则需要通过跨头、跨层共享或残差补偿降低碎片成本。这个权衡贯穿后续所有路线。
公平比较需要固定什么
比较稠密与稀疏模型时,应固定模型规模、训练 token、上下文长度分布、数据顺序和优化器预算。比较现有模型的推理近似时,应使用同一 checkpoint、同一请求集合和相同精度,并把索引时间计入端到端延迟。若方法允许回退到 full attention,还要分别报告固定预算与固定质量两种结果:前者回答同样成本下谁更准,后者回答保持质量需要付出多少成本。
质量侧不能只看平均任务分数。路由类方法应记录候选 recall、保留 attention mass、关键证据覆盖与最坏样本;状态压缩要观察信息写入后的干扰与遗忘;缓存驱逐要记录被删除 token 后续重新变重要的比例。系统侧则至少拆开 kernel 时间、索引时间、数据搬运、峰值显存、KV 容量和并发吞吐。单点加速比离开 batch、长度、硬件和回退策略就没有可比性。
成本模型从公式和数值算例展开这些量,并把训练、Prefill、Decode 分别代入。后续阅读静态拓扑、动态路由、KV 选择与有限状态路线时,可以反复追问四件事:它删掉了哪类对象,索引本身花了什么,逻辑稀疏怎样落到物理 tile,以及质量损失能否沿候选、状态或传播路径定位。回答清楚以后,“稀疏”才从一个形容词变成可核算的设计。
一个容易混淆的数值例子
假设历史长度为六万四千,每个 query 逻辑上保留四千个 token。token 级选择的名义比例是十六分之一;若这些 token 均匀散落在一千个页面中,每页六十四个位置,设备仍要读取接近完整历史。另一个选择器只保留三千五百个 token,却集中在一百个页面中,逻辑召回可能略低,物理流量反而小得多。因而预算表应同时列出 token 数、独立页面数、独立 tile 数和真实字节,不能只写 top-k。
同样的区别也出现在层轴。若三十二层中只有四层使用 full attention,其余层用窗口,完整 KV 可能只在四层增长;若实现为了接口统一仍为所有层保存完整 KV,理论状态压缩并没有兑现。有限状态层若每层保留矩阵状态,还要把状态维度、读写次数与少数 full 层的 cache 一并计算。模型结构图和内存分配表必须能够互相对上。
成本估算最终应落到可证伪的预测。例如,若判断 Decode 受带宽约束,那么减少一半实际 KV 字节应明显降低单步时间,而单纯削减少量算术不会同比生效;若实测相反,应检查 cache 命中、并发、索引流水或 attention 是否已经不再是主导项。复杂度分析的价值就在这里:它给出需要被测量验证的因果预测,而不是替代真实测量。