Las longitudes de contexto de los modelos de lenguaje (LLM) están creciendo rápidamente, y las arquitecturas evolucionan hacia esquemas de atención complejos como la Atención Latente Multi-Cabeza (MLA) y la Atención de Consulta Agrupada (GQA). Por lo tanto, la «velocidad de pensamiento» de la inteligencia artificial está cada vez más determinada no solo por el alto rendimiento de las multiplicaciones de matrices, sino también por la matemática trascendental de la función softmax.
Las funciones trascendentales no se pueden expresar como raíces de ecuaciones polinómicas con coeficientes racionales, lo que significa que trascienden operaciones algebraicas básicas como la suma y la multiplicación, que son las que dominan los Tensor Cores. En el contexto específico de softmax, la función exponencial natural es la más costosa computacionalmente y se ejecuta en las Unidades de Función Especial (SFUs). En las instrucciones de ensamblador de NVIDIA (SASS), esta función se invoca mediante la instrucción MUFU.EX2. Esta separación arquitectónica crea un cuello de botella en softmax dentro del bloque de atención, donde los potentes motores de matriz deben esperar mientras los caminos de datos de la SFU normalizan las puntuaciones de atención.
Optimización de softmax con Blackwell Ultra
NVIDIA Blackwell Ultra alivia este cuello de botella al duplicar el rendimiento de la SFU en comparación con la arquitectura estándar de NVIDIA Blackwell.
Este blog profundiza en la mecánica de softmax dentro del bucle de atención, explora cómo las optimizaciones de hardware de Blackwell Ultra eliminan los cuellos de botella en la tubería, y proporciona un benchmark para medir la velocidad de MUFU.EX2 por uno mismo.
Funcionamiento de la atención
Un componente fundamental de los modelos de lenguaje grandes modernos es el mecanismo de atención, que permite al modelo transformar dinámicamente vectores de tokens estáticos en representaciones contextuales. En su esencia, es un proceso de re-pesado de la información al permitir que los tokens ajusten su importancia entre sí. Para facilitar esta interacción, cada token en una secuencia se proyecta en tres roles funcionales:
- Consulta: Representa lo que el token actual busca entender sobre su contexto.
- Clave: Representa el perfil de un token que otros utilizan para hacer coincidencias.
- Valor: Contiene el contenido informativo real. Una vez que se confirma una coincidencia entre una consulta y una clave, el Valor es la información que se transfiere al token original.
La figura 1 a continuación muestra la atención en acción. En dos oraciones se utiliza la palabra “perro” en diferentes definiciones. Inicialmente, las incrustaciones de ambas menciones de “perro” son idénticas.
La atención opera calculando un producto punto entre la consulta de “perro” y las claves de cada otro token en la secuencia.
Si la consulta para “perro” se alinea bien con la clave para “perezoso”, indica un alto grado de relevancia. Esta interacción permite que la palabra “perro” obtenga el valor específico de su vecino. Al final de este ciclo, el vector original para “perro” se actualiza físicamente con el contenido de sus vecinos, evolucionando de una definición genérica a una incrustación contextualizada que “comprende” si se refiere a un animal letárgico o al punto culminante de una temporada calurosa.
Relación entre softmax y atención
Softmax es la fase crítica de toma de decisiones que convierte las puntuaciones de compatibilidad en pesos utilizables. Una vez que se calculan los productos punto entre consultas y claves, las puntuaciones resultantes se pasan a través de la función softmax para ser normalizadas en probabilidades que suman exactamente uno. Este paso determina el “rango de atención” del modelo, decidiendo efectivamente qué tokens priorizar y cuáles ignorar. Sin softmax, el modelo no tendría forma objetiva de pesar la información que recopila, lo que llevaría a una mezcla de datos incontrolable y ruidosa.
Sin embargo, la operación softmax es la principal fuente del “desnivel de rendimiento” observado en la IA de contexto largo. Cada token en una secuencia debe compararse con todos los demás, creando una matriz de atención de [8,192 x 8,192] para una secuencia de 8,192 tokens. Normalizar esta matriz requiere miles de millones de cálculos trascendentales y crece cuadráticamente con la longitud de la secuencia. Esto crea un cuello de botella, donde el volumen de matemáticas trascendentales puede estancar toda la tubería de inferencia.
Blackwell Ultra se centra en acelerar estos cálculos exponenciales específicamente para aliviar este cuello de botella matemático y asegurar que el sistema pueda manejar la normalización masiva requerida para grandes ventanas de contexto sin sacrificar el rendimiento.
Aliviando el cuello de botella de softmax en Blackwell Ultra
Al duplicar el rendimiento de la SFU para exponenciales en la arquitectura Blackwell Ultra, NVIDIA alivia este cuello de botella y permite un pipeline de procesamiento más equilibrado y eficiente. Esto resulta en un rendimiento general más rápido, especialmente para tareas que dependen en gran medida de los mecanismos de atención.
La figura 2 a continuación ilustra la dependencia secuencial inherente al mecanismo de atención estándar, a menudo denominado bucle de atención, ejecutado en la generación anterior de NVIDIA Blackwell (GB200). Note que el Multiprocesador de Streaming (SM) carga dos bloques de hilos que ejecutan bucles de atención simultáneamente. Estos bucles de atención separados se denotan en dos tonos diferentes de verde.
Este pipeline consta de tres fases distintas que deben ejecutarse en orden:
- BMM1 (cálculo de puntuaciones): Los Tensor Cores realizan una multiplicación de matrices para calcular las puntuaciones de atención brutas.
- Softmax (normalización): La tubería pasa a las SFUs para normalizar estas puntuaciones en probabilidades utilizando funciones exponenciales.
- BMM2 (agregación de contexto): La tubería regresa a los Tensor Cores para multiplicar las probabilidades por los vectores de valor.
La línea de tiempo ilustra las limitaciones de latencia inherentes a la GPU Blackwell durante la ejecución del kernel de atención. Dado que la segunda multiplicación de matrices (BMM2) actúa sobre la salida del softmax, no puede comenzar hasta que la normalización esté completa.
El menor rendimiento de las SFUs de la GPU Blackwell obliga a los Tensor Cores a estar inactivos entre el cálculo de puntuaciones (BMM1) y la agregación de contexto (BMM2). Esta dependencia impide que la tubería aproveche completamente los recursos de cómputo y prolonga la duración de la operación softmax.
La siguiente línea de tiempo, como se muestra en la figura 3, demuestra el impacto directo del aumento del rendimiento de las GPUs Blackwell Ultra en los sistemas NVIDIA GB300 NVL72 y NVIDIA HGX B300, que duplicaron el rendimiento de la SFU en la misma secuencia de instrucciones.
Visualmente, el ancho de los bloques de softmax se reduce casi un 50%, reflejando la capacidad del hardware para procesar instrucciones MUFU al doble de la tasa.
Esta reducción en la latencia de softmax ajusta toda la tubería. La brecha entre BMM1 y BMM2 se minimiza drásticamente, permitiendo que los Tensor Cores cambien entre la multiplicación de consultas-claves y la multiplicación de probabilidades-valores con un mínimo de espera. El resultado es un bucle principal más denso donde los motores de matriz de alto rendimiento pasan un mayor porcentaje del tiempo total de ejecución activos, lo que se traduce directamente en un mayor rendimiento de inferencia.
Evaluación del rendimiento de MUFU.EX2
Para verificar empíricamente el rendimiento teórico del pipeline MUFU, podemos construir un micro-benchmark sintético. El siguiente código de kernel aísla las instrucciones exponenciales para medir el conteo de ciclos crudos sin interferencias de la latencia de memoria global u otras operaciones aritméticas.
Este conjunto de pruebas lanza una cuadrícula de hilos donde cada hilo realiza un bucle denso de instrucciones MUFU.EX2. Al cronometrar la ejecución y compararla con la frecuencia del reloj, puedes calcular directamente el rendimiento de las instrucciones efectivas y validar el punto de saturación de ancho de banda mencionado anteriormente.
Paso 1: Clona el siguiente repositorio para obtener el benchmark exp2-bg300.cu.
git clone https://github.com/jamieliNVIDIA/mufu_ex2_bench.git cd mufu_ex2_bench
Paso 2: Compila con (Usando sm100f para GB300 o sm103a para GB200).
nvcc -O3 -gencode=arch=compute_103a,code=sm_103a --extended-lambda -o /tmp/exp2-gb300.out exp2-gb300.cu
Resultados de muestra
Observamos que GB300 presenta un rendimiento de FLOPs aproximadamente 2 veces superior al de GB200 para todos los tipos de datos probados, en línea con el rendimiento duplicado de la SFU.
Blackwell (GB200)
exp2 BF16x2 2454 Gop/s (4908 GFLOPS) exp2 BF16 4938 Gop/s exp2 FP32 4943 Gop/s
Blackwell Ultra (GB300)
exp2 BF16x2 4996 Gop/s (9992 GFLOPS) exp2 BF16 9738 Gop/s exp2 FP32 10024 Gop/s
Rendimiento de propagación hacia adelante en Blackwell vs Blackwell Ultra
La transición de Blackwell a Blackwell Ultra proporciona un aumento dirigido en el rendimiento de cálculo impulsado por un aumento de 2x en el rendimiento de la SFU. Esta actualización de hardware acelera directamente la tubería de propagación hacia adelante (FPROP) para modelos como DeepSeek-V3.
FPROP es el proceso donde los datos de entrada viajan “hacia adelante” a través de la red neuronal, desde la capa de entrada, pasando por las capas ocultas, hasta la capa de salida, para generar una predicción. Cada vez que el modelo produce una nueva palabra, debe ejecutar un pase completo de FPROP.
La figura 4 a continuación muestra que, al duplicar el rendimiento de las SFUs, el GB300 reduce drásticamente el tiempo de ejecución de las capas softmax dentro de los bloques de atención. Esta normalización más rápida significa que la GPU pasa menos tiempo procesando las puntuaciones de atención y más tiempo utilizando los motores de matriz de alta velocidad para el cálculo de la siguiente capa, aumentando directamente la velocidad general del pase hacia adelante.
Los resultados de benchmark destacan un aumento del ~35% en la capacidad de FPROP para operaciones FP8. Esta mejora es especialmente notable en FP8 porque las matemáticas matriciales son ya extremadamente rápidas. En este régimen de baja precisión, el tiempo dedicado a softmax representa un mayor porcentaje del paso total.
Iniciando con Blackwell Ultra
Las dinámicas de rendimiento de DeepSeek-V3 en Blackwell Ultra destacan un cuello de botella crítico, pero a menudo pasado por alto, en la inferencia: el costo computacional de las operaciones no lineales.
Al optimizar y comprimir el mecanismo de atención, los modelos de última generación aumentan efectivamente la densidad de las operaciones softmax en relación con los cálculos lineales estándar, exponiendo a las SFUs como reguladoras del rendimiento total.
Blackwell Ultra aborda directamente este cuello de botella. Al duplicar el rendimiento de estas unidades especializadas, Blackwell Ultra desbloquea el embotellamiento trascendental que anteriormente forzaba a los potentes Tensor Cores a estar inactivos. Los resultados de benchmark confirman el impacto, demostrando una ganancia del 35% en la propagación hacia adelante en FP8.
Para arquitecturas modernas y altamente optimizadas, el camino hacia una inferencia más rápida no se trata solo de Tensor Cores más rápidos, sino también de asegurar que las unidades de matemáticas no lineales sean lo suficientemente rápidas para mantener el ritmo.
Visita el repositorio de NVIDIA trtllm-gen para más benchmarks e información sobre cómo utilizar esta aceleración de velocidad de SFU en cargas de trabajo. Duplicar el rendimiento de las SFUs para MUFU.EX2 es solo una de las muchas características que permiten la rápida velocidad de atención de Blackwell Ultra. La co-diseño extremo de hardware y software de NVIDIA acelera todo el bucle de atención a través de tecnologías como:
- Descargar reducciones críticas de “find-max” al controlador de memoria Tensor mediante LDTM.STAT.
- Optimización del rendimiento usando CUDNN.
- Optimización de movimientos de datos de KVCache usando NVFP4.
Esté atento al blog técnico de NVIDIA para futuras publicaciones.
Agradecimientos
Agradecimientos especiales al equipo de ingeniería de cuDNN por crear los benchmarks y construir las optimizaciones de software que hacen posible este rendimiento de vanguardia.
