Resumo Técnico: StreamKL

Definição do Problema

A destilação de atenção, uma técnica utilizada para treinar uma distribuição de atenção para corresponder a outra minimizando a divergência de Kullback-Leibler (KL), é crítica para destilação de conhecimento, poda de modelos (pruning), aprendizado contínuo e treinamento de LLMs com atenção esparsa. No entanto, as implementações existentes enfrentam gargalos severos de escalabilidade ao lidar com comprimentos de contexto longos ($N_Q, N_K$).

As abordagens atuais "materializam" ambas as distribuições de atenção ($P_1, P_2 \in \mathbb{R}^{N_Q \times N_K}$) na Memória de Alta Largura de Banda (HBM) antes de computar a redução KL elemento a elemento. Isso resulta em:

1. Pegada de Memória Quadrática: Armazenar duas matrizes $N_Q \times N_K$ requer $O(N_Q N_K)$ de memória. Para um contexto de 64K tokens com 32 cabeças em BF16, isso sozinho consome ~512 GB, excedendo a capacidade de 141 GB de uma única GPU NVIDIA H200.
2. Custos de IO Quadráticos: Escrever essas distribuições na HBM e lê-las de volta para a redução KL domina o tempo de execução, escalando quadraticamente com o comprimento do contexto.
3. Ineficiência do Chunking: Soluções de contorno comuns, como o processamento em blocos (chunked processing), trocam latência por memória; blocos pequenos subutilizam o paralelismo da GPU, enquanto blocos grandes disparam erros de Falta de Memória (OOM).

A computação de atenção padrão enfrentou problemas semelhantes, os quais o FlashAttention resolveu via kernels fundidos e online. No entanto, aplicar isso à destilação de atenção não é trivial porque a divergência KL acopla duas distribuições independentes através de um termo de diferença de logit ponderado, tornando a formulação online de passagem única matematicamente complexa.

Metodologia: StreamKL

O StreamKL introduz o primeiro primitivo de GPU fundido de passagem única para a divergência KL de atenção que elimina a materialização $O(N_Q N_K)$ das distribuições intermediárias.

1. Formulação Online (Passagem Forward)

O StreamKL deriva um novo esquema de atualização online para a redução KL de duas distribuições acopladas. Ao reformular a divergência KL por linha, os autores expressam o objetivo em termos de cinco escalares de execução por linha:

• $m_1, m_2$: Máximos correntes para as duas distribções de logit.
• $l_1, l_2$: Somas não normalizadas de exponenciais (semelhante ao softmax online).
• $acc$: Um acumulador para a diferença de logit ponderada.

A formulação reformulada é:
$$L = \frac{acc}{l_1} + \text{LSE}_2 - \text{LSE}_1$$
onde $\text{LSE}_t = m_t + \log l_t$.

Esta formulação permite que o kernel processe fluxos de tiles de query-key através da SRAM on-chip, atualizando esses cinco escalares incrementalmente. Crucialmente, o acumulador $acc$ é um vetor de tamanho $B_Q$ (tamanho do bloco), não uma matriz completa, reduzindo o overhead de reescalonamento aritmético em comparação ao FlashAttention padrão. O algoritmo lida com o mascaramento causal saltando tiles de chave inteiros além da fronteira causal.

2. Estratégia de Recomputação (Passagem Backward)

Para evitar o armazenamento das matrizes quadráticas $P_1$ e $P_2$ durante a retropropagação, o StreamKL emprega uma estratégia de recomputação:

• Estado Salvo: Apenas os valores KL finais e os vetores de Log-Sum-Exp ($\text{LSE}_1, \text{LSE}_2$) são armazenados na HBM ($O(N_Q)$).
• Regeneração: Durante a passagem backward, as probabilidades de atenção são regeneradas tile-a-tile usando os LSEs salvos e as matrizes originais $Q, K$ ($P_i = e^{S_i - \text{LSE}}$).
• Computação de Gradiente: O kernel computa gradientes para duas configurações:
  • Configuração 1: $P_1$ fixo, otimizando $P_2$ (destilação padrão).
  • Configuração 2: $P_2$ fixo, otimizando $P_1$.
