Escribiendo Código GPU Rápido: De la Teoría a la Práctica

El Manual de Optimización

Los artículos 1–5 cubrieron el hardware de la GPU, el modelo de programación CUDA, la pila de software y el análisis de rendimiento. Este artículo conecta esos conceptos con la optimización práctica — las técnicas que los autores de kernels (y compiladores como Triton) utilizan para escribir código GPU rápido. Si los artículos anteriores construyeron el modelo mental, este lo pone en acción.

El manual de optimización se apoya en cuatro pilares:

• Maximizar la intensidad aritmética — realizar más cómputo por byte transferido. Esto empuja tu kernel hacia el lado limitado por cómputo del roofline, donde realmente puede utilizar los FLOPS que el hardware provee.
• Maximizar la ocupación — mantener suficientes warps activos para ocultar la latencia de memoria e instrucciones. Cuando un warp se detiene en una solicitud de memoria, otro warp debería estar listo para ejecutar.
• Usar la memoria eficientemente — fusionar accesos a memoria global, organizar datos en bloques dentro de la memoria compartida y minimizar las transferencias entre host y dispositivo. El ancho de banda de memoria es casi siempre el cuello de botella.
• Minimizar la sobrecarga — fusionar kernels para reducir la sobrecarga de lanzamiento, usar CUDA Graphs para inferencia y evitar puntos de sincronización innecesarios entre host y dispositivo.

Cada una de las siguientes secciones profundiza en una de estas técnicas. Juntas, representan el conjunto de herramientas que frameworks como PyTorch, compiladores como Triton y bibliotecas como cuBLAS utilizan internamente para extraer el máximo rendimiento del hardware.

Fusión de Kernels: La Mayor Ganancia

La fusión de kernels es probablemente la optimización individual más impactante para cargas de trabajo de deep learning. La idea es sencilla: en lugar de lanzar kernels separados para cada operación — digamos, una multiplicación de matrices, una suma de sesgo y una activación ReLU — se fusionan en un único kernel que realiza las tres operaciones en una sola pasada sobre los datos.

¿Por qué la fusión ayuda tanto? Tres razones:

• Menos viajes de ida y vuelta a memoria: sin fusión, los resultados intermedios deben escribirse en HBM después de cada kernel y leerse de vuelta para el siguiente. Con fusión, esos intermedios permanecen en registros o memoria compartida — sin tocar HBM en absoluto. Para una secuencia de 3 operaciones, esto puede reducir el tráfico de HBM en aproximadamente 2/3.
• Menos lanzamientos de kernels: cada lanzamiento de kernel incurre en aproximadamente 5–10 μs de sobrecarga del lado de la CPU (validación de argumentos, llamadas al driver, planificación de la GPU). Fusionar 10 kernels pequeños en 1 ahorra alrededor de 50 μs — lo cual importa mucho cuando tu modelo ejecuta miles de estas secuencias por segundo.
• Mayor intensidad aritmética: el kernel fusionado realiza más FLOPs por byte cargado desde HBM, moviéndolo más cerca del (o más allá del) punto de cresta del roofline. Un kernel que era limitado por memoria como tres lanzamientos separados puede volverse limitado por cómputo cuando se fusiona.

Cuantifiquemos el ahorro de memoria para un ejemplo concreto: un matmul + sesgo + ReLU fusionado versus tres kernels separados.

import numpy as np

batch, dim = 1024, 4096
elem_size = 2  # bytes per FP16 element
hbm_bw = 3.35e12  # H100 HBM bandwidth (bytes/s)

# Unfused: 3 separate kernels
# Op 1: matmul (read W, X; write Y)
# Op 2: add bias (read Y, bias; write Z)
# Op 3: relu (read Z; write out)
unfused_reads = (batch * dim + dim * dim + batch * dim + dim + batch * dim) * elem_size
unfused_writes = (batch * dim + batch * dim + batch * dim) * elem_size
unfused_total = unfused_reads + unfused_writes

# Fused: 1 kernel
# Read W, X, bias once; write out once
fused_reads = (batch * dim + dim * dim + dim) * elem_size
fused_writes = batch * dim * elem_size
fused_total = fused_reads + fused_writes

savings = 1 - fused_total / unfused_total

# Collect rows for aligned comparison
rows = [
    ("Unfused (3 kernels)", unfused_total),
    ("Fused (1 kernel)",    fused_total),
]
w_name = max(len(r[0]) for r in rows)
w_mb   = max(len(f"{r[1]/1e6:.1f}") for r in rows)
w_ms   = max(len(f"{r[1]/hbm_bw*1000:.3f}") for r in rows)

print("Kernel Fusion: matmul + bias + relu")
print(f"  Shape: ({batch}, {dim}) × ({dim}, {dim})")
print()
for name, total in rows:
    mb = total / 1e6
    ms = total / hbm_bw * 1000
    print(f"  {name:<{w_name}}  HBM traffic: {mb:>{w_mb}.1f} MB  Time at peak BW: {ms:>{w_ms}.3f} ms")
