PISA 用金字塔 Top-K 选择把块稀疏注意力的选择阶段降到 O(N log N)，在 256K 长度上比 BSA 快 9.95 倍

PISA 把块稀疏注意力的 Top-K 块选择从“对每个查询扫描全部键块”改为“在键块金字塔上由粗到细逐层收窄候选”。 传统块稀疏注意力（如 BSA）需要为每个查询给全部 N/C 个键块打分，选择阶段整体仍是 O(N²/C)；PISA 通过均值池化构造 O(log N) 层键表示，每层只对不超过 gK 个候选块用 LogSumExp 打分并保留 Top-K，再展开到下一层，使每查询路由复杂度为 O(log N)、整序列为 O(N log N)。 论文给出完整复杂度推导：金字塔层数 L=O(log N)，每层候选数受 gK 约束，并给出训练/预填充与解码两种核设计的 IO 成本对比（GQ=16、Qtile=4、C=64 时 GQ+C/Qtile=32<64=C）。

作者为训练/预填充和解码分别实现了硬件感知的 Triton 核，融合分层路由与 LogSumExp 打分，不物化查询-键打分矩阵。 训练/预填充采用两阶段核：第一阶段并行处理中间层选择并按 GQA 组对查询头分数求和，第二阶段把需要同一候选键块的查询聚成簇、以 Qtile=4 分块复用键块；解码用单阶段核，缓存均值金字塔并只更新新 token 的叶均值及其祖先路径。 论文给出两阶段与单阶段在叶层 Q/K IO 成本上的解析比较，并说明在 GQ=16、Qtile=4、C=64 的配置下两阶段 IO 更低，而解码时每簇只有一个查询、无跨查询复用，故用单阶段。

在 418M、1.47B、2.67B 三个规模上，PISA 的语言建模与常识推理表现与 BSA、NSA、HiLS 相当，在检索类任务上更好。 论文报告在 100B token、4K 序列预训练后，PISA 在六项 containment 任务（SWDE、SQuAD、FDA、TQA、NQ、DROP）上取得稀疏方法中最高平均准确率（418M 为 41.77、1.47B 为 49.83、2.67B 为 52.14），但全注意力平均仍更高；在 16K 继续预训练后的 RULER 检索上，2.67B 模型 PISA 平均 62.80，高于 BSA 的 54.99 与 NSA 的 61.07，低于全注意力的 66.24。 三个规模使用相同解码器骨干与匹配训练设置，稀疏方法统一 C=64、K=8（继续预训练 K=32），并报告训练损失、困惑度、多项选择与 containment 准确率；继续预训练后 PISA 的平均 containment 准确率低于 BSA。

块选择效率随序列增长的优势明显：PISA 的 Q4 实现在 64K、128K、256K 上分别比 BSA 快 2.86 倍、5.31 倍、9.95 倍。 在 4K 到 16K 区间 BSA 更快，但从 32K 起 Q4 延迟最低；Q4 相对逐查询实现也在 64K、128K、256K 上分别有 1.35 倍、1.30 倍、1.33 倍加速。 延迟基准使用随机 BF16 查询与键、固定种子、batch size 1、32 个查询头、2 个 KV 头、d=64、C=64、K=8、g=2，计时包含摘要构造与全部选择阶段但不含稀疏注意力本身。

这项工作面向需要长上下文语言建模与检索的训练与推理场景，适用于采用解码器骨干、GQA 分组和固定块大小/预算的模型；它给出的是选择阶段的复杂度与延迟改进，以及在此设置下与稀疏基线的可比性能。对希望降低长序列预填充与解码选择开销的工程与研究团队，论文提供了可复用的金字塔选择思路、LogSumExp 打分形式与 Triton 核设计（两阶段用于训练/预填充、单阶段用于解码），并给出块选择质量诊断方法（Recall@K、捕获质量、注意力质量比）供后续比较。

论文自述计算资源限制了所探索的模型规模与预训练数据预算，并指出这两者都可能显著影响性能与相对增益，因此更大规模与更长训练下的表现仍是开放问题。继续预训练后 PISA 的平均 containment 准确率低于 BSA，而全注意力在多项选择与 containment 平均上仍更高，说明稀疏化带来的取舍在不同阶段并不一致。块选择质量诊断显示 PISA 的 Recall@8 为 90.95、注意力质量比 99.46，但该诊断在固定全注意力权重上进行、参考集与 PISA 的分数聚合方式不同，二者可能给出不同排序。此外，本次读取的是论文全文文本，图与部分表格以文字形式呈现，若需核对具体曲线数值仍应回到原文图表。