🤖 machine learning

StreamKL: Fast and Memory-Efficient KL Divergence for Boosting Attention Distillation

StreamKL 引入了一种新型的融合 GPU 原语，通过单次传递流式传输查询-键（query-key）分块，消除了注意力蒸馏中的二次方内存和 I/O 瓶颈，实现了显著的加速，并将内存占用从 $O(N_QN_K)$ 降低至 $O(1)$ ，从而使在单 GPU 上进行长上下文蒸馏成为可能。

原作者： Guangda Liu, Yiquan Wang, Chengwei Li, Wenhao Chen, Jing Lin, Yiwu Yao, Danning Ke, Wenchao Ding, Jieru Zhao

发布于 2026-06-19

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Guangda Liu, Yiquan Wang, Chengwei Li, Wenhao Chen, Jing Lin, Yiwu Yao, Danning Ke, Wenchao Ding, Jieru Zhao

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图教一个小巧、快速的学生（“学生模型”）像一位博学、缓慢的老师（“教师模型”）那样思考。在人工智能的世界里，他们通过比较两者的“注意力”如何分配到故事或句子的不同部分来进行学习。这个过程被称为注意力蒸馏（Attention Distillation）。

为了进行这种比较，计算机需要计算一个特定的数值，叫做 KL 散度（KL Divergence）。你可以把它看作是一个“距离得分”，它会告诉你学生的注意力与老师的注意力之间有多大的差异。目标是让这个得分尽可能小。

问题所在：“内存爆炸”

论文解释说，对于长篇故事（比如一本有 10 万词的小说）进行这种比较，目前是计算机内存的噩梦。

这里有一个类比：
想象你有两个巨大的白板，一个记录老师的注意力，一个记录学生的注意力。为了进行比较，旧的方法要求你在这些白板上写下每一对可能的单词组合。

如果有 64,000 个词，你就必须写下 64,000 × 64,000 个配对。这可是超过 40 亿个数字。
这样做需要的白板极其巨大，以至于无法放入计算机的主内存（HBM）中。这就像是试图把一座图书馆的书存进一个鞋盒里。
因为计算机无法一次性看到全貌，它不得不把故事切成碎片，处理一段，再拼凑一段。这非常慢，就像通过逐个查看字母并记录下来，然后再移动到下一个字母来读一本书一样。

解决方案：StreamKL（“流式”方法）

研究人员创造了一个名为 StreamKL 的新工具。与其先在巨大的白板上写下所有内容，StreamKL 使用了一个聪明的技巧，像传送带一样实时计算“距离得分”。

创意类比：工厂流水线
想象一个工厂，你正在比较两条传送带上的产品（老师的注意力和学生的注意力）。

旧方法： 你停止流水线，把所有产品倾倒在一个巨大的仓库地面上（HBM），测量完所有产品后再清理现场。这占用了整个仓库，而且速度很慢。
StreamKL 方法： 你让产品保持在传送带上移动。当每一对物品经过传感器（GPU 芯片）时，你会立即对它们进行比较，计算差异，并将结果放入一个微小的口袋（SRAM）中，然后等待下一对物品到达。你从不停止流水线，也不需要仓库。你只需要一个口袋。

它是如何工作的（神奇的魔术）

论文描述了这一魔术的两个主要部分：

前向传播（计算得分）：
研究人员发明了一种新的数学公式，允许计算机增量地更新“距离得分”。当它流经数据时，它只保留一些运行中的统计量（如运行最大值和总和），而不是保留整个列表。这意味着它可以处理任何长度的故事，而不会耗尽内存。
反向传播（从错误中学习）：
当计算机需要根据得分来改进学生时，它通常需要回看数据。旧的方法会保存整个巨大的数据列表以便回看。StreamKL 则更聪明：它扔掉列表，但记住了一些“秘密钥匙”（称为 LSE 值）。当它需要回看时，它会利用这些密钥重建它当前需要的特定数据片段，就在那一刻进行计算，然后再次将其丢弃。这就像是记住了制作蛋糕的食谱，以便在需要品尝时随时烤出一片，而不是把整块蛋糕都烤好后存在冰箱里。

结果：速度与空间

论文在强大的 NVIDIA GPU（H200 和 A100）上测试了这些功能，使用了极长的上下文（高达 512,000 个词）。

