🤖 machine learning

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

StreamKL introduces a novel fused GPU primitive that eliminates the quadratic memory and I/O bottlenecks of attention distillation by streaming query-key tiles in a single pass, achieving significant speedups and reducing memory footprint from $O(N_QN_K)$ to $O(1)$ to enable long-context distillation on a single GPU.

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

Published 2026-06-19

📖 5 min read🧠 Deep dive

CC BY 4.0

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

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a small, fast student (a "student model") to think exactly like a brilliant, slow teacher (a "teacher model"). In the world of AI, they do this by comparing how they both "pay attention" to different parts of a story or sentence. This process is called Attention Distillation.

To make this comparison, the computer calculates a specific number called KL Divergence. Think of this as a "distance score" that tells you how different the student's attention is from the teacher's. The goal is to make this score as small as possible.

The Problem: The "Memory Explosion"

The paper explains that doing this comparison for long stories (like a novel with 100,000 words) is currently a nightmare for computer memory.

Here is the analogy:
Imagine you have two giant whiteboards, one for the teacher's attention and one for the student's. To compare them, the old method requires you to write down every single possible pairing of words on these boards.

If you have 64,000 words, you have to write down 64,000 × 64,000 pairs. That's over 4 billion numbers.
Doing this requires a whiteboard so huge it doesn't fit in the computer's main memory (HBM). It's like trying to store a library of books in a shoebox.
Because the computer can't fit the whole picture at once, it has to chop the story into tiny pieces, process them, and then put them back together. This is slow, like trying to read a book by looking at one letter at a time and writing it down before moving to the next.

The Solution: StreamKL (The "Streaming" Approach)

The authors created a new tool called StreamKL. Instead of writing everything down on a giant whiteboard first, StreamKL uses a clever trick to calculate the "distance score" on the fly, like a conveyor belt.

The Creative Analogy: The Factory Assembly Line
Imagine a factory where you are comparing two conveyor belts of products (the teacher's attention and the student's attention).

The Old Way: You stop the line, dump every single product onto a massive warehouse floor (HBM), measure them all, and then clean up. This takes up the whole warehouse and is slow.
The StreamKL Way: You keep the products moving on the conveyor belt. As each pair of items passes a sensor (the GPU chip), you instantly compare them, calculate the difference, and throw the result into a tiny pocket (SRAM) before the next pair arrives. You never stop the line, and you never need a warehouse. You only need a pocket.

How It Works (The Magic Trick)

The paper describes two main parts of this magic:

The Forward Pass (Calculating the Score):
The researchers invented a new mathematical formula that allows the computer to update the "distance score" incrementally. As it streams through the data, it keeps a running tally of just a few numbers (like a running maximum and a sum) instead of the whole list. This means it can handle stories of any length without running out of memory.
The Backward Pass (Learning from Mistakes):
When the computer needs to learn from the score to improve the student, it usually needs to look back at the data. The old way saves the whole giant list of data to look back at. StreamKL is smarter: it throws away the list but remembers a few "secret keys" (called LSE values). When it needs to look back, it uses these keys to rebuild the specific piece of data it needs, right then and there, calculates the lesson, and then forgets it again. It's like remembering the recipe for a cake so you can bake a slice whenever you need to taste it, instead of baking the whole cake and storing it in the fridge.

The Results: Speed and Space

The paper tested this on powerful NVIDIA GPUs (H200 and A100) with very long contexts (up to 512,000 words).

Memory Savings: StreamKL reduced the extra memory needed from "quadratic" (exploding to terabytes) to "constant" (staying tiny). It went from needing 512 GB of memory for a 64k context to needing almost nothing extra. This allows a single GPU to handle tasks that previously required a supercomputer or were impossible.
Speed: Because it doesn't have to write and read massive amounts of data back and forth, it is incredibly fast.
- In some tests, it was 43 times faster than the standard method for calculating the score.
- In the learning phase, it was 14 times faster.

Summary

StreamKL is a new way to teach AI models to pay attention. It solves the problem of "running out of memory" when dealing with long texts by stopping the computer from writing down the entire comparison list. Instead, it streams the data through a tiny, efficient pipeline, calculating the result instantly. This makes it possible to train and run AI models on single computers that were previously too large to handle.

Technical Summary: StreamKL

Problem Statement

Attention distillation, a technique used to train one attention distribution to match another by minimizing Kullback-Leibler (KL) divergence, is critical for knowledge distillation, model pruning, continual learning, and training sparse-attention Large Language Models (LLMs). However, existing implementations face severe scalability bottlenecks when handling long context lengths ( $N_Q, N_K$ ).

Current approaches "materialize" both attention distributions ( $P_1, P_2 \in \mathbb{R}^{N_Q \times N_K}$ ) in High Bandwidth Memory (HBM) before computing the element-wise KL reduction. This results in:

Quadratic Memory Footprint: Storing two $N_Q \times N_K$ matrices requires $O(N_Q N_K)$ memory. For a context of 64K tokens with 32 heads in BF16, this alone consumes ~512 GB, exceeding the 141 GB capacity of a single NVIDIA H200 GPU.
Quadratic IO Costs: Writing these distributions to HBM and reading them back for the KL reduction dominates runtime, scaling quadratically with context length.
Inefficiency of Chunking: Common workarounds like chunked processing trade latency for memory; small chunks underutilize GPU parallelism, while large chunks trigger Out-of-Memory (OOM) errors.

