LiteTopK：用于长上下文稀疏注意力的融合索引与Top-k算子

Indexer-TopK 是计算分数并选择前 k 个候选者的操作，广泛应用于 大语言模型 中的稀疏注意算法以及推荐系统和向量数据库中的向量检索。然而，由于全局内存流量过多、同步成本高昂以及内存开销过高，现有基于 GPU 的 Indexer-TopK 内核（例如 DeepSeek Sparse Attention (DSA)）仍然效率低下。在这项研究中，受到维数诅咒现象的启发，我们首先观察到稀疏的注意力分数表现出分数集中的现象，分数往往落在一个狭窄的范围内。基于这一观察，我们提出了 LITETOPK，一种高效的融合 Indexer-TopK 内核。 LITETOPK 首先对一小部分数据进行采样以估计查询数据分数范围，然后相应地将候选数据划分到 bin 中。这种组织允许 LITETOPK 内核在线维持严格的近似阈值，仅写回有希望的候选者，减少不必要的 I/O 和内存开销，同时保持精确的 Top-k 正确性。在 LITETOPK 的基础上，我们进一步提出 LITEDSA，它利用了相邻 token 之间 top-k 候选集的相似性。 LITEDSA 打包相邻词元的候选者以进行联合计算，并屏蔽每个查询的额外分数，从而减少内存流量，同时保持正确性。在具有 8 个 B200 GPU 的实际部署环境中的实验结果表明，LITETOPK+LITEDSA 将 GLM 5.2 的预填充阶段加速了 1.35 倍，且没有性能损失且内存开销更低。