Autoatención, kernels CUDA y FlashAttention

Los Transformers modernos se apoyan en el cálculo matricial de autoatención, compiladores de GPU de alto nivel como Triton y kernels CUDA que tienen en cuenta la memoria, como FlashAttention, para procesar enormes secuencias de tokens a la máxima velocidad del hardware.

1. Cómo funciona matemáticamente la autoatención en un transformador de visión

En un Vision Transformer (ViT), una imagen (p. ej., de 224×224 píxeles) se divide en una cuadrícula de regiones de 16×16, lo que produce N = 196 regiones aplanadas. Cada región se proyecta linealmente en un vector de dimensión D (p. ej., 768) y se combina con una representación posicional, formando una matriz de entrada X ∈ R^(N × D).

• Creación de consultas, claves y valores. El modelo multiplica la matriz de regiones X por tres matrices de pesos aprendidas (W_Q, W_K, W_V ∈ R^(D × d_k)) para crear tres representaciones distintas de cada región de la imagen:

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

• Consulta (Q): representa lo que busca una región concreta de la imagen (p. ej., "Soy una textura de pelaje y busco ojos u orejas cerca").
• Clave (K): anuncia qué información contiene una región (p. ej., "Contengo un borde curvo y nítido").
• Valor (V): contiene la carga real de características visuales que se transmitirá hacia delante si se produce una coincidencia.

• Cálculo de las puntuaciones de atención. Para medir la intensidad de la relación de cada región con todas las demás, el modelo calcula el producto escalar entre todas las consultas y claves:

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

Dividir por sqrt(d_(k)) evita que los productos escalares grandes lleven la función softmax posterior a regiones con gradientes diminutos. Esto produce una matriz de atención N × N en la que la entrada (i, j) puntúa cuánta atención debe prestar la región i a la región j.

• Ponderación y agregación con Softmax. Aplicar Softmax por filas convierte las puntuaciones sin procesar en probabilidades que suman 1. Multiplicar estas probabilidades por la matriz de valores V produce una representación combinada para cada región:

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

En la autoatención de múltiples cabezas, este mismo proceso se ejecuta en paralelo en varias cabezas (p. ej., 12 cabezas), lo que permite que distintas cabezas se centren simultáneamente en diferentes distancias espaciales, colores o partes de objetos.

2. Comparación práctica de código: CUDA C++ directo frente a OpenAI Triton

Para ver cómo ha evolucionado la programación de GPU, considera una sencilla suma de vectores elemento a elemento (C = A + B) para N elementos.

CUDA C++ directo (gestión a nivel de hilo)

En CUDA C++, programas desde la perspectiva de un único hilo escalar. Debes calcular manualmente los identificadores de los hilos, comprobar los límites de los arrays y escribir las configuraciones de lanzamiento repetitivas del lado del 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 (programación en Python a nivel de bloque)

En Triton, programas desde la perspectiva de un bloque de datos directamente en Python. El compilador de Triton gestiona automáticamente por dentro la sincronización de hilos, la coalescencia de memoria y la vectorización.

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: solución al cuello de botella de memoria de los Transformers

La autoatención estándar se encuentra con una barrera importante al procesar imágenes de alta resolución o secuencias de texto largas, porque la matriz de atención S = QK^T crece de forma cuadrática (N × N) con la longitud de secuencia N.

• El cuello de botella del ancho de banda de memoria (HBM frente a SRAM). Las GPU modernas de NVIDIA tienen dos niveles principales de memoria:

• HBM (memoria de alto ancho de banda / VRAM): grande (p. ej., 80 GB), pero relativamente lenta para leer y escribir.
• SRAM (memoria compartida en el chip): diminuta (~20 MB en total entre todos los Streaming Multiprocessors), pero aproximadamente de 10× a 20× más rápida que HBM.

En la atención estándar, la GPU calcula S = QK^T, escribe la enorme matriz N × N en la lenta HBM, vuelve a leerla para aplicar Softmax, escribe las probabilidades de nuevo en HBM y las lee otra vez para multiplicarlas por V. Los Tensor Cores pasan la mayor parte del tiempo inactivos, esperando las transferencias de memoria.

• Procesamiento por bloques (sin materializar nunca la matriz N × N). Presentado por Tri Dao et al. (2022), FlashAttention usa la fusión de kernels para mantener todo el cálculo de atención dentro de la rápida SRAM del chip. Divide Q, K y V en bloques pequeños que caben en la SRAM, calcula la atención bloque a bloque y escribe solo la matriz de salida final de vuelta en HBM. La enorme matriz intermedia N × N nunca se guarda en VRAM.

• El truco de Softmax en línea. Normalmente, Softmax necesita conocer el valor máximo y la suma de exponenciales de una fila completa antes de normalizar cualquier elemento individual. FlashAttention usa un algoritmo Softmax en línea que lleva un registro de los máximos acumulados y los factores de normalización bloque a bloque, y reescala matemáticamente las salidas parciales conforme se procesa cada nuevo bloque.

• Recálculo durante la retropropagación. Como FlashAttention no guarda la matriz de atención N × N durante el paso hacia delante, emplea un ingenioso compromiso durante el entrenamiento: guarda solo las pequeñas estadísticas de normalización en HBM y vuelve a calcular los bloques de atención sobre la marcha dentro de la rápida SRAM durante la retropropagación. Realizar cálculos adicionales en SRAM resulta mucho más rápido que leer matrices grandes desde HBM, lo que produce aceleraciones de 2× a 4× en el tiempo real de ejecución y reduce el uso de memoria de cuadrático O(N²) a lineal O(N).