AWS apresenta otimizações de kernel para modelos Mixture-of-Experts no Amazon SageMaker

A Amazon Web Services (AWS) anunciou novas otimizações de kernel para modelos Mixture-of-Experts (MoE) no Amazon SageMaker. Essas melhorias visam aumentar a eficiência operacional e reduzir a latência em tarefas de processamento de linguagem natural, especialmente em GPUs H200. As otimizações incluem a utilização do Swap-AB para maximizar o throughput do WGMMA e o roteamento do backend Triton para quantização FP8, além da eliminação da recompilação JIT do Triton.

Os modelos MoE são arquiteturas que ativam apenas uma parte dos especialistas em cada operação, permitindo que a capacidade dos grandes modelos de linguagem (LLMs) seja escalada sem aumentar significativamente o custo computacional por token. A AWS destacou que a otimização das operações MoE por meio de engenharia de kernel pode resultar em ganhos significativos de desempenho, com aumentos de até 10% no throughput e até 9% na velocidade do tempo até o primeiro token.

A técnica de MoE fundido, que combina o roteamento de tokens, a troca de dados e cálculos de rede especializados em uma única operação em lote na GPU, é central para essas otimizações. Embora essa abordagem ofereça vantagens em termos de latência, também introduz complexidade na implementação e na depuração dos kernels. Portanto, a utilização eficiente da arquitetura de GPU requer otimizações específicas de hardware.

Essas inovações podem impactar desenvolvedores e empresas que utilizam o Amazon SageMaker para construir e treinar modelos de linguagem, oferecendo uma maneira mais eficiente de lidar com grandes volumes de dados e melhorar a experiência do usuário final. "As otimizações de kernel são uma oportunidade significativa para ganhos de desempenho", afirmaram os engenheiros da AWS.