FairyFuse: Multiplication-Free LLM Inference on CPUs via Fused Ternary Kernels

1. El Problema

La implementación de Grandes Modelos de Lenguaje (LLM) en plataformas basadas únicamente en CPU (como asistentes en dispositivos, servidores de borde y copilotos de escritorio) enfrenta un cuello de botella crítico: el ancho de banda de memoria.

• Limitaciones actuales: Aunque la cuantización de pesos a 4 bits o menos reduce la presión de memoria, los sistemas de inferencia existentes (como llama.cpp) aún descuantizan los pesos comprimidos y ejecutan multiplicaciones en punto flotante. Esto no aprovecha plenamente los ahorros de ancho de banda.
• El desafío de los pesos ternarios: Los pesos ternarios $\{-1, 0, +1\}$ ofrecen una compresión de 16x frente a FP32 y teóricamente permiten reemplazar las multiplicaciones por sumas, restas o operaciones nulas. Sin embargo, sistemas anteriores como Fairy2i (que entrena modelos ternarios complejos de alta calidad) o BitNet no han logrado ejecutar esta lógica a nivel de instrucción en hardware real; siguen utilizando rutinas estándar de punto flotante o tablas de búsqueda (LUT) que generan sobrecarga de memoria y latencia.
• Complejidad de los modelos complejos: Los modelos como Fairy2i utilizan una representación "ampliamente lineal" (widely-linear) que descompone cada capa lineal compleja en ocho sub-operaciones GEMV reales. Una implementación ingenua ejecutaría estas ocho operaciones por separado, perdiendo las ganancias de ancho de banda debido a la repetición de cargas de activación y la sobrecarga de sincronización.

2. Metodología

FairyFuse es un sistema de inferencia diseñado para ejecutar modelos LLM cuantizados a ternarios en CPUs comerciales (x86) sin realizar ninguna multiplicación en punto flotante.

A. Representación y Decodificación de Pesos

• Empaquetado: Se empaquetan 16 pesos ternarios en una sola palabra de 32 bits (2 bits por peso).
  • (1, 0) = +1 (Sumar).
  • (0, 1) = -1 (Restar).
  • (0, 0) = 0 (No-op).
• Decodificación SIMD: Utilizando la instrucción de arquitectura BMI2 (_pext_u32), se extraen máscaras de bits para los valores positivos y negativos en un solo ciclo de reloj. Esto evita la descuantización a punto flotante.

B. Kernel Fused 8-GEMV (Núcleo de la Innovación)

Para manejar la estructura de los modelos complejos de Fairy2i, FairyFuse fusiona las ocho sub-operaciones GEMV en un único bucle paralelo AVX-512. Esto se logra mediante cuatro decisiones de diseño clave:

1. Reutilización de Máscaras (O1): Cada par de máscaras decodificadas se aplica tanto al camino de activación real como al imaginario, amortizando el costo de decodificación.
2. Reutilización de Vectores de Entrada (O2): Las activaciones de entrada ($x_{re}$ y $x_{im}$) se cargan una sola vez en registros AVX-512 y se comparten entre las ocho sub-operaciones, reduciendo el tráfico de memoria a la mitad.
3. Truco de Cambio de Signo (O3): En lugar de negar activaciones en tiempo de ejecución para los términos conjugados, se precalcula un alias con signo invertido ($neg\_x\_im$), eliminando operaciones de negación por canal.
4. Acumulación en Registro (O4): Todas las sumas parciales se mantienen en registros AVX-512 durante todo el bucle interno, evitando el desbordamiento (spilling) a la memoria DRAM hasta la reducción final.

C. Ejecución

El bucle interno utiliza exclusivamente instrucciones de suma y resta enmascaradas (_mm512_mask_add_ps, _mm512_mask_sub_ps) y extracción de bits. No hay instrucciones de multiplicación (vmulps) ni tablas de búsqueda (LUT) en el camino crítico.

