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

Research Problem

Attention distillation is a critical technique for training modern transformer models, including knowledge distillation, model pruning, and training sparse-attention LLMs. The process involves minimizing the Kullback-Leibler (KL) divergence between two attention distributions. Traditionally, this requires materializing two quadratic attention matrices in high-bandwidth memory (HBM). As context lengths scale to 128K tokens or beyond, this O(N_Q*N_K) memory footprint becomes prohibitive, often exceeding the capacity of modern GPUs and creating significant I/O bottlenecks.

The StreamKL Approach

StreamKL introduces a novel online formulation for the coupled two-distribution KL reduction. By deriving a mathematical framework that allows the KL divergence to be computed incrementally, the authors enable a fully fused, one-pass forward kernel. This kernel streams query-key tiles through on-chip SRAM, maintaining running statistics—such as row-wise maxima and unnormalized sums—to compute the final KL divergence without ever storing the full attention matrices in HBM. For the backward pass, StreamKL utilizes a recomputation strategy, regenerating attention probabilities tile-by-tile from saved log-sum-exp (LSE) values to avoid storing quadratic intermediates.

Performance and Impact

By eliminating the quadratic materialization of attention distributions, StreamKL significantly reduces memory usage and I/O traffic. Experimental results on NVIDIA H200 and A100 GPUs demonstrate that StreamKL achieves up to 43x speedup in the forward pass and 14x in the backward pass compared to baseline methods. Most importantly, it reduces the extra HBM footprint from O(N_Q*N_K) to O(1), making long-context attention distillation feasible on a single GPU where it was previously impossible due to out-of-memory (OOM) errors.