Tres optimizaciones para una inferencia MoE más rápida en NVIDIA DGX Spark

Los ingenieros que optimizan modelos de lenguaje grandes a menudo se topan con un muro donde añadir más potencia de cálculo no acelera la generación. Una reciente evaluación técnica del modelo Ling-3.0-flash en hardware NVIDIA DGX Spark identificó que el rendimiento de la decodificación estaba limitado por el ancho de banda de memoria, y no por la potencia bruta de cálculo. Al atacar cuellos de botella específicos en la lectura de pesos y la verificación de tokens, los investigadores lograron reducciones significativas de latencia sin sacrificar la precisión.

Qué ocurrió

La evaluación se centró en Ling-3.0-flash, un modelo de mezcla de expertos (MoE) de Inclusion AI con 124 mil millones de parámetros totales, pero solo 5,1 mil millones de parámetros activos por token. Al ejecutarlo con un tamaño de lote de uno, las herramientas de perfilado mostraron que el tiempo dedicado a generar tokens estaba dominado por dos operaciones: el propio cómputo MoE y la proyección de la capa de salida, conocida como LM head. Estos pasos requerían leer grandes cantidades de datos de pesos desde la memoria, creando un cuello de botella que las técnicas de optimización estándar no conseguían resolver por completo.

Para resolver esto, el equipo probó tres optimizaciones distintas. Primero, sustituyeron el backend FlashInfer MXFP4 MoE predeterminado por la librería de kernels Humming. Segundo, aplicaron cuantización FP8 online específicamente al LM head para reducir su huella de memoria. Tercero, implementaron decodificación especulativa de predicción multi-token utilizando la capa NEXTN integrada del modelo. Cada cambio apuntaba a un aspecto diferente del pipeline de inferencia, desde la eficiencia de ejecución del kernel hasta el volumen de datos y la reducción de pasos.

Los resultados se midieron tanto en cargas de trabajo sintéticas como en el conjunto de pruebas completo GSM8K. Mientras que las pruebas sintéticas mostraban enormes ganancias debido a patrones de salida repetitivos, la evaluación en el mundo real sobre 1.319 preguntas de GSM8K demostró un aumento del 42,5 % en el rendimiento de decodificación y una reducción del 24,5 % en la latencia media. Crucialmente, estas mejoras vinieron acompañadas de ningún cambio estadísticamente detectable en la precisión del modelo, manteniendo una puntuación del 96,82 % frente al 96,89 % de la línea base.

Cómo funciona

La primera optimización aprovecha Humming, una librería de kernels independiente que utiliza la Compilación Runtime CUDA de NVIDIA para generar código adaptado a las dimensiones específicas del modelo. A diferencia de la ruta FlashInfer de la línea base, que compila muchos objetivos genéricos, Humming crea binarios especializados para las formas de matriz reales utilizadas por Ling-3.0-flash. Esta especialización reduce el tiempo por token de salida asegurando que la GPU ejecute solo las instrucciones necesarias para el tamaño intermedio estrecho del modelo.

La segunda técnica aborda el LM head, que proyecta estados ocultos a un vocabulario de más de 157.000 tokens. En la configuración de la línea base, esta capa usa precisión BF16, lo que requiere un ancho de banda de memoria significativo para cada token generado. Al cambiar a cuantización FP8 online solo para esta capa, el sistema reduce a la mitad el volumen de datos de pesos leídos durante cada paso. Este cambio es mutuamente exclusivo con la opción de LM head FP32 de la línea base y requiere una configuración cuidadosa para evitar problemas de compatibilidad con otras partes del modelo.

El tercer método utiliza decodificación especulativa para generar múltiples tokens en una sola ronda de verificación. Ling-3.0-flash incluye una capa NEXTN nativa que actúa como modelo borrador, prediciendo tokens subsiguientes que el modelo principal luego verifica. En lugar de generar un token a la vez, el sistema redacta tres tokens y verifica cuatro posiciones simultáneamente. Este enfoque reduce el número total de pasos de decodificación necesarios, aunque aumenta la complejidad de cada paso de verificación debido a mayores recuentos de activación de expertos.

Detalles clave

• Ling-3.0-flash tiene 124 mil millones de parámetros totales con 5,1 mil millones de parámetros activos por token a través de 42 capas.
• Sustituir el backend FlashInfer por Humming redujo el tiempo por token de salida de 27,71 ms a 25,35 ms, una ganancia del 9,3 %.
• La cuantización FP8 online del LM head redujo a la mitad el volumen de lectura de pesos para esa capa, contribuyendo a una mejora acumulada del rendimiento del 21,7 %.
• La predicción multi-token con NEXTN logró 3,90 tokens aceptados por ronda en datos sintéticos, impulsando el rendimiento un 185,1 % en esa prueba específica.
• En el benchmark GSM8K, las optimizaciones combinadas aumentaron el rendimiento de decodificación un 42,5 % y redujeron la latencia media un 24,5 %.
• La precisión permaneció estable, cayendo solo 0,08 puntos porcentuales del 96,89 % al 96,82 %, lo cual no es estadísticamente significativo.

Por qué importa

Para los desarrolladores que construyen flujos de trabajo agénticos o aplicaciones en tiempo real, la latencia suele ser la restricción principal. Este análisis confirma que el ancho de banda de memoria, y no solo la potencia de cálculo, es el factor limitante para muchos modelos MoE en tamaños de lote bajos. Al enfocarse en reducir el tráfico de pesos y aumentar los tokens por paso, los ingenieros pueden lograr ganancias sustanciales de rendimiento sin necesitar clústeres de hardware más grandes. Esto es particularmente relevante para despliegues en el edge o entornos cloud sensibles a costos donde cada milisegundo cuenta.

Sin embargo, el estudio también destaca los riesgos de combinar optimizaciones sin un conocimiento profundo de la integración. Los investigadores encontraron que habilitar la cuantización FP8 del LM head y la predicción multi-token juntos requería que el modelo borrador compartiera todo el módulo LM-head. Sin esta interfaz específica, el sistema caía silenciosamente a una configuración más lenta, interpretando los pesos FP8 como BF16. Esto subraya la importancia de validar cada optimización al inicio en lugar de asumir compatibilidad basada en pruebas de componentes individuales.

Qué puedes hacer

• Perfilá tu pipeline actual de inferencia MoE para identificar si la decodificación está limitada por memoria o por cómputo antes de aplicar optimizaciones.
• Prueba la librería de kernels Humming como reemplazo directo de FlashInfer si tu modelo tiene dimensiones de matriz especializadas.
• Habilita la cuantización FP8 online para el LM head usando la variable de entorno SGLANG_ENABLE_FP8_LM_HEAD para reducir el uso de ancho de banda de memoria.
• Implementa la predicción multi-token usando capas borrador nativas como NEXTN para reducir el número de pasos secuenciales de decodificación.
• Verifica que los modelos borrador y objetivo compartan los módulos correctos cuando combines cuantización con decodificación especulativa para evitar degradaciones silenciosas de rendimiento.
• Haz benchmarks de los cambios en cargas de trabajo realistas como GSM8K en lugar de depender únicamente de pruebas sintéticas, que pueden sobreestimar los beneficios de la decodificación especulativa.