内存节省： StreamKL 将所需的额外内存从“二次方级”（爆炸式增长至数 TB）降低到了“常数级”（保持极小）。它从 64k 上下文需要 512 GB 内存，降到了几乎不需要额外内存。这使得单个 GPU 能够处理以前需要超级计算机才能完成、甚至根本无法完成的任务。
速度： 因为它不需要在后台频繁地读写海量数据，所以速度极快。
- 在某些测试中，它的计算得分速度比标准方法快了 43 倍。
- 在学习阶段，它快了 14 倍。

总结

StreamKL 是一种教 AI 模型如何分配注意力的新方法。它通过阻止计算机记录完整的对比列表，解决了处理长文本时“内存耗尽”的问题。相反，它让数据流经一个微小且高效的管道，实现即时计算。这使得在单台计算机上训练和运行以前过于庞大、无法处理的 AI 模型成为了可能。

技术摘要：StreamKL

问题陈述

注意力蒸馏（Attention Distillation）是一种通过最小化 KL 散度使一个注意力分布匹配另一个分布的技术，在知识蒸馏、模型剪枝、持续学习以及训练稀疏注意力大语言模型（LLM）中至关重要。然而，现有的实现方案在处理长上下文长度（ $N_Q, N_K$ ）时面临严重的扩展性瓶颈。

当前的方法在计算逐元素 KL 归约之前，会在高带宽内存（HBM）中“实例化”（materialize）两个注意力分布（ $P_1, P_2 \in \mathbb{R}^{N_Q \times N_K}$ ）。这导致了：

二次方内存占用： 存储两个 $N_Q \times N_K$ 矩阵需要 $O(N_Q N_K)$ 的内存。对于 64K token 的上下文，使用 32 个头且采用 BF16 格式，仅此一项就会消耗约 512 GB 内存，超过了单张 NVIDIA H200 GPU 141 GB 的容量。
二次方 IO 开销： 将这些分布写入 HBM 并将其读回进行 KL 归约的过程主导了运行时间，其规模随上下文长度呈二次方增长。
分块处理的低效性： 常用的折中方案（如分块处理）是用延迟换取内存；较小的块无法充分利用 GPU 并行性，而较大的块则会触发显存溢出（OOM）错误。

标准注意力计算曾面临类似问题，FlashAttention 通过在线、融合的算子解决了这一问题。然而，将此技术应用于注意力蒸馏并非易事，因为 KL 散度通过一个加权的对数差异项将两个独立的 softmax 分布耦合在一起，使得单次遍历（single-pass）的在线公式在数学上变得非常复杂。

方法论：StreamKL

StreamKL 引入了首个用于计算注意力 KL 散度的融合、单次遍历 GPU 原语，消除了 $O(N_Q N_K)$ 中间分布的实例化过程。

1. 在线公式（前向传递）

StreamKL 推导了一种全新的用于耦合双分布 KL 归约的在线更新方案。通过重构行向 KL 散度，作者将其表示为五个每行运行标量：

$m_1, m_2$ ：两个对数分布的运行最大值。
$l_1, l_2$ ：未归一化的指数和（类似于在线 softmax）。
$acc$：用于加权对数差异的累加器。

重构后的目标函数为：
$L = \frac{acc}{l_1} + \text{LSE}_2 - \text{LSE}_1$
其中 $\text{LSE}_t = m_t + \log l_t$ 。

该公式允许内核将查询-键（query-key）分块流式传输到片上 SRAM 中，并增量更新这五个标量。至关重要的是，累加器 $acc $是大小为$ B_Q$（块大小）的向量，而非完整的矩阵，这降低了相比于标准 FlashAttention 更高的算术缩放开销。该算法通过跳过因果边界之外的键分块来处理因果掩码（causal masking）。

2. 重计算策略（反向传递）

为了避免在反向传播期间存储二次方的 $P_1$ 和 $P_2$ 矩阵，StreamKL 采用了重计算策略：

