StreamKL：快速且内存高效的 KL 散度，用于提高注意力 蒸馏

注意力 蒸馏 通过最小化 Kullback-Leibler (KL) 散度来训练一种注意力分布以匹配另一种注意力分布，广泛应用于 知识蒸馏、模型压缩、持续学习和稀疏注意力 LLM 训练。然而，现有方法在计算 KL 缩减之前实现了两种注意力分布，从而导致 $O(N_QN_K)$ 内存和 IO 成本，这在长上下文长度下变得令人望而却步。我们提出了 StreamKL，这是第一个用于注意力 KL 散度的融合 GPU 原语，消除了这种二次实体化。 StreamKL 为耦合的两个分布 KL 缩减导出了一种新颖的在线公式，支持通过片上 SRAM 流式传输查询密钥图块的单个一次性前向内核。对于后向传递，StreamKL 逐块重新计算注意力概率，避免存储二次中间值。我们通过专门的优化进一步设计和实现高效的 GPU 内核。实验表明，与基线方法相比，StreamKL 在前向和后向传递中分别提供高达 43\times$ 和 $14\times$ 的加速。最重要的是，StreamKL 将注意力 蒸馏 的额外 HBM 占用空间从 $O(N_QN_K)$ 减少到 $O(1)$，从而在单个 GPU 上实现长上下文 蒸馏。