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

arXiv: https://arxiv.org/abs/2606.20005 · HTML: https://mlsanity.com/paper/2606.20005/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in General knowledge distillation methods, one of 12 directions in Knowledge distillation.
- Builds on 12 of its 18 references that are in the map.
- First posted on arXiv 18 Jun 2026.

## Placed in

- Knowledge distillation / [General knowledge distillation methods](https://mlsanity.com/distill/general-knowledge-distillation-methods/)

## Builds on (most cited first, 12 of 12)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [MiniLM: Deep Self-Attention Distillation for Task-Agnostic Compression of Pre-Trained Transformers](https://mlsanity.com/paper/2002.10957/) (2020)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Diffuse, Attend, and Segment: Unsupervised Zero-Shot Segmentation using Stable Diffusion](https://mlsanity.com/paper/2308.12469/) (2023)
- [MimiQ: Low-Bit Data-Free Quantization of Vision Transformers with Encouraging Inter-Head Attention Similarity](https://mlsanity.com/paper/2407.20021/) (2024)
- [TRACE: Your Diffusion Model is Secretly an Instance Edge Detector](https://mlsanity.com/paper/2503.07982/) (2025)
- [ConceptSplit: Decoupled Multi-Concept Personalization of Diffusion Models via Token-wise Adaptation and Attention Disentanglement](https://mlsanity.com/paper/2510.04668/) (2025)
- [D2Prune: Sparsifying Large Language Models via Dual Taylor Expansion and Attention Distribution Awareness](https://mlsanity.com/paper/2601.09176/) (2026)

## Abstract (arXiv)

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(NQNK) 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× and 14× speedups over baseline methods in the forward and backward passes, respectively. Most importantly, StreamKL reduces the extra HBM footprint of attention distillation from O(NQNK) to O(1), enabling long-context distillation on a single GPU.