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

2026-06-18

AI Summary

This work addresses the prohibitive memory and I/O overhead of traditional attention distillation, which explicitly stores two attention distributions and incurs O(N_QN_K) space complexity, hindering scalability to long contexts. The authors propose StreamKL, a novel approach that leverages an online dual-distribution KL divergence reduction formula to enable single-pass, streaming forward computation. During backpropagation, attention probabilities are recomputed on-demand in blocks, eliminating the need to store quadratic intermediate results. StreamKL is the first method to compute attention-based KL divergence without additional O(N_QN_K) GPU memory, reducing HBM usage to O(1). Combined with GPU streaming tiling, efficient on-chip SRAM utilization, and custom kernel optimizations, it achieves up to 43× and 14× speedups in forward and backward passes, respectively, drastically lowering memory consumption and enabling long-context attention distillation on a single GPU.

Abstract

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.

Problem

• attention distillation
• KL divergence
• memory efficiency
• long-context
• GPU memory

Innovation

• StreamKL
• attention distillation
• KL divergence
• memory efficiency
• GPU kernel