🤖 machine learning

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

StreamKL introduce una novedosa primitiva de GPU fusionada que elimina los cuellos de botella cuadráticos de memoria y E/S de la destilación de atención al transmitir teselas de consulta-clave en un solo paso, logrando aceleraciones significativas y reduciendo la huella de memoria de $O(N_QN_K)$ a $O(1)$ para permitir la destilación de contexto largo en una sola GPU.

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

Publicado 2026-06-19

📖 5 min de lectura🧠 Análisis profundo

CC BY 4.0

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

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Imagina que estás intentando enseñar a un estudiante pequeño y rápido (un "modelo estudiante") a pensar exactamente como un profesor brillante y lento (un "modelo profesor"). En el mundo de la IA, esto lo hacen comparando cómo ambos "prestan atención" a diferentes partes de una historia o frase. Este proceso se llama Destilación de Atención.

Para realizar esta comparación, la computadora calcula un número específico llamado Divergencia KL. Piensa en esto como una "puntuación de distancia" que te dice qué tan diferente es la atención del estudiante de la del profesor. El objetivo es hacer que esta puntuación sea lo más pequeña posible.

El Problema: La "Explosión de Memoria"

El artículo explica que realizar esta comparación para historias largas (como una novela de 100,000 palabras) es actualmente una pesadilla para la memoria de la computadora.

Aquí está la analogía:
Imagina que tienes dos pizarras gigantes, una para la atención del profesor y otra para la del estudiante. Para compararlas, el método antiguo requiere que escribas cada una de las posibles combinaciones de palabras en estas pizarrras.

Si tienes 64,000 palabras, tienes que escribir 64,000 × 64,000 pares. Eso es más de 4 mil millones de números.
Hacer esto requiere una pizarra tan grande que no cabe en la memoria principal de la computadora (HMB). Es como intentar guardar una biblioteca de libros en un zapato.
Debido a que la computadora no puede visualizar toda la imagen a la vez, tiene que trocear la historia en piezas diminutas, procesarlas y luego volver a juntarlas. Esto es lento, como intentar leer un libro mirando una letra a la vez y escribiéndola antes de pasar a la siguiente.

La Solución: StreamKL (El enfoque de "Streaming")

Los autores crearon una nueva herramienta llamada StreamKL. En lugar de escribir todo en una pizarra gigante primero, StreamKL utiliza un truco inteligente para calcular la "puntuación de distancia" sobre la marcha, como una cinta transportadora.

La Analogía Creativa: La Línea de Ensamblaje de una Fábrica
Imagina una fábrica donde estás comparando dos cintas transportadoras de productos (la atención del profesor y la atención del estudiante).

La forma antigua: Detienes la línea, viertes todos los productos en un enorme almacén (HBM), los mides todos y luego limpias. Esto ocupa todo el almacén y es lento.
La forma de StreamKL: Mantienes los productos moviéndose en la cinta transportadora. A medida que cada par de artículos pasa por un sensor (el chip GPU), los comparas instantáneamente, calculas la diferencia y lanzas el resultado a un bolsillo diminuto (SRAM) antes de que llegue el siguiente par. Nunca detienes la línea y nunca necesitas un almacén. Solo necesitas un bolsillo.

Cómo Funciona (El Truco de Magia)

El artículo describe las dos partes principales de esta magia:

El Paso hacia Adelante (Calculando la Puntuación): Los investigadores inventaron una nueva fórmula matemática que permite a la computadora actualizar la "puntuación de distancia" de forma incremental. Mientras recorre los datos, mantiene un recuento acumulativo de solo unos pocos números (como un máximo acumulado y una suma) en lugar de toda la lista. Esto significa que puede manejar historias de cualquier longitud sin quedarse sin memoria.
El Paso hacia Atrás (Aprendiendo de los Errores): Cuando la computadora necesita aprender de la puntuación para mejorar al estudiante, normalmente necesita volver a mirar los datos. El método antiguo guarda toda la lista gigante de datos para volver a verlos. StreamKL es más inteligente: desecha la lista pero recuerda algunas "llaves secretas" (llamadas valores LSE). Cuando necesita volver a mirar, utiliza estas llaves para reconstruir la pieza específica de datos que necesita, justo en ese momento, calcula la lección y luego la olvida de nuevo. Es como recordar la receta de un pastel para poder hornear una rebanada cada vez que quieras probarlo, en lugar de hornear todo el pastel y guardarlo en el refrigerador.

Los Resultados: Velocidad y Espacio

El artículo probó esto en potentes GPUs NVIDIA (H200 y A100) con contextos muy largos (hasta 512,000 palabras).