保存状态： 仅将最终的 KL 值和 Log-Sum-Exp (LSE) 向量（ $\text{LSE}_1, \text{LSE}_2$ ）存储在 HBM 中（ $O(N_Q)$ ）。
再生： 在反向传递期间，利用保存的 LSE 和原始的 $Q, K$ 矩阵重新生成注意力概率（ $P_i = e^{S_i - \text{LSE}}$ ）。
梯度计算： 内核针对两种设置计算梯度：
- 设置 1： $P_1$ 固定，优化 $P_2$ （标准蒸馏）。
- 设置 2： $P_2$ 固定，优化 $P_1$ 。
内核融合： 作者设计了一个融合的反向内核，该内核仅需访问一次每个 $QK$ 分块对，从而减少了 HBM 流量。为了管理累积 $dQ $（沿$ K $方向）与$ dK $（沿$ Q $方向）之间的冲突，内核在缓冲区能够放入 L2 缓存时使用原子操作处理$ dQ$，在竞争过于激烈时则退回到独立的内核。

3. GPU 内核优化

Split-K 变体： 对于查询数量较少的负载（例如自回归解码，其中 $N_Q=1$ ），Split-K 变体通过将键维度划分到不同的线程块（thread blocks）来饱和流式多处理器（SMs）。
TMA 加速： 在 Hopper GPU 上，实现利用张量内存加速器（TMA）进行异步批量内存拷贝，以隐藏内存延迟。
数值稳定性： 设置 2 中的对数比值计算是使用对数差异和保存的 LSE 进行的，而不是直接使用概率对数，以避免数值不稳定（ $-\infty$ 或 NaN）。

核心贡献

识别瓶颈： 本文识别了阻碍注意力 KL 散度扩展到长上下文的二次方内存和 IO 瓶颈。
在线重构： 引入了首个注意力 KL 散度的在线重构方案，实现了完全融合的前向原语和基于重计算的反向传递，消除了二次方中间变量。
高效实现： 作者为前向和反向传递设计并实现了高效的 GPU 内核，并包含专门的优化（Split-K、TMA、自适应融合）。
可扩展性： StreamKL 将注意力蒸馏的额外 HBM 足迹从 $O(N_Q N_K)$ 降低到 $O(1)$ ，使得在单张 GPU 上进行长上下文蒸馏成为可能。

实验结果

在 NVIDIA A100 和 H200 GPU 上，针对 4K 到 512K 的上下文长度进行了评估：

内存效率： StreamKL 保持了平坦的内存足迹，不受上下文长度影响。相比之下，基准测试（PyTorch, torch.compile, FLA）呈二次方增长，并在 16K–64K 时触发 OOM 错误。在 64K 长度下，StreamKL 使用的 HBM 比 torch.compile/FLA 少 16,000 倍以上，比 PyTorch 少 135,000 倍以上。
前向延迟：
- 在非因果设置下，StreamKL 比 PyTorch 快 13–18 倍，比 torch.compile/FLA 快 2.4–3.7 倍。
- 在因果掩码下，由于 StreamKL 完全跳过了被掩码的计算，而基准测试仍会实例化完整矩阵，其加速效果扩大至比 PyTorch 快 26–43 倍，比基准测试快 4.2–7.1 倍。
反向延迟：
- 在设置 1 中，StreamKL 在非因果设置下比 PyTorch 快 2.6–4.0 倍，在因果设置下快 5.3–9.9 倍。
- 在设置 2 中，StreamKL 在非因果设置下比 PyTorch 快 3.8–6.5 倍，在因果设置下快 9.1–14.0 倍。
- 与 torch.compile 相比，StreamKL 具有竞争力或更快，特别是在因果设置下，通过跳过掩码分块，重计算成本减半。
消融实验： Split-K 变体为小 $N_Q$ 负载提供了高达 8.4 倍 的加速。融合反向内核在 $N_Q$ 较小时比独立内核快达 2.0 倍，证明了自适应选择策略的合理性。

重要意义

本文声称 StreamKL 是第一个能够在单张 GPU 上支持 64K 及以上上下文长度的注意力蒸馏工作负载的解决方案。通过消除注意力分布的二次方实例化，StreamKL 消除了阻碍现代 LLM 训练（例如 DeepSeek V3.2 和 GLM-5 等稀疏注意力模型）中注意力蒸馏扩展的主要系统瓶颈。它证明了融合的在线原语可以成功地从标准注意力扩展到更复杂的耦合分布（如 KL 散度）。