COREY: A Prototype Study of Entropy-Guided Operator Fusion with Hadamard Reparameterization for Selective State Space Models

Resumen Técnico: COREY

1. El Problema

Los Modelos de Espacio de Estados (SSM), representados por la familia Mamba, ofrecen una complejidad lineal $O(L)$ para el modelado de secuencias, lo que los hace ideales para contextos largos y cargas de trabajo de streaming. Sin embargo, su despliegue práctico en hardware (como GPUs) enfrenta dos limitaciones críticas:

• Cuello de botella de ancho de banda de memoria: Las actualizaciones de estado selectivo a menudo se descomponen en múltiples kernels fragmentados (punto a punto y reducción), lo que genera una materialización repetida de tensores intermedios en la memoria global. Esto convierte la ruta crítica en un problema limitado por memoria, no por cómputo.
• Inestabilidad en cuantización y fusión: Las activaciones internas de los SSM suelen tener colas pesadas (heavy-tailed), con valores atípicos (outliers) que desestabilizan la cuantización de baja precisión (ej. 4 o 8 bits) y reducen la profundidad segura de fusión de kernels debido a la sensibilidad numérica amplificada.

2. Metodología

El artículo presenta COREY, un marco de trabajo prototipo que unifica la optimización de sistemas con la reestructuración matemática de las activaciones. Se basa en dos pilares principales:

A. Fusión de Operadores Guiada por Entropía
En lugar de usar límites de fusión estáticos o heurísticos fijos, COREY utiliza la entropía de las activaciones como una señal en tiempo real para decidir los límites de fusión y el tamaño de las baldosas (tiles).

• Estimación de Entropía: Se calcula la entropía de Shannon utilizando histogramas de ancho fijo sobre las activaciones.
• Estrategia de Fusión: Se define una puntuación de fusión $S(R)$ que combina la entropía normalizada, la intensidad aritmética y el tráfico de memoria estimado.
  • Si la entropía es alta, las activaciones están bien distribuidas, permitiendo baldosas más grandes y una fusión más agresiva.
  • Si la entropía es baja (zonas con outliers), se reducen los tamaños de las baldosas para evitar la presión de registros y la inestabilidad.
• Algoritmo: Se utiliza un algoritmo de selección de límites que evalúa si fusionar un operador adicional mejora la puntuación sin violar las restricciones de memoria compartida y registros.

B. Reparametrización Hadamard Fundida
Para abordar los valores atípicos en las activaciones, el método integra transformadas de Hadamard dentro de las capas de proyección lineal.

• Mecanismo: Se inserta una transformación $H$ (matriz de Hadamard normalizada) y su inversa en la proyección lineal $y = Wx$, reescribiéndola como $y = W'Hx$.
• Efecto: La transformación Hadamard redistribuye la energía de las colas pesadas a través de los canales, suavizando la distribución de activaciones.
• Beneficio Sistémico: Esto aumenta la entropía empírica de las activaciones, lo que a su vez expande la "zona de operación segura" para permitir una fusión de kernels más profunda y una ejecución en baja precisión sin desestabilizar el modelo.
• Eficiencia: La transformación se fusiona en el mismo kernel que la proyección y los operadores elementales adyacentes, evitando pasadas adicionales de memoria.

3. Contribuciones Clave

1. Criterio de Fusión Dinámico: Introducción de un criterio impulsado por la entropía para seleccionar dinámicamente los límites de fusión en los kernels de SSM, conectando la geometría de las activaciones con la viabilidad de la fusión.
2. Capa Hadamard Fundida: Diseño de una capa que incorpora transformaciones hacia adelante/inverso en bloques de proyección, preservando la equivalencia funcional exacta mientras reduce la concentración de coordenadas pico.
3. Vinculación Teórica-Empírica: Demostración teórica de que la mezcla Hadamard (bajo ciertas condiciones de matrices doblemente estocásticas) aumenta la entropía del histograma y reduce la magnitud de los picos de coordenadas, lo cual es crucial para la ejecución limitada por recorte (clipping) en baja precisión.
4. Validación de Prototipo Controlado: Un estudio exhaustivo que valida el diseño mediante un prototipo reproducible, separando claramente las métricas de rendimiento del sistema de las afirmaciones de calidad de nivel de punto de control (checkpoint).

4. Resultados

El estudio se realizó en un entorno de prototipo controlado (sin kernels Triton finales, simulando la planificación) sobre activaciones sintéticas de colas pesadas y validado con checkpoints reales de Mamba en un entorno WSL2.

• Rendimiento y Latencia:
  • COREY reduce consistentemente la latencia y mejora el rendimiento (tokens/segundo) en comparación con las líneas base sin fusión y con fusión estática.
  • Las mejoras son más pronunciadas en contextos largos (hasta 64k tokens), donde la reducción del tráfico de DRAM es crítica.
  • En la categoría "Ultra-long" (64k), COREY logró una latencia de 77.97 ms frente a 104.54 ms (No Fusion) y 87.22 ms (Fusión Estática), con un aumento de rendimiento de 840.5 k tok/s.
• Tráfico de Memoria:
  • Se observó una reducción significativa en los bytes de DRAM por token, confirmando que la mejora proviene de la reducción de la materialización de memoria intermedia.
• Estabilidad en Cuantización (W8A8 / W4A8):
  • La reparametrización Hadamard redujo la degradación del "proxy de calidad" en modos de baja precisión.
  • En W4A8, el proxy de degradación para COREY fue 0.0871, significativamente mejor que la fusión estática (0.2522), indicando una mejor estabilidad numérica.
• Validación de Checkpoint:
  • Se validó la reproducibilidad cruzada en hardware (RTX 3070 y RTX 3090) utilizando modelos Mamba-370M y Mamba-1.4B en tareas de LongBench, confirmando que la calidad del modelo se mantiene independientemente de la plataforma.

5. Significado e Impacto

• Puente entre Teoría y Sistemas: El trabajo demuestra que las propiedades estadísticas de las activaciones (entropía) pueden utilizarse directamente como señales de planificación para optimizar el hardware, cerrando la brecha entre la teoría de la distribución de datos y la ingeniería de kernels.
• Habilitador de Baja Precisión: Al suavizar las activaciones mediante Hadamard, se hace viable la ejecución de SSMs en precisión mixta o baja (4/8 bits) sin sacrificar la estabilidad, lo cual es esencial para el despliegue eficiente en dispositivos con recursos limitados.
• Enfoque Práctico: Aunque el estudio actual es un prototipo y no una implementación final de kernels de alto rendimiento, establece un marco metodológico sólido para futuros co-diseños de aceleradores que puedan exponer primitivas nativas de Hadamard y planificación de memoria sensible a la entropía.

Conclusión:
COREY propone que la gestión inteligente de la fusión de operadores, guiada por la entropía de las activaciones y facilitada por la reestructuración matemática (Hadamard), es una vía prometedora para superar las limitaciones de ancho de banda de memoria en los SSMs modernos, mejorando tanto la eficiencia energética como la viabilidad de la inferencia en baja precisión.