Ahorro de Memoria: StreamKL redujo la memoria adicional necesaria de "cuadrática" (explotando a terabytes) a "constante" (manteniéndose diminuta). Pasó de necesitar 512 GB de memoria para un contexto de 64k a necesitar casi nada extra. Esto permite que una sola GPU maneje tareas que antes requerían una supercomputadora o que eran imposibles.
Velocidad: Debido a que no tiene que escribir y leer cantidades masivas de datos de ida y vuelta, es increíblemente rápido.
- En algunas pruebas, fue 43 veces más rápido que el método estándar para calcular la puntuación.
- En la fase de aprendizaje, fue 14 veces más rápido.

Resumen

StreamKL es una nueva forma de enseñar a los modelos de IA a prestar atención. Resuelve el problema de "quedarse sin memoria" al tratar con textos largos impidiendo que la computadora escriba la lista completa de comparaciones. En su lugar, transmite los datos a través de un flujo de trabajo diminuto y eficiente, calculando el resultado instantáneamente. Esto hace posible entrenar y ejecutar modelos de IA en computadoras individuales que anteriormente eran demasiado grandes para manejarse.

Resumen Técnico: StreamKL

Declaración del Problema

La destilación de atención, una técnica utilizada para entrenar una distribución de atención para que coincida con otra minimizando la divergencia de Kullback-Leibler (KL), es crítica para la destilación de conocimiento, la poda de modelos, el aprendizaje continuo y el entrenamiento de modelos de lenguaje de gran escala (LLMs) con atención dispersa. Sin embargo, las implementaciones existentes enfrentan severos cuellos de botella de escalabilidad al manejar longitudes de contexto largas ( $N_Q, N_K$ ).

Los enfoques actuales "materializan" ambas distribuciones de atención ( $P_1, P_2 \in \mathbb{R}^{N_Q \times N_K}$ ) en la Memoria de Alto Ancho de Banda (HBM) antes de computar la reducción KL elemento a elemento. Esto resulta en:

Huella de Memoria Cuadrática: Almacenar dos matrices $N_Q \times N_K$ requiere $O(N_Q N_K)$ de memoria. Para un contexto de 64K tokens con 32 cabezales en BF16, esto por sí solo consume ~512 GB, excediendo la capacidad de 141 GB de una sola GPU NVIDIA H200.
Costos de E/S Cuadráticos: La escritura de estas distribuciones en HBM y su posterior lectura para la reducción KL domina el tiempo de ejecución, escalando cuadráticamente con la longitud del contexto.
Ineficiencia del Fragmentado (Chunking): Los métodos de trabajo comunes como el procesamiento por fragmentos intercambian latencia por memoria; los fragmentos pequeños subutilizan el paralelismo de la GPU, mientras que los fragmentos grandes activan errores de falta de memoria (OOM).

El cómputo de la atención estándar enfrentó problemas similares, los cuales FlashAttention resolvió mediante kernels fusionados y en línea (online). Sin embargo, aplicar esto a la destilación de atención no es trivial porque la divergencia KL acopla dos distribuciones independientes a través de un término de diferencia de logit ponderado, lo que hace que una formulación online de un solo paso sea matemáticamente compleja.

Metodología: StreamKL

StreamKL introduce la primera primitiva de GPU fusionada de un solo paso para la divergencia KL de atención que elimina la materialización $O(N_Q N_K)$ de las distribuciones intermedias.

1. Formulación Online (Paso hacia adelante / Forward Pass)

StreamKL deriva un nuevo esquema de actualización online para la reducción KL de dos distribuciones acopladas. Al reformular la divergencia KL por fila, los autores expresan el objetivo en términos de cinco escalares de ejecución por fila:

$m_1, m_2$ : Máximos corrientes para las dos distribuciones de logit.
$l_1, l_2$ : Sumas de exponenciales no normalizadas (similares al softmax online).
$acc$: Un acumulador para la diferencia de logit ponderada.

La formulación del objetivo es:
$L = \frac{acc}{l_1} + \text{LSE}_2 - \text{LSE}_1$
donde $\text{LSE}_t = m_t + \log l_t$ .

Esta formulación permite que el kernel transmita (stream) teselas de query-key a través de la SRAM en el chip, actualizando estos cinco escalares incrementalmente. Crucialmente, el acumulador $acc$ es un vector de tamaño $B_Q$ (tamaño de bloque), no una matriz completa, reduciendo la sobrecarga de reescalado aritmético en comparación con el FlashAttention estándar. El algoritmo maneja el enmascaramiento causal saltándose completamente las teselas de clave más allá de la frontera causal.

2. Estrategia de Recomputación (Paso hacia atrás / Backward Pass)

Para evitar almacenar las matrices cuadráticas $P_1$ y $P_2$ durante la retropropagación, StreamKL emplea una estrategia de recomputación:

Estado Guardado: Solo se almacenan en HBM los valores KL finales y los vectores de Log-Sum-Exp ( $\text{LSE}_1, \text{LSE}_2$ ) ( $O(N_Q)$ ).
Regeneración: Durante el paso hacia atrás, las probabilidades de atención se regeneran tesela por tesela utilizando los LSE guardados y las matrices originales $Q, K$ ( $P_i = e^{S_i - \text{LSE}}$ ).
Cómputo de Gradientes: El kernel computa gradientes para dos configuraciones:
- Configuración 1: $P_1$ fijo, optimizando $P_2$ (destilación estándar).
- Configuración 2: $P_2$ fijo, optimizando $P_1$ .
Fusión de Kernel: Los autores diseñan un kernel de retropropagación fusionado que visita cada par de teselas $QK$ exactamente una vez, reduciendo el tráfico de HBM. Para gestionar el conflicto entre la acumulación de $dQ$ (a lo largo de $K$ ) y $dK$ (a lo largo de $Q$ ), el kernel utiliza operaciones atómicas para $dQ$ cuando el búfer cabe en la caché L2, recurriendo a kernels separados cuando la contención se vuelve demasiado alta.

3. Optimizaciones de Kernel de GPU

Variante Split-K: Para cargas de trabajo con pocos conteos de query (por ejemplo, decodificación autorregresiva donde $N_Q=1$ ), una variante Split-K particiona la dimensión de la clave entre los bloques de hilos para saturar los Multi-Processors de Flujo (SMs).
Aceleración por TMA: En las GPUs Hopper, la implementación aprovecha el Tensor Memory Accelerator (TMA) para copias de memoria masivas asíncronas, ocultando la latencia de memoria.
Estabilidad Numérica: El cómputo de la razón de logaritmos en la Configuración 2 se realiza utilizando diferencias de logits y LSE guardados en lugar de logaritmos de probabilidad directa para evitar la inestabilidad numérica ( $-\infty$ o NaN).

Contribuciones Clave

Identificación de Cuellos de Botella: El artículo identifica los cuellos de botella cuadráticos de memoria e IO que impiden que la divergencia KL de atención escale a contextos largos.
Reformulación Online: Introduce la primera reformulación online de la divergencia KL de atención, permitiendo un primitivo de paso hacia adelante totalmente fusionado y un paso hacia atrás basado en recomputación que elimina los intermediarios cuadráticos.
Implementación Eficiente: Los autores diseñan e implementan kernels de GPU eficientes con optimizaciones dedicadas (Split-K, TMA, fusión adaptativa) tanto para el paso hacia adelante como para el paso hacia atrás.
Escalabilidad: StreamKL reduce la huella de memoria adicional de la destilación de atención de $O(N_Q N_K)$ a $O(1)$ , permitiendo la destilación de contexto largo en una sola GPU.

Resultados Experimentales

Evaluado en GPUs NVIDIA A100 y H200 a través de longitudes de contexto desde 4K hasta 512K:

Eficiencia de Memoria: StreamKL mantiene una huella de memoria plana independientemente de la longitud del contexto. En contraste, los modelos de referencia (PyTorch, torch.compile, FLA) escalan cuadráticamente y sufren errores de OOM en 16K–64K. A 64K, StreamKL utiliza más de 16,000× menos HBM que torch.compile/FLA y más de 135,000× menos que PyTorch.
Latencia de Paso hacia Adelante (Forward Latency):
- StreamKL logra una aceleración de 13–18× sobre PyTorch y de 2.4–3.7× sobre torch.compile/FLA en entornos no causales.
- Bajo enmascaramiento causal, las aceleraciones se amplían a 26–43× sobre PyTorch y 4.2–7.1× sobre los modelos de referencia, ya que StreamKL omite los cálculos enmascarados por completo mientras que los otros aún materializan matrices completas.
Latencia de Paso hacia Atrás (Backward Latency):
- En la Configuración 1, StreamKL es de 2.6–4.0× más rápido que PyTorch (no causal) y 5.3–9.9× más rápido (causal).
- En la Configuración 2, StreamKL logra una aceleración de 3.8–6.5× sobre PyTorch (no causal) y de 9.1–14.0× (causal).
- Comparado con torch.compile, StreamKL es competitivo o más rápido, particularmente en entornos causales donde el costo de recomputación se reduce a la mitad al saltarse las teselas enmascaradas.
Ablación: La variante Split-K proporciona hasta 8.4× de aceleración para cargas de trabajo con $N_Q$ pequeño. El kernel de retropropagación fusionado es hasta 2.0× más rápido que los kernels separados cuando $N_Q$ es pequeño, justificando la estrategia de selección adaptativa.

Significancia

El artículo afirma que StreamKL es la primera solución capaz de sostener cargas de trabajo de destilación de atención con longitudes de contexto de 64K en adelante en una sola GPU. Al eliminar la materialización cuadrática de las distribuciones de atención, StreamKL elimina el principal cuello de botella del sistema que impide el escalado de la destilación de atención en el entrenamiento de LLMs modernos (por ejemplo, para modelos de atención dispersa como DeepSeek V3.2 y GLM-5). Demuestra que los primitivos fusionados y online pueden extenderse con éxito más allá de la atención estándar hacia distribuciones acopladas complejas como la divergencia KL.