Resumen Técnico: ¡Busca tus Escalas de Punto Flotante por Bloques!

Declaración del Problema

La cuantización es una técnica estándar para acelerar la inferencia de modelos generativos, pero mantener la precisión en anchos de bits ultra-bajos sigue siendo un desafío. Si bien los formatos de Punto Flotante por Bloques (BFP) ofrecen un equilibrio entre rango dinámico, eficiencia de memoria y rendimiento, los algoritmos BFP estándar dependen de una escala fija determinada por la magnitud absoluta máxima dentro de un bloque. Los autores observan que esta heurística de "escalado por máximo" a menudo es subóptima, lo que conduce a errores de cuantización innecesarios. Este problema es particularmente relevante con la llegada de la arquitectura Blackwell de NVIDIA, que introduce soporte de hardware para formatos BFP de microescala como NVFP4. A diferencia del BFP tradicional, NVFP4 incluye bits de mantisa dentro de su escala de bloque de 8 bits (formato E4M3), proporcionando una granularidad más fina para la selección de escala que el escalado máximo estándar no logra aprovechar.

Metodología

ScaleSearch

La contribución central es ScaleSearch, un algoritmo que aprovecha los bits de mantisa en formatos de microescala para realizar una búsqueda de granularidad fina de la escala de bloque óptima.

• Mecanismo: En lugar de seleccionar rígida la escala basándose en el valor máximo del bloque ($s = \text{round}(\|x\|_\infty / 6.0)$), ScaleSearch explora un vecindario de escalas alrededor de este valor predeterminado. Trata la escala como un desplazamiento entero $f$ aplicado a la representación en bits de la escala predeterminada: $s(f) = \text{reinterpret}(\text{reinterpret}(s, \text{int8}) + f, \text{fp8}_{UE4M3})$.
• Estrategia de Búsqueda: El algoritmo busca exhaustivamente un pequeño rango de desplazamientos (por ejemplo, $f \in [-2, 6]$) para encontrar la escala que minimiza el Error Cuadrático Medio (MSE) para la distribución del bloque dada.
• Fundamento: El análisis empírico en datos gaussianos sintéticos y en activaciones reales de redes neuronales (por ejemplo, estados Key de Llama 3.1) revela una distribución bimodal de desplazamientos óptimos. Mientras que la escala predeterminada (desplazamiento 0) mapea la mayor magnitud al valor FP4 6, una porción significativa de bloques se representa mejor mapeando la mayor magnitud al valor FP4 4, lo que requiere un factor de escala aproximadamente 1.5 veces mayor. ScaleSearch identifica estos mapeos óptimos.

ScaleSearchAttention

Para extender estos beneficios al mecanismo de atención y a la memoria caché KV, los autores proponen ScaleSearchAttention. Este es un algoritmo de atención FP4 de extremo a extremo diseñado para el modelado de lenguaje causal que evita sobrecargas de descuantización.

