Resumen Técnico: Axon: Un superoptimizador de síntesis para programas tensoriales

Planteamiento del problema

Escribir kernels de alto rendimiento para aceleradores de IA (como el Trainium de Amazon) requiere una profunda experiencia en detalles específicos del hardware, incluyendo jerarquías de memoria, estrategias de tiling (segmentación), selección de instrucciones, disposición de datos (layout) y fusión de operadores. Aunque los lenguajes basados en tiles como NKI (Neuron Kernel Interface) abstraen algunos detalles de bajo nivel, alcanzar la utilización máxima del hardware sigue siendo un desafío para la mayoría de los ingenieros de ML. Los enfoques de optimización existentes sufren de varias limitaciones:

• Reglas diseñadas a mano: Muchos sistemas dependen de reglas de reescritura escritas manualmente para transformaciones algebraicas, las cuales son propensas a errores y difíciles de escalar a nuevas combinaciones de operadores.
• Generación basada en plantillas: Algunos enfoques utilizan plantillas de generación de código escritas a mano, lo que limita la flexibilidad.
• Alcance limitado: Muchos optimizadores se centran solo en CPUs o GPUs, fallando al abordar las arquitecturas únicas de múltiples motores y dirigidas por DMA de los aceleradores de IA modernos.
• Falta de corrección demostrable: Las pruebas probabilísticas o el reordenamiento basado en heurísticas pueden omitir configuraciones óptimas o introducir errores semánticos.

Metodología

Axon es un superoptimizador de síntesis diseñado para programas de aceleradores de IA basados en tiles. Acepta programas tensoriales de alto nivel (escritos en una interfaz similar a NumPy) y genera automáticamente instrucciones de destino optimizadas (NKI) a través de un pipeline de múltiples etapas. La innovación central es el $\nu$Graph ("nu-graph"), una representación interna persistente que rastrea todas las variantes de programa semánticamente equivalentes simultáneamente durante todo el proceso de síntesis, postergando la elección de optimización a la selección empírica.

El pipeline consta de las siguientes etapas:

1. Transformaciones algebraicas mediante propagación de operadores:

   • Axon descubre grafos de computación semánticamente equivalentes propagando operadores a través del grafo de computación.
   • Intenta reordenar operadores adyacentes (por ejemplo, mover una multiplicación elemento a elemento más allá de una multiplicación de matrices) para exponer paralelismo y oportunidades de fusión.
   • Garantía de corrección: En lugar de usar reglas de reescritura diseñadas a mano, Axon descompone las operaciones tensoriales en computaciones granulares y utiliza un solucionador SMT (Satisfiability Modulo Theories) sobre tensores no acotados para verificar de manera demostrable que el reordenamiento preserva la semántica. Esto maneja formas de tensores arbitrarias y no requiere definiciones de propiedades manuales.

2. Tiling Simbólico:

   • El sistema segmenta (tiles) cada operación tensorial utilizando parámetros simbólicos para dimensiones de strip, block y tile.
   • Estos parámetros están restringidos por descripciones de hardware (por ejemplo, capacidades de búfer en chip y límites específicos de instrucciones).
   • Axon genera múltiples versiones con diferentes configuraciones de tiling, permitiendo la optimización independiente de ramas paralelas (por ejemplo, separando las ramas de RMSNorm y MatMul para que se ejecuten en diferentes motores de cómputo).

3. Síntesis de ISA (Selección de Instrucciones):

   • Utilizando un enfoque de síntesis de programas basado en sketches, Axon reduce las operaciones tensoriales de alto nivel a instrucciones ISA de destino (por ejemplo, nc_matmul y dma_copy de Trainium).
   • Explora secuencias de instrucciones válidas y transformaciones de disposición de datos (como transposiciones) requeridas por instrucciones de hardware específicas.
   • La equivalencia entre la operación de alto nivel y la secuencia de instrucciones sintetizada es verificada nuevamente mediante SMT.

4. Fusión de ISA y Operadores:

   • Se realiza una segunda ronda de transformaciones algebraicas en el grafo de ISA sintetizado para exponer nuevas oportunidades de fusión creadas por el proceso de síntesis (por ejemplo, fusionar una transposición con una multiplicación de matrices).
   • Fusión de Nodos: Combina operaciones consecutivas con dimensiones de tiling compatibles para mantener los resultados intermedios en los búferes en chip (SBUF/PSUM), reduciendo el tráfico de HBM.
   • Fusión de Instrucciones: Reemplaza secuencias de instrucciones ISA con variantes fusionadas (por ejemplo, dma_copy + nc_transpose $\rightarrow$ dma_transpose).

