Autoatenção, kernels CUDA e FlashAttention

Os Transformers modernos dependem dos cálculos matriciais de autoatenção, de compiladores de GPU de alto nível como o Triton e de kernels CUDA atentos à memória como o FlashAttention para processar sequências enormes de tokens na velocidade máxima do hardware.

1. Como a autoatenção funciona matematicamente em um Vision Transformer

Em um Vision Transformer (ViT), uma imagem (por exemplo, de 224×224 pixels) é dividida em uma grade de regiões de 16×16, produzindo N = 196 regiões achatadas. Cada região é projetada linearmente em um vetor de dimensão D (por exemplo, 768) e combinada com um embedding posicional, formando uma matriz de entrada X ∈ R^(N × D).

• Criação de consultas, chaves e valores. O modelo multiplica a matriz de regiões X por três matrizes de pesos aprendidas (W_Q, W_K, W_V ∈ R^(D × d_k)) para criar três representações distintas para cada região da imagem:

Q=XW_(Q), K=XW_(K), V=XW_(V)

• Query (Q): representa o que uma região específica da imagem está procurando (por exemplo, "Sou uma textura de pelo, procurando olhos ou orelhas por perto").
• Key (K): informa quais dados uma região contém (por exemplo, "Contenho uma borda curva bem definida").
• Value (V): contém as características visuais reais que serão transmitidas adiante se houver uma correspondência.

• Cálculo dos escores de atenção. Para medir com que intensidade cada região se relaciona com todas as outras, o modelo calcula o produto escalar entre todas as consultas e chaves:

Scores=(QK^(T))/(sqrt(d_(k)))

Dividir por sqrt(d_(k)) impede que produtos escalares grandes empurrem a função softmax seguinte para regiões com gradientes minúsculos. Isso produz uma matriz de atenção N × N em que a entrada (i, j) pontua quanto a região i deve prestar atenção à região j.

• Ponderação Softmax e agregação. Aplicar Softmax linha a linha converte os escores brutos em probabilidades cuja soma é 1. Multiplicar essas probabilidades pela matriz de valores V produz uma representação combinada para cada região:

Attention(Q,K,V)=softmax((QK^(T))/(sqrt(d_(k))))V

Na autoatenção com múltiplas cabeças, esse mesmo processo é executado em paralelo em várias cabeças (por exemplo, 12), permitindo que diferentes cabeças se concentrem simultaneamente em diferentes distâncias espaciais, cores ou partes de objetos.

2. Comparação prática de código: CUDA C++ direto versus OpenAI Triton

Para ver como a programação de GPUs evoluiu, considere uma simples adição de vetores elemento a elemento (C = A + B) para N elementos.

CUDA C++ direto (gerenciamento no nível da thread)

Em CUDA C++, você programa da perspectiva de uma única thread escalar. É preciso calcular manualmente os IDs das threads, verificar os limites dos arrays e escrever as configurações padronizadas de lançamento no lado do host.

#include <cuda_runtime.h>

// Device kernel: executed by thousands of individual threads
__global__ void vector_add_cuda(const float* A, const float* B, float* C, int N) {
    // Manually compute global thread index from block and thread IDs
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Guard against out-of-bounds memory access
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// Host launch setup
void launch_cuda(const float* d_A, const float* d_B, float* d_C, int N) {
    int threads_per_block = 256;
    int blocks_per_grid = (N + threads_per_block - 1) / threads_per_block;
    vector_add_cuda<<<blocks_per_grid, threads_per_block>>>(d_A, d_B, d_C, N);
}

OpenAI Triton (programação Python no nível do bloco)

No Triton, você programa da perspectiva de um bloco de dados diretamente em Python. O compilador do Triton cuida automaticamente, nos bastidores, da sincronização de threads, da coalescência de memória e da vetorização.

import triton
import triton.language as tl

@triton.jit
def vector_add_triton(A_ptr, B_ptr, C_ptr, N, BLOCK_SIZE: tl.constexpr):
    # Identify which block of data this program instance handles
    pid = tl.program_id(axis=0)
    
    # Create a vector of offsets for the entire block at once
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    
    # Load chunks of memory into fast SRAM, add them, and store back
    a = tl.load(A_ptr + offsets, mask=mask)
    b = tl.load(B_ptr + offsets, mask=mask)
    tl.store(C_ptr + offsets, a + b, mask=mask)

# Launch directly from PyTorch tensors
def launch_triton(A, B, C):
    N = A.numel()
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
    vector_add_triton[grid](A, B, C, N, BLOCK_SIZE=1024)

3. FlashAttention: solução para o gargalo de memória dos Transformers

A autoatenção padrão encontra uma barreira séria ao processar imagens de alta resolução ou longas sequências de texto, porque a matriz de atenção S = QK^T cresce quadraticamente (N × N) com o comprimento N da sequência.

• O gargalo da largura de banda de memória (HBM versus SRAM). As GPUs NVIDIA modernas têm dois níveis principais de memória:

• HBM (High Bandwidth Memory / VRAM): grande (por exemplo, 80 GB), mas relativamente lenta para ler e gravar.
• SRAM (memória compartilhada no chip): minúscula (~20 MB no total em todos os Streaming Multiprocessors), mas aproximadamente de 10× a 20× mais rápida que a HBM.

Na atenção padrão, a GPU calcula S = QK^T, grava a enorme matriz N × N na HBM lenta, lê a matriz de volta para aplicar Softmax, grava as probabilidades novamente na HBM e as lê outra vez para multiplicá-las por V. Os Tensor Cores passam a maior parte do tempo ociosos, esperando transferências de memória.

• Divisão em blocos (sem jamais materializar a matriz N × N). Introduzido por Tri Dao e outros (2022), o FlashAttention usa a fusão de kernels para manter todo o cálculo da atenção dentro da SRAM rápida no chip. Ele divide Q, K e V em pequenos blocos que cabem na SRAM, calcula a atenção bloco a bloco e grava somente a matriz de saída final de volta na HBM. A enorme matriz intermediária N × N nunca é salva na VRAM.

• O truque da Softmax online. Normalmente, Softmax exige conhecer o valor máximo e a soma das exponenciais de uma linha inteira antes de normalizar qualquer elemento individual. O FlashAttention usa um algoritmo Softmax online que acompanha máximos acumulados e fatores de normalização bloco a bloco, redimensionando matematicamente as saídas parciais à medida que cada novo bloco é processado.

• Recálculo durante a retropropagação. Como o FlashAttention não salva a matriz de atenção N × N durante a passagem para a frente, ele adota uma troca inteligente durante o treinamento: salva apenas as pequenas estatísticas de normalização na HBM e recalcula os blocos de atenção em tempo real dentro da SRAM rápida durante a retropropagação. Fazer cálculos adicionais na SRAM acaba sendo muito mais rápido do que ler grandes matrizes da HBM, proporcionando ganhos de velocidade de 2× a 4× no tempo real de execução e reduzindo o uso de memória de quadrático O(N²) para linear O(N).