5 · 内核：把逻辑稀疏变成真实加速

Sparse Attention 写在算法图里很简洁: 每个 query 只连接少量 key, 其余边从计算图中消失。但 GPU 不会因为 mask 里有大量零就自动少做工作。若实现仍然生成完整的 $QK^T$ , 再把不允许访问的位置改成负无穷, 点积、logits 写回和大部分数据读取都已经发生。真实加速要求无效的 query-key tile 在调度之前就不存在。邻接集合必须进一步变成设备可执行的块表、索引、访存与同步顺序。

这项转换决定了前几章讨论的稀疏模式能否进入训练和在线推理。固定窗口、全局 token 与规则块可以预先编排, 内容路由和 learned top-k 却要在运行时生成选择结果。两类方法在数学上都减少注意力边, 在 kernel 中承担的元数据成本、负载不均和控制流却完全不同。内核设计不能只接收一个抽象的稀疏率 $s$ ; 它还需要知道有效边是否连续、每行块数是否一致、索引何时生成, 以及同一个选择能否跨 head、跨层或跨 decode step 复用。

1. 从注意力公式到设备工作

1.1. 被省掉的究竟是哪一段

设序列长度为 $N$ , 单个 head 的维度为 $d$ , query $i$ 允许访问的 key 集合为 $S_i$ 。稠密注意力对每个 $i$ 都读取 $N$ 个 key 和 value; 稀疏形式只需读取 $|S_i|$ 个候选。算法层的点积量从 $O(N^2d)$ 降到 $O(\sum_i |S_i|d)$ 。这个表达式只计算乘加次数, 没有包含索引读取、候选排序、数据重排、线程空闲和 kernel 启动。若 $S_i$ 分散在显存的许多位置, 读取量减少以后, 带宽利用率也可能同时下降。

因此需要分开记录三类工作。第一类是从稠密版本中确实删除的 QK 点积和 PV 聚合; 第二类是仍需执行的投影、位置编码、在线 softmax 与输出写回; 第三类是稀疏方案新增的路由、索引和布局转换。只有第一类节省量大于第三类附加量, 稀疏结构才越过实际性能交叉点。 这个交叉点随序列长度、batch、head 维度、数据类型和硬件变化, 不能由模型配置中的稀疏率单独确定。

FlashAttention 提供的参照系。 FlashAttention 已经避免把完整 logits 矩阵写入 HBM。它按 tile 读取 Q、K、V, 在片上维护每行最大值 $m_i$ 、指数和 $l_i$ 与部分输出, 随后完成归一化。Sparse Attention 要保留这种 online softmax 的数值稳定性, 同时只遍历有效 KV tile。一个 query block 若选择多个不连续 KV block, 每次载入后都要用新的局部最大值更新 $m_i$ , 并按比例修正此前累积的 $l_i$ 与输出。跳过 tile 并不意味着可以改变 softmax 的归一化集合。

稠密 FlashAttention 是必要基线。它的计算量更大, 但 tile 规则、访问连续, tensor core 利用率高。稀疏实现若只与生成稠密 logits 的朴素 attention 比较, 得到的加速不能说明它胜过当前常用内核。以优化后的稠密实现为参照并拆出路由时间，才能区分算法节省与弱基线造成的表面收益。

1.2. 块布局、索引与负载

块大小在两种浪费之间取值。 GPU 适合处理形状稳定的矩阵 tile。把 token 级选择提升为 block 级选择, 可以让一个索引对应一整块 QK 乘法, 也能让多个线程连续读取 K、V。代价是块内填充: 只要块里有一个 key 有效, 整块都可能被计算。块越大, tensor core 越容易吃满, 但逻辑稀疏率和物理稀疏率的差距越大; 块越小, 无效乘加减少, 索引数量、分支和调度开销却会上升。

块大小不能脱离模式讨论。Sliding Window 的相邻块天然连续, 边界上少量填充通常可控。动态路由的候选可能散落, 还可能让某些 query block 选择很多 KV block、另一些只选很少。此时固定网格会让同一 thread block 中的线程工作量不同。实现可以按候选数分桶, 也可以使用 persistent kernel 从任务队列中持续取块; 前者增加排序与重排, 后者要求仔细处理队列竞争和输出归并。

