Engenharia de kernels Mixture-of-Experts na Amazon SageMaker AI
aws ·
Este artigo detalha otimizações de engenharia de kernels para modelos Mixture-of-Experts (MoE) no Amazon SageMaker AI, com foco nos ganhos de desempenho na inferência de grandes modelos de linguagem. Os autores descrevem três técnicas principais: transposição Swap-AB para alinhar as dimensões das matrizes aos requisitos de hardware WGMMA da NVIDIA Hopper; ajustes no roteamento do backend para utilizar o Triton em vez do CUTLASS para quantização FP8; e a eliminação de recompilações JIT desnecessárias do Triton para reduzir a latência. Essas otimizações abordam gargalos específicos nas arquiteturas MoE, como lotes finos durante a decodificação e contagens variáveis de tokens. A implementação dessas mudanças resultou em até 10% maior throughput e uma melhoria de 9% no Tempo para o Primeiro Token (TTFT) em GPUs H200. Vários desses aprimoramentos foram incorporados à comunidade de código aberto vLLM. O trabalho demonstra como o ajuste fino específico ao hardware pode melhorar significativamente a eficiência em modelos de ativação esparsa.
aws