• Ejecución Nativa NVFP4: Realiza multiplicaciones de matrices QKT y PV directamente en los Tensor Cores utilizando precisión NVFP4 con acumulación FP32.
• Procesamiento de Incoherencia (IP) y Reducción de Magnitud: Para mitigar valores atípicos y reducir el error de cuantización en los tensores Q y K, el método emplea transformadas de Hadamard (siguiendo QuIP#) y un par adicional de transformaciones lineales ($Q' = QR^{-T}, K' = KR$) que preservan las puntuaciones de atención mientras reducen la magnitud cuadrática promedio de los tensores proyectados.
• Almacenamiento en Caché de Precisión Mixta Consciente del Sumidero de Atención: Reconociendo que las puntuaciones de atención se concentran en los tokens iniciales y recientes (Sumidero de Atención), el método almacena los primeros y últimos bloques de la memoria caché KV en precisión completa, mientras que los bloques intermedios se almacenan en NVFP4 comprimido. Esto asegura alta precisión para tokens críticos sin escalar el uso de memoria con la longitud de la secuencia.

Contribuciones Clave

1. Algoritmo ScaleSearch: Una técnica de búsqueda de escala para la cuantización de Punto Flotante por Bloques que minimiza el error de cuantización a nivel de bloque explorando escalas representables vecinas. Aprovecha específicamente la resolución de mantisa del formato de escala NVFP4 E4M3.
2. ScaleSearchAttention: Un mecanismo de atención FP4 co-optimizado que permite la inferencia de extremo a extremo con una degradación de precisión cercana a cero. Integra ScaleSearch, procesamiento de incoherencia y almacenamiento en caché KV de precisión mixta para realizar operaciones QKT y PV directamente en NVFP4 en los Tensor Cores.
3. Validación Empírica: Experimentos extensos que demuestran ganancias consistentes en perplejidad y precisión en pruebas de referencia sobre líneas base FP4 anteriores (como SageAttention3 y Naive-FP4) en modelos de lenguaje a gran escala y modelos de difusión.

Resultados Experimentales

Cuantización Post-Entrenamiento (PTQ)

• Modelos de Lenguaje: En el modelo Qwen3-8B, ScaleSearch mejoró el rendimiento hasta en 15 puntos en la prueba de referencia MATH-500 en comparación con la cuantización NVFP4 estándar. También cerró significativamente la brecha en las pruebas de referencia GPQA y MMLU.
• Modelos de Difusión: Al aplicarse a los modelos Mochi y CogVideoX utilizando SageAttention3, ScaleSearch mejoró las métricas VQA-a, VQA-t y FScore, al tiempo que igualó las puntuaciones de alineación de precisión completa (CLIPSIM, CLIP-T).

Atención y Perplejidad

• Perplejidad (Wikitext-2): ScaleSearchAttention redujo la perplejidad hasta en un 22% para Llama 3.1 70B en comparación con la cuantización FP4 estándar. Para Llama 3.1 70B, logró una PPL de 2.6348, una mejora significativa sobre SageAttention3 (3.3899) y Naive-FP4 (3.4000).
• GPQA Diamond: En el modelo Llama 3.1 8B Instruct, ScaleSearchAttention logró una precisión del 32.32%, superando a SageAttention3 (26.26%) e igualando la línea base de precisión completa (31.81%).

Eficiencia y Sobrecarga

• Sobrecarga de Cuantización: ScaleSearch introduce una sobrecarga mínima de 1.74x durante la fase de cuantización de FP32 a NVFP4 (0.0449 ms frente a 0.0258 ms de línea base) para un rango de búsqueda restringido.
• Rendimiento: En los núcleos de atención, ScaleSearch alcanza el 98.3% del rendimiento de la línea base (SageAttention3) para atención no causal con longitudes de secuencia de 32K, demostrando un impacto insignificante en el tiempo de ejecución.
• Latencia: La latencia de extremo a extremo para la generación de texto a video aumentó solo marginalmente (por ejemplo, de 353.40s a 364.68s para Mochi), confirmando la practicidad del método para cargas de trabajo del mundo real.

Significado

El artículo afirma que ScaleSearch aborda una suboptimalidad fundamental en las estrategias de escalado BFP estándar al utilizar los bits de mantisa previamente subutilizados en formatos de microescala. Al buscar sistemáticamente la escala óptima en lugar de depender de un máximo heurístico, el método reduce significativamente el error de cuantización. La introducción de ScaleSearchAttention demuestra que estas ganancias pueden extenderse al mecanismo de atención intensivo en memoria y cómputo, permitiendo inferencia de alta velocidad y baja precisión con una pérdida de precisión cercana a cero. Los autores posicionan este trabajo como un paso hacia el aprovechamiento completo de las capacidades emergentes de hardware de 4 bits (como NVIDIA Blackwell) para una IA generativa eficiente.