索引也是计算路径的一部分。规则拓扑的块坐标可以在模型加载时或 shape 确定后生成, 随后被多个层复用。内容相关拓扑必须先计算路由分数, 再做 top-k 或阈值选择, 最终产生压缩行格式、块坐标表或 ragged offset。索引存储本身通常不大, 但索引读取会打断规则地址计算; 选择结果若先落到 CPU, 再回传 GPU, 同步成本足以抵消 attention 的节省。

较好的实现让路由、选择和 sparse attention 留在同一设备流中, 并尽量复用选择结果。不过复用有明确边界: 跨 head 共享会减少元数据和排序, 也会限制每个 head 的独立检索; 跨层共享节省路由开销, 但后续层看到的表示已经变化; decode step 之间复用历史选择能减少重复搜索, 却可能漏掉刚出现的重要 token。索引缓存同时改变性能与模型可表达的连接, 因此不能被当作纯工程优化。

2. Prefill、Decode 与训练

2.1. Prefill 和 Decode 不是同一种 shape

Prefill 一次处理很多 query, Q 与 KV 两侧都有可供矩阵乘法利用的长度。规则 block-sparse 在这个阶段最容易形成足够大的 tile, 路由开销也能被大量 query 分摊。若动态选择按 query token 独立执行, 候选集合会高度 ragged; 工程上常把相邻 query 合成一组, 用共同候选换取更规则的执行。合并提高了实际访问的 key 数, 因而报告性能时必须给出物理执行块数, 不能只给路由器输出的 token 级 top-k。

Decode 每一步通常只有一个新 query, 历史 KV 很长。此时瓶颈往往从矩阵乘法转到 KV 读取, 稀疏选择可以直接减少显存流量, 但小 Q 也更难填满 GPU。多请求 continuous batching 能提供并行度, 前提是各请求的候选长度和 cache 布局允许合批。Page-based KV cache、候选 page gather 和 block table 的关系会直接影响访问是否合并。一个 prefill 很快的内核不能据此推出低 batch decode 也快。

反向传播必须重现前向语义。训练不仅要重算或保存 softmax 统计量, 还要决定选择操作是否接收梯度。固定拓扑的反向可以遍历相同块表, 为 Q、K、V 累积梯度。动态 hard top-k 通常只让被选边进入 attention 梯度, 路由器另有代理目标、straight-through estimator 或可微近似。若 backward 重新执行路由而结果与 forward 不同, 梯度就不再对应前向计算; 因而选择结果、随机种子或足以重建选择的状态必须保留。

保存完整选择会增加训练显存, 重算则增加计算。稀疏 backward 还面临写冲突: 多个 query block 可能同时更新同一 KV block 的梯度。原子加、分阶段归约和重新排列任务各有不同代价。训练性能需要同时报告 forward、backward 与优化器步的占比, 只展示 forward kernel 会高估端到端收益。

数值精度也不能从稠密实现直接假定。被选块的遍历顺序改变以后, online softmax 的归约顺序随之变化, FP16 或 BF16 下会产生不同舍入误差。Backward 使用原子加时, 并发写入顺序还可能带来非确定性。实现验收既要与高精度参照比较输出和梯度误差, 也要在不同布局、长度与随机选择下覆盖空候选、重复索引和越界索引。追求 bitwise 一致会限制部分高性能归约, 但允许的误差范围必须在训练稳定性实验中确定。

2.2. 如何阅读这一章

可编程掩码从另一端进入内核问题：研究者用 mask_mod 与 score_mod 写清注意力语义，编译器生成融合前向、反向和decode路径。FlexAttention进一步拆解BlockMask、online softmax、分页KV、动态selector接入与完整基准。它和稀疏执行形成两条互补路线：一条解释专用工作表怎样喂给kernel，一条解释通用编译器怎样避免为每种组合重写算子。