Standard attention computation faced similar issues, which FlashAttention solved via online, fused kernels. However, applying this to attention distillation is non-trivial because KL divergence couples two independent softmax distributions through a weighted logit-difference term, making a single-pass online formulation mathematically complex.

Methodology: StreamKL

StreamKL introduces the first fused, one-pass GPU primitive for attention KL divergence that eliminates the $O(N_Q N_K)$ materialization of intermediate distributions.

1. Online Formulation (Forward Pass)

StreamKL derives a novel online-update scheme for the coupled two-distribution KL reduction. By reformulating the row-wise KL divergence, the authors express the objective in terms of five per-row running scalars:

$m_1, m_2$ : Running maximums for the two logit distributions.
$l_1, l_2$ : Unnormalized sums of exponentials (similar to online softmax).
$acc$: An accumulator for the weighted logit difference.

The reformulated objective is:
$L = \frac{acc}{l_1} + \text{LSE}_2 - \text{LSE}_1$
where $\text{LSE}_t = m_t + \log l_t$ .

This formulation allows the kernel to stream query-key tiles through on-chip SRAM, updating these five scalars incrementally. Crucially, the accumulator $acc$ is a vector of size $B_Q$ (block size), not a full matrix, reducing the arithmetic rescaling overhead compared to standard FlashAttention. The algorithm handles causal masking by skipping key tiles entirely beyond the causal frontier.

2. Recomputation Strategy (Backward Pass)

To avoid storing the quadratic $P_1$ and $P_2$ matrices during backpropagation, StreamKL employs a recomputation strategy:

State Saved: Only the final KL values and the Log-Sum-Exp (LSE) vectors ( $\text{LSE}_1, \text{LSE}_2$ ) are stored in HBM ( $O(N_Q)$ ).
Regeneration: During the backward pass, attention probabilities are regenerated tile-by-tile using the saved LSEs and the original $Q, K$ matrices ( $P_i = e^{S_i - \text{LSE}}$ ).
Gradient Computation: The kernel computes gradients for two settings:
- Setting 1: $P_1$ fixed, optimizing $P_2$ (standard distillation).
- Setting 2: $P_2$ fixed, optimizing $P_1$ .
Kernel Fusion: The authors design a fused backward kernel that visits each $QK$ tile pair exactly once, reducing HBM traffic. To manage the conflict between accumulating $dQ$ (along $K$ ) and $dK$ (along $Q$ ), the kernel uses atomic operations for $dQ$ when the buffer fits in L2 cache, falling back to separate kernels when contention becomes too high.

3. GPU Kernel Optimizations

Split-K Variant: For workloads with small query counts (e.g., autoregressive decoding where $N_Q=1$ ), a split-K variant partitions the key dimension across thread blocks to saturate Streaming Multiprocessors (SMs).
TMA Acceleration: On Hopper GPUs, the implementation leverages the Tensor Memory Accelerator (TMA) for asynchronous bulk memory copies, hiding memory latency.
Numerical Stability: The log-ratio computation in Setting 2 is performed using logit differences and saved LSEs rather than direct probability logs to avoid numerical instability ( $-\infty$ or NaN).

Key Contributions

Identification of Bottlenecks: The paper identifies the quadratic memory and IO bottlenecks that prevent attention KL divergence from scaling to long contexts.
Online Reformulation: It introduces the first online reformulation of attention KL divergence, enabling a fully fused forward primitive and a recomputation-based backward pass that eliminates quadratic intermediates.
Efficient Implementation: The authors design and implement efficient GPU kernels with dedicated optimizations (Split-K, TMA, adaptive fusion) for both forward and backward passes.
Scalability: StreamKL reduces the extra HBM footprint of attention distillation from $O(N_Q N_K)$ to $O(1)$ , enabling long-context distillation on a single GPU.

Experimental Results

Evaluated on NVIDIA A100 and H200 GPUs across context lengths from 4K to 512K:

Memory Efficiency: StreamKL maintains a flat memory footprint regardless of context length. In contrast, baselines (PyTorch, torch.compile, FLA) scale quadratically and hit OOM errors at 16K–64K. At 64K, StreamKL uses over 16,000× less HBM than torch.compile/FLA and over 135,000× less than PyTorch.
Forward Latency:
- StreamKL achieves 13–18× speedup over PyTorch and 2.4–3.7× over torch.compile/FLA in non-causal settings.
- Under causal masking, speedups widen to 26–43× over PyTorch and 4.2–7.1× over baselines, as StreamKL skips masked computations entirely while baselines still materialize full matrices.
Backward Latency:
- In Setting 1, StreamKL is 2.6–4.0× faster than PyTorch (non-causal) and 5.3–9.9× faster (causal).
- In Setting 2, StreamKL achieves 3.8–6.5× speedup over PyTorch (non-causal) and 9.1–14.0× (causal).
- Compared to torch.compile, StreamKL is competitive or faster, particularly in causal settings where the recomputation cost is halved by skipping masked tiles.
Ablation: The Split-K variant provides up to 8.4× speedup for small $N_Q$ workloads. The fused backward kernel is up to 2.0× faster than separate kernels when $N_Q$ is small, justifying the adaptive selection strategy.

Significance

The paper claims that StreamKL is the first solution capable of sustaining attention distillation workloads with context lengths of 64K and beyond on a single GPU. By eliminating the quadratic materialization of attention distributions, StreamKL removes the primary system bottleneck preventing the scaling of attention distillation in modern LLM training (e.g., for sparse-attention models like DeepSeek V3.2 and GLM-5). It demonstrates that fused, online primitives can be successfully extended beyond standard attention to complex coupled distributions like KL divergence.