El equipo de MoonMath AI ha lanzado un kernel de atención directa bf16 para la GPU MI300X de AMD. Está escrito en HIP, no en ensamblaje escrito a mano. El código es de código abierto bajo la licencia MIT. El equipo de MoonMath.ai informa que supera a AITER v3, el núcleo optimizado de AMD, en todas las formas probadas. El acceso bare-metal provino de HotAisle, un proveedor de nube de AMD.
La atención es la operación fusionada softmax(QKᵀ/√d)·V dentro de cada transformador. El MI300X es la GPU para centros de datos CDNA3 de AMD, con objetivo ISA (gfx942). Este kernel se ejecuta únicamente en ese hardware.
TL;DR
MoonMath.ai ofrece un núcleo de atención directa bf16 de código abierto para AMD MI300X, escrito en HIP, no en ensamblador (MIT). Supera al AITER v3 de AMD en todas las formas y modos de redondeo: media geométrica 1,18×/1,15×/1,08×, hasta 1,26×. El truco principal: los envoltorios ASM de una instrucción le permiten elegir el código de operación mientras el compilador asigna los registros. La mayor parte de la aceleración se debe a la ubicación de la memoria: K en LDS, V hot en L1, Q y acumuladores en los registros. Un PR real de SGLang lo utilizó para acelerar la difusión de vídeo Wan2.1 en 1,23×, sin regresión de calidad.
Entendiendo el núcleo
Un kernel es un pequeño programa que se ejecuta directamente en los numerosos núcleos de la GPU para realizar un cálculo específico (en este caso, las matemáticas de atención) tan rápido como lo permita el hardware. El kernel calcula la atención directa en bf16 solo en MI300X. Toma entradas en formato BSHD o BHSD, sin transposición. La dimensión de la cabeza está fijada en 128. Admite cualquier longitud de secuencia, incluida la atención cruzada.
Hay límites reales. No existe una máscara causal, ni GQA, ni lotes varlen. Las salidas son bf16 y se ejecuta exclusivamente en hardware gfx942.
Los números están estrictamente controlados. Los tres modos de redondeo coinciden con la regla de redondeo por modo de AITER. Cada producción finita se encuentra dentro de 1 bf16 ULP de AITER. El manejo de NaN e Inf es idéntico en bits y los resultados son deterministas.
El truco principal: envoltorios de ensamblaje de una sola instrucción
La técnica central evita un dilema familiar. Los intrínsecos del compilador mantienen el código ordenado pero permiten que el compilador reordene o cambie el nombre de los operandos. El ensamblaje en línea sin formato proporciona control pero obliga a la gestión manual de registros y direcciones.
MoonMath envuelve exactamente una instrucción en una función __device__ __forceinline__. Las restricciones de ASM extendidas describen los operandos. El equipo de investigación elige el código de operación. El compilador aún asigna registros y rastrea el flujo de datos.
La restricción “+v”(c) vincula la entrada y salida del acumulador al mismo VGPR. No se emite ninguna instrucción de copia. Esto mantiene el kernel cerca del HIP normal. Todavía dirige la máquina una instrucción a la vez.
La arquitectura: ocho olas, dos grupos, dos barreras
Una unidad de cómputo CDNA3 tiene cuatro unidades SIMD. El bloque de libros de texto es de cuatro ondas. En cambio, MoonMath ejecuta ocho ondas por bloque, en dos grupos de cuatro.
Los dos grupos ejecutan la misma secuencia Q*K, softmax, O += P*V. Están compensados por una fase. Mientras un grupo satura el núcleo de la matriz, el otro ejecuta softmax y emite cargas. Luego se intercambian, por lo que el núcleo de la matriz nunca queda inactivo.
Hay dos s_barriers por iteración. Uno se sienta en el traspaso de fase. Uno se sitúa en el límite de la iteración. Las esperas por contador se encargan del resto de la sincronización.
Esto se hace eco de la alternancia matmul y softmax de FlashAttention-3. No copia la división warp de productor y consumidor de FA3. En CDNA3, cada movimiento de memoria ya es asíncrono, por lo que no es necesaria una onda de productor dedicada.
Dónde viven los datos y por qué 16×16×16
La mayor parte de la aceleración proviene de la ubicación de la memoria. K fluye desde HBM a LDS, con doble buffer, compartido por las ocho ondas. V permanece caliente en L1, lea en cada matmul PV. Q y acumuladores viven en registros.
El equipo de investigación eligió el MFMA de 16×16×16 en lugar de 32×32×8. Ambas formas tienen un rendimiento idéntico. El mosaico más pequeño se acumula en 4 elementos fp32 por carril, frente a 16. La menor presión del acumulador deja espacio para una captación previa más profunda y un tercer mosaico Q.
Dos victorias posteriores cierran la brecha. Un tercer mosaico Q (3Q) aumenta la reutilización de datos por mosaico K y V cargado. Una división KV de cola estilo Flash-Decoding rescata la ronda fraccionaria varada en las 304 CU del MI300X. Estas victorias se suceden en cascada. Mover V a L1 liberó el LDS que luego llena el tercer mosaico Q.
Punto de referencia
Las pruebas se realizaron en MI300X en bf16, dimensión de cabeza 128. Cada forma se midió en tres modos de redondeo. RTNE redondea al par más cercano. RTNA redondea al más cercano, se aleja de cero. RTZ se trunca hacia cero.
Las geomedias a lo largo del barrido favorecen MoonMath. Frente a AITER, obtiene una puntuación de 1,18 × (RTNE), 1,15 × (RTNA) y 1,08 × (RTZ). En comparación con Modular MAX, las geomedias van de 1,44× a 1,49× y las aceleraciones por forma alcanzan 1,59×.
RTZ es el modo más rápido de AITER y la carrera más reñida. La forma RTZ (4, 16, 16384) pasó de 0,95× a 1,07×. La división KV de la cola es lo que cerró esa brecha final.