Los núcleos de atención rápida y la fusión de operadores reorganizan el cálculo de los transformers para limitar los desplazamientos de datos y los lanzamientos de núcleos en el acelerador. Estas optimizaciones dejan en principio inalterados los pesos del modelo y su función matemática, aunque el orden de las operaciones y la precisión pueden producir pequeñas diferencias numéricas. El tutorial publicado por Machine Learning Mastery detalla estas técnicas, con ejemplos de código.
Una capa de transformer se basa en operaciones comunes: proyecciones lineales, atención, normalización, funciones de activación y adiciones residuales. Una implementación directa lanza un núcleo de cálculo distinto, llamado «kernel», para cada operación, y escribe cada tensor intermedio en la memoria del acelerador. El cálculo sigue siendo correcto, pero el tiempo de ejecución puede dedicarse más al desplazamiento de valores y al lanzamiento de núcleos que a los cálculos útiles.
Fusión de operadores
El bloque SwiGLU, utilizado por muchos modelos de solo decodificador, ilustra este problema. Una implementación directa lanza multiplicaciones de matrices para las proyecciones, un núcleo para la función SiLU, otro para la multiplicación elemento por elemento, y luego una última multiplicación de matrices. En cada frontera, un tensor intermedio puede escribirse en la memoria de alto ancho de banda (HBM) y luego releerse. La fusión consiste en combinar estas etapas en un solo núcleo. Ahorra tiempo cuando los desplazamientos de memoria y los lanzamientos pesan mucho, y mucho menos cuando una gran multiplicación de matrices ya domina la ejecución. La elección del núcleo sigue siendo empírica: hay que perfilar toda la carga de trabajo.
Otras fusiones son comunes en los transformers, como añadir una conexión residual seguida de una normalización RMSNorm o LayerNorm, o combinar la desquantización con una multiplicación de matrices. Sin embargo, la fusión tiene sus límites: las operaciones con organizaciones de datos incompatibles o un estado intermedio muy voluminoso se combinan mal. Un núcleo demasiado grande puede aumentar la presión sobre los registros, reducir la tasa de ocupación y volverse más lento que varios núcleos bien ajustados.
Atención tipo FlashAttention y atención paginada
Una atención clásica calcula una matriz de puntuaciones de tamaño T² por cabeza, siendo T la longitud de la secuencia. En una GPU, producirla, escribirla en la memoria global y luego releerla consume mucho ancho de banda, lo que puede convertirse en un cuello de botella durante los largos prellenados («prefill»). FlashAttention procesa las consultas, las claves y los valores por bloques. Para cada uno, calcula una parte de las puntuaciones, actualiza las estadísticas del softmax y acumula la salida correspondiente en la memoria situada en el chip, sin almacenar la matriz completa. PyTorch expone esta optimización mediante la función scaled_dot_product_attention(). El tutorial precisa que imponer un backend fijo puede hacer fallar el código si ningún backend admite las entradas o el hardware.
La caché KV, que conserva las claves y valores ya calculados, crece de un token a la vez, para una longitud final impredecible. Una asignación contigua por petición deja por tanto espacio reservado y sin usar. La caché paginada divide el almacenamiento en bloques de tamaño fijo y se apoya en una tabla de bloques que vincula las posiciones lógicas con los bloques físicos. El núcleo de atención debe entonces traducir cada posición a través de esta tabla. La paginación no reduce el número de elementos lógicos, pero limita la fragmentación y la capacidad reservada.
En el ejemplo proporcionado, vLLM utiliza el modelo Qwen/Qwen2.5-1.5B-Instruct con bloques de 16 tokens (block_size=16) y el backend «TRITON_ATTN». Según el tutorial, vLLM verifica la compatibilidad del backend con el modelo y el hardware durante la inicialización, y señala un error en caso de incompatibilidad.
Fusión y CUDA Graphs con torch.compile
La función torch.compile() captura una secuencia de operaciones con TorchDynamo, y luego la compila con TorchInductor, su backend por defecto. TorchInductor puede generar un solo núcleo Triton para operaciones compatibles que el modo estándar lanzaría por separado. Las interrupciones del grafo, debidas por ejemplo a un control de flujo dependiente de los datos o a conversiones de tensores a escalares Python, limitan las zonas que el compilador puede optimizar.
El modo «reduce-overhead» pide a TorchInductor que utilice CUDA Graphs cuando el grafo capturado sea elegible. Un CUDA Graph registra la secuencia de operaciones GPU y la reproduce con un solo envío al procesador. No crea fusión adicional, complementa la de TorchInductor. Exige direcciones de memoria estables y formas compatibles, y puede consumir más memoria. El tutorial recomienda verificar los resultados con tolerancias adaptadas al tipo de datos, perfilar la ejecución y medir el rendimiento de extremo a extremo.
