Ingeniería de kernels Mixture-of-Experts en Amazon SageMaker AI
aws ·
Este artículo detalla las optimizaciones de ingeniería de kernels para modelos de Mezcla de Expertos (MoE) en Amazon SageMaker AI, centrándose en las mejoras de rendimiento para la inferencia de modelos de lenguaje grandes. Los autores describen tres técnicas clave: la transposición Swap-AB para alinear las dimensiones de las matrices con los requisitos de hardware WGMMA de NVIDIA Hopper; ajustes en el enrutamiento del backend para utilizar Triton en lugar de CUTLASS para la cuantización FP8; y la eliminación de recompilaciones JIT innecesarias de Triton para reducir la latencia. Estas optimizaciones abordan cuellos de botella específicos en las arquitecturas MoE, como los tamaños de lote reducidos durante la decodificación y los conteos variables de tokens. La implementación de estos cambios resultó en un aumento de hasta el 10% en el rendimiento y una mejora del 9% en el Tiempo hasta el Primer Token (TTFT) en GPUs H200. Varias de estas mejoras se han incorporado a la comunidad de código abierto vLLM. El trabajo demuestra cómo la sintonización específica del hardware puede mejorar significativamente la eficiencia en los modelos de activación dispersa.
aws