Sparse-K 如何在 llama.cpp 中削减数百万次 Attention 计算
Source: Dev.to
什么是注意力（Attention）？
在深入了解 Sparse‑K 之前，先弄清楚注意力机制到底是干什么的，以及它为何如此耗算力是很重要的。
当语言模型收到一句话，例如：
“The cat sat on the gray carpet,”
它并不是仅仅读取这些词——它需要理解句子的含义、词语之间的关系，以及每个词如何与其他词相连。为此，每个词都会评估它与句子中所有其他词的关联强度。这就形成了一张密集的连接网，每个 token 都要对所有其他 token 计算一次得分。
公式：
(N) 个 token → (N \times N) 次注意力得分计算。
这种二次增长正是注意力成本高昂的根本原因，尤其在处理长序列时。
为什么这会成为性能瓶颈？
因为计算量随序列长度呈二次方增长——(O(N^{2}))。
一个 2048 token 的序列在单个注意力层中就需要超过 400 万次得分计算。
深入观察这些矩阵会发现，大多数连接对最终预测并没有实质性贡献；很多连接要么很弱，要么根本不相关。直观上看，与其让每个 token 与所有其他 token 形成密集的网络，我们往往只需要少数几个有意义的连接。
那 Sparse‑K 实际上做了什么？
Sparse‑K 引入了一个简单却强大的思路：不再评估 token 之间的所有成对关系，而是仅保留每个 token 前 K（top‑K）个最有意义的连接，其余全部忽略。
- 注意力机制的结构保持不变；仅仅是处理的数据量被削减。
- 这大幅降低了计算次数，减轻了内存压力，并加速了整个推理流程。
从注意力得分到稀疏掩码
在标准注意力中，计算相似度得分（(Q \times K)）后会得到一个完整的矩阵。Sparse‑K 在此基础上增加了一个步骤：
- Top‑K 选择 – 对每个 token，仅保留最高的 K 个得分。
- 掩码（Masking） – 将其余所有值通过掩码“移除”。
- 负无穷 – 被掩码的值被设为一个极大的负数（(-\infty)），使其在 Softmax 时不产生影响。
于是，Softmax 以及随后的所有计算只会在真正相关的连接上进行。
该图展示了每个 token 与所有其他 token 之间的注意力得分计算。对每个 token，最高的 K 个得分被高亮显示。
此图展示了在完成 Top‑K 选择后得到的最终 Sparse‑K 注意力掩码。
为什么掩码是合适的解决方案？
掩码本身已经是注意力机制的核心组成部分。每个 Transformer 模型都会使用多种掩码，例如：
- 因果掩码（Causal masks） – 防止模型关注未来的 token。
- 序列分离掩码（Sequence separation masks） – 区分不同的输入片段。
Sparse‑K 并未引入全新的机制；它自然地融入了已有的掩码系统：
- 构建一个额外的掩码来表示 Top‑K 选择。
- 将其与已有的基础掩码合并。
- 在已有的注意力例程中使用该合并后的掩码。
这种做法：
- 能够无缝嵌入当前流水线，不破坏兼容性。
- 保持计算图结构不变。
- 允许复用高度优化的注意力实现（例如 Flash Attention），因为这些实现已经接受掩码作为接口的一部分。
因此，使用掩码使得 Sparse‑K 能够利用现有的优化，同时保持高效、兼容和易维护。
在 llama.cpp 中的集成方式
Sparse‑K 直接插入到注意力块的自然位置：
- 计算相似度得分 (Q \times K)。
- 构建 Sparse‑K 掩码（Top‑K 选择）。
- 在后续注意力阶段之前应用合并后的掩码。
模型的其他部分——嵌入层、MLP 计算以及多头组合——保持不变。这样，Sparse‑K 针对模型中最耗算力的部分进行优化，而不影响整体架构。