print()
print(f"  Memory traffic reduction: {savings:.0%}")
print(f"  This is why torch.compile exists — it fuses operations automatically.")

Los números hablan por sí mismos. La fusión no cambia la cantidad de cómputo útil — los mismos FLOPs ocurren de cualquier manera — pero reduce drásticamente el tráfico de memoria necesario para realizarlos. Para operaciones limitadas por memoria (que es lo que son la mayoría de las operaciones element-wise y reducciones), esta es la diferencia entre un kernel que se atasca en el ancho de banda y uno que corre cerca del rendimiento máximo.

Multiplicación de Matrices con Tiling

Las multiplicaciones de matrices grandes no caben completamente en registros o memoria compartida — las matrices son simplemente demasiado grandes. El enfoque estándar es tiling : cargar un bloque pequeño (tile) de cada matriz de entrada en memoria compartida, calcular el resultado parcial para ese tile, acumular el resultado parcial en registros y repetir hasta completar el cálculo completo.

Aquí está la estructura conceptual de un kernel de matmul con tiling:

# Pseudocode for tiled matmul C = A @ B
# A is (M, K), B is (K, N), C is (M, N)
# Tile size: BLOCK_M × BLOCK_K and BLOCK_K × BLOCK_N

for tile_k in range(0, K, BLOCK_K):
    # 1. Load tiles from global memory (HBM) → shared memory
    A_tile = A[row_start:row_end, tile_k:tile_k+BLOCK_K]  # → shared mem
    B_tile = B[tile_k:tile_k+BLOCK_K, col_start:col_end]  # → shared mem
    __syncthreads()  # wait for all threads to finish loading

    # 2. Compute partial matmul from shared memory (fast!)
    C_partial += A_tile @ B_tile  # stays in registers

    __syncthreads()  # wait before loading next tile

# 3. Write accumulated result to global memory
C[row_start:row_end, col_start:col_end] = C_partial

¿Por qué funciona tan bien el tiling? La memoria compartida tiene aproximadamente 5 ciclos de latencia comparado con aproximadamente 300 ciclos para HBM . Al cargar un tile una vez en memoria compartida y reutilizarlo para muchas operaciones de multiplicación-acumulación, aumentamos drásticamente la intensidad aritmética. Un tile de $K \times K$ se carga una vez pero se usa para $K$ operaciones MAC por elemento, dando una intensidad que crece proporcionalmente con la dimensión del tile.

El bucle externo sobre tile_k recorre la dimensión K en bloques de BLOCK_K . En cada paso, ambos tiles son cargados cooperativamente por todos los hilos del bloque, una barrera __syncthreads() asegura que las cargas estén completas, y luego cada hilo calcula su porción del producto parcial usando la rápida memoria compartida. El resultado se acumula en registros a lo largo de todas las iteraciones, y solo el resultado final se escribe de vuelta a HBM.

Precisión Mixta: Más Cómputo, Menos Memoria

Las GPUs modernas son dramáticamente más rápidas en matemáticas FP16/BF16 que en FP32. La H100, por ejemplo, entrega aproximadamente 990 TFLOPS en FP16 pero solo alrededor de 67 TFLOPS en FP32 — una diferencia de aproximadamente 15×. El entrenamiento con precisión mixta explota esta brecha usando aritmética de menor precisión donde es seguro hacerlo, mientras retiene precisión completa donde importa.

La receta estándar de precisión mixta tiene tres partes:

• Almacenar pesos en FP32 — la copia maestra de cada parámetro permanece en precisión completa. Esto asegura que las pequeñas actualizaciones de gradiente (que podrían perderse en la precisión limitada de FP16) se acumulen con precisión a lo largo de muchos pasos de entrenamiento.
• Convertir a FP16/BF16 para los pases forward y backward — aquí es donde viene la velocidad. Los Tensor Cores operan con entradas FP16/BF16, y el tamaño reducido de datos significa menos tráfico de memoria (FP16 tiene la mitad de bytes que FP32, duplicando efectivamente el ancho de banda de memoria).
• Acumular en FP32 — los Tensor Cores realizan nativamente operaciones de multiplicación-acumulación con entradas FP16 y acumuladores FP32. Esto preserva la estabilidad numérica durante las grandes reducciones que ocurren en las multiplicaciones de matrices.

Los beneficios son dobles: cómputo más rápido (los Tensor Cores operan con datos FP16/BF16) y menor tráfico de memoria (la mitad de bytes significa, en efecto, el doble del ancho de banda de memoria utilizable). En PyTorch, habilitar precisión mixta es sencillo:

# PyTorch automatic mixed precision
with torch.autocast(device_type='cuda', dtype=torch.float16):
    output = model(input)    # forward in FP16
    loss = criterion(output, target)

scaler = torch.amp.GradScaler()
scaler.scale(loss).backward()  # backward in FP16
scaler.step(optimizer)          # update in FP32
scaler.update()

