MoonMath AI lança kernel de atenção open-source para AMD MI300X que supera AITER v3 em todos os cenários

MoonMath AI lança kernel de atenção open-source para AMD MI300X que supera o AITER v3 em todos os cenários

Um time de pesquisadores do laboratório MoonMath AI acaba de liberar como código aberto um kernel de atenção em bf16 para a GPU AMD MI300X — e os números impressionam. Escrito em HIP (não em assembly manual), o kernel supera o AITER v3, a própria implementação otimizada da AMD, em todos os formatos e modos de arredondamento testados, com ganho médio geométrico de 1,18× e picos de até 1,26×.

O código está disponível sob licença MIT e já foi testado em ambiente real: um pull request no SGLang acelerou a difusão de vídeo do modelo Wan2.1 em 1,23×, sem qualquer regressão de qualidade.

Por que isso importa

O kernel de atenção — a operação softmax(QKᵀ/√d)·V — é o coração computacional de todo transformer. Ter uma implementação eficiente e aberta para GPUs AMD reduz a dependência do ecossistema NVIDIA e amplia as opções de hardware para treinamento e inferência de modelos de linguagem.

A MI300X é a GPU de data center da AMD baseada na arquitetura CDNA3 (ISA gfx942), e este kernel roda exclusivamente nesse hardware. O acesso bare-metal foi fornecido pela HotAisle, provedora de nuvem especializada em AMD.

O truque central: wrappers assembly de uma instrução

A grande sacada técnica está em como o time evita um dilema conhecido: usar intrínsecos do compilador mantém o código limpo, mas permite que o compilador reordene ou renomeie operandos; usar assembly inline puro dá controle total, mas exige gerenciamento manual de registradores.

A abordagem do MoonMath é envolver exatamente uma instrução em uma função __device__ __forceinline__, usando restrições extended asm para descrever os operandos:

__device__ __forceinline__ void asm_mfma(bf16x4_t a, bf16x4_t b, fp32x4_t& c) {
 asm volatile("v_mfma_f32_16x16x16_bf16 %0, %1, %2, %0"
  : "+v"(c) : "v"(a), "v"(b));
}

A restrição "+v"(c) vincula a entrada e saída do acumulador ao mesmo registrador VGPR — sem gerar instruções de cópia desnecessárias.

Arquitetura: oito waves, dois grupos, duas barreiras

Cada unidade computacional CDNA3 tem quatro unidades SIMD. O bloco tradicional usaria quatro waves. O MoonMath escala para oito waves por bloco, divididas em dois grupos de quatro, operando em fases alternadas:

• Enquanto um grupo satura o núcleo de matriz com Q*K, o outro executa softmax e emite loads
• Depois trocam — o núcleo de matriz nunca fica ocioso

Duas s_barrier por iteração gerenciam a troca de fase e o limite de iteração. A abordagem ecoa a alternância matmul/softmax do FlashAttention-3, mas sem copiar a divisão produtor/consumidor do FA3 (na CDNA3, todo movimento de memória já é assíncrono).

Onde os dados vivem — e por que 16×16×16

A maior parte do ganho de velocidade vem do posicionamento inteligente dos dados na hierarquia de memória:

O time escolheu o tile 16×16×16 (MFMA) em vez de 32×32×8. Ambos têm a mesma vazão teórica, mas o tile menor acumula em 4 elementos fp32 por lane (contra 16), liberando pressão dos registradores para prefetch mais profundo e um terceiro tile de Q.

Duas otimizações adicionais fecham a conta: um terceiro tile Q (3Q) aumenta o reuso de dados por tile K/V carregado, e uma divisão de cauda KV no estilo Flash-Decoding resolve o problema de unidades computacionais fracionárias nas 304 CUs da MI300X.

Limitações

O kernel é focado e não pretende cobrir todos os casos de uso. Não há suporte para:

• Causal mask
• Grouped Query Attention (GQA)
• Varlen batching
• Hardware que não seja gfx942

As saídas são sempre em bf16, e a dimensão de cabeça é fixa em 128.

O que isso significa para o ecossistema

Este lançamento é mais um sinal de que o ecossistema AMD para IA está amadurecendo rapidamente. Kernels de atenção eficientes e abertos eram uma das principais barreiras para adoção de GPUs AMD em cargas de transformer. Com o código disponível sob MIT, a comunidade pode integrar, adaptar e melhorar — e o teste no SGLang já mostra resultados práticos.

O repositório está disponível no GitHub do MoonMath AI. Para quem trabalha com inferência ou treinamento na MI300X, é uma adição que merece atenção imediata.

---

Fonte: MarkTechPost