算法输出的 $S_i$ 要先压成块表或其他设备索引, prefill kernel与decode gather再据此读取KV。实现时还要确认有效 QK 与 PV 工作如何调度, 分别处理单 token decode、训练梯度和动态 shape。各篇共享的判断标准是「哪些字节没有读、哪些乘加没有发射、为此新增了多少工作」。

稀疏执行从索引格式、tile 与变长批处理进入设备路径；索引到 Kernel沿着候选整理、地址生成、KV 装载、online softmax 与输出归并逐步展开，并给出 roofline 手算和 profiler 定位方法。

它与前一章「训练」的关系在于, 稀疏模式的学习目标会决定索引是否动态、选择是否需要保存以及梯度走向; 与后一章「评测」的关系在于, kernel 吞吐只是评测的一层。读完这一章, 应当能够从一张稀疏注意力结构图继续追到块大小、索引格式、online softmax 状态、KV cache 布局和端到端关键路径, 并判断论文报告的理论复杂度是否可能兑现为实际训练或推理收益。

3. Online softmax 为什么能分块合并

一行 attention 的 logits 被拆成若干块。对已处理集合 $A$ ，保存最大值 $m_A$ 、指数和 $l_A$ 与未归一化 value 累积 $u_A$ ：

m_A=\max_{j\in A}s_j,\qquad l_A=\sum_{j\in A}e^{s_j-m_A},\qquad u_A=\sum_{j\in A}e^{s_j-m_A}v_j.

新块 $B$ 到来后，令 $m=\max(m_A,m_B)$ ，则

l=e^{m_A-m}l_A+e^{m_B-m}l_B,\qquad u=e^{m_A-m}u_A+e^{m_B-m}u_B.

最终输出为 $u/l$ 。这组状态可以按任意块顺序合并，不需要把完整 score 行写回 HBM。FlashAttention 的关键收益来自这种分块归约与片上复用，而非近似 softmax。Sparse Attention 只枚举允许的 KV 块，也能沿用同一合并律；若候选出现重复索引，重复块会被计算两次并改变概率，必须在进入 kernel 前去重或由格式保证唯一。

空候选也需要明确定义。因果窗口通常至少包含自身，动态选择却可能因阈值或边界错误产生空行。对空集合，最大值和分母没有合法实数，直接执行会出现 NaN。实现可强制加入局部位置或 sink，也可以令空行输出零并由残差接管；训练和推理必须采用相同规则。

3.1. 块化带来的计算放大

设逻辑选择器每个 query 选中 $k$ 个 token，kernel 的 query tile 与 KV tile 分别为 $B_q,B_k$ 。只要一个 token 对落入某个 tile 对，整个块就可能被执行。定义逻辑边数 $E$ 与实际发射乘加覆盖的边数 $\widetilde E$ ，块化放大率为

\rho_{\text{tile}}=\frac{\widetilde E}{E}\ge1.

局部窗口具有连续边界，只有窗口两侧的部分块产生浪费， $\rho_{\text{tile}}$ 通常接近一。随机 token top- $k$ 会把少量边散到许多块，极端情况下每个选中 token 都触发一个 $B_qB_k$ tile，逻辑稀疏率很高，物理工作却很密。若选择语义就是整块，块级路由会把块内合法token一同纳入softmax；若块只是token候选的装载单位，仍要用成员mask排除未选位置。两者读取同样的块，也可能得到不同输出。

选择器训练若只惩罚 token 数，可能学出硬件不友好的散点。可以把候选先映射到块集合，使用块数或估计字节作为成本项；也可以让同一 query tile 的相邻 query 共享候选，提高块填充率。共享会改变模型连接，因此成本正则必须与质量损失联合训练，不能等部署时静默合并。

4. 索引格式对应不同遍历方向

前向计算以 query 行为中心，常用每个 query block 对应的 KV block 列表，形式接近 BCSR。计算 $dK,dV$ 时，多条 query 行会写向同一 KV block。若仍沿前向顺序遍历，需要原子累积；若预先构造转置索引，可让每个任务负责一个 KV block，顺序读取所有相关 query。转置表增加元数据与构造成本，却能改善确定性和写合并。