• Fusão de Kernel: Os autores projetam um kernel backward fundido que visita cada par de tiles $QK$ exatamente uma vez, reduzindo o tráfego de HBM. Para gerenciar o conflito entre acumular $dQ$ (ao longo de $K$) e $dK$ (ao longo de $Q$), o kernel usa operações atômicas para $dQ$ quando o buffer cabe no cache L2, recorrendo a kernels separados quando a contenção se torna muito alta.

3. Otimizações de Kernel de GPU

• Variante Split-K: Para cargas de trabalho com contagens de query pequenas (ex: decodificação autorregressiva onde $N_Q=1$), uma variante Split-K particiona a dimensão de chave entre blocos de threads para saturar os Multiprocessadores de Fluxo (SMs).
• Aceleração por TMA: Em GPUs Hopper, a implementação aproveita o Tensor Memory Accelerator (TMA) para cópias de memória em massa assíncronas, ocultando a latência de memória.
• Estabilidade Numérica: A computação da razão de logaritmo na Configuração 2 é realizada usando diferenças de logit e LSEs salvos em vez de logs diretos de probabilidade para evitar instabilidade numérica ($-\infty$ ou NaN).

Principais Contribuições

1. Identificação de Gargalos: O artigo identifica os gargalos quadráticos de memória e IO que impedem a escala da divergência KL de atenção para contextos longos.
2. Reformulação Online: Introduz a primeira reformulação online da divergência KL de atenção, permitindo um primitivo forward totalmente fundido e uma passagem backward baseada em recomputação que elimina intermediários quadráticos.
3. Implementação Eficiente: Os autores projetam e implementam kernels de GPU eficientes com otimizações dedicadas (Split-K, TMA, fusão adaptativa) para as passagens forward e backward.
4. Escalabilidade: O StreamKL reduz a pegada adicional de HBM da destilação de atenção de $O(N_Q N_K)$ para $O(1)$, permitindo a destilação de contexto longo em uma única GPU.

Resultados Experimentais

Avaliado em GPUs NVIDIA A100 e H200 em comprimentos de contexto de 4K a 512K:

• Eficiência de Memória: O StreamKL mantém uma pegada de memória constante, independentemente do comprimento do contexto. Em contraste, as linhas de base (PyTorch, torch.compile, FLA) escalam quadraticamente e atingem erros de OOM em 16K–64K. Em 64K, o StreamKL usa mais de 16.000× menos HBM que o torch.compile/FLA e mais de 135.000× menos que o PyTorch.
• Latência Forward:
  • O StreamKL alcança um aceleração de 13–18× sobre o PyTorch e 2.4–3.7× sobre o torch.compile/FLA em configurações não causais.
  • Sob mascaramento causal, as acelerações aumentam para 26–43× sobre o PyTorch, como o StreamKL pula completamente as computações mascaradas enquanto as linhas de base ainda materializam matrizes completas.
• Latência Backward:
  • Na Configuração 1, o StreamKL é 2.6–4.0× mais rápido que o PyTorch (não causal) e 5.3–9.9× mais rápido (causal).
  • Na Configuração 2, o StreamKL alcança 3.8–6.5× de aceleração sobre o PyTorch (não causal) e 9.1–14.0× (causal).
  • Comparado ao torch.compile, o StreamKL é competitivo ou mais rápido, particularmente em cenários causais onde o custo de recomputação é reduzido pela metade ao pular tiles mascarados.
• Ablação: A variante Split-K fornece até 8.4× de aceleração para cargas de trabalho com $N_Q$ pequeno. O kernel backward fundido é até 2.0× mais rápido que kernels separados quando $N_Q$ é pequeno, justificando a estratégia de seleção adaptativa.

Significância

O artigo afirma que o StreamKL é a primeira solução capaz de sustentar cargas de trabalho de destilação de atenção com comprimentos de contexto de 64K e além em uma única GPU. Ao eliminar a materialização quadrática das distribuições de atenção, o StreamKL remove o principal gargalo de sistema que impede a escala da destilação de atenção no treinamento de LLMs modernos (ex: para modelos de atenção esparsa como DeepSeek V3.2 e GLM-5). Ele demonstra que primitivos fundidos e online podem ser estendidos com sucesso além da atenção padrão para distribuições acopladas complexas como a divergência KL.