Un nuevo trabajo de ingeniería de Meta, publicado en el blog de PyTorch, presenta varias técnicas de fusión de kernels para operaciones de normalización comunes como LayerNorm y RMSNorm. La promesa: reducir la sobrecarga de memoria de estos kernels, que son intensamente limitados por el ancho de banda, y con eso acelerar de forma significativa el entrenamiento.
¿Cuánto se puede acelerar?
Los autores muestran que estas técnicas pueden ocultar hasta un 90% de la latencia de un kernel de normalización al fusionarlo con las multiplicaciones de matrices (GEMM). Al final presentan el algoritmo FlashNormAttention, que aplica la fusión a varias normalizaciones alrededor de un kernel de atención como GDPA, logrando hasta un 35% de aceleración del kernel.
El trabajo usa principalmente dos lenguajes de dominio específico (DSL) para kernels: TLX, un conjunto de extensiones de Triton con soporte de bajo nivel para controlar la ejecución en GPU, y Helion, un DSL de alto nivel que destaca por la velocidad de desarrollo, la portabilidad y el autotuning exhaustivo. Los benchmarks se hicieron con el tipo de dato bfloat16, sobre GPU NVIDIA B200 en los centros de datos de Meta, con un límite de potencia de 750 W.
¿Por qué la normalización es un cuello de botella?
Las técnicas de normalización se volvieron indispensables en casi toda arquitectura de aprendizaje profundo, porque estabilizan el entrenamiento y aceleran la convergencia. Pero su ubicuidad trae un problema de rendimiento: son intensamente limitadas por memoria y no usan los Tensor Cores. Eso impide saturar la capacidad de cómputo del hardware.
Usando la arquitectura Kunlun del modelo de recomendación GEM de Meta como ejemplo, la normalización se lleva alrededor del 20% de la latencia total de entrenamiento. Es decir, se pierde de entrada un 20% del throughput de cómputo del hardware si no se optimiza. En un LLM típico, más limitado por cómputo, la normalización todavía puede representar cerca del 10% de la latencia total.
Como la mayoría de las operaciones de normalización siguen o preceden a una multiplicación de matrices (en la MLP o en la atención), el trabajo se concentra en cómo fusionar eficientemente ambas.
El problema del tiling
A diferencia de la fusión de activaciones estándar (por ejemplo GEMM más ReLU), el desafío de fondo con la normalización es la diferencia en el tiling. La normalización es, por naturaleza, una operación de reducción que necesita acceder a los datos a lo largo de toda una dimensión. Un GEMM, en cambio, se divide en tiles en ambas dimensiones, de modo que cada tile no abarca una fila completa, lo que vuelve imposible una normalización posterior por filas.
La solución más directa sería estirar el tile del GEMM para que abarque toda la dimensión interna. Pero eso trae dos problemas: se aleja del tiling óptimo para un GEMM puro, degradando su rendimiento, e impone un límite duro al tamaño de la entrada. Con algo de cálculo aproximado sobre una GPU Blackwell con 228 KB de memoria compartida y datos en bfloat16, la restricción deja el tamaño de tile, y por lo tanto la dimensión N, en un máximo de 512 para que el kernel siquiera pueda ejecutarse.
Aun así, esta estrategia rinde para valores pequeños de N. Para formas como 64 y 128, la fusión logra un ahorro de latencia del 17% al 32% en el kernel de LayerNorm. Pero a medida que N crece más allá de 128, la ganancia desaparece e incluso se convierte en una gran regresión, porque forzar el tamaño de tile se aleja demasiado del óptimo del GEMM sin fusionar.
Lazy Pre-Norm: el truco matemático
La segunda técnica, llamada Lazy Pre-Norm, fusiona la pre-normalización con el GEMM siguiente como una fusión de prólogo, para el caso especial de RMSNorm sin parámetros afines. La clave es que existe una dependencia cíclica: se necesita el resultado de la reducción para procesar cada tile, pero ese resultado solo está disponible al final del bucle.
El rescate es una propiedad matemática: en una RMSNorm sin afines, la multiplicación es por filas, y todos los elementos de la misma fila se multiplican por el mismo factor. Eso permite reordenar el cálculo, de modo que la multiplicación de matrices se hace primero y el escalado de normalización se aplica después, rompiendo la dependencia cíclica sin perder precisión.
