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

📄 arXiv: 2606.20005v1 📥 PDF

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

分类: cs.LG, cs.AI

发布日期: 2026-06-18

💡 一句话要点

提出StreamKL以解决注意力蒸馏中的内存和速度问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 注意力蒸馏 Kullback-Leibler散度 模型压缩 深度学习 GPU优化

📋 核心要点

  1. 现有的注意力蒸馏方法在计算KL散度时需要物化两个注意力分布，导致高昂的内存和IO成本，尤其在长上下文时表现不佳。
  2. 本文提出StreamKL，通过引入在线形式的KL散度计算，避免了二次物化，从而实现了高效的内存使用和计算速度。
  3. 实验结果显示，StreamKL在前向和后向传递中分别比传统方法快$43 imes$和$14 imes$，显著提升了性能和效率。

📝 摘要（中文）

注意力蒸馏通过最小化Kullback-Leibler（KL）散度来训练一个注意力分布以匹配另一个分布，广泛应用于知识蒸馏、模型压缩、持续学习和稀疏注意力大语言模型训练。然而，现有方法在计算KL散度时需要先物化两个注意力分布，导致$O(N_QN_K)$的内存和IO成本，尤其在长上下文长度时变得不可承受。本文提出StreamKL，这是第一个融合的GPU原语，消除了这种二次物化。StreamKL推导出了一种新的在线形式，允许通过片段流式处理查询-键块，从而实现单次前向内核。后向过程中，StreamKL逐块重新计算注意力概率，避免存储二次中间结果。实验表明，StreamKL在前向和后向传递中分别比基线方法快$43 imes$和$14 imes$，并将注意力蒸馏的额外HBM占用从$O(N_QN_K)$降低到$O(1)$，使得在单个GPU上进行长上下文蒸馏成为可能。

🔬 方法详解

问题定义：本文解决的是注意力蒸馏中KL散度计算的内存和速度问题。现有方法需要物化两个注意力分布，导致$O(N_QN_K)$的内存和IO成本，尤其在长上下文时变得不可承受。

核心思路：StreamKL的核心思路是通过推导出一种新的在线KL散度计算形式，避免了二次物化的需求。通过流式处理查询-键块，StreamKL实现了高效的内存使用和计算速度。

技术框架：StreamKL的整体架构包括前向和后向两个阶段。在前向阶段，使用单次内核流式处理查询-键块；在后向阶段，逐块重新计算注意力概率，避免存储中间结果。

关键创新：StreamKL的主要创新在于其融合的GPU原语设计，首次实现了注意力KL散度的高效计算，显著降低了内存占用和计算复杂度。与现有方法相比，StreamKL在内存使用和计算速度上具有本质区别。

关键设计：StreamKL在实现中采用了高效的GPU内核，并进行了专门的优化设计，确保在计算过程中最大限度地减少内存占用和提高计算速度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，StreamKL在前向传递中比基线方法快$43 imes$，在后向传递中快$14 imes$，同时将注意力蒸馏的额外HBM占用从$O(N_QN_K)$降低到$O(1)$，极大提升了效率。

🎯 应用场景

该研究的潜在应用领域包括知识蒸馏、模型压缩和大规模语言模型的训练。通过降低内存占用和提高计算速度，StreamKL能够支持更长上下文的处理，推动深度学习模型在实际应用中的发展，尤其是在资源受限的环境中。

📄 摘要（原文）

Attention distillation, which trains one attention distribution to match another by minimizing their Kullback-Leibler (KL) divergence, is widely used in knowledge distillation, model compression, continual learning, and sparse-attention LLM training. However, existing approaches materialize both attention distributions before computing the KL reduction, incurring $O(N_QN_K)$ memory and IO costs that become prohibitive at long context lengths. We present StreamKL, the first fused GPU primitive for attention KL divergence that eliminates this quadratic materialization. StreamKL derives a novel online formulation for the coupled two-distribution KL reduction, enabling a single one-pass forward kernel that streams query-key tiles through on-chip SRAM. For the backward pass, StreamKL recomputes attention probabilities tile-by-tile, avoiding storage of quadratic intermediates. We further design and implement efficient GPU kernels with dedicated optimizations. Experiments show StreamKL delivers up to $43\times$ and $14\times$ speedups over baseline methods in the forward and backward passes, respectively. Most importantly, StreamKL reduces the extra HBM footprint of attention distillation from $O(N_QN_K)$ to $O(1)$, enabling long-context distillation on a single GPU.