当现代大语言模型（LLM）将上下文长度推向 128K、256K 甚至百万 token 级别时，注意力矩阵的存储会使一种训练操作成本迅速上升：注意力蒸馏（Attention Distillation）。它用于知识蒸馏、模型压缩、持续学习和稀疏注意力训练。朴素实现需要将两个 O(N²) 的注意力矩阵完整写入 HBM，再逐元素计算 KL 散度——在 64K 上下文、32 注意力头、BF16 精度下，仅存储两个注意力分布就需要 512 GB，是单张 NVIDIA H200（141 GB HBM）的 3.6 倍。
上海交通大学、华为和复旦大学联合团队提出的 StreamKL，将注意力 KL 散度计算融合为 GPU 算子，按块读取数据、在线累计结果，无需在 HBM 中保存两个完整注意力矩阵。论文在算子测试中报告，前向相对基线最高加速 43 倍，反向最高加速 14 倍，并使部分长上下文配置能够在单 GPU 上运行。
注意力蒸馏用于对齐教师、学生与稀疏索引器
注意力蒸馏的做法是：训练一个注意力分布去匹配另一个，损失函数是它们之间的 KL 散度。常见应用包括：
- 知识蒸馏：MiniLM 等工作通过蒸馏教师模型与学生的注意力分布来迁移知识；
- 模型压缩：剪枝或量化后的模型通过注意力蒸馏保持原始模型的注意力模式；
- 持续学习：通过将当前模型的注意力分布与历史快照对齐，缓解灾难性遗忘；
- 稀疏注意力训练：DeepSeek V3.2 和 GLM-5 等模型在训练稀疏注意力时，使用注意力蒸馏将轻量级 indexer 的分布与原始密集注意力分布对齐。
这些场景的共同点是：训练目标都是最小化两个 N_Q × N_K 的注意力矩阵 P₁ 和 P₂ 之间的 KL 散度。当 N_Q = N_K = 64K 时，这就是一个 64K × 64K 的矩阵，且有两个。
双分布 KL 计算需要处理两套在线归一化统计
FlashAttention 已通过融合 kernel 避免保存标准注意力的完整矩阵。注意力蒸馏可以借鉴分块和在线更新的思路，但计算形式不同，不能直接沿用原 kernel。
标准注意力计算只需要处理一个 softmax 分布，然后乘以 V 矩阵。而 KL 散度耦合了两个独立的 softmax 分布，通过一个加权的 logit 差值项连接。如果直接使用在线 softmax，要么仍然需要物化 O(N²) 的中间结果，要么需要两次独立的遍历——计算量和 HBM 读取量翻倍。
这里的难点是：当两个分布的 running maximum 各自变化时，如何正确地对累积的加权 logit 差值进行 rescale？StreamKL 为此重新推导了累积量的更新公式。
在线累积 KL，并在反向传播时重算概率
在线公式：耦合双分布 KL 的增量计算
StreamKL 将注意力 KL 散度改写为可逐块更新的形式。对于行级 KL 散度 L = Σᵢ P₁ⁱ log(P₁ⁱ/P₂ⁱ)，将 softmax 概率用在线统计量（running maximum m 和未归一化和 l）表示后，经过推导可以得到：
L = acc / l₁ + LSE₂ − LSE₁
其中 acc = Σᵢ e^(S₁ⁱ − m₁) · (S₁ⁱ − S₂ⁱ) 是一个可以在流式处理中增量更新的累加器。关键的是，这个累加器和标准化常数 l₁, l₂ 都可以随着 key tile 的流入，使用类似在线 softmax 的 rescale 机制逐 tile 更新——且累加器的 rescale 只是一个 B_Q 维向量，而非 FlashAttention 中需要 rescale 的 B_Q × d 输出矩阵，仅这一重缩放操作的算术开销降低了 d 倍（通常 ≥128），并非整个算子的计算量减少同样倍数。
前向：单遍流式 Kernel
StreamKL 的前向 kernel 将 query 分块，每个 thread block 处理一个 query tile，在内部循环中依次流式加载 K₁ 和 K₂ 的 tile。每处理一个 tile：
- 在 SRAM 中计算 logit tile S₁ = Q₁K₁ᵀ 和 S₂ = Q₂K₂ᵀ
- 更新 running maximum m₁, m₂
- 计算 correction factor 进行 rescale
- 更新累加器 acc
前向结束后，仅将 KL 值和两个 LSE 向量（O(N_Q) 个标量）写回 HBM。论文将避免完整注意力矩阵后的额外工作区概括为 O(1)；若计入这些逐行统计量、输入输出和反向梯度，整体显存仍会随规模增长。
反向：基于重计算的 Tile-by-Tile 梯度
反向传播不存储任何 O(N²) 中间结果。通过保存的 LSE₁ 和 LSE₂，backward kernel 可以随时按需重建 P₁ⁱ = exp(S₁ⁱ − LSE₁) 和 P₂ⁱ = exp(S₂ⁱ − LSE₂)。StreamKL 支持两种梯度设置：优化 P₂（标准蒸馏）和优化 P₁，并对后者设计了数值稳定的 log-ratio 计算，避免概率趋近于零时产生 NaN。
GPU Kernel 工程优化
实现上，StreamKL 还针对 GPU 的并行度和访存作了优化：
- Split-K 变体：当 batch size 小或 N_Q 少导致 SM 利用率不足时，沿 key 维度分片，将部分统计量写回 HBM 后由轻量级 reduce kernel 合并，提高 GPU 利用率；
- 融合反向 Kernel：通过原子操作将 dQ 和 dK 的计算合并到一个 kernel 中，每个 QK tile 对只访问一次，减少相对分离 dQ/dK kernel 的重复访问；
- Hopper 架构适配：利用 Tensor Memory Accelerator（TMA）进行异步 bulk copy，并用 exp₂ 替代自然指数以利用单周期 SFU 指令。
H200 与 A100 算子测试：速度和显存收益
论文在 NVIDIA H200 和 A100 上比较 PyTorch、torch.compile 和 FLA 实现，报告以下最高加速比。长上下文下部分基线因显存不足无法运行，相关数据按已完成运行外推，并在原图中标记：
| 指标 | 前向 | 反向 |
|---|---|---|
| 对比 PyTorch（causal） | 43× 加速 | 14× 加速 |
| 对比 PyTorch（non-causal） | 18× 加速 | 6.5× 加速 |
| 对比 torch.compile（causal） | 7.0× 加速 | — |
| 对比 FLA（causal） | 7.1× 加速 | — |
| 完整注意力矩阵存储 | 无需保存 | 按块重算，无需保存 |
显存收益在长上下文下更突出。论文的受测实现中，只有 StreamKL 能持续运行 64K 及更长上下文，测试范围最高到 512K。另一个分块基线实验使用 H200、batch 32 和 128K 上下文，query chunk 为 1024 时反向计算即触发 OOM；这并非所有分块实现的通用长度上限。
改变计算顺序，减少完整矩阵的存储与读写
StreamKL 展示了公式推导与 GPU 实现如何配合：先把损失改写成可在线更新的形式，再让分块数据在片上 SRAM 中完成计算，减少向 HBM 写入和读回中间矩阵。
这一思路与 FlashAttention 相近，都通过在线计算避免存储完整注意力矩阵。StreamKL 将它应用于双分布 KL 散度，但不会把所有输入读取或完整模型训练的复杂度变成常数。
对于稀疏注意力训练、知识蒸馏和长上下文模型压缩，StreamKL 提供了一个可替换注意力 KL 计算的实现。它对完整训练吞吐的影响，还取决于该算子在总耗时中的占比及其他显存开销。论文已于 2026 年 6 月 18 日发布在 arXiv 上。