5. Emisión de Código y Selección Empírica:

   • Axon instancia los parámetros simbólicos con valores concretos para emitir programas NKI concretos.
   • Cada candidato es compilado por el compilador de Neuron y ejecutado en el hardware de destino.
   • El kernel con mejor desempeño es seleccionado basándose en mediciones empíricas de latencia y validación de corrección numérica contra una implementación de referencia.

Contribuciones Clave

El artículo presenta tres contribuciones primarias:

1. Transformaciones Algebraicas Demostrables: Una metodología para descubrir automáticamente grafos de computación semánticamente equivalentes mediante la propagación de operadores y la verificación de equivalencia usando SMT sobre tensores no acotados. Esto elimina la necesidad de reglas de reescritura diseñadas a mano.
2. Pipeline Basado en Síntesis: Un pipeline completo para aceleradores de IA basados en tiles que explora conjuntamente transformaciones algebraicas, selección de instrucciones, configuraciones de tiling y fusión. Mantiene todas las variantes válidas en un $\nu$Graph y selecciona el kernel óptimo de forma empírica.
3. Evaluación Exhaustiva: Una evaluación en el Trainium de Amazon a través de 20 benchmarks, que incluyen operadores individuales y kernels complejos de múltiples operadores de LLM (por ejemplo, Group Query Attention, Gated MLP).

Resultados Experimentales

Axon fue evaluado contra tres líneas base: el compilador Neuron de Amazon, kernels NKI optimizados a mano y Mirage (un compilador de búsqueda de vanguardia).

• Vs. Compilador Neuron:

  • Operadores Individuales: Axon logra aceleraciones de hasta 3.7x (ej. SiLU en dimensiones grandes) con una aceleración de media geométrica de 1.23x–1.90x a través de varios operadores.
  • Kernels de Múltiples Operadores: Axon logra aceleraciones de hasta 19x en kernels como Transpose+MatMul y hasta 1.47x en RMSNorm+MatMul. Las aceleraciones de media geométrica oscilan entre 1.07x y 1.32x.
  • Razonamiento: Las ganancias se atribuyen a mejores estrategias de tiling, utilización eficiente de DMA y la capacidad de fusionar operaciones a través de diferentes motores de cómputo (por ejemplo, ejecutando RMSNorm en motores de Escalar/Vector mientras MatMul se ejecuta en el Motor de Tensores en paralelo).

• Vs. NKI Optimizado a Mano:

  • Axon iguala o supera a los kernels escritos por expertos hasta en 1.35x en configuraciones específicas (ej. Cumsum y RoPE), demostrando que la exploración automatizada puede rivalizar con el ajuste manual sin requerir una profunda experiencia en hardware.

• Vs. Mirage:

  • En los benchmarks soportados por ambos, Axon logra aceleraciones de media geométrica del 2–10%.
  • Crucialmente, Axon sintetiza con éxito kernels para 16 benchmarks adicionales que Mirage no puede manejar debido a su dependencia de plantillas diseñadas a mano.

• Sobrecarga de Compilación:

  • El espacio de búsqueda puede generar hasta ~30,000 candidatos para kernels complejos. La Fase 1 (síntesis simbólica) es rápida (<82 segundos). La Fase 2 (compilación y ejecución) domina el tiempo (hasta ~1.7 horas para capas complejas de LLM). El artículo señala que este es un costo de una sola vez, comparable a otros compiladores basados en esquemas, y los kernels optimizados pueden almacenarse en caché para su reutilización.

Significancia y Reivindicaciones

El artículo afirma que Axon representa un cambio de la optimización basada en heurísticas o plantillas hacia un enfoque de síntesis impulsado por la semántica. Al aprovechar los solucionadores SMT para garantizar la equivalencia semántica, Axon puede descubrir automáticamente transformaciones algebraicas complejas (como el reordenamiento de operaciones a través de diferentes motores de cómputo) que son difíciles de codificar manualmente.

La significancia radica en:

• Automatización: Reducir la carga de los programadores para ajustar manualmente los kernels para hardware específico.
• Corrección: Proporcionar garantías demostrables para las transformaciones, a diferencia de los métodos de prueba probabilísticos.
• Desempeño: Lograr aceleraciones significativas en aceleradores de IA modernos mediante la optimización conjunta de tiling, selección de instrucciones y fusión en un espacio de búsqueda unificado.
• Generalidad: La capacidad de manejar kernels de LLM complejos de múltiples operadores que los compiladores de búsqueda existentes (como Mirage) no pueden sintetizar.

Los autores enfatizan que, si bien la sobrecarga de compilación no es despreciable, es un costo estándar de una sola vez para los compiladores de optimización, y las ganancias de rendimiento resultantes justifican la inversión, particularmente para cargas de trabajo de LLM a gran escala.