3. Contribuciones Clave

1. Primer Kernel GEMV Ternario sin Multiplicación en x86: Implementación que utiliza únicamente sumas/restas enmascaradas AVX-512 y decodificación BMI2, verificada mediante inspección de ensamblador (cero multiplicaciones en el bucle interno).
2. Fusión de Capas Ampliamente Lineales: Consolidación de las ocho sub-GEMV complejas en un único bucle SIMD, logrando un 1.55x de aceleración sobre una ejecución no fusionada al eliminar barreras de sincronización y tráfico de memoria redundante.
3. Validación de la Arquitectura CPU para Cuantización Extrema: Demostración de que, gracias a la compresión de datos, las CPUs (con ancho de banda limitado) se benefician más de la cuantización ternaria que las GPUs (con ancho de banda masivo), invirtiendo la lógica tradicional de despliegue.

4. Resultados

Las pruebas se realizaron en un servidor con un procesador Intel Xeon 8558P (48 núcleos, ~200 GB/s de ancho de banda) y un modelo LLaMA-2-7B.

• Rendimiento del Kernel (GEMV):
  • En un solo hilo, se logra una aceleración de 2.0x a 6.6x frente a FP32.
  • Con 48 hilos, la aceleración alcanza 29.6x a 54.4x frente a FP32, debido a que el conjunto de trabajo empaquetado cabe en la caché L3, reduciendo drásticamente el tráfico de DRAM.
• Rendimiento de Extremo a Extremo (Throughput):
  • 32.4 tokens/segundo en un solo socket.
  • Supera a llama.cpp en formato Q4_K_M en un 1.24x (26.15 tok/s) y a Q2_K en un 1.61x, a pesar de usar un modelo de 2 bits (3.3 GB) en lugar de 4 bits.
• Calidad del Modelo:
  • Perplejidad en WikiText-2: 5.52 (vs 5.47 de FP16 y 5.68 de Q4_K_M). La pérdida es casi nula (0.05 de diferencia con FP16).
  • Precisión Promedio en Tareas Posteriores: 66.0% (vs 67.3% de FP16 y 65.1% de Q4_K_M).
• Análisis Roofline:
  • La compresión 16x aumenta la intensidad aritmética de 0.25 a 8.0 OP/byte.
  • En CPU, esto mueve la operación desde el régimen limitado por memoria hacia el límite de cómputo, habilitando la gran aceleración.
  • En GPU (H200), el ancho de banda masivo hace que la compresión sea menos efectiva, y la falta de instrucciones eficientes de extracción de bits (como pext) hace que la implementación ternaria sea 130x más lenta que FP16 en GPU.

5. Significado e Impacto

FairyFuse demuestra que es posible desplegar modelos LLM de miles de millones de parámetros en hardware de consumo (solo CPU) con latencia interactiva y calidad casi sin pérdidas, eliminando por completo la necesidad de multiplicaciones en punto flotante.

• Cambio de Paradigma: Desafía la noción de que las GPUs son siempre superiores para la inferencia de LLM. Para modelos cuantizados extremadamente (ternarios/binares), las CPUs con extensiones de instrucciones específicas (AVX-512, BMI2) son el objetivo natural debido a la eficiencia en el manejo de ancho de banda limitado.
• Privacidad y Despliegue: Facilita la ejecución de modelos potentes en entornos donde las GPUs no están disponibles o no son viables por razones de privacidad, energía o soberanía de datos, manteniendo la inferencia totalmente offline y en el host.
• Eficiencia Energética: Al eliminar las unidades de multiplicación y reducir el movimiento de datos, se espera una mejora significativa en la eficiencia energética, crucial para dispositivos de borde.

En resumen, FairyFuse cierra la brecha entre la teoría de los pesos ternarios y la práctica de la inferencia de alto rendimiento en hardware estándar, ofreciendo una ruta viable para la inferencia de LLMs en la "nube de borde" y dispositivos locales.