FlexAttention 的 BlockMask 同时保存非空块与完全未遮挡块。完全块不必逐元素调用 mask_mod，边界块才执行细粒度谓词。规则窗口、因果和文档边界组合后，一行可能同时含完整块、部分块与空块；分开编码可避免在内部完整块上重复判断。动态 mask 每层变化时，BlockMask 构造本身可能成为关键路径；跨层共享索引正好能摊薄这项成本。

分页 KV 又增加一次地址转换。逻辑块号先经页表映射到物理页，再得到 K/V 地址。若先把候选 KV 拷贝成连续临时张量，会多一次完整读写；融合的间接访问直接让 attention kernel 读取物理页。两者数学输出相同，带宽成本不同。页表、稀疏块表和 GQA head 映射应在索引阶段组合，避免内层循环反复做除法、取模和多级查表。

4.1. Roofline 下的稀疏转折点

一个 tile 的算术强度近似为 FLOPs 与从 HBM 读取字节之比。Prefill 中同一 KV tile 可被多个 query 行复用， $B_q$ 增大时强度上升；decode 只有少量 query，KV 基本只读一次，通常受带宽限制。删掉 KV 块会近似同比减少 decode 字节，却未必同比减少 prefill 时间，因为后者可能仍受矩阵计算或索引调度限制。

稀疏 kernel 的时间可粗分为

T=T_{\text{route}}+T_{\text{index}}+ \max\left(\frac{F_{\text{tile}}}{P_{\text{eff}}}, \frac{D_{\text{HBM}}}{W_{\text{eff}}}\right)+T_{\text{launch}}.

$F_{\text{tile}}$ 按物理 tile 而非逻辑边计算， $D_{\text{HBM}}$ 包含索引和 KV， $P_{\text{eff}},W_{\text{eff}}$ 是该 shape 下实测有效吞吐。序列较短时，路由、索引和启动项占比高；长度增长后，减少的数据与乘加才压过固定开销。转折点必须按 prefill、decode、forward 和 backward 分别测量。

5. 动态负载怎样排程

不同 query block 的候选数差异很大时，一块对应一个程序任务会导致长尾。把长行拆成多个任务可提高并行度，但它们要分别计算局部 softmax 状态，随后再做一次合并。短行合并进同一任务能减少启动，却需要 ragged offset 和边界判断。FlashInfer 使用负载均衡调度处理请求与 KV 长度的动态性，其问题与动态稀疏行相同：任务粒度既要足够均匀，也要保留数据局部性。

CUDA Graph 偏好稳定的任务形状，动态 top- $k$ 却可能每步变化。固定最大任务表再用计数屏蔽空槽，可以保留静态启动图，但过大的上限造成空任务；按长度桶维护多套图能缩小浪费，代价是更多缓存和调度逻辑。这里的固定是执行容量上限，候选内容仍可变化，不应与静态 attention 拓扑混为一谈。

5.1. 正确性基线怎样构造

稀疏 kernel 的参照不能只用另一个稀疏实现。先把块表展开成逐元素布尔 mask，在高精度下运行普通 attention，得到同一邻接集合的输出与梯度；这一步隔离实现错误。随后再与 full attention 比较，第二个差值才包含删边造成的模型近似。两个差值混在一起时，数值错误可能被误认成稀疏质量代价。

测试样例要覆盖块边界前后一位、长度不足一个块、只有一个合法 key、全局块与窗口重叠、候选重复以及不同请求共用物理页。反向还要分别检查 $dQ,dK,dV$ ，因为前向正确并不保证转置索引和并发累积正确。数值容差按数据类型与归约长度设定，并记录最大绝对误差和相对误差分布；只用单个 allclose 布尔值无法显示误差随长度增长的趋势。

确定性测试应固定候选表，重复运行并比较输出与梯度。若原子累积导致低位差异，可接受范围要通过训练回归确定。候选排序并列、NaN logits 和极大正负分数也应进入用例，它们容易触发不同设备或编译版本之间的分支差异。