Sparse attention：只看需要的历史

稀疏 mask 在 softmax 之前生效

第 i 个 query 允许读取的集合记作 Si。先选集合，再在集合内部重新归一化。

不能先做完整 softmax，再把没选中的权重置零而不归一化。那会让权重和小于 1，输出幅值也改变。

Sᵢ 必须非空；通常至少保留当前 token 或局部窗口。实现用 masked fill(−∞)，再减去行最大值求 softmax。

同样的 V，换集合就会换输出

logits = [0, ln 2, ln 3]，V = [2, 4, 8]。完整权重是 [1/6, 2/6, 3/6]，输出 17/3 ≈ 5.667。

若把完整权重的中间项直接清零，得到 1/6 × 2 + 3/6 × 8 = 4.333；它并不是这个稀疏 mask 定义的 attention。

kernel 的正确性应与“相同 mask 的 dense 参考”比较；任务质量再与原始 full attention 比较。

窗口的精确配对数，不只是 O(nw)

含自己的因果窗口宽 w。前 w−1 行还没填满，后面每行恰好读 w 个位置；假设 n ≥ w。

n = 4096、w = 256：配对数 = 1,048,576 − 32,640 = 1,015,936。full causal 是 4096 × 4097 / 2 = 8,390,656。

这组算术约少 8.26 倍配对。实际延迟还包括投影、softmax、访存和 kernel 调度。

多层窗口扩大的是间接感受野

L = 4、w = 3：最后一层位置 10 可以间接受到位置 2 的影响，范围宽 9；每层仍只直接读取 3 个位置。

这条路径经过多次压缩与变换，不等同于最后一层直接拿到位置 2 的原始 KV。精确复制远处数字，可能仍需要远程读取。

推导假设每层都是相同因果局部窗口、没有全局 token 或跨层跳接扩大可见性。

选择器先要赚回自己的时间

假想测量：full attention 2.0 ms；选择 0.35 ms、gather 0.25 ms、稀疏计算 0.50 ms，总计 1.10 ms，attention 加速 1.82×。

若其他网络层还花 6 ms，端到端从 8.0 降到 7.1 ms，只快 1.13×。小上下文中 full 只需 0.4 ms 时，选择开销可能已不划算。

选择器若扫描低维索引，仍然可能随 n 线性增长；“主 attention 只读 k 个”不是整个算法 O(k) 的保证。

top-k 也必须遵守因果约束

valid = key_position <= query_position
index_score = where(valid, index_score, -inf)
chosen = topk(index_score, k=min(k, valid_count))
chosen = sorted(unique(chosen + local_positions))
y = attention(q, gather(K, chosen), gather(V, chosen))

先 mask 再 top-k，避免未来高分把过去候选挤掉。全局位置与局部窗口取并集时去重，否则同一个 token 会被重复计权。

多 query 的稀疏索引可不同。若为了方便共用一组块，块内每个 query 仍需自己的精确 causal mask。

选择、重归一化、读取内容，顺序不能乱

已缩放 logits 为 [0, ln 2, ln 3]，选择位置 {0, 2}，对应 V 为 [2, 8]。

所有数值为手算教学示例。

选了 3 个 token，为什么读了 12 个？

历史长度 16，块宽 4，目标位置 {1, 6, 14}（从 0 编号）。每个目标落在不同块里。

这里区分“选块后块内全部参与”与“块内仍做 token mask”；两者语义不同，读取成本可以相同。

把稀疏模式写成 kernel 能消费的索引

按 query block 存储可见 key blocks。row_ptr 指出每行在 col_idx 中的范围，避免保存巨大布尔矩阵。

row_ptr = [0, 1, 3, 5]     # 3 个 query blocks
col_idx = [0, 0, 1, 0, 2]  # 依次存每行选中的 key blocks
# query block 1: col_idx[1:3] == [0, 1]
for key_block in col_idx[row_ptr[r]:row_ptr[r+1]]:
    load(K[key_block], V[key_block])
    update_online_softmax(q_block, k_block, v_block)

query / key block 指的是逻辑位置块。分页 KV 还需要块表把逻辑位置映射到物理显存地址。

块可见，不等于块内所有元素都可见

query block 覆盖位置 4…7，key block 也覆盖 4…7。它们处于因果矩阵对角线上。

跳过块内 mask 会泄露未来。只在整个 key block 的最大位置不超过最早 query 时，才可以整块免 mask。

把复杂度换成 FLOP 与读取字节

单个 query head，n = 32768、k = 1024、d = 128、BF16。点积与 value 加权各约 2kd FLOP。

如果索引维度是 16 且每步全扫，光索引就要读 1 MiB。还未计索引计算、输出 top-k、非连续访问与多个 head 的共享。

这是一个明确的教学配置。选择器与主 attention 的精度、头共享方式可以不同。

漏选的损失比“命中几个”更有意义

假设完整 attention 在三个关键位置上的权重为 0.80、0.15、0.05。选中两个位置的 recall 都是 2/3，但漏掉谁差别很大。

attention mass 仍只是诊断指标；最终要检查模型输出与任务质量。采样部分层和 query 做 full 参考，避免评估本身太贵。

训练阶段还要决定选择器怎么学

固定窗口不需要学习索引；动态 top-k 对“哪些位置被选中”通常没有普通的连续梯度。

把预训练 full 模型临时套上稀疏 mask，是一种近似推理策略；它与从训练阶段就采用该稀疏结构不是同一件事。

稀疏后更慢，先看三段时间

假想 trace：原 attention 1.4 ms；新方案选择 0.7 ms、gather 0.6 ms、稀疏计算 0.4 ms，共 1.7 ms。

每次只改一个变量，并用相同 mask 对齐正确性。若质量下降，先查 causal、去重、位置编码和归一化，再调选择预算。

长度 8，窗口 3，省掉多少配对？

先写出假设和单位，再代入数字。

块读取能兑现多少理论节省？

先写出假设和单位，再代入数字。

未来信息从哪里漏进来了？

先写出假设和单位，再代入数字。

让 sparse gather 对上 dense masked 参考

用 NumPy 在 CPU 上跑一个可检查的参考实现。小尺寸便于把中间张量打印出来。

python3 -m pip install numpy
python3 lab.py
python3 lab.py --seed 7

masked dense vs gather: PASS
future perturbation: PASS
window pairs n=8,w=3: 21
3 targets / 12 loaded = 25.0%

把窗口改成局部 + 固定全局位置，先去重；再故意让对角块免 mask，观察未来扰动检查失败。

这是正确性与成本的教学实验；耗时不代表 GPU kernel 性能。