El GradScaler merece una breve explicación. FP16 tiene un rango dinámico limitado — aproximadamente de $6 \times 10^{-8}$ a $65504$ — lo que significa que gradientes pequeños pueden sufrir underflow a cero durante el pase backward. El GradScaler resuelve esto multiplicando la pérdida por un factor de escala grande antes de llamar a .backward() , lo que mantiene los valores de gradiente dentro del rango representable de FP16. Antes del paso del optimizador, los gradientes se dividen por el mismo factor de escala para restaurar su magnitud real. Si se detecta overflow (cualquier gradiente se convierte en inf o NaN ), ese paso se omite por completo y el factor de escala se reduce para la siguiente iteración. Con el tiempo, el scaler encuentra un factor de escala estable que evita tanto el underflow como el overflow.

CUDA Graphs: Amortizando la Sobrecarga de Lanzamiento

Cada lanzamiento de kernel incurre en aproximadamente 5–10 μs de sobrecarga del lado de la CPU: validación de argumentos, llamadas al driver y planificación de la GPU. Para kernels grandes que se ejecutan durante milisegundos, esta sobrecarga es despreciable. Pero para cargas de trabajo de inferencia con muchos kernels pequeños — proyecciones de atención, normalizaciones de capa, sumas de sesgo, funciones de activación — la sobrecarga de lanzamiento puede convertirse en una fracción significativa del tiempo total de ejecución.

Los CUDA Graphs resuelven esto grabando una secuencia de lanzamientos de kernels una vez, capturando todo el grafo de dependencias, y luego reproduciendo toda la secuencia con un solo comando de lanzamiento. La sobrecarga de CPU baja de $N \times 5\mu s$ a $1 \times 5\mu s$ — una ganancia sustancial cuando $N$ está en los cientos.

Aquí está el patrón de uso típico en PyTorch:

# Record a CUDA Graph
static_input = torch.randn(32, 512, device='cuda')
static_output = torch.empty(32, 512, device='cuda')

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    static_output = model(static_input)

# Replay — no CPU-side kernel launch overhead
for batch in dataloader:
    static_input.copy_(batch)  # copy new data into the recorded buffer
    g.replay()                  # replays all kernels in one go
    result = static_output      # output is in the pre-allocated buffer

Observa que los tensores de entrada y salida están preasignados — durante la grabación, CUDA captura las direcciones de memoria exactas utilizadas, y durante la reproducción reutiliza esas mismas direcciones. Los nuevos datos se alimentan copiándolos en el búfer de entrada preasignado con .copy_() , y el resultado aparece en el búfer de salida preasignado después de .replay() .

Conectando de Vuelta con PyTorch

Todo en este track converge en torch.compile . Cuando llamas a torch.compile(model) , TorchDynamo traza tu código Python, TorchInductor lo convierte en kernels optimizados, y el modelo compilado resultante incorpora muchas de las técnicas que hemos discutido:

• Fusión de kernels (artículo 4, este artículo): TorchInductor fusiona operaciones element-wise, reducciones y operaciones pointwise en kernels Triton combinados, eliminando viajes intermedios de ida y vuelta a HBM.
• Matmul con tiling: cuBLAS maneja las multiplicaciones de matrices estándar con kernels con tiling altamente optimizados; Triton genera kernels con tiling para operaciones fusionadas que incluyen matmuls.
• Precisión mixta: torch.autocast habilita la computación FP16/BF16 con acumulación FP32, desbloqueando el rendimiento de los Tensor Cores y reduciendo a la mitad el tráfico de memoria.
• CUDA Graphs: torch.cuda.CUDAGraph captura y reproduce secuencias de kernels, eliminando la sobrecarga de lanzamiento por kernel para cargas de trabajo de inferencia con formas fijas.
• Coalescing de memoria, memoria compartida, gestión de warps: Triton maneja estos aspectos automáticamente al generar kernels a partir de la salida de TorchInductor. El programador especifica qué computar; Triton decide cómo mapearlo al hardware.

El track de GPU te ha dado el modelo mental para entender por qué funcionan estas optimizaciones. La fusión de kernels reduce el tráfico de HBM, moviendo operaciones limitadas por memoria hacia una mayor intensidad aritmética. La precisión mixta duplica el ancho de banda efectivo y habilita los Tensor Cores. Los CUDA Graphs eliminan la sobrecarga de lanzamiento. El tiling coloca datos en memoria compartida donde pueden reutilizarse muchas veces antes de ser desalojados. Nada de esto es magia — son consecuencias directas de la arquitectura de hardware cubierta en los artículos 1 al 5.

Cuando alguien dice que un modelo está "optimizado para GPU," lo que realmente quiere decir es que el código está estructurado para explotar estas propiedades: alta intensidad aritmética, alta ocupación, patrones de acceso a memoria eficientes y sobrecarga mínima. Ahora sabes lo que significa cada una — y, más importante aún, por qué cada una importa.

Quiz

Pon a prueba tu comprensión de las técnicas prácticas de optimización GPU — fusión de kernels, matmul con tiling, precisión mixta y